From fe7882b598d1583d09b97ff16a68c82806805467 Mon Sep 17 00:00:00 2001
From: Folkert <folkert@folkertdev.nl>
Date: Wed, 13 Jul 2022 16:23:39 +0200
Subject: [PATCH] implement expect collector

---
 crates/compiler/can/src/expr.rs           | 48 +++++++++++++++++++++++
 crates/compiler/can/src/module.rs         |  2 +-
 crates/compiler/can/src/traverse.rs       |  1 +
 crates/compiler/load_internal/src/file.rs |  4 +-
 4 files changed, 53 insertions(+), 2 deletions(-)

diff --git a/crates/compiler/can/src/expr.rs b/crates/compiler/can/src/expr.rs
index b548db2372..1fbf0ade19 100644
--- a/crates/compiler/can/src/expr.rs
+++ b/crates/compiler/can/src/expr.rs
@@ -10,6 +10,7 @@ use crate::num::{
 use crate::pattern::{canonicalize_pattern, BindingsFromPattern, Pattern};
 use crate::procedure::References;
 use crate::scope::Scope;
+use crate::traverse::Visitor;
 use roc_collections::soa::Index;
 use roc_collections::{SendMap, VecMap, VecSet};
 use roc_error_macros::internal_error;
@@ -2243,6 +2244,38 @@ impl Declarations {
                 Some((length_so_far, *e))
             })
     }
+
+    pub fn expects(&self) -> VecMap<Region, Vec<(Symbol, Variable)>> {
+        let mut collector = ExpectCollector {
+            expects: VecMap::default(),
+        };
+
+        let var = Variable::EMPTY_RECORD;
+
+        for index in 0..self.len() {
+            use crate::expr::DeclarationTag::*;
+
+            match self.declarations[index] {
+                Value | Function(_) | Recursive(_) | TailRecursive(_) | Destructure(_) => {
+                    // def pattern has no default expressions, so skip
+                    let loc_expr = &self.expressions[index];
+
+                    collector.visit_expr(&loc_expr.value, loc_expr.region, var);
+                }
+                MutualRecursion { .. } => {
+                    // the self of this group will be treaded individually by later iterations
+                }
+                Expectation => {
+                    let loc_expr =
+                        toplevel_expect_to_inline_expect(self.expressions[index].clone());
+
+                    collector.visit_expr(&loc_expr.value, loc_expr.region, var);
+                }
+            }
+        }
+
+        collector.expects
+    }
 }
 
 roc_error_macros::assert_sizeof_default!(DeclarationTag, 8);
@@ -2470,3 +2503,18 @@ pub fn toplevel_expect_to_inline_expect(mut loc_expr: Loc<Expr>) -> Loc<Expr> {
 
     loc_expr
 }
+
+struct ExpectCollector {
+    expects: VecMap<Region, Vec<(Symbol, Variable)>>,
+}
+
+impl crate::traverse::Visitor for ExpectCollector {
+    fn visit_expr(&mut self, expr: &Expr, region: Region, _var: Variable) {
+        if let Expr::Expect {
+            lookups_in_cond, ..
+        } = expr
+        {
+            self.expects.insert(region, lookups_in_cond.to_vec());
+        }
+    }
+}
diff --git a/crates/compiler/can/src/module.rs b/crates/compiler/can/src/module.rs
index 42d1f7fcda..217d544466 100644
--- a/crates/compiler/can/src/module.rs
+++ b/crates/compiler/can/src/module.rs
@@ -737,7 +737,7 @@ pub fn canonicalize_module_defs<'a>(
         }
     }
 
-    let loc_expects = Default::default();
+    let loc_expects = declarations.expects();
 
     ModuleOutput {
         scope,
diff --git a/crates/compiler/can/src/traverse.rs b/crates/compiler/can/src/traverse.rs
index 4c1d03676d..5f3eb4c50a 100644
--- a/crates/compiler/can/src/traverse.rs
+++ b/crates/compiler/can/src/traverse.rs
@@ -1,5 +1,6 @@
 //! Traversals over the can ast.
 
+use roc_collections::VecMap;
 use roc_module::{ident::Lowercase, symbol::Symbol};
 use roc_region::all::{Loc, Region};
 use roc_types::subs::Variable;
diff --git a/crates/compiler/load_internal/src/file.rs b/crates/compiler/load_internal/src/file.rs
index 9e2f721cce..99443083b5 100644
--- a/crates/compiler/load_internal/src/file.rs
+++ b/crates/compiler/load_internal/src/file.rs
@@ -4137,7 +4137,9 @@ fn run_solve<'a>(
     // TODO remove when we write builtins in roc
     let aliases = module.aliases.clone();
 
+    let mut module = module;
     let loc_expects = std::mem::take(&mut module.loc_expects);
+    let module = module;
 
     let (solved_subs, solved_specializations, exposed_vars_by_symbol, problems, abilities_store) = {
         if module_id.is_builtin() {
@@ -4403,7 +4405,7 @@ fn canonicalize_and_constrain<'a>(
         aliases,
         rigid_variables: module_output.rigid_variables,
         abilities_store: module_output.scope.abilities_store,
-        loc_expects,
+        loc_expects: module_output.loc_expects,
     };
 
     let constrained_module = ConstrainedModule {