Add typing rules for functions/calls

tree_unique_args/src/function_inlining.egg

@@ -3,7 +3,7 @@
 (rule
     (
         (= call (Call f arg))
-        (Function f out)
+        (Function f out in-ty out-ty)
         (ExprIsValid call)
     )
     (

tree_unique_args/src/purity_analysis.rs

@@ -55,7 +55,7 @@ fn purity_rule_for_ctor(ctor: Constructor) -> Option<String> {
 pub(crate) fn purity_analysis_rules() -> Vec<String> {
     ESort::iter()
         .map(|sort| "(relation *IsPure (*))".replace('*', sort.name()))
-        .chain(iter::once( "(relation FunctionIsPure (IdSort))\n(rule ((Function id out) (ExprIsPure out)) ((FunctionIsPure id)):ruleset always-run)".to_string()))
+        .chain(iter::once( "(relation FunctionIsPure (IdSort))\n(rule ((Function id out in-ty out-ty) (ExprIsPure out)) ((FunctionIsPure id)):ruleset always-run)".to_string()))
         .chain(Constructor::iter().filter_map(purity_rule_for_ctor))
         .collect::<Vec<_>>()
 }
@@ -112,7 +112,8 @@ fn test_purity_function() -> Result<(), egglog::Error> {
     (Function id_fun1
             (Add 
                 (Get (Arg id_fun1) 0) 
-                (Get (Arg id_fun1) 0))))
+                (Get (Arg id_fun1) 0))
+            (TupleT (TCons (IntT) (TNil))) (IntT)))
 ;; f2 is impure
 (let f2
     (Function id_fun2
@@ -123,7 +124,8 @@ fn test_purity_function() -> Result<(), egglog::Error> {
                     (Add 
                         (Get (Arg id_fun2) 0) 
                         (Get (Arg id_fun2) 0))))
-            1)))
+            1)
+        (TupleT (TCons (IntT) (TNil))) (IntT)))
 (let pure-loop
     (Loop id1
         (All id-outer (Parallel) (Pair (Num id-outer 0) (Num id-outer 0)))

tree_unique_args/src/schema.egg

@@ -80,8 +80,8 @@
 ; Functions
 ; ==========================
 
-;                   f      output
-(function Function (IdSort Expr) Expr)
+;                   f    output in-type out-type
+(function Function (IdSort Expr Type Type) Expr)
 
 ;               f      arg
 (function Call (IdSort Expr) Expr)

tree_unique_args/src/type_analysis.egg

@@ -211,4 +211,25 @@
        (!= (TypeList-length lst) 2)
       )
       ((panic "Loop did not get two arguments (predicate and output)"))
+      :ruleset type-analysis)
+
+; Functions
+(rule ((= f (Function id body in-ty out-ty)))
+      ((HasType (Arg id) in-ty))
+      :ruleset type-analysis)
+
+(rule (
+       (= f (Function id body in-ty out-ty))
+       (HasType body out-ty)
+      )
+      ((HasType f (FuncT in-ty out-ty)))
+      :ruleset type-analysis)
+
+(rule (
+       (= c (Call f-id arg))
+       (= f (Function f-id body in-ty out-ty))
+       (HasType f (FuncT in-ty out-ty))
+       (HasType arg in-ty)
+      )
+      ((HasType c out-ty))
       :ruleset type-analysis)

tree_unique_args/src/type_analysis.rs

@@ -231,3 +231,27 @@ fn read_write() -> Result<(), egglog::Error> {
     );
     crate::run_test(build, &check)
 }
+
+#[test]
+fn func() -> Result<(), egglog::Error> {
+    let build = "
+  (let f-id (Id (i64-fresh!)))
+  (let ctx  (Id (i64-fresh!)))
+
+(let f (Function f-id (Switch (Get (Arg f-id) 1) 
+(Cons (Add (Get (Arg f-id) 0) (Num f-id 4))
+      (Cons (Get (Arg f-id) 0) (Nil))))
+      (TupleT (TCons (IntT) (TCons (BoolT) (TNil))))
+      (IntT)))
+  (let call (Call f-id (All ctx (Parallel) (Cons (Num ctx 2) (Cons (Boolean ctx true) (Nil))))))
+
+  ";
+    let check = format!(
+        "
+  {SCHED}
+  (check (HasType call (IntT)))
+  (check (HasType f (FuncT (TupleT (TCons (IntT) (TCons (BoolT) (TNil)))) (IntT))))
+  "
+    );
+    crate::run_test(build, &check)
+}