Merge pull request #2645 from rtfeldman/solve-fully-tail-recursive

make solve fully tail recursive

compiler/can/src/constraint.rs

@@ -259,16 +259,15 @@ impl Constraints {
         Slice::new(start as _, length as _)
     }
 
-    fn def_types_slice<I>(&mut self, it: I) -> Slice<(Symbol, Loc<Type>)>
+    fn def_types_slice<I>(&mut self, it: I) -> Slice<(Symbol, Loc<Index<Type>>)>
     where
         I: IntoIterator<Item = (Symbol, Loc<Type>)>,
     {
         let start = self.def_types.len();
 
         for (symbol, loc_type) in it {
-            let type_index = Index::new(self.types.len() as _);
             let Loc { region, value } = loc_type;
-            self.types.push(value);
+            let type_index = Index::push_new(&mut self.types, value);
 
             self.def_types.push((symbol, Loc::at(region, type_index)));
         }
@@ -469,7 +468,7 @@ pub enum Constraint {
 pub struct LetConstraint {
     pub rigid_vars: Slice<Variable>,
     pub flex_vars: Slice<Variable>,
-    pub def_types: Slice<(Symbol, Loc<Type>)>,
+    pub def_types: Slice<(Symbol, Loc<Index<Type>>)>,
     pub defs_and_ret_constraint: Index<(Constraint, Constraint)>,
 }
 

compiler/collections/src/soa.rs

@@ -1,11 +1,22 @@
 use std::usize;
 
-#[derive(Clone, Copy, PartialEq, Eq)]
+#[derive(PartialEq, Eq)]
 pub struct Index<T> {
     index: u32,
     _marker: std::marker::PhantomData<T>,
 }
 
+impl<T> Clone for Index<T> {
+    fn clone(&self) -> Self {
+        Self {
+            index: self.index,
+            _marker: self._marker,
+        }
+    }
+}
+
+impl<T> Copy for Index<T> {}
+
 impl<T> std::fmt::Debug for Index<T> {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         write!(f, "Index({})", self.index)
@@ -33,13 +44,25 @@ impl<T> Index<T> {
     }
 }
 
-#[derive(Clone, Copy, PartialEq, Eq)]
+#[derive(PartialEq, Eq)]
 pub struct Slice<T> {
     start: u32,
     length: u16,
     _marker: std::marker::PhantomData<T>,
 }
 
+impl<T> Clone for Slice<T> {
+    fn clone(&self) -> Self {
+        Self {
+            start: self.start,
+            length: self.length,
+            _marker: self._marker,
+        }
+    }
+}
+
+impl<T> Copy for Slice<T> {}
+
 impl<T> std::fmt::Debug for Slice<T> {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         write!(f, "Slice(start = {}, length = {})", self.start, self.length)

compiler/solve/src/solve.rs

@@ -1,7 +1,9 @@
 use bumpalo::Bump;
-use roc_can::constraint::{Constraint, Constraints};
+use roc_can::constraint::Constraint::{self, *};
+use roc_can::constraint::{Constraints, LetConstraint};
 use roc_can::expected::{Expected, PExpected};
 use roc_collections::all::MutMap;
+use roc_collections::soa::{Index, Slice};
 use roc_module::ident::TagName;
 use roc_module::symbol::Symbol;
 use roc_region::all::{Loc, Region};
@@ -213,13 +215,28 @@ pub fn run_in_place(
     state.env
 }
 
-enum SolveWork<'a> {
+enum Work<'a> {
     Constraint {
         env: &'a Env,
         rank: Rank,
         constraint: &'a Constraint,
     },
     CheckForInfiniteTypes(LocalDefVarsVec<(Symbol, Loc<Variable>)>),
+    /// The ret_con part of a let constraint that does NOT introduces rigid and/or flex variables
+    LetConNoVariables {
+        env: &'a Env,
+        rank: Rank,
+        let_con: &'a LetConstraint,
+    },
+    /// The ret_con part of a let constraint that introduces rigid and/or flex variables
+    ///
+    /// These introduced variables must be generalized, hence this variant
+    /// is more complex than `LetConNoVariables`.
+    LetConIntroducesVariables {
+        env: &'a Env,
+        rank: Rank,
+        let_con: &'a LetConstraint,
+    },
 }
 
 #[allow(clippy::too_many_arguments)]
@@ -235,35 +252,163 @@ fn solve(
     subs: &mut Subs,
     constraint: &Constraint,
 ) -> State {
-    use Constraint::*;
+    let initial = Work::Constraint {
-
-    let initial = SolveWork::Constraint {
         env,
         rank,
         constraint,
     };
+
     let mut stack = vec![initial];
 
     while let Some(work_item) = stack.pop() {
         let (env, rank, constraint) = match work_item {
-            SolveWork::CheckForInfiniteTypes(def_vars) => {
+            Work::Constraint {
-                for (symbol, loc_var) in def_vars.iter() {
-                    check_for_infinite_type(subs, problems, *symbol, *loc_var);
-                }
-                // No constraint to be solved
-                continue;
-            }
-            SolveWork::Constraint {
                 env,
                 rank,
                 constraint,
-            } => (env, rank, constraint),
+            } => {
+                // the default case; actually solve this constraint
+                (env, rank, constraint)
+            }
+            Work::CheckForInfiniteTypes(def_vars) => {
+                // after a LetCon, we must check if any of the variables that we introduced
+                // loop back to themselves after solving the ret_constraint
+                for (symbol, loc_var) in def_vars.iter() {
+                    check_for_infinite_type(subs, problems, *symbol, *loc_var);
+                }
+
+                continue;
+            }
+            Work::LetConNoVariables { env, rank, let_con } => {
+                // NOTE be extremely careful with shadowing here
+                let offset = let_con.defs_and_ret_constraint.index();
+                let ret_constraint = &constraints.constraints[offset + 1];
+
+                // Add a variable for each def to new_vars_by_env.
+                let local_def_vars = LocalDefVarsVec::from_def_types(
+                    constraints,
+                    rank,
+                    pools,
+                    cached_aliases,
+                    subs,
+                    let_con.def_types,
+                );
+
+                let mut new_env = env.clone();
+                for (symbol, loc_var) in local_def_vars.iter() {
+                    new_env.insert_symbol_var_if_vacant(*symbol, loc_var.value);
+                }
+
+                stack.push(Work::CheckForInfiniteTypes(local_def_vars));
+                stack.push(Work::Constraint {
+                    env: arena.alloc(new_env),
+                    rank,
+                    constraint: ret_constraint,
+                });
+
+                continue;
+            }
+            Work::LetConIntroducesVariables { env, rank, let_con } => {
+                // NOTE be extremely careful with shadowing here
+                let offset = let_con.defs_and_ret_constraint.index();
+                let ret_constraint = &constraints.constraints[offset + 1];
+
+                let next_rank = rank.next();
+
+                let mark = state.mark;
+                let saved_env = state.env;
+
+                let young_mark = mark;
+                let visit_mark = young_mark.next();
+                let final_mark = visit_mark.next();
+
+                // Add a variable for each def to local_def_vars.
+                let local_def_vars = LocalDefVarsVec::from_def_types(
+                    constraints,
+                    next_rank,
+                    pools,
+                    cached_aliases,
+                    subs,
+                    let_con.def_types,
+                );
+
+                debug_assert_eq!(
+                    {
+                        let offenders = pools
+                            .get(next_rank)
+                            .iter()
+                            .filter(|var| {
+                                subs.get_rank(**var).into_usize() > next_rank.into_usize()
+                            })
+                            .collect::<Vec<_>>();
+
+                        let result = offenders.len();
+
+                        if result > 0 {
+                            dbg!(&subs, &offenders, &let_con.def_types);
+                        }
+
+                        result
+                    },
+                    0
+                );
+
+                // pop pool
+                generalize(subs, young_mark, visit_mark, next_rank, pools);
+
+                pools.get_mut(next_rank).clear();
+
+                // check that things went well
+                debug_assert!({
+                    // NOTE the `subs.redundant` check is added for the uniqueness
+                    // inference, and does not come from elm. It's unclear whether this is
+                    // a bug with uniqueness inference (something is redundant that
+                    // shouldn't be) or that it just never came up in elm.
+                    let rigid_vars = &constraints.variables[let_con.rigid_vars.indices()];
+
+                    let failing: Vec<_> = rigid_vars
+                        .iter()
+                        .filter(|&var| !subs.redundant(*var) && subs.get_rank(*var) != Rank::NONE)
+                        .collect();
+
+                    if !failing.is_empty() {
+                        println!("Rigids {:?}", &rigid_vars);
+                        println!("Failing {:?}", failing);
+                    }
+
+                    failing.is_empty()
+                });
+
+                let mut new_env = env.clone();
+                for (symbol, loc_var) in local_def_vars.iter() {
+                    new_env.insert_symbol_var_if_vacant(*symbol, loc_var.value);
+                }
+
+                // Note that this vars_by_symbol is the one returned by the
+                // previous call to solve()
+                let state_for_ret_con = State {
+                    env: saved_env,
+                    mark: final_mark,
+                };
+
+                // Now solve the body, using the new vars_by_symbol which includes
+                // the assignments' name-to-variable mappings.
+                stack.push(Work::CheckForInfiniteTypes(local_def_vars));
+                stack.push(Work::Constraint {
+                    env: arena.alloc(new_env),
+                    rank,
+                    constraint: ret_constraint,
+                });
+
+                state = state_for_ret_con;
+
+                continue;
+            }
         };
 
         state = match constraint {
             True => state,
             SaveTheEnvironment => {
-                // NOTE deviation: elm only copies the env into the state on SaveTheEnvironment
                 let mut copy = state;
 
                 copy.env = env.clone();
@@ -418,7 +563,7 @@ fn solve(
             And(slice) => {
                 let it = constraints.constraints[slice.indices()].iter().rev();
                 for sub_constraint in it {
-                    stack.push(SolveWork::Constraint {
+                    stack.push(Work::Constraint {
                         env,
                         rank,
                         constraint: sub_constraint,
@@ -486,209 +631,74 @@ fn solve(
                 let flex_vars = &constraints.variables[let_con.flex_vars.indices()];
                 let rigid_vars = &constraints.variables[let_con.rigid_vars.indices()];
 
-                let def_types = &constraints.def_types[let_con.def_types.indices()];
+                if matches!(&ret_constraint, True) && let_con.rigid_vars.is_empty() {
+                    introduce(subs, rank, pools, flex_vars);
 
-                match &ret_constraint {
+                    // If the return expression is guaranteed to solve,
-                    True if let_con.rigid_vars.is_empty() => {
+                    // solve the assignments themselves and move on.
-                        introduce(subs, rank, pools, flex_vars);
+                    stack.push(Work::Constraint {
+                        env,
+                        rank,
+                        constraint: defs_constraint,
+                    });
 
-                        // If the return expression is guaranteed to solve,
+                    state
-                        // solve the assignments themselves and move on.
+                } else if let_con.rigid_vars.is_empty() && let_con.flex_vars.is_empty() {
-                        stack.push(SolveWork::Constraint {
+                    // items are popped from the stack in reverse order. That means that we'll
-                            env,
+                    // first solve then defs_constraint, and then (eventually) the ret_constraint.
-                            rank,
+                    //
-                            constraint: defs_constraint,
+                    // Note that the LetConSimple gets the current env and rank,
-                        });
+                    // and not the env/rank from after solving the defs_constraint
-                        state
+                    stack.push(Work::LetConNoVariables { env, rank, let_con });
+                    stack.push(Work::Constraint {
+                        env,
+                        rank,
+                        constraint: defs_constraint,
+                    });
+
+                    state
+                } else {
+                    // work in the next pool to localize header
+                    let next_rank = rank.next();
+
+                    // introduce variables
+                    for &var in rigid_vars.iter().chain(flex_vars.iter()) {
+                        subs.set_rank(var, next_rank);
                     }
-                    ret_con if let_con.rigid_vars.is_empty() && let_con.flex_vars.is_empty() => {
-                        // TODO: make into `WorkItem` with `After`
-                        let state = solve(
-                            arena,
-                            constraints,
-                            env,
-                            state,
-                            rank,
-                            pools,
-                            problems,
-                            cached_aliases,
-                            subs,
-                            defs_constraint,
-                        );
 
-                        // Add a variable for each def to new_vars_by_env.
+                    // determine the next pool
-                        let mut local_def_vars =
+                    if next_rank.into_usize() < pools.len() {
-                            LocalDefVarsVec::with_length(let_con.def_types.len());
+                        // Nothing to do, we already accounted for the next rank, no need to
-
+                        // adjust the pools
-                        for (symbol, loc_type_index) in def_types.iter() {
+                    } else {
-                            let typ = &constraints.types[loc_type_index.value.index()];
+                        // we should be off by one at this point
-                            let var = type_to_var(subs, rank, pools, cached_aliases, typ);
+                        debug_assert_eq!(next_rank.into_usize(), 1 + pools.len());
-
+                        pools.extend_to(next_rank.into_usize());
-                            local_def_vars.push((
-                                *symbol,
-                                Loc {
-                                    value: var,
-                                    region: loc_type_index.region,
-                                },
-                            ));
-                        }
-
-                        let mut new_env = env.clone();
-                        for (symbol, loc_var) in local_def_vars.iter() {
-                            new_env.insert_symbol_var_if_vacant(*symbol, loc_var.value);
-                        }
-
-                        stack.push(SolveWork::CheckForInfiniteTypes(local_def_vars));
-                        stack.push(SolveWork::Constraint {
-                            env: arena.alloc(new_env),
-                            rank,
-                            constraint: ret_con,
-                        });
-
-                        state
                     }
-                    ret_con => {
-                        // work in the next pool to localize header
-                        let next_rank = rank.next();
 
-                        // introduce variables
+                    let pool: &mut Vec<Variable> = pools.get_mut(next_rank);
-                        for &var in rigid_vars.iter().chain(flex_vars.iter()) {
-                            subs.set_rank(var, next_rank);
-                        }
 
-                        // determine the next pool
+                    // Replace the contents of this pool with rigid_vars and flex_vars
-                        if next_rank.into_usize() < pools.len() {
+                    pool.clear();
-                            // Nothing to do, we already accounted for the next rank, no need to
+                    pool.reserve(rigid_vars.len() + flex_vars.len());
-                            // adjust the pools
+                    pool.extend(rigid_vars.iter());
-                        } else {
+                    pool.extend(flex_vars.iter());
-                            // we should be off by one at this point
-                            debug_assert_eq!(next_rank.into_usize(), 1 + pools.len());
-                            pools.extend_to(next_rank.into_usize());
-                        }
 
-                        let pool: &mut Vec<Variable> = pools.get_mut(next_rank);
+                    // run solver in next pool
 
-                        // Replace the contents of this pool with rigid_vars and flex_vars
+                    // items are popped from the stack in reverse order. That means that we'll
-                        pool.clear();
+                    // first solve then defs_constraint, and then (eventually) the ret_constraint.
-                        pool.reserve(rigid_vars.len() + flex_vars.len());
+                    //
-                        pool.extend(rigid_vars.iter());
+                    // Note that the LetConSimple gets the current env and rank,
-                        pool.extend(flex_vars.iter());
+                    // and not the env/rank from after solving the defs_constraint
+                    stack.push(Work::LetConIntroducesVariables { env, rank, let_con });
+                    stack.push(Work::Constraint {
+                        env,
+                        rank: next_rank,
+                        constraint: defs_constraint,
+                    });
 
-                        // run solver in next pool
+                    state
-
-                        // Add a variable for each def to local_def_vars.
-                        let mut local_def_vars =
-                            LocalDefVarsVec::with_length(let_con.def_types.len());
-
-                        for (symbol, loc_type) in def_types.iter() {
-                            let def_type = &constraints.types[loc_type.value.index()];
-
-                            let var = type_to_var(subs, next_rank, pools, cached_aliases, def_type);
-
-                            local_def_vars.push((
-                                *symbol,
-                                Loc {
-                                    value: var,
-                                    region: loc_type.region,
-                                },
-                            ));
-                        }
-
-                        // Solve the assignments' constraints first.
-                        // TODO: make into `WorkItem` with `After`
-                        let State {
-                            env: saved_env,
-                            mark,
-                        } = solve(
-                            arena,
-                            constraints,
-                            env,
-                            state,
-                            next_rank,
-                            pools,
-                            problems,
-                            cached_aliases,
-                            subs,
-                            defs_constraint,
-                        );
-
-                        let young_mark = mark;
-                        let visit_mark = young_mark.next();
-                        let final_mark = visit_mark.next();
-
-                        debug_assert_eq!(
-                            {
-                                let offenders = pools
-                                    .get(next_rank)
-                                    .iter()
-                                    .filter(|var| {
-                                        let current_rank =
-                                            subs.get_rank(roc_types::subs::Variable::clone(var));
-
-                                        current_rank.into_usize() > next_rank.into_usize()
-                                    })
-                                    .collect::<Vec<_>>();
-
-                                let result = offenders.len();
-
-                                if result > 0 {
-                                    dbg!(&subs, &offenders, &let_con.def_types);
-                                }
-
-                                result
-                            },
-                            0
-                        );
-
-                        // pop pool
-                        generalize(subs, young_mark, visit_mark, next_rank, pools);
-
-                        pools.get_mut(next_rank).clear();
-
-                        // check that things went well
-                        debug_assert!({
-                            // NOTE the `subs.redundant` check is added for the uniqueness
-                            // inference, and does not come from elm. It's unclear whether this is
-                            // a bug with uniqueness inference (something is redundant that
-                            // shouldn't be) or that it just never came up in elm.
-                            let failing: Vec<_> = rigid_vars
-                                .iter()
-                                .filter(|&var| {
-                                    !subs.redundant(*var) && subs.get_rank(*var) != Rank::NONE
-                                })
-                                .collect();
-
-                            if !failing.is_empty() {
-                                println!("Rigids {:?}", &rigid_vars);
-                                println!("Failing {:?}", failing);
-                            }
-
-                            failing.is_empty()
-                        });
-
-                        let mut new_env = env.clone();
-                        for (symbol, loc_var) in local_def_vars.iter() {
-                            new_env.insert_symbol_var_if_vacant(*symbol, loc_var.value);
-                        }
-
-                        // Note that this vars_by_symbol is the one returned by the
-                        // previous call to solve()
-                        let state_for_ret_con = State {
-                            env: saved_env,
-                            mark: final_mark,
-                        };
-
-                        // Now solve the body, using the new vars_by_symbol which includes
-                        // the assignments' name-to-variable mappings.
-                        stack.push(SolveWork::CheckForInfiniteTypes(local_def_vars));
-                        stack.push(SolveWork::Constraint {
-                            env: arena.alloc(new_env),
-                            rank,
-                            constraint: ret_con,
-                        });
-
-                        state_for_ret_con
-                    }
                 }
             }
             IsOpenType(type_index) => {
@@ -802,6 +812,36 @@ impl<T> LocalDefVarsVec<T> {
     }
 }
 
+impl LocalDefVarsVec<(Symbol, Loc<Variable>)> {
+    fn from_def_types(
+        constraints: &Constraints,
+        rank: Rank,
+        pools: &mut Pools,
+        cached_aliases: &mut MutMap<Symbol, Variable>,
+        subs: &mut Subs,
+        def_types_slice: Slice<(Symbol, Loc<Index<Type>>)>,
+    ) -> Self {
+        let def_types = &constraints.def_types[def_types_slice.indices()];
+
+        let mut local_def_vars = Self::with_length(def_types.len());
+
+        for (symbol, loc_type_index) in def_types.iter() {
+            let typ = &constraints.types[loc_type_index.value.index()];
+            let var = type_to_var(subs, rank, pools, cached_aliases, typ);
+
+            local_def_vars.push((
+                *symbol,
+                Loc {
+                    value: var,
+                    region: loc_type_index.region,
+                },
+            ));
+        }
+
+        local_def_vars
+    }
+}
+
 use std::cell::RefCell;
 std::thread_local! {
     /// Scratchpad arena so we don't need to allocate a new one all the time