Skip to content

Commit 07f7f4d

Browse files
committed
feat: support conditional statements inside @equations,
`@parameters`,`@variables`, `@components` in the `@mtkmodel` and `@components`.
1 parent 397ab12 commit 07f7f4d

File tree

1 file changed

+226
-73
lines changed

1 file changed

+226
-73
lines changed

src/systems/model_parsing.jl

Lines changed: 226 additions & 73 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@ struct Model{F, S}
1616
(:structural_parameters) and equations (:equations).
1717
"""
1818
structure::S
19-
"""
19+
"""
2020
This flag is `true` when the Model is a connector and is `false` when it is
2121
a component
2222
"""
@@ -25,7 +25,7 @@ end
2525
(m::Model)(args...; kw...) = m.f(args...; kw...)
2626

2727
for f in (:connector, :mtkmodel)
28-
isconnector = f == :connector ? true : false
28+
isconnector = f == :connector ? true : false
2929
@eval begin
3030
macro $f(name::Symbol, body)
3131
esc($(:_model_macro)(__module__, name, body, $isconnector))
@@ -41,7 +41,11 @@ function _model_macro(mod, name, expr, isconnector)
4141
ext = Ref{Any}(nothing)
4242
eqs = Expr[]
4343
icon = Ref{Union{String, URI}}()
44-
kwargs, ps, sps, vs, = [], [], [], []
44+
ps, sps, vs, = [], [], []
45+
kwargs = Set()
46+
47+
push!(exprs.args, :(systems = ODESystem[]))
48+
push!(exprs.args, :(equations = Equation[]))
4549

4650
for arg in expr.args
4751
arg isa LineNumberNode && continue
@@ -54,7 +58,7 @@ function _model_macro(mod, name, expr, isconnector)
5458
# Connectors can have variables listed without `@variables` prefix or
5559
# begin block.
5660
parse_variable_arg!(exprs, vs, dict, mod, arg, :variables, kwargs)
57-
else
61+
else
5862
error("$arg is not valid syntax. Expected a macro call.")
5963
end
6064
end
@@ -64,11 +68,14 @@ function _model_macro(mod, name, expr, isconnector)
6468
iv = dict[:independent_variable] = variable(:t)
6569
end
6670

71+
push!(exprs.args, :(push!(systems, $(comps...))))
72+
push!(exprs.args, :(push!(equations, $(eqs...))))
73+
6774
gui_metadata = isassigned(icon) > 0 ? GUIMetadata(GlobalRef(mod, name), icon[]) :
6875
GUIMetadata(GlobalRef(mod, name))
6976

70-
sys = :($ODESystem($Equation[$(eqs...)], $iv, [$(vs...)], [$(ps...)];
71-
name, systems = [$(comps...)], gui_metadata = $gui_metadata))
77+
sys = :($ODESystem($Equation[equations...], $iv, [$(vs...)], [$(ps...)];
78+
name, systems, gui_metadata = $gui_metadata))
7279

7380
if ext[] === nothing
7481
push!(exprs.args, :(var"#___sys___" = $sys))
@@ -213,10 +220,33 @@ function parse_default(mod, a)
213220
end
214221
(expr, nothing)
215222
end
216-
_ => error("Cannot parse default $a")
223+
#=Expr(:if, condition::Expr, x, y) => begin
224+
@info 212
225+
if condition.args[1] in (:(==), :(<), :(>))
226+
op = compare_op(condition.args[1])
227+
expr = Expr(:call)
228+
push!(expr.args, op)
229+
for cond in condition.args[2:end]
230+
# cond isa Symbol ? push!(expr.args, :($getdefault($cond))) :
231+
push!(expr.args, cond)
232+
end
233+
a.args[1] = expr
234+
end
235+
(a, nothing)
236+
end=#
237+
Expr(:if, condition, x, y) => (a, nothing)
238+
_ => error("Cannot parse default $a $(typeof(a))")
217239
end
218240
end
219241

242+
compare_op(a) = if a == :(==)
243+
:isequal
244+
elseif a == :(<)
245+
:isless
246+
elseif a == :(>)
247+
:(Base.isgreater)
248+
end
249+
220250
function parse_metadata(mod, a)
221251
MLStyle.@match a begin
222252
Expr(:vect, eles...) => Dict(parse_metadata(mod, e) for e in eles)
@@ -247,7 +277,7 @@ function parse_model!(exprs, comps, ext, eqs, icon, vs, ps, sps,
247277
mname = arg.args[1]
248278
body = arg.args[end]
249279
if mname == Symbol("@components")
250-
parse_components!(exprs, comps, dict, body, kwargs)
280+
parse_components!(mod, exprs, comps, dict, body, kwargs)
251281
elseif mname == Symbol("@extend")
252282
parse_extend!(exprs, ext, dict, mod, body, kwargs)
253283
elseif mname == Symbol("@variables")
@@ -284,68 +314,6 @@ function parse_structural_parameters!(exprs, sps, dict, mod, body, kwargs)
284314
end
285315
end
286316

287-
function parse_components!(exprs, cs, dict, body, kwargs)
288-
expr = Expr(:block)
289-
varexpr = Expr(:block)
290-
push!(exprs, varexpr)
291-
comps = Vector{Symbol}[]
292-
for arg in body.args
293-
arg isa LineNumberNode && continue
294-
MLStyle.@match arg begin
295-
Expr(:(=), a, b) => begin
296-
push!(cs, a)
297-
push!(comps, [a, b.args[1]])
298-
arg = deepcopy(arg)
299-
b = deepcopy(arg.args[2])
300-
301-
component_args!(a, b, dict, expr, varexpr, kwargs)
302-
303-
arg.args[2] = b
304-
push!(expr.args, arg)
305-
end
306-
_ => error("`@components` only takes assignment expressions. Got $arg")
307-
end
308-
end
309-
310-
push!(exprs, :(@named $expr))
311-
312-
dict[:components] = comps
313-
end
314-
315-
function _rename(compname, varname)
316-
compname = Symbol(compname, :__, varname)
317-
end
318-
319-
function component_args!(a, b, dict, expr, varexpr, kwargs)
320-
# Whenever `b` is a function call, skip the first arg aka the function name.
321-
# Whenever it is a kwargs list, include it.
322-
start = b.head == :call ? 2 : 1
323-
for i in start:lastindex(b.args)
324-
arg = b.args[i]
325-
arg isa LineNumberNode && continue
326-
MLStyle.@match arg begin
327-
x::Symbol || Expr(:kw, x) => begin
328-
_v = _rename(a, x)
329-
b.args[i] = Expr(:kw, x, _v)
330-
push!(varexpr.args, :((@isdefined $x) && ($_v = $x)))
331-
push!(kwargs, Expr(:kw, _v, nothing))
332-
dict[:kwargs][_v] = nothing
333-
end
334-
Expr(:parameters, x...) => begin
335-
component_args!(a, arg, dict, expr, varexpr, kwargs)
336-
end
337-
Expr(:kw, x, y) => begin
338-
_v = _rename(a, x)
339-
b.args[i] = Expr(:kw, x, _v)
340-
push!(varexpr.args, :($_v = $_v === nothing ? $y : $_v))
341-
push!(kwargs, Expr(:kw, _v, nothing))
342-
dict[:kwargs][_v] = nothing
343-
end
344-
_ => error("Could not parse $arg of component $a")
345-
end
346-
end
347-
end
348-
349317
function extend_args!(a, b, dict, expr, kwargs, varexpr, has_param = false)
350318
# Whenever `b` is a function call, skip the first arg aka the function name.
351319
# Whenver it is a kwargs list, include it.
@@ -371,6 +339,14 @@ function extend_args!(a, b, dict, expr, kwargs, varexpr, has_param = false)
371339
dict[:kwargs][x] = nothing
372340
end
373341
Expr(:kw, x, y) => begin
342+
#= _v = _rename(a, x)
343+
push!(expr.args, :($_v = $y))
344+
def = Expr(:kw)
345+
push!(def.args, x)
346+
push!(def.args, :($getdefault($_v)))
347+
b.args[i] = def
348+
# b.args[i] = Expr(:kw, x, _v)
349+
push!(kwargs, Expr(:kw, _v, nothing))=#
374350
b.args[i] = Expr(:kw, x, x)
375351
push!(varexpr.args, :($x = $x === nothing ? $y : $x))
376352
push!(kwargs, Expr(:kw, x, nothing))
@@ -463,13 +439,50 @@ function parse_variables!(exprs, vs, dict, mod, body, varclass, kwargs)
463439
end
464440
end
465441

442+
function handle_if_x_equations!(ifexpr, condition, x, dict)
443+
push!(ifexpr.args, condition, :(push!(equations, $(x.args...))))
444+
# push!(dict[:equations], [:if, readable_code(condition), readable_code.(x.args)])
445+
readable_code.(x.args)
446+
end
447+
448+
function handle_if_y_equations!(ifexpr, y, dict)
449+
if y.head == :elseif
450+
elseifexpr = Expr(:elseif)
451+
eq_entry = [:elseif, readable_code.(y.args[1].args)...]
452+
push!(eq_entry, handle_if_x_equations!(elseifexpr, y.args[1], y.args[2], dict))
453+
get(y.args, 3, nothing) !== nothing && push!(eq_entry, handle_if_y_equations!(elseifexpr, y.args[3], dict))
454+
push!(ifexpr.args, elseifexpr)
455+
(eq_entry...,)
456+
else
457+
push!(ifexpr.args, :(push!(equations, $(y.args...))))
458+
readable_code.(y.args)
459+
end
460+
end
461+
466462
function parse_equations!(exprs, eqs, dict, body)
463+
dict[:equations] = []
464+
Base.remove_linenums!(body)
467465
for arg in body.args
468-
arg isa LineNumberNode && continue
469-
push!(eqs, arg)
466+
MLStyle.@match arg begin
467+
Expr(:if, condition, x) => begin
468+
ifexpr = Expr(:if)
469+
eq_entry = handle_if_x_equations!(ifexpr, condition, x, dict)
470+
push!(exprs, ifexpr)
471+
push!(dict[:equations], [:if, condition, eq_entry])
472+
end
473+
Expr(:if, condition, x, y) => begin
474+
ifexpr = Expr(:if)
475+
xeq_entry = handle_if_x_equations!(ifexpr, condition, x, dict)
476+
yeq_entry = handle_if_y_equations!(ifexpr, y, dict)
477+
push!(exprs, ifexpr)
478+
push!(dict[:equations], [:if, condition, xeq_entry, yeq_entry])
479+
# push!(dict[:equations], yeq_entry...)
480+
end
481+
_ => push!(eqs, arg)
482+
end
470483
end
471484
# TODO: does this work with TOML?
472-
dict[:equations] = readable_code.(eqs)
485+
push!(dict[:equations], readable_code.(eqs)...)
473486
end
474487

475488
function parse_icon!(icon, dict, body::String)
@@ -494,3 +507,143 @@ end
494507
function parse_icon!(icon, dict, body::Expr)
495508
parse_icon!(icon, dict, eval(body))
496509
end
510+
511+
### Parsing Components:
512+
513+
function component_args!(a, b, expr, varexpr, kwargs)
514+
# Whenever `b` is a function call, skip the first arg aka the function name.
515+
# Whenever it is a kwargs list, include it.
516+
start = b.head == :call ? 2 : 1
517+
for i in start:lastindex(b.args)
518+
arg = b.args[i]
519+
arg isa LineNumberNode && continue
520+
MLStyle.@match arg begin
521+
x::Symbol || Expr(:kw, x) => begin
522+
_v = _rename(a, x)
523+
b.args[i] = Expr(:kw, x, _v)
524+
push!(varexpr.args, :((@isdefined $x) && ($_v = $x)))
525+
push!(kwargs, Expr(:kw, _v, nothing))
526+
# dict[:kwargs][_v] = nothing
527+
end
528+
Expr(:parameters, x...) => begin
529+
component_args!(a, arg, expr, varexpr, kwargs)
530+
end
531+
Expr(:kw, x, y) => begin
532+
_v = _rename(a, x)
533+
b.args[i] = Expr(:kw, x, _v)
534+
push!(varexpr.args, :($_v = $_v === nothing ? $y : $_v))
535+
push!(kwargs, Expr(:kw, _v, nothing))
536+
# dict[:kwargs][_v] = nothing
537+
end
538+
_ => error("Could not parse $arg of component $a")
539+
end
540+
end
541+
end
542+
543+
function _parse_components!(exprs, body, kwargs)
544+
expr = Expr(:block)
545+
varexpr = Expr(:block)
546+
# push!(exprs, varexpr)
547+
comps = Vector{Symbol}[]
548+
comp_names = []
549+
550+
for arg in body.args
551+
arg isa LineNumberNode && continue
552+
MLStyle.@match arg begin
553+
Expr(:(=), a, b) => begin
554+
arg = deepcopy(arg)
555+
b = deepcopy(arg.args[2])
556+
557+
component_args!(a, b, expr, varexpr, kwargs)
558+
559+
# push!(b.args, Expr(:kw, :name, Meta.quot(a)))
560+
# arg.args[2] = b
561+
562+
push!(expr.args, arg)
563+
push!(comp_names, a)
564+
push!(comps, [a, b.args[1]])
565+
end
566+
_ => @info "Couldn't parse the component body: $arg"
567+
end
568+
end
569+
return comp_names, comps, expr, varexpr
570+
end
571+
572+
function push_conditional_component!(ifexpr, expr_vec, comp_names, varexpr)
573+
blk = Expr(:block)
574+
push!(blk.args, varexpr)
575+
push!(blk.args, :(@named begin $(expr_vec.args...) end))
576+
push!(blk.args, :($push!(systems, $(comp_names...))))
577+
push!(ifexpr.args, blk)
578+
end
579+
580+
function handle_if_x!(mod, exprs, ifexpr, x, kwargs, condition = nothing)
581+
@info 576 condition typeof(condition)
582+
# push!(ifexpr.args, :($substitute_defaults($condition)))
583+
#= if condition isa Symbol
584+
@info 579 condition
585+
push!(ifexpr.args, :($getdefault($condition)))
586+
elseif condition isa Num
587+
push!(ifexpr.args, :($substitute_defaults($condition)))
588+
elseif condition isa Expr
589+
push!(ifexpr.args, morph_with_default!(condition))
590+
else
591+
@info "Don't know what to do with $(typeof(condition))"
592+
end =#
593+
push!(ifexpr.args, condition)
594+
comp_names, comps, expr_vec, varexpr = _parse_components!(ifexpr, x, kwargs)
595+
push_conditional_component!(ifexpr, expr_vec, comp_names, varexpr)
596+
comps
597+
end
598+
599+
function handle_if_y!(exprs, ifexpr, y, kwargs)
600+
Base.remove_linenums!(y)
601+
if Meta.isexpr(y, :elseif)
602+
comps = [:elseif, y.args[1]]
603+
elseifexpr = Expr(:elseif)
604+
push!(comps, handle_if_x!(mod, exprs, elseifexpr, y.args[2], kwargs, y.args[1]))
605+
get(y.args, 3, nothing) !== nothing && push!(comps, handle_if_y!(exprs, elseifexpr, y.args[3], kwargs))
606+
push!(ifexpr.args, elseifexpr)
607+
(comps...,)
608+
else
609+
comp_names, comps, expr_vec, varexpr, = _parse_components!(exprs, y, kwargs)
610+
push_conditional_component!(ifexpr, expr_vec, comp_names, varexpr)
611+
comps
612+
end
613+
end
614+
615+
function parse_components!(mod, exprs, cs, dict, compbody, kwargs)
616+
dict[:components] = []
617+
Base.remove_linenums!(compbody)
618+
for arg in compbody.args
619+
MLStyle.@match arg begin
620+
Expr(:if, condition, x) => begin
621+
ifexpr = Expr(:if)
622+
comps = handle_if_x!(mod, exprs, ifexpr, x, kwargs, condition)
623+
push!(exprs, ifexpr)
624+
push!(dict[:components], (:if, condition, comps, []))
625+
end
626+
Expr(:if, condition, x, y) => begin
627+
ifexpr = Expr(:if)
628+
comps = handle_if_x!(mod, exprs, ifexpr, x, kwargs, condition)
629+
ycomps = handle_if_y!(exprs, ifexpr, y, kwargs)
630+
push!(exprs, ifexpr)
631+
push!(dict[:components], (:if, condition, comps, ycomps))
632+
end
633+
Expr(:(=), a, b) => begin
634+
comp_names, comps, expr_vec, varexpr = _parse_components!(exprs, :(begin $arg end), kwargs)
635+
push!(cs, comp_names...)
636+
push!(dict[:components], comps...)
637+
push!(exprs, varexpr, :(@named begin $(expr_vec.args...) end))
638+
end
639+
_ => @info "410 Couldn't parse the component body $arg"
640+
end
641+
end
642+
643+
### zzz
644+
# push!(exprs, :(@named $expr))
645+
end
646+
647+
function _rename(compname, varname)
648+
compname = Symbol(compname, :__, varname)
649+
end

0 commit comments

Comments
 (0)