kleene, vprove, abs_sum, more ir

Author: philzook58

src/kdrag/contrib/ir/__init__.py

@@ -18,6 +18,7 @@
 from dataclasses import dataclass, field
 import kdrag as kd
 import kdrag.smt as smt
+from collections import defaultdict
 
 
 def pp_sort(s: smt.SortRef) -> str:
@@ -100,13 +101,44 @@ def __repr__(self) -> str:
             res.append(f"%{i} = {rhs}")
         return "\n".join(res)
 
+    def succ_calls(self) -> list[smt.ExprRef]:
+        jmp = self.insns[-1]
+        if smt.is_if(jmp):
+            return jmp.children()
+        else:
+            return [jmp]
+
+
+type Label = str
+
 
 @dataclass
 class Function:
     """ """
 
-    entry: str  # smt.FuncDeclRef?
-    blocks: dict[str, Block]  # 0th block is entry. Or "entry" is entry? Naw. 0th.
+    entry: Label  # smt.FuncDeclRef?
+    blocks: dict[Label, Block]  # 0th block is entry. Or "entry" is entry? Naw. 0th.
+
+    def calls_of(self) -> dict[str, list[tuple[Label, smt.ExprRef]]]:
+        """
+        Returns a mapping from labels to a list of calls to that label
+        """
+        p = defaultdict(list)
+        for label, blk in self.blocks.items():
+            for call in blk.succ_calls():
+                p[call.decl().name()].append((label, call))
+        return p
+
+    def phis(self):
+        """
+        Return the analog a mapping from labels to phi nodes in that block
+        """
+
+        preds = self.calls_of()
+        phis = {}
+        for label, blk in self.blocks.items():
+            phis[label] = zip(*[call.children() for _, call in preds[label]])
+        return phis
 
 
 @dataclass

src/kdrag/solvers/__init__.py

@@ -343,8 +343,8 @@ def check(self):
             cmd = [
                 self.options["solver_path"],
                 fp.name,
-                "--mode",
-                "casc",
+                # "--mode", # This adds portfolio mode, but it was slower. Maybe more useful for hard questions?
+                # "casc",
                 "--input_syntax",
                 "smtlib2",
                 # "--ignore_unrecognized_logic", "on",

src/kdrag/tactics.py

@@ -324,6 +324,17 @@ def prove(
         raise e
 
 
+def vprove(thm: smt.BoolRef, **kwargs) -> kd.kernel.Proof:
+    """
+    Prove a theorem using an egraph-based solver.
+
+    >>> x = smt.Int("x")
+    >>> vprove(x + 1 > x)
+    |= x + 1 > x
+    """
+    return prove(thm, solver=solvers.VampireSolver, **kwargs)
+
+
 def simp(t: smt.ExprRef, by: list[kd.kernel.Proof] = [], **kwargs) -> kd.kernel.Proof:
     rules = [kd.rewrite.rewrite_of_expr(lem.thm) for lem in by]
     t1 = kd.rewrite.rewrite_once(t, rules)

src/kdrag/theories/algebra/kleene.py

@@ -0,0 +1,89 @@
+"""
+Kleene Algebra
+
+
+- https://www.philipzucker.com/bryzzowski_kat/
+- https://www.cs.uoregon.edu/research/summerschool/summer24/topics.php#Silva
+
+"""
+
+import kdrag as kd
+import kdrag.smt as smt
+
+K = smt.DeclareSort("Kleene")
+zero = smt.Const("_0", K)
+one = smt.Const("_1", K)
+x, y, z, w, e, f = smt.Consts("x y z w e f", K)
+par = smt.Function("par", K, K, K)
+seq = smt.Function("seq", K, K, K)
+star = smt.Function("star", K, K)
+kd.notation.or_.register(K, par)
+kd.notation.add.register(K, par)
+kd.notation.matmul.register(K, seq)
+kd.notation.mul.register(K, seq)
+kd.notation.le.register(K, lambda x, y: smt.Eq(par(x, y), y))
+
+
+par_assoc = kd.axiom(smt.ForAll([x, y, z], x + (y + z) == (x + y) + z))
+par_comm = kd.axiom(smt.ForAll([x, y], x + y == y + x))
+par_idem = kd.axiom(smt.ForAll([x], x + x == x))
+par_zero = kd.axiom(smt.ForAll([x], x + zero == x))
+
+zero_par = kd.prove(smt.ForAll([x], zero + x == x), by=[par_comm, par_zero])
+
+seq_assoc = kd.axiom(smt.ForAll([x, y, z], x * (y * z) == (x * y) * z))
+seq_zero = kd.axiom(smt.ForAll([x], x * zero == zero))
+seq_one = kd.axiom(smt.ForAll([x], x * one == x))
+
+rdistrib = kd.axiom(smt.ForAll([x, y, z], x * (y + z) == x * y + x * z))
+ldistrib = kd.axiom(smt.ForAll([x, y, z], (y + z) * x == y * x + z * x))
+
+unfold = kd.axiom(smt.ForAll([e], star(e) == one + e * star(e)))
+
+# If a set shrinks, star(e) is less than it
+
+least_fix = kd.axiom(smt.ForAll([x, e, f], f + e * x <= x, star(e) * f <= x))
+
+KLEENE = [
+    par_assoc,
+    par_comm,
+    par_idem,
+    par_zero,
+    seq_assoc,
+    seq_zero,
+    seq_one,
+    rdistrib,
+    ldistrib,
+    unfold,
+    least_fix,
+]
+par_monotone = kd.prove(
+    smt.ForAll([x, y, z, w], x <= y, z <= w, x + z <= y + w),
+    by=[par_assoc, par_comm],
+)
+
+seq_monotone = kd.prove(
+    smt.ForAll([x, y, z, w], x <= y, z <= w, x * z <= y * w),
+    by=[rdistrib, ldistrib, par_assoc],
+)
+
+star_seq_star = kd.tactics.vprove(
+    smt.ForAll([x], star(x) * star(x) == star(x)), by=KLEENE
+)
+# z3 takes 0.5 seconds. Vampire is actually faster despite all the overhead
+star_star = kd.tactics.vprove(smt.ForAll([x], star(star(x)) == star(x)), by=KLEENE)
+
+Test = smt.DeclareSort("Test")
+guard = smt.Function("guard", Test, K, K)
+and_ = smt.Function("and", Test, Test, Test)
+or_ = smt.Function("or", Test, Test, Test)
+not_ = smt.Function("not", Test, Test)
+true = smt.Const("true", Test)
+false = smt.Const("false", Test)
+kd.notation.invert.register(Test, not_)
+
+a, b, c, d = smt.Consts("a b c d", Test)
+
+
+def while_(t, c):
+    return star(guard(t, c)) * guard(~t, one)

src/kdrag/theories/real/__init__.py

@@ -137,6 +137,9 @@ def abstract_arith(t: smt.ExprRef) -> smt.ExprRef:
     ForAll([x, y], smt.Implies(y != 0, abs(x / y) == abs(x) / abs(y))),
     by=[abs.defn],
 )
+abs_Lipschitz = kd.prove(
+    ForAll([x, y], abs(abs(x) - abs(y)) <= abs(x - y)), by=[abs.defn]
+)
 abs_triangle = kd.prove(
     ForAll([x, y, z], abs(x - y) <= abs(x - z) + abs(z - y)), by=[abs.defn]
 )
@@ -220,6 +223,10 @@ def abstract_arith(t: smt.ExprRef) -> smt.ExprRef:
 floor_le = kd.prove(ForAll([x], floor(x) <= x), by=[floor.defn])
 floor_gt = kd.prove(ForAll([x], x < floor(x) + 1), by=[floor.defn])
 
+
+ceil = kd.define("ceil", [x], -floor(-x))
+ceil_le = kd.prove(ForAll([x], x <= ceil(x)), by=[ceil.defn, floor_le])
+ceil_gt = kd.prove(ForAll([x], ceil(x) - 1 < x), by=[ceil.defn, floor_gt])
 # c = kd.Calc([n, x], smt.ToReal(n) <= x)
 # c.eq(n <= smt.ToInt(x))
 # c.eq(smt.ToReal(n) <= floor(x), by=[floor.defn])
@@ -237,6 +244,23 @@ def abstract_arith(t: smt.ExprRef) -> smt.ExprRef:
 kd.kernel.prove(smt.Implies(x > 0, x**0 == 1))
 # pow_pos = kd.prove(kd.QForAll([x, y], x > 0, pow(x, y) > 0), by=[pow.defn])
 
+pownat = smt.Function("pownat", smt.RealSort(), smt.IntSort(), smt.RealSort())
+pownat = kd.define(
+    "pownat",
+    [x, n],
+    kd.cond((n == 0, 1), (n > 0, x * pownat(x, n - 1)), default=pownat(x, n + 1) / x),
+)
+
+# Basic lemmas for pownat.
+pownat_zero = kd.prove(
+    smt.ForAll([x], pownat(x, 0) == 1),
+    by=[pownat.defn],
+)
+pownat_succ = kd.prove(
+    kd.QForAll([x, n], n >= 0, pownat(x, n + 1) == x * pownat(x, n)),
+    by=[pownat.defn],
+)
+
 sqr = kd.define("sqr", [x], x * x)
 
 
@@ -385,6 +409,29 @@ def abstract_arith(t: smt.ExprRef) -> smt.ExprRef:
 )
 
 
+# x, y = smt.Reals("x y")
+N = smt.Int("N")
+
+
+@kd.Theorem(
+    smt.ForAll(
+        [x, y],
+        smt.Implies(
+            smt.And(x > 0, y > 0),
+            smt.Exists([N], y < x * N),
+        ),
+    )
+)
+def archimedes(l):
+    # https://en.wikipedia.org/wiki/Archimedean_property
+    x, y = l.fixes()
+    l.intros()
+
+    # Choose N = floor(y/x) + 1.
+    l.exists(smt.ToInt(y / x) + 1)
+    l.auto()
+
+
 # smt.Function("cont_at", RFun, R, smt.BoolSort())
 
 is_diff = kd.define("is_diff", [f], smt.ForAll([x], diff_at(f, x)))

src/kdrag/theories/real/seq.py

@@ -2,6 +2,7 @@
 import kdrag as kd
 import kdrag.smt as smt
 
+import kdrag.theories.int as int_
 import kdrag.theories.real as real
 
 
@@ -209,6 +210,72 @@ def max_suffix_bound(l):
 cos = kd.define("cos", [a], smt.Map(real.cos, a))
 
 abs = kd.define("abs", [a], smt.Map(real.abs, a))
+abs_ext = kd.prove(smt.ForAll([a, n], abs(a)[n] == real.abs(a[n])), by=[abs.defn])
+abs_ge_0 = kd.prove(
+    kd.QForAll([a, n], abs(a)[n] >= 0),
+    by=[abs_ext, real.abs_pos],
+)
+
+
+@kd.Theorem(
+    kd.QForAll(
+        [a, N],
+        N >= 0,
+        kd.QForAll([n], n >= 0, n <= N, abs(a)[n] <= finsum(abs(a), N)),
+    )
+)
+def abs_le_finsum_abs(l):
+    a, N = l.fixes()
+
+    l.induct(N, using=int_.induct_nat)
+    # base case: N = 0
+    l.intros()
+    n = l.fix()
+    l.intros()
+    l.split(at=-1)
+    l.have(n == 0, by=[])
+    l.rw(-1)
+    l.unfold(finsum)
+    l.unfold(cumsum)
+    l.simp()
+    l.auto()
+
+    # step: N >= 0, IH -> N + 1
+    N = l.fix()
+    l.intros()
+    l.intros()
+
+    n = l.fix()
+    l.intros()
+    l.split(at=0)
+
+    # Case split: n <= N or n = N + 1
+    l.cases(n <= N)
+
+    # n <= N: use IH and monotonicity of finsum
+    l.have(abs(a)[n] <= finsum(abs(a), N), by=[])
+    l.have(N + 1 > 0, by=[])
+    l.have(
+        finsum(abs(a), N + 1) == finsum(abs(a), N) + abs(a)[N + 1],
+        by=[finsum.defn, cumsum.defn],
+    )
+    l.have(abs(a)[N + 1] >= 0, by=[abs_ge_0])
+    l.have(finsum(abs(a), N) <= finsum(abs(a), N + 1), by=[])
+    l.auto()
+
+    # n > N: then n = N + 1
+    l.have(n == N + 1, by=[])
+    l.rw(-1)
+    l.have(N + 1 > 0, by=[])
+    l.have(
+        finsum(abs(a), N + 1) == finsum(abs(a), N) + abs(a)[N + 1],
+        by=[finsum.defn, cumsum.defn],
+    )
+    l.have(abs(a)[N + 1] >= 0, by=[abs_ge_0])
+    l.have(abs(a)[0] <= finsum(abs(a), N), by=[])
+    l.have(abs(a)[0] >= 0, by=[abs_ge_0])
+    l.have(finsum(abs(a), N + 1) >= abs(a)[N + 1], by=[])
+    l.auto()
 
 
 has_bound = kd.define("has_bound", [a, M], kd.QForAll([n], real.abs(a[n]) <= M))

src/kdrag/theories/real/seq_extra.py

@@ -0,0 +1,2 @@
+import kdrag as kd
+import kdrag.smt as smt