From dddd69884244be31e9f8d9d3a83b9b126eba16ce Mon Sep 17 00:00:00 2001
From: Venkateshprasad <32921645+ven-k@users.noreply.github.com>
Date: Wed, 30 Oct 2024 18:27:34 +0000
Subject: [PATCH] fix: unpack components of the base sys in implicit extend

---
 src/systems/model_parsing.jl | 22 +++++++++++-----------
 test/model_parsing.jl        | 26 ++++++++++++++++++++++++++
 2 files changed, 37 insertions(+), 11 deletions(-)

diff --git a/src/systems/model_parsing.jl b/src/systems/model_parsing.jl
index ab70646824..e8955b5b84 100644
--- a/src/systems/model_parsing.jl
+++ b/src/systems/model_parsing.jl
@@ -772,18 +772,17 @@ function Base.names(model::Model)
         map(first, get(model.structure, :components, EMPTY_VoVoSYMBOL))))
 end
 
-function _parse_extend!(ext, a, b, dict, expr, kwargs, vars, additional_args)
+function _parse_extend!(ext, a, b, dict, expr, kwargs, vars, implicit_arglist)
     extend_args!(a, b, dict, expr, kwargs)
 
-    # `additional_args` doubles as a flag to check the mode of `@extend`. It is
+    # `implicit_arglist` doubles as a flag to check the mode of `@extend`. It is
     # `nothing` for explicit destructuring.
     # The following block modifies the arguments of both base and higher systems
     # for the implicit extend statements.
-    if additional_args !== nothing
+    if implicit_arglist !== nothing
         b.args = [b.args[1]]
-        allvars = [additional_args.args..., vars.args...]
         push!(b.args, Expr(:parameters))
-        for var in allvars
+        for var in implicit_arglist.args
             push!(b.args[end].args, var)
             if !haskey(dict[:kwargs], var)
                 push!(dict[:kwargs], var => Dict(:value => NO_VALUE))
@@ -814,8 +813,8 @@ function parse_extend!(exprs, ext, dict, mod, body, kwargs)
                 end
                 a, b = b.args
                 # This doubles as a flag to identify the mode of `@extend`
-                additional_args = nothing
-                _parse_extend!(ext, a, b, dict, expr, kwargs, vars, additional_args)
+                implicit_arglist = nothing
+                _parse_extend!(ext, a, b, dict, expr, kwargs, vars, implicit_arglist)
             else
                 error("When explicitly destructing in `@extend` please use the syntax: `@extend a, b = oneport = OnePort()`.")
             end
@@ -825,11 +824,12 @@ function parse_extend!(exprs, ext, dict, mod, body, kwargs)
             b = body
             if (model = getproperty(mod, b.args[1])) isa Model
                 vars = Expr(:tuple)
-                append!(vars.args, _arguments(model))
-                additional_args = Expr(:tuple)
-                append!(additional_args.args,
+                append!(vars.args, names(model))
+                implicit_arglist = Expr(:tuple)
+                append!(implicit_arglist.args, _arguments(model))
+                append!(implicit_arglist.args,
                     keys(get(model.structure, :structural_parameters, EMPTY_DICT)))
-                _parse_extend!(ext, a, b, dict, expr, kwargs, vars, additional_args)
+                _parse_extend!(ext, a, b, dict, expr, kwargs, vars, implicit_arglist)
             else
                 error("Cannot infer the exact `Model` that `@extend $(body)` refers." *
                       " Please specify the names that it brings into scope by:" *
diff --git a/test/model_parsing.jl b/test/model_parsing.jl
index 2dcb8e3f11..9cdd8712d4 100644
--- a/test/model_parsing.jl
+++ b/test/model_parsing.jl
@@ -932,3 +932,29 @@ end
     @test getdefault(main_sys.p2) == 12
     @test getdefault(main_sys.v1) == 13
 end
+
+@mtkmodel InnerModel begin
+    @parameters begin
+        p
+    end
+end
+
+@mtkmodel MidModel begin
+    @components begin
+        inmodel = InnerModel()
+    end
+end
+
+@mtkmodel OuterModel begin
+    @extend MidModel()
+    @equations begin
+        inmodel.p ~ 0
+    end
+end
+
+# The base system is fetched from the module while extending implicitly. This
+# way of defining fails when defined inside the `@testset`. So, it is moved out.
+@testset "Test unpacking of components in implicit extend" begin
+    @named out = OuterModel()
+    @test OuterModel.structure[:extend][1] == [:inmodel]
+end