feat: support for forge-std assert cheatcodes

src/halmos/__main__.py

@@ -762,6 +762,7 @@ def future_callback(future_model):
         ) or is_global_fail_set(ex.context):
             if args.verbose >= 1:
                 print(f"Found potential path (id: {idx+1})")
+                print(f"{ex.context.output.error}")
 
             if args.verbose >= VERBOSITY_TRACE_COUNTEREXAMPLE:
                 traces[idx] = rendered_trace(ex.context)

src/halmos/assertions.py

@@ -0,0 +1,245 @@
+# SPDX-License-Identifier: AGPL-3.0
+
+from dataclasses import dataclass
+from z3 import *
+
+from .utils import *
+
+
+@dataclass(frozen=True)
+class VmAssertion:
+    """
+    Forge Standard Assertions
+    """
+
+    cond: BitVecRef
+    msg: Optional
+
+
+def mk_cond(bop, v1, v2):
+    v1 = bytes_to_bv_value(v1) if isinstance(v1, bytes) else v1
+    v2 = bytes_to_bv_value(v2) if isinstance(v2, bytes) else v2
+
+    if not is_bv(v1):
+        raise ValueError(f"mk_cond: not bv: {v1}")
+    if not is_bv(v2):
+        raise ValueError(f"mk_cond: not bv: {v2}")
+
+    # for Eq and NotEq, the bitsize can be arbitrary, e.g., arrays
+
+    if v1.size() != v2.size():
+        if bop == "Eq":
+            return BoolVal(False)
+        elif bop == "NotEq":
+            return BoolVal(True)
+        else:
+            raise ValueError(f"mk_cond: incompatible size: {v1} {bop} {v2}")
+
+    if bop == "Eq":
+        return v1 == v2
+    elif bop == "NotEq":
+        return v1 != v2
+
+    # for comparison operators, the bitsize must be 256-bit
+
+    if v1.size() != 256 or v2.size() != 256:
+        raise ValueError(f"mk_cond: incompatible size: {v1} {bop} {v2}")
+
+    if bop == "ULt":
+        return ULT(v1, v2)
+    elif bop == "UGt":
+        return UGT(v1, v2)
+    elif bop == "ULe":
+        return ULE(v1, v2)
+    elif bop == "UGe":
+        return UGE(v1, v2)
+    elif bop == "SLt":
+        return v1 < v2
+    elif bop == "SGt":
+        return v1 > v2
+    elif bop == "SLe":
+        return v1 <= v2
+    elif bop == "SGe":
+        return v1 >= v2
+    else:
+        raise ValueError(f"mk_cond: unknown bop: {bop}")
+
+
+def vm_assert_binary(bop: str, typ: str, log: bool = False):
+    arr = typ.endswith("[]")
+    typ = typ.replace("[]", "")
+
+    is_bytes = typ in ["bytes", "string"]
+
+    if not arr:
+        # bool, uint256, int256, address, bytes32
+        if not is_bytes:
+
+            def _f(arg):
+                v1 = extract_bytes(arg, 4, 32)
+                v2 = extract_bytes(arg, 36, 32)
+                cond = mk_cond(bop, v1, v2)
+                msg = extract_string_argument(arg, 2) if log else None
+                return VmAssertion(cond, msg)
+
+            return _f
+
+        # bytes, string
+        else:
+
+            def _f(arg):
+                v1 = extract_bytes_argument(arg, 0)
+                v2 = extract_bytes_argument(arg, 1)
+                cond = mk_cond(bop, v1, v2)
+                msg = extract_string_argument(arg, 2) if log else None
+                return VmAssertion(cond, msg)
+
+            return _f
+
+    else:
+        # bool[], uint256[], int256[], address[], bytes32[]
+        if not is_bytes:
+
+            def _f(arg):
+                v1 = extract_bytes32_array_argument(arg, 0)
+                v2 = extract_bytes32_array_argument(arg, 1)
+                cond = mk_cond(bop, v1, v2)
+                msg = extract_string_argument(arg, 2) if log else None
+                return VmAssertion(cond, msg)
+
+            return _f
+
+        # bytes[], string[]
+        else:
+
+            def _f(arg):
+                # TODO: implement extract_bytes_array
+                raise NotImplementedError(f"assert {bop} {typ}[]")
+
+            return _f
+
+
+def vm_assert_unary(expected: bool, log: bool = False):
+    def _f(arg):
+        actual = uint256(arg.get_word(4))
+        cond = test(actual, expected)
+        msg = extract_string_argument(arg, 1) if log else None
+        return VmAssertion(cond, msg)
+
+    return _f
+
+
+def mk_assert_handler(signature):
+    # pattern: assert<operator>(<params>)
+    match = re.search(r"assert([^(]+)\(([^)]+)\)", signature)
+    if match:
+        operator = match.group(1)
+        params = match.group(2).split(",")
+    else:
+        raise HalmosException(f"not supported signatures: {signature}")
+
+    # operator type:
+    # - binary: compare two arguments
+    # - unary: check given condition
+    is_binary = operator not in ["True", "False"]
+
+    # whether it includes log message or not
+    has_log = len(params) > (2 if is_binary else 1)
+
+    if is_binary:
+        typ = params[0]  # params[0] == params[1]
+        if operator in ["Eq", "NotEq"]:
+            bop = operator
+        else:
+            # for comparison operators, we need to identify whether they are unsigned or signed
+            sign = "U" if typ == "uint256" else "S"
+            bop = sign + operator
+        return vm_assert_binary(bop, typ, has_log)
+    else:
+        return vm_assert_unary(operator == "True", has_log)
+
+
+assert_cheatcode_handler = {
+    # assertTrue/False
+    0x0C9FD581: mk_assert_handler("assertTrue(bool)"),
+    0xA34EDC03: mk_assert_handler("assertTrue(bool,string)"),
+    0xA5982885: mk_assert_handler("assertFalse(bool)"),
+    0x7BA04809: mk_assert_handler("assertFalse(bool,string)"),
+    # assertEq(T, T)
+    0xF7FE3477: mk_assert_handler("assertEq(bool,bool)"),
+    0x4DB19E7E: mk_assert_handler("assertEq(bool,bool,string)"),
+    0x98296C54: mk_assert_handler("assertEq(uint256,uint256)"),
+    0x88B44C85: mk_assert_handler("assertEq(uint256,uint256,string)"),
+    0xFE74F05B: mk_assert_handler("assertEq(int256,int256)"),
+    0x714A2F13: mk_assert_handler("assertEq(int256,int256,string)"),
+    0x515361F6: mk_assert_handler("assertEq(address,address)"),
+    0x2F2769D1: mk_assert_handler("assertEq(address,address,string)"),
+    0x7C84C69B: mk_assert_handler("assertEq(bytes32,bytes32)"),
+    0xC1FA1ED0: mk_assert_handler("assertEq(bytes32,bytes32,string)"),
+    0xF320D963: mk_assert_handler("assertEq(string,string)"),
+    0x36F656D8: mk_assert_handler("assertEq(string,string,string)"),
+    0x97624631: mk_assert_handler("assertEq(bytes,bytes)"),
+    0xE24FED00: mk_assert_handler("assertEq(bytes,bytes,string)"),
+    # assertEq(T[], T[])
+    0x707DF785: mk_assert_handler("assertEq(bool[],bool[])"),
+    0xE48A8F8D: mk_assert_handler("assertEq(bool[],bool[],string)"),
+    0x975D5A12: mk_assert_handler("assertEq(uint256[],uint256[])"),
+    0x5D18C73A: mk_assert_handler("assertEq(uint256[],uint256[],string)"),
+    0x711043AC: mk_assert_handler("assertEq(int256[],int256[])"),
+    0x191F1B30: mk_assert_handler("assertEq(int256[],int256[],string)"),
+    0x3868AC34: mk_assert_handler("assertEq(address[],address[])"),
+    0x3E9173C5: mk_assert_handler("assertEq(address[],address[],string)"),
+    0x0CC9EE84: mk_assert_handler("assertEq(bytes32[],bytes32[])"),
+    0xE03E9177: mk_assert_handler("assertEq(bytes32[],bytes32[],string)"),
+    0xCF1C049C: mk_assert_handler("assertEq(string[],string[])"),
+    0xEFF6B27D: mk_assert_handler("assertEq(string[],string[],string)"),
+    0xE5FB9B4A: mk_assert_handler("assertEq(bytes[],bytes[])"),
+    0xF413F0B6: mk_assert_handler("assertEq(bytes[],bytes[],string)"),
+    # assertNotEq(T, T)
+    0x236E4D66: mk_assert_handler("assertNotEq(bool,bool)"),
+    0x1091A261: mk_assert_handler("assertNotEq(bool,bool,string)"),
+    0xB7909320: mk_assert_handler("assertNotEq(uint256,uint256)"),
+    0x98F9BDBD: mk_assert_handler("assertNotEq(uint256,uint256,string)"),
+    0xF4C004E3: mk_assert_handler("assertNotEq(int256,int256)"),
+    0x4724C5B9: mk_assert_handler("assertNotEq(int256,int256,string)"),
+    0xB12E1694: mk_assert_handler("assertNotEq(address,address)"),
+    0x8775A591: mk_assert_handler("assertNotEq(address,address,string)"),
+    0x898E83FC: mk_assert_handler("assertNotEq(bytes32,bytes32)"),
+    0xB2332F51: mk_assert_handler("assertNotEq(bytes32,bytes32,string)"),
+    0x6A8237B3: mk_assert_handler("assertNotEq(string,string)"),
+    0x78BDCEA7: mk_assert_handler("assertNotEq(string,string,string)"),
+    0x3CF78E28: mk_assert_handler("assertNotEq(bytes,bytes)"),
+    0x9507540E: mk_assert_handler("assertNotEq(bytes,bytes,string)"),
+    # assertNotEq(T[], T[])
+    0x286FAFEA: mk_assert_handler("assertNotEq(bool[],bool[])"),
+    0x62C6F9FB: mk_assert_handler("assertNotEq(bool[],bool[],string)"),
+    0x56F29CBA: mk_assert_handler("assertNotEq(uint256[],uint256[])"),
+    0x9A7FBD8F: mk_assert_handler("assertNotEq(uint256[],uint256[],string)"),
+    0x0B72F4EF: mk_assert_handler("assertNotEq(int256[],int256[])"),
+    0xD3977322: mk_assert_handler("assertNotEq(int256[],int256[],string)"),
+    0x46D0B252: mk_assert_handler("assertNotEq(address[],address[])"),
+    0x72C7E0B5: mk_assert_handler("assertNotEq(address[],address[],string)"),
+    0x0603EA68: mk_assert_handler("assertNotEq(bytes32[],bytes32[])"),
+    0xB873634C: mk_assert_handler("assertNotEq(bytes32[],bytes32[],string)"),
+    0xBDFACBE8: mk_assert_handler("assertNotEq(string[],string[])"),
+    0xB67187F3: mk_assert_handler("assertNotEq(string[],string[],string)"),
+    0xEDECD035: mk_assert_handler("assertNotEq(bytes[],bytes[])"),
+    0x1DCD1F68: mk_assert_handler("assertNotEq(bytes[],bytes[],string)"),
+    # assertLt/Gt/Le/Ge
+    0xB12FC005: mk_assert_handler("assertLt(uint256,uint256)"),
+    0x65D5C135: mk_assert_handler("assertLt(uint256,uint256,string)"),
+    0x3E914080: mk_assert_handler("assertLt(int256,int256)"),
+    0x9FF531E3: mk_assert_handler("assertLt(int256,int256,string)"),
+    0xDB07FCD2: mk_assert_handler("assertGt(uint256,uint256)"),
+    0xD9A3C4D2: mk_assert_handler("assertGt(uint256,uint256,string)"),
+    0x5A362D45: mk_assert_handler("assertGt(int256,int256)"),
+    0xF8D33B9B: mk_assert_handler("assertGt(int256,int256,string)"),
+    0x8466F415: mk_assert_handler("assertLe(uint256,uint256)"),
+    0xD17D4B0D: mk_assert_handler("assertLe(uint256,uint256,string)"),
+    0x95FD154E: mk_assert_handler("assertLe(int256,int256)"),
+    0x4DFE692C: mk_assert_handler("assertLe(int256,int256,string)"),
+    0xA8D4D1D9: mk_assert_handler("assertGe(uint256,uint256)"),
+    0xE25242C0: mk_assert_handler("assertGe(uint256,uint256,string)"),
+    0x0A30B771: mk_assert_handler("assertGe(int256,int256)"),
+    0xA84328DD: mk_assert_handler("assertGe(int256,int256,string)"),
+}

src/halmos/cheatcodes.py

@@ -11,6 +11,7 @@
 from .bytevec import ByteVec
 from .exceptions import FailCheatcode, HalmosException, InfeasiblePath
 from .utils import *
+from .assertions import *
 
 
 # f_vmaddr(key) -> address
@@ -330,12 +331,23 @@ class hevm_cheat_code:
     label_sig: int = 0xC657C718
 
     @staticmethod
-    def handle(sevm, ex, arg: ByteVec) -> Optional[ByteVec]:
+    def handle(sevm, ex, arg: ByteVec, stack, step_id) -> Optional[ByteVec]:
         funsig: int = int_of(arg[:4].unwrap(), "symbolic hevm cheatcode")
         ret = ByteVec()
 
+        if funsig in assert_cheatcode_handler:
+            vm_assert = assert_cheatcode_handler[funsig](arg)
+            not_cond = simplify(Not(vm_assert.cond))
+
+            if ex.check(not_cond) != unsat:
+                new_ex = sevm.create_branch(ex, not_cond, ex.pc)
+                new_ex.halt(data=ByteVec(), error=FailCheatcode(f"{vm_assert}"))
+                stack.push(new_ex, step_id)
+
+            return ret
+
         # vm.assume(bool)
-        if funsig == hevm_cheat_code.assume_sig:
+        elif funsig == hevm_cheat_code.assume_sig:
             assume_cond = simplify(is_non_zero(arg.get_word(4)))
             if is_false(assume_cond):
                 raise InfeasiblePath("vm.assume(false)")

src/halmos/sevm.py

@@ -698,6 +698,9 @@ def halt(
         output.error = error
         output.return_scheme = self.current_opcode()
 
+    def is_halted(self) -> bool:
+        return self.context.output.data is not None
+
     def emit_log(self, log: EventLog):
         self.context.trace.append(log)
 
@@ -1741,7 +1744,7 @@ def call_unknown() -> None:
             # vm cheat code
             elif eq(to, hevm_cheat_code.address):
                 exit_code = con(1)
-                ret = hevm_cheat_code.handle(self, ex, arg)
+                ret = hevm_cheat_code.handle(self, ex, arg, stack, step_id)
 
             # console
             elif eq(to, console.address):
@@ -2143,6 +2146,11 @@ def finalize(ex: Exec):
                 if not ex.path.is_activated():
                     ex.path.activate()
 
+                # PathEndingException may not be immediately raised; it could be delayed until it comes out of the worklist
+                # see the assert cheatcode hanlder logic for the delayed case
+                if isinstance(ex.context.output.error, PathEndingException):
+                    raise ex.context.output.error
+
                 if ex.context.depth > MAX_CALL_DEPTH:
                     raise MessageDepthLimitError(ex.context)
 
@@ -2545,8 +2553,9 @@ def finalize(ex: Exec):
                 continue
 
             except FailCheatcode as err:
-                # return data shouldn't be None, as it is considered being stuck
-                ex.halt(data=ByteVec(), error=err)
+                if not ex.is_halted():
+                    # return data shouldn't be None, as it is considered being stuck
+                    ex.halt(data=ByteVec(), error=err)
                 yield ex  # early exit; do not call finalize()
                 continue
 

src/halmos/utils.py

@@ -190,6 +190,22 @@ def create_solver(logic="QF_AUFBV", ctx=None, timeout=0, max_memory=0):
     return solver
 
 
+def extract_bytes32_array_argument(calldata: BitVecRef, arg_idx: int):
+    """Extracts idx-th argument of bytes32[] from calldata"""
+    offset = int_of(
+        extract_bytes(calldata, 4 + arg_idx * 32, 32),
+        "symbolic offset for bytes argument",
+    )
+    length = int_of(
+        extract_bytes(calldata, 4 + offset, 32),
+        "symbolic size for bytes argument",
+    )
+    if length == 0:
+        return b""
+
+    return extract_bytes(calldata, 4 + offset + 32, length * 32)
+
+
 def extract_bytes_argument(calldata: BitVecRef, arg_idx: int) -> bytes:
     """Extracts idx-th argument of string from calldata"""
     offset = int_of(