Merge pull request #5592 from roc-lang/simplify-solving-spike

Break up solving and put common values in a shared struct

crates/compiler/can/src/constraint.rs

@@ -64,6 +64,26 @@ pub type PExpectedTypeIndex = Index<PExpected<TypeOrVar>>;
 pub type TypeOrVar = EitherIndex<TypeTag, Variable>;
 
 impl Constraints {
+    pub fn empty() -> Self {
+        Self {
+            constraints: Default::default(),
+            type_slices: Default::default(),
+            variables: Default::default(),
+            loc_symbols: Default::default(),
+            let_constraints: Default::default(),
+            categories: Default::default(),
+            pattern_categories: Default::default(),
+            expectations: Default::default(),
+            pattern_expectations: Default::default(),
+            includes_tags: Default::default(),
+            strings: Default::default(),
+            sketched_rows: Default::default(),
+            eq: Default::default(),
+            pattern_eq: Default::default(),
+            cycles: Default::default(),
+        }
+    }
+
     pub fn new() -> Self {
         let constraints = Vec::new();
         let type_slices = Vec::with_capacity(16);

crates/compiler/late_solve/src/lib.rs

@@ -5,20 +5,22 @@ use std::sync::{Arc, RwLock};
 
 use bumpalo::Bump;
 use roc_can::abilities::AbilitiesStore;
+use roc_can::constraint::Constraints;
 use roc_can::module::ExposedByModule;
 use roc_collections::MutMap;
 use roc_derive::SharedDerivedModule;
 use roc_error_macros::internal_error;
 use roc_module::symbol::ModuleId;
 use roc_module::symbol::Symbol;
 use roc_solve::ability::AbilityResolver;
-use roc_solve::solve::Pools;
-use roc_solve::specialize::{compact_lambda_sets_of_vars, DerivedEnv, Phase};
+use roc_solve::specialize::{compact_lambda_sets_of_vars, Phase};
+use roc_solve::Pools;
+use roc_solve::{DerivedEnv, Env};
 use roc_types::subs::{get_member_lambda_sets_at_region, Content, FlatType, LambdaSet};
 use roc_types::subs::{ExposedTypesStorageSubs, Subs, Variable};
 use roc_types::types::Polarity;
 use roc_unify::unify::MetaCollector;
-use roc_unify::unify::{Env, Mode, Unified};
+use roc_unify::unify::{Env as UEnv, Mode, Unified};
 
 pub use roc_solve::ability::{ResolveError, Resolved};
 pub use roc_types::subs::instantiate_rigids;
@@ -340,6 +342,19 @@ impl MetaCollector for ChangedVariableCollector {
     }
 }
 
+std::thread_local! {
+    static SCRATCHPAD_FOR_OCCURS: std::cell::RefCell<Option<Constraints>> = std::cell::RefCell::new(Some(Constraints::empty()));
+}
+
+fn with_empty_solve_constraints<T>(f: impl FnOnce(&Constraints) -> T) -> T {
+    SCRATCHPAD_FOR_OCCURS.with(|cell| {
+        let constr = cell.take().unwrap();
+        let result = f(&constr);
+        cell.replace(Some(constr));
+        result
+    })
+}
+
 /// Unifies two variables and performs lambda set compaction.
 /// Ranks and other ability demands are disregarded.
 #[allow(clippy::too_many_arguments)]
@@ -359,7 +374,7 @@ pub fn unify(
         "derived module can only unify its subs in its own context!"
     );
     let unified = roc_unify::unify::unify_with_collector::<ChangedVariableCollector>(
-        &mut Env::new(subs),
+        &mut UEnv::new(subs),
         left,
         right,
         Mode::EQ,
@@ -381,14 +396,17 @@ pub fn unify(
                 exposed_types: exposed_by_module,
             };
 
-            let must_implement_constraints = compact_lambda_sets_of_vars(
-                subs,
-                &derived_env,
-                arena,
-                &mut pools,
-                lambda_sets_to_specialize,
-                &late_phase,
-            );
+            let must_implement_constraints = with_empty_solve_constraints(|c| {
+                let mut env = Env {
+                    constraints: c,
+                    subs,
+                    derived_env: &derived_env,
+                    arena,
+                    pools: &mut pools,
+                };
+
+                compact_lambda_sets_of_vars(&mut env, lambda_sets_to_specialize, &late_phase)
+            });
             // At this point we can't do anything with must-implement constraints, since we're no
             // longer solving. We must assume that they were totally caught during solving.
             // After we land https://github.com/roc-lang/roc/issues/3207 this concern should totally

crates/compiler/load_internal/src/file.rs

@@ -56,7 +56,7 @@ use roc_region::all::{LineInfo, Loc, Region};
 #[cfg(not(target_family = "wasm"))]
 use roc_reporting::report::to_https_problem_report_string;
 use roc_reporting::report::{to_file_problem_report_string, Palette, RenderTarget};
-use roc_solve::module::{extract_module_owned_implementations, Solved, SolvedModule};
+use roc_solve::module::{extract_module_owned_implementations, SolveConfig, Solved, SolvedModule};
 use roc_solve_problem::TypeError;
 use roc_target::TargetInfo;
 use roc_types::subs::{CopiedImport, ExposedTypesStorageSubs, Subs, VarStore, Variable};
@@ -5061,6 +5061,14 @@ fn import_variable_for_symbol(
     }
 }
 
+struct SolveResult {
+    solved: Solved<Subs>,
+    solved_implementations: ResolvedImplementations,
+    exposed_vars_by_symbol: Vec<(Symbol, Variable)>,
+    problems: Vec<TypeError>,
+    abilities_store: AbilitiesStore,
+}
+
 #[allow(clippy::complexity)]
 fn run_solve_solve(
     exposed_for_module: ExposedForModule,
@@ -5071,13 +5079,7 @@ fn run_solve_solve(
     var_store: VarStore,
     module: Module,
     derived_module: SharedDerivedModule,
-) -> (
-    Solved<Subs>,
-    ResolvedImplementations,
-    Vec<(Symbol, Variable)>,
-    Vec<TypeError>,
-    AbilitiesStore,
-) {
+) -> SolveResult {
     let Module {
         exposed_symbols,
         aliases,
@@ -5111,30 +5113,34 @@ fn run_solve_solve(
         &import_variables,
     );
 
-    let mut solve_aliases = roc_solve::solve::Aliases::with_capacity(aliases.len());
+    let mut solve_aliases = roc_solve::Aliases::with_capacity(aliases.len());
     for (name, (_, alias)) in aliases.iter() {
         solve_aliases.insert(&mut types, *name, alias.clone());
     }
 
-    let (solved_subs, solved_implementations, exposed_vars_by_symbol, problems, abilities_store) = {
+    let (solve_output, solved_implementations, exposed_vars_by_symbol) = {
         let module_id = module.module_id;
 
-        let (solved_subs, solved_env, problems, abilities_store) = roc_solve::module::run_solve(
-            module_id,
+        let solve_config = SolveConfig {
+            home: module_id,
             types,
-            &constraints,
-            actual_constraint,
+            constraints: &constraints,
+            root_constraint: actual_constraint,
+            pending_derives,
+            exposed_by_module: &exposed_for_module.exposed_by_module,
+            derived_module,
+        };
+
+        let solve_output = roc_solve::module::run_solve(
+            solve_config,
             rigid_variables,
             subs,
             solve_aliases,
             abilities_store,
-            pending_derives,
-            &exposed_for_module.exposed_by_module,
-            derived_module,
         );
 
         let solved_implementations =
-            extract_module_owned_implementations(module_id, &abilities_store);
+            extract_module_owned_implementations(module_id, &solve_output.resolved_abilities_store);
 
         let is_specialization_symbol = |sym| {
             solved_implementations
@@ -5147,7 +5153,8 @@ fn run_solve_solve(
 
         // Expose anything that is explicitly exposed by the header, or is a specialization of an
         // ability.
-        let exposed_vars_by_symbol: Vec<_> = solved_env
+        let exposed_vars_by_symbol: Vec<_> = solve_output
+            .scope
             .vars_by_symbol()
             .filter(|(k, _)| {
                 exposed_symbols.contains(k)
@@ -5156,22 +5163,23 @@ fn run_solve_solve(
             })
             .collect();
 
-        (
-            solved_subs,
-            solved_implementations,
-            exposed_vars_by_symbol,
-            problems,
-            abilities_store,
-        )
+        (solve_output, solved_implementations, exposed_vars_by_symbol)
     };
 
-    (
-        solved_subs,
+    let roc_solve::module::SolveOutput {
+        subs,
+        scope: _,
+        errors,
+        resolved_abilities_store,
+    } = solve_output;
+
+    SolveResult {
+        solved: subs,
         solved_implementations,
         exposed_vars_by_symbol,
-        problems,
-        abilities_store,
-    )
+        problems: errors,
+        abilities_store: resolved_abilities_store,
+    }
 }
 
 fn run_solve<'a>(
@@ -5201,7 +5209,7 @@ fn run_solve<'a>(
     let loc_dbgs = std::mem::take(&mut module.loc_dbgs);
     let module = module;
 
-    let (solved_subs, solved_implementations, exposed_vars_by_symbol, problems, abilities_store) = {
+    let solve_result = {
         if module_id.is_builtin() {
             match cached_types.lock().remove(&module_id) {
                 None => run_solve_solve(
@@ -5219,13 +5227,13 @@ fn run_solve<'a>(
                     exposed_vars_by_symbol,
                     abilities,
                     solved_implementations,
-                }) => (
-                    Solved(subs),
+                }) => SolveResult {
+                    solved: Solved(subs),
                     solved_implementations,
                     exposed_vars_by_symbol,
-                    vec![],
-                    abilities,
-                ),
+                    problems: vec![],
+                    abilities_store: abilities,
+                },
             }
         } else {
             run_solve_solve(
@@ -5241,7 +5249,14 @@ fn run_solve<'a>(
         }
     };
 
-    let mut solved_subs = solved_subs;
+    let SolveResult {
+        solved: mut solved_subs,
+        solved_implementations,
+        exposed_vars_by_symbol,
+        problems,
+        abilities_store,
+    } = solve_result;
+
     let exposed_types = roc_solve::module::exposed_types_storage_subs(
         module_id,
         &mut solved_subs,

crates/compiler/solve/src/ability.rs

@@ -18,11 +18,11 @@ use roc_types::subs::{
     TupleElems, Variable,
 };
 use roc_types::types::{AliasKind, Category, MemberImpl, PatternCategory, Polarity, Types};
-use roc_unify::unify::{Env, MustImplementConstraints};
+use roc_unify::unify::{Env as UEnv, MustImplementConstraints};
 use roc_unify::unify::{MustImplementAbility, Obligated};
 
-use crate::solve::type_to_var;
-use crate::solve::{Aliases, Pools};
+use crate::env::Env;
+use crate::{aliases::Aliases, to_var::type_to_var};
 
 #[derive(Debug, Clone)]
 pub enum AbilityImplError {
@@ -56,7 +56,7 @@ pub struct PendingDerivesTable(
 
 impl PendingDerivesTable {
     pub fn new(
-        subs: &mut Subs,
+        env: &mut Env,
         types: &mut Types,
         aliases: &mut Aliases,
         pending_derives: PendingDerives,
@@ -81,17 +81,16 @@ impl PendingDerivesTable {
                 // Neither rank nor pools should matter here.
                 let typ = types.from_old_type(&typ);
                 let opaque_var = type_to_var(
-                    subs,
+                    env,
                     Rank::toplevel(),
                     problems,
                     abilities_store,
                     obligation_cache,
-                    &mut Pools::default(),
                     types,
                     aliases,
                     typ,
                 );
-                let real_var = match subs.get_content_without_compacting(opaque_var) {
+                let real_var = match env.subs.get_content_without_compacting(opaque_var) {
                     Content::Alias(_, _, real_var, AliasKind::Opaque) => real_var,
                     _ => internal_error!("Non-opaque in derives table"),
                 };
@@ -1284,7 +1283,7 @@ impl DerivableVisitor for DeriveEq {
 
         // Of the floating-point types,
         // only Dec implements Eq.
-        let mut env = Env::new(subs);
+        let mut env = UEnv::new(subs);
         let unified = unify(
             &mut env,
             content_var,
@@ -1419,7 +1418,7 @@ pub fn resolve_ability_specialization<R: AbilityResolver>(
 
     instantiate_rigids(subs, signature_var);
     let (_vars, must_implement_ability, _lambda_sets_to_specialize, _meta) = unify(
-        &mut Env::new(subs),
+        &mut UEnv::new(subs),
         specialization_var,
         signature_var,
         Mode::EQ,