Apply eveem-org#16

pano/vm.py

@@ -429,7 +429,7 @@ def handle_jumps(self, trace, line, condition):
             trace.append(("jump", n))
             return trace
 
-        if op == "jumpi":
+        elif op == "jumpi":
             target = stack.pop()
             if_condition = simplify_bool(stack.pop())
 
@@ -480,19 +480,7 @@ def handle_jumps(self, trace, line, condition):
             logger.debug("jumpi -> if %s", trace[-1])
             return trace
 
-        if op == "selfdestruct":
-            trace.append(("selfdestruct", stack.pop(),))
-            return trace
-
-        if op in ["stop", "assert_fail", "invalid"]:
-            trace.append((op,))
-            return trace
-
-        if op == "UNKNOWN":
-            trace.append(("invalid",))
-            return trace
-
-        if op in ["return", "revert"]:
+        elif op in ["return", "revert"]:
             p = stack.pop()
             n = stack.pop()
 
@@ -504,6 +492,18 @@ def handle_jumps(self, trace, line, condition):
 
             return trace
 
+        elif op == "selfdestruct":
+            trace.append(("selfdestruct", stack.pop(),))
+            return trace
+
+        elif op in ["stop", "assert_fail", "invalid"]:
+            trace.append((op,))
+            return trace
+
+        elif op == "UNKNOWN":
+            trace.append(("invalid",))
+            return trace
+
         return None
 
     def apply_stack(self, ret, line):
@@ -538,16 +538,6 @@ def trace(exp, *format_args):
                 else:
                     trace("[{}] {} {}", line[0], C.asm(op), C.asm(str(line[2])))
 
-        assert op not in [
-            "jump",
-            "jumpi",
-            "revert",
-            "return",
-            "stop",
-            "jumpdest",
-            "UNKNOWN",
-        ]
-
         param = 0
         if len(line) > 2:
             param = line[2]
@@ -569,32 +559,53 @@ def trace(exp, *format_args):
         ]:
             stack.append(arithmetic.eval((op, stack.pop(), stack.pop(),)))
 
-        if op in ["mulmod", "addmod"]:
-            stack.append(("mulmod", stack.pop(), stack.pop(), stack.pop()))
+        elif op[:4] == "push":
+            stack.append(param)
+
+        elif op == "pop":
+            stack.pop()
+
+        elif op == "dup":
+            stack.dup(param)
 
-        if op == "mul":
+        elif op == "mul":
             stack.append(mul_op(stack.pop(), stack.pop()))
 
-        if op == "or":
+        elif op == "or":
             stack.append(or_op(stack.pop(), stack.pop()))
 
-        if op == "shl":
+        elif op == "add":
+            stack.append(add_op(stack.pop(), stack.pop()))
+
+        elif op == "sub":
+            left = stack.pop()
+            right = stack.pop()
+
+            if type(left) == int and type(right) == int:
+                stack.append(arithmetic.sub(left, right))
+            else:
+                stack.append(sub_op(left, right))
+
+        elif op in ["mulmod", "addmod"]:
+            stack.append(("mulmod", stack.pop(), stack.pop(), stack.pop()))
+
+        elif op == "shl":
             off = stack.pop()
             exp = stack.pop()
             if all_concrete(off, exp):
                 stack.append(exp << off)
             else:
                 stack.append(mask_op(exp, shl=off))
 
-        if op == "shr":
+        elif op == "shr":
             off = stack.pop()
             exp = stack.pop()
             if all_concrete(off, exp):
                 stack.append(exp >> off)
             else:
                 stack.append(mask_op(exp, offset=minus_op(off), shr=off))
 
-        if op == "sar":
+        elif op == "sar":
             off = stack.pop()
             exp = stack.pop()
             if all_concrete(off, exp):
@@ -613,18 +624,6 @@ def trace(exp, *format_args):
                 # FIXME: This won't give the right result...
                 stack.append(mask_op(exp, offset=minus_op(off), shr=off))
 
-        if op == "add":
-            stack.append(add_op(stack.pop(), stack.pop()))
-
-        if op == "sub":
-            left = stack.pop()
-            right = stack.pop()
-
-            if type(left) == int and type(right) == int:
-                stack.append(arithmetic.sub(left, right))
-            else:
-                stack.append(sub_op(left, right))
-
         elif op in ["not", "iszero"]:
             stack.append((op, stack.pop()))
 
@@ -883,18 +882,9 @@ def trace(exp, *format_args):
 
             stack.append("create2.new_address")
 
-        elif op[:4] == "push":
-            stack.append(param)
-
         elif op == "pc":
             stack.append(line[0])
 
-        elif op == "pop":
-            stack.pop()
-
-        elif op == "dup":
-            stack.dup(param)
-
         elif op == "msize":
             self.counter += 1
             vname = f"_{self.counter}"
@@ -922,6 +912,18 @@ def trace(exp, *format_args):
         ]:
             stack.append(op)
 
+        else:
+            # TODO: Maybe raise an error directly?
+            assert op not in [
+                "jump",
+                "jumpi",
+                "revert",
+                "return",
+                "stop",
+                "jumpdest",
+                "UNKNOWN",
+            ]
+
         if stack.len() - previous_len != opcode_dict.stack_diffs[op]:
             logger.error("line: %s", line)
             logger.error("stack: %s", stack)