Break up SolveEnv

crates/compiler/solve/src/ability.rs

@@ -21,7 +21,7 @@ use roc_types::types::{AliasKind, Category, MemberImpl, PatternCategory, Polarit
 use roc_unify::unify::{Env as UEnv, MustImplementConstraints};
 use roc_unify::unify::{MustImplementAbility, Obligated};
 
-use crate::env::Env;
+use crate::env::InferenceEnv;
 use crate::{aliases::Aliases, to_var::type_to_var};
 
 #[derive(Debug, Clone)]
@@ -56,7 +56,7 @@ pub struct PendingDerivesTable(
 
 impl PendingDerivesTable {
     pub fn new(
-        env: &mut Env,
+        env: &mut InferenceEnv,
         types: &mut Types,
         aliases: &mut Aliases,
         pending_derives: PendingDerives,

crates/compiler/solve/src/aliases.rs

@@ -9,7 +9,7 @@ use roc_types::{
 };
 
 use crate::to_var::type_to_var_help;
-use crate::{ability::ObligationCache, env::Env};
+use crate::{ability::ObligationCache, env::InferenceEnv};
 
 #[derive(Debug, Clone, Copy)]
 struct DelayedAliasVariables {
@@ -121,7 +121,7 @@ impl Aliases {
     }
 
     fn instantiate_result_result(
-        env: &mut Env,
+        env: &mut InferenceEnv,
         rank: Rank,
         alias_variables: AliasVariables,
     ) -> Variable {
@@ -143,7 +143,7 @@ impl Aliases {
 
     /// Build an alias of the form `Num range := range`
     fn build_num_opaque(
-        env: &mut Env,
+        env: &mut InferenceEnv,
         rank: Rank,
         symbol: Symbol,
         range_var: Variable,
@@ -160,7 +160,7 @@ impl Aliases {
 
     fn instantiate_builtin_aliases_real_var(
         &mut self,
-        env: &mut Env,
+        env: &mut InferenceEnv,
         rank: Rank,
         symbol: Symbol,
         alias_variables: AliasVariables,
@@ -228,7 +228,7 @@ impl Aliases {
 
     pub fn instantiate_real_var(
         &mut self,
-        env: &mut Env,
+        env: &mut InferenceEnv,
         rank: Rank,
         problems: &mut Vec<TypeError>,
         abilities_store: &AbilitiesStore,

crates/compiler/solve/src/deep_copy.rs

@@ -10,10 +10,15 @@ use roc_types::{
     types::{RecordField, Uls},
 };
 
-use crate::env::Env;
+use crate::env::SolveEnv;
 
 // TODO: eventually, we could possibly use the arena in Env instead.
-pub(crate) fn deep_copy_var_in(env: &mut Env, rank: Rank, var: Variable, arena: &Bump) -> Variable {
+pub(crate) fn deep_copy_var_in(
+    env: &mut SolveEnv,
+    rank: Rank,
+    var: Variable,
+    arena: &Bump,
+) -> Variable {
     let mut visited = bumpalo::collections::Vec::with_capacity_in(256, arena);
 
     let pool = env.pools.get_mut(rank);

crates/compiler/solve/src/env.rs

@@ -4,7 +4,7 @@ use roc_derive::SharedDerivedModule;
 use roc_types::subs::{Content, Descriptor, Mark, OptVariable, Rank, Subs, Variable};
 use roc_unify::unify::Env as UEnv;
 
-use crate::Pools;
+use crate::{FunctionKind, Pools};
 
 pub struct DerivedEnv<'a> {
     pub derived_module: &'a SharedDerivedModule,
@@ -12,15 +12,38 @@ pub struct DerivedEnv<'a> {
     pub exposed_types: &'a ExposedByModule,
 }
 
-pub struct Env<'a> {
+/// Environment necessary for solving and specialization.
+pub struct SolveEnv<'a> {
     pub arena: &'a Bump,
-    pub constraints: &'a Constraints,
     pub derived_env: &'a DerivedEnv<'a>,
     pub subs: &'a mut Subs,
     pub pools: &'a mut Pools,
 }
 
-impl<'a> Env<'a> {
+/// Environment necessary for inference.
+pub struct InferenceEnv<'a> {
+    pub constraints: &'a Constraints,
+    pub function_kind: FunctionKind,
+    pub arena: &'a Bump,
+    pub derived_env: &'a DerivedEnv<'a>,
+    pub subs: &'a mut Subs,
+    pub pools: &'a mut Pools,
+}
+
+impl<'a> SolveEnv<'a> {
+    /// Introduce some variables to Pools at the given rank.
+    /// Also, set each of their ranks in Subs to be the given rank.
+    pub fn introduce(&mut self, rank: Rank, vars: &[Variable]) {
+        introduce(self.subs, self.pools, rank, vars);
+    }
+
+    /// Retrieves an environment for unification.
+    pub fn uenv(&mut self) -> UEnv {
+        UEnv::new(self.subs)
+    }
+}
+
+impl<'a> InferenceEnv<'a> {
     #[inline(always)]
     pub fn register(&mut self, rank: Rank, content: Content) -> Variable {
         let descriptor = Descriptor {
@@ -40,13 +63,7 @@ impl<'a> Env<'a> {
     /// Introduce some variables to Pools at the given rank.
     /// Also, set each of their ranks in Subs to be the given rank.
     pub fn introduce(&mut self, rank: Rank, vars: &[Variable]) {
-        let pool: &mut Vec<Variable> = self.pools.get_mut(rank);
-
-        for &var in vars.iter() {
-            self.subs.set_rank(var, rank);
-        }
-
-        pool.extend(vars);
+        introduce(self.subs, self.pools, rank, vars);
     }
 
     #[inline(always)]
@@ -78,4 +95,25 @@ impl<'a> Env<'a> {
     pub fn uenv(&mut self) -> UEnv {
         UEnv::new(self.subs)
     }
+
+    pub fn as_solve_env(&mut self) -> SolveEnv {
+        SolveEnv {
+            arena: self.arena,
+            derived_env: self.derived_env,
+            subs: self.subs,
+            pools: self.pools,
+        }
+    }
+}
+
+/// Introduce some variables to Pools at the given rank.
+/// Also, set each of their ranks in Subs to be the given rank.
+fn introduce(subs: &mut Subs, pools: &mut Pools, rank: Rank, vars: &[Variable]) {
+    let pool: &mut Vec<Variable> = pools.get_mut(rank);
+
+    for &var in vars.iter() {
+        subs.set_rank(var, rank);
+    }
+
+    pool.extend(vars);
 }

crates/compiler/solve/src/kinds.rs

@@ -1,5 +1,5 @@
 /// How function kinds should be represented in the type system.
-#[derive(Debug)]
+#[derive(Debug, Clone, Copy)]
 pub enum FunctionKind {
     /// Function values are solved to lambda sets; lambda sets are the kind.
     LambdaSet,

crates/compiler/solve/src/lib.rs

@@ -19,6 +19,6 @@ mod pools;
 mod to_var;
 
 pub use aliases::Aliases;
-pub use env::{DerivedEnv, Env};
+pub use env::{DerivedEnv, SolveEnv};
 pub use kinds::FunctionKind;
 pub use pools::Pools;

crates/compiler/solve/src/module.rs

@@ -1,3 +1,4 @@
+use crate::FunctionKind;
 use crate::{aliases::Aliases, solve};
 use roc_can::abilities::{AbilitiesStore, ResolvedImpl};
 use roc_can::constraint::{Constraint, Constraints};
@@ -61,6 +62,8 @@ pub struct SolveConfig<'a> {
     /// All types introduced in the module. Canonicalized, but not necessarily yet associated with
     /// a variable substitution.
     pub types: Types,
+    /// How functions should be kinded.
+    pub function_kind: FunctionKind,
     /// Table of types introduced in this module that claim to derive an ability implementation.
     /// Due for checking and instantiation after the solver runs over the module.
     pub pending_derives: PendingDerives,

crates/compiler/solve/src/solve.rs

@@ -3,7 +3,7 @@ use crate::ability::{
     CheckedDerives, ObligationCache, PendingDerivesTable, Resolved,
 };
 use crate::deep_copy::deep_copy_var_in;
-use crate::env::{DerivedEnv, Env};
+use crate::env::{DerivedEnv, InferenceEnv};
 use crate::module::{SolveConfig, Solved};
 use crate::pools::Pools;
 use crate::specialize::{
@@ -122,6 +122,7 @@ fn run_in_place(
         pending_derives,
         exposed_by_module,
         derived_module,
+        function_kind,
     } = config;
 
     let mut pools = Pools::default();
@@ -141,9 +142,10 @@ fn run_in_place(
         exposed_types: exposed_by_module,
     };
 
-    let mut env = Env {
+    let mut env = InferenceEnv {
         arena: &arena,
         constraints,
+        function_kind,
         derived_env: &derived_env,
         subs,
         pools: &mut pools,
@@ -218,7 +220,7 @@ enum Work<'a> {
 }
 
 fn solve(
-    env: &mut Env,
+    env: &mut InferenceEnv,
     mut can_types: Types,
     mut state: State,
     rank: Rank,
@@ -580,7 +582,11 @@ fn solve(
                         // then we copy from that module's Subs into our own. If the value
                         // is being looked up in this module, then we use our Subs as both
                         // the source and destination.
-                        let actual = deep_copy_var_in(env, rank, var, env.arena);
+                        let actual = {
+                            let mut solve_env = env.as_solve_env();
+                            let solve_env = &mut solve_env;
+                            deep_copy_var_in(solve_env, rank, var, solve_env.arena)
+                        };
                         let expectation = &env.constraints.expectations[expectation_index.index()];
 
                         let expected = either_type_index_to_var(
@@ -1426,7 +1432,7 @@ fn chase_alias_content(subs: &Subs, mut var: Variable) -> (Variable, &Content) {
 }
 
 fn compact_lambdas_and_check_obligations(
-    env: &mut Env,
+    env: &mut InferenceEnv,
     problems: &mut Vec<TypeError>,
     abilities_store: &mut AbilitiesStore,
     obligation_cache: &mut ObligationCache,
@@ -1437,7 +1443,7 @@ fn compact_lambdas_and_check_obligations(
         obligations,
         awaiting_specialization: new_awaiting,
     } = compact_lambda_sets_of_vars(
-        env,
+        &mut env.as_solve_env(),
         lambda_sets_to_specialize,
         &SolvePhase { abilities_store },
     );
@@ -1450,7 +1456,7 @@ fn compact_lambdas_and_check_obligations(
     awaiting_specialization.union(new_awaiting);
 }
 
-fn open_tag_union(env: &mut Env, var: Variable) {
+fn open_tag_union(env: &mut InferenceEnv, var: Variable) {
     let mut stack = vec![var];
     while let Some(var) = stack.pop() {
         use {Content::*, FlatType::*};
@@ -1582,7 +1588,7 @@ fn close_pattern_matched_tag_unions(subs: &mut Subs, var: Variable) {
 // Aggressive but necessary - there aren't many usages.
 #[inline(always)]
 fn check_ability_specialization(
-    env: &mut Env,
+    env: &mut InferenceEnv,
     rank: Rank,
     abilities_store: &mut AbilitiesStore,
     obligation_cache: &mut ObligationCache,
@@ -1606,8 +1612,16 @@ fn check_ability_specialization(
 
         // We need to freshly instantiate the root signature so that all unifications are reflected
         // in the specialization type, but not the original signature type.
-        let root_signature_var =
-            deep_copy_var_in(env, Rank::toplevel(), root_signature_var, env.arena);
+        let root_signature_var = {
+            let mut solve_env = env.as_solve_env();
+            let solve_env = &mut solve_env;
+            deep_copy_var_in(
+                solve_env,
+                Rank::toplevel(),
+                root_signature_var,
+                solve_env.arena,
+            )
+        };
         let snapshot = env.subs.snapshot();
         let unified = unify_introduced_ability_specialization(
 <<<<<<< HEAD
@@ -1817,7 +1831,7 @@ impl<T> LocalDefVarsVec<T> {
 
 impl LocalDefVarsVec<(Symbol, Loc<Variable>)> {
     fn from_def_types(
-        env: &mut Env,
+        env: &mut InferenceEnv,
         rank: Rank,
         problems: &mut Vec<TypeError>,
         abilities_store: &mut AbilitiesStore,
@@ -1851,7 +1865,7 @@ impl LocalDefVarsVec<(Symbol, Loc<Variable>)> {
 }
 
 fn check_for_infinite_type(
-    env: &mut Env,
+    env: &mut InferenceEnv,
     problems: &mut Vec<TypeError>,
     symbol: Symbol,
     loc_var: Loc<Variable>,
@@ -1914,7 +1928,7 @@ fn circular_error(
 /// Ensures that variables introduced at the `young_rank`, but that should be
 /// stuck at a lower level, are marked at that level and not generalized at the
 /// present `young_rank`. See [adjust_rank].
-fn generalize(env: &mut Env, young_mark: Mark, visit_mark: Mark, young_rank: Rank) {
+fn generalize(env: &mut InferenceEnv, young_mark: Mark, visit_mark: Mark, young_rank: Rank) {
     let subs = &mut env.subs;
     let pools = &mut env.pools;
 

crates/compiler/solve/src/specialize.rs

@@ -22,7 +22,7 @@ use roc_unify::unify::{unify, Mode, MustImplementConstraints};
 use crate::{
     ability::builtin_module_with_unlisted_ability_impl,
     deep_copy::deep_copy_var_in,
-    env::{DerivedEnv, Env},
+    env::{DerivedEnv, SolveEnv},
 };
 
 /// What phase in the compiler is reaching out to specialize lambda sets?
@@ -295,7 +295,7 @@ fn unique_unspecialized_lambda(subs: &Subs, c_a: Variable, uls: &[Uls]) -> Optio
 
 #[must_use]
 pub fn compact_lambda_sets_of_vars<P: Phase>(
-    env: &mut Env,
+    env: &mut SolveEnv,
     uls_of_var: UlsOfVar,
     phase: &P,
 ) -> CompactionResult {
@@ -464,7 +464,7 @@ enum OneCompactionResult {
 #[must_use]
 #[allow(clippy::too_many_arguments)]
 fn compact_lambda_set<P: Phase>(
-    env: &mut Env,
+    env: &mut SolveEnv,
     resolved_concrete: Variable,
     this_lambda_set: Variable,
     phase: &P,

crates/compiler/solve/src/to_var.rs

@@ -22,7 +22,7 @@ use roc_unify::unify::{unify, Mode, Unified};
 use crate::{
     ability::{AbilityImplError, ObligationCache},
     deep_copy::deep_copy_var_in,
-    env::Env,
+    env::InferenceEnv,
     Aliases,
 };
 
@@ -42,7 +42,7 @@ fn put_scratchpad(scratchpad: bumpalo::Bump) {
 }
 
 pub(crate) fn either_type_index_to_var(
-    env: &mut Env,
+    env: &mut InferenceEnv,
     rank: Rank,
     problems: &mut Vec<TypeError>,
     abilities_store: &mut AbilitiesStore,
@@ -82,7 +82,7 @@ pub(crate) fn either_type_index_to_var(
 }
 
 pub(crate) fn type_to_var(
-    env: &mut Env,
+    env: &mut InferenceEnv,
     rank: Rank,
     problems: &mut Vec<TypeError>,
     abilities_store: &mut AbilitiesStore,
@@ -126,7 +126,7 @@ enum RegisterVariable {
 
 impl RegisterVariable {
     fn from_type(
-        env: &mut Env,
+        env: &mut InferenceEnv,
         rank: Rank,
         arena: &'_ bumpalo::Bump,
         types: &mut Types,
@@ -149,7 +149,7 @@ impl RegisterVariable {
                         reserved
                     } else {
                         // for any other rank, we need to copy; it takes care of adjusting the rank
-                        deep_copy_var_in(env, rank, reserved, arena)
+                        deep_copy_var_in(&mut env.as_solve_env(), rank, reserved, arena)
                     };
                     // Safety: the `destination` will become the source-of-truth for the type index, since it
                     // was not already transformed before (if it was, we'd be in the Variable branch!)
@@ -165,7 +165,7 @@ impl RegisterVariable {
 
     #[inline(always)]
     fn with_stack(
-        env: &mut Env,
+        env: &mut InferenceEnv,
         rank: Rank,
         arena: &'_ bumpalo::Bump,
         types: &mut Types,
@@ -253,7 +253,7 @@ enum TypeToVar {
 
 #[allow(clippy::too_many_arguments)]
 pub(crate) fn type_to_var_help(
-    env: &mut Env,
+    env: &mut InferenceEnv,
     rank: Rank,
     problems: &mut Vec<TypeError>,
     abilities_store: &AbilitiesStore,
@@ -911,7 +911,7 @@ pub(crate) fn type_to_var_help(
 
 #[inline(always)]
 fn roc_result_to_var(
-    env: &mut Env,
+    env: &mut InferenceEnv,
     rank: Rank,
     arena: &'_ bumpalo::Bump,
     types: &mut Types,
@@ -1090,7 +1090,7 @@ fn find_tag_name_run(slice: &[TagName], subs: &mut Subs) -> Option<SubsSlice<Tag
 
 #[inline(always)]
 fn register_tag_arguments(
-    env: &mut Env,
+    env: &mut InferenceEnv,
     rank: Rank,
     arena: &'_ bumpalo::Bump,
     types: &mut Types,
@@ -1114,7 +1114,7 @@ fn register_tag_arguments(
 
 /// Assumes that the tags are sorted and there are no duplicates!
 fn insert_tags_fast_path(
-    env: &mut Env,
+    env: &mut InferenceEnv,
     rank: Rank,
     arena: &'_ bumpalo::Bump,
     types: &mut Types,
@@ -1185,7 +1185,7 @@ fn insert_tags_fast_path(
 }
 
 fn insert_tags_slow_path(
-    env: &mut Env,
+    env: &mut InferenceEnv,
     rank: Rank,
     arena: &'_ bumpalo::Bump,
     types: &mut Types,
@@ -1215,7 +1215,7 @@ fn insert_tags_slow_path(
 }
 
 fn type_to_union_tags(
-    env: &mut Env,
+    env: &mut InferenceEnv,
     rank: Rank,
     arena: &'_ bumpalo::Bump,
     types: &mut Types,
@@ -1274,7 +1274,7 @@ fn type_to_union_tags(
 }
 
 fn create_union_lambda(
-    env: &mut Env,
+    env: &mut InferenceEnv,
     rank: Rank,
     arena: &'_ bumpalo::Bump,
     types: &mut Types,