Merge pull request #297 from egraphs-good/ajpal-type-annotations

Add Read/Write to Type analysis

tree_unique_args/src/ir.rs

@@ -9,6 +9,7 @@ pub(crate) enum Sort {
     IdSort,
     I64,
     Bool,
+    Type,
 }
 
 impl Sort {
@@ -20,6 +21,7 @@ impl Sort {
             Sort::IdSort => "IdSort",
             Sort::I64 => "i64",
             Sort::Bool => "bool",
+            Sort::Type => "Type",
         }
     }
 }
@@ -161,7 +163,7 @@ impl Constructor {
             Constructor::Not => vec![f(SubExpr, "x")],
             Constructor::Get => vec![f(SubExpr, "tup"), f(Static(Sort::I64), "i")],
             Constructor::Print => vec![f(SubExpr, "printee")],
-            Constructor::Read => vec![f(SubExpr, "addr")],
+            Constructor::Read => vec![f(SubExpr, "addr"), f(Static(Sort::Type), "type")],
             Constructor::Write => vec![f(SubExpr, "addr"), f(SubExpr, "data")],
             Constructor::All => vec![
                 f(ReferencingId, "id"),

tree_unique_args/src/schedule.egg

@@ -10,4 +10,6 @@
     function-inlining
     interval-analysis
     ivt
+    ; can't add the type-analysis rules here yet until we fix the correctness bug
+    ; described in https://github.com/egraphs-good/eggcc/issues/292
     ))

tree_unique_args/src/schema.egg

@@ -5,6 +5,19 @@
 (datatype Expr)
 (datatype ListExpr (Cons Expr ListExpr) (Nil))
 
+(sort TypeList)
+
+(datatype Type
+  (IntT)
+  (BoolT)
+  (FuncT Type Type)
+  (TupleT TypeList)
+  (UnitT)
+)
+
+(function TNil () TypeList)
+(function TCons (Type TypeList) TypeList)
+
 ; ==========================
 ; Pure operators
 ; ==========================
@@ -29,8 +42,8 @@
 
 ;                value      unit
 (function Print (Expr)      Expr)
-;                addr       value
-(function Read  (Expr)      Expr)
+;                addr type      value
+(function Read  (Expr Type)      Expr)
 ;                addr value unit
 (function Write (Expr Expr) Expr)
 

tree_unique_args/src/switch_rewrites.rs

@@ -134,10 +134,10 @@ fn test_constant_condition() -> Result<(), egglog::Error> {
 fn switch_pull_in_below() -> Result<(), egglog::Error> {
     let build = "
     (let id (Id (i64-fresh!)))
-    (let c (Read (Num id 3)))
-    (let s1 (Read (Num id 4)))
-    (let s2 (Read (Num id 5)))
-    (let s3 (Read (Num id 6)))
+    (let c (Read (Num id 3) (IntT)))
+    (let s1 (Read (Num id 4) (IntT)))
+    (let s2 (Read (Num id 5) (IntT)))
+    (let s3 (Read (Num id 6) (IntT)))
 
     (let switch (Switch c (Cons s1 (Cons s2 (Nil)))))
     (let lhs (All id (Sequential) (Cons switch (Cons s3 (Nil)))))

tree_unique_args/src/type_analysis.egg

@@ -1,30 +1,5 @@
 (ruleset type-analysis)
 
-(sort TypeList)
-
-(datatype Type
-  (IntT)
-  (BoolT)
-  (FuncT Type Type)
-  (TupleT TypeList)
-)
-
-(function TNil () TypeList)
-(function TCons (Type TypeList) TypeList)
-
-(function TypeList-length (TypeList) i64)
-(function TypeList-ith (TypeList i64) Type :unextractable)
-(function TypeList-suffix (TypeList i64) TypeList :unextractable)
-
-(rule ((TupleT tylist)) ((union (TypeList-suffix tylist 0) tylist)) :ruleset type-analysis)
-
-(rule ((= (TypeList-suffix top n) (TCons hd tl)))
-    ((union (TypeList-ith top n) hd)
-     (union (TypeList-suffix top (+ n 1)) tl)) :ruleset type-analysis)
-
-(rule ((= (TypeList-suffix list n) (TNil)))
-    ((set (TypeList-length list) n)) :ruleset type-analysis)
-
 (relation HasType (Expr Type))
 
 ; Primitives
@@ -110,7 +85,12 @@
 (rule ((Print e))
       ((HasType (Print e) (TupleT (TNil))))
       :ruleset type-analysis)
-; TODO: Read and Write (requires type annotations)
+(rule ((= lhs (Read addr ty)))
+      ((HasType lhs ty))
+      :ruleset type-analysis)
+(rule ((= lhs (Write addr val)))
+      ((HasType lhs (UnitT)))
+      :ruleset type-analysis)
 
 ; Switch 
 ; if the condition is a boolean, it must have exactly two branches

tree_unique_args/src/type_analysis.rs

@@ -1,3 +1,6 @@
+#[cfg(test)]
+const SCHED: &str = "(run-schedule (repeat 5 (saturate type-analysis) (saturate always-run)))";
+
 #[test]
 fn simple_types() -> Result<(), egglog::Error> {
     let build = "
@@ -10,17 +13,18 @@ fn simple_types() -> Result<(), egglog::Error> {
         (let y (Not x))
         (let z (And x (Or y y)))
         ";
-    let check = "
-    (run-schedule (saturate type-analysis))
-
+    let check = format!(
+        "
+    {SCHED}
     (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)
+    "
+    );
+    crate::run_test(build, &check)
 }
 
 #[test]
@@ -35,13 +39,14 @@ fn switch_boolean() -> Result<(), egglog::Error> {
   (let wrong_switch
     (Switch b1 (Cons n1 (Cons n2 (Cons n1 (Nil))))))
   ";
-    let check = "
-  (run-schedule (saturate type-analysis))
-
+    let check = format!(
+        "
+    {SCHED}
   (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)
+  "
+    );
+    crate::run_test(build, &check)
 }
 
 #[test]
@@ -59,14 +64,15 @@ fn switch_int() -> Result<(), egglog::Error> {
   (let s3
     (Switch (Sub n2 n2) (Cons (Print n1) (Cons (Print n4) (Cons (Print n3) (Nil))))))  
   ";
-    let check = "
-  (run-schedule (saturate type-analysis))
-
+    let check = format!(
+        "
+    {SCHED}
   (check (HasType s1 (IntT)))
   (check (HasType s2 (BoolT)))
   (check (HasType s3 (TupleT (TNil))))
-  ";
-    crate::run_test(build, check)
+  "
+    );
+    crate::run_test(build, &check)
 }
 
 #[test]
@@ -89,8 +95,9 @@ fn tuple() -> Result<(), egglog::Error> {
   (let get2 (Get tup3 1))
   (let get3 (Get (Get tup4 1) 1))
   ";
-    let check = "
-  (run-schedule (saturate type-analysis))
+    let check = format!(
+        "
+    {SCHED}
   (check (HasType tup1 (TupleT (TNil))))
   (check (HasType tup2 (TupleT (TCons (BoolT) (TNil)))))
   (check (HasType tup3 (TupleT (TCons (BoolT) (TCons (IntT) (TNil))))))
@@ -103,8 +110,9 @@ fn tuple() -> Result<(), egglog::Error> {
   (check (HasType get1 (BoolT)))
   (check (HasType get2 (IntT)))
   (check (HasType get3 (IntT)))
-  ";
-    crate::run_test(build, check)
+  "
+    );
+    crate::run_test(build, &check)
 }
 
 #[test]
@@ -121,12 +129,14 @@ fn lets() -> Result<(), egglog::Error> {
                  (Let inner (All ctx (Parallel) (Cons (Arg outer) (Cons (Num outer 2) (Nil))))
                             (Add (Get (Arg inner) 0) (Get (Arg inner) 1)))))
   ";
-    let check = "
-    (run-schedule (saturate type-analysis))
+    let check = format!(
+        "
+  {SCHED}
     (check (HasType l (IntT)))
     (check (HasType nested (IntT)))
-  ";
-    crate::run_test(build, check)
+  "
+    );
+    crate::run_test(build, &check)
 }
 
 #[test]
@@ -141,11 +151,13 @@ fn loops() -> Result<(), egglog::Error> {
                              (Cons (Num loop-id 4) (Cons (Num loop-id 5) (Nil))))
                      (Nil))))))
   ";
-    let check = "
-  (run-schedule (saturate type-analysis))
+    let check = format!(
+        "
+    {SCHED}
   (check (HasType l (IntT)))
-  ";
-    crate::run_test(build, check)
+  "
+    );
+    crate::run_test(build, &check)
 }
 
 #[test]
@@ -159,11 +171,13 @@ fn loop_pred_boolean() {
         (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))))))
-  (run-schedule (saturate type-analysis))";
-    let check = "";
+                    (Nil))))))";
+    let check = format!(
+        "
+    {SCHED}"
+    );
 
-    let _ = crate::run_test(build, check);
+    let _ = crate::run_test(build, &check);
 }
 
 #[test]
@@ -172,11 +186,13 @@ fn loop_args1() {
     let build = "
   (let ctx (Id 0))
   (let loop-id (Id 1))
-  (let l (Loop loop-id (Num ctx 1) (All loop-id (Sequential) (Nil))))
-  (run-schedule (saturate type-analysis))";
-    let check = "";
+  (let l (Loop loop-id (Num ctx 1) (All loop-id (Sequential) (Nil))))";
+    let check = format!(
+        "
+  {SCHED}"
+    );
 
-    let _ = crate::run_test(build, check);
+    let _ = crate::run_test(build, &check);
 }
 
 #[test]
@@ -190,9 +206,28 @@ fn loop_args3() {
         (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)))))))
-  (run-schedule (saturate type-analysis))";
-    let check = "";
+                    (Cons (Num loop-id 1) (Nil)))))))";
+    let check = format!(
+        "
+                    {SCHED}"
+    );
+
+    let _ = crate::run_test(build, &check);
+}
 
-    let _ = crate::run_test(build, check);
+#[test]
+fn read_write() -> Result<(), egglog::Error> {
+    let build = "
+    (let id (Id 0))
+    (let r (Read (Num id 4) (IntT)))
+    (let w (Write (Num id 2) (Num id 45)))
+  ";
+    let check = format!(
+        "
+  {SCHED}
+    (check (HasType r (IntT)))
+    (check (HasType w (UnitT)))
+  "
+    );
+    crate::run_test(build, &check)
 }

tree_unique_args/src/util.egg

@@ -21,3 +21,16 @@
 
 (function Expr-size (Expr) i64 :merge (min old new))
 (function ListExpr-size (ListExpr) i64 :merge (min old new))
+
+(function TypeList-length (TypeList) i64)
+(function TypeList-ith (TypeList i64) Type :unextractable)
+(function TypeList-suffix (TypeList i64) TypeList :unextractable)
+
+(rule ((TupleT tylist)) ((union (TypeList-suffix tylist 0) tylist)) :ruleset always-run)
+
+(rule ((= (TypeList-suffix top n) (TCons hd tl)))
+    ((union (TypeList-ith top n) hd)
+     (union (TypeList-suffix top (+ n 1)) tl)) :ruleset always-run)
+
+(rule ((= (TypeList-suffix list n) (TNil)))
+    ((set (TypeList-length list) n)) :ruleset always-run)