A Lean Proof Printing Python Union Find
It is somewhat strange, but I haven’t really spent much time working on proof production or recording out of egraphs.
It’s a useful capability for all sorts of reasons. If you want to use an e-graph as an theorem prover backend for another system it’s table stakes. Something I found interesting is that even if you are using e-graphs for optimization, a proof certificate might be necessary in some applications. Sam has mentioned that in the hardware industry, if you don’t have a certificate understandable to the tools with the golden seal of approval, it’s a no go. Makes sense, because mistakes in hardware can be very very expensive.
There is a paper Proof Producing Congruence Closure https://www.cs.upc.edu/~roberto/papers/rta05.pdf Nieuwenhuis and Oliveras which is a touchstone of this topic.
Before you consider anything about e-graphs, it’s useful to first think about the simpler problem of union finds.
A union find is useful for efficiently tracking connected components of a graph https://en.wikipedia.org/wiki/Disjoint-set_data_structure https://en.wikipedia.org/wiki/Kruskal's_algorithm . A proof that two vertices are in the same connected component is an explicit path between them. There could be more than one path but they all serve the purpose of demonstrating connectivity.
A spanning tree is a nice compact way of storing available paths.
The union find is maintaining a forest. A confusing thing is that this forest does not reflect a spanning tree in the usual implementation of union finds. This is because union(a,b) first calls c = find(a) and d = find(b) and then attaches c-d in the union find. But the edge in the graph was a-b, not c-d.
One way of fixing this is to perform a rerooting operations akin to what is done in semipersistent data structures https://www.philipzucker.com/toa_unionfind/ to make a or b the root. You may want to maintain multiple trees, one that is an ordinary union find with compression, and another one with no compression or with rerooting.
But this is all kind of overwrought for a 0th attempt.
A principle I love is:
Proofs as Traces of Search
Now, I could have said “Proofs as traces of proof search” which would see a little circular. But it is the case that we often have some intuitive understanding of a search process like SAT solving, prolog backtracking, mixed integer programming, egraphs, term rewriting etc.
If we want to design a proof certificate for a process that we believe demonstrates something, the data of the proof has to be in the trace. It is fairly obvious how to remove redundant data to some degree. I don’t need to know how every register changed or about some failures. There are usually pretty obvious high level statements. These may correspond roughly to logging function calls or successful function calls.
A canonical example of this sort of thing is UNSAT proofs https://github.com/marijnheule/drat-trim . It is a trace of learned clauses from the sat solver which is basically a resolution proof. It also include deletions which is a useful but operational looking detail.
The union find and the e-graph are intuitively speaking manipulating equations. Path compression in the union find is applying transitivity on e1 = e14. The memo table is an equation between nodes and eids f(e1,e2) = e7. We can basically just stream out everything that is happening as Lean as we go along.
Here is a basic union find without proofs. Unusually I am doing path compression and using a recursive find rather than iterative because I found it easier later.
from dataclasses import dataclass, field
type Id = int
@dataclass
class UF:
parents : list[Id] = field(default_factory=list)
def makeset(self) -> Id:
id = len(self.parents)
self.parents.append(id)
return id
def find(self, id0 : Id) -> Id: # or just regturn Id and let user lookup proof if interested?
assert isinstance(id0, int), "id0 must be an integer" + str(id0)
id = self.parents[id0]
if id == id0:
return id
else:
id1 = self.find(id)
self.parents[id0] = id1 # path compression
return id1
def union(self, id0: Id, id1: Id) -> None:
root0 = self.find(id0)
root1 = self.find(id1)
if root0 != root1:
self.parents[root0] = root1
uf = UF()
a,b,c = [uf.makeset() for _ in range(3)]
uf.union(a,b)
uf.union(b,c)
assert uf.find(a) == uf.find(c)
Then here is it filled in with enough junk to talk about proofs.
log contains the lines of a lean proof. I use the length of log as a freshness counter for allocating new nodes of the proof tree. It’s a RecExpr of proofs kind of. I think there isn’t much point of hash consing proofs as a first cut.
find records path compression via Eq.trans. There are possibly more efficient ways of implementing it.
- each edge in
parentshas a proof associated with it. This is an index to a proof object identifier memohas a proof associated with it about how we defined externally named constants to internally namedletbounde42ids
from dataclasses import dataclass, field
type PId = int
type Id = int
type FatId = tuple[PId, Id]
@dataclass
class UF:
parents : list[FatId] = field(default_factory=list)
memo : dict[str, FatId] = field(default_factory=dict) # named identifiers
log : list[str] = field(default_factory=list) # a list of lean lines
def makeset(self, name : str) -> Id:
if name in self.memo:
return self.memo[name][1]
else:
id = len(self.parents)
pid = len(self.log)
self.log.append(f"let e{id} := {name}") # fresh name via let
self.log.append(f"let p{pid} : {name} = e{id} := Eq.refl e{id}") # let expansion. Could elide?
self.memo[name] = (pid, id)
fid = (self.identity(id), id)
self.parents.append(fid)
return id
def comp(self, pid1 : PId, pid2 : PId) -> PId:
pid = len(self.log)
self.log.append(f"let p{pid} := Eq.trans p{pid1} p{pid2}")
return pid
def identity(self, id : Id) -> PId:
pid = len(self.log)
self.log.append(f"let p{pid} : e{id} = e{id} := Eq.refl e{id}")
return pid
def inv(self, pid : PId) -> PId: # Or maybe use negative PId to avoid doing this so much.
pid_new = len(self.log)
self.log.append(f"let p{pid_new} := Eq.symm p{pid}")
return pid_new
def find(self, id0 : Id) -> Id: # or just regturn Id and let user lookup proof if interested?
assert isinstance(id0, int), "id0 must be an integer" + str(id0)
pid, id1 = self.parents[id0]
if id1 == id0:
return id0
else:
root = self.find(id1)
pid1 = self.parents[id1][0]
pid2 = self.comp(pid, pid1)
self.parents[id0] = (pid2, root)
return root
# we could have a `rebuild``, which normalizes memo. There isn't much point here.
def union(self, id0: Id, id1: Id, pname : str) -> None:
root0 = self.find(id0)
root1 = self.find(id1)
if root0 != root1:
pid0, pid1 = self.parents[id0][0], self.parents[id1][0]
pid2 = len(self.log)
self.log.append(f"let p{pid2} : e{root0} = e{root1} := trans (Eq.symm p{pid0}) (trans {pname} p{pid1})")
self.parents[root0] = (pid2, root1)
#self.log.append(f"let p{pid2} := Eq.trans p{pid} (Eq.symm p{pid1}")
def proof(self, name0 : str, name1: str) -> str:
(pid0, id0), (pid1, id1) = self.memo[name0], self.memo[name1]
root0 = self.find(id0)
root1 = self.find(id1)
pid2, pid3 = self.parents[id0][0], self.parents[id1][0]
if root0 != root1:
raise ValueError(f"No proof: {id0} and {id1} are in different sets")
else:
return " " + "\n ".join(self.log) + "\n " + f"exact trans (Eq.trans p{pid0} p{pid2}) (Eq.symm (trans p{pid1} p{pid3}))"
We can call lean as subprocess to check our proofs
import subprocess
def check(s):
result = subprocess.run(["lake", "env", "lean", "--stdin"], input=s, text=True, capture_output=True)
if result.returncode != 0:
raise RuntimeError(result.stdout)
return result.stdout
#succeeds
check("example : True := by trivial")
''
This fails
check("example : Not True := by trivial") # fails
---------------------------------------------------------------------------
RuntimeError Traceback (most recent call last)
Cell In[88], line 1
----> 1 check("example : Not True := by trivial") # fails
Cell In[85], line 5, in check(s)
2 def check(s):
3 result = subprocess.run(["lake", "env", "lean", "--stdin"], input=s, text=True, capture_output=True)
4 if result.returncode != 0:
----> 5 raise RuntimeError(result.stdout)
6 return result.stdout
RuntimeError: <stdin>:1:25: error: Tactic `assumption` failed
⊢ ¬True
And here is using the union find trivially and a printout of the proof term. I’m building the term in tactic mode (note the exact) since there might be perf wins to this?
uf = UF()
a,b,c = uf.makeset("a"), uf.makeset("b"), uf.makeset("c")
#print(uf.proof("a", "a"))
uf.union(a, b, "pfab")
uf.union(b, c, "pfbc")
print(uf.proof("a", "c"))
let e0 := a
let p0 : a = e0 := Eq.refl e0
let p2 : e0 = e0 := Eq.refl e0
let e1 := b
let p3 : b = e1 := Eq.refl e1
let p5 : e1 = e1 := Eq.refl e1
let e2 := c
let p6 : c = e2 := Eq.refl e2
let p8 : e2 = e2 := Eq.refl e2
let p9 : e0 = e1 := trans (Eq.symm p2) (trans pfab p5)
let p10 : e1 = e2 := trans (Eq.symm p5) (trans pfbc p8)
let p11 := Eq.trans p10 p8
let p12 := Eq.trans p9 p11
exact trans (Eq.trans p0 p12) (Eq.symm (trans p6 p8))
Hey it passes. Neato. You need to somehow pass references to the names of the reasons for the unions. Here I’m doing so manually. In equality saturation, you’d pass in the rules and instantiate them internally in the egraph / proof.
check("example (a b c : Int) (pfab : a = b) (pfbc : b = c) : a = c := by \n" + uf.proof("a", "c"))
''
Bits and Bobbles
Other experiments
The full path union find. You can store all paths as loops in the root The 2-union find - string knuth bendix on paths. So kind of the full path union find but it allows new relations / quotients on the path objects.
Fat Ids. https://www.philipzucker.com/union-find-groupoid/ Rudi has been working on proofs as annotations. This post was partially inspired by that
Outputting lean proofs is also nice because I messed up a bunch. Then i could copy and paste the output into a file and get nice little squiggles and hovers.
https://www.philipzucker.com/proof_objects/ Other proofs objects I love besides traces of search
https://arxiv.org/abs/2209.03398 small proofs from congruence https://dl.acm.org/doi/10.1007/978-3-032-15700-3_1 Producing Shorter Congruence Closure Proofs in a State-of-the-Art SMT Solver https://arxiv.org/abs/2504.10246 Simplified and Verified: A Second Look at a Proof-Producing Union-Find Algorithm
Rebuilding is repair of the memo table. normalizing the arguments of an enode is applying cong f(e1,e2) = f(e3,e4) but then also replacement of f(e1,e2) -> e7 in the memo table is a transitivity chaining tran (cong find(e))
The techniques in the proof producing congruence closure correspond to trying to reduce the amount of data saved, and pruning the proof objects / proof terms to avoid unused elements basically. Dead code elimination. If instead of opaque lines in log I actually stored which subproofs things depend on, it would not be hard to traverse log backwards only actually collecting the relevant proofs. Lean might have built in dead code elimination facilities to textually process a file? Maybe not. There are some zeta unfolding options or something.
https://github.com/Z3Prover/z3/discussions/4881 z3 proof certs discussion. There is a drat like congruence format I saw https://github.com/Z3Prover/z3/discussions/5000
There is some difference in performance of driectly producing proof terms vs producing them in by tactic mode. This may have something to do with HaveTelescope https://lean-lang.org/doc/api/Lean/Meta/HaveTelescope.html .
AI ate shit on proof production in lambda microegg, which was another inspriation for this post. Graham just vibe coded a proof producing egraph right into aufbau https://github.com/gleachkr/Aufbau though. Graham’s real smart though.
It seems like it is probably better to unbundle / struct of arrays the parents fatid. Proofs have a cost, it’s easier to avoid that cost if I just have it under a bool flag.
It is somewhat amusing to consider tuple[object, PId] as a python analog of the lean type {a | p a}. One can write intrinsically proof producing python programs that when run produce a lean certificate that the result satisfies some desire property. This does not prove the algorithm in python will generally produce correct results though. The ole ten penny refinement types.
Concretized to that degree, I would tend to expect p a to be decidable (in which case there isn’t much point to outputting a cert. Just run a checker), but maybe sometimes it isn’t.
from dataclasses import dataclass
type Id = int
type Path = tuple[str, ...] # but maybe (Id,str,Id) would be appropriate? Then we can check for correctness, derive codomain etc.
type FatId = tuple[Path, Id]
@dataclass
class UF:
parents : list[FatId]
loops : list[set[Path]] # "Lattice" (set) of loops
def find(self, x: FatId) -> FatId:
pf, id = x
while True:
p, id2 = self.parents[id]
if id == id2:
return (pf, id)
pf = pf + p
id = id2
def makeset(self):
id = len(self.parents)
self.parents.append(((), id))
return ((),), id
def union(self, x: FatId, y: FatId): # why/explain parameter?
x,y = self.find(x), self.find(y)
xpf, xid = x
ypf, yid = y
if xid != yid:
self.parents[xid] = (xpf + reversed(ypf), yid)
else:
self.loops.append({xpf + reversed(ypf)})
def move(p : Path, x : FatId):
return (p + x[0], x[1])
Hmm. https://www.cs.upc.edu/~roberto/papers/rta05.pdf I don’t need to restore the FatIds. unions is kind of a separate storage of FatId, factoring into two arenas. Maybe that’s a nice design generally?
That you have to record the direction is interesting. That is UnionResult
Hmm. If two paths share a lot, reducing that is prefix matching which is an LCA computation
rerooting looks like theory of arrays union find?
In the union find forest, nearly all the edges can be not ones in the graph. Is it a problem that we may normalize outside of the vision of the proof recorder? That does make fatids seem interesting Proofs being ordered sounds nice.
It doesn’t have to carry around the full chain, and why should it? We could construct the chain if fatid is (rawid,rawid) pairs. the latter being where we’re at, and the former being where we came from
Yes so keeping a normal union find around that counts sizes (or depth) helps us avoid too many inversion.
We always want both normalized and unnormalized. Intensional and extensionsal?
Relation to “Theory of Array” ideas. Semi persistent proof term hash cons. Destructively destroy term but keep proofs in semipersistent annotations.
(a,root,b) = (a,root2,b) we are equating two proofs and picking one as more canonical. It’s a new square. (a,root) would be more like a rewrite proof, or a normalization derivation?
(a,root,b) if we erase a,b we are conflating distinct proofs, akin to erasing context is conflating. (a,a,a) refl (a,b,b) or (a,a,b) as one rule.
Rule gives us edge we didn’t even have. so isn’t unioning anything? (a,a,b) there is no other (a,?,b)
cong(a,a’, b, b’) # binary cong
Strong pruning view is important for ematching though f((a,root,b), (c, root2, d)) enode is cong node f(a,root,b) # keep around previous stages egraphs? Match was done on previous stage, so we’re good.
f((a, root), (b, root)) # a is root from previous stage rather than maintaining separate renaming HashMap<Id, NewId>
Vec
f((a, rule(x,y), b), …)
makeset(src,dst, reason) - play the role of ctx in thinning uf rather than union(a,b,reason)
a hash cons modulo theories for proofs = proof producing union find
from dataclasses import dataclass, field
type Reason = tuple[object, bool] # reason and reversal
type FatId = tuple[Reason, int]
# type FatId = tuple[int, int] # src, dst
@dataclass
class UF:
parents : list[FatId]
def makeset(self) -> FatId: # hmm interesting. Doesn't work because what fatness. Actual path?
self.parents.append(((None, False), len(self.parents)))
return len(self.parents) - 1
@dataclass
class UF:
parents : list[int] = field(default_factory=list)
reasons : list[Reason] = field(default_factory=list)
def makeset(self) -> int:
self.parents.append(len(self.parents))
self.reasons.append((None, False)) # refl
return len(self.parents) - 1
def find(self, x: int) -> int:
while self.parents[x] != x:
x = self.parents[x]
return x # also return count for omre effiicnet rerooting choice?
def find_chain(self, x: int) -> list[int]:
chain = [x]
while self.parents[x] != x:
x = self.parents[x]
chain.append(x)
return chain
def reroot(self, x : int):
# hmm. so there is a version that reverses as you go up.
chain = self.find_chain(x)
N = len(chain)
for i in range(N-1): # reverse direction
x,y = chain[N-i-2], chain[N-i-1]
self.parents[y] = x
reason,rev = self.reasons[x]
self.reasons[y] = (reason, not rev)
def union(self, x : int, y : int, reason) -> bool:
x1,y1 = self.find(x), self.find(y) # cheap find first
if x1 != y1:
self.reroot(x)
self.parents[x] = y
self.reasons[x] = (reason, False)
def explain(self, x : int, y : int) -> tuple[list[Reason], list[Reason]]:
assert self.find(x) == self.find(y)
chain_x = self.find_chain(x)
chain_y = self.find_chain(y)
reasons_x = [self.reasons[node] for node in chain_x[:-1]] # [:-1] to ignore refl at top? Eh.
reasons_y = [self.reasons[node] for node in chain_y[:-1]]
reasons = reasons_x + [(r, not rev) for r, rev in reversed(reasons_y)]
# could also pop True, False pairs
return reasons
#return reasons_x, reasons_y
#reason = [self.reasons[node] for node in chain_x[:-1]]
#for r in reversedchain_y[:-1]:
# reason.append((r, not rev))
uf = UF()
x,y,z = uf.makeset(), uf.makeset(), uf.makeset()
uf.union(x,y, "x->y")
uf.union(x, z, "x->z")
uf.explain(x,y)
[('x->z', False), ('x->z', True), ('x->y', False)]
from dataclasses import dataclass, field
@dataclass
class ProofUF:
parents: list[int] = field(default_factory=list)
unions : list[tuple[int,int] | None] = field(default_factory=list)
def makeset(self):
id = len(self.parents)
self.parents.append(id)
self.unions.append(None)
return id
def union(self, x: int, y: int):
x1, y1 = self.find(x), self.find(y)
if x1 != y1:
self.parents[x1] = y1
self.unions[x1] = (x, y)
def find(self, x: int) -> int:
while self.parents[x] != x:
x = self.parents[x]
return x
def find1(self, x : int) -> tuple[int, int]:
# return root and one just below root
if self.parents[x] == x:
return
prevx = x
while True:
x = self.parents[x]
if self.parents[x] == x:
return prevx, x
prevx = x
def explain(self, x : int, y : int) -> set[int]:
proof = set()
z = self.find(z)
todo = [(x,z), (y,z)]
while todo:
(a,b) = todo.pop()
a1 , b1= self.find1(a), self.find1(b)
if a1 is None:
continue
(x,y) = a1
if a1 not in proof:
self.proofs.add(a1) # unions index
todo.append((a, a1))
todo.append((b, b1))
return proof
uf = ProofUF()
x,y,z = [uf.makeset() for _ in range(3)]
uf.union(x,y)
uf.find1(x)
(0, 1)
Combine string KB at proof level with union find at bottom level.
https://smimram.github.io/ocaml-alg/
squier completion https://webusers.imj-prg.fr/~yves.guiraud/articles/polybook.pdf
UF isn’t collapsing proofs really. We could do that. Proof irrelevance I guess. Compare with a more ordinary proof producing union find
FatID with proof kind of reminds me of dirac belt trick or jordan wigner string or anyon paths. We have the thing here, yes, but there is a worldline saying where it came from.
class UF2:
string_rewrites : list[tuple[Path, Path]]
parents : list[FatId]
loops : list[set[Path]]
def rebuild(self):
# string KB
# normalize loops according to the string rewrites
# also loops shoul prune anything that is just a multiple. loop_generators
# we could have done this pruning in UF1 also
The group union find is not
You can kind of make a null groupoid out of a group and a set of objects. (ob, grp) pairs and there are morphisms (dom, grp, cod)
Prearrows / Arr def infer_cod(Ob, Arr) -> Option(Ob) # cod def infer_dom(Ob, Arr) ->
infer_cod(a, id) = a
Inferarrow, object, prearrow -> we can infer cod
id(a)
linear maps - pullbacks
pullback is unique up to isomorphisms. isomorphisms will abound in this thing. https://en.wikipedia.org/wiki/Pullback_(category_theory)#Least_common_multiple least common miltiple is pullback
A groupoid with just copies of the group on every edge. linear transformations. Again dots are different 1-d vector spaces. Copies
- ->
type Arr = object
type Obj = int
type FatId = tuple[Arr, Obj]
class PBArr(Protocol):
def pb(f, g) -> (Option Self, Option Self) # allow for the case where pullback is equalizer / can reuse a/b
# forall a b c f g, hom(f,a,c), hom(g,b,c) -> exists h d, pb1,pb2, f. pb1 = g. pb2 and d best
# The exists d is skolemized by sometimesn eeding makeset
class PullbackUF:
def union(self, a : FatId, b : FatId) -> None:
# there is an implicit c object a -annota> c <annotb- b
# infer_cod(a : FatId) == infer_cod(b : FatId)
#
Finset thinnings. Thinnings + permutations / injective finmaps Actually, yeah. Specializing the relation of alpha pemruation doesn’t seem that important compared to scope? Alpha permuted things aren’t equal unless the theory says they are.
Noninjective thinmaps also let us notice narrowing f(X,Y) narrows f(X,X) . A curious kind of equality assumption kind of. lam xy , f(X,Y) = g(X,Y) implies lam x, f(X,X) == g(X,X) . This is true in interpretations I’ve been considering, but has not been specialized out
thinnings f(t(X), s(Y)). Lifting that out would make every variable unique in an expressoin with equality joining at the top. Maybe not feasible (?) because expressions oculd have inf vars?
subst(t, x, y)
Substitution pounded in there.
from dataclasses import dataclass
class Term:
pass
@dataclass
class App(Term): # Comp(f, args) but specialize like Cons compared to append. ConsComp
f : str
args : MultiTerm
def shift(self, d : int) -> App:
return App(f=self.f, args=self.args.shift(d))
def substitute(self, t : MultiTerm) -> App:
return App(f=self.f, args=self.args.substitute(t))
@dataclass
class Var(Term):
n : int
def shift(self, d : int) -> Var:
return Var(n=self.n + d)
def substitute(self, t : MultiTerm) -> Var:
return t.ts[self.n]
@dataclass
class MultiTerm:
d : int
ts : list[Term]
def __or__(self, other : MultiTerm) -> MultiTerm: # truly independet paralle compose
return MultiTerm(d=self.d + other.d, ts=self.ts + other.shift(self.d).ts)
def __mul__(self, other : MultiTerm) -> MultiTerm: #
assert self.d == other.d
return MultiTerm(d=self.d, ts=self.ts + other.ts)
def __matmul__(self, other : MultiTerm) -> MultiTerm: # sequential compose
return self.substitute(other)
def substitute(self, t : MultiTerm) -> MultiTerm:
assert self.d == len(t.ts)
return MultiTerm(d=self.d, ts=[s.substitute(t) for s in self.ts])
def shift(self, d : int) -> MultiTerm:
return MultiTerm(d=self.d + d, ts=[t.shift(d) for t in self.ts])
def weaken(self, n : int) -> MultiTerm: # I guess we could weaken using a thinning.
return MultiTerm(d=self.d + n, ts=self.ts)
def app(f: str, *args : MultiTerm) -> MultiTerm:
d = args[0].d
assert all(a.d == d for a in args) and all(len(a.ts) == 1 for a in args)
return MultiTerm(d=d, ts=[App(f=f, args=MultiTerm(d=args[0].d, ts=list(args)))])
def const(f, d : int) -> MultiTerm:
return MultiTerm(d=d, ts=[App(f=f, args=MultiTerm(d=d, ts=[]))])
def var(n : int, d : int) -> MultiTerm:
return MultiTerm(d=d, ts=[Var(n)])
id_ = var(0, 1)
v0 = var(0, 1)
app("f", v0).substitute(app("g", v0))
swap = MultiTerm(d=2, ts=[var(1,2),var(0,2)])
swap
f = app("f", var(0,1)) # f(X) as representing f
g = app("g", var(0,2), var(1,2)) # g(X,Y) as representing g
fst = var(0,2)
snd = var(1,2)
proj = var
from pprint import pprint
pprint(g @ ((f @ f) * v0))
MultiTerm(d=2,
ts=[App(f='g',
args=MultiTerm(d=2,
ts=[MultiTerm(d=2,
ts=[App(f='f',
args=MultiTerm(d=1,
ts=[MultiTerm(d=1,
ts=[App(f='f',
args=MultiTerm(d=1,
ts=[MultiTerm(d=1,
ts=[Var(n=0)])]))])]))]),
MultiTerm(d=2, ts=[Var(n=0)])]))])
from dataclasses import dataclass
class Cat: ...
@dataclass
class Comp(Cat):
f : Cat
g : Cat
# or could flatten associativity?
@dataclass
class Id(Cat): # var
pass
@dataclass
class Decl(Cat):
f : str
arity : int
@dataclass
class Fork(Cat):
# same domain. but product over. Hmm. That's just the produce?
f : Cat
g : Cat
Cell In[10], line 21
class
^
SyntaxError: invalid syntax
t1 <- vs > t1 product diagram
exists vs -> t12 with projections. Yes
- is literally the categorical product. (but without giving the projects)
Multiterm flavored egraph
type FatId = (thin, SmallVec[rawId]) type FatId = Vec[ThinId] yes. Node = | {f : String, FatId} | Var
FatId.prod()
rgsvd
Just a version without ematching. Rebuild
or makeset / reason could be an obligation that goes into the example … : result := header
from dataclasses import dataclass, field
type PId = int
type Id = int
type FatId = tuple[PId, Id]
@dataclass
class UF:
parents : list[FatId] = field(default_factory=list)
memo : dict[str, FatId] = field(default_factory=dict)
log : list[str] = field(default_factory=list)
def makeset(self, name : str) -> Id:
if name in self.memo:
return self.memo[name]
else:
id = len(self.parents)
pid = len(self.log)
self.log.append(f"let e{id} := {name}") # fresh name via let
self.log.append(f"let p{pid} : {name} = e{id} := Eq.refl e{id}") # let expansion. Could elide?
self.memo[name] = (pid, id)
fid = (self.identity(id), id)
self.parents.append(fid)
return id
def comp(self, pid1 : PId, pid2 : PId) -> PId:
pid = len(self.log)
self.log.append(f"let p{pid} : := Eq.trans p{pid1} p{pid2}")
return pid
def identity(self, id : Id) -> PId:
pid = len(self.log)
self.log.append(f"let p{pid} : e{id} = e{id} := Eq.refl e{id}")
return pid
def inv(self, pid : PId) -> PId: # Or maybe use negative PId to avoid doing this so much.
pid_new = len(self.log)
self.log.append(f"let p{pid_new} := Eq.symm p{pid}")
return pid_new
"""
def find(self, id0: Id) -> FatId:
# maybe recusive find would be nice.
pid, id = self.parents[id0]
if id == id0:
return (pid, id1)
else:
while True:
pid1, id1 = self.parents[id]
if id == id1:
self.parents[id0] = (pid1, id1)
return (pid, id1)
pid = self.comp(pid, pid1)
id = id1
"""
def find(self, id0 : Id) -> Id: # or just regturn Id and let user lookup proof if interested?
assert isinstance(id0, int), "id0 must be an integer" + str(id0)
pid, id = self.parents[id0]
if id == id0:
return id
else:
id1 = self.find(id)
pid1 = self.parents[id1][0]
pid2 = self.comp(pid, pid1)
self.parents[id0] = (pid2, id1)
return id1
def union(self, id0: Id, id1: Id, pname : str) -> None:
root0 = self.find(id0)
root1 = self.find(id1)
if root0 != root1:
pid0, pid1 = self.parents[root0][0], self.parents[root1][0]
pid2 = len(self.log)
self.log.append(f"let p{pid2} : e{root0} = e{root1} := trans (symm p{pid0} (trans {pname} p{pid1}))")
self.parents[root0] = (pid2, root1)
#self.log.append(f"let p{pid2} := Eq.trans p{pid} (Eq.symm p{pid1}")
def proof(self, name0 : str, name1: str) -> str:
(pid0, id0), (pid1, id1) = self.memo[name0], self.memo[name1]
root0 = self.find(id0)
root1 = self.find(id1)
if root0 != root1:
raise ValueError(f"No proof: {id0} and {id1} are in different sets")
else:
return "\n".join(self.log) + "\n" + f"Eq.trans p{pid0} (Eq.symm p{pid1})"
uf = UF()
a,b = uf.makeset("a"), uf.makeset("b")
#print(uf.proof("a", "a"))
uf.union(a, b,"pfab")
print(uf.proof("a", "b"))
let e0 := a
let p0 : a = e0 := Eq.refl e0
let p2 : e0 = e0 := Eq.refl e0
let e1 := b
let p3 : b = e1 := Eq.refl e1
let p5 : e1 = e1 := Eq.refl e1
let p6 : e0 = e1 := trans (symm p2 (trans pfab p5))
let p7 : := Eq.trans p6 p5
Eq.trans p0 (Eq.symm p3)
Hmm. What if (pid, userterm)
type FatId = tuple[Pid, str]
Are some things unreferrable? Maybe.
Ooh the quantifier union find + proofs could be interesting.
# Based on https://github.com/mwillsey/microegg
from dataclasses import dataclass, field
from typing import Callable
type Id = int
type PId = int
type FatId = tuple[PId, Id]
class Term:...
@dataclass(frozen=True)
class App(Term):
f: object
args: tuple[Term, ...]
@dataclass(frozen=True)
class Var(Term): name: str
type Subst = dict[str, Id]
@dataclass(frozen=True)
class Node:
f: object
args: tuple[Id, ...]
@dataclass
class EGraph:
memo: dict[Node, FatId] = field(default_factory=dict) # also should have proof id. p : f x = e which follows from e := f x
uf: list[Id] = field(default_factory=list)
reason : list[Id] = field(default_factory=list)
#pf_memo : dict[Node, Id] = field(default_factory=dict)
pf_index : list[Node] = field(default_factory=list)
def _add_node(self, node: Node) -> Id:
id = self.memo.get(node)
if id is not None:
return self.find(id)
else:
id = len(self.uf)
self.uf.append(id)
self.memo[node] = id
print(f"let e{str(id)} := {node.f} {" ".join([f"e{n}" for n in node.args])}")
pid = len(self.pf_index)
pid = len(self.pf_index)
#self.pf_memo[node] = pid
self.pf_index.append(("refl",)) # e := f x
self.reason.append(pid)
print(f"let p{pid} : e{str(id)} = e{str(id)} := Eq.refl")
return id
def add_term(self, term: Term, subst={}) -> Id:
match term:
case Var(name):
return subst[name]
case App(f, args):
arg_ids = tuple(self.add_term(arg, subst) for arg in args)
node = Node(f, arg_ids)
return self._add_node(node)
case _:
raise ValueError(f"Unexpected term: {term}")
def find(self, id: Id) -> Id:
trail = [id]
while self.uf[id] != id:
id = self.uf[id]
trail.append(id)
root = id
# path compression
for i in range(len(trail) - 3, -1, -1):
id, nextid = trail[i], trail[i+1]
pid = len(self.pf_index)
print(f"let p{pid} : e{id} = e{root} := Eq.trans p{self.reason[id]} p{self.reason[nextid]}")
self.pf_index.append(("trans", self.reason[id], self.reason[nextid]))
self.uf[id] = root
self.reason[id] = pid
return root
def _union(self, id1: Id, id2: Id, reason : PId):
a, b = self.find(id1), self.find(id2)
if a != b:
self.uf[a] = b
self.reason[a] = reason
return True
return False
def union(self, t: Term, u: Term):
a,b = self.add_term(t), self.add_term(u)
if self._union(a,b):
pid = len(self.pf_index)
self.pf_index.append(("user_union", a, b))
print(f"fun p{pid} : e{a} = e{b} =>")
def nodes_in_class(self, id: Id) -> list[Node]:
id = self.find(id)
return [obj for obj, obj_id in self.memo.items() if self.find(obj_id) == id]
def _is_eq(self, a: Id, b: Id) -> bool:
return self.find(a) == self.find(b)
def is_eq(self, a: Term, b: Term) -> bool:
return self._is_eq(self.add_term(a), self.add_term(b))
def canonize_node(self, node: Node) -> Node:
return Node(f=node.f, args=tuple(self.find(arg) for arg in node.args))
def rebuild(self):
while True:
l = len(self.memo)
copy_memo = self.memo.copy()
self.memo = {}
for node, id in copy_memo.items():
id = self.find(id)
new_node = self.canonize_node(node)
new_id = self.memo.get(new_node)
if new_id is not None:
self._union(new_id, id)
else:
self.memo[new_node] = id
if l == len(self.memo):
return
# Based on https://github.com/mwillsey/microegg
from dataclasses import dataclass, field
from typing import Callable
type Id = int
type PId = int
type FatId = tuple[PId, Id]
class Term:...
@dataclass(frozen=True)
class App(Term):
f: object
args: tuple[Term, ...]
@dataclass(frozen=True)
class Var(Term): name: str
type Subst = dict[str, Id]
@dataclass(frozen=True)
class Node:
f: object
args: tuple[Id, ...]
@dataclass
class EGraph:
memo: dict[Node, FatId] = field(default_factory=dict) # also should have proof id. p : f x = e which follows from e := f x
uf: list[Id] = field(default_factory=list)
reason : list[Id] = field(default_factory=list)
#pf_memo : dict[Node, Id] = field(default_factory=dict)
pf_index : list[Node] = field(default_factory=list)
def _add_node(self, node: Node) -> Id:
id = self.memo.get(node)
if id is not None:
return self.find(id)
else:
id = len(self.uf)
self.uf.append(id)
self.memo[node] = id
print(f"let e{str(id)} := {node.f} {" ".join([f"e{n}" for n in node.args])}")
pid = len(self.pf_index)
pid = len(self.pf_index)
#self.pf_memo[node] = pid
self.pf_index.append(("refl",)) # e := f x
self.reason.append(pid)
print(f"let p{pid} : e{str(id)} = e{str(id)} := Eq.refl")
return id
def add_term(self, term: Term, subst={}) -> Id:
match term:
case Var(name):
return subst[name]
case App(f, args):
arg_ids = tuple(self.add_term(arg, subst) for arg in args)
node = Node(f, arg_ids)
return self._add_node(node)
case _:
raise ValueError(f"Unexpected term: {term}")
def find(self, id: Id) -> Id:
trail = [id]
while self.uf[id] != id:
id = self.uf[id]
trail.append(id)
root = id
# path compression
for i in range(len(trail) - 3, -1, -1):
id, nextid = trail[i], trail[i+1]
pid = len(self.pf_index)
print(f"let p{pid} : e{id} = e{root} := Eq.trans p{self.reason[id]} p{self.reason[nextid]}")
self.pf_index.append(("trans", self.reason[id], self.reason[nextid]))
self.uf[id] = root
self.reason[id] = pid
return root
def _union(self, id1: Id, id2: Id, reason : PId):
a, b = self.find(id1), self.find(id2)
if a != b:
self.uf[a] = b
self.reason[a] = reason
return True
return False
def union(self, t: Term, u: Term):
a,b = self.add_term(t), self.add_term(u)
if self._union(a,b):
pid = len(self.pf_index)
self.pf_index.append(("user_union", a, b))
print(f"fun p{pid} : e{a} = e{b} =>")
def nodes_in_class(self, id: Id) -> list[Node]:
id = self.find(id)
return [obj for obj, obj_id in self.memo.items() if self.find(obj_id) == id]
def _is_eq(self, a: Id, b: Id) -> bool:
return self.find(a) == self.find(b)
def is_eq(self, a: Term, b: Term) -> bool:
return self._is_eq(self.add_term(a), self.add_term(b))
def canonize_node(self, node: Node) -> Node:
return Node(f=node.f, args=tuple(self.find(arg) for arg in node.args))
def rebuild(self):
while True:
l = len(self.memo)
copy_memo = self.memo.copy()
self.memo = {}
for node, id in copy_memo.items():
id = self.find(id)
new_node = self.canonize_node(node)
new_id = self.memo.get(new_node)
if new_id is not None:
self._union(new_id, id)
else:
self.memo[new_node] = id
if l == len(self.memo):
return
def ematch(self, pattern: Term, id: Id) -> list[Subst]:
return self.ematch_rec(pattern, id, {})
def ematch_rec(self, pattern: Term, id: Id, subst: Subst) -> list[Subst]:
id = self.find(id)
match pattern:
case Var(name):
if name in subst:
if self._is_eq(subst[name], id):
return [subst]
else:
return []
else:
return [{**subst, name: id}]
case App(f, args):
results = []
for node in self.nodes_in_class(id):
if node.f == f and len(node.args) == len(args):
todo = [subst]
for arg_pattern, arg_id in zip(args, node.args):
todo = [
subst1
for subst0 in todo
for subst1 in self.ematch_rec(
arg_pattern, arg_id, subst0
)
]
results.extend(todo)
return results
case _:
raise ValueError(f"Unexpected pattern: {pattern}")
def rw(self, lhs: Term, rhs: Term):
substs = []
for id in range(len(self.uf)):
substs.extend(self.ematch(lhs, id))
# It is surprisingly important to do ematching in two stages of search and apply. Avoids blowup during ematching.
for subst in substs:
lhs_id = self.add_term(lhs, subst)
rhs_id = self.add_term(rhs, subst)
self._union(lhs_id, rhs_id)
self.rebuild()
E = EGraph()
a = App("a", ())
b = App("b", ())
f = lambda x: App("f", (x,))
#a = E.add_term(App("a", ()))
#b = E.add_term(App("b", ()))
#fa = E.add_term(App("f", (a,)))
fa = E.add_term(f(a))
fb = E.add_term(f(b))
E.union(a, b)
E.rebuild()
let e0 := a
let p0 : e0 = e0 := Eq.refl
let e1 := f e0
let p1 : e1 = e1 := Eq.refl
let e2 := b
let p2 : e2 = e2 := Eq.refl
let e3 := f e2
let p3 : e3 = e3 := Eq.refl
fun p4 : e0 = e2 =>