New data structure for storing sorted abilities bound to a variable

crates/compiler/can/src/annotation.rs

@@ -9,8 +9,8 @@ use roc_problem::can::ShadowKind;
 use roc_region::all::{Loc, Region};
 use roc_types::subs::{VarStore, Variable};
 use roc_types::types::{
-    name_type_var, Alias, AliasCommon, AliasKind, AliasVar, LambdaSet, OptAbleType, OptAbleVar,
-    Problem, RecordField, Type, TypeExtension,
+    name_type_var, AbilitySet, Alias, AliasCommon, AliasKind, AliasVar, LambdaSet, OptAbleType,
+    OptAbleVar, Problem, RecordField, Type, TypeExtension,
 };
 
 #[derive(Clone, Debug)]
@@ -105,7 +105,7 @@ impl OwnedNamedOrAble {
         }
     }
 
-    pub fn opt_abilities(&self) -> Option<&VecSet<Symbol>> {
+    pub fn opt_abilities(&self) -> Option<&AbilitySet> {
         match self {
             OwnedNamedOrAble::Named(_) => None,
             OwnedNamedOrAble::Able(av) => Some(&av.abilities),
@@ -127,9 +127,7 @@ pub struct NamedVariable {
 pub struct AbleVariable {
     pub variable: Variable,
     pub name: Lowercase,
-    /// Abilities bound to this type variable.
-    /// INVARIANT: sorted and de-duplicated.
-    pub abilities: VecSet<Symbol>,
+    pub abilities: AbilitySet,
     // NB: there may be multiple occurrences of a variable
     pub first_seen: Region,
 }
@@ -168,7 +166,7 @@ impl IntroducedVariables {
         self.named.insert(named_variable);
     }
 
-    pub fn insert_able(&mut self, name: Lowercase, var: Loc<Variable>, abilities: VecSet<Symbol>) {
+    pub fn insert_able(&mut self, name: Lowercase, var: Loc<Variable>, abilities: AbilitySet) {
         self.debug_assert_not_already_present(var.value);
 
         let able_variable = AbleVariable {
@@ -544,7 +542,7 @@ fn can_annotation_help(
                 introduced_variables.insert_able(
                     fresh_ty_var,
                     Loc::at(region, var),
-                    VecSet::singleton(symbol),
+                    AbilitySet::singleton(symbol),
                 );
                 return Type::Variable(var);
             }
@@ -936,7 +934,7 @@ fn canonicalize_has_clause(
     );
     let var_name = Lowercase::from(var_name);
 
-    let mut can_abilities = VecSet::with_capacity(abilities.len());
+    let mut can_abilities = AbilitySet::with_capacity(abilities.len());
     for &Loc {
         region,
         value: ability,
@@ -986,11 +984,6 @@ fn canonicalize_has_clause(
 
     let var = var_store.fresh();
 
-    let can_abilities = {
-        let mut vec = can_abilities.into_vec();
-        vec.sort();
-        VecSet::from_iter(vec)
-    };
     introduced_variables.insert_able(var_name, Loc::at(region, var), can_abilities);
 
     Ok(())

crates/compiler/can/src/def.rs

@@ -1345,7 +1345,7 @@ fn resolve_abilities<'a>(
                 .introduced_variables
                 .able
                 .iter()
-                .partition(|av| av.abilities.iter().any(|ab| *ab == ability));
+                .partition(|av| av.abilities.contains(&ability));
 
             let var_bound_to_ability = match variables_bound_to_ability.as_slice() {
                 [one] => one.variable,

crates/compiler/can/src/module.rs

@@ -18,7 +18,7 @@ use roc_parse::pattern::PatternType;
 use roc_problem::can::{Problem, RuntimeError};
 use roc_region::all::{Loc, Region};
 use roc_types::subs::{ExposedTypesStorageSubs, Subs, VarStore, Variable};
-use roc_types::types::{Alias, AliasKind, AliasVar, Type};
+use roc_types::types::{AbilitySet, Alias, AliasKind, AliasVar, Type};
 
 /// The types of all exposed values/functions of a collection of modules
 #[derive(Clone, Debug, Default)]
@@ -136,7 +136,7 @@ pub struct Module {
 #[derive(Debug, Default)]
 pub struct RigidVariables {
     pub named: MutMap<Variable, Lowercase>,
-    pub able: MutMap<Variable, (Lowercase, VecSet<Symbol>)>,
+    pub able: MutMap<Variable, (Lowercase, AbilitySet)>,
     pub wildcards: VecSet<Variable>,
 }
 

crates/compiler/solve/src/module.rs

@@ -71,8 +71,7 @@ pub fn run_solve(
     }
 
     for (var, (name, abilities)) in rigid_variables.able {
-        // TODO(abilities)
-        subs.rigid_able_var(var, name, abilities[0]);
+        subs.rigid_able_var(var, name, abilities);
     }
 
     for var in rigid_variables.wildcards {

crates/compiler/solve/src/solve.rs

@@ -2609,7 +2609,7 @@ fn type_to_variable<'a>(
                                         match *subs.get_content_without_compacting(var) {
                                             FlexVar(opt_name) => {
                                                 // TODO(multi-abilities): check run cache
-                                                let abilities_slice = SubsSlice::extend_new(&mut subs.symbol_names, abilities.iter().copied());
+                                                let abilities_slice = SubsSlice::extend_new(&mut subs.symbol_names, abilities.sorted_iter().copied());
                                                 subs.set_content(var, FlexAbleVar(opt_name, abilities_slice))
                                             },
                                             RigidVar(..) => internal_error!("Rigid var in type arg for {:?} - this is a bug in the solver, or our understanding", actual),

crates/compiler/types/src/subs.rs

@@ -1,6 +1,7 @@
 #![deny(unsafe_op_in_unsafe_fn)]
 use crate::types::{
-    name_type_var, AliasKind, ErrorType, Problem, RecordField, RecordFieldsError, TypeExt, Uls,
+    name_type_var, AbilitySet, AliasKind, ErrorType, Problem, RecordField, RecordFieldsError,
+    TypeExt, Uls,
 };
 use roc_collections::all::{FnvMap, ImMap, ImSet, MutSet, SendMap};
 use roc_collections::{VecMap, VecSet};
@@ -1807,10 +1808,9 @@ impl Subs {
         self.set(var, desc);
     }
 
-    pub fn rigid_able_var(&mut self, var: Variable, name: Lowercase, ability: Symbol) {
+    pub fn rigid_able_var(&mut self, var: Variable, name: Lowercase, abilities: AbilitySet) {
         let name_index = SubsIndex::push_new(&mut self.field_names, name);
-        // TODO(multi-abilities)
-        let abilities = SubsSlice::extend_new(&mut self.symbol_names, [ability]);
+        let abilities = SubsSlice::extend_new(&mut self.symbol_names, abilities.into_sorted_iter());
         let content = Content::RigidAbleVar(name_index, abilities);
         let desc = Descriptor::from(content);
 

crates/compiler/types/src/types.rs

@@ -4,7 +4,6 @@ use crate::subs::{
     GetSubsSlice, RecordFields, Subs, UnionTags, VarStore, Variable, VariableSubsSlice,
 };
 use roc_collections::all::{HumanIndex, ImMap, ImSet, MutMap, MutSet, SendMap};
-use roc_collections::VecSet;
 use roc_error_macros::internal_error;
 use roc_module::called_via::CalledVia;
 use roc_module::ident::{ForeignSymbol, Ident, Lowercase, TagName};
@@ -245,10 +244,57 @@ pub struct AliasCommon {
     pub lambda_set_variables: Vec<LambdaSet>,
 }
 
+/// Represents a collection of abilities bound to a type variable.
+///
+/// Enforces the invariants
+///   - There are no duplicate abilities (like a [VecSet][roc_collections::VecSet])
+///   - Inserted abilities are in sorted order; they can be extracted with
+///     [AbilitySet::into_sorted_iter]
+///
+/// This is useful for inserting into [Subs][crate::subs::Subs], so that the set need not be
+/// re-sorted.
+///
+/// In the future we might want to do some small-vec optimizations, though that may be trivialized
+/// away with a SoA representation of canonicalized types.
+#[derive(Clone, Debug, Default, PartialEq, PartialOrd, Eq, Ord)]
+pub struct AbilitySet(Vec<Symbol>);
+
+impl AbilitySet {
+    pub fn with_capacity(cap: usize) -> Self {
+        Self(Vec::with_capacity(cap))
+    }
+
+    pub fn singleton(ability: Symbol) -> Self {
+        Self(vec![ability])
+    }
+
+    pub fn insert(&mut self, ability: Symbol) -> bool {
+        match self.0.binary_search(&ability) {
+            Ok(_) => true,
+            Err(insert_index) => {
+                self.0.insert(insert_index, ability);
+                false
+            }
+        }
+    }
+
+    pub fn contains(&self, ability: &Symbol) -> bool {
+        self.0.contains(ability)
+    }
+
+    pub fn sorted_iter(&self) -> impl Iterator<Item = &Symbol> {
+        self.0.iter()
+    }
+
+    pub fn into_sorted_iter(self) -> impl Iterator<Item = Symbol> {
+        self.0.into_iter()
+    }
+}
+
 #[derive(Clone, Debug)]
 pub struct OptAbleVar {
     pub var: Variable,
-    pub opt_abilities: Option<VecSet<Symbol>>,
+    pub opt_abilities: Option<AbilitySet>,
 }
 
 impl OptAbleVar {
@@ -263,7 +309,7 @@ impl OptAbleVar {
 #[derive(PartialEq, Eq, Debug)]
 pub struct OptAbleType {
     pub typ: Type,
-    pub opt_abilities: Option<VecSet<Symbol>>,
+    pub opt_abilities: Option<AbilitySet>,
 }
 
 impl OptAbleType {
@@ -2113,8 +2159,7 @@ pub struct AliasVar {
     pub name: Lowercase,
     pub var: Variable,
     /// `Some` if this variable is bound to abilities; `None` otherwise.
-    /// INVARIANT: if abilities are present, they are sorted and de-duplicated.
-    pub opt_bound_abilities: Option<VecSet<Symbol>>,
+    pub opt_bound_abilities: Option<AbilitySet>,
 }
 
 impl AliasVar {