schema BUGFIX processing extension instance argument

Due to differences in YANG and YIN, the extension instance's argument is
processed differently. So far, the argument was missing when the source
format was YIN and extension was not compiled. This was breaking
printers and limiting even other use cases of the argument.

This patch solves the problem by getting the extension definition when
needed and finishing the argument processing not finished by YIN parser.
diff --git a/src/plugins_exts_yangdata.c b/src/plugins_exts_yangdata.c
index 75c4e57..ec4e8fc 100644
--- a/src/plugins_exts_yangdata.c
+++ b/src/plugins_exts_yangdata.c
@@ -58,13 +58,6 @@
                 p_ext->name, ly_stmt2str(c_ext->parent_stmt));
         return LY_ENOT;
     }
-    /* check mandatory argument */
-    if (!c_ext->argument) {
-        lyext_log(c_ext, LY_LLERR, LY_EVALID, cctx->path,
-                "Extension %s is instantiated without mandatory argument representing YANG data template name.",
-                p_ext->name);
-        return LY_EVALID;
-    }
 
     mod_c = (struct lysc_module *)c_ext->parent;
 
diff --git a/src/printer_yang.c b/src/printer_yang.c
index fbaa2e6..ccf321b 100644
--- a/src/printer_yang.c
+++ b/src/printer_yang.c
@@ -209,12 +209,13 @@
     LY_ARRAY_COUNT_TYPE u;
     struct lysp_stmt *stmt;
     ly_bool child_presence;
-    const char *argument;
 
     if (!count && ext) {
         count = LY_ARRAY_COUNT(ext);
     }
     LY_ARRAY_FOR(ext, u) {
+        struct lysp_ext *ext_def = NULL;
+
         if (!count) {
             break;
         }
@@ -224,16 +225,17 @@
             continue;
         }
 
-        ypr_open(ctx->out, flag);
-        argument = NULL;
-        if (ext[u].compiled) {
-            argument = ext[u].compiled->argument;
-        } else {
-            argument = ext[u].argument;
+        lysp_ext_find_definition(ctx->module->ctx, &ext[u], NULL, &ext_def);
+        if (!ext_def) {
+            continue;
         }
-        if (argument) {
+
+        ypr_open(ctx->out, flag);
+
+        if (ext_def->argument) {
             ly_print_(ctx->out, "%*s%s \"", INDENT, ext[u].name);
-            ypr_encode(ctx->out, argument, -1);
+            lysp_ext_instance_resolve_argument(ctx->module->ctx, &ext[u], ext_def);
+            ypr_encode(ctx->out, ext[u].argument, -1);
             ly_print_(ctx->out, "\"");
         } else {
             ly_print_(ctx->out, "%*s%s", INDENT, ext[u].name);
diff --git a/src/printer_yin.c b/src/printer_yin.c
index e1154d0..208e60a 100644
--- a/src/printer_yin.c
+++ b/src/printer_yin.c
@@ -1248,14 +1248,14 @@
 {
     LY_ARRAY_COUNT_TYPE u;
     struct lysp_stmt *stmt;
-    const char *argument;
-    const char *ext_argument;
     int8_t inner_flag = 0;
 
     if (!count && ext) {
         count = LY_ARRAY_COUNT(ext);
     }
     LY_ARRAY_FOR(ext, u) {
+        struct lysp_ext *ext_def = NULL;
+
         if (!count) {
             break;
         }
@@ -1265,21 +1265,38 @@
             continue;
         }
 
-        ypr_close_parent(ctx, flag);
-        inner_flag = 0;
-        argument = NULL;
-        ext_argument = NULL;
-
-        if (ext[u].compiled) {
-            argument = ext[u].compiled->argument;
-            ext_argument = ext[u].compiled->def->argument;
-        } else {
-            argument = ext[u].argument;
+        lysp_ext_find_definition(ctx->module->ctx, &ext[u], NULL, &ext_def);
+        if (!ext_def) {
+            continue;
         }
 
-        ypr_open(ctx, ext[u].name, ext_argument, argument, inner_flag);
+        ypr_close_parent(ctx, flag);
+        inner_flag = 0;
+
+        if (ext_def->argument) {
+            lysp_ext_instance_resolve_argument(ctx->module->ctx, &ext[u], ext_def);
+        }
+
+        ypr_open(ctx, ext[u].name, (ext_def->flags & LYS_YINELEM_TRUE) ? NULL : ext_def->argument, ext[u].argument, inner_flag);
         LEVEL++;
+        if (ext_def->flags & LYS_YINELEM_TRUE) {
+            const char *prefix, *name, *id;
+            size_t prefix_len, name_len;
+
+            ypr_close_parent(ctx, &inner_flag);
+
+            /* we need to use the same namespace as for the extension instance element */
+            id = ext[u].name;
+            ly_parse_nodeid(&id, &prefix, &prefix_len, &name, &name_len);
+            ly_print_(ctx->out, "%*s<%.*s:%s>", INDENT, prefix_len, prefix, ext_def->argument);
+            lyxml_dump_text(ctx->out, ext[u].argument, 0);
+            ly_print_(ctx->out, "</%.*s:%s>\n", prefix_len, prefix, ext_def->argument);
+        }
         LY_LIST_FOR(ext[u].child, stmt) {
+            if (stmt->flags & (LYS_YIN_ATTR | LYS_YIN_ARGUMENT)) {
+                continue;
+            }
+
             ypr_close_parent(ctx, &inner_flag);
             yprp_stmt(ctx, stmt);
         }
diff --git a/src/schema_compile.c b/src/schema_compile.c
index 5db58e0..39ba42b 100644
--- a/src/schema_compile.c
+++ b/src/schema_compile.c
@@ -92,13 +92,8 @@
 lys_compile_ext(struct lysc_ctx *ctx, struct lysp_ext_instance *ext_p, struct lysc_ext_instance *ext, void *parent,
         const struct lys_module *ext_mod)
 {
-    LY_ERR r, ret = LY_SUCCESS;
-    const char *tmp, *name, *prefix;
-    size_t pref_len, name_len;
-    LY_ARRAY_COUNT_TYPE v;
-
-    DUP_STRING(ctx->ctx, ext_p->argument, ext->argument, ret);
-    LY_CHECK_RET(ret);
+    LY_ERR ret = LY_SUCCESS;
+    struct lysp_ext *ext_def;
 
     ext->parent_stmt = ext_p->parent_stmt;
     ext->parent_stmt_index = ext_p->parent_stmt_index;
@@ -108,79 +103,15 @@
     lysc_update_path(ctx, LY_STMT_IS_NODE(ext->parent_stmt) ? ((struct lysc_node *)ext->parent)->module : NULL, "{extension}");
     lysc_update_path(ctx, NULL, ext_p->name);
 
-    /* parse the prefix */
-    tmp = ext_p->name;
-    r = ly_parse_nodeid(&tmp, &prefix, &pref_len, &name, &name_len);
-    /* it was parsed already */
-    assert(!r && !tmp[0]);
+    LY_CHECK_GOTO(ret = lysp_ext_find_definition(ctx->ctx, ext_p, &ext_mod, &ext_def), cleanup);
+    LY_CHECK_GOTO(ret = lys_compile_extension(ctx, ext_mod, ext_def, &ext->def), cleanup);
 
-    /* get module where the extension definition should be placed */
-    if (!ext_mod) {
-        ext_mod = ly_resolve_prefix(ctx->ctx, prefix, pref_len, ext_p->format, ext_p->prefix_data);
-        if (!ext_mod) {
-            LOGVAL(ctx->ctx, LYVE_REFERENCE, "Invalid prefix \"%.*s\" used for extension instance identifier.",
-                    pref_len, prefix);
-            ret = LY_EVALID;
-            goto cleanup;
-        } else if (!ext_mod->parsed->extensions) {
-            LOGVAL(ctx->ctx, LYVE_REFERENCE,
-                    "Extension instance \"%s\" refers \"%s\" module that does not contain extension definitions.",
-                    ext_p->name, ext_mod->name);
-            ret = LY_EVALID;
-            goto cleanup;
-        }
+    if (ext_def->argument) {
+        LY_CHECK_GOTO(ret = lysp_ext_instance_resolve_argument(ctx->ctx, ext_p, ext_def), cleanup);
     }
 
-    /* find the parsed extension definition there */
-    LY_ARRAY_FOR(ext_mod->parsed->extensions, v) {
-        if (!strcmp(name, ext_mod->parsed->extensions[v].name)) {
-            /* compile extension definition and assign it */
-            LY_CHECK_GOTO(ret = lys_compile_extension(ctx, ext_mod, &ext_mod->parsed->extensions[v], &ext->def), cleanup);
-            break;
-        }
-    }
-    if (!ext->def) {
-        LOGVAL(ctx->ctx, LYVE_REFERENCE, "Extension definition of extension instance \"%s\" not found.", ext_p->name);
-        ret = LY_EVALID;
-        goto cleanup;
-    }
-
-    /* unify the parsed extension from YIN and YANG sources. Without extension definition, it is not possible
-     * to get extension's argument from YIN source, so it is stored as one of the substatements. Here we have
-     * to find it, mark it with LYS_YIN_ARGUMENT and store it in the compiled structure. */
-    if ((ext_p->format == LY_PREF_XML) && ext->def->argument && !ext->argument) {
-        /* Schema was parsed from YIN and an argument is expected, ... */
-        struct lysp_stmt *stmt = NULL;
-
-        if (ext->def->flags & LYS_YINELEM_TRUE) {
-            /* ... argument was the first XML child element */
-            if (ext_p->child && !(ext_p->child->flags & LYS_YIN_ATTR)) {
-                /* TODO check namespace of the statement */
-                if (!strcmp(ext_p->child->stmt, ext->def->argument)) {
-                    stmt = ext_p->child;
-                }
-            }
-        } else {
-            /* ... argument was one of the XML attributes which are represented as child stmt
-             * with LYS_YIN_ATTR flag */
-            for (stmt = ext_p->child; stmt && (stmt->flags & LYS_YIN_ATTR); stmt = stmt->next) {
-                if (!strcmp(stmt->stmt, ext->def->argument)) {
-                    /* this is the extension's argument */
-                    break;
-                }
-            }
-        }
-        if (!stmt) {
-            /* missing extension's argument */
-            LOGVAL(ctx->ctx, LYVE_REFERENCE, "Extension instance \"%s\" misses argument \"%s\".",
-                    ext_p->name, ext->def->argument);
-            ret = LY_EVALID;
-            goto cleanup;
-
-        }
-        LY_CHECK_GOTO(ret = lydict_insert(ctx->ctx, stmt->arg, 0, &ext->argument), cleanup);
-        stmt->flags |= LYS_YIN_ARGUMENT;
-    }
+    DUP_STRING(ctx->ctx, ext_p->argument, ext->argument, ret);
+    LY_CHECK_RET(ret);
 
     if (ext->def->plugin && ext->def->plugin->compile) {
         if (ext->argument) {
@@ -195,7 +126,6 @@
         }
         LY_CHECK_GOTO(ret, cleanup);
     }
-    ext_p->compiled = ext;
 
 cleanup:
     lysc_update_path(ctx, NULL, NULL);
diff --git a/src/tree_schema.h b/src/tree_schema.h
index c4279b7..dc7772f 100644
--- a/src/tree_schema.h
+++ b/src/tree_schema.h
@@ -512,8 +512,6 @@
     void *parent;                           /**< pointer to the parent element holding the extension instance(s), use
                                                  ::lysp_ext_instance#parent_stmt to access the schema element */
     struct lysp_stmt *child;                /**< list of the extension's substatements (linked list) */
-    struct lysc_ext_instance *compiled;     /**< pointer to the compiled data if any - in case the source format is YIN,
-                                                 some of the information (argument) are available only after compilation */
     enum ly_stmt parent_stmt;               /**< value identifying placement of the extension instance */
     LY_ARRAY_COUNT_TYPE parent_stmt_index;  /**< in case the instance is in a substatement, this identifies
                                                  the index of that substatement in its [sized array](@ref sizedarrays) (if any) */
diff --git a/src/tree_schema_helpers.c b/src/tree_schema_helpers.c
index 99f97e4..2845c69 100644
--- a/src/tree_schema_helpers.c
+++ b/src/tree_schema_helpers.c
@@ -1818,6 +1818,125 @@
     return result;
 }
 
+LY_ERR
+lysp_ext_find_definition(const struct ly_ctx *ctx, const struct lysp_ext_instance *ext, const struct lys_module **ext_mod,
+        struct lysp_ext **ext_def)
+{
+    LY_ERR r;
+    const char *tmp, *name, *prefix;
+    size_t pref_len, name_len;
+    LY_ARRAY_COUNT_TYPE v;
+    const struct lys_module *mod = NULL;
+
+    *ext_def = NULL;
+    if (ext_mod) {
+        *ext_mod = NULL;
+    }
+
+    /* parse the prefix */
+    tmp = ext->name;
+    r = ly_parse_nodeid(&tmp, &prefix, &pref_len, &name, &name_len);
+    /* it was parsed already */
+    assert(!r && !tmp[0]);
+
+    /* get module where the extension definition should be placed */
+    mod = ly_resolve_prefix(ctx, prefix, pref_len, ext->format, ext->prefix_data);
+    if (!mod) {
+        LOGVAL(ctx, LYVE_REFERENCE, "Invalid prefix \"%.*s\" used for extension instance identifier.", pref_len, prefix);
+        return LY_EVALID;
+    } else if (!mod->parsed->extensions) {
+        LOGVAL(ctx, LYVE_REFERENCE, "Extension instance \"%s\" refers \"%s\" module that does not contain extension definitions.",
+                ext->name, mod->name);
+        return LY_EVALID;
+    }
+
+    /* find the parsed extension definition there */
+    LY_ARRAY_FOR(mod->parsed->extensions, v) {
+        if (!strcmp(name, mod->parsed->extensions[v].name)) {
+            *ext_def = &mod->parsed->extensions[v];
+            break;
+        }
+    }
+
+    if (!ext_def) {
+        LOGVAL(ctx, LYVE_REFERENCE, "Extension definition of extension instance \"%s\" not found.", ext->name);
+        return LY_EVALID;
+    }
+
+    if (ext_mod) {
+        *ext_mod = mod;
+    }
+    return LY_SUCCESS;
+}
+
+LY_ERR
+lysp_ext_instance_resolve_argument(struct ly_ctx *ctx, struct lysp_ext_instance *ext_p, struct lysp_ext *ext_def)
+{
+    if (!ext_def->argument || ext_p->argument) {
+        /* nothing to do */
+        return LY_SUCCESS;
+    }
+
+    if (ext_p->format == LY_PREF_XML) {
+        /* Schema was parsed from YIN and an argument is expected, ... */
+        struct lysp_stmt *stmt = NULL;
+
+        if (ext_def->flags & LYS_YINELEM_TRUE) {
+            /* ... argument was the first XML child element */
+            for (stmt = ext_p->child; stmt && (stmt->flags & LYS_YIN_ATTR); stmt = stmt->next) {}
+            if (stmt) {
+                const char *arg, *ext, *name_arg, *name_ext, *prefix_arg, *prefix_ext;
+                size_t name_arg_len, name_ext_len, prefix_arg_len, prefix_ext_len;
+
+                stmt = ext_p->child;
+
+                arg = stmt->stmt;
+                ly_parse_nodeid(&arg, &prefix_arg, &prefix_arg_len, &name_arg, &name_arg_len);
+                if (ly_strncmp(ext_def->argument, name_arg, name_arg_len)) {
+                    LOGVAL(ctx, LYVE_SEMANTICS, "Extension instance \"%s\" expects argument element \"%s\" as its first XML child, "
+                            "but \"%.*s\" element found.", ext_p->name, ext_def->argument, name_arg_len, name_arg);
+                    return LY_EVALID;
+                }
+
+                /* check namespace - all the extension instances must be qualified and argument element is expected in the same
+                 * namespace. Do not check just prefixes, there can be different prefixes pointing to the same namespace */
+                ext = ext_p->name; /* include prefix */
+                ly_parse_nodeid(&ext, &prefix_ext, &prefix_ext_len, &name_ext, &name_ext_len);
+
+                if (ly_resolve_prefix(ctx, prefix_ext, prefix_ext_len, ext_p->format, ext_p->prefix_data) !=
+                        ly_resolve_prefix(ctx, prefix_arg, prefix_arg_len, stmt->format, stmt->prefix_data)) {
+                    LOGVAL(ctx, LYVE_SEMANTICS, "Extension instance \"%s\" element and its argument element \"%s\" are "
+                            "expected in the same namespace, but they differ.", ext_p->name, ext_def->argument);
+                    return LY_EVALID;
+                }
+            }
+        } else {
+            /* ... argument was one of the XML attributes which are represented as child stmt
+             * with LYS_YIN_ATTR flag */
+            for (stmt = ext_p->child; stmt && (stmt->flags & LYS_YIN_ATTR); stmt = stmt->next) {
+                if (!strcmp(stmt->stmt, ext_def->argument)) {
+                    /* this is the extension's argument */
+                    break;
+                }
+            }
+        }
+
+        if (stmt) {
+            LY_CHECK_RET(lydict_insert(ctx, stmt->arg, 0, &ext_p->argument));
+            stmt->flags |= LYS_YIN_ARGUMENT;
+        }
+    }
+
+    if (!ext_p->argument) {
+        /* missing extension's argument */
+        LOGVAL(ctx, LYVE_SEMANTICS, "Extension instance \"%s\" misses argument %s\"%s\".",
+                ext_p->name, (ext_def->flags & LYS_YINELEM_TRUE) ? "element " : "", ext_def->argument);
+        return LY_EVALID;
+    }
+
+    return LY_SUCCESS;
+}
+
 LY_ARRAY_COUNT_TYPE
 lysp_ext_instance_iter(struct lysp_ext_instance *ext, LY_ARRAY_COUNT_TYPE index, enum ly_stmt substmt)
 {
diff --git a/src/tree_schema_internal.h b/src/tree_schema_internal.h
index 1e4a9c6..062e1d5 100644
--- a/src/tree_schema_internal.h
+++ b/src/tree_schema_internal.h
@@ -447,6 +447,34 @@
 struct lysc_must **lysc_node_musts_p(const struct lysc_node *node);
 
 /**
+ * @brief Find parsed extension definition for the given extension instance.
+ *
+ * @param[in] ctx libyang context.
+ * @param[in] ext Extension instance for which the definition will be searched.
+ * @param[in, out] ext_mod Pointer to the module where the extension definition of the @p ext to correctly resolve prefixes.
+ * @param[out] ext_def Pointer to return found extension definition.
+ * @return LY_SUCCESS when the definition was found.
+ * @return LY_EVALID when the extension instance is invalid and/or the definition not found.
+ */
+LY_ERR lysp_ext_find_definition(const struct ly_ctx *ctx, const struct lysp_ext_instance *ext, const struct lys_module **ext_mod,
+        struct lysp_ext **ext_def);
+
+/**
+ * @brief When the module comes from YIN format, the argument name is unknown because of missing extension definition
+ * (it might come from import modules which is not yet parsed at that time). Therefore, all the attributes are stored
+ * as substatements and resolving argument is postponed.
+ *
+ * There are 3 places which need the argument, so they resolve it when missing - YIN and YANG printers and extension instance
+ * compiler.
+ *
+ * @param[in] ctx libyang context
+ * @param[in] ext_p Parsed extension to be updated.
+ * @param[in] ext_def Extension definition, found with ::lysp_ext_find_definition().
+ * @return LY_ERR value.
+ */
+LY_ERR lysp_ext_instance_resolve_argument(struct ly_ctx *ctx, struct lysp_ext_instance *ext_p, struct lysp_ext *ext_def);
+
+/**
  * @brief Iterate over the specified type of the extension instances
  *
  * @param[in] ext ([Sized array](@ref sizedarrays)) of extensions to explore