refactor: Generalize Substitutable::apply_ref

libflux/flux-core/src/errors.rs

@@ -12,10 +12,7 @@ use derive_more::Display;
 
 use crate::{
     ast,
-    semantic::{
-        sub::{Substitutable, Substituter},
-        types::Tvar,
-    },
+    semantic::sub::{Substitutable, Substituter},
 };
 
 #[derive(Clone, Debug, Eq, PartialEq, Hash)]
@@ -205,15 +202,12 @@ impl<E> Substitutable for Located<E>
 where
     E: Substitutable,
 {
-    fn apply_ref(&self, sub: &dyn Substituter) -> Option<Self> {
-        self.error.apply_ref(sub).map(|error| Located {
+    fn walk(&self, sub: &dyn Substituter) -> Option<Self> {
+        self.error.visit(sub).map(|error| Located {
             location: self.location.clone(),
             error,
         })
     }
-    fn free_vars(&self, vars: &mut Vec<Tvar>) {
-        self.error.free_vars(vars)
-    }
 }
 
 pub(crate) trait AsDiagnostic {

libflux/flux-core/src/semantic/bootstrap.rs

@@ -297,7 +297,7 @@ fn add_record_to_map(
     for field in r.fields() {
         let new_vars = {
             let new_vars = CollectBoundVars(RefCell::new(Vec::new()));
-            field.v.apply_ref(&new_vars);
+            field.v.visit(&new_vars);
             new_vars.0.into_inner()
         };
 

libflux/flux-core/src/semantic/env.rs

@@ -4,7 +4,7 @@ use std::{fmt, mem};
 use crate::semantic::{
     nodes::Symbol,
     sub::{apply2, Substitutable, Substituter},
-    types::{PolyType, PolyTypeHashMap, PolyTypeMap, Tvar},
+    types::{PolyType, PolyTypeHashMap, PolyTypeMap},
     PackageExports,
 };
 
@@ -34,7 +34,7 @@ impl fmt::Display for Environment<'_> {
 }
 
 impl Substitutable for Environment<'_> {
-    fn apply_ref(&self, sub: &dyn Substituter) -> Option<Self> {
+    fn walk(&self, sub: &dyn Substituter) -> Option<Self> {
         match (self.readwrite, &self.parent) {
             // This is a performance optimization where false implies
             // this is the top-level of the type environment and apply
@@ -43,7 +43,7 @@ impl Substitutable for Environment<'_> {
             // Even though this is the top-level of the type environment
             // and apply should be a no-op, readwrite is set to true so
             // we apply anyway.
-            (true, None) => self.values.apply_ref(sub).map(|values| Environment {
+            (true, None) => self.values.visit(sub).map(|values| Environment {
                 external: self.external,
                 parent: None,
                 values,
@@ -59,16 +59,6 @@ impl Substitutable for Environment<'_> {
             }
         }
     }
-    fn free_vars(&self, vars: &mut Vec<Tvar>) {
-        match (self.readwrite, &self.parent) {
-            (false, None) | (false, _) => (),
-            (true, None) => self.values.free_vars(vars),
-            (true, Some(env)) => {
-                env.free_vars(vars);
-                self.values.free_vars(vars);
-            }
-        }
-    }
 }
 
 // Derive a type environment from a hash map

libflux/flux-core/src/semantic/infer.rs

@@ -110,15 +110,12 @@ pub struct Error {
 impl std::error::Error for Error {}
 
 impl Substitutable for Error {
-    fn apply_ref(&self, sub: &dyn Substituter) -> Option<Self> {
-        self.err.apply_ref(sub).map(|err| Error {
+    fn walk(&self, sub: &dyn Substituter) -> Option<Self> {
+        self.err.visit(sub).map(|err| Error {
             loc: self.loc.clone(),
             err,
         })
     }
-    fn free_vars(&self, vars: &mut Vec<Tvar>) {
-        self.err.free_vars(vars)
-    }
 }
 
 // Solve a set of type constraints
@@ -222,7 +219,7 @@ pub(crate) fn temporary_generalize(
     }
 
     let generalize = Generalize {
-        env_free_vars: env.mk_free_vars(),
+        env_free_vars: env.free_vars(),
         vars: Default::default(),
     };
     let t = t.apply(&generalize);
@@ -282,7 +279,7 @@ pub fn generalize(env: &Environment, sub: &mut Substitution, t: MonoType) -> Pol
     }
 
     let generalize = Generalize {
-        env_free_vars: env.mk_free_vars(),
+        env_free_vars: env.free_vars(),
         sub,
         vars: Default::default(),
     };

libflux/flux-core/src/semantic/nodes.rs

@@ -80,9 +80,9 @@ impl AsDiagnostic for ErrorKind {
 }
 
 impl Substitutable for ErrorKind {
-    fn apply_ref(&self, sub: &dyn Substituter) -> Option<Self> {
+    fn walk(&self, sub: &dyn Substituter) -> Option<Self> {
         match self {
-            Self::Inference(err) => err.apply_ref(sub).map(Self::Inference),
+            Self::Inference(err) => err.visit(sub).map(Self::Inference),
             Self::UndefinedBuiltin(_)
             | Self::UndefinedIdentifier(_)
             | Self::InvalidBinOp(_)
@@ -93,19 +93,6 @@ impl Substitutable for ErrorKind {
             | Self::Bug(_) => None,
         }
     }
-    fn free_vars(&self, vars: &mut Vec<Tvar>) {
-        match self {
-            Self::Inference(err) => err.free_vars(vars),
-            Self::UndefinedBuiltin(_)
-            | Self::UndefinedIdentifier(_)
-            | Self::InvalidBinOp(_)
-            | Self::InvalidUnaryOp(_)
-            | Self::InvalidImportPath(_)
-            | Self::UnableToVectorize(_)
-            | Self::InvalidReturn
-            | Self::Bug(_) => (),
-        }
-    }
 }
 
 impl From<types::Error> for ErrorKind {

libflux/flux-core/src/semantic/sub.rs

@@ -1,7 +1,7 @@
 //! Substitutions during type inference.
 use std::{borrow::Cow, cell::RefCell, iter::FusedIterator};
 
-use crate::semantic::types::{union, Error, MonoType, SubstitutionMap, Tvar, TvarKinds};
+use crate::semantic::types::{union, Error, MonoType, PolyType, SubstitutionMap, Tvar, TvarKinds};
 
 /// A substitution defines a function that takes a monotype as input
 /// and returns a monotype as output. The output type is interpreted
@@ -154,15 +154,15 @@ pub trait Substitutable {
     where
         Self: Sized,
     {
-        self.apply_ref(sub).unwrap_or(self)
+        self.visit(sub).unwrap_or(self)
     }
 
     /// Apply a substitution to a type variable.
     fn apply_mut(&mut self, sub: &dyn Substituter)
     where
         Self: Sized,
     {
-        if let Some(new) = self.apply_ref(sub) {
+        if let Some(new) = self.visit(sub) {
             *self = new;
         }
     }
@@ -172,38 +172,71 @@ pub trait Substitutable {
     where
         Self: Clone + Sized,
     {
-        match self.apply_ref(sub) {
+        match self.visit(sub) {
             Some(t) => Cow::Owned(t),
             None => Cow::Borrowed(self),
         }
     }
 
     /// Apply a substitution to a type variable. Should return `None` if there was nothing to apply
     /// which allows for optimizations.
-    fn apply_ref(&self, sub: &dyn Substituter) -> Option<Self>
+    fn visit(&self, sub: &dyn Substituter) -> Option<Self>
+    where
+        Self: Sized,
+    {
+        self.walk(sub)
+    }
+
+    /// Apply a substitution to a type variable. Should return `None` if there was nothing to apply
+    /// which allows for optimizations.
+    fn walk(&self, sub: &dyn Substituter) -> Option<Self>
     where
         Self: Sized;
 
     /// Get all free type variables in a type.
-    fn mk_free_vars(&self) -> Vec<Tvar> {
-        let mut vars = Vec::new();
-        self.free_vars(&mut vars);
-        vars
-    }
+    fn free_vars(&self) -> Vec<Tvar>
+    where
+        Self: Sized,
+    {
+        #[derive(Default)]
+        struct FreeVars {
+            vars: RefCell<Vec<Tvar>>,
+        }
 
-    /// Get all free type variables in a type.
-    fn free_vars(&self, vars: &mut Vec<Tvar>);
+        impl Substituter for FreeVars {
+            fn try_apply(&self, var: Tvar) -> Option<MonoType> {
+                let mut vars = self.vars.borrow_mut();
+                if let Err(i) = vars.binary_search(&var) {
+                    vars.insert(i, var);
+                }
+                None
+            }
+
+            fn visit_poly_type_spec(
+                &self,
+                sub: &dyn Substituter,
+                typ: &PolyType,
+            ) -> Option<PolyType> {
+                typ.expr.visit(sub);
+                self.vars.borrow_mut().retain(|v| !typ.vars.contains(v));
+                None
+            }
+        }
+
+        let free_vars = FreeVars::default();
+
+        self.visit(&free_vars);
+
+        free_vars.vars.into_inner()
+    }
 }
 
 impl<T> Substitutable for Box<T>
 where
     T: Substitutable,
 {
-    fn apply_ref(&self, sub: &dyn Substituter) -> Option<Self> {
-        T::apply_ref(self, sub).map(Box::new)
-    }
-    fn free_vars(&self, vars: &mut Vec<Tvar>) {
-        T::free_vars(self, vars)
+    fn walk(&self, sub: &dyn Substituter) -> Option<Self> {
+        T::visit(self, sub).map(Box::new)
     }
 }
 
@@ -218,6 +251,32 @@ pub trait Substituter {
         let _ = var;
         None
     }
+
+    // Hack to allow `visit_poly_type_spec` to be implemented both here as a default and in `impl`
+    // blocks. `self` and `sub` should refer to the same object, but passing `sub` lets us call
+    // `walk` without needing a `Self: Sized` bound.
+    #[doc(hidden)]
+    fn visit_poly_type_spec(&self, sub: &dyn Substituter, typ: &PolyType) -> Option<PolyType> {
+        typ.walk(sub)
+    }
+
+    /// Apply a substitution to a type, returning None if there is no substitution for the
+    /// type.
+    fn visit_type(&self, typ: &MonoType) -> Option<MonoType> {
+        match *typ {
+            MonoType::Var(var) => self.try_apply(var),
+            MonoType::BoundVar(var) => self.try_apply_bound(var),
+            _ => None,
+        }
+    }
+}
+
+impl<'a> dyn Substituter + 'a {
+    /// Apply a substitution to a polytype, returning None if there is no substitution for the
+    /// type.
+    pub fn visit_poly_type(&self, typ: &PolyType) -> Option<PolyType> {
+        self.visit_poly_type_spec(self, typ)
+    }
 }
 
 impl<F> Substituter for F
@@ -233,12 +292,48 @@ impl Substituter for SubstitutionMap {
     fn try_apply(&self, var: Tvar) -> Option<MonoType> {
         self.get(&var).cloned()
     }
+
+    fn visit_poly_type_spec(&self, sub: &dyn Substituter, typ: &PolyType) -> Option<PolyType> {
+        // `vars` defines new distinct variables for `expr` so any substitutions applied on a
+        // variable named the same must not be applied in `expr`
+        typ.expr
+            .visit(&|var| {
+                if typ.vars.contains(&var) {
+                    None
+                } else {
+                    sub.try_apply(var)
+                }
+            })
+            .map(|expr| PolyType {
+                vars: typ.vars.clone(),
+                cons: typ.cons.clone(),
+                expr,
+            })
+    }
 }
 
 impl Substituter for Substitution {
     fn try_apply(&self, var: Tvar) -> Option<MonoType> {
         Substitution::try_apply(self, var)
     }
+
+    fn visit_poly_type_spec(&self, sub: &dyn Substituter, typ: &PolyType) -> Option<PolyType> {
+        // `vars` defines new distinct variables for `expr` so any substitutions applied on a
+        // variable named the same must not be applied in `expr`
+        typ.expr
+            .visit(&|var| {
+                if typ.vars.contains(&var) {
+                    None
+                } else {
+                    sub.try_apply(var)
+                }
+            })
+            .map(|expr| PolyType {
+                vars: typ.vars.clone(),
+                cons: typ.cons.clone(),
+                expr,
+            })
+    }
 }
 
 pub(crate) struct BindVars<'a> {
@@ -285,13 +380,13 @@ where
 {
     merge4(
         a,
-        a.apply_ref(sub),
+        a.visit(sub),
         b,
-        b.apply_ref(sub),
+        b.visit(sub),
         c,
-        c.apply_ref(sub),
+        c.visit(sub),
         d,
-        d.apply_ref(sub),
+        d.visit(sub),
     )
 }
 
@@ -301,22 +396,15 @@ where
     B: Substitutable + Clone,
     C: Substitutable + Clone,
 {
-    merge3(
-        a,
-        a.apply_ref(sub),
-        b,
-        b.apply_ref(sub),
-        c,
-        c.apply_ref(sub),
-    )
+    merge3(a, a.visit(sub), b, b.visit(sub), c, c.visit(sub))
 }
 
 pub(crate) fn apply2<A, B>(a: &A, b: &B, sub: &dyn Substituter) -> Option<(A, B)>
 where
     A: Substitutable + Clone,
     B: Substitutable + Clone,
 {
-    merge(a, a.apply_ref(sub), b, b.apply_ref(sub))
+    merge(a, a.visit(sub), b, b.visit(sub))
 }
 
 #[allow(clippy::too_many_arguments, clippy::type_complexity)]