From d070fc80d9f5b9d736e46fa8d0df1ffcd4d2fbd4 Mon Sep 17 00:00:00 2001
From: Anjali Pal <anjalip@cs.washington.edu>
Date: Tue, 23 Jan 2024 22:20:40 -0800
Subject: [PATCH 01/20] add context + let

---
 tree_unique_args/src/type_analysis.egg | 208 +++++++++++++------------
 tree_unique_args/src/type_analysis.rs  | 100 +++++++-----
 2 files changed, 169 insertions(+), 139 deletions(-)

diff --git a/tree_unique_args/src/type_analysis.egg b/tree_unique_args/src/type_analysis.egg
index 3412183d3..a2f81b2e9 100644
--- a/tree_unique_args/src/type_analysis.egg
+++ b/tree_unique_args/src/type_analysis.egg
@@ -26,225 +26,241 @@
 (rule ((= (TypeList-suffix list n) (TNil)))
     ((set (TypeList-length list) n)) :ruleset type-analysis)
 
-(relation HasTypeDemand (Expr))
+(relation HasTypeDemand (IdSort Expr))
 
-(relation HasType (Expr Type))
+(relation HasType (IdSort Expr Type))
 
 ; Primitives
-(rule ((HasTypeDemand (Num id n)))
-      ((HasType (Num id n) (IntT)))
+(rule ((HasTypeDemand ctx (Num ctx n)))
+      ((HasType ctx (Num ctx n) (IntT)))
       :ruleset type-analysis)
 
-(rule ((HasTypeDemand (Boolean id b)))
-      ((HasType (Boolean id b) (BoolT)))
+(rule ((HasTypeDemand ctx (Boolean ctx b)))
+      ((HasType ctx (Boolean ctx b) (BoolT)))
       :ruleset type-analysis)
 
-(rule ((HasTypeDemand (UnitExpr id)))
-      ((HasType (UnitExpr id) (UnitT)))
+(rule ((HasTypeDemand ctx (UnitExpr ctx)))
+      ((HasType ctx (UnitExpr ctx) (UnitT)))
       :ruleset type-analysis)
 
 
 ; Pure Op Demand
-(rule ((HasTypeDemand (Add x y)))
+(rule ((HasTypeDemand ctx (Add x y)))
       (
-       (HasTypeDemand x)
-       (HasTypeDemand y)
+       (HasTypeDemand ctx x)
+       (HasTypeDemand ctx y)
       )
       :ruleset type-analysis)
-(rule ((HasTypeDemand (Sub x y)))
+(rule ((HasTypeDemand ctx (Sub x y)))
       (
-       (HasTypeDemand x)
-       (HasTypeDemand y)
+       (HasTypeDemand ctx x)
+       (HasTypeDemand ctx y)
       )
       :ruleset type-analysis)
-(rule ((HasTypeDemand (Mul x y)))
+(rule ((HasTypeDemand ctx (Mul x y)))
       (
-       (HasTypeDemand x)
-       (HasTypeDemand y)
+       (HasTypeDemand ctx x)
+       (HasTypeDemand ctx y)
       )
       :ruleset type-analysis)
-(rule ((HasTypeDemand (LessThan x y)))
+(rule ((HasTypeDemand ctx (LessThan x y)))
       (
-       (HasTypeDemand x)
-       (HasTypeDemand y)
+       (HasTypeDemand ctx x)
+       (HasTypeDemand ctx y)
       )
       :ruleset type-analysis)
-(rule ((HasTypeDemand (And x y)))
+(rule ((HasTypeDemand ctx (And x y)))
       (
-       (HasTypeDemand x)
-       (HasTypeDemand y)
+       (HasTypeDemand ctx x)
+       (HasTypeDemand ctx y)
       )
       :ruleset type-analysis)
-(rule ((HasTypeDemand (Or x y)))
+(rule ((HasTypeDemand ctx (Or x y)))
       (
-       (HasTypeDemand x)
-       (HasTypeDemand y)
+       (HasTypeDemand ctx x)
+       (HasTypeDemand ctx y)
       )
       :ruleset type-analysis)
-(rule ((HasTypeDemand (Not x)))
-      ((HasTypeDemand x))
+(rule ((HasTypeDemand ctx (Not x)))
+      ((HasTypeDemand ctx x))
       :ruleset type-analysis)
-(rule ((HasTypeDemand (Get e idx)))
-      ((HasTypeDemand e))
+(rule ((HasTypeDemand ctx (Get e idx)))
+      ((HasTypeDemand ctx e))
       :ruleset type-analysis)
 
 ; Pure Op Compute
 (rule (
-       (HasTypeDemand (Add x y))
-       (HasType x (IntT))
-       (HasType y (IntT))
+       (HasTypeDemand ctx (Add x y))
+       (HasType ctx x (IntT))
+       (HasType ctx y (IntT))
       )
       (
-       (HasType (Add x y) (IntT))
+       (HasType ctx (Add x y) (IntT))
       )
       :ruleset type-analysis)
 (rule (
-       (HasTypeDemand (Sub x y))
-       (HasType x (IntT))
-       (HasType y (IntT))
+       (HasTypeDemand ctx (Sub x y))
+       (HasType ctx x (IntT))
+       (HasType ctx y (IntT))
       )
       (
-       (HasType (Sub x y) (IntT))
+       (HasType ctx (Sub x y) (IntT))
       )
       :ruleset type-analysis)
 (rule (
-       (HasTypeDemand (Mul x y))
-       (HasType x (IntT))
-       (HasType y (IntT))
+       (HasTypeDemand ctx (Mul x y))
+       (HasType ctx x (IntT))
+       (HasType ctx y (IntT))
       )
       (
-       (HasType (Mul x y) (IntT))
+       (HasType ctx (Mul x y) (IntT))
       )
       :ruleset type-analysis)
 (rule (
-       (HasTypeDemand (LessThan x y))
-       (HasType x (IntT))
-       (HasType y (IntT))
+       (HasTypeDemand ctx (LessThan x y))
+       (HasType ctx x (IntT))
+       (HasType ctx y (IntT))
       )
       (
-       (HasType (LessThan x y) (BoolT))
+       (HasType ctx (LessThan x y) (BoolT))
       )
       :ruleset type-analysis)
 (rule (
-       (HasTypeDemand (And x y))
-       (HasType x (BoolT))
-       (HasType y (BoolT))
+       (HasTypeDemand ctx (And x y))
+       (HasType ctx x (BoolT))
+       (HasType ctx y (BoolT))
       )
       (
-       (HasType (And x y) (BoolT))
+       (HasType ctx (And x y) (BoolT))
       )
       :ruleset type-analysis)
 (rule (
-       (HasTypeDemand (Or x y))
-       (HasType x (BoolT))
-       (HasType y (BoolT))
+       (HasTypeDemand ctx (Or x y))
+       (HasType ctx x (BoolT))
+       (HasType ctx y (BoolT))
       )
       (
-       (HasType (Or x y) (BoolT))
+       (HasType ctx (Or x y) (BoolT))
       )
       :ruleset type-analysis)
 (rule (
-       (HasTypeDemand (Not x))
-       (HasType x (BoolT))
+       (HasTypeDemand ctx (Not x))
+       (HasType ctx x (BoolT))
       )
       (
-       (HasType (Not x) (BoolT))
+       (HasType ctx (Not x) (BoolT))
       )
       :ruleset type-analysis)
 (rule (
-       (HasTypeDemand (Get e n))
-       (HasType e (TupleT tylist))
+       (HasTypeDemand ctx (Get e n))
+       (HasType ctx e (TupleT tylist))
       )
-      ((HasType (Get e n) (TypeList-ith tylist n)))
+      ((HasType ctx (Get e n) (TypeList-ith tylist n)))
       :ruleset type-analysis)
 
 ; Effectful Ops
-(rule ((HasTypeDemand (Print e)))
-      ((HasType (Print e) (UnitT)))
+(rule ((HasTypeDemand ctx (Print e)))
+      ((HasType ctx (Print e) (UnitT)))
       :ruleset type-analysis)
 ; TODO: Read and Write (requires type annotations)
 
 ; Switch 
 ; if the condition is a boolean, it must have exactly two branches
-(rule ((HasTypeDemand (Switch cond (Cons A (Cons B (Nil))))))
+(rule ((HasTypeDemand ctx (Switch cond (Cons A (Cons B (Nil))))))
       (
-       (HasTypeDemand cond)
-       (HasTypeDemand A)
-       (HasTypeDemand B)
+       (HasTypeDemand ctx cond)
+       (HasTypeDemand ctx A)
+       (HasTypeDemand ctx B)
       )
       :ruleset type-analysis)
 (rule (
        (= switch (Switch cond (Cons A (Cons B (Nil)))))
-       (HasTypeDemand switch)
-       (HasType cond (BoolT))
-       (HasType A ty)
-       (HasType B ty)
+       (HasTypeDemand ctx switch)
+       (HasType ctx cond (BoolT))
+       (HasType ctx A ty)
+       (HasType ctx B ty)
       )
-      ((HasType switch ty))
+      ((HasType ctx switch ty))
       :ruleset type-analysis)
 
 ; Otherwise, the condition must be an integer, and we can have any number of branches.
 
 ; peel off a branch and demand type
-(rule ((HasTypeDemand (Switch cond (Cons branch rest))))
+(rule ((HasTypeDemand ctx (Switch cond (Cons branch rest))))
       (
-       (HasTypeDemand branch)
-       (HasTypeDemand (Switch cond rest))
+       (HasTypeDemand ctx branch)
+       (HasTypeDemand ctx (Switch cond rest))
       )
       :ruleset type-analysis)
 ; base case- demand the type of the condition
-(rule ((HasTypeDemand (Switch cond Nil)))
-      ((HasTypeDemand cond))
+(rule ((HasTypeDemand ctx (Switch cond Nil)))
+      ((HasTypeDemand ctx cond))
       :ruleset type-analysis)
 
 ; base case- single branch switch has type of branch
 (rule (
-       (HasTypeDemand (Switch cond (Cons branch (Nil))))
+       (HasTypeDemand ctx (Switch cond (Cons branch (Nil))))
        ; boolean condition handled above, now we must have an integer condition
-       (HasType cond (IntT))
-       (HasType branch ty)
+       (HasType ctx cond (IntT))
+       (HasType ctx branch ty)
       )
-      ((HasType (Switch cond (Cons branch (Nil))) ty))
+      ((HasType ctx (Switch cond (Cons branch (Nil))) ty))
       :ruleset type-analysis)
 ; recursive case
 (rule (
-       (HasTypeDemand (Switch cond (Cons branch rest)))
-       (HasType (Switch cond rest) ty)
+       (HasTypeDemand ctx (Switch cond (Cons branch rest)))
+       (HasType ctx (Switch cond rest) ty)
        ; make sure the condition is an integer
        ; (prevents us from typing boolean switches with >2 branches)
-       (HasType cond (IntT))
-       (HasType branch ty)
+       (HasType ctx cond (IntT))
+       (HasType ctx branch ty)
       )
-      ((HasType (Switch cond (Cons branch rest)) ty))
+      ((HasType ctx (Switch cond (Cons branch rest)) ty))
       :ruleset type-analysis)
 
 ; Sequencing
-(rule ((HasTypeDemand (All ord (Cons hd tl))))
+(rule ((HasTypeDemand ctx (All ord (Cons hd tl))))
       (
-       (HasTypeDemand hd)
-       (HasTypeDemand (All ord tl))
+       (HasTypeDemand ctx hd)
+       (HasTypeDemand ctx (All ord tl))
       )
       :ruleset type-analysis)
 ; base case: Nil
 (rule (
-       (HasTypeDemand (All ord (Nil)))
+       (HasTypeDemand ctx (All ord (Nil)))
       )
-      ((HasType (All ord (Nil)) (TupleT (TNil))))
+      ((HasType ctx (All ord (Nil)) (TupleT (TNil))))
       :ruleset type-analysis)
 ; rec case
 (rule (
-       (HasTypeDemand (All ord (Cons hd tl)))
-       (HasType hd ty)
-       (HasType (All ord tl) (TupleT tylist))
+       (HasTypeDemand ctx (All ord (Cons hd tl)))
+       (HasType ctx hd ty)
+       (HasType ctx (All ord tl) (TupleT tylist))
       )
-      ((HasType (All ord (Cons hd tl)) (TupleT (TCons ty tylist))))
+      ((HasType ctx (All ord (Cons hd tl)) (TupleT (TCons ty tylist))))
       :ruleset type-analysis)
 
 ; If an expr has two different types, panic
 (rule (
-       (HasType e t1)
-       (HasType e t2)
+       (HasType ctx e t1)
+       (HasType ctx e t2)
        (!= t1 t2)
       )
       ((panic "Type Mismatch!"))
+      :ruleset type-analysis)
+
+
+; Lets
+(rule ((HasTypeDemand ctx (Let id in out)))
+      ((HasTypeDemand ctx in))
+      :ruleset type-analysis)
+
+(rule (
+       (HasTypeDemand ctx (Let id in out))
+       (HasType ctx in ty)
+      )
+      (
+       (HasType id in ty) ; assert in has type ty in the let's context
+       (HasTypeDemand id out) ; demand the type of out in the let's context
+      )
       :ruleset type-analysis)
\ No newline at end of file
diff --git a/tree_unique_args/src/type_analysis.rs b/tree_unique_args/src/type_analysis.rs
index c601dda23..6eeefcd79 100644
--- a/tree_unique_args/src/type_analysis.rs
+++ b/tree_unique_args/src/type_analysis.rs
@@ -1,49 +1,50 @@
 #[test]
 fn simple_types() -> Result<(), egglog::Error> {
     let build = "
-        (let id1 (Id (i64-fresh!)))
-        (let id2 (Id (i64-fresh!)))
-        (let n (Add (Num id1 1) (Num id2 2)))
+        (let id (Id (i64-fresh!)))
+        (let n (Add (Num id 1) (Num id 2)))
         (let m (Mul n n))
         (let s (Sub n m))
         (let x (LessThan m n))
         (let y (Not x))
         (let z (And x (Or y y)))
-        (HasTypeDemand s)
-        (HasTypeDemand z)
+        (HasTypeDemand id s)
+        (HasTypeDemand id z)
         ";
     let check = "
     (run-schedule (saturate type-analysis))
 
-    (check (HasType n (IntT)))
-    (check (HasType m (IntT)))
-    (check (HasType s (IntT)))
-    (check (HasType x (BoolT)))
-    (check (HasType y (BoolT)))
-    (check (HasType z (BoolT)))
+    (check (HasType id n (IntT)))
+    (check (HasType id m (IntT)))
+    (check (HasType id s (IntT)))
+    (check (HasType id x (BoolT)))
+    (check (HasType id y (BoolT)))
+    (check (HasType id z (BoolT)))
     ";
     crate::run_test(build, check)
 }
 
+
 #[test]
 fn switch_boolean() -> Result<(), egglog::Error> {
     let build = "
-  (let b1 (Boolean (Id (i64-fresh!)) true))
-  (let n1 (Num (Id (i64-fresh!)) 1))
-  (let n2 (Num (Id (i64-fresh!)) 3))
+  (let id (Id (i64-fresh!)))
+  (let b1 (Boolean id true))
+  (let n1 (Num id 1))
+  (let n2 (Num id 3))
   (let switch
     (Switch (Not (LessThan n1 n2))
             (Cons (Add n1 n1) (Cons (Sub n1 n2) (Nil)))))
   (let wrong_switch
     (Switch b1 (Cons n1 (Cons n2 (Cons n1 (Nil))))))
-  (HasTypeDemand switch)
-  (HasTypeDemand wrong_switch)
+  (HasTypeDemand id switch)
+  (HasTypeDemand id wrong_switch)
   ";
     let check = "
   (run-schedule (saturate type-analysis))
 
-  (check (HasType switch (IntT)))
-  (fail (check (HasType wrong_switch ty))) ; should not be able to type a boolean swith with 3 cases
+  (check (HasType id switch (IntT)))
+  (fail (check (HasType id wrong_switch ty))) ; should not be able to type a boolean swith with 3 cases
   ";
     crate::run_test(build, check)
 }
@@ -51,10 +52,11 @@ fn switch_boolean() -> Result<(), egglog::Error> {
 #[test]
 fn switch_int() -> Result<(), egglog::Error> {
     let build = "
-  (let n1 (Num (Id (i64-fresh!)) 1))
-  (let n2 (Num (Id (i64-fresh!)) 2))
-  (let n3 (Num (Id (i64-fresh!)) 3))
-  (let n4 (Num (Id (i64-fresh!)) 4))
+  (let id (Id (i64-fresh!)))
+  (let n1 (Num id 1))
+  (let n2 (Num id 2))
+  (let n3 (Num id 3))
+  (let n4 (Num id 4))
   (let s1
     (Switch n1
             (Cons (Add n1 n1) (Cons (Sub n1 n2) (Nil)))))
@@ -62,16 +64,16 @@ fn switch_int() -> Result<(), egglog::Error> {
     (Switch (Mul n1 n2) (Cons (LessThan n3 n4) (Nil))))
   (let s3
     (Switch (Sub n2 n2) (Cons (Print n1) (Cons (Print n4) (Cons (Print n3) (Nil))))))  
-  (HasTypeDemand s1)
-  (HasTypeDemand s2)
-  (HasTypeDemand s3)
+  (HasTypeDemand id s1)
+  (HasTypeDemand id s2)
+  (HasTypeDemand id s3)
   ";
     let check = "
   (run-schedule (saturate type-analysis))
 
-  (check (HasType s1 (IntT)))
-  (check (HasType s2 (BoolT)))
-  (check (HasType s3 (UnitT)))
+  (check (HasType id s1 (IntT)))
+  (check (HasType id s2 (BoolT)))
+  (check (HasType id s3 (UnitT)))
   ";
     crate::run_test(build, check)
 }
@@ -79,7 +81,8 @@ fn switch_int() -> Result<(), egglog::Error> {
 #[test]
 fn tuple() -> Result<(), egglog::Error> {
     let build = "
-  (let n (Add (Num (Id (i64-fresh!)) 1) (Num (Id (i64-fresh!)) 2)))
+  (let id (Id (i64-fresh!)))
+  (let n (Add (Num id 1) (Num id 2)))
         (let m (Mul n n))
         (let s (Sub n m))
         (let x (LessThan m n))
@@ -90,32 +93,43 @@ fn tuple() -> Result<(), egglog::Error> {
   (let tup2 (All (Sequential) (Cons z (Nil))))
   (let tup3 (All (Parallel) (Cons x (Cons m (Nil)))))
   (let tup4 (All (Parallel) (Cons tup2 (Cons tup3 (Nil)))))
-  (HasTypeDemand tup1)
-  (HasTypeDemand tup2)
-  (HasTypeDemand tup3)
-  (HasTypeDemand tup4)
+  (HasTypeDemand id tup1)
+  (HasTypeDemand id tup2)
+  (HasTypeDemand id tup3)
+  (HasTypeDemand id tup4)
 
   (let get1 (Get tup3 0))
   (let get2 (Get tup3 1))
   (let get3 (Get (Get tup4 1) 1))
-  (HasTypeDemand get1)
-  (HasTypeDemand get2)
-  (HasTypeDemand get3)
+  (HasTypeDemand id get1)
+  (HasTypeDemand id get2)
+  (HasTypeDemand id get3)
   ";
     let check = "
   (run-schedule (saturate type-analysis))
-  (check (HasType tup1 (TupleT (TNil))))
-  (check (HasType tup2 (TupleT (TCons (BoolT) (TNil)))))
-  (check (HasType tup3 (TupleT (TCons (BoolT) (TCons (IntT) (TNil))))))
-  (check (HasType tup4
+  (check (HasType id tup1 (TupleT (TNil))))
+  (check (HasType id tup2 (TupleT (TCons (BoolT) (TNil)))))
+  (check (HasType id tup3 (TupleT (TCons (BoolT) (TCons (IntT) (TNil))))))
+  (check (HasType id tup4
     (TupleT (TCons (TupleT (TCons (BoolT) (TNil)))
     (TCons (TupleT (TCons (BoolT) (TCons (IntT) (TNil))))
           (TNil))))))
 
   
-  (check (HasType get1 (BoolT)))
-  (check (HasType get2 (IntT)))
-  (check (HasType get3 (IntT)))
+  (check (HasType id get1 (BoolT)))
+  (check (HasType id get2 (IntT)))
+  (check (HasType id get3 (IntT)))
   ";
     crate::run_test(build, check)
 }
+
+#[test]
+fn lets() -> Result<(), egglog::Error> {
+  let build = "
+    (let let-id (Id (i64-fresh!)))
+    (let outer-ctx (Id (i64-fresh!)))
+    (let l (Let let-id (Num outer-ctx 5) (Add (Arg let-id) (Arg let-id))))
+  ";
+  let check = "";
+  crate::run_test(build, check)
+}
\ No newline at end of file

From 1ca4ef39a86c0a3c5854c8ab01ae8ee4160703f4 Mon Sep 17 00:00:00 2001
From: oflatt <oflatt@gmail.com>
Date: Wed, 24 Jan 2024 09:14:25 -0800
Subject: [PATCH 02/20] working on test

---
 src/rvsdg/tree_unique/mod.rs     | 10 +----
 src/rvsdg/tree_unique/to_tree.rs | 66 ++++++++++++++++++++++++++++++++
 2 files changed, 67 insertions(+), 9 deletions(-)
 create mode 100644 src/rvsdg/tree_unique/to_tree.rs

diff --git a/src/rvsdg/tree_unique/mod.rs b/src/rvsdg/tree_unique/mod.rs
index eaa944fa5..525e29854 100644
--- a/src/rvsdg/tree_unique/mod.rs
+++ b/src/rvsdg/tree_unique/mod.rs
@@ -1,9 +1 @@
-//! Convert RVSDG programs to the tree
-//! encoding of programs.
-//! RVSDGs are close to this encoding,
-//! but use a DAG-based semantics.
-//! This means that nodes that are shared
-//! are only computed once.
-//! These shared nodes need to be let-bound so that they are only
-//! computed once in the tree encoded
-//! program.
+pub mod to_tree;
diff --git a/src/rvsdg/tree_unique/to_tree.rs b/src/rvsdg/tree_unique/to_tree.rs
new file mode 100644
index 000000000..54d4eca7f
--- /dev/null
+++ b/src/rvsdg/tree_unique/to_tree.rs
@@ -0,0 +1,66 @@
+//! Convert RVSDG programs to the tree
+//! encoding of programs.
+//! RVSDGs are close to this encoding,
+//! but use a DAG-based semantics.
+//! This means that nodes that are shared
+//! are only computed once.
+//! These shared nodes need to be let-bound so that they are only
+//! computed once in the tree encoded
+//! program.
+
+#[cfg(test)]
+use crate::{cfg::program_to_cfg, rvsdg::cfg_to_rvsdg, util::parse_from_string};
+
+use crate::rvsdg::{RvsdgFunction, RvsdgProgram};
+use tree_unique_args::Expr;
+
+struct FreshIdGen {
+    counter: i64,
+}
+
+impl FreshIdGen {
+    fn new() -> Self {
+        Self { counter: 0 }
+    }
+
+    fn fresh(&mut self) -> i64 {
+        let id = self.counter;
+        self.counter += 1;
+        id
+    }
+}
+
+impl RvsdgProgram {
+    pub fn to_tree_encoding(&self) -> Expr {
+        self.to_tree_encoding_helper(&mut FreshIdGen::new())
+    }
+
+    fn to_tree_encoding_helper(&self, idgen: &mut FreshIdGen) -> Expr {
+        unimplemented!()
+    }
+}
+
+impl RvsdgFunction {
+    pub fn to_tree_encoding(&self) -> Expr {
+        unimplemented!()
+    }
+}
+
+#[test]
+fn simple_translation() {
+    const PROGRAM: &str = r#"
+    @sub() {
+        v0: int = const 1;
+        v1: int = const 2;
+        v2: int = add v0 v1;
+        print v2;
+        print v1;
+    }
+    "#;
+
+    let prog = parse_from_string(PROGRAM);
+    let cfg = program_to_cfg(&prog);
+    let rvsdg = cfg_to_rvsdg(&cfg).unwrap();
+
+    assert_eq!(rvsdg.to_tree_encoding(), Expr::Num(3));
+}

From dfee6409aa4005496d1d58ace9fc00e8d1d3440b Mon Sep 17 00:00:00 2001
From: oflatt <oflatt@gmail.com>
Date: Wed, 24 Jan 2024 10:42:52 -0800
Subject: [PATCH 03/20] add ast builder utility

---
 tree_unique_args/src/ast.rs         | 204 ++++++++++++++++++++++++++++
 tree_unique_args/src/interpreter.rs |  24 ++++
 tree_unique_args/src/lib.rs         |  54 ++++++++
 3 files changed, 282 insertions(+)
 create mode 100644 tree_unique_args/src/ast.rs

diff --git a/tree_unique_args/src/ast.rs b/tree_unique_args/src/ast.rs
new file mode 100644
index 000000000..def00792d
--- /dev/null
+++ b/tree_unique_args/src/ast.rs
@@ -0,0 +1,204 @@
+use crate::{Expr, Expr::*, Id, Order};
+
+pub fn give_fresh_ids(expr: &mut Expr) {
+    let mut id = 1;
+    give_fresh_ids_helper(expr, 0, &mut id);
+}
+
+fn give_fresh_ids_helper(expr: &mut Expr, current_id: i64, fresh_id: &mut i64) {
+    match expr {
+        Loop(id, input, body) => {
+            let new_id = *fresh_id;
+            *fresh_id += 1;
+            *id = Id(new_id);
+            give_fresh_ids_helper(input, current_id, fresh_id);
+            give_fresh_ids_helper(body, new_id, fresh_id);
+        }
+        Let(id, arg, body) => {
+            let new_id = *fresh_id;
+            *fresh_id += 1;
+            *id = Id(new_id);
+            give_fresh_ids_helper(arg, current_id, fresh_id);
+            give_fresh_ids_helper(body, new_id, fresh_id);
+        }
+        Arg(id) => {
+            *id = Id(current_id);
+        }
+        Function(id, body) => {
+            let new_id = *fresh_id;
+            *fresh_id += 1;
+            *id = Id(new_id);
+            give_fresh_ids_helper(body, new_id, fresh_id);
+        }
+        Call(id, arg) => {
+            *id = Id(current_id);
+            give_fresh_ids_helper(arg, current_id, fresh_id);
+        }
+        _ => expr.for_each_child(move |child| give_fresh_ids_helper(child, current_id, fresh_id)),
+    }
+}
+
+pub fn program(args: Vec<Expr>) -> Expr {
+    let mut prog = Program(args);
+    give_fresh_ids(&mut prog);
+    prog
+}
+
+pub fn num(n: i64) -> Expr {
+    Num(n)
+}
+
+pub fn ttrue() -> Expr {
+    Boolean(true)
+}
+pub fn tfalse() -> Expr {
+    Boolean(false)
+}
+
+pub fn unit() -> Expr {
+    Unit
+}
+
+pub fn add(a: Expr, b: Expr) -> Expr {
+    Add(Box::new(a), Box::new(b))
+}
+
+pub fn sub(a: Expr, b: Expr) -> Expr {
+    Sub(Box::new(a), Box::new(b))
+}
+
+pub fn mul(a: Expr, b: Expr) -> Expr {
+    Mul(Box::new(a), Box::new(b))
+}
+
+pub fn lessthan(a: Expr, b: Expr) -> Expr {
+    LessThan(Box::new(a), Box::new(b))
+}
+
+pub fn and(a: Expr, b: Expr) -> Expr {
+    And(Box::new(a), Box::new(b))
+}
+
+pub fn or(a: Expr, b: Expr) -> Expr {
+    Or(Box::new(a), Box::new(b))
+}
+
+pub fn not(a: Expr) -> Expr {
+    Not(Box::new(a))
+}
+
+pub fn get(a: Expr, i: usize) -> Expr {
+    Get(Box::new(a), i)
+}
+
+pub fn concat(a: Expr, b: Expr) -> Expr {
+    Concat(Box::new(a), Box::new(b))
+}
+
+pub fn print(a: Expr) -> Expr {
+    Print(Box::new(a))
+}
+
+pub fn sequence(args: Vec<Expr>) -> Expr {
+    All(Order::Sequential, args)
+}
+
+pub fn parallel(args: Vec<Expr>) -> Expr {
+    All(Order::Parallel, args)
+}
+
+pub fn switch(arg: Expr, cases: Vec<Expr>) -> Expr {
+    Switch(Box::new(arg), cases)
+}
+
+pub fn tloop(input: Expr, body: Expr) -> Expr {
+    Loop(Id(0), Box::new(input), Box::new(body))
+}
+
+pub fn tlet(arg: Expr, body: Expr) -> Expr {
+    Let(Id(0), Box::new(arg), Box::new(body))
+}
+
+pub fn arg() -> Expr {
+    Arg(Id(0))
+}
+
+pub fn function(arg: Expr) -> Expr {
+    Function(Id(0), Box::new(arg))
+}
+
+pub fn call(arg: Expr) -> Expr {
+    Call(Id(0), Box::new(arg))
+}
+
+#[test]
+fn test_gives_nested_ids() {
+    let mut prog = tlet(num(0), tlet(num(1), num(2)));
+    give_fresh_ids(&mut prog);
+    assert_eq!(
+        prog,
+        Let(
+            Id(1),
+            Box::new(Num(0)),
+            Box::new(Let(Id(2), Box::new(Num(1)), Box::new(Num(2))))
+        )
+    );
+}
+
+#[test]
+fn test_gives_loop_ids() {
+    let mut prog = tlet(num(0), tloop(num(1), num(2)));
+    give_fresh_ids(&mut prog);
+    assert_eq!(
+        prog,
+        Let(
+            Id(1),
+            Box::new(Num(0)),
+            Box::new(Loop(Id(2), Box::new(Num(1)), Box::new(Num(2))))
+        )
+    );
+}
+
+#[test]
+fn test_complex_program_ids() {
+    // test a program that includes
+    // a let, a loop, a switch, and a call
+    let prog = program(vec![function(tlet(
+        num(0),
+        tloop(
+            num(1),
+            switch(
+                arg(),
+                vec![
+                    num(2),
+                    call(num(3)),
+                    tlet(num(4), num(5)),
+                    tloop(num(6), num(7)),
+                ],
+            ),
+        ),
+    ))]);
+    assert_eq!(
+        prog,
+        Program(vec![Function(
+            Id(1),
+            Box::new(Let(
+                Id(2),
+                Box::new(Num(0)),
+                Box::new(Loop(
+                    Id(3),
+                    Box::new(Num(1)),
+                    Box::new(Switch(
+                        Box::new(Arg(Id(3))),
+                        vec![
+                            Num(2),
+                            Call(Id(3), Box::new(Num(3))),
+                            Let(Id(4), Box::new(Num(4)), Box::new(Num(5))),
+                            Loop(Id(5), Box::new(Num(6)), Box::new(Num(7))),
+                        ]
+                    ))
+                ))
+            ))
+        )])
+    );
+}
diff --git a/tree_unique_args/src/interpreter.rs b/tree_unique_args/src/interpreter.rs
index 464394f94..b91d04e45 100644
--- a/tree_unique_args/src/interpreter.rs
+++ b/tree_unique_args/src/interpreter.rs
@@ -18,6 +18,7 @@ pub fn typecheck(e: &Expr, arg_ty: &Option<Type>) -> Result<Type, TypeError> {
         }
     };
     match e {
+        Expr::Program(_) => panic!("Found non top level program."),
         Expr::Num(_) => Ok(Type::Num),
         Expr::Boolean(_) => Ok(Type::Boolean),
         Expr::Unit => Ok(Type::Unit),
@@ -50,6 +51,19 @@ pub fn typecheck(e: &Expr, arg_ty: &Option<Type>) -> Result<Type, TypeError> {
                 )),
             }
         }
+        Expr::Concat(tuple_1, tuple_2) => {
+            let ty_tuple_1 = typecheck(tuple_1, arg_ty)?;
+            let ty_tuple_2 = typecheck(tuple_2, arg_ty)?;
+            match (ty_tuple_1.clone(), ty_tuple_2) {
+                (Type::Tuple(tys_1), Type::Tuple(tys_2)) => {
+                    Ok(Type::Tuple(tys_1.into_iter().chain(tys_2).collect()))
+                }
+                _ => Err(TypeError::ExpectedTupleType(
+                    *tuple_1.clone(),
+                    ty_tuple_1.clone(),
+                )),
+            }
+        }
         Expr::Print(e) => {
             // right now, only print nums
             expect_type(e, Type::Num)?;
@@ -117,6 +131,7 @@ pub struct VirtualMachine {
 // assumes e typechecks and that memory is written before read
 pub fn interpret(e: &Expr, arg: &Option<Value>, vm: &mut VirtualMachine) -> Value {
     match e {
+        Expr::Program(_) => todo!("interpret programs"),
         Expr::Num(x) => Value::Num(*x),
         Expr::Boolean(x) => Value::Boolean(*x),
         Expr::Unit => Value::Unit,
@@ -186,6 +201,15 @@ pub fn interpret(e: &Expr, arg: &Option<Value>, vm: &mut VirtualMachine) -> Valu
             };
             vals[*i].clone()
         }
+        Expr::Concat(tuple_1, tuple_2) => {
+            let Value::Tuple(t1) = interpret(tuple_1, arg, vm) else {
+                panic!("concat")
+            };
+            let Value::Tuple(t2) = interpret(tuple_2, arg, vm) else {
+                panic!("concat")
+            };
+            Value::Tuple(t1.into_iter().chain(t2).collect())
+        }
         Expr::Print(e) => {
             let Value::Num(n) = interpret(e, arg, vm) else {
                 panic!("print")
diff --git a/tree_unique_args/src/lib.rs b/tree_unique_args/src/lib.rs
index 12df6f746..2d7ee51a6 100644
--- a/tree_unique_args/src/lib.rs
+++ b/tree_unique_args/src/lib.rs
@@ -1,3 +1,4 @@
+pub mod ast;
 pub(crate) mod body_contains;
 pub(crate) mod conditional_invariant_code_motion;
 pub(crate) mod deep_copy;
@@ -35,6 +36,11 @@ pub enum Expr {
     Or(Box<Expr>, Box<Expr>),
     Not(Box<Expr>),
     Get(Box<Expr>, usize),
+    /// Concat is a convenient built-in way
+    /// to put two tuples together.
+    /// It's not strictly necessary, but
+    /// doing it by constructing a new big tuple is tedius and slow.
+    Concat(Box<Expr>, Box<Expr>),
     Print(Box<Expr>),
     Read(Box<Expr>),
     Write(Box<Expr>, Box<Expr>),
@@ -44,9 +50,57 @@ pub enum Expr {
     Let(Id, Box<Expr>, Box<Expr>),
     Arg(Id),
     Function(Id, Box<Expr>),
+    /// A list of functions, with the first
+    /// being the main function.
+    Program(Vec<Expr>),
     Call(Id, Box<Expr>),
 }
 
+impl Expr {
+    pub fn for_each_child(&mut self, mut func: impl FnMut(&mut Expr)) {
+        match self {
+            Expr::Num(_) | Expr::Boolean(_) | Expr::Unit | Expr::Arg(_) => {}
+            Expr::Add(a, b)
+            | Expr::Sub(a, b)
+            | Expr::Mul(a, b)
+            | Expr::LessThan(a, b)
+            | Expr::And(a, b)
+            | Expr::Or(a, b)
+            | Expr::Concat(a, b)
+            | Expr::Write(a, b) => {
+                func(a);
+                func(b);
+            }
+            Expr::Not(a) | Expr::Print(a) | Expr::Read(a) => {
+                func(a);
+            }
+            Expr::Get(a, _) | Expr::Function(_, a) | Expr::Call(_, a) => {
+                func(a);
+            }
+            Expr::All(_, children) => {
+                for child in children {
+                    func(child);
+                }
+            }
+            Expr::Switch(input, children) => {
+                func(input);
+                for child in children {
+                    func(child);
+                }
+            }
+            Expr::Loop(_, pred, output) | Expr::Let(_, pred, output) => {
+                func(pred);
+                func(output);
+            }
+            Expr::Program(functions) => {
+                for function in functions {
+                    func(function);
+                }
+            }
+        }
+    }
+}
+
 #[derive(Clone, Debug, PartialEq)]
 pub enum Value {
     Num(i64),

From 3e26718844699eeb283e7ac262e9be45adced8c4 Mon Sep 17 00:00:00 2001
From: oflatt <oflatt@gmail.com>
Date: Wed, 24 Jan 2024 10:44:07 -0800
Subject: [PATCH 04/20] remove to_tree changes

---
 src/rvsdg/tree_unique/to_tree.rs | 57 --------------------------------
 1 file changed, 57 deletions(-)

diff --git a/src/rvsdg/tree_unique/to_tree.rs b/src/rvsdg/tree_unique/to_tree.rs
index 54d4eca7f..eaa944fa5 100644
--- a/src/rvsdg/tree_unique/to_tree.rs
+++ b/src/rvsdg/tree_unique/to_tree.rs
@@ -7,60 +7,3 @@
 //! These shared nodes need to be let-bound so that they are only
 //! computed once in the tree encoded
 //! program.
-
-#[cfg(test)]
-use crate::{cfg::program_to_cfg, rvsdg::cfg_to_rvsdg, util::parse_from_string};
-
-use crate::rvsdg::{RvsdgFunction, RvsdgProgram};
-use tree_unique_args::Expr;
-
-struct FreshIdGen {
-    counter: i64,
-}
-
-impl FreshIdGen {
-    fn new() -> Self {
-        Self { counter: 0 }
-    }
-
-    fn fresh(&mut self) -> i64 {
-        let id = self.counter;
-        self.counter += 1;
-        id
-    }
-}
-
-impl RvsdgProgram {
-    pub fn to_tree_encoding(&self) -> Expr {
-        self.to_tree_encoding_helper(&mut FreshIdGen::new())
-    }
-
-    fn to_tree_encoding_helper(&self, idgen: &mut FreshIdGen) -> Expr {
-        unimplemented!()
-    }
-}
-
-impl RvsdgFunction {
-    pub fn to_tree_encoding(&self) -> Expr {
-        unimplemented!()
-    }
-}
-
-#[test]
-fn simple_translation() {
-    const PROGRAM: &str = r#"
-    @sub() {
-        v0: int = const 1;
-        v1: int = const 2;
-        v2: int = add v0 v1;
-        print v2;
-        print v1;
-    }
-    "#;
-
-    let prog = parse_from_string(PROGRAM);
-    let cfg = program_to_cfg(&prog);
-    let rvsdg = cfg_to_rvsdg(&cfg).unwrap();
-
-    assert_eq!(rvsdg.to_tree_encoding(), Expr::Num(3));
-}

From a60406a8ac5fb13f3b3b1e11a26ff6d176ab81ee Mon Sep 17 00:00:00 2001
From: oflatt <oflatt@gmail.com>
Date: Wed, 24 Jan 2024 15:28:38 -0800
Subject: [PATCH 05/20] functions are expr

---
 tree_unique_args/src/schema.egg | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tree_unique_args/src/schema.egg b/tree_unique_args/src/schema.egg
index bda9b308c..c0f553385 100644
--- a/tree_unique_args/src/schema.egg
+++ b/tree_unique_args/src/schema.egg
@@ -68,7 +68,7 @@
 ; ==========================
 
 ;                   f      output
-(relation Function (IdSort Expr))
+(function Function (IdSort Expr) Expr)
 
 ;               f      arg
 (function Call (IdSort Expr) Expr)

From e469d487b5436a64a385719068e956a9cce0c261 Mon Sep 17 00:00:00 2001
From: Anjali Pal <anjalip@cs.washington.edu>
Date: Wed, 24 Jan 2024 19:50:52 -0800
Subject: [PATCH 06/20] finish let

---
 tree_unique_args/src/type_analysis.egg | 10 +++++++++-
 tree_unique_args/src/type_analysis.rs  |  8 +++++++-
 2 files changed, 16 insertions(+), 2 deletions(-)

diff --git a/tree_unique_args/src/type_analysis.egg b/tree_unique_args/src/type_analysis.egg
index a2f81b2e9..c5ce79038 100644
--- a/tree_unique_args/src/type_analysis.egg
+++ b/tree_unique_args/src/type_analysis.egg
@@ -260,7 +260,15 @@
        (HasType ctx in ty)
       )
       (
-       (HasType id in ty) ; assert in has type ty in the let's context
+      ;  (HasType id in ty) ; assert in has type ty in the let's context ; TODO- is this needed?
+       (HasType id (Arg id) ty) ; assert the let's argument has type ty in the let's context
        (HasTypeDemand id out) ; demand the type of out in the let's context
       )
+      :ruleset type-analysis)
+
+(rule (
+       (HasTypeDemand ctx (Let id in out))
+       (HasType id out ty)
+      )
+      ((HasType ctx (Let id in out) ty))
       :ruleset type-analysis)
\ No newline at end of file
diff --git a/tree_unique_args/src/type_analysis.rs b/tree_unique_args/src/type_analysis.rs
index 6eeefcd79..767dbdf16 100644
--- a/tree_unique_args/src/type_analysis.rs
+++ b/tree_unique_args/src/type_analysis.rs
@@ -119,6 +119,7 @@ fn tuple() -> Result<(), egglog::Error> {
   (check (HasType id get1 (BoolT)))
   (check (HasType id get2 (IntT)))
   (check (HasType id get3 (IntT)))
+  
   ";
     crate::run_test(build, check)
 }
@@ -129,7 +130,12 @@ fn lets() -> Result<(), egglog::Error> {
     (let let-id (Id (i64-fresh!)))
     (let outer-ctx (Id (i64-fresh!)))
     (let l (Let let-id (Num outer-ctx 5) (Add (Arg let-id) (Arg let-id))))
+
+    (HasTypeDemand outer-ctx l)
+  ";
+  let check = "
+    (run-schedule (saturate type-analysis))
+    (check (HasType outer-ctx l (IntT)))
   ";
-  let check = "";
   crate::run_test(build, check)
 }
\ No newline at end of file

From ba0f264630e72f918ba675678196f898f3d43467 Mon Sep 17 00:00:00 2001
From: Anjali Pal <anjalip@cs.washington.edu>
Date: Wed, 24 Jan 2024 19:56:03 -0800
Subject: [PATCH 07/20] formatting

---
 tree_unique_args/src/type_analysis.rs | 9 ++++-----
 1 file changed, 4 insertions(+), 5 deletions(-)

diff --git a/tree_unique_args/src/type_analysis.rs b/tree_unique_args/src/type_analysis.rs
index 767dbdf16..bfec5f642 100644
--- a/tree_unique_args/src/type_analysis.rs
+++ b/tree_unique_args/src/type_analysis.rs
@@ -24,7 +24,6 @@ fn simple_types() -> Result<(), egglog::Error> {
     crate::run_test(build, check)
 }
 
-
 #[test]
 fn switch_boolean() -> Result<(), egglog::Error> {
     let build = "
@@ -126,16 +125,16 @@ fn tuple() -> Result<(), egglog::Error> {
 
 #[test]
 fn lets() -> Result<(), egglog::Error> {
-  let build = "
+    let build = "
     (let let-id (Id (i64-fresh!)))
     (let outer-ctx (Id (i64-fresh!)))
     (let l (Let let-id (Num outer-ctx 5) (Add (Arg let-id) (Arg let-id))))
 
     (HasTypeDemand outer-ctx l)
   ";
-  let check = "
+    let check = "
     (run-schedule (saturate type-analysis))
     (check (HasType outer-ctx l (IntT)))
   ";
-  crate::run_test(build, check)
-}
\ No newline at end of file
+    crate::run_test(build, check)
+}

From 35abd937ed479e8170a4671bffb41337ba1bc2ec Mon Sep 17 00:00:00 2001
From: Anjali Pal <anjalip@cs.washington.edu>
Date: Wed, 24 Jan 2024 20:03:06 -0800
Subject: [PATCH 08/20] Another let test

---
 tree_unique_args/src/type_analysis.rs | 11 +++++++++++
 1 file changed, 11 insertions(+)

diff --git a/tree_unique_args/src/type_analysis.rs b/tree_unique_args/src/type_analysis.rs
index bfec5f642..a74655c7c 100644
--- a/tree_unique_args/src/type_analysis.rs
+++ b/tree_unique_args/src/type_analysis.rs
@@ -131,10 +131,21 @@ fn lets() -> Result<(), egglog::Error> {
     (let l (Let let-id (Num outer-ctx 5) (Add (Arg let-id) (Arg let-id))))
 
     (HasTypeDemand outer-ctx l)
+
+    (let outer (Id (i64-fresh!)))
+    (let inner (Id (i64-fresh!)))
+    (let ctx (Id (i64-fresh!)))
+    (let nested
+      (Let outer (Num ctx 3)
+                 (Let inner (All (Parallel) (Cons (Arg outer) (Cons (Num outer 2) (Nil))))
+                            (Add (Get (Arg inner) 0) (Get (Arg inner) 1)))))
+    (HasTypeDemand ctx nested)
   ";
     let check = "
     (run-schedule (saturate type-analysis))
     (check (HasType outer-ctx l (IntT)))
+
+    (check (HasType ctx nested (IntT)))
   ";
     crate::run_test(build, check)
 }

From 98a523131eb92266d60d7c37caa3ca91115969f2 Mon Sep 17 00:00:00 2001
From: Anjali Pal <anjalip@cs.washington.edu>
Date: Wed, 24 Jan 2024 20:38:40 -0800
Subject: [PATCH 09/20] loops!

---
 tree_unique_args/src/type_analysis.egg | 27 +++++++++++++++++++++++++-
 tree_unique_args/src/type_analysis.rs  | 20 +++++++++++++++++++
 2 files changed, 46 insertions(+), 1 deletion(-)

diff --git a/tree_unique_args/src/type_analysis.egg b/tree_unique_args/src/type_analysis.egg
index c5ce79038..e4a2df43d 100644
--- a/tree_unique_args/src/type_analysis.egg
+++ b/tree_unique_args/src/type_analysis.egg
@@ -271,4 +271,29 @@
        (HasType id out ty)
       )
       ((HasType ctx (Let id in out) ty))
-      :ruleset type-analysis)
\ No newline at end of file
+      :ruleset type-analysis)
+
+; Loops
+(rule ((HasTypeDemand ctx (Loop id in pred-out)))
+      ((HasTypeDemand ctx in))
+      :ruleset type-analysis)
+
+(rule (
+       (HasTypeDemand ctx (Loop id in pred-out))
+       (HasType ctx in ty)
+      )
+      (
+       (HasType id (Arg id) ty) ; assert the argument has type ty in the loop's context
+       (HasTypeDemand id pred-out) ; demand the types of the predicate and output in the loop's context
+      )
+      :ruleset type-analysis)
+
+(rule (
+       (HasTypeDemand ctx (Loop id in pred-out))
+       (HasType ctx in ty) ; input type
+       ; pred-out must be a two-element tuple.
+       ; pred must be boolean, output type must match input type
+       (HasType id pred-out (TupleT (TCons BoolT (TCons ty (TNil)))))
+      )
+      ((HasType ctx (Loop id in pred-out) ty)) ; whole loop has type of output
+      :ruleset type-analysis)
diff --git a/tree_unique_args/src/type_analysis.rs b/tree_unique_args/src/type_analysis.rs
index a74655c7c..97723d556 100644
--- a/tree_unique_args/src/type_analysis.rs
+++ b/tree_unique_args/src/type_analysis.rs
@@ -149,3 +149,23 @@ fn lets() -> Result<(), egglog::Error> {
   ";
     crate::run_test(build, check)
 }
+
+#[test]
+fn loops() -> Result<(), egglog::Error> {
+    let build = "
+  (let ctx (Id 0))
+  (let loop-id (Id 1))
+  (let l (Loop loop-id (Num ctx 1)
+    (All (Sequential)
+         (Cons (LessThan (Num loop-id 2) (Num loop-id 3))
+               (Cons (Switch (Boolean loop-id true)
+                             (Cons (Num loop-id 4) (Cons (Num loop-id 5) (Nil))))
+                     (Nil))))))
+  (HasTypeDemand ctx l)
+  ";
+    let check = "
+  (run-schedule (saturate type-analysis))
+  (check (HasType ctx l (IntT)))
+  ";
+    crate::run_test(build, check)
+}

From 6e10397f2fa85d5fe637524cd1db5f4364288f17 Mon Sep 17 00:00:00 2001
From: zhiyuan yan <yanzhiyuan2019@gmail.com>
Date: Thu, 25 Jan 2024 03:48:25 -0800
Subject: [PATCH 10/20] port

---
 tree_unique_args/src/lib.rs    |   2 +-
 tree_unique_args/src/sugar.egg |  12 ++++
 tree_unique_args/src/util.egg  |  11 +++
 tree_unique_args/src/util.rs   | 128 +++++++++++++++++++++++++++++++++
 4 files changed, 152 insertions(+), 1 deletion(-)

diff --git a/tree_unique_args/src/lib.rs b/tree_unique_args/src/lib.rs
index 63116482b..c35df9b8b 100644
--- a/tree_unique_args/src/lib.rs
+++ b/tree_unique_args/src/lib.rs
@@ -84,7 +84,7 @@ pub fn run_test(build: &str, check: &str) -> Result {
             &subst::subst_rules().join("\n"),
             &deep_copy::deep_copy_rules().join("\n"),
             include_str!("sugar.egg"),
-            include_str!("util.egg"),
+            &util::rules().join("\n"),
             &id_analysis::id_analysis_rules().join("\n"),
             // optimizations
             include_str!("simple.egg"),
diff --git a/tree_unique_args/src/sugar.egg b/tree_unique_args/src/sugar.egg
index bf53ed250..eda9206bf 100644
--- a/tree_unique_args/src/sugar.egg
+++ b/tree_unique_args/src/sugar.egg
@@ -4,6 +4,18 @@
          (Cons a (Cons b (Nil)))
          :ruleset always-run)
 
+(function list3 (Expr Expr Expr) ListExpr)
+(rewrite (list3 a b c)
+         (Cons a (Cons b (Cons c (Nil)))) :ruleset always-run)
+
+(function list4 (Expr Expr Expr Expr) ListExpr)
+(rewrite (list4 a b c d)
+         (Cons a (Cons b (Cons c (Cons d (Nil))))) :ruleset always-run)
+
+(function list5 (Expr Expr Expr Expr Expr) ListExpr)
+(rewrite (list5 a b c d e)
+         (Cons a (Cons b (Cons c (Cons d (Cons e (Nil)))))) :ruleset always-run)
+
 (function IgnoreFirst (Expr Expr) Expr)
 (rewrite (IgnoreFirst a b)
          (Get
diff --git a/tree_unique_args/src/util.egg b/tree_unique_args/src/util.egg
index 693b2e691..ed78d5648 100644
--- a/tree_unique_args/src/util.egg
+++ b/tree_unique_args/src/util.egg
@@ -18,3 +18,14 @@
 (rewrite (Append (Nil) e)
     (Cons e (Nil))
     :ruleset always-run)
+
+;; get the ith output of a loop
+(function get-loop-outputs-ith (Expr i64) Expr :unextractable)
+(rule ((= loop (Loop id inputs pred-outputs))
+       (= pred-outputs (All ord1 pred-out-list))
+       (= (All ord2 outputs-list) (ListExpr-ith pred-out-list 1))
+       (= ith-outputs (ListExpr-ith outputs-list i)))
+    ((union (get-loop-outputs-ith loop i) ith-outputs)) :ruleset always-run)
+
+(function Expr-size (Expr) i64 :merge(min old new))
+(function ListExpr-size (ListExpr) i64 :merge(min old new))
diff --git a/tree_unique_args/src/util.rs b/tree_unique_args/src/util.rs
index 146204abe..dbb10dfaf 100644
--- a/tree_unique_args/src/util.rs
+++ b/tree_unique_args/src/util.rs
@@ -1,3 +1,50 @@
+use crate::ir::{Constructor, Purpose};
+use std::iter;
+use strum::IntoEnumIterator;
+
+fn ast_size_for_ctor(ctor: Constructor) -> String {
+    let ctor_pattern = ctor.construct(|field| field.var());
+    let ruleset = " :ruleset always-run";
+    match ctor {
+        // List itself don't count size
+        Constructor::Nil =>  format!("(rule ({ctor_pattern}) ((set (ListExpr-size {ctor_pattern}) 0)) {ruleset})"),
+        Constructor::Cons => format!("(rule ((= list (Cons expr xs)) (= a (Expr-size expr)) (= b (ListExpr-size xs))) ((set (ListExpr-size list) (+ a b))){ruleset})"), 
+        // let Get and All's size = children's size (I prefer not +1 here)
+        Constructor::Get => format!("(rule ((= expr (Get tup i)) (= n (Expr-size tup))) ((set (Expr-size expr) n)) {ruleset})"),
+        Constructor::All => format!("(rule ((= expr (All ord list)) (= n (ListExpr-size list))) ((set (Expr-size expr) n)) {ruleset})"),
+        _ => {
+            let field_pattern = ctor.fields().iter().filter_map(|field| {
+                let sort = field.sort().name();
+                let var = field.var();
+                match field.purpose {
+                    Purpose::CapturedExpr
+                    | Purpose::SubExpr
+                    | Purpose::SubListExpr =>
+                        Some(format!("({sort}-size {var})")),
+                    _ => None
+                }
+            }).collect::<Vec<_>>();
+
+            let len = field_pattern.len();
+            let result_str = field_pattern.join(" ");
+
+            match len {
+                // Num, Bool Arg, UnitExpr for 0
+                0 => format!("(rule ((= expr {ctor_pattern})) ((set (Expr-size expr) 1)) {ruleset})"),
+                1 => format!("(rule ((= expr {ctor_pattern}) (= n {result_str})) ((set (Expr-size expr) (+ 1 n))){ruleset})"),
+                2 => format!("(rule ((= expr {ctor_pattern}) (= sum (+ {result_str}))) ((set (Expr-size expr) (+ 1 sum))){ruleset})"),
+                _ => panic!("Unimplemented") // we don't have ast take three Expr
+            }
+        },
+    }
+}
+
+pub(crate) fn rules() -> Vec<String> {
+    iter::once(include_str!("util.egg").to_string())
+        .chain(Constructor::iter().map(ast_size_for_ctor))
+        .collect::<Vec<_>>()
+}
+
 #[test]
 fn test_list_util() -> Result<(), egglog::Error> {
     let build = &*"
@@ -35,3 +82,84 @@ fn append_test() -> Result<(), egglog::Error> {
 
     crate::run_test(build, check)
 }
+
+#[test]
+fn get_loop_output_ith_test() -> Result<(), egglog::Error> {
+    let build = "
+    (let id1 (Id (i64-fresh!)))
+    (let id-outer (Id (i64-fresh!)))
+    (let loop1
+        (Loop id1
+            (All (Parallel) (Pair (Arg id-outer) (Num id-outer 0)))
+            (All (Sequential) (Pair
+                ; pred
+                (LessThan (Get (Arg id1) 0) (Get (Arg id1) 1))
+                ; output
+                (All (Parallel) (Pair
+                    (Add (Get (Arg id1) 0) (Num id1 1))
+                    (Sub (Get (Arg id1) 1) (Num id1 1))))))))
+    (let out0 (Add (Get (Arg id1) 0) (Num id1 1)))
+    (let out1 (Sub (Get (Arg id1) 1) (Num id1 1)))
+    ";
+
+    let check = "
+        (check (
+            =
+            (get-loop-outputs-ith loop 0)
+            out0
+        ))
+        (check (
+            =
+            (get-loop-outputs-ith loop 1)
+            out1
+        ))
+    ";
+
+    crate::run_test(build, check)
+}
+
+#[test]
+fn ast_size_test() -> Result<(), egglog::Error> {
+    let build = "
+    (let id1 (Id (i64-fresh!)))
+    (let id-outer (Id (i64-fresh!)))
+    (let inv 
+        (Sub (Get (Arg id1) 4)
+            (Mul (Get (Arg id1) 2) 
+                (Switch (Num id1 1) (list4 (Num id1 1)
+                                        (Num id1 2)
+                                        (Num id1 3)
+                                        (Num id1 4))
+                )
+            )
+        ))
+    
+    (let loop
+        (Loop id1
+            (All (Parallel) (list5 (Num id-outer 0)
+                                    (Num id-outer 1)
+                                    (Num id-outer 2)
+                                    (Num id-outer 3)
+                                    (Num id-outer 4)))
+            (All (Sequential) (Pair
+                ; pred
+                (LessThan (Get (Arg id1) 0) (Get (Arg id1) 4))
+                ; output
+                (All (Parallel) 
+                    (list5
+                        (Add (Get (Arg id1) 0) 
+                            inv
+                        )
+                        (Get (Arg id1) 1)
+                        (Get (Arg id1) 2)
+                        (Get (Arg id1) 3)
+                        (Get (Arg id1) 4) ))))))
+    ";
+
+    let check = "
+       (check (= 10 (Expr-size inv)))
+       (check (= 25 (Expr-size loop)))
+    ";
+
+    crate::run_test(build, check)
+}

From 4d6c7db76bebe61c4caa96a28e1668376b38dae4 Mon Sep 17 00:00:00 2001
From: zhiyuan yan <yanzhiyuan2019@gmail.com>
Date: Thu, 25 Jan 2024 03:53:33 -0800
Subject: [PATCH 11/20] port

---
 tree_unique_args/src/util.egg |  8 --------
 tree_unique_args/src/util.rs  | 34 ----------------------------------
 2 files changed, 42 deletions(-)

diff --git a/tree_unique_args/src/util.egg b/tree_unique_args/src/util.egg
index ed78d5648..8e48b56f0 100644
--- a/tree_unique_args/src/util.egg
+++ b/tree_unique_args/src/util.egg
@@ -19,13 +19,5 @@
     (Cons e (Nil))
     :ruleset always-run)
 
-;; get the ith output of a loop
-(function get-loop-outputs-ith (Expr i64) Expr :unextractable)
-(rule ((= loop (Loop id inputs pred-outputs))
-       (= pred-outputs (All ord1 pred-out-list))
-       (= (All ord2 outputs-list) (ListExpr-ith pred-out-list 1))
-       (= ith-outputs (ListExpr-ith outputs-list i)))
-    ((union (get-loop-outputs-ith loop i) ith-outputs)) :ruleset always-run)
-
 (function Expr-size (Expr) i64 :merge(min old new))
 (function ListExpr-size (ListExpr) i64 :merge(min old new))
diff --git a/tree_unique_args/src/util.rs b/tree_unique_args/src/util.rs
index dbb10dfaf..2f32e2b08 100644
--- a/tree_unique_args/src/util.rs
+++ b/tree_unique_args/src/util.rs
@@ -83,40 +83,6 @@ fn append_test() -> Result<(), egglog::Error> {
     crate::run_test(build, check)
 }
 
-#[test]
-fn get_loop_output_ith_test() -> Result<(), egglog::Error> {
-    let build = "
-    (let id1 (Id (i64-fresh!)))
-    (let id-outer (Id (i64-fresh!)))
-    (let loop1
-        (Loop id1
-            (All (Parallel) (Pair (Arg id-outer) (Num id-outer 0)))
-            (All (Sequential) (Pair
-                ; pred
-                (LessThan (Get (Arg id1) 0) (Get (Arg id1) 1))
-                ; output
-                (All (Parallel) (Pair
-                    (Add (Get (Arg id1) 0) (Num id1 1))
-                    (Sub (Get (Arg id1) 1) (Num id1 1))))))))
-    (let out0 (Add (Get (Arg id1) 0) (Num id1 1)))
-    (let out1 (Sub (Get (Arg id1) 1) (Num id1 1)))
-    ";
-
-    let check = "
-        (check (
-            =
-            (get-loop-outputs-ith loop 0)
-            out0
-        ))
-        (check (
-            =
-            (get-loop-outputs-ith loop 1)
-            out1
-        ))
-    ";
-
-    crate::run_test(build, check)
-}
 
 #[test]
 fn ast_size_test() -> Result<(), egglog::Error> {

From bf4cc3f641bb8f6a1b16ab34969fab9e494bf995 Mon Sep 17 00:00:00 2001
From: zhiyuan yan <yanzhiyuan2019@gmail.com>
Date: Thu, 25 Jan 2024 03:54:02 -0800
Subject: [PATCH 12/20] fmt......

---
 tree_unique_args/src/util.rs | 1 -
 1 file changed, 1 deletion(-)

diff --git a/tree_unique_args/src/util.rs b/tree_unique_args/src/util.rs
index 2f32e2b08..fb7f2b382 100644
--- a/tree_unique_args/src/util.rs
+++ b/tree_unique_args/src/util.rs
@@ -83,7 +83,6 @@ fn append_test() -> Result<(), egglog::Error> {
     crate::run_test(build, check)
 }
 
-
 #[test]
 fn ast_size_test() -> Result<(), egglog::Error> {
     let build = "

From 89b93a775bb2887495df78bcb9826c547d3b2999 Mon Sep 17 00:00:00 2001
From: oflatt <oflatt@gmail.com>
Date: Thu, 25 Jan 2024 11:32:44 -0800
Subject: [PATCH 13/20] fix type error

---
 tree_unique_args/src/interpreter.rs | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/tree_unique_args/src/interpreter.rs b/tree_unique_args/src/interpreter.rs
index b91d04e45..4158a827b 100644
--- a/tree_unique_args/src/interpreter.rs
+++ b/tree_unique_args/src/interpreter.rs
@@ -54,10 +54,14 @@ pub fn typecheck(e: &Expr, arg_ty: &Option<Type>) -> Result<Type, TypeError> {
         Expr::Concat(tuple_1, tuple_2) => {
             let ty_tuple_1 = typecheck(tuple_1, arg_ty)?;
             let ty_tuple_2 = typecheck(tuple_2, arg_ty)?;
-            match (ty_tuple_1.clone(), ty_tuple_2) {
+            match (ty_tuple_1.clone(), ty_tuple_2.clone()) {
                 (Type::Tuple(tys_1), Type::Tuple(tys_2)) => {
                     Ok(Type::Tuple(tys_1.into_iter().chain(tys_2).collect()))
                 }
+                (Type::Tuple(_tys_1), _) => Err(TypeError::ExpectedTupleType(
+                    *tuple_2.clone(),
+                    ty_tuple_2.clone(),
+                )),
                 _ => Err(TypeError::ExpectedTupleType(
                     *tuple_1.clone(),
                     ty_tuple_1.clone(),

From f63856828849f850d1dfed63183b297d4b902b01 Mon Sep 17 00:00:00 2001
From: oflatt <oflatt@gmail.com>
Date: Thu, 25 Jan 2024 11:35:00 -0800
Subject: [PATCH 14/20] add doc comment about func

---
 tree_unique_args/src/lib.rs | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tree_unique_args/src/lib.rs b/tree_unique_args/src/lib.rs
index 2d7ee51a6..b9723b3aa 100644
--- a/tree_unique_args/src/lib.rs
+++ b/tree_unique_args/src/lib.rs
@@ -57,6 +57,7 @@ pub enum Expr {
 }
 
 impl Expr {
+    /// Runs `func` on every child of this expression.
     pub fn for_each_child(&mut self, mut func: impl FnMut(&mut Expr)) {
         match self {
             Expr::Num(_) | Expr::Boolean(_) | Expr::Unit | Expr::Arg(_) => {}

From f237bd6ba2a96a28b119151fff428fb9938177a4 Mon Sep 17 00:00:00 2001
From: Anjali Pal <anjalip@cs.washington.edu>
Date: Thu, 25 Jan 2024 14:10:07 -0800
Subject: [PATCH 15/20] rip out context

---
 tree_unique_args/src/type_analysis.egg | 211 ++++++++++++-------------
 tree_unique_args/src/type_analysis.rs  |  11 +-
 2 files changed, 109 insertions(+), 113 deletions(-)

diff --git a/tree_unique_args/src/type_analysis.egg b/tree_unique_args/src/type_analysis.egg
index c5ce79038..1049e8fba 100644
--- a/tree_unique_args/src/type_analysis.egg
+++ b/tree_unique_args/src/type_analysis.egg
@@ -26,224 +26,224 @@
 (rule ((= (TypeList-suffix list n) (TNil)))
     ((set (TypeList-length list) n)) :ruleset type-analysis)
 
-(relation HasTypeDemand (IdSort Expr))
+(relation HasTypeDemand (Expr))
 
-(relation HasType (IdSort Expr Type))
+(relation HasType (Expr Type))
 
 ; Primitives
-(rule ((HasTypeDemand ctx (Num ctx n)))
-      ((HasType ctx (Num ctx n) (IntT)))
+(rule ((HasTypeDemand (Num id n)))
+      ((HasType (Num id n) (IntT)))
       :ruleset type-analysis)
 
-(rule ((HasTypeDemand ctx (Boolean ctx b)))
-      ((HasType ctx (Boolean ctx b) (BoolT)))
+(rule ((HasTypeDemand (Boolean id b)))
+      ((HasType (Boolean id b) (BoolT)))
       :ruleset type-analysis)
 
-(rule ((HasTypeDemand ctx (UnitExpr ctx)))
-      ((HasType ctx (UnitExpr ctx) (UnitT)))
+(rule ((HasTypeDemand (UnitExpr id)))
+      ((HasType (UnitExpr id) (UnitT)))
       :ruleset type-analysis)
 
 
 ; Pure Op Demand
-(rule ((HasTypeDemand ctx (Add x y)))
+(rule ((HasTypeDemand (Add x y)))
       (
-       (HasTypeDemand ctx x)
-       (HasTypeDemand ctx y)
+       (HasTypeDemand x)
+       (HasTypeDemand y)
       )
       :ruleset type-analysis)
-(rule ((HasTypeDemand ctx (Sub x y)))
+(rule ((HasTypeDemand (Sub x y)))
       (
-       (HasTypeDemand ctx x)
-       (HasTypeDemand ctx y)
+       (HasTypeDemand x)
+       (HasTypeDemand y)
       )
       :ruleset type-analysis)
-(rule ((HasTypeDemand ctx (Mul x y)))
+(rule ((HasTypeDemand (Mul x y)))
       (
-       (HasTypeDemand ctx x)
-       (HasTypeDemand ctx y)
+       (HasTypeDemand x)
+       (HasTypeDemand y)
       )
       :ruleset type-analysis)
-(rule ((HasTypeDemand ctx (LessThan x y)))
+(rule ((HasTypeDemand (LessThan x y)))
       (
-       (HasTypeDemand ctx x)
-       (HasTypeDemand ctx y)
+       (HasTypeDemand x)
+       (HasTypeDemand y)
       )
       :ruleset type-analysis)
-(rule ((HasTypeDemand ctx (And x y)))
+(rule ((HasTypeDemand (And x y)))
       (
-       (HasTypeDemand ctx x)
-       (HasTypeDemand ctx y)
+       (HasTypeDemand x)
+       (HasTypeDemand y)
       )
       :ruleset type-analysis)
-(rule ((HasTypeDemand ctx (Or x y)))
+(rule ((HasTypeDemand (Or x y)))
       (
-       (HasTypeDemand ctx x)
-       (HasTypeDemand ctx y)
+       (HasTypeDemand x)
+       (HasTypeDemand y)
       )
       :ruleset type-analysis)
-(rule ((HasTypeDemand ctx (Not x)))
-      ((HasTypeDemand ctx x))
+(rule ((HasTypeDemand (Not x)))
+      ((HasTypeDemand x))
       :ruleset type-analysis)
-(rule ((HasTypeDemand ctx (Get e idx)))
-      ((HasTypeDemand ctx e))
+(rule ((HasTypeDemand (Get e idx)))
+      ((HasTypeDemand e))
       :ruleset type-analysis)
 
 ; Pure Op Compute
 (rule (
-       (HasTypeDemand ctx (Add x y))
-       (HasType ctx x (IntT))
-       (HasType ctx y (IntT))
+       (HasTypeDemand (Add x y))
+       (HasType x (IntT))
+       (HasType y (IntT))
       )
       (
-       (HasType ctx (Add x y) (IntT))
+       (HasType (Add x y) (IntT))
       )
       :ruleset type-analysis)
 (rule (
-       (HasTypeDemand ctx (Sub x y))
-       (HasType ctx x (IntT))
-       (HasType ctx y (IntT))
+       (HasTypeDemand (Sub x y))
+       (HasType x (IntT))
+       (HasType y (IntT))
       )
       (
-       (HasType ctx (Sub x y) (IntT))
+       (HasType (Sub x y) (IntT))
       )
       :ruleset type-analysis)
 (rule (
-       (HasTypeDemand ctx (Mul x y))
-       (HasType ctx x (IntT))
-       (HasType ctx y (IntT))
+       (HasTypeDemand (Mul x y))
+       (HasType x (IntT))
+       (HasType y (IntT))
       )
       (
-       (HasType ctx (Mul x y) (IntT))
+       (HasType (Mul x y) (IntT))
       )
       :ruleset type-analysis)
 (rule (
-       (HasTypeDemand ctx (LessThan x y))
-       (HasType ctx x (IntT))
-       (HasType ctx y (IntT))
+       (HasTypeDemand (LessThan x y))
+       (HasType x (IntT))
+       (HasType y (IntT))
       )
       (
-       (HasType ctx (LessThan x y) (BoolT))
+       (HasType (LessThan x y) (BoolT))
       )
       :ruleset type-analysis)
 (rule (
-       (HasTypeDemand ctx (And x y))
-       (HasType ctx x (BoolT))
-       (HasType ctx y (BoolT))
+       (HasTypeDemand (And x y))
+       (HasType x (BoolT))
+       (HasType y (BoolT))
       )
       (
-       (HasType ctx (And x y) (BoolT))
+       (HasType (And x y) (BoolT))
       )
       :ruleset type-analysis)
 (rule (
-       (HasTypeDemand ctx (Or x y))
-       (HasType ctx x (BoolT))
-       (HasType ctx y (BoolT))
+       (HasTypeDemand (Or x y))
+       (HasType x (BoolT))
+       (HasType y (BoolT))
       )
       (
-       (HasType ctx (Or x y) (BoolT))
+       (HasType (Or x y) (BoolT))
       )
       :ruleset type-analysis)
 (rule (
-       (HasTypeDemand ctx (Not x))
-       (HasType ctx x (BoolT))
+       (HasTypeDemand (Not x))
+       (HasType x (BoolT))
       )
       (
-       (HasType ctx (Not x) (BoolT))
+       (HasType (Not x) (BoolT))
       )
       :ruleset type-analysis)
 (rule (
-       (HasTypeDemand ctx (Get e n))
-       (HasType ctx e (TupleT tylist))
+       (HasTypeDemand (Get e n))
+       (HasType e (TupleT tylist))
       )
-      ((HasType ctx (Get e n) (TypeList-ith tylist n)))
+      ((HasType (Get e n) (TypeList-ith tylist n)))
       :ruleset type-analysis)
 
 ; Effectful Ops
-(rule ((HasTypeDemand ctx (Print e)))
-      ((HasType ctx (Print e) (UnitT)))
+(rule ((HasTypeDemand (Print e)))
+      ((HasType (Print e) (UnitT)))
       :ruleset type-analysis)
 ; TODO: Read and Write (requires type annotations)
 
 ; Switch 
 ; if the condition is a boolean, it must have exactly two branches
-(rule ((HasTypeDemand ctx (Switch cond (Cons A (Cons B (Nil))))))
+(rule ((HasTypeDemand (Switch cond (Cons A (Cons B (Nil))))))
       (
-       (HasTypeDemand ctx cond)
-       (HasTypeDemand ctx A)
-       (HasTypeDemand ctx B)
+       (HasTypeDemand cond)
+       (HasTypeDemand A)
+       (HasTypeDemand B)
       )
       :ruleset type-analysis)
 (rule (
        (= switch (Switch cond (Cons A (Cons B (Nil)))))
-       (HasTypeDemand ctx switch)
-       (HasType ctx cond (BoolT))
-       (HasType ctx A ty)
-       (HasType ctx B ty)
+       (HasTypeDemand switch)
+       (HasType cond (BoolT))
+       (HasType A ty)
+       (HasType B ty)
       )
-      ((HasType ctx switch ty))
+      ((HasType switch ty))
       :ruleset type-analysis)
 
 ; Otherwise, the condition must be an integer, and we can have any number of branches.
 
 ; peel off a branch and demand type
-(rule ((HasTypeDemand ctx (Switch cond (Cons branch rest))))
+(rule ((HasTypeDemand (Switch cond (Cons branch rest))))
       (
-       (HasTypeDemand ctx branch)
-       (HasTypeDemand ctx (Switch cond rest))
+       (HasTypeDemand branch)
+       (HasTypeDemand (Switch cond rest))
       )
       :ruleset type-analysis)
 ; base case- demand the type of the condition
-(rule ((HasTypeDemand ctx (Switch cond Nil)))
-      ((HasTypeDemand ctx cond))
+(rule ((HasTypeDemand (Switch cond Nil)))
+      ((HasTypeDemand cond))
       :ruleset type-analysis)
 
 ; base case- single branch switch has type of branch
 (rule (
-       (HasTypeDemand ctx (Switch cond (Cons branch (Nil))))
+       (HasTypeDemand (Switch cond (Cons branch (Nil))))
        ; boolean condition handled above, now we must have an integer condition
-       (HasType ctx cond (IntT))
-       (HasType ctx branch ty)
+       (HasType cond (IntT))
+       (HasType branch ty)
       )
-      ((HasType ctx (Switch cond (Cons branch (Nil))) ty))
+      ((HasType (Switch cond (Cons branch (Nil))) ty))
       :ruleset type-analysis)
 ; recursive case
 (rule (
-       (HasTypeDemand ctx (Switch cond (Cons branch rest)))
-       (HasType ctx (Switch cond rest) ty)
+       (HasTypeDemand (Switch cond (Cons branch rest)))
+       (HasType (Switch cond rest) ty)
        ; make sure the condition is an integer
        ; (prevents us from typing boolean switches with >2 branches)
-       (HasType ctx cond (IntT))
-       (HasType ctx branch ty)
+       (HasType cond (IntT))
+       (HasType branch ty)
       )
-      ((HasType ctx (Switch cond (Cons branch rest)) ty))
+      ((HasType (Switch cond (Cons branch rest)) ty))
       :ruleset type-analysis)
 
 ; Sequencing
-(rule ((HasTypeDemand ctx (All ord (Cons hd tl))))
+(rule ((HasTypeDemand (All ord (Cons hd tl))))
       (
-       (HasTypeDemand ctx hd)
-       (HasTypeDemand ctx (All ord tl))
+       (HasTypeDemand hd)
+       (HasTypeDemand (All ord tl))
       )
       :ruleset type-analysis)
 ; base case: Nil
 (rule (
-       (HasTypeDemand ctx (All ord (Nil)))
+       (HasTypeDemand (All ord (Nil)))
       )
-      ((HasType ctx (All ord (Nil)) (TupleT (TNil))))
+      ((HasType (All ord (Nil)) (TupleT (TNil))))
       :ruleset type-analysis)
 ; rec case
 (rule (
-       (HasTypeDemand ctx (All ord (Cons hd tl)))
-       (HasType ctx hd ty)
-       (HasType ctx (All ord tl) (TupleT tylist))
+       (HasTypeDemand (All ord (Cons hd tl)))
+       (HasType hd ty)
+       (HasType (All ord tl) (TupleT tylist))
       )
-      ((HasType ctx (All ord (Cons hd tl)) (TupleT (TCons ty tylist))))
+      ((HasType (All ord (Cons hd tl)) (TupleT (TCons ty tylist))))
       :ruleset type-analysis)
 
 ; If an expr has two different types, panic
 (rule (
-       (HasType ctx e t1)
-       (HasType ctx e t2)
+       (HasType e t1)
+       (HasType e t2)
        (!= t1 t2)
       )
       ((panic "Type Mismatch!"))
@@ -251,24 +251,23 @@
 
 
 ; Lets
-(rule ((HasTypeDemand ctx (Let id in out)))
-      ((HasTypeDemand ctx in))
+(rule ((HasTypeDemand (Let id in out)))
+      ((HasTypeDemand in))
       :ruleset type-analysis)
 
 (rule (
-       (HasTypeDemand ctx (Let id in out))
-       (HasType ctx in ty)
+       (HasTypeDemand (Let id in out))
+       (HasType in ty)
       )
       (
-      ;  (HasType id in ty) ; assert in has type ty in the let's context ; TODO- is this needed?
-       (HasType id (Arg id) ty) ; assert the let's argument has type ty in the let's context
-       (HasTypeDemand id out) ; demand the type of out in the let's context
+       (HasType (Arg id) ty) ; assert the let's argument has type ty in the let's context
+       (HasTypeDemand out) ; demand the type of out in the let's context
       )
       :ruleset type-analysis)
 
 (rule (
-       (HasTypeDemand ctx (Let id in out))
-       (HasType id out ty)
+       (HasTypeDemand (Let id in out))
+       (HasType out ty)
       )
-      ((HasType ctx (Let id in out) ty))
+      ((HasType (Let id in out) ty))
       :ruleset type-analysis)
\ No newline at end of file
diff --git a/tree_unique_args/src/type_analysis.rs b/tree_unique_args/src/type_analysis.rs
index a74655c7c..9683e5272 100644
--- a/tree_unique_args/src/type_analysis.rs
+++ b/tree_unique_args/src/type_analysis.rs
@@ -129,9 +129,7 @@ fn lets() -> Result<(), egglog::Error> {
     (let let-id (Id (i64-fresh!)))
     (let outer-ctx (Id (i64-fresh!)))
     (let l (Let let-id (Num outer-ctx 5) (Add (Arg let-id) (Arg let-id))))
-
-    (HasTypeDemand outer-ctx l)
-
+    (HasTypeDemand l)
     (let outer (Id (i64-fresh!)))
     (let inner (Id (i64-fresh!)))
     (let ctx (Id (i64-fresh!)))
@@ -139,13 +137,12 @@ fn lets() -> Result<(), egglog::Error> {
       (Let outer (Num ctx 3)
                  (Let inner (All (Parallel) (Cons (Arg outer) (Cons (Num outer 2) (Nil))))
                             (Add (Get (Arg inner) 0) (Get (Arg inner) 1)))))
-    (HasTypeDemand ctx nested)
+    (HasTypeDemand nested)
   ";
     let check = "
     (run-schedule (saturate type-analysis))
-    (check (HasType outer-ctx l (IntT)))
-
-    (check (HasType ctx nested (IntT)))
+    (check (HasType l (IntT)))
+    (check (HasType nested (IntT)))
   ";
     crate::run_test(build, check)
 }

From 732f2fcb5090a974c456f7105f3f6ad30d7b1d4d Mon Sep 17 00:00:00 2001
From: Anjali Pal <anjalip@cs.washington.edu>
Date: Thu, 25 Jan 2024 14:11:53 -0800
Subject: [PATCH 16/20] rip out of tests too

---
 tree_unique_args/src/type_analysis.rs | 89 +++++++++++++--------------
 1 file changed, 43 insertions(+), 46 deletions(-)

diff --git a/tree_unique_args/src/type_analysis.rs b/tree_unique_args/src/type_analysis.rs
index 9683e5272..070f18f88 100644
--- a/tree_unique_args/src/type_analysis.rs
+++ b/tree_unique_args/src/type_analysis.rs
@@ -1,25 +1,26 @@
 #[test]
 fn simple_types() -> Result<(), egglog::Error> {
     let build = "
-        (let id (Id (i64-fresh!)))
-        (let n (Add (Num id 1) (Num id 2)))
+        (let id1 (Id (i64-fresh!)))
+        (let id2 (Id (i64-fresh!)))
+        (let n (Add (Num id1 1) (Num id2 2)))
         (let m (Mul n n))
         (let s (Sub n m))
         (let x (LessThan m n))
         (let y (Not x))
         (let z (And x (Or y y)))
-        (HasTypeDemand id s)
-        (HasTypeDemand id z)
+        (HasTypeDemand s)
+        (HasTypeDemand z)
         ";
     let check = "
     (run-schedule (saturate type-analysis))
 
-    (check (HasType id n (IntT)))
-    (check (HasType id m (IntT)))
-    (check (HasType id s (IntT)))
-    (check (HasType id x (BoolT)))
-    (check (HasType id y (BoolT)))
-    (check (HasType id z (BoolT)))
+    (check (HasType n (IntT)))
+    (check (HasType m (IntT)))
+    (check (HasType s (IntT)))
+    (check (HasType x (BoolT)))
+    (check (HasType y (BoolT)))
+    (check (HasType z (BoolT)))
     ";
     crate::run_test(build, check)
 }
@@ -27,23 +28,22 @@ fn simple_types() -> Result<(), egglog::Error> {
 #[test]
 fn switch_boolean() -> Result<(), egglog::Error> {
     let build = "
-  (let id (Id (i64-fresh!)))
-  (let b1 (Boolean id true))
-  (let n1 (Num id 1))
-  (let n2 (Num id 3))
+  (let b1 (Boolean (Id (i64-fresh!)) true))
+  (let n1 (Num (Id (i64-fresh!)) 1))
+  (let n2 (Num (Id (i64-fresh!)) 3))
   (let switch
     (Switch (Not (LessThan n1 n2))
             (Cons (Add n1 n1) (Cons (Sub n1 n2) (Nil)))))
   (let wrong_switch
     (Switch b1 (Cons n1 (Cons n2 (Cons n1 (Nil))))))
-  (HasTypeDemand id switch)
-  (HasTypeDemand id wrong_switch)
+  (HasTypeDemand switch)
+  (HasTypeDemand wrong_switch)
   ";
     let check = "
   (run-schedule (saturate type-analysis))
 
-  (check (HasType id switch (IntT)))
-  (fail (check (HasType id wrong_switch ty))) ; should not be able to type a boolean swith with 3 cases
+  (check (HasType switch (IntT)))
+  (fail (check (HasType wrong_switch ty))) ; should not be able to type a boolean swith with 3 cases
   ";
     crate::run_test(build, check)
 }
@@ -51,11 +51,10 @@ fn switch_boolean() -> Result<(), egglog::Error> {
 #[test]
 fn switch_int() -> Result<(), egglog::Error> {
     let build = "
-  (let id (Id (i64-fresh!)))
-  (let n1 (Num id 1))
-  (let n2 (Num id 2))
-  (let n3 (Num id 3))
-  (let n4 (Num id 4))
+  (let n1 (Num (Id (i64-fresh!)) 1))
+  (let n2 (Num (Id (i64-fresh!)) 2))
+  (let n3 (Num (Id (i64-fresh!)) 3))
+  (let n4 (Num (Id (i64-fresh!)) 4))
   (let s1
     (Switch n1
             (Cons (Add n1 n1) (Cons (Sub n1 n2) (Nil)))))
@@ -63,16 +62,16 @@ fn switch_int() -> Result<(), egglog::Error> {
     (Switch (Mul n1 n2) (Cons (LessThan n3 n4) (Nil))))
   (let s3
     (Switch (Sub n2 n2) (Cons (Print n1) (Cons (Print n4) (Cons (Print n3) (Nil))))))  
-  (HasTypeDemand id s1)
-  (HasTypeDemand id s2)
-  (HasTypeDemand id s3)
+  (HasTypeDemand s1)
+  (HasTypeDemand s2)
+  (HasTypeDemand s3)
   ";
     let check = "
   (run-schedule (saturate type-analysis))
 
-  (check (HasType id s1 (IntT)))
-  (check (HasType id s2 (BoolT)))
-  (check (HasType id s3 (UnitT)))
+  (check (HasType s1 (IntT)))
+  (check (HasType s2 (BoolT)))
+  (check (HasType s3 (UnitT)))
   ";
     crate::run_test(build, check)
 }
@@ -80,8 +79,7 @@ fn switch_int() -> Result<(), egglog::Error> {
 #[test]
 fn tuple() -> Result<(), egglog::Error> {
     let build = "
-  (let id (Id (i64-fresh!)))
-  (let n (Add (Num id 1) (Num id 2)))
+  (let n (Add (Num (Id (i64-fresh!)) 1) (Num (Id (i64-fresh!)) 2)))
         (let m (Mul n n))
         (let s (Sub n m))
         (let x (LessThan m n))
@@ -92,33 +90,32 @@ fn tuple() -> Result<(), egglog::Error> {
   (let tup2 (All (Sequential) (Cons z (Nil))))
   (let tup3 (All (Parallel) (Cons x (Cons m (Nil)))))
   (let tup4 (All (Parallel) (Cons tup2 (Cons tup3 (Nil)))))
-  (HasTypeDemand id tup1)
-  (HasTypeDemand id tup2)
-  (HasTypeDemand id tup3)
-  (HasTypeDemand id tup4)
+  (HasTypeDemand tup1)
+  (HasTypeDemand tup2)
+  (HasTypeDemand tup3)
+  (HasTypeDemand tup4)
 
   (let get1 (Get tup3 0))
   (let get2 (Get tup3 1))
   (let get3 (Get (Get tup4 1) 1))
-  (HasTypeDemand id get1)
-  (HasTypeDemand id get2)
-  (HasTypeDemand id get3)
+  (HasTypeDemand get1)
+  (HasTypeDemand get2)
+  (HasTypeDemand get3)
   ";
     let check = "
   (run-schedule (saturate type-analysis))
-  (check (HasType id tup1 (TupleT (TNil))))
-  (check (HasType id tup2 (TupleT (TCons (BoolT) (TNil)))))
-  (check (HasType id tup3 (TupleT (TCons (BoolT) (TCons (IntT) (TNil))))))
-  (check (HasType id tup4
+  (check (HasType tup1 (TupleT (TNil))))
+  (check (HasType tup2 (TupleT (TCons (BoolT) (TNil)))))
+  (check (HasType tup3 (TupleT (TCons (BoolT) (TCons (IntT) (TNil))))))
+  (check (HasType tup4
     (TupleT (TCons (TupleT (TCons (BoolT) (TNil)))
     (TCons (TupleT (TCons (BoolT) (TCons (IntT) (TNil))))
           (TNil))))))
 
   
-  (check (HasType id get1 (BoolT)))
-  (check (HasType id get2 (IntT)))
-  (check (HasType id get3 (IntT)))
-  
+  (check (HasType get1 (BoolT)))
+  (check (HasType get2 (IntT)))
+  (check (HasType get3 (IntT)))
   ";
     crate::run_test(build, check)
 }

From cc8d174c913dc42bcb57b19808b59773eb615f96 Mon Sep 17 00:00:00 2001
From: Anjali Pal <anjalip@cs.washington.edu>
Date: Thu, 25 Jan 2024 14:18:56 -0800
Subject: [PATCH 17/20] fix up loops

---
 tree_unique_args/src/type_analysis.egg | 20 ++++++++++----------
 tree_unique_args/src/type_analysis.rs  |  4 ++--
 2 files changed, 12 insertions(+), 12 deletions(-)

diff --git a/tree_unique_args/src/type_analysis.egg b/tree_unique_args/src/type_analysis.egg
index 79a8ceeeb..97102357f 100644
--- a/tree_unique_args/src/type_analysis.egg
+++ b/tree_unique_args/src/type_analysis.egg
@@ -273,26 +273,26 @@
       :ruleset type-analysis)
 
 ; Loops
-(rule ((HasTypeDemand ctx (Loop id in pred-out)))
-      ((HasTypeDemand ctx in))
+(rule ((HasTypeDemand (Loop id in pred-out)))
+      ((HasTypeDemand in))
       :ruleset type-analysis)
 
 (rule (
-       (HasTypeDemand ctx (Loop id in pred-out))
-       (HasType ctx in ty)
+       (HasTypeDemand (Loop id in pred-out))
+       (HasType in ty)
       )
       (
-       (HasType id (Arg id) ty) ; assert the argument has type ty in the loop's context
-       (HasTypeDemand id pred-out) ; demand the types of the predicate and output in the loop's context
+       (HasType (Arg id) ty) ; assert the argument has type ty in the loop's context
+       (HasTypeDemand pred-out) ; demand the types of the predicate and output in the loop's context
       )
       :ruleset type-analysis)
 
 (rule (
-       (HasTypeDemand ctx (Loop id in pred-out))
-       (HasType ctx in ty) ; input type
+       (HasTypeDemand (Loop id in pred-out))
+       (HasType in ty) ; input type
        ; pred-out must be a two-element tuple.
        ; pred must be boolean, output type must match input type
-       (HasType id pred-out (TupleT (TCons BoolT (TCons ty (TNil)))))
+       (HasType pred-out (TupleT (TCons BoolT (TCons ty (TNil)))))
       )
-      ((HasType ctx (Loop id in pred-out) ty)) ; whole loop has type of output
+      ((HasType (Loop id in pred-out) ty)) ; whole loop has type of output
       :ruleset type-analysis)
diff --git a/tree_unique_args/src/type_analysis.rs b/tree_unique_args/src/type_analysis.rs
index 5d43a3ea7..b4fff817c 100644
--- a/tree_unique_args/src/type_analysis.rs
+++ b/tree_unique_args/src/type_analysis.rs
@@ -155,11 +155,11 @@ fn loops() -> Result<(), egglog::Error> {
                (Cons (Switch (Boolean loop-id true)
                              (Cons (Num loop-id 4) (Cons (Num loop-id 5) (Nil))))
                      (Nil))))))
-  (HasTypeDemand ctx l)
+  (HasTypeDemand l)
   ";
     let check = "
   (run-schedule (saturate type-analysis))
-  (check (HasType ctx l (IntT)))
+  (check (HasType l (IntT)))
   ";
     crate::run_test(build, check)
 }

From ff75c8a7c04a59759fb89d17f1e0e22e8ac2dd16 Mon Sep 17 00:00:00 2001
From: Anjali Pal <anjalip@cs.washington.edu>
Date: Thu, 25 Jan 2024 14:34:21 -0800
Subject: [PATCH 18/20] loop type errors

---
 tree_unique_args/src/type_analysis.egg | 19 +++++++++-
 tree_unique_args/src/type_analysis.rs  | 52 ++++++++++++++++++++++++++
 2 files changed, 70 insertions(+), 1 deletion(-)

diff --git a/tree_unique_args/src/type_analysis.egg b/tree_unique_args/src/type_analysis.egg
index 97102357f..02ae13917 100644
--- a/tree_unique_args/src/type_analysis.egg
+++ b/tree_unique_args/src/type_analysis.egg
@@ -292,7 +292,24 @@
        (HasType in ty) ; input type
        ; pred-out must be a two-element tuple.
        ; pred must be boolean, output type must match input type
-       (HasType pred-out (TupleT (TCons BoolT (TCons ty (TNil)))))
+       (HasType pred-out (TupleT (TCons (BoolT) (TCons ty (TNil)))))
       )
       ((HasType (Loop id in pred-out) ty)) ; whole loop has type of output
       :ruleset type-analysis)
+
+(rule (
+       (HasTypeDemand (Loop id in pred-out))
+       (HasType pred-out (TupleT (TCons pred-ty rest)))
+       (!= pred-ty (BoolT))
+      )
+      ((panic "Loop predicate was not a boolean"))
+      :ruleset type-analysis)
+
+(rule (
+       (HasTypeDemand (Loop id in pred-out))
+       (HasType in in-ty)
+       (HasType pred-out (TupleT lst))
+       (!= (TypeList-length lst) 2)
+      )
+      ((panic "Loop did not get two arguments (predicate and output)"))
+      :ruleset type-analysis)
\ No newline at end of file
diff --git a/tree_unique_args/src/type_analysis.rs b/tree_unique_args/src/type_analysis.rs
index b4fff817c..04f95de8c 100644
--- a/tree_unique_args/src/type_analysis.rs
+++ b/tree_unique_args/src/type_analysis.rs
@@ -163,3 +163,55 @@ fn loops() -> Result<(), egglog::Error> {
   ";
     crate::run_test(build, check)
 }
+
+#[test]
+#[should_panic]
+fn loop_pred_boolean() {
+    let build = "
+  (let ctx (Id 0))
+  (let loop-id (Id 1))
+  (let l (Loop loop-id (Num ctx 1)
+    (All (Sequential)
+        (Cons (Add (Num loop-id 2) (Num loop-id 3))
+              (Cons (Switch (Boolean loop-id true)
+                            (Cons (Num loop-id 4) (Cons (Num loop-id 5) (Nil))))
+                    (Nil))))))
+  (HasTypeDemand l)
+  (run-schedule (saturate type-analysis))";
+    let check = "";
+
+    let _ = crate::run_test(build, check);
+}
+
+#[test]
+#[should_panic]
+fn loop_args1() {
+    let build = "
+  (let ctx (Id 0))
+  (let loop-id (Id 1))
+  (let l (Loop loop-id (Num ctx 1) (All (Sequential) (Nil))))
+  (HasTypeDemand l)
+  (run-schedule (saturate type-analysis))";
+    let check = "";
+
+    let _ = crate::run_test(build, check);
+}
+
+#[test]
+#[should_panic]
+fn loop_args3() {
+    let build = "
+  (let ctx (Id 0))
+  (let loop-id (Id 1))
+  (let l (Loop loop-id (Num ctx 1)
+    (All (Sequential)
+        (Cons (LessThan (Num loop-id 2) (Num loop-id 3))
+              (Cons (Switch (Boolean loop-id true)
+                            (Cons (Num loop-id 4) (Cons (Num loop-id 5) (Nil))))
+                    (Cons (Num loop-id 1) (Nil)))))))
+  (HasTypeDemand l)
+  (run-schedule (saturate type-analysis))";
+    let check = "";
+
+    let _ = crate::run_test(build, check);
+}

From 14f70dca1495c42185a0a8892c6f29493dcc474d Mon Sep 17 00:00:00 2001
From: Anjali Pal <anjalip@cs.washington.edu>
Date: Thu, 25 Jan 2024 14:51:27 -0800
Subject: [PATCH 19/20] rip out HasTypeDemand

---
 tree_unique_args/src/type_analysis.egg | 128 +++++--------------------
 tree_unique_args/src/type_analysis.rs  |  20 ----
 2 files changed, 26 insertions(+), 122 deletions(-)

diff --git a/tree_unique_args/src/type_analysis.egg b/tree_unique_args/src/type_analysis.egg
index 02ae13917..eaab2f245 100644
--- a/tree_unique_args/src/type_analysis.egg
+++ b/tree_unique_args/src/type_analysis.egg
@@ -26,71 +26,24 @@
 (rule ((= (TypeList-suffix list n) (TNil)))
     ((set (TypeList-length list) n)) :ruleset type-analysis)
 
-(relation HasTypeDemand (Expr))
-
 (relation HasType (Expr Type))
 
 ; Primitives
-(rule ((HasTypeDemand (Num id n)))
+(rule ((Num id n))
       ((HasType (Num id n) (IntT)))
       :ruleset type-analysis)
 
-(rule ((HasTypeDemand (Boolean id b)))
+(rule ((Boolean id b))
       ((HasType (Boolean id b) (BoolT)))
       :ruleset type-analysis)
 
-(rule ((HasTypeDemand (UnitExpr id)))
+(rule ((UnitExpr id))
       ((HasType (UnitExpr id) (UnitT)))
       :ruleset type-analysis)
 
-
-; Pure Op Demand
-(rule ((HasTypeDemand (Add x y)))
-      (
-       (HasTypeDemand x)
-       (HasTypeDemand y)
-      )
-      :ruleset type-analysis)
-(rule ((HasTypeDemand (Sub x y)))
-      (
-       (HasTypeDemand x)
-       (HasTypeDemand y)
-      )
-      :ruleset type-analysis)
-(rule ((HasTypeDemand (Mul x y)))
-      (
-       (HasTypeDemand x)
-       (HasTypeDemand y)
-      )
-      :ruleset type-analysis)
-(rule ((HasTypeDemand (LessThan x y)))
-      (
-       (HasTypeDemand x)
-       (HasTypeDemand y)
-      )
-      :ruleset type-analysis)
-(rule ((HasTypeDemand (And x y)))
-      (
-       (HasTypeDemand x)
-       (HasTypeDemand y)
-      )
-      :ruleset type-analysis)
-(rule ((HasTypeDemand (Or x y)))
-      (
-       (HasTypeDemand x)
-       (HasTypeDemand y)
-      )
-      :ruleset type-analysis)
-(rule ((HasTypeDemand (Not x)))
-      ((HasTypeDemand x))
-      :ruleset type-analysis)
-(rule ((HasTypeDemand (Get e idx)))
-      ((HasTypeDemand e))
-      :ruleset type-analysis)
-
 ; Pure Op Compute
 (rule (
-       (HasTypeDemand (Add x y))
+       (Add x y)
        (HasType x (IntT))
        (HasType y (IntT))
       )
@@ -99,7 +52,7 @@
       )
       :ruleset type-analysis)
 (rule (
-       (HasTypeDemand (Sub x y))
+       (Sub x y)
        (HasType x (IntT))
        (HasType y (IntT))
       )
@@ -108,7 +61,7 @@
       )
       :ruleset type-analysis)
 (rule (
-       (HasTypeDemand (Mul x y))
+       (Mul x y)
        (HasType x (IntT))
        (HasType y (IntT))
       )
@@ -117,7 +70,7 @@
       )
       :ruleset type-analysis)
 (rule (
-       (HasTypeDemand (LessThan x y))
+       (LessThan x y)
        (HasType x (IntT))
        (HasType y (IntT))
       )
@@ -126,7 +79,7 @@
       )
       :ruleset type-analysis)
 (rule (
-       (HasTypeDemand (And x y))
+       (And x y)
        (HasType x (BoolT))
        (HasType y (BoolT))
       )
@@ -135,7 +88,7 @@
       )
       :ruleset type-analysis)
 (rule (
-       (HasTypeDemand (Or x y))
+       (Or x y)
        (HasType x (BoolT))
        (HasType y (BoolT))
       )
@@ -144,7 +97,7 @@
       )
       :ruleset type-analysis)
 (rule (
-       (HasTypeDemand (Not x))
+       (Not x)
        (HasType x (BoolT))
       )
       (
@@ -152,30 +105,22 @@
       )
       :ruleset type-analysis)
 (rule (
-       (HasTypeDemand (Get e n))
+       (Get e n)
        (HasType e (TupleT tylist))
       )
       ((HasType (Get e n) (TypeList-ith tylist n)))
       :ruleset type-analysis)
 
 ; Effectful Ops
-(rule ((HasTypeDemand (Print e)))
+(rule ((Print e))
       ((HasType (Print e) (UnitT)))
       :ruleset type-analysis)
 ; TODO: Read and Write (requires type annotations)
 
 ; Switch 
 ; if the condition is a boolean, it must have exactly two branches
-(rule ((HasTypeDemand (Switch cond (Cons A (Cons B (Nil))))))
-      (
-       (HasTypeDemand cond)
-       (HasTypeDemand A)
-       (HasTypeDemand B)
-      )
-      :ruleset type-analysis)
 (rule (
        (= switch (Switch cond (Cons A (Cons B (Nil)))))
-       (HasTypeDemand switch)
        (HasType cond (BoolT))
        (HasType A ty)
        (HasType B ty)
@@ -185,21 +130,13 @@
 
 ; Otherwise, the condition must be an integer, and we can have any number of branches.
 
-; peel off a branch and demand type
-(rule ((HasTypeDemand (Switch cond (Cons branch rest))))
-      (
-       (HasTypeDemand branch)
-       (HasTypeDemand (Switch cond rest))
-      )
-      :ruleset type-analysis)
-; base case- demand the type of the condition
-(rule ((HasTypeDemand (Switch cond Nil)))
-      ((HasTypeDemand cond))
+(rule ((Switch cond (Cons branch rest)))
+      ((Switch cond rest))                ; peel off a branch for type checking
       :ruleset type-analysis)
 
 ; base case- single branch switch has type of branch
 (rule (
-       (HasTypeDemand (Switch cond (Cons branch (Nil))))
+       (Switch cond (Cons branch (Nil)))
        ; boolean condition handled above, now we must have an integer condition
        (HasType cond (IntT))
        (HasType branch ty)
@@ -208,7 +145,7 @@
       :ruleset type-analysis)
 ; recursive case
 (rule (
-       (HasTypeDemand (Switch cond (Cons branch rest)))
+       (Switch cond (Cons branch rest))
        (HasType (Switch cond rest) ty)
        ; make sure the condition is an integer
        ; (prevents us from typing boolean switches with >2 branches)
@@ -219,21 +156,16 @@
       :ruleset type-analysis)
 
 ; Sequencing
-(rule ((HasTypeDemand (All ord (Cons hd tl))))
-      (
-       (HasTypeDemand hd)
-       (HasTypeDemand (All ord tl))
-      )
+(rule ((All ord (Cons hd tl)))
+      ((All ord tl))                ; peel off a layer for type checking
       :ruleset type-analysis)
 ; base case: Nil
-(rule (
-       (HasTypeDemand (All ord (Nil)))
-      )
+(rule ((All ord (Nil)))
       ((HasType (All ord (Nil)) (TupleT (TNil))))
       :ruleset type-analysis)
 ; rec case
 (rule (
-       (HasTypeDemand (All ord (Cons hd tl)))
+       (All ord (Cons hd tl))
        (HasType hd ty)
        (HasType (All ord tl) (TupleT tylist))
       )
@@ -251,44 +183,36 @@
 
 
 ; Lets
-(rule ((HasTypeDemand (Let id in out)))
-      ((HasTypeDemand in))
-      :ruleset type-analysis)
 
 (rule (
-       (HasTypeDemand (Let id in out))
+       (Let id in out)
        (HasType in ty)
       )
       (
        (HasType (Arg id) ty) ; assert the let's argument has type ty in the let's context
-       (HasTypeDemand out) ; demand the type of out in the let's context
       )
       :ruleset type-analysis)
 
 (rule (
-       (HasTypeDemand (Let id in out))
+       (Let id in out)
        (HasType out ty)
       )
       ((HasType (Let id in out) ty))
       :ruleset type-analysis)
 
 ; Loops
-(rule ((HasTypeDemand (Loop id in pred-out)))
-      ((HasTypeDemand in))
-      :ruleset type-analysis)
 
 (rule (
-       (HasTypeDemand (Loop id in pred-out))
+       (Loop id in pred-out)
        (HasType in ty)
       )
       (
        (HasType (Arg id) ty) ; assert the argument has type ty in the loop's context
-       (HasTypeDemand pred-out) ; demand the types of the predicate and output in the loop's context
       )
       :ruleset type-analysis)
 
 (rule (
-       (HasTypeDemand (Loop id in pred-out))
+       (Loop id in pred-out)
        (HasType in ty) ; input type
        ; pred-out must be a two-element tuple.
        ; pred must be boolean, output type must match input type
@@ -298,7 +222,7 @@
       :ruleset type-analysis)
 
 (rule (
-       (HasTypeDemand (Loop id in pred-out))
+       (Loop id in pred-out)
        (HasType pred-out (TupleT (TCons pred-ty rest)))
        (!= pred-ty (BoolT))
       )
@@ -306,7 +230,7 @@
       :ruleset type-analysis)
 
 (rule (
-       (HasTypeDemand (Loop id in pred-out))
+       (Loop id in pred-out)
        (HasType in in-ty)
        (HasType pred-out (TupleT lst))
        (!= (TypeList-length lst) 2)
diff --git a/tree_unique_args/src/type_analysis.rs b/tree_unique_args/src/type_analysis.rs
index 04f95de8c..982425038 100644
--- a/tree_unique_args/src/type_analysis.rs
+++ b/tree_unique_args/src/type_analysis.rs
@@ -9,8 +9,6 @@ fn simple_types() -> Result<(), egglog::Error> {
         (let x (LessThan m n))
         (let y (Not x))
         (let z (And x (Or y y)))
-        (HasTypeDemand s)
-        (HasTypeDemand z)
         ";
     let check = "
     (run-schedule (saturate type-analysis))
@@ -36,8 +34,6 @@ fn switch_boolean() -> Result<(), egglog::Error> {
             (Cons (Add n1 n1) (Cons (Sub n1 n2) (Nil)))))
   (let wrong_switch
     (Switch b1 (Cons n1 (Cons n2 (Cons n1 (Nil))))))
-  (HasTypeDemand switch)
-  (HasTypeDemand wrong_switch)
   ";
     let check = "
   (run-schedule (saturate type-analysis))
@@ -62,9 +58,6 @@ fn switch_int() -> Result<(), egglog::Error> {
     (Switch (Mul n1 n2) (Cons (LessThan n3 n4) (Nil))))
   (let s3
     (Switch (Sub n2 n2) (Cons (Print n1) (Cons (Print n4) (Cons (Print n3) (Nil))))))  
-  (HasTypeDemand s1)
-  (HasTypeDemand s2)
-  (HasTypeDemand s3)
   ";
     let check = "
   (run-schedule (saturate type-analysis))
@@ -90,17 +83,10 @@ fn tuple() -> Result<(), egglog::Error> {
   (let tup2 (All (Sequential) (Cons z (Nil))))
   (let tup3 (All (Parallel) (Cons x (Cons m (Nil)))))
   (let tup4 (All (Parallel) (Cons tup2 (Cons tup3 (Nil)))))
-  (HasTypeDemand tup1)
-  (HasTypeDemand tup2)
-  (HasTypeDemand tup3)
-  (HasTypeDemand tup4)
 
   (let get1 (Get tup3 0))
   (let get2 (Get tup3 1))
   (let get3 (Get (Get tup4 1) 1))
-  (HasTypeDemand get1)
-  (HasTypeDemand get2)
-  (HasTypeDemand get3)
   ";
     let check = "
   (run-schedule (saturate type-analysis))
@@ -126,7 +112,6 @@ fn lets() -> Result<(), egglog::Error> {
     (let let-id (Id (i64-fresh!)))
     (let outer-ctx (Id (i64-fresh!)))
     (let l (Let let-id (Num outer-ctx 5) (Add (Arg let-id) (Arg let-id))))
-    (HasTypeDemand l)
     (let outer (Id (i64-fresh!)))
     (let inner (Id (i64-fresh!)))
     (let ctx (Id (i64-fresh!)))
@@ -134,7 +119,6 @@ fn lets() -> Result<(), egglog::Error> {
       (Let outer (Num ctx 3)
                  (Let inner (All (Parallel) (Cons (Arg outer) (Cons (Num outer 2) (Nil))))
                             (Add (Get (Arg inner) 0) (Get (Arg inner) 1)))))
-    (HasTypeDemand nested)
   ";
     let check = "
     (run-schedule (saturate type-analysis))
@@ -155,7 +139,6 @@ fn loops() -> Result<(), egglog::Error> {
                (Cons (Switch (Boolean loop-id true)
                              (Cons (Num loop-id 4) (Cons (Num loop-id 5) (Nil))))
                      (Nil))))))
-  (HasTypeDemand l)
   ";
     let check = "
   (run-schedule (saturate type-analysis))
@@ -176,7 +159,6 @@ fn loop_pred_boolean() {
               (Cons (Switch (Boolean loop-id true)
                             (Cons (Num loop-id 4) (Cons (Num loop-id 5) (Nil))))
                     (Nil))))))
-  (HasTypeDemand l)
   (run-schedule (saturate type-analysis))";
     let check = "";
 
@@ -190,7 +172,6 @@ fn loop_args1() {
   (let ctx (Id 0))
   (let loop-id (Id 1))
   (let l (Loop loop-id (Num ctx 1) (All (Sequential) (Nil))))
-  (HasTypeDemand l)
   (run-schedule (saturate type-analysis))";
     let check = "";
 
@@ -209,7 +190,6 @@ fn loop_args3() {
               (Cons (Switch (Boolean loop-id true)
                             (Cons (Num loop-id 4) (Cons (Num loop-id 5) (Nil))))
                     (Cons (Num loop-id 1) (Nil)))))))
-  (HasTypeDemand l)
   (run-schedule (saturate type-analysis))";
     let check = "";
 

From dd10eaa5f90ab0184ac08773439ef11197f7302d Mon Sep 17 00:00:00 2001
From: zhiyuan yan <yanzhiyuan2019@gmail.com>
Date: Fri, 26 Jan 2024 03:33:43 -0800
Subject: [PATCH 20/20] small changes

---
 tree_unique_args/src/util.egg | 4 ++--
 tree_unique_args/src/util.rs  | 4 ++--
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/tree_unique_args/src/util.egg b/tree_unique_args/src/util.egg
index 8e48b56f0..e9f802c29 100644
--- a/tree_unique_args/src/util.egg
+++ b/tree_unique_args/src/util.egg
@@ -19,5 +19,5 @@
     (Cons e (Nil))
     :ruleset always-run)
 
-(function Expr-size (Expr) i64 :merge(min old new))
-(function ListExpr-size (ListExpr) i64 :merge(min old new))
+(function Expr-size (Expr) i64 :merge (min old new))
+(function ListExpr-size (ListExpr) i64 :merge (min old new))
diff --git a/tree_unique_args/src/util.rs b/tree_unique_args/src/util.rs
index fb7f2b382..a8df7d90e 100644
--- a/tree_unique_args/src/util.rs
+++ b/tree_unique_args/src/util.rs
@@ -13,7 +13,7 @@ fn ast_size_for_ctor(ctor: Constructor) -> String {
         Constructor::Get => format!("(rule ((= expr (Get tup i)) (= n (Expr-size tup))) ((set (Expr-size expr) n)) {ruleset})"),
         Constructor::All => format!("(rule ((= expr (All ord list)) (= n (ListExpr-size list))) ((set (Expr-size expr) n)) {ruleset})"),
         _ => {
-            let field_pattern = ctor.fields().iter().filter_map(|field| {
+            let field_pattern = ctor.filter_map_fields(|field| {
                 let sort = field.sort().name();
                 let var = field.var();
                 match field.purpose {
@@ -23,7 +23,7 @@ fn ast_size_for_ctor(ctor: Constructor) -> String {
                         Some(format!("({sort}-size {var})")),
                     _ => None
                 }
-            }).collect::<Vec<_>>();
+            });
 
             let len = field_pattern.len();
             let result_str = field_pattern.join(" ");