copy bitvecfunc logic to generic handlers

mythril/laser/smt/bitvec.py

@@ -6,7 +6,7 @@ import z3
 
 from mythril.laser.smt.bool import Bool, And, Or
 from mythril.laser.smt.expression import Expression
-
+from itertools import product
 Annotations = Set[Any]
 
 # fmt: off
@@ -292,8 +292,33 @@ def _comparison_helper(
         ):
             return Bool(z3.BoolVal(default_value), annotations=annotations)
 
+        condition = True
+        for a_nest, b_nest in product(a.nested_functions, b.nested_functions):
+            if a_nest.func_name != b_nest.func_name:
+                continue
+            if a_nest.func_name == "Hybrid":
+                continue
+            # a.input (eq/neq) b.input ==> a == b
+            if inputs_equal:
+                condition = z3.And(
+                    condition,
+                    z3.Or(
+                        z3.Not((a_nest.input_ == b_nest.input_).raw),
+                        (a_nest.raw == b_nest.raw),
+                    ),
+                )
+            else:
+                condition = z3.And(
+                    condition,
+                    z3.Or(
+                        z3.Not((a_nest.input_ != b_nest.input_).raw),
+                        (a_nest.raw == b_nest.raw),
+                    ),
+                )
+
         return And(
-            Bool(operation(a.raw, b.raw), annotations=annotations),
+            Bool(cast(z3.BoolRef, operation(a.raw, b.raw)), annotations=annotations),
+            Bool(condition) if b.nested_functions else Bool(True),
             a.input_ == b.input_ if inputs_equal else a.input_ != b.input_,
         )
 
@@ -307,8 +332,20 @@ def _arithmetic_helper(a: BitVec, b: BitVec, operation: Callable) -> BitVec:
     if isinstance(a, BitVecFunc) and isinstance(b, BitVecFunc):
         return BitVecFunc(raw=raw, func_name=None, input_=None, annotations=union)
     elif isinstance(a, BitVecFunc):
+        if isinstance(b, BitVecFunc):
+            return BitVecFunc(
+                raw=raw,
+                func_name="Hybrid",
+                input_=BitVec(z3.BitVec("", 256), annotations=union),
+                nested_functions=a.nested_functions + b.nested_functions + [a, b],
+            )
+
         return BitVecFunc(
-            raw=raw, func_name=a.func_name, input_=a.input_, annotations=union
+            raw=raw,
+            func_name=a.func_name,
+            input_=a.input_,
+            annotations=union,
+            nested_functions=a.nested_functions + [a],
         )
     elif isinstance(b, BitVecFunc):
         return BitVecFunc(
@@ -405,18 +442,17 @@ def Concat(*args: Union[BitVec, List[BitVec]]) -> BitVec:
     nraw = z3.Concat([a.raw for a in bvs])
     annotations = set()  # type: Annotations
     bitvecfunc = False
+    nested_functions = []
     for bv in bvs:
         annotations = annotations.union(bv.annotations)
         if isinstance(bv, BitVecFunc):
             bitvecfunc = True
+            nested_functions += bv.nested_functions
 
     if bitvecfunc:
-        # Added this not so good and misleading NOTATION to help with this
-        str_hash = ",".join(["hashed({})".format(hash(bv)) for bv in bvs])
-        input_string = "MisleadingNotationConcat({})".format(str_hash)
-
         return BitVecFunc(
-            raw=nraw, func_name="Hybrid", input_=BitVec(z3.BitVec(input_string, 256), annotations=annotations)
+            raw=nraw, func_name="Hybrid", input_=BitVec(z3.BitVec("", 256), annotations=annotations),
+            nested_functions=nested_functions
         )
 
     return BitVec(nraw, annotations)
@@ -432,10 +468,11 @@ def Extract(high: int, low: int, bv: BitVec) -> BitVec:
     """
     raw = z3.Extract(high, low, bv.raw)
     if isinstance(bv, BitVecFunc):
-        input_string = "MisleadingNotationExtract({}, {}, hashed({}))".format(high, low, hash(bv))
+        input_string = ""
         # Is there a better value to set func_name and input to in this case?
         return BitVecFunc(
-            raw=raw, func_name="Hybrid", input_=BitVec(z3.BitVec(input_string, 256), annotations=bv.annotations)
+            raw=raw, func_name="Hybrid", input_=BitVec(z3.BitVec(input_string, 256), annotations=bv.annotations),
+            nested_functions=bv.nested_functions
         )
 
     return BitVec(raw, annotations=bv.annotations)
@@ -486,13 +523,15 @@ def Sum(*args: BitVec) -> BitVec:
             bitvecfuncs.append(bv)
 
     if len(bitvecfuncs) >= 2:
-        return BitVecFunc(raw=raw, func_name="Hybrid", input_=None, annotations=annotations)
+        return BitVecFunc(raw=raw, func_name="Hybrid", input_=None, annotations=annotations,
+                          nested_functions=[nf for func in bitvecfuncs for nf in func.nested_functions])
     elif len(bitvecfuncs) == 1:
         return BitVecFunc(
             raw=raw,
             func_name=bitvecfuncs[0].func_name,
             input_=bitvecfuncs[0].input_,
             annotations=annotations,
+            nested_functions=[nf for func in bitvecfuncs for nf in func.nested_functions]
         )
 
     return BitVec(raw, annotations)