Make type_to_variable manually tail-recursive

by using a work stack and reserving variables. fun stuff

compiler/solve/src/solve.rs

@@ -865,8 +865,12 @@ fn type_to_var(
     _: &mut MutMap<Symbol, Variable>,
     typ: &Type,
 ) -> Variable {
+    if let Type::Variable(var) = typ {
+        *var
+    } else {
         let mut arena = take_scratchpad();
 
+        // let var = type_to_variable(subs, rank, pools, &arena, typ);
         let var = type_to_variable(subs, rank, pools, &arena, typ);
 
         arena.reset();
@@ -874,6 +878,53 @@ fn type_to_var(
 
         var
     }
+}
+
+enum RegisterVariable {
+    /// Based on the Type, we already know what variable this will be
+    Direct(Variable),
+    /// This Type needs more complicated Content. We reserve a Variable
+    /// for it, but put a placeholder Content in subs
+    Deferred,
+}
+
+impl RegisterVariable {
+    fn from_type(
+        subs: &mut Subs,
+        rank: Rank,
+        pools: &mut Pools,
+        arena: &'_ bumpalo::Bump,
+        typ: &Type,
+    ) -> Self {
+        use RegisterVariable::*;
+
+        match typ {
+            Variable(var) => Direct(*var),
+            EmptyRec => Direct(Variable::EMPTY_RECORD),
+            EmptyTagUnion => Direct(Variable::EMPTY_TAG_UNION),
+            Type::Alias { symbol, .. } => {
+                if let Some(reserved) = Variable::get_reserved(*symbol) {
+                    if rank.is_none() {
+                        // reserved variables are stored with rank NONE
+                        return Direct(reserved);
+                    } else {
+                        // for any other rank, we need to copy; it takes care of adjusting the rank
+                        let copied = deep_copy_var_in(subs, rank, pools, reserved, arena);
+                        return Direct(copied);
+                    }
+                }
+
+                Deferred
+            }
+            _ => Deferred,
+        }
+    }
+}
+
+#[derive(Debug)]
+enum TypeToVar<'a> {
+    Defer(&'a Type, Variable),
+}
 
 fn type_to_variable<'a>(
     subs: &mut Subs,
@@ -884,29 +935,47 @@ fn type_to_variable<'a>(
 ) -> Variable {
     use bumpalo::collections::Vec;
 
+    let mut stack = Vec::with_capacity_in(8, arena);
+
+    macro_rules! helper {
+        ($typ:expr) => {{
+            match RegisterVariable::from_type(subs, rank, pools, arena, $typ) {
+                RegisterVariable::Direct(var) => var,
+                RegisterVariable::Deferred => {
+                    let var = subs.fresh_unnamed_flex_var();
+                    stack.push(TypeToVar::Defer($typ, var));
+                    var
+                }
+            }
+        }};
+    }
+
+    let result = helper!(typ);
+
+    while let Some(TypeToVar::Defer(typ, destination)) = stack.pop() {
         match typ {
-        Variable(var) => *var,
+            Variable(_) | EmptyRec | EmptyTagUnion => {
+                unreachable!("This variant should never be deferred!")
+            }
             RangedNumber(typ, vars) => {
-            let ty_var = type_to_variable(subs, rank, pools, arena, typ);
+                let ty_var = helper!(typ);
                 let vars = VariableSubsSlice::insert_into_subs(subs, vars.iter().copied());
                 let content = Content::RangedNumber(ty_var, vars);
 
-            register(subs, rank, pools, content)
+                register_with_known_var(subs, destination, rank, pools, content)
             }
             Apply(symbol, arguments, _) => {
                 let new_arguments = VariableSubsSlice::reserve_into_subs(subs, arguments.len());
                 for (target_index, var_index) in (new_arguments.indices()).zip(arguments) {
-                let var = type_to_variable(subs, rank, pools, arena, var_index);
+                    let var = helper!(var_index);
                     subs.variables[target_index] = var;
                 }
 
                 let flat_type = FlatType::Apply(*symbol, new_arguments);
                 let content = Content::Structure(flat_type);
 
-            register(subs, rank, pools, content)
+                register_with_known_var(subs, destination, rank, pools, content)
             }
-        EmptyRec => Variable::EMPTY_RECORD,
-        EmptyTagUnion => Variable::EMPTY_TAG_UNION,
 
             ClosureTag { name, ext } => {
                 let tag_name = TagName::Closure(*name);
@@ -919,22 +988,22 @@ fn type_to_variable<'a>(
 
                 let content = Content::Structure(FlatType::TagUnion(union_tags, *ext));
 
-            register(subs, rank, pools, content)
+                register_with_known_var(subs, destination, rank, pools, content)
             }
-
             // This case is important for the rank of boolean variables
             Function(arguments, closure_type, ret_type) => {
                 let new_arguments = VariableSubsSlice::reserve_into_subs(subs, arguments.len());
                 for (target_index, var_index) in (new_arguments.indices()).zip(arguments) {
-                let var = type_to_variable(subs, rank, pools, arena, var_index);
+                    let var = helper!(var_index);
                     subs.variables[target_index] = var;
                 }
 
-            let ret_var = type_to_variable(subs, rank, pools, arena, ret_type);
+                let ret_var = helper!(ret_type);
-            let closure_var = type_to_variable(subs, rank, pools, arena, closure_type);
+                let closure_var = helper!(closure_type);
-            let content = Content::Structure(FlatType::Func(new_arguments, closure_var, ret_var));
+                let content =
+                    Content::Structure(FlatType::Func(new_arguments, closure_var, ret_var));
 
-            register(subs, rank, pools, content)
+                register_with_known_var(subs, destination, rank, pools, content)
             }
             Record(fields, ext) => {
                 // An empty fields is inefficient (but would be correct)
@@ -944,13 +1013,19 @@ fn type_to_variable<'a>(
                 let mut field_vars = Vec::with_capacity_in(fields.len(), arena);
 
                 for (field, field_type) in fields {
-                let field_var =
+                    let field_var = {
-                    field_type.map(|typ| type_to_variable(subs, rank, pools, arena, typ));
+                        use roc_types::types::RecordField::*;
+                        match &field_type {
+                            Optional(t) => Optional(helper!(t)),
+                            Required(t) => Required(helper!(t)),
+                            Demanded(t) => Demanded(helper!(t)),
+                        }
+                    };
 
                     field_vars.push((field.clone(), field_var));
                 }
 
-            let temp_ext_var = type_to_variable(subs, rank, pools, arena, ext);
+                let temp_ext_var = helper!(ext);
 
                 let (it, new_ext_var) =
                     gather_fields_unsorted_iter(subs, RecordFields::empty(), temp_ext_var)
@@ -967,8 +1042,9 @@ fn type_to_variable<'a>(
 
                 let content = Content::Structure(FlatType::Record(record_fields, new_ext_var));
 
-            register(subs, rank, pools, content)
+                register_with_known_var(subs, destination, rank, pools, content)
             }
+
             TagUnion(tags, ext) => {
                 // An empty tags is inefficient (but would be correct)
                 // If hit, try to turn the value into an EmptyTagUnion in canonicalization
@@ -977,10 +1053,10 @@ fn type_to_variable<'a>(
                 let (union_tags, ext) = type_to_union_tags(subs, rank, pools, arena, tags, ext);
                 let content = Content::Structure(FlatType::TagUnion(union_tags, ext));
 
-            register(subs, rank, pools, content)
+                register_with_known_var(subs, destination, rank, pools, content)
             }
             FunctionOrTagUnion(tag_name, symbol, ext) => {
-            let temp_ext_var = type_to_variable(subs, rank, pools, arena, ext);
+                let temp_ext_var = helper!(ext);
 
                 let (it, ext) = roc_types::types::gather_tags_unsorted_iter(
                     subs,
@@ -997,7 +1073,7 @@ fn type_to_variable<'a>(
 
                 let content = Content::Structure(FlatType::FunctionOrTagUnion(slice, *symbol, ext));
 
-            register(subs, rank, pools, content)
+                register_with_known_var(subs, destination, rank, pools, content)
             }
             RecursiveTagUnion(rec_var, tags, ext) => {
                 // An empty tags is inefficient (but would be correct)
@@ -1008,7 +1084,8 @@ fn type_to_variable<'a>(
                 let content =
                     Content::Structure(FlatType::RecursiveTagUnion(*rec_var, union_tags, ext));
 
-            let tag_union_var = register(subs, rank, pools, content);
+                let tag_union_var = destination;
+                register_with_known_var(subs, tag_union_var, rank, pools, content);
 
                 register_with_known_var(
                     subs,
@@ -1031,33 +1108,41 @@ fn type_to_variable<'a>(
                 lambda_set_variables,
                 kind,
             } => {
-            if let Some(reserved) = Variable::get_reserved(*symbol) {
+                debug_assert!(Variable::get_reserved(*symbol).is_none());
-                if rank.is_none() {
+
-                    // reserved variables are stored with rank NONE
+                let alias_variables = {
-                    return reserved;
+                    let length = type_arguments.len() + lambda_set_variables.len();
-                } else {
+                    let new_variables = VariableSubsSlice::reserve_into_subs(subs, length);
-                    // for any other rank, we need to copy; it takes care of adjusting the rank
+
-                    return deep_copy_var_in(subs, rank, pools, reserved, arena);
+                    for (target_index, (_, arg_type)) in
-                }
+                        (new_variables.indices()).zip(type_arguments)
+                    {
+                        let copy_var = helper!(arg_type);
+                        subs.variables[target_index] = copy_var;
                     }
 
-            let alias_variables = alias_to_var(
+                    let it = (new_variables.indices().skip(type_arguments.len()))
-                subs,
+                        .zip(lambda_set_variables);
-                rank,
+                    for (target_index, ls) in it {
-                pools,
+                        let copy_var = helper!(&ls.0);
-                arena,
+                        subs.variables[target_index] = copy_var;
-                type_arguments,
+                    }
-                lambda_set_variables,
+
-            );
+                    AliasVariables {
+                        variables_start: new_variables.start,
+                        type_variables_len: type_arguments.len() as _,
+                        all_variables_len: length as _,
+                    }
+                };
 
                 let alias_variable = if let Symbol::RESULT_RESULT = *symbol {
                     roc_result_to_var(subs, rank, pools, arena, actual)
                 } else {
-                type_to_variable(subs, rank, pools, arena, actual)
+                    helper!(actual)
                 };
                 let content = Content::Alias(*symbol, alias_variables, alias_variable, *kind);
 
-            register(subs, rank, pools, content)
+                register_with_known_var(subs, destination, rank, pools, content)
             }
             HostExposedAlias {
                 name: symbol,
@@ -1067,15 +1152,32 @@ fn type_to_variable<'a>(
                 lambda_set_variables,
                 ..
             } => {
-            let alias_variables = alias_to_var(
+                let alias_variables = {
-                subs,
+                    let length = type_arguments.len() + lambda_set_variables.len();
-                rank,
+                    let new_variables = VariableSubsSlice::reserve_into_subs(subs, length);
-                pools,
-                arena,
-                type_arguments,
-                lambda_set_variables,
-            );
 
+                    for (target_index, (_, arg_type)) in
+                        (new_variables.indices()).zip(type_arguments)
+                    {
+                        let copy_var = helper!(arg_type);
+                        subs.variables[target_index] = copy_var;
+                    }
+
+                    let it = (new_variables.indices().skip(type_arguments.len()))
+                        .zip(lambda_set_variables);
+                    for (target_index, ls) in it {
+                        let copy_var = helper!(&ls.0);
+                        subs.variables[target_index] = copy_var;
+                    }
+
+                    AliasVariables {
+                        variables_start: new_variables.start,
+                        type_variables_len: type_arguments.len() as _,
+                        all_variables_len: length as _,
+                    }
+                };
+
+                // cannot use helper! here because this variable may be involved in unification below
                 let alias_variable = type_to_variable(subs, rank, pools, arena, alias_type);
                 // TODO(opaques): I think host-exposed aliases should always be structural
                 // (when does it make sense to give a host an opaque type?)
@@ -1085,7 +1187,8 @@ fn type_to_variable<'a>(
                     alias_variable,
                     AliasKind::Structural,
                 );
-            let result = register(subs, rank, pools, content);
+                // let result = register(subs, rank, pools, content);
+                let result = register_with_known_var(subs, destination, rank, pools, content);
 
                 // We only want to unify the actual_var with the alias once
                 // if it's already redirected (and therefore, redundant)
@@ -1100,39 +1203,12 @@ fn type_to_variable<'a>(
             Erroneous(problem) => {
                 let content = Content::Structure(FlatType::Erroneous(Box::new(problem.clone())));
 
-            register(subs, rank, pools, content)
+                register_with_known_var(subs, destination, rank, pools, content)
-        }
             }
+        };
     }
 
-#[inline(always)]
+    result
-fn alias_to_var<'a>(
-    subs: &mut Subs,
-    rank: Rank,
-    pools: &mut Pools,
-    arena: &'a bumpalo::Bump,
-    type_arguments: &[(roc_module::ident::Lowercase, Type)],
-    lambda_set_variables: &[roc_types::types::LambdaSet],
-) -> AliasVariables {
-    let length = type_arguments.len() + lambda_set_variables.len();
-    let new_variables = VariableSubsSlice::reserve_into_subs(subs, length);
-
-    for (target_index, (_, arg_type)) in (new_variables.indices()).zip(type_arguments) {
-        let copy_var = type_to_variable(subs, rank, pools, arena, arg_type);
-        subs.variables[target_index] = copy_var;
-    }
-
-    let it = (new_variables.indices().skip(type_arguments.len())).zip(lambda_set_variables);
-    for (target_index, ls) in it {
-        let copy_var = type_to_variable(subs, rank, pools, arena, &ls.0);
-        subs.variables[target_index] = copy_var;
-    }
-
-    AliasVariables {
-        variables_start: new_variables.start,
-        type_variables_len: type_arguments.len() as _,
-        all_variables_len: length as _,
-    }
 }
 
 #[inline(always)]
@@ -2167,7 +2243,7 @@ fn register_with_known_var(
     rank: Rank,
     pools: &mut Pools,
     content: Content,
-) {
+) -> Variable {
     let descriptor = Descriptor {
         content,
         rank,
@@ -2178,4 +2254,6 @@ fn register_with_known_var(
     subs.set(var, descriptor);
 
     pools.get_mut(rank).push(var);
+
+    var
 }