diff --git a/src/tree_schema_helpers.c b/src/tree_schema_helpers.c
index 7f82482..a3b201d 100644
--- a/src/tree_schema_helpers.c
+++ b/src/tree_schema_helpers.c
@@ -1289,6 +1289,9 @@
     struct lys_module *context_mod = (struct lys_module*)private;
     unsigned int u;
 
+    if (context_mod == mod) {
+        return context_mod->prefix;
+    }
     LY_ARRAY_FOR(context_mod->compiled->imports, u) {
         if (context_mod->compiled->imports[u].module == mod) {
             /* match */