From 7fcccfa9e00b3beeddbafa75948bd0b6bd33e935 Mon Sep 17 00:00:00 2001
From: oflatt <oflatt@gmail.com>
Date: Wed, 7 Feb 2024 14:25:40 -0800
Subject: [PATCH] top down assume list experiment

---
 tree_assume/src/schema.egg         |  2 +
 tree_assume/src/utility/assume.egg | 91 ++++++++++++++++++++++++++++++
 2 files changed, 93 insertions(+)
 create mode 100644 tree_assume/src/utility/assume.egg

diff --git a/tree_assume/src/schema.egg b/tree_assume/src/schema.egg
index 90b94a1c0..32b80b0ca 100644
--- a/tree_assume/src/schema.egg
+++ b/tree_assume/src/schema.egg
@@ -173,6 +173,7 @@
   ; The term is in a loop with `input` and `pred_output`.
   ;      input    pred_output
   (InLoop Expr     Expr)
+  (InFunc String Expr)
   ; Other assumptions are possible, but not supported yet.
   ; For example:
   ;      A boolean predicate is true.
@@ -183,6 +184,7 @@
 ; Assume allows creating context-specific terms.
 ; e.g. (Assume (InLet (Const (Int 2))) (Arg)) is equal to (Const (Int 2))
 (function Assume (Assumption Expr) Expr)
+(function AssumeList (Assumption ListExpr) ListExpr)
 
 
 ; =================================
diff --git a/tree_assume/src/utility/assume.egg b/tree_assume/src/utility/assume.egg
new file mode 100644
index 000000000..5cfe20d42
--- /dev/null
+++ b/tree_assume/src/utility/assume.egg
@@ -0,0 +1,91 @@
+(ruleset assume)
+
+(sort AssumeList)
+;; We collect multiple assumes into one AssumeMulti.
+;; When it gets to the bottom, we flatten it back out to multiple `Assume` nodes
+;; for the convenience of other analyses.
+(function AssumeMulti (AssumeList Expr) Expr :unextractable)
+
+(function AssumeCons (Assumption AssumeList) AssumeList :unextractable)
+(function AssumeNil () AssumeList :unextractable)
+
+; ################### start top-down assumptions
+
+(rewrite
+ (Function name in_ty out_ty output)
+ (Function name in_ty out_ty
+   (AssumeList (AssumeCons (InFunction name) (AssumeNil))
+               output))
+ :ruleset assume)
+
+
+; ################### operations
+(rewrite (AssumeList asum (Bop op c1 c2))
+         (Bop op (AssumeList asum c1) (AssumeList asum c2))
+         :ruleset assume)
+(rewrite (AssumeList assum (Uop op c1))
+         (Uop op (AssumeList assum c1))
+         :ruleset assume)
+(rewrite (AssumeList assum (Get expr index))
+         (Get (AssumeList assum expr) index)
+         :ruleset assume)
+(rewrite (AssumeList assum (Alloc expr ty))
+         (Alloc (AssumeList assum expr) ty)
+         :ruleset assume)
+(rewrite (AssumeList assum (Call name expr))
+         (Call name (AssumeList assum expr))
+         :ruleset assume)
+
+; ################### tuple operations
+(rewrite (AssumeList assum (Single expr))
+         (Single (AssumeList assum expr))
+         :ruleset assume)
+(rewrite (AssumeList assum (Concat order e1 e2))
+         (Concat order (AssumeList assum e1) (AssumeList assum e2))
+         :ruleset assume)
+
+; #################### control flow
+
+(rewrite (AssumeList assum (Switch pred cases))
+         (Switch (AssumeList assum pred)
+                 (AssumeList
+
+(rule ((= lhs (Switch pred cases))
+       (AssumeDepth pred d1)
+       (ListAssumeDepth cases d2))
+      ((set (AssumeDepth lhs) (max d1 d2)))
+      :ruleset assume)
+(rule ((= lhs (If pred then else))
+       (AssumeDepth pred d1)
+       (AssumeDepth then d2)
+       (AssumeDepth else d3))
+      ((set (AssumeDepth lhs) (max d1 (max d2 d3))))
+      :ruleset assume)
+
+
+; Let and DoWhile are the only complex cases
+; When the depth is at the limit, don't substitute
+; into the body and reset the depth.
+
+(rule ((= lhs (Let inputs body))
+       (AssumeDepth inputs inputdepth)
+       (< inputdepth max-assumedepth))
+      ((let newlet
+         (Let inputs (Assume (InLet inputs) body)))
+       (set (AssumeDepth newlet) 
+       
+
+
+
+
+; ####################### lists
+
+(rule ((= lhs (Nil)))
+      ((set (ListAssumeDepth lhs) 0))
+      :ruleset assume)
+(rule ((= lhs (Cons head tail))
+       (AssumeDepth head d1)
+       (ListAssumeDepth tail d2))
+      ((set (ListAssumeDepth lhs) (max d1 d2)))
+      :ruleset assume)
+