feat: support array of components in @mtkmodel

- for loop or a list comprehension can be used to declare component arrays

Author: ven-k

src/systems/abstractsystem.jl

@@ -1402,7 +1402,7 @@ function _named(name, call, runtime = false)
     end
 end
 
-function _named_idxs(name::Symbol, idxs, call)
+function _named_idxs(name::Symbol, idxs, call; extra_args = "")
     if call.head !== :->
         throw(ArgumentError("Not an anonymous function"))
     end
@@ -1413,7 +1413,10 @@ function _named_idxs(name::Symbol, idxs, call)
     ex = Base.Cartesian.poplinenum(ex)
     ex = _named(:(Symbol($(Meta.quot(name)), :_, $sym)), ex, true)
     ex = Base.Cartesian.poplinenum(ex)
-    :($name = $map($sym -> $ex, $idxs))
+    :($name = map($sym -> begin
+            $extra_args
+            $ex
+        end, $idxs))
 end
 
 function single_named_expr(expr)

src/systems/model_parsing.jl

@@ -40,7 +40,7 @@ function _model_macro(mod, name, expr, isconnector)
         :kwargs => Dict{Symbol, Dict}(),
         :structural_parameters => Dict{Symbol, Dict}()
     )
-    comps = Symbol[]
+    comps = Union{Symbol, Expr}[]
     ext = Ref{Any}(nothing)
     eqs = Expr[]
     icon = Ref{Union{String, URI}}()
@@ -745,7 +745,7 @@ end
 
 ### Parsing Components:
 
-function component_args!(a, b, expr, varexpr, kwargs)
+function component_args!(a, b, varexpr, kwargs; index_name = nothing)
     # Whenever `b` is a function call, skip the first arg aka the function name.
     # Whenever it is a kwargs list, include it.
     start = b.head == :call ? 2 : 1
@@ -754,73 +754,115 @@ function component_args!(a, b, expr, varexpr, kwargs)
         arg isa LineNumberNode && continue
         MLStyle.@match arg begin
             x::Symbol || Expr(:kw, x) => begin
-                _v = _rename(a, x)
-                b.args[i] = Expr(:kw, x, _v)
-                push!(varexpr.args, :((@isdefined $x) && ($_v = $x)))
-                push!(kwargs, Expr(:kw, _v, nothing))
-                # dict[:kwargs][_v] = nothing
+                varname, _varname = _rename(a, x)
+                b.args[i] = Expr(:kw, x, _varname)
+                push!(varexpr.args, :((if $varname !== nothing
+                    $_varname = $varname
+                elseif @isdefined $x
+                    # Allow users to define a var in `structural_parameters` and set
+                    # that as positional arg of subcomponents; it is useful for cases
+                    # where it needs to be passed to multiple subcomponents.
+                    $_varname = $x
+                end)))
+                push!(kwargs, Expr(:kw, varname, nothing))
+                # dict[:kwargs][varname] = nothing
             end
             Expr(:parameters, x...) => begin
-                component_args!(a, arg, expr, varexpr, kwargs)
+                component_args!(a, arg, varexpr, kwargs)
             end
             Expr(:kw, x, y) => begin
-                _v = _rename(a, x)
-                b.args[i] = Expr(:kw, x, _v)
-                push!(varexpr.args, :($_v = $_v === nothing ? $y : $_v))
-                push!(kwargs, Expr(:kw, _v, nothing))
-                # dict[:kwargs][_v] = nothing
+                varname, _varname = _rename(a, x)
+                b.args[i] = Expr(:kw, x, _varname)
+                if isnothing(index_name)
+                    push!(varexpr.args, :($_varname = $varname === nothing ? $y : $varname))
+                else
+                    push!(varexpr.args,
+                        :($_varname = $varname === nothing ? $y : $varname[$index_name]))
+                end
+                push!(kwargs, Expr(:kw, varname, nothing))
+                # dict[:kwargs][varname] = nothing
             end
             _ => error("Could not parse $arg of component $a")
         end
     end
 end
 
-function _parse_components!(exprs, body, kwargs)
-    expr = Expr(:block)
+model_name(name, range) = Symbol.(name, :_, collect(range))
+
+function _parse_components!(body, kwargs)
+    local expr
     varexpr = Expr(:block)
-    # push!(exprs, varexpr)
-    comps = Vector{Union{Symbol, Expr}}[]
+    comps = Vector{Union{Union{Expr, Symbol}, Expr}}[]
     comp_names = []
 
-    for arg in body.args
-        arg isa LineNumberNode && continue
-        MLStyle.@match arg begin
-            Expr(:block) => begin
-                # TODO: Do we need this?
-                error("Multiple `@components` block detected within a single block")
-            end
-            Expr(:(=), a, b) => begin
-                arg = deepcopy(arg)
-                b = deepcopy(arg.args[2])
+    Base.remove_linenums!(body)
+    arg = body.args[end]
+
+    MLStyle.@match arg begin
+        Expr(:(=), a, Expr(:comprehension, Expr(:generator, b, Expr(:(=), c, d)))) => begin
+            array_varexpr = Expr(:block)
 
-                component_args!(a, b, expr, varexpr, kwargs)
+            push!(comp_names, :($a...))
+            push!(comps, [a, b.args[1], d])
+            b = deepcopy(b)
 
-                arg.args[2] = b
-                push!(expr.args, arg)
-                push!(comp_names, a)
+            component_args!(a, b, array_varexpr, kwargs; index_name = c)
+
+            expr = _named_idxs(a, d, :($c -> $b); extra_args = array_varexpr)
+        end
+        Expr(:(=), a, Expr(:comprehension, Expr(:generator, b, Expr(:filter, e, Expr(:(=), c, d))))) => begin
+            error("List comprehensions with conditional statements aren't supported.")
+        end
+        Expr(:(=), a, Expr(:comprehension, Expr(:generator, b, Expr(:(=), c, d), e...))) => begin
+            # Note that `e` is of the form `Tuple{Expr(:(=), c, d)}`
+            error("More than one index isn't supported while building component array")
+        end
+        Expr(:block) => begin
+            # TODO: Do we need this?
+            error("Multiple `@components` block detected within a single block")
+        end
+        Expr(:(=), a, Expr(:for, Expr(:(=), c, d), b)) => begin
+            Base.remove_linenums!(b)
+            array_varexpr = Expr(:block)
+            push!(array_varexpr.args, b.args[1:(end - 1)]...)
+            push!(comp_names, :($a...))
+            push!(comps, [a, b.args[end].args[1], d])
+            b = deepcopy(b)
+
+            component_args!(a, b.args[end], array_varexpr, kwargs; index_name = c)
+
+            expr = _named_idxs(a, d, :($c -> $(b.args[end])); extra_args = array_varexpr)
+        end
+        Expr(:(=), a, b) => begin
+            arg = deepcopy(arg)
+            b = deepcopy(arg.args[2])
+
+            component_args!(a, b, varexpr, kwargs)
+
+            arg.args[2] = b
+            expr = :(@named $arg)
+            push!(comp_names, a)
                 if (isa(b.args[1], Symbol) || Meta.isexpr(b.args[1], :.))
-                    push!(comps, [a, b.args[1]])
+                push!(comps, [a, b.args[1]])
                 end
-            end
-            _ => error("Couldn't parse the component body: $arg")
         end
+        _ => error("Couldn't parse the component body: $arg")
     end
+
     return comp_names, comps, expr, varexpr
 end
 
 function push_conditional_component!(ifexpr, expr_vec, comp_names, varexpr)
     blk = Expr(:block)
     push!(blk.args, varexpr)
-    push!(blk.args, :(@named begin
-        $(expr_vec.args...)
-    end))
+    push!(blk.args, expr_vec)
     push!(blk.args, :($push!(systems, $(comp_names...))))
     push!(ifexpr.args, blk)
 end
 
 function handle_if_x!(mod, exprs, ifexpr, x, kwargs, condition = nothing)
     push!(ifexpr.args, condition)
-    comp_names, comps, expr_vec, varexpr = _parse_components!(ifexpr, x, kwargs)
+    comp_names, comps, expr_vec, varexpr = _parse_components!(x, kwargs)
     push_conditional_component!(ifexpr, expr_vec, comp_names, varexpr)
     comps
 end
@@ -836,7 +878,7 @@ function handle_if_y!(exprs, ifexpr, y, kwargs)
         push!(ifexpr.args, elseifexpr)
         (comps...,)
     else
-        comp_names, comps, expr_vec, varexpr = _parse_components!(exprs, y, kwargs)
+        comp_names, comps, expr_vec, varexpr = _parse_components!(y, kwargs)
         push_conditional_component!(ifexpr, expr_vec, comp_names, varexpr)
         comps
     end
@@ -861,25 +903,23 @@ function parse_components!(exprs, cs, dict, compbody, kwargs)
             Expr(:if, condition, x, y) => begin
                 handle_conditional_components(condition, dict, exprs, kwargs, x, y)
             end
-            Expr(:(=), a, b) => begin
-                comp_names, comps, expr_vec, varexpr = _parse_components!(exprs,
-                    :(begin
+            # Either the arg is top level component declaration or an invalid cause - both are handled by `_parse_components`
+            _ => begin
+                comp_names, comps, expr_vec, varexpr = _parse_components!(:(begin
                         $arg
                     end),
                     kwargs)
                 push!(cs, comp_names...)
                 push!(dict[:components], comps...)
-                push!(exprs, varexpr, :(@named begin
-                    $(expr_vec.args...)
-                end))
+                push!(exprs, varexpr, expr_vec)
             end
-            _ => error("Couldn't parse the component body $compbody")
         end
     end
 end
 
 function _rename(compname, varname)
     compname = Symbol(compname, :__, varname)
+    (compname, Symbol(:_, compname))
 end
 
 # Handle top level branching

test/model_parsing.jl

@@ -650,3 +650,41 @@ end
 @named m = MyModel()
 @variables x___(t)
 @test isequal(x___, _b[])
+
+@testset "Component array" begin
+    @mtkmodel SubComponent begin
+        @parameters begin
+            sc
+        end
+    end
+
+    @mtkmodel Component begin
+        @structural_parameters begin
+            N = 2
+        end
+        @components begin
+            comprehension = [SubComponent(sc = i) for i in 1:N]
+            written_out_for = for i in 1:N
+                sc = i + 1
+                SubComponent(; sc)
+            end
+            single_sub_component = SubComponent()
+        end
+    end
+
+    @named component = Component()
+    component = complete(component)
+
+    @test nameof.(ModelingToolkit.get_systems(component)) == [
+        :comprehension_1,
+        :comprehension_2,
+        :written_out_for_1,
+        :written_out_for_2,
+        :single_sub_component,
+    ]
+
+    @test getdefault(component.comprehension_1.sc) == 1
+    @test getdefault(component.comprehension_2.sc) == 2
+    @test getdefault(component.written_out_for_1.sc) == 2
+    @test getdefault(component.written_out_for_2.sc) == 3
+end