Merge pull request #4855 from roc-lang/weakening-1

Print weak variables in tests and mark constraints that can't be generalized

crates/compiler/can/src/annotation.rs

@@ -520,7 +520,7 @@ pub fn find_type_def_symbols(
 }
 
 fn find_fresh_var_name(introduced_variables: &IntroducedVariables) -> Lowercase {
-    name_type_var(0, &mut introduced_variables.iter_named(), |var, str| {
+    name_type_var("", 0, &mut introduced_variables.iter_named(), |var, str| {
         var.name().as_str() == str
     })
     .0

crates/compiler/can/src/constraint.rs

@@ -390,6 +390,8 @@ impl Constraints {
         }
     }
 
+    /// A constraint that some type variables exist and should be introduced into the environment
+    /// at a given rank.
     #[inline(always)]
     pub fn exists<I>(&mut self, flex_vars: I, defs_constraint: Constraint) -> Constraint
     where
@@ -405,6 +407,10 @@ impl Constraints {
             flex_vars: self.variable_slice(flex_vars),
             def_types: DefTypes::default(),
             defs_and_ret_constraint,
+
+            // We're just marking that the variables exist, not that they should be generalized
+            // (in fact there is no return constraint, so nothing to generalize at this level)
+            generalizable: Generalizable(false),
         };
 
         let let_index = Index::new(self.let_constraints.len() as _);
@@ -431,6 +437,7 @@ impl Constraints {
             flex_vars: self.variable_slice(flex_vars),
             def_types: DefTypes::default(),
             defs_and_ret_constraint,
+            generalizable: Generalizable(false),
         };
 
         let let_index = Index::new(self.let_constraints.len() as _);
@@ -447,6 +454,7 @@ impl Constraints {
         def_types: I3,
         defs_constraint: Constraint,
         ret_constraint: Constraint,
+        generalizable: Generalizable,
     ) -> Constraint
     where
         I1: IntoIterator<Item = Variable>,
@@ -465,6 +473,7 @@ impl Constraints {
             flex_vars: self.variable_slice(flex_vars),
             def_types: self.def_types_slice(def_types),
             defs_and_ret_constraint,
+            generalizable,
         };
 
         let let_index = Index::new(self.let_constraints.len() as _);
@@ -513,6 +522,9 @@ impl Constraints {
             flex_vars: self.variable_slice(flex_vars),
             def_types: self.def_types_slice(def_types),
             defs_and_ret_constraint,
+            // If the module these variables were solved in solved them as generalized,
+            // they should be generalized here too.
+            generalizable: Generalizable(true),
         };
 
         let let_index = Index::new(self.let_constraints.len() as _);
@@ -769,12 +781,22 @@ pub struct DefTypes {
     pub loc_symbols: Slice<(Symbol, Region)>,
 }
 
+#[derive(Debug, Clone, Copy)]
+pub struct Generalizable(pub bool);
+
 #[derive(Debug, Clone)]
 pub struct LetConstraint {
     pub rigid_vars: Slice<Variable>,
     pub flex_vars: Slice<Variable>,
     pub def_types: DefTypes,
     pub defs_and_ret_constraint: Index<(Constraint, Constraint)>,
+
+    /// Whether the defs introduces in the let-binding can be generalized.
+    /// Only
+    ///   - syntactic functions
+    ///   - syntactic numbers
+    /// may be eligible for generalization, though there are other restrictions too.
+    pub generalizable: Generalizable,
 }
 
 #[derive(Debug, Clone)]

crates/compiler/constrain/src/expr.rs

@@ -8,7 +8,7 @@ use crate::builtins::{
 use crate::pattern::{constrain_pattern, PatternState};
 use roc_can::annotation::IntroducedVariables;
 use roc_can::constraint::{
-    Constraint, Constraints, ExpectedTypeIndex, OpportunisticResolve, TypeOrVar,
+    Constraint, Constraints, ExpectedTypeIndex, Generalizable, OpportunisticResolve, TypeOrVar,
 };
 use roc_can::def::Def;
 use roc_can::exhaustive::{sketch_pattern_to_rows, sketch_when_branches, ExhaustiveContext};
@@ -176,6 +176,7 @@ fn constrain_untyped_closure(
             pattern_state.headers,
             pattern_state_constraints,
             ret_constraint,
+            Generalizable(true),
         ),
         constraints.equal_types_with_storage(
             function_type,
@@ -1054,6 +1055,7 @@ pub fn constrain_expr(
                 pattern_headers,
                 pattern_constraints,
                 body_constraints,
+                Generalizable(false),
             );
 
             let result_con =
@@ -1874,6 +1876,8 @@ fn constrain_function_def(
                     argument_pattern_state.headers,
                     defs_constraint,
                     ret_constraint,
+                    // This is a syntactic function, it can be generalized
+                    Generalizable(true),
                 ),
                 constraints.equal_types_var(
                     closure_var,
@@ -1927,6 +1931,7 @@ fn constrain_function_def(
                 loc_symbol,
                 expr_var,
                 expr_type,
+                Generalizable(true), // this is a syntactic function
             )
         }
     }
@@ -2057,6 +2062,8 @@ fn constrain_value_def(
     let expr_var = declarations.variables[index];
     let opt_annotation = &declarations.annotations[index];
 
+    let generalizable = Generalizable(is_generalizable_expr(&loc_expr.value));
+
     match opt_annotation {
         Some(annotation) => {
             let arity = 1;
@@ -2122,6 +2129,7 @@ fn constrain_value_def(
                 loc_symbol,
                 expr_var,
                 signature_index,
+                generalizable,
             )
         }
         None => {
@@ -2148,6 +2156,7 @@ fn constrain_value_def(
                 loc_symbol,
                 expr_var,
                 expr_type,
+                generalizable,
             )
         }
     }
@@ -2267,7 +2276,14 @@ fn constrain_when_branch_help(
             // must introduce the headers from the pattern before constraining the guard
             let delayed_is_open_constraints = state.delayed_is_open_constraints;
             let state_constraints = constraints.and_constraint(state.constraints);
-            let inner = constraints.let_constraint([], [], [], guard_constraint, ret_constraint);
+            let inner = constraints.let_constraint(
+                [],
+                [],
+                [],
+                guard_constraint,
+                ret_constraint,
+                Generalizable(false),
+            );
 
             (state_constraints, delayed_is_open_constraints, inner)
         } else {
@@ -2377,7 +2393,14 @@ pub fn constrain_decls(
                     expected,
                 );
 
-                constraint = constraints.let_constraint([], [], [], expect_constraint, constraint)
+                constraint = constraints.let_constraint(
+                    [],
+                    [],
+                    [],
+                    expect_constraint,
+                    constraint,
+                    Generalizable(false),
+                )
             }
             ExpectationFx => {
                 let loc_expr = &declarations.expressions[index];
@@ -2398,7 +2421,14 @@ pub fn constrain_decls(
                     expected,
                 );
 
-                constraint = constraints.let_constraint([], [], [], expect_constraint, constraint)
+                constraint = constraints.let_constraint(
+                    [],
+                    [],
+                    [],
+                    expect_constraint,
+                    constraint,
+                    Generalizable(false),
+                )
             }
             Function(function_def_index) => {
                 constraint = constrain_function_def(
@@ -2653,6 +2683,8 @@ fn constrain_typed_def(
                     argument_pattern_state.headers,
                     defs_constraint,
                     ret_constraint,
+                    // This is a syntactic function, it can be generalized
+                    Generalizable(true),
                 ),
                 constraints.equal_types_var(
                     closure_var,
@@ -2675,6 +2707,7 @@ fn constrain_typed_def(
                 expr_con,
                 body_con,
                 def_pattern_state,
+                Generalizable(true),
             )
         }
 
@@ -2698,6 +2731,8 @@ fn constrain_typed_def(
             );
             let expr_con = attach_resolution_constraints(constraints, env, ret_constraint);
 
+            let generalizable = Generalizable(is_generalizable_expr(&def.loc_expr.value));
+
             constrain_def_make_constraint(
                 constraints,
                 new_rigid_variables.into_iter(),
@@ -2705,6 +2740,7 @@ fn constrain_typed_def(
                 expr_con,
                 body_con,
                 def_pattern_state,
+                generalizable,
             )
         }
     }
@@ -3011,6 +3047,8 @@ fn constrain_def(
             );
             let expr_con = attach_resolution_constraints(constraints, env, expr_con);
 
+            let generalizable = Generalizable(is_generalizable_expr(&def.loc_expr.value));
+
             constrain_def_make_constraint(
                 constraints,
                 std::iter::empty(),
@@ -3018,6 +3056,7 @@ fn constrain_def(
                 expr_con,
                 body_con,
                 def_pattern_state,
+                generalizable,
             )
         }
     }
@@ -3032,6 +3071,7 @@ pub(crate) fn constrain_def_make_constraint(
     def_expr_con: Constraint,
     after_def_con: Constraint,
     def_pattern_state: PatternState,
+    generalizable: Generalizable,
 ) -> Constraint {
     let all_flex_variables = (def_pattern_state.vars.into_iter()).chain(annotation_infer_variables);
 
@@ -3044,6 +3084,7 @@ pub(crate) fn constrain_def_make_constraint(
         def_pattern_state.headers,
         def_pattern_and_body_con,
         after_def_con,
+        generalizable,
     )
 }
 
@@ -3056,6 +3097,7 @@ fn constrain_value_def_make_constraint(
     symbol: Loc<Symbol>,
     expr_var: Variable,
     expr_type: TypeOrVar,
+    generalizable: Generalizable,
 ) -> Constraint {
     let def_con = constraints.let_constraint(
         [],
@@ -3063,11 +3105,19 @@ fn constrain_value_def_make_constraint(
         [], // empty, because our functions have no arguments!
         Constraint::True,
         expr_con,
+        generalizable,
     );
 
     let headers = [(symbol.value, Loc::at(symbol.region, expr_type))];
 
-    constraints.let_constraint(new_rigid_variables, [expr_var], headers, def_con, body_con)
+    constraints.let_constraint(
+        new_rigid_variables,
+        [expr_var],
+        headers,
+        def_con,
+        body_con,
+        generalizable,
+    )
 }
 
 fn constrain_function_def_make_constraint(
@@ -3086,6 +3136,7 @@ fn constrain_function_def_make_constraint(
         [], // empty, because our functions have no arguments!
         and_constraint,
         expr_con,
+        Generalizable(true),
     );
 
     constraints.let_constraint(
@@ -3094,6 +3145,7 @@ fn constrain_function_def_make_constraint(
         def_pattern_state.headers,
         def_con,
         body_con,
+        Generalizable(true),
     )
 }
 
@@ -3469,6 +3521,8 @@ fn constraint_recursive_function(
                     argument_pattern_state.headers,
                     state_constraints,
                     expr_con,
+                    // Syntactic function can be generalized
+                    Generalizable(true),
                 ),
                 constraints.equal_types(fn_type, annotation_expected, Category::Lambda, region),
                 // "fn_var is equal to the closure's type" - fn_var is used in code gen
@@ -3512,6 +3566,7 @@ fn constraint_recursive_function(
                     def_pattern_state.headers.clone(),
                     def_con,
                     Constraint::True,
+                    Generalizable(true),
                 )
             });
             rigid_info.def_types.extend(def_pattern_state.headers);
@@ -3540,6 +3595,18 @@ pub fn rec_defs_help_simple(
     let mut loc_symbols = Vec::with_capacity(length);
     let mut expr_regions = Vec::with_capacity(length);
 
+    // Rec defs are generalizable only if everything in the cycle is generalizable.
+    let generalizable = {
+        let generalizable = range.clone().all(|i| match declarations.declarations[i] {
+            DeclarationTag::Value => {
+                let loc_expr = &declarations.expressions[i];
+                is_generalizable_expr(&loc_expr.value)
+            }
+            _ => true, // this must be a function
+        });
+        Generalizable(generalizable)
+    };
+
     for index in range {
         // Clear the rigids from the previous iteration.
         // rigids are not shared between top-level definitions
@@ -3654,6 +3721,7 @@ pub fn rec_defs_help_simple(
                                 [], // no headers introduced (at this level)
                                 def_con,
                                 Constraint::True,
+                                generalizable,
                             ));
                             rigid_info.def_types.insert(loc_symbol.value, loc_type);
                         }
@@ -3712,6 +3780,7 @@ pub fn rec_defs_help_simple(
         hybrid_and_flex_info.def_types.clone(),
         Constraint::True,
         untyped_body_constraints,
+        generalizable,
     );
 
     // an extra constraint that propagates information to the solver to check for invalid recursion
@@ -3730,6 +3799,7 @@ pub fn rec_defs_help_simple(
         hybrid_and_flex_info.def_types,
         untyped_def_symbols_constr,
         typed_body_and_final_constr,
+        generalizable,
     );
 
     // 1. Let-generalize annotations we know.
@@ -3739,9 +3809,24 @@ pub fn rec_defs_help_simple(
         rigid_info.def_types,
         Constraint::True,
         inner,
+        generalizable,
     )
 }
 
+/// A let-bound expression is generalizable if it is
+///   - a syntactic function under an opaque wrapper
+///   - a number literal under an opaque wrapper
+fn is_generalizable_expr(mut expr: &Expr) -> bool {
+    loop {
+        match expr {
+            Num(..) | Int(..) | Float(..) => return true,
+            Closure(_) => return true,
+            OpaqueRef { argument, .. } => expr = &argument.1.value,
+            _ => return false,
+        }
+    }
+}
+
 fn constrain_recursive_defs(
     types: &mut Types,
     constraints: &mut Constraints,
@@ -3768,6 +3853,14 @@ fn rec_defs_help(
     let mut rigid_info = Info::with_capacity(defs.len());
     let mut hybrid_and_flex_info = Info::with_capacity(defs.len());
 
+    // Rec defs are generalizable only if everything in the cycle is generalizable.
+    let generalizable = {
+        let generalizable = defs
+            .iter()
+            .all(|d| is_generalizable_expr(&d.loc_expr.value));
+        Generalizable(generalizable)
+    };
+
     for def in defs {
         let expr_var = def.expr_var;
         let expr_type_index = constraints.push_variable(expr_var);
@@ -3927,6 +4020,7 @@ fn rec_defs_help(
                                 state.headers,
                                 state_constraints,
                                 expr_con,
+                                generalizable,
                             ),
                             constraints.equal_types(
                                 fn_type_index,
@@ -3968,6 +4062,7 @@ fn rec_defs_help(
                                 [], // no headers introduced (at this level)
                                 def_con,
                                 Constraint::True,
+                                generalizable,
                             ));
                             rigid_info.def_types.extend(def_pattern_state.headers);
                         }
@@ -4013,6 +4108,7 @@ fn rec_defs_help(
                                 [], // no headers introduced (at this level)
                                 def_con,
                                 Constraint::True,
+                                generalizable,
                             ));
                             rigid_info.def_types.extend(def_pattern_state.headers);
                         }
@@ -4056,6 +4152,7 @@ fn rec_defs_help(
         hybrid_and_flex_info.def_types.clone(),
         Constraint::True,
         untyped_body_constraints,
+        generalizable,
     );
 
     // an extra constraint that propagates information to the solver to check for invalid recursion
@@ -4082,6 +4179,7 @@ fn rec_defs_help(
         untyped_def_symbols_constr,
         // 4 + 5. Solve the typed body defs, and the rest of the program.
         typed_body_and_final_constr,
+        generalizable,
     );
 
     // 1. Let-generalize annotations we know.
@@ -4091,6 +4189,7 @@ fn rec_defs_help(
         rigid_info.def_types,
         Constraint::True,
         inner,
+        generalizable,
     )
 }
 

crates/compiler/constrain/src/module.rs

@@ -1,6 +1,6 @@
 use crate::expr::{constrain_def_make_constraint, constrain_def_pattern, Env};
 use roc_can::abilities::{PendingAbilitiesStore, PendingMemberType};
-use roc_can::constraint::{Constraint, Constraints};
+use roc_can::constraint::{Constraint, Constraints, Generalizable};
 use roc_can::expected::Expected;
 use roc_can::expr::Declarations;
 use roc_can::pattern::Pattern;
@@ -76,6 +76,7 @@ fn constrain_symbols_from_requires(
                     Constraint::True,
                     constraint,
                     def_pattern_state,
+                    Generalizable(true),
                 )
             } else {
                 // Otherwise, this symbol comes from an app module - we want to check that the type
@@ -158,6 +159,8 @@ pub fn frontload_ability_constraints(
                 Constraint::True,
                 constraint,
                 def_pattern_state,
+                // Ability signatures can always be generalized since they're prototypical
+                Generalizable(true),
             );
         }
     }

crates/compiler/solve/tests/solve_expr.rs

@@ -332,6 +332,7 @@ mod solve_expr {
                     print_lambda_sets: true,
                     print_only_under_alias: options.print_only_under_alias,
                     ignore_polarity: true,
+                    print_weakened_vars: true,
                 },
             );
             subs.rollback_to(snapshot);

crates/compiler/types/src/pretty_print.rs

@@ -56,6 +56,7 @@ pub struct DebugPrint {
     pub print_lambda_sets: bool,
     pub print_only_under_alias: bool,
     pub ignore_polarity: bool,
+    pub print_weakened_vars: bool,
 }
 
 impl DebugPrint {
@@ -63,6 +64,7 @@ impl DebugPrint {
         print_lambda_sets: false,
         print_only_under_alias: false,
         ignore_polarity: false,
+        print_weakened_vars: false,
     };
 }
 
@@ -380,11 +382,7 @@ struct NamedResult {
     recursion_structs_to_expand: Vec<Variable>,
 }
 
-fn name_all_type_vars(
-    variable: Variable,
-    subs: &mut Subs,
-    find_names_under_alias: bool,
-) -> NamedResult {
+fn name_all_type_vars(variable: Variable, subs: &mut Subs, debug_print: DebugPrint) -> NamedResult {
     let mut roots = Vec::new();
     let mut letters_used = 0;
     let mut appearances = MutMap::default();
@@ -397,7 +395,7 @@ fn name_all_type_vars(
         &mut roots,
         &mut appearances,
         &mut taken,
-        find_names_under_alias,
+        debug_print.print_only_under_alias,
     );
 
     let mut recursion_structs_to_expand = vec![];
@@ -407,14 +405,14 @@ fn name_all_type_vars(
         // set_root_name(root, (format!("<{:?}>", root).into()), subs);
         match appearances.get(&root) {
             Some(Appearances::Multiple) => {
-                letters_used = name_root(letters_used, root, subs, &mut taken);
+                letters_used = name_root(letters_used, root, subs, &mut taken, debug_print);
             }
             Some(Appearances::Single) => {
                 if let Content::RecursionVar { structure, .. } =
                     subs.get_content_without_compacting(root)
                 {
                     recursion_structs_to_expand.push(*structure);
-                    letters_used = name_root(letters_used, root, subs, &mut taken);
+                    letters_used = name_root(letters_used, root, subs, &mut taken, debug_print);
                 }
             }
             _ => {}
@@ -426,12 +424,29 @@ fn name_all_type_vars(
     }
 }
 
+fn is_weakened_unbound(subs: &Subs, var: Variable) -> bool {
+    use Content::*;
+    let desc = subs.get_without_compacting(var);
+    !desc.rank.is_none()
+        && !matches!(
+            desc.content,
+            FlexVar(_) | RigidVar(_) | FlexAbleVar(..) | RigidAbleVar(..)
+        )
+}
+
 fn name_root(
     letters_used: u32,
     root: Variable,
     subs: &mut Subs,
     taken: &mut MutMap<Lowercase, Variable>,
+    debug_print: DebugPrint,
 ) -> u32 {
+    let prefix = if debug_print.print_weakened_vars && is_weakened_unbound(subs, root) {
+        "w_" // weakened variable
+    } else {
+        ""
+    };
+
     let (generated_name, new_letters_used) = match subs.get_content_unchecked(root) {
         Content::FlexVar(Some(name))
         | Content::RigidVar(name)
@@ -444,19 +459,21 @@ fn name_root(
             let name_hint = &subs[*name];
             if name_hint.as_str() == "*" {
                 // Give a proper name to named wildcards!
-                name_type_var(letters_used, &mut taken.keys(), |var, str| {
+                name_type_var(prefix, letters_used, &mut taken.keys(), |var, str| {
                     var.as_str() == str
                 })
             } else {
-                let generated =
-                    name_type_var_with_hint(name_hint.as_str(), &mut taken.keys(), |var, str| {
-                        var.as_str() == str
-                    });
+                let generated = name_type_var_with_hint(
+                    prefix,
+                    name_hint.as_str(),
+                    &mut taken.keys(),
+                    |var, str| var.as_str() == str,
+                );
 
                 (generated, letters_used)
             }
         }
-        _ => name_type_var(letters_used, &mut taken.keys(), |var, str| {
+        _ => name_type_var(prefix, letters_used, &mut taken.keys(), |var, str| {
             var.as_str() == str
         }),
     };
@@ -568,7 +585,7 @@ pub fn name_and_print_var(
     interns: &Interns,
     debug_print: DebugPrint,
 ) -> String {
-    let named_result = name_all_type_vars(var, subs, debug_print.print_only_under_alias);
+    let named_result = name_all_type_vars(var, subs, debug_print);
     let content = subs.get_content_without_compacting(var);
     content_to_string(
         content,

crates/compiler/types/src/subs.rs

@@ -2185,6 +2185,7 @@ pub struct Rank(u32);
 impl Rank {
     pub const NONE: Rank = Rank(0);
 
+    /// The generalized rank
     pub fn is_none(&self) -> bool {
         *self == Self::NONE
     }
@@ -4095,10 +4096,12 @@ fn get_fresh_error_var_name(state: &mut ErrorTypeState) -> Lowercase {
     //
     // We want to claim both the "#name" and "name" forms, because if "#name" appears multiple
     // times during error type reporting, we'll use "name" for display.
-    let (name, new_index) =
-        name_type_var(state.letters_used, &mut state.taken.iter(), |var, str| {
-            var.as_str() == str
-        });
+    let (name, new_index) = name_type_var(
+        "",
+        state.letters_used,
+        &mut state.taken.iter(),
+        |var, str| var.as_str() == str,
+    );
 
     state.letters_used = new_index;
 

crates/compiler/types/src/types.rs

@@ -4095,13 +4095,16 @@ static THE_LETTER_A: u32 = 'a' as u32;
 
 /// Generates a fresh type variable name, composed of lowercase alphabetic characters in sequence.
 pub fn name_type_var<I, F: FnMut(&I, &str) -> bool>(
+    prefix: &str,
     letters_used: u32,
     taken: &mut impl Iterator<Item = I>,
     mut predicate: F,
 ) -> (Lowercase, u32) {
     // TODO we should arena-allocate this String,
     // so all the strings in the entire pass only require ~1 allocation.
-    let mut buf = String::with_capacity((letters_used as usize) / 26 + 1);
+    let mut buf = String::with_capacity(prefix.len() + (letters_used as usize) / 26 + 1);
+
+    buf.push_str(prefix);
 
     let is_taken = {
         let mut remaining = letters_used as i32;
@@ -4118,7 +4121,7 @@ pub fn name_type_var<I, F: FnMut(&I, &str) -> bool>(
 
     if is_taken {
         // If the generated name is already taken, try again.
-        name_type_var(letters_used + 1, taken, predicate)
+        name_type_var(prefix, letters_used + 1, taken, predicate)
     } else {
         (buf.into(), letters_used + 1)
     }
@@ -4127,18 +4130,19 @@ pub fn name_type_var<I, F: FnMut(&I, &str) -> bool>(
 /// Generates a fresh type variable name given a hint, composed of the hint as a prefix and a
 /// number as a suffix. For example, given hint `a` we'll name the variable `a`, `a1`, or `a27`.
 pub fn name_type_var_with_hint<I, F: FnMut(&I, &str) -> bool>(
+    prefix: &str,
     hint: &str,
     taken: &mut impl Iterator<Item = I>,
     mut predicate: F,
 ) -> Lowercase {
     if !taken.any(|item| predicate(&item, hint)) {
-        return hint.into();
+        return format!("{prefix}{hint}").into();
     }
 
     let mut i = 0;
     loop {
         i += 1;
-        let cand = format!("{}{}", hint, i);
+        let cand = format!("{prefix}{hint}{i}");
 
         if !taken.any(|item| predicate(&item, &cand)) {
             return cand.into();