Clarify control flow

Author: HarrisonGrodin

src/simplify.jl

@@ -1,4 +1,4 @@
-function simplify_constants(O::Operation,shorten_tree = true)
+function simplify_constants(O::Operation, shorten_tree = true)
     O_last = nothing
     _O = O
     while _O != O_last
@@ -11,79 +11,62 @@ function simplify_constants(O::Operation,shorten_tree = true)
     _O
 end
 
-const TREE_SHRINK_OPS = [*, +]
+const AC_OPERATORS = [*, +]
 
-function _simplify_constants(O,shorten_tree = true)
+function _simplify_constants(O, shorten_tree = true)
     # Tree shrinking
-    if shorten_tree
-        for cur_op in TREE_SHRINK_OPS
-            if O.op == cur_op
-                # Shrink tree
-                if any(x -> is_operation(x) && x.op == cur_op, O.args)
-                    idxs = findall(x -> is_operation(x) && x.op == cur_op, O.args)
-                    keep_idxs = 1:length(O.args) .∉ (idxs,)
-                    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
-                idxs = findall(is_constant, O.args)
-                if length(idxs) > 1
-                    other_idxs = 1:length(O.args) .∉ (idxs,)
-                    if cur_op == (*)
-                        new_var = Constant(prod(get, O.args[idxs]))
-                    elseif cur_op == (+)
-                        new_var = Constant(sum(get, O.args[idxs]))
-                    end
-                    new_args = O.args[other_idxs]
-                    push!(new_args,new_var)
-                    if length(new_args) > 1
-                        return Operation(O.op,new_args)
-                    else
-                        return new_args[1]
-                    end
-                end
-            end
+    if shorten_tree && O.op ∈ AC_OPERATORS
+        # Flatten tree
+        idxs = findall(x -> is_operation(x) && x.op == O.op, O.args)
+        if !isempty(idxs)
+            keep_idxs = eachindex(O.args) .∉ Ref(idxs)
+            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
+        idxs = findall(is_constant, O.args)
+        if length(idxs) > 1
+            other_idxs = eachindex(O.args) .∉ (idxs,)
+            new_var = Constant(mapreduce(get, O.op, O.args[idxs]))
+            new_args = O.args[other_idxs]
+            push!(new_args,new_var)
+
+            return length(new_args) > 1 ? Operation(O.op, new_args) : first(new_args)
         end
     end
 
-    if O.op == (*)
+    if O.op === (*)
         # If any variable is `Constant(0)`, zero the whole thing
+        any(iszero, O.args) && return Constant(0)
+
         # 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(iszero, O.args)
-            return Constant(0)
-        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)
-            elseif length(_O.args) == 1
-                return _O.args[1]
-            else
-                return _O
-            end
-        else
-            return O
-        end
-    elseif O.op == (+) && any(iszero, O.args)
-        # If there are Constant(0)s in a big `+` expression, get rid of them
-        idxs = findall(iszero, O.args)
-        _O = Operation(O.op,O.args[1:length(O.args) .∉ (idxs,)])
-        if isempty(_O.args)
-            return Constant(0)
-        elseif length(_O.args) == 1
-            return _O.args[1]
-        else
-            return O
+        if any(isone, O.args)
+            args = filter(!isone, O.args)
+
+            isempty(args)     && return Constant(1)
+            length(args) == 1 && return first(args)
+            return Operation(O.op, args)
         end
-    elseif O.op == identity
-        return O.args[1]
-    elseif O.op == (-) && length(O.args) == 1
-        return Operation(*,Expression[-1,O.args[1]])
-    else
+
         return O
     end
+
+    if O.op === (+) && any(iszero, O.args)
+        # If there are Constant(0)s in a big `+` expression, get rid of them
+        args = filter(!iszero, O.args)
+
+        isempty(args)     && return Constant(0)
+        length(args) == 1 && return first(args)
+        return Operation(O.op, args)
+    end
+
+    (O.op, length(O.args)) === (identity, 1) && return O.args[1]
+
+    (O.op, length(O.args)) === (-, 1) && return Operation(*, Expression[-1, O.args[1]])
+
     return O
 end
 simplify_constants(x::Variable,y=false) = x