fix: improve handling of expression in metadata

ven-k · ven-k · commit fc32fe902029 · 2024-04-22T14:48:59.000+05:30
The variables in mtkmodel can have complicated expressions as `guess` and other
metadata. Instead of evaluating them while defining the mtk-model, make those exprs
part of the the definition.

Along with this, the metadata (in structure dict) is kept in expr form.
diff --git a/docs/src/basics/MTKModel_Connector.md b/docs/src/basics/MTKModel_Connector.md
@@ -298,7 +298,7 @@ end
 ```
 
 !!! note
-    
+
     For more examples of usage, checkout [ModelingToolkitStandardLibrary.jl](https://github.com/SciML/ModelingToolkitStandardLibrary.jl/)
 
 ## More on `Model.structure`
@@ -408,12 +408,13 @@ The conditional parts are reflected in the `structure`. For `BranchOutsideTheBlo
 
 ```julia
 julia> BranchOutsideTheBlock.structure
-Dict{Symbol, Any} with 6 entries:
+Dict{Symbol, Any} with 7 entries:
   :components            => Any[(:if, :flag, Vector{Union{Expr, Symbol}}[[:sys1, :C]], Any[])]
   :kwargs                => Dict{Symbol, Dict}(:flag=>Dict{Symbol, Bool}(:value=>1))
   :structural_parameters => Dict{Symbol, Dict}(:flag=>Dict{Symbol, Bool}(:value=>1))
   :independent_variable  => t
-  :parameters            => Dict{Symbol, Dict{Symbol, Any}}(:a2 => Dict(:type => AbstractArray{Real}, :condition => (:if, :flag, Dict{Symbol, Any}(:kwargs => Dict{Any, Any}(:a1 => Dict{Symbol, Union{Nothing, DataType}}(:value => nothing, :type => Real)), :parameters => Any[Dict{Symbol, Dict{Symbol, Any}}(:a1 => Dict(:type => AbstractArray{Real}))]), Dict{Symbol, Any}(:variables => Any[Dict{Symbol, Dict{Symbol, Any}}()], :kwargs => Dict{Any, Any}(:a2 => Dict{Symbol, Union{Nothing, DataType}}(:value => nothing, :type => Real)), :parameters => Any[Dict{Symbol, Dict{Symbol, Any}}(:a2 => Dict(:type => AbstractArray{Real}))]))), :a1 => Dict(:type => AbstractArray{Real}, :condition => (:if, :flag, Dict{Symbol, Any}(:kwargs => Dict{Any, Any}(:a1 => Dict{Symbol, Union{Nothing, DataType}}(:value => nothing, :type => Real)), :parameters => Any[Dict{Symbol, Dict{Symbol, Any}}(:a1 => Dict(:type => AbstractArray{Real}))]), Dict{Symbol, Any}(:variables => Any[Dict{Symbol, Dict{Symbol, Any}}()], :kwargs => Dict{Any, Any}(:a2 => Dict{Symbol, Union{Nothing, DataType}}(:value => nothing, :type => Real)), :parameters => Any[Dict{Symbol, Dict{Symbol, Any}}(:a2 => Dict(:type => AbstractArray{Real}))]))))
+  :parameters            => Dict{Symbol, Dict{Symbol, Any}}(:a2 => Dict(:type=>AbstractArray{Real}, :condition=>(:if, :flag, Dict{Symbol, Any}(:kwargs=>Dict{Any, Any}(:a1=>Dict{Symbol, Union{Nothing, DataType}}(:value=>nothing, :type=>Real)), :parameters=>Any[Dict{Symbol, Dict{Symbol, Any}}(:a1=>Dict(:type=>AbstractArray{Real}))]), Dict{Symbol, Any}(:variables=>Any[Dict{Symbol, Dict{Symbol, Any}}()], :kwargs=>Dict{Any, Any}(:a2=>Dict{Symbol, Union{Nothing, DataType}}(:value=>nothing, :type=>Real)), :parameters=>Any[Dict{Symbol, Dict{Symbol, Any}}(:a2=>Dict(:type=>AbstractArray{Real}))]))), :a1 => Dict(:type=>AbstractArray{Real}, :condition=>(:if, :flag, Dict{Symbol, Any}(:kwargs=>Dict{Any, Any}(:a1=>Dict{Symbol, Union{Nothing, DataType}}(:value=>nothing, :type=>Real)), :parameters=>Any[Dict{Symbol, Dict{Symbol, Any}}(:a1=>Dict(:type=>AbstractArray{Real}))]), Dict{Symbol, Any}(:variables=>Any[Dict{Symbol, Dict{Symbol, Any}}()], :kwargs=>Dict{Any, Any}(:a2=>Dict{Symbol, Union{Nothing, DataType}}(:value=>nothing, :type=>Real)), :parameters=>Any[Dict{Symbol, Dict{Symbol, Any}}(:a2=>Dict(:type=>AbstractArray{Real}))]))))
+  :defaults              => Dict{Symbol, Any}(:a1=>10)
   :equations             => Any[(:if, :flag, ["a1 ~ 0"], ["a2 ~ 0"])]
 ```
 
diff --git a/src/systems/model_parsing.jl b/src/systems/model_parsing.jl
@@ -187,7 +187,7 @@ function parse_variable_def!(dict, mod, arg, varclass, kwargs, where_types;
             var = generate_var!(dict, a, varclass; indices, type)
             update_kwargs_and_metadata!(dict, kwargs, a, def, indices, type, var,
                 varclass, where_types)
-            (var, def)
+            return var, def, Dict()
         end
         Expr(:(::), a, type) => begin
             type = getfield(mod, type)
@@ -201,13 +201,13 @@ function parse_variable_def!(dict, mod, arg, varclass, kwargs, where_types;
         Expr(:call, a, b) => begin
             var = generate_var!(dict, a, b, varclass, mod; indices, type)
             update_kwargs_and_metadata!(dict, kwargs, a, def, indices, type, var,
-                varclass, where_types)
-            (var, def)
+            varclass, where_types)
+            return var, def, Dict()
         end
         Expr(:(=), a, b) => begin
             Base.remove_linenums!(b)
             def, meta = parse_default(mod, b)
-            var, def = parse_variable_def!(
+            var, def, _ = parse_variable_def!(
                 dict, mod, a, varclass, kwargs, where_types; def, type)
             if dict[varclass] isa Vector
                 dict[varclass][1][getname(var)][:default] = def
@@ -225,12 +225,13 @@ function parse_variable_def!(dict, mod, arg, varclass, kwargs, where_types;
                         end
                     end
                 end
-                var = set_var_metadata(var, meta)
+                var, metadata_with_exprs = set_var_metadata(var, meta)
+                return var, def, metadata_with_exprs
             end
-            (var, def)
+            return var, def, Dict()
         end
         Expr(:tuple, a, b) => begin
-            var, def = parse_variable_def!(
+            var, def, _ = parse_variable_def!(
                 dict, mod, a, varclass, kwargs, where_types; type)
             meta = parse_metadata(mod, b)
             if meta !== nothing
@@ -244,9 +245,10 @@ function parse_variable_def!(dict, mod, arg, varclass, kwargs, where_types;
                         end
                     end
                 end
-                var = set_var_metadata(var, meta)
+                var, metadata_with_exprs = set_var_metadata(var, meta)
+                return var, def, metadata_with_exprs
             end
-            (var, def)
+            return var, def, Dict()
         end
         Expr(:ref, a, b...) => begin
             indices = map(i -> UnitRange(i.args[2], i.args[end]), b)
@@ -350,21 +352,27 @@ function parse_metadata(mod, a)
     end
 end
 
+function _set_var_metadata!(metadata_with_exprs, a, m, v::Expr)
+    push!(metadata_with_exprs, m => v)
+    a
+end
+_set_var_metadata!(metadata_with_exprs, a, m, v) = wrap(set_scalar_metadata(unwrap(a), m, v))
+
 function set_var_metadata(a, ms)
+    metadata_with_exprs = Dict{DataType, Expr}()
     for (m, v) in ms
-        a = wrap(set_scalar_metadata(unwrap(a), m, v))
+        m == VariableGuess && v isa Symbol && (v = quote $v end)
+        a = _set_var_metadata!(metadata_with_exprs, a, m, v)
     end
-    a
+    a, metadata_with_exprs
 end
 
 function get_var(mod::Module, b)
     if b isa Symbol
-        getproperty(mod, b)
-    elseif b isa Expr
-        Core.eval(mod, b)
-    else
-        b
+        isdefined(mod, b) && return getproperty(mod, b)
+        isdefined(@__MODULE__, b) && return getproperty(@__MODULE__, b)
     end
+    b
 end
 
 function parse_model!(exprs, comps, ext, eqs, icon, vs, ps, sps, c_evts, d_evts,
@@ -595,10 +603,24 @@ function parse_variable_arg!(exprs, vs, dict, mod, arg, varclass, kwargs, where_
 end
 
 function parse_variable_arg(dict, mod, arg, varclass, kwargs, where_types)
-    vv, def = parse_variable_def!(dict, mod, arg, varclass, kwargs, where_types)
+    vv, def, metadata_with_exprs = parse_variable_def!(dict, mod, arg, varclass, kwargs, where_types)
     name = getname(vv)
-    return vv isa Num ? name : :($name...),
-    :($name = $name === nothing ? $setdefault($vv, $def) : $setdefault($vv, $name))
+
+    varexpr = quote
+        $name = if $name === nothing
+            $setdefault($vv, $def)
+        else
+            $setdefault($vv, $name)
+        end
+    end
+
+    metadata_expr = Expr(:block)
+    for (k,v) in metadata_with_exprs
+        push!(metadata_expr.args, :($name = $wrap($set_scalar_metadata($unwrap($name), $k, $v))))
+    end
+
+    push!(varexpr.args, metadata_expr)
+    return vv isa Num ? name : :($name...), varexpr
 end
 
 function handle_conditional_vars!(
@@ -882,7 +904,7 @@ function _parse_components!(body, kwargs)
     Base.remove_linenums!(body)
     arg = body.args[end]
 
-    MLStyle.@match arg begin
+        MLStyle.@match arg begin
         Expr(:(=), a, Expr(:comprehension, Expr(:generator, b, Expr(:(=), c, d)))) => begin
             array_varexpr = Expr(:block)
 
@@ -901,10 +923,10 @@ function _parse_components!(body, kwargs)
             # 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(: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)
@@ -917,21 +939,21 @@ function _parse_components!(body, kwargs)
 
             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])
+            Expr(:(=), a, b) => begin
+                arg = deepcopy(arg)
+                b = deepcopy(arg.args[2])
 
             component_args!(a, b, varexpr, kwargs)
 
-            arg.args[2] = b
+                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!(comp_names, a)
+                if (isa(b.args[1], Symbol) || Meta.isexpr(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
diff --git a/test/model_parsing.jl b/test/model_parsing.jl
@@ -376,8 +376,8 @@ end
 
 @testset "Metadata in variables" begin
     metadata = Dict(:description => "Variable to test metadata in the Model.structure",
-        :input => true, :bounds => (-1, 1), :connection_type => :Flow,
-        :tunable => false, :disturbance => true, :dist => Normal(1, 1))
+        :input => true, :bounds => :((-1, 1)), :connection_type => :Flow,
+        :tunable => false, :disturbance => true, :dist => :(Normal(1, 1)))
 
     @connector MockMeta begin
         m(t),
@@ -473,7 +473,7 @@ using ModelingToolkit: getdefault, scalarize
 
     @named model_with_component_array = ModelWithComponentArray()
 
-    @test ModelWithComponentArray.structure[:parameters][:r][:unit] == u"Ω"
+    @test eval(ModelWithComponentArray.structure[:parameters][:r][:unit]) == eval(u"Ω")
     @test lastindex(parameters(model_with_component_array)) == 3
 
     # Test the constant `k`. Manually k's value should be kept in sync here
