Break redundancy in simplification into helper functions

Author: HarrisonGrodin

src/simplify.jl

@@ -25,14 +25,15 @@ function _simplify_constants(O,shorten_tree = true)
               args = Vector{Expression}[O.args[i].args for i in idxs]
               push!(args,O.args[keep_idxs])
               return Operation(O.op,vcat(args...))
+          end
           # Collapse constants
-          elseif length(findall(x->typeof(x)<:Variable && x.subtype == :Constant ,O.args)) > 1
-              idxs = findall(x->typeof(x)<:Variable && x.subtype == :Constant ,O.args)
+          idxs = findall(is_constant, O.args)
+          if length(idxs) > 1
               other_idxs = 1:length(O.args) .∉ (idxs,)
               if cur_op == :*
-                new_var = Constant(prod(x->x.value,O.args[idxs]))
+                new_var = Constant(prod(get, O.args[idxs]))
               elseif cur_op == :+
-                new_var = Constant(sum(x->x.value,O.args[idxs]))
+                new_var = Constant(sum(get, O.args[idxs]))
               end
               new_args = O.args[other_idxs]
               push!(new_args,new_var)
@@ -50,10 +51,10 @@ function _simplify_constants(O,shorten_tree = true)
         # If any variable is `Constant(0)`, zero the whole thing
         # If any variable is `Constant(1)`, remove that `Constant(1)` unless
         # they are all `Constant(1)`, in which case simplify to a single variable
-        if any(x->typeof(x)<:Variable && (isequal(x,Constant(0)) || isequal(x,Constant(-0))),O.args)
+        if any(iszero, O.args)
             return Constant(0)
-        elseif any(x->typeof(x)<:Variable && isequal(x,Constant(1)),O.args)
-            idxs = findall(x->typeof(x)<:Variable && isequal(x,Constant(1)),O.args)
+        elseif any(isone, O.args)
+            idxs = findall(isone, O.args)
             _O = Operation(O.op,O.args[1:length(O.args) .∉ (idxs,)])
             if isempty(_O.args)
                 return Constant(1)
@@ -65,10 +66,9 @@ function _simplify_constants(O,shorten_tree = true)
         else
             return O
         end
-    elseif Symbol(O.op) == :+ && any(x->typeof(x)<:Variable &&
-           (isequal(x,Constant(0)) || isequal(x,Constant(-0))),O.args)
+    elseif Symbol(O.op) == :+ && any(iszero, O.args)
         # If there are Constant(0)s in a big `+` expression, get rid of them
-        idxs = findall(x->typeof(x)<:Variable && (isequal(x,Constant(0)) || isequal(x,Constant(-0))),O.args)
+        idxs = findall(iszero, O.args)
         _O = Operation(O.op,O.args[1:length(O.args) .∉ (idxs,)])
         if isempty(_O.args)
             return Constant(0)

src/utils.jl

@@ -18,3 +18,6 @@ function flatten_expr!(x)
 end
 
 toexpr(ex) = MacroTools.postwalk(x->x isa Union{Expression,Operation} ? Expr(x) : x, ex)
+
+is_constant(x::Variable) = x.subtype === :Constant
+is_constant(::Any) = false

src/variables.jl

@@ -63,6 +63,15 @@ end
 export Variable,Parameter,Constant,DependentVariable,IndependentVariable,JumpVariable,NoiseVariable,
        @Var, @DVar, @IVar, @Param, @Const
 
+
+Base.get(x::Variable) = x.value
+
+Base.iszero(::Expression) = false
+Base.iszero(c::Variable) = get(c) isa Number && iszero(get(c))
+Base.isone(::Expression) = false
+Base.isone(c::Variable) = get(c) isa Number && isone(get(c))
+
+
 # Variables use isequal for equality since == is an Operation
 function Base.:(==)(x::Variable,y::Variable)
     x.name == y.name && x.subtype == y.subtype && x.value == y.value &&