Fix lambda if

Author: saulshanabrook

python/egg_smol/examples/lambda.py

@@ -160,12 +160,24 @@ def freer(t: Term) -> StringSet:
     ),
 )
 
+result = egraph.relation("result")
+
+
+def l(fn: Callable[[Term], Term]) -> Term:  # noqa
+    """
+    Create a lambda term from a function
+    """
+    # Use first var name from fn
+    x = fn.__code__.co_varnames[0]
+    return lam(Var(x), fn(Term.var(Var(x))))
+
 
 def assert_simplifies(left: BaseExpr, right: BaseExpr) -> None:
     """
     Simplify and print
     """
-    res = egraph.simplify(left, 30)
+    with egraph:
+        res = egraph.simplify(left, 30)
     print(f"{left} ➡  {res}")
     assert expr_parts(res) == expr_parts(right), f"{res} != {right}"
 
@@ -174,18 +186,6 @@ def assert_simplifies(left: BaseExpr, right: BaseExpr) -> None:
 assert_simplifies((Term.val(Val(1)) + Term.val(Val(2))).eval(), Val(3))
 
 
-result = egraph.relation("result")
-
-
-def l(fn: Callable[[Term], Term]) -> Term:  # noqa
-    """
-    Create a lambda term from a function
-    """
-    # Use first var name from fn
-    x = fn.__code__.co_varnames[0]
-    return lam(Var(x), fn(Term.var(Var(x))))
-
-
 # lambda under
 assert_simplifies(
     l(lambda x: Term.val(Val(4)) + l(lambda y: y)(Term.val(Val(4)))),
@@ -270,7 +270,7 @@ def l(fn: Callable[[Term], Term]) -> Term:  # noqa
 assert_simplifies(if_(Term.val(Val(1)) == Term.val(Val(1)), Term.val(Val(7)), Term.val(Val(9))), Term.val(Val(7)))
 
 
-# lambda_compose_many
+# # lambda_compose_many
 assert_simplifies(
     let_(
         compose,
@@ -298,7 +298,7 @@ def l(fn: Callable[[Term], Term]) -> Term:  # noqa
     let_(
         zeroone,
         l(lambda x: if_(x == Term.val(Val(0)), Term.val(Val(0)), Term.val(Val(1)))),
-        Term.val(Val(0)) + Term.var(zeroone)(Term.val(Val(0))) + Term.var(zeroone)(Term.val(Val(10))),
+        Term.var(zeroone)(Term.val(Val(0))) + Term.var(zeroone)(Term.val(Val(10))),
     ),
     Term.val(Val(1)),
 )