GenericDispatch and filter example

Author: philzook58

src/kdrag/contrib/pcode/asmspec.py

@@ -465,6 +465,24 @@ def execute_insn(
     ]
 
 
+def execute_spec_and_insn(
+    tracestate0: TraceState,
+    spec: AsmSpec,
+    ctx: pcode.BinaryContext,
+    verbose=True,
+) -> tuple[list[TraceState], list[VerificationCondition]]:
+    """
+    Execute spec statements and then one instruction.
+    """
+    addr = tracestate0.state.pc[0]
+    specstmts = spec.addrmap.get(addr, [])
+    tracestate, vcs = execute_spec_stmts(specstmts, tracestate0, ctx, verbose=verbose)
+    if tracestate is None:
+        return [], vcs
+    new_tracestates = execute_insn(tracestate, ctx, verbose=verbose)
+    return new_tracestates, vcs
+
+
 def init_trace_states(
     ctx: pcode.BinaryContext, mem: pcode.MemState, spec: AsmSpec, verbose=True
 ) -> tuple[list[TraceState], list[VerificationCondition]]:
@@ -538,16 +556,95 @@ def run_all_paths(
                 )
             )
             continue
-        addr = tracestate.state.pc[0]
-        specstmts = spec.addrmap.get(addr, [])
-        tracestate, new_vcs = execute_spec_stmts(specstmts, tracestate, ctx)
+        new_tracestates, new_vcs = execute_spec_and_insn(
+            tracestate, spec, ctx, verbose=verbose
+        )
         vcs.extend(new_vcs)
-        if tracestate is not None:
-            todo.extend(execute_insn(tracestate, ctx, verbose=verbose))
+        todo.extend(new_tracestates)
+
+        # addr = tracestate.state.pc[0]
+        # specstmts = spec.addrmap.get(addr, [])
+        # tracestate, new_vcs = execute_spec_stmts(specstmts, tracestate, ctx)
+        # vcs.extend(new_vcs)
+        # if tracestate is not None:
+        #    todo.extend(execute_insn(tracestate, ctx, verbose=verbose))
 
     return vcs
 
 
+class Debug:
+    def __init__(self, ctx: pcode.BinaryContext, spec: AsmSpec):
+        self.ctx = ctx
+        self.spec: AsmSpec = spec
+        self.tracestates: list[TraceState] = []
+        self.vcs: list[VerificationCondition] = []
+        self.breakpoints = set()
+
+    def spec_file(self, filename: str):
+        self.spec = AsmSpec.of_file(filename, self.ctx)
+
+    def add_entry(self, name, precond=smt.BoolVal(True)):
+        assert self.ctx.loader is not None, (
+            "BinaryContext must be loaded before disassembling"
+        )
+        sym = self.ctx.loader.find_symbol(name)
+        if sym is None:
+            raise Exception(f"Symbol {name} not found in binary {self.ctx.filename}")
+        self.spec.add_entry(name, sym.rebased_addr, precond)
+
+    def start(self, mem=None):
+        if mem is None:
+            mem = self.ctx.init_mem()
+        tracestates, vcs = init_trace_states(self.ctx, mem, self.spec)
+        self.tracestates = tracestates
+        self.vcs = vcs
+
+    def breakpoint(self, addr):
+        self.breakpoints.add(addr)
+
+    def step(self, n=1):
+        for _ in range(n):
+            tracestate = self.pop()
+            new_tracestates, new_vcs = execute_spec_and_insn(
+                tracestate, self.spec, self.ctx
+            )
+            self.vcs.extend(new_vcs)
+            self.tracestates.extend(new_tracestates)
+
+    def run(self):
+        while self.tracestates:
+            if self.addr() in self.breakpoints:
+                break
+            self.step()
+
+    def pop(self):
+        return self.tracestates.pop()
+
+    def addr(self):
+        return self.tracestates[-1].state.pc[0]
+
+    def ghost(self, name):
+        return self.tracestates[-1].ghost_env[name]
+
+    def reg(self, name):
+        reg = self.ctx._subst_decls[name]
+        return smt.simplify(
+            self.ctx.substitute(self.tracestates[-1].state.memstate, reg)
+        )
+
+    def ram(self, addr, size=None):
+        if size is None:
+            size = self.ctx.bits // 8
+        return smt.simplify(
+            self.tracestates[-1].state.memstate.getvalue_ram(addr, size)
+        )
+
+    def insn(self):
+        return self.ctx.disassemble(self.addr())
+
+    def model(self): ...  # TODO.
+
+
 @dataclass
 class Results:
     successes: list[str] = dataclasses.field(default_factory=list)

src/kdrag/notation.py

@@ -59,11 +59,12 @@ class SortDispatch:
     Not(True)
     """
 
-    def __init__(self, name=None, default=None, pointwise=False):
+    def __init__(self, name=None, default=None, pointwise=False, default_factory=None):
         self.methods = {}
         self.default = default
         self.name = name
         self.pointwise = pointwise
+        self.default_factory = default_factory
 
     def register(self, sort, func):
         self.methods[sort] = func
@@ -96,30 +97,36 @@ def __call__(self, *args, **kwargs):
         """
         if not args:
             raise TypeError("No arguments provided")
+        x0 = args[0]
         sort = args[0].sort()
-        res = self.methods.get(sort, self.default)
-        if res is None:
-            x0 = args[0]
-            if self.pointwise and smt.is_func(x0):
-                doms = smt.domains(x0)
-                vs = [smt.FreshConst(d, prefix=f"x{n}") for n, d in enumerate(doms)]
-                sort = x0.sort()
-                # sorts are same or attempt coerce/lift
-                return smt.Lambda(
-                    vs,
-                    self(
-                        *[
-                            arg(*vs)
-                            if isinstance(arg, smt.ExprRef) and arg.sort() == x0.sort()
-                            else arg
-                            for arg in args
-                        ]
-                    ),
-                )
-            else:
-                raise NotImplementedError(
-                    f"No implementation of {self.name} for sort {sort}. Register a definition via {self.name}.register({sort}, your_code)",
-                )
+        res = self.methods.get(sort)
+        if res is not None:
+            return res(*args, **kwargs)
+        elif self.default is not None:
+            return self.default(*args, **kwargs)
+        elif self.default_factory is not None:
+            res = self.default_factory(sort)
+            self.register(sort, res)
+            return res(*args, **kwargs)
+        elif self.pointwise and smt.is_func(x0):
+            doms = smt.domains(x0)
+            vs = [smt.FreshConst(d, prefix=f"x{n}") for n, d in enumerate(doms)]
+            # sorts are same or attempt coerce/lift
+            return smt.Lambda(
+                vs,
+                self(
+                    *[
+                        arg(*vs)
+                        if isinstance(arg, smt.ExprRef) and arg.sort() == x0.sort()
+                        else arg
+                        for arg in args
+                    ]
+                ),
+            )
+        else:
+            raise NotImplementedError(
+                f"No implementation of {self.name} for sort {sort}. Register a definition via {self.name}.register({sort}, your_code)",
+            )
         return res(*args, **kwargs)
 
     def define(self, args, body):
@@ -132,6 +139,21 @@ def define(self, args, body):
         return defn
 
 
+def GenericDispatch(default_factory) -> SortDispatch:
+    """
+    A decorator version of SortDispatch with a default factory.
+    This is useful for definition Sort generic definitions.
+
+    >>> @GenericDispatch
+    ... def id(sort):
+    ...    x = smt.Const("x", sort)
+    ...    return kd.define("id", [x], x)
+    >>> id(smt.IntVal(3))
+    id(3)
+    """
+    return SortDispatch(default_factory.__name__, default_factory=default_factory)
+
+
 call = SortDispatch(name="call")
 """Sort based dispatch for `()` call syntax"""
 smt.ExprRef.__call__ = lambda x, *y, **kwargs: call(x, *y, **kwargs)  # type: ignore

src/kdrag/theories/algebra/filter.py

@@ -0,0 +1,50 @@
+import functools
+import kdrag.smt as smt
+import kdrag as kd
+import kdrag.theories.set as set_
+
+
+@functools.cache
+def Filter(T: smt.SortRef):
+    """
+    A sort generic filter over sets of type T.
+
+    >>> Filter(smt.RealSort())
+    Filter_Real
+    """
+    dt = kd.Struct(f"Filter_{T}", ("sets", smt.SetSort(smt.SetSort(T))))
+    A, B = smt.Consts("A B", smt.SetSort(T))
+    F = smt.Const("F", dt)
+    kd.notation.call.register(dt, lambda F, A: dt.sets(F)[A])
+    kd.notation.wf.define(
+        [F],
+        smt.And(
+            F(smt.FullSet(T)),
+            smt.ForAll([A, B], F(A), A <= B, F(B)),
+            smt.ForAll([A, B], F(A), F(B), F(A & B)),
+        ),
+    )
+    return dt
+
+
+class FilterMod:
+    """
+    A module encapsulating filter theory over sets of type T.
+
+    >>> FilterMod(smt.RealSort()).filter_full
+    |= ForAll(F, Implies(wf(F), sets(F)[K(Real, True)]))
+    """
+
+    def __init__(self, T: smt.SortRef):
+        self.T = T
+        self.S = Filter(T)
+        F = smt.Const("F", self.S)
+        set_.Set(T)
+        A, B = smt.Consts("A B", smt.SetSort(T))
+        self.filter_full = kd.prove(kd.QForAll([F], F(smt.FullSet(T))), unfold=1)
+        self.filter_mono = kd.prove(
+            kd.QForAll([F], smt.ForAll([A, B], F(A), A <= B, F(B))), unfold=1
+        )
+        self.filter_inter = kd.prove(
+            kd.QForAll([F], smt.ForAll([A, B], F(A), F(B), F(A & B))), unfold=1
+        )

tests/test_pcode.py

@@ -1,6 +1,6 @@
 
 from kdrag.all import *
-from kdrag.contrib.pcode.asmspec import assemble_and_check_str, AsmSpec, run_all_paths, kd_macro
+from kdrag.contrib.pcode.asmspec import assemble_and_check_str, AsmSpec, run_all_paths, kd_macro, Debug
 import kdrag.contrib.pcode as pcode
 
 import pytest
@@ -175,6 +175,15 @@ def test_cli():
     res = subprocess.run("python3 -m kdrag.contrib.pcode /tmp/mov42.s", shell=True, capture_output=True, text=True)
     assert "verification conditions failed. ❌❌❌❌" in res.stdout
 
+def riscv_as(code):
+    with open("/tmp/mov42.s", "w") as f:
+        f.write(code)
+        f.flush()
+    res = subprocess.run("""nix-shell -p pkgsCross.riscv32-embedded.buildPackages.gcc \
+        --run "riscv32-none-elf-as /tmp/mov42.s -o /tmp/mov42.o" """, shell=True, check=True, capture_output=True, text=True)
+
+
+
 @pytest.mark.slow
 def test_riscv32():
     code = """
@@ -221,4 +230,38 @@ def test_riscv32_write():
     vcs = run_all_paths(ctx, spec)
     assert len(vcs) == 1
     for vc in vcs:
-        vc.verify(ctx)
+        vc.verify(ctx)
+
+@pytest.mark.slow
+def test_debugger():
+    code = """
+    .include "/tmp/knuckle.s"
+    .option norvc
+    .text
+    .globl _start
+    _start:
+    li   t0, 42          # load immediate 42 into t0
+    li   t1, 1000        # load address 1000 into t1
+    sw   t0, 0(t1)       # store t0 at memory address
+    ret
+    """
+    riscv_as(code)
+    ctx = pcode.BinaryContext("/tmp/mov42.o", langid="RISCV:LE:32:default")
+    spec = AsmSpec.of_file("/tmp/mov42.s", ctx)
+    d = Debug(ctx, spec)
+    d.add_entry("_start")
+    d.start()
+    print(d.insn())
+    print(d.addr())
+    assert d.reg("t0").eq(smt.BitVecVal(42,32))
+    assert not d.reg("t1").eq(smt.BitVecVal(1000,32))
+    d.step()
+    assert d.reg("t1").eq(smt.BitVecVal(1000,32))
+    d.step()
+    assert d.ram(smt.BitVecVal(1000,32), 4).eq(smt.BitVecVal(42,32))
+
+
+    
+
+
+