From 6571d829a116277d45aed64350140c984b2e5fd5 Mon Sep 17 00:00:00 2001
From: rvcas <x@rvcas.dev>
Date: Thu, 22 Apr 2021 21:31:02 -0400
Subject: [PATCH] fix: use bump for Constraint enum instead of pool

---
 editor/src/lang/constrain.rs | 68 ++++++++++++++++++------------------
 editor/tests/solve_expr2.rs  |  1 +
 2 files changed, 35 insertions(+), 34 deletions(-)

diff --git a/editor/src/lang/constrain.rs b/editor/src/lang/constrain.rs
index 2574541e11..03a3ee8d7f 100644
--- a/editor/src/lang/constrain.rs
+++ b/editor/src/lang/constrain.rs
@@ -1,4 +1,6 @@
-use crate::lang::pool::{NodeId, Pool, PoolVec};
+use bumpalo::{collections::Vec as BumpVec, Bump};
+
+use crate::lang::pool::{Pool, PoolVec};
 use crate::lang::{ast::Expr2, expr::Env, types::Type2};
 
 use roc_can::expected::Expected;
@@ -11,70 +13,68 @@ use roc_types::{
 };
 
 #[derive(Debug)]
-pub enum Constraint {
+pub enum Constraint<'a> {
     Eq(Type2, Expected<Type2>, Category, Region),
     // Store(Type, Variable, &'static str, u32),
     // Lookup(Symbol, Expected<Type>, Region),
     // Pattern(Region, PatternCategory, Type, PExpected<Type>),
-    And(PoolVec<Constraint>),
-    Let(NodeId<LetConstraint>),
+    And(BumpVec<'a, Constraint<'a>>),
+    Let(&'a LetConstraint<'a>),
     // SaveTheEnvironment,
     True, // Used for things that always unify, e.g. blanks and runtime errors
 }
 
 #[derive(Debug)]
-pub struct LetConstraint {
-    pub rigid_vars: PoolVec<Variable>,
-    pub flex_vars: PoolVec<Variable>,
+pub struct LetConstraint<'a> {
+    pub rigid_vars: BumpVec<'a, Variable>,
+    pub flex_vars: BumpVec<'a, Variable>,
     pub def_types: SendMap<Symbol, Located<Type2>>,
-    pub defs_constraint: Constraint,
-    pub ret_constraint: Constraint,
+    pub defs_constraint: Constraint<'a>,
+    pub ret_constraint: Constraint<'a>,
 }
 
-pub fn constrain_expr(
+pub fn constrain_expr<'a>(
+    arena: &'a Bump,
     env: &mut Env,
     expr: &Expr2,
     expected: Expected<Type2>,
     region: Region,
-) -> Constraint {
+) -> Constraint<'a> {
     use Constraint::*;
 
     match expr {
         Expr2::Str(_) => Eq(str_type(env.pool), expected, Category::Str, region),
         Expr2::EmptyRecord => Eq(Type2::EmptyRec, expected, Category::Record, region),
-        Expr2::SmallInt { number, var, .. } => {
-            let flex_vars = PoolVec::with_capacity(1, env.pool);
-            let rigid_vars = PoolVec::empty(env.pool);
+        Expr2::SmallInt { var, .. } => {
+            let mut flex_vars = BumpVec::with_capacity_in(1, arena);
+            let rigid_vars = BumpVec::new_in(arena);
 
-            for flex_var_node_id in flex_vars.iter_node_ids() {
-                env.pool[flex_var_node_id] = *var;
-            }
+            let mut and_constraints = BumpVec::with_capacity_in(2, arena);
+
+            flex_vars.push(*var);
 
             let num_type = Type2::Variable(*var);
 
-            let and_nodes = vec![
-                Eq(
-                    num_type,
-                    Expected::ForReason(
-                        Reason::IntLiteral,
-                        Type2::Alias(
-                            Symbol::NUM_INT,
-                            PoolVec::empty(env.pool),
-                            env.pool.add(Type2::EmptyRec),
-                        ),
-                        region,
+            and_constraints.push(Eq(
+                num_type,
+                Expected::ForReason(
+                    Reason::IntLiteral,
+                    Type2::Alias(
+                        Symbol::NUM_INT,
+                        PoolVec::empty(env.pool),
+                        env.pool.add(Type2::EmptyRec),
                     ),
-                    Category::Int,
                     region,
                 ),
-                Eq(Type2::Variable(*var), expected, Category::Int, region),
-            ];
+                Category::Int,
+                region,
+            ));
 
-            let and_constraints = PoolVec::new(and_nodes.into_iter(), env.pool);
+            and_constraints.push(Eq(Type2::Variable(*var), expected, Category::Int, region));
 
             let defs_constraint = And(and_constraints);
 
-            let node_id = env.pool.add(LetConstraint {
+            let let_constraint = arena.alloc(LetConstraint {
                 rigid_vars,
                 flex_vars,
                 def_types: SendMap::default(),
@@ -82,7 +82,7 @@ pub fn constrain_expr(
                 ret_constraint: Constraint::True,
             });
 
-            Let(node_id)
+            Let(let_constraint)
         }
         _ => todo!("implement constaints for {:?}", expr),
     }
diff --git a/editor/tests/solve_expr2.rs b/editor/tests/solve_expr2.rs
index a0e9290a00..9bb8bd03f6 100644
--- a/editor/tests/solve_expr2.rs
+++ b/editor/tests/solve_expr2.rs
@@ -83,6 +83,7 @@ fn infer_eq(actual: &str, expected_str: &str) {
     match expr2_result {
         Ok((expr, _)) => {
             let constraint = constrain_expr(
+                &code_arena,
                 &mut env,
                 &expr,
                 Expected::NoExpectation(Type2::Variable(var)),