From d1e5d4129f03259eaa747826c8616cff0fb6f944 Mon Sep 17 00:00:00 2001
From: Jeremy Wall <jeremy@marzhillstudios.com>
Date: Thu, 28 Sep 2023 19:29:18 -0400
Subject: [PATCH] feat: Properly handle closure logic in function inference.

---
 src/ast/typecheck/mod.rs  | 16 +++++++++--
 src/ast/typecheck/test.rs | 60 +++++++++++++++++++++++++++++++++------
 2 files changed, 66 insertions(+), 10 deletions(-)

diff --git a/src/ast/typecheck/mod.rs b/src/ast/typecheck/mod.rs
index 309d28a..5189d91 100644
--- a/src/ast/typecheck/mod.rs
+++ b/src/ast/typecheck/mod.rs
@@ -35,14 +35,21 @@ pub trait DeriveShape {
 impl DeriveShape for FuncDef {
     fn derive_shape(&self, symbol_table: &mut BTreeMap<Rc<str>, Shape>) -> Shape {
         // 1. First set up our symbols.
-        let mut table = self
+        let mut sym_table = self
             .argdefs
             .iter()
             .map(|sym| (sym.val.clone(), Shape::Hole(sym.clone())))
             .collect::<BTreeMap<Rc<str>, Shape>>();
+        sym_table.append(&mut symbol_table.clone());
         // 2.Then determine the shapes of those symbols in our expression.
-        let shape = self.fields.derive_shape(&mut table);
+        let shape = self.fields.derive_shape(&mut sym_table);
         // 3. Finally determine what the return shape can be.
+        // only include the closed over shapes.
+        let table = self
+            .argdefs
+            .iter()
+            .map(|sym| (sym.val.clone(), sym_table.get(&sym.val).unwrap().clone()))
+            .collect::<BTreeMap<Rc<str>, Shape>>();
         Shape::Func(FuncShapeDef {
             args: table,
             ret: shape.into(),
@@ -278,6 +285,11 @@ impl Checker {
         };
     }
 
+    pub fn with_symbol_table(mut self, symbol_table: BTreeMap<Rc<str>, Shape>) -> Self {
+        self.symbol_table = symbol_table;
+        self
+    }
+
     pub fn pop_shape(&mut self) -> Option<Shape> {
         self.shape_stack.pop()
     }
diff --git a/src/ast/typecheck/test.rs b/src/ast/typecheck/test.rs
index 66fd573..3e3a9b1 100644
--- a/src/ast/typecheck/test.rs
+++ b/src/ast/typecheck/test.rs
@@ -24,19 +24,18 @@ macro_rules! assert_type_fail {
 }
 
 macro_rules! assert_type_success {
-    ($e:expr, $shap:expr) => {{
-        let mut checker = Checker::new();
+    ($e:expr, $shape:expr) => {{
+        assert_type_success!($e, $shape, BTreeMap::new());
+    }};
+    ($e:expr, $shape:expr, $sym_table:expr) => {{
+        let mut checker = Checker::new().with_symbol_table($sym_table);
         let mut expr = parse($e.into(), None).unwrap();
         checker.walk_statement_list(expr.iter_mut().collect());
         let maybe_shape = checker.pop_shape();
         let result = checker.result();
         assert!(result.is_ok(), "We expect this to typecheck successfully.");
-        assert!(
-            result.unwrap().is_empty(),
-            "We don't expect a symbol table entry."
-        );
         assert!(maybe_shape.is_some(), "We got a shape out of it");
-        assert_eq!(maybe_shape.unwrap(), $shap);
+        assert_eq!(maybe_shape.unwrap(), $shape);
     }};
 }
 
@@ -224,8 +223,10 @@ fn infer_symbol_type_test() {
 #[test]
 fn infer_func_type_test() {
     let mut args = BTreeMap::new();
+    let foo = Into::<Rc<str>>::into("foo");
+    let bar = Into::<Rc<str>>::into("bar");
     args.insert(
-        Into::<Rc<str>>::into("foo"),
+        foo.clone(),
         Shape::Int(PositionedItem {
             pos: Position {
                 file: None,
@@ -250,4 +251,47 @@ fn infer_func_type_test() {
         "func(foo) => foo + 1;",
         Shape::Func(FuncShapeDef { args, ret })
     );
+    let mut symbol_table = BTreeMap::new();
+    symbol_table.insert(
+        bar.clone(),
+        Shape::Int(PositionedItem {
+            pos: Position {
+                file: None,
+                line: 1,
+                column: 20,
+                offset: 19,
+            },
+            val: 1,
+        }),
+    );
+    let mut args = BTreeMap::new();
+    args.insert(
+        foo.clone(),
+        Shape::Int(PositionedItem {
+            pos: Position {
+                file: None,
+                line: 1,
+                column: 6,
+                offset: 5,
+            },
+            val: 1,
+        }),
+    );
+    assert_type_success!(
+        "func(foo) => foo + bar;",
+        Shape::Func(FuncShapeDef {
+            args: args,
+            ret: Shape::Int(PositionedItem {
+                pos: Position {
+                    file: None,
+                    line: 1,
+                    column: 20,
+                    offset: 19,
+                },
+                val: 1,
+            })
+            .into()
+        }),
+        symbol_table
+    );
 }