diff --git a/src/tree_schema.c b/src/tree_schema.c
index 07a661d..655e877 100644
--- a/src/tree_schema.c
+++ b/src/tree_schema.c
@@ -40,6 +40,7 @@
 
     LY_CHECK_ARG_RET(NULL, parent || module, NULL);
 
+next:
     if (!last) {
         /* first call */
 
@@ -47,17 +48,15 @@
         if (parent) {
             /* schema subtree */
             if (parent->nodetype == LYS_CHOICE && (options & LYS_GETNEXT_WITHCASE)) {
-                if (!((struct lysc_node_choice*)parent)->cases) {
-                    return NULL;
+                if (((struct lysc_node_choice*)parent)->cases) {
+                    next = last = (const struct lysc_node*)&((struct lysc_node_choice*)parent)->cases[0];
                 }
-                next = last = (const struct lysc_node*)&((struct lysc_node_choice*)parent)->cases[0];
             } else {
                 snode = lysc_node_children_p(parent, (options & LYS_GETNEXT_OUTPUT) ? LYS_CONFIG_R : LYS_CONFIG_W);
                 /* do not return anything if the node does not have any children */
-                if (!snode || !(*snode)) {
-                    return NULL;
+                if (snode && *snode) {
+                    next = last = *snode;
                 }
-                next = last = *snode;
             }
         } else {
             /* top level data */
@@ -115,7 +114,7 @@
         /* possibly go back to parent */
         if (last && last->parent != parent) {
             last = last->parent;
-            next = last->next;
+            goto next;
         } else if (!action_flag) {
             action_flag = 1;
             next = parent ? (struct lysc_node*)lysc_node_actions(parent) : (struct lysc_node*)module->rpcs;
diff --git a/tests/src/test_tree_schema.c b/tests/src/test_tree_schema.c
index 48fdfc1..7dd29c1 100644
--- a/tests/src/test_tree_schema.c
+++ b/tests/src/test_tree_schema.c
@@ -196,6 +196,15 @@
     assert_non_null(node = lys_getnext(NULL, NULL, mod->compiled, LYS_GETNEXT_NOSTATECHECK));
     assert_string_equal("a", node->name);
 
+    assert_non_null(mod = lys_parse_mem(ctx, "module c {namespace urn:c;prefix c; rpc c;}", LYS_IN_YANG));
+    assert_non_null(node = lys_getnext(NULL, NULL, mod->compiled, 0));
+    assert_string_equal("c", node->name);
+    assert_null(node = lys_getnext(node, NULL, mod->compiled, LYS_GETNEXT_NOSTATECHECK));
+
+    assert_non_null(mod = lys_parse_mem(ctx, "module d {namespace urn:d;prefix d; notification d;}", LYS_IN_YANG));
+    assert_non_null(node = lys_getnext(NULL, NULL, mod->compiled, 0));
+    assert_string_equal("d", node->name);
+    assert_null(node = lys_getnext(node, NULL, mod->compiled, LYS_GETNEXT_NOSTATECHECK));
 
     *state = NULL;
     ly_ctx_destroy(ctx, NULL);