Merge pull request #135 from neonWhiteout/ternary_op

Add ternary operation in stockflow macro syntax

src/StockFlow.jl

@@ -39,6 +39,10 @@ vectorify(n) = [n]
 
 state_dict(n) = Dict(s=>i for (i, s) in enumerate(n))
 
+cond(c,t,f) = c ? t : f
+or(a,b) = Bool(a) || Bool(b)
+and(a,b) = Bool(a) && Bool(b)
+
 #= operators definition
 # Operators:
 
@@ -66,7 +70,7 @@ https://docs.julialang.org/en/v1/manual/mathematical-operations/
 Operators = Dict(2 => [:+, :-, :*, :/, :÷, :^, :%, :log, Symbol("")],
                  1 => [:log, :exp, :sqrt, Symbol("")]) #:NN is for NONE, which is used in create the special diagram using graph rewriting
 
-
+math_expr(op, op1, op2, op3) = Expr(:call, op, op1, op2, op3)
 math_expr(op, op1, op2) = Expr(:call, op, op1, op2)
 math_expr(op, op1) = Expr(:call, op, op1)
 #math_expr(op) = Expr(:call, op)

src/Syntax.jl

@@ -194,8 +194,13 @@ end
 struct Ref <: DyvarExprT
     ref::Pair{Symbol,Symbol}
 end
+struct Cond{C,T,F} <: DyvarExprT
+    cond::Tuple{C,T,F}
+end
+
 get(r::Ref) = r.ref
 get(b::Binop) = b.binop
+get(c::Cond) = (c.cond => :cond)
 struct StockAndFlowArguments
     stocks::Vector{
         Pair{
@@ -353,7 +358,7 @@ function is_recursive_dyvar(dyvar_name, dyvar_def)
     @match dyvar_def begin
         ::Symbol => dyvar_def == dyvar_name
         :($f()) => f == dyvar_name
-        Expr(:call, args...) => true in map(arg -> is_recursive_dyvar(dyvar_name, arg), args)
+        Expr(:call, args...) || Expr(:if, args...) || Expr(:||, args...) || Expr(:&&, args...)  => true in map(arg -> is_recursive_dyvar(dyvar_name, arg), args)
     end
 end
 """
@@ -376,6 +381,24 @@ function parse_dyvar!(dyvars::Vector{Tuple{Symbol,Expr}}, dyvar::Expr)
 
  function parse_dyvar(dyvar::Expr)
      @match dyvar begin
+         :($dyvar_name = $(cond::Union{Symbol, Expr}) ? $(true_val::Union{Symbol, Expr}) : $(false_val::Union{Symbol, Expr})) =>
+             if !is_recursive_dyvar(dyvar_name, cond) && !is_recursive_dyvar(dyvar_name, true_val) && !is_recursive_dyvar(dyvar_name, false_val)
+                 return (dyvar_name, :($cond ? $true_val : $false_val))
+             else
+                 error("Recursive dyvar detected in " * String(dyvar_name))
+             end
+         :($dyvar_name = $A || $B) => 
+             if !is_recursive_dyvar(dyvar_name, A) && !is_recursive_dyvar(dyvar_name, B)
+                 return (dyvar_name, :(or($A,$B)))
+             else
+                error("Recursive dyvar detected in " * String(dyvar_name))
+             end
+         :($dyvar_name = $A && $B) =>
+             if !is_recursive_dyvar(dyvar_name, A) && !is_recursive_dyvar(dyvar_name, B)
+                 return (dyvar_name, :(and($A,$B)))
+             else
+                error("Recursive dyvar detected in " * String(dyvar_name))
+             end
          :($dyvar_name = $dyvar_def) =>
              if !is_recursive_dyvar(dyvar_name, dyvar_def)
                  return (dyvar_name, dyvar_def)
@@ -588,6 +611,9 @@ function dyvar_exprs_to_symbolic_repr(dyvars::Vector{Tuple{Symbol,Expr}})
                 Expr(:call, op, a, b) => begin
                     push!(syms, (dyvar_name => Binop((a, b) => op)))
                 end
+                Expr(:if, c, t, f) => begin
+                    push!(syms, (dyvar_name => Cond((c, t, f),)))
+                end
                 Expr(c, _, _) || Expr(c, _, _, _) => error(
                     "Unhandled expression in dynamic variable definition " * String(c),
                 )
@@ -816,6 +842,8 @@ function infix_expression_to_binops(
     function loop(e)
         @match e begin
             ::Symbol || ::Float32 || ::Float64 || ::Int || ::String => e
+            Expr(:||, args...) => loop(Expr(:call, :or, args...))
+            Expr(:&&, args...) => loop(Expr(:call, :and, args...))
             Expr(:call, f, a) => begin
                 asym = loop(a)
                 varname = gensym(gensymbase)
@@ -846,9 +874,17 @@ function infix_expression_to_binops(
                 end
                 lastsym
             end
+            Expr(:if, c, t, f) => begin
+               csym = loop(c)
+               tsym = loop(t)
+               fsym = loop(f)
+               varname = gensym(gensymbase)
+               push!(exprs, (varname, :($csym ? $tsym : $fsym)))
+               varname
+            end
             Expr(en, _, _, _) || Expr(en, _, _) => begin
                 error(
-                    "Unhandled expression type " * String(en) * " cannot be converted into form f(a, b)",
+                    "Unhandled expression type " * String(en) * " cannot be converted into form f(a, b)"
                 )
             end
         end
@@ -887,7 +923,7 @@ Check if a Julia expression is a call of the form `op(a, b)` or `a op b`, where
 - `e` -- a Julia expression
 
 ### Output
-A boolean indicating if the given julia expression is a function call of non-expression parameter(s).
+A boolean indicating if the given julia expression is a function call or if statement of non-expression parameter(s).
 
 ### Examples
 ```julia-repl
@@ -903,12 +939,17 @@ julia> is_simple_dyvar(:(f(a, b, c)))
 false
 julia> is_simple_dyvar(:f(a + b, c + d))
 false
+julia> is_simple_dyvar(:(A ? B : C))
+true
+julia> is_simple_dyvar(:(A < B ? X : Y))
+false
 ```
 """
 function is_simple_dyvar(e::Expr)
     @match e begin
         Expr(:call, f::Symbol, a) => !(typeof(a) <: Expr)
         Expr(:call, f::Symbol, a, b) => !(typeof(a) <: Expr) && !(typeof(b) <: Expr)
+        Expr(:if, c, t, f) => !(typeof(c) <: Expr) && !(typeof(t) <: Expr) && !(typeof(f) <: Expr)
         _ => false
     end
 end

test/Syntax.jl

@@ -597,3 +597,17 @@ end
     end) == CausalLoopPM([:A, :B], [:A => :B, :B => :A], [POL_NEGATIVE, POL_POSITIVE])
 
 end
+
+
+
+@testset "Ternary" begin
+        @test ((@stock_and_flow begin
+                       :stocks
+                       A
+                       B
+                       C
+                       :dynamic_variables
+                       cond = A == B
+                       v = cond ? A : B
+              end) == StockAndFlowF([:A => (:F_NONE, :F_NONE, :SV_NONE), :B => (:F_NONE, :F_NONE, :SV_NONE), :C => (:F_NONE, :F_NONE, :SV_NONE)], [], [:cond => ((:A, :B) => :(==)), :v => ((:cond, :A, :B) => :cond)], [], []))
+end