Track canonical names in CheckedModule

Author: purefunctor

compiler-core/checking2/src/core.rs

@@ -52,8 +52,6 @@ pub struct ForallBinder {
     pub visible: bool,
     /// The unique identity attached to the type variable.
     pub name: Name,
-    /// The source-level text of the type variable.
-    pub text: SmolStrId,
     /// The kind of the type variable.
     pub kind: TypeId,
 }

compiler-core/checking2/src/core/generalise.rs

@@ -216,9 +216,7 @@ where
             }
         };
 
-        let text = context.queries.intern_smol_str(name.as_text());
-
-        let binder = ForallBinder { visible: false, name, text, kind };
+        let binder = ForallBinder { visible: false, name, kind };
         let binder = context.intern_forall_binder(binder);
         quantified = context.intern_forall(binder, quantified);
     }
@@ -321,15 +319,13 @@ where
     for implicit_unification in implicits_unifications {
         match implicit_unification {
             ImplicitOrUnification::Implicit(name, kind) => {
-                let text = context.queries.intern_smol_str(name.as_text());
-                binders.push(ForallBinder { visible: false, name, text, kind })
+                binders.push(ForallBinder { visible: false, name, kind })
             }
             ImplicitOrUnification::Unification(id, kind) => {
                 let name = state.names.fresh();
-                let text = context.queries.intern_smol_str(name.as_text());
                 let rigid = context.intern_rigid(name, depth, kind);
                 state.unifications.solve(id, rigid);
-                binders.push(ForallBinder { visible: false, name, text, kind })
+                binders.push(ForallBinder { visible: false, name, kind })
             }
         }
     }

compiler-core/checking2/src/core/pretty.rs

@@ -6,24 +6,24 @@ use pretty::{Arena, DocAllocator, DocBuilder};
 use rustc_hash::FxHashMap;
 use smol_str::{SmolStr, SmolStrBuilder};
 
-use crate::ExternalQueries;
 use crate::core::{
     ForallBinder, ForallBinderId, Name, RowField, RowType, RowTypeId, SmolStrId, Synonym,
     SynonymId, Type, TypeId,
 };
+use crate::{CheckedModule, ExternalQueries};
 
 type Doc<'a> = DocBuilder<'a, Arena<'a>, ()>;
 
 pub struct Pretty<'a, Q> {
     queries: &'a Q,
     width: usize,
     signature: Option<&'a str>,
-    names: FxHashMap<Name, SmolStr>,
+    checked: &'a CheckedModule,
 }
 
 impl<'a, Q: ExternalQueries> Pretty<'a, Q> {
-    pub fn new(queries: &'a Q) -> Self {
-        Pretty { queries, width: 100, signature: None, names: FxHashMap::default() }
+    pub fn new(queries: &'a Q, checked: &'a CheckedModule) -> Self {
+        Pretty { queries, width: 100, signature: None, checked }
     }
 
     pub fn width(mut self, width: usize) -> Self {
@@ -36,15 +36,9 @@ impl<'a, Q: ExternalQueries> Pretty<'a, Q> {
         self
     }
 
-    pub fn names(mut self, names: impl IntoIterator<Item = (Name, SmolStr)>) -> Self {
-        self.names.extend(names);
-        self
-    }
-
     pub fn render(self, id: TypeId) -> SmolStr {
         let arena = Arena::new();
-        let mut printer = Printer::new(&arena, self.queries);
-        printer.names = self.names;
+        let mut printer = Printer::new(&arena, self.queries, &self.checked.names);
 
         let document = if let Some(name) = self.signature {
             printer.signature(name, id)
@@ -75,15 +69,19 @@ where
 {
     arena: &'a Arena<'a>,
     queries: &'a Q,
-    names: FxHashMap<Name, SmolStr>,
+    names: &'a FxHashMap<Name, SmolStrId>,
 }
 
 impl<'a, Q> Printer<'a, Q>
 where
     Q: ExternalQueries,
 {
-    fn new(arena: &'a Arena<'a>, queries: &'a Q) -> Printer<'a, Q> {
-        Printer { arena, queries, names: FxHashMap::default() }
+    fn new(
+        arena: &'a Arena<'a>,
+        queries: &'a Q,
+        names: &'a FxHashMap<Name, SmolStrId>,
+    ) -> Printer<'a, Q> {
+        Printer { arena, queries, names }
     }
 
     fn lookup_type(&self, id: TypeId) -> Type {
@@ -239,10 +237,12 @@ where
             }
 
             Type::Rigid(name, _, kind) => {
-                let name = self.names.entry(name).or_insert_with(|| name.as_text());
-                let name = SmolStr::clone(name);
+                let text = match self.names.get(&name) {
+                    Some(&id) => self.lookup_smol_str(id),
+                    None => name.as_text(),
+                };
                 let kind = self.traverse(Precedence::Top, kind);
-                self.arena.text(format!("({name} :: ")).append(kind).append(self.arena.text(")"))
+                self.arena.text(format!("({text} :: ")).append(kind).append(self.arena.text(")"))
             }
 
             Type::Unification(unification_id) => self.arena.text(format!("?{unification_id}")),
@@ -358,16 +358,13 @@ where
             inner = next_inner;
         }
 
-        // Register source-level names so rigid variables in the body
-        // display their original names instead of synthetic ones.
-        for binder in &binders {
-            self.names.insert(binder.name, self.lookup_smol_str(binder.text));
-        }
-
         let binders = binders
             .iter()
             .map(|binder| {
-                let text = self.lookup_smol_str(binder.text);
+                let text = match self.names.get(&binder.name) {
+                    Some(&id) => self.lookup_smol_str(id),
+                    None => binder.name.as_text(),
+                };
                 let kind = self.traverse(Precedence::Top, binder.kind);
                 self.arena
                     .text(format!("({} :: ", text))

compiler-core/checking2/src/lib.rs

@@ -19,8 +19,8 @@ use rustc_hash::FxHashMap;
 use smol_str::SmolStr;
 
 use crate::core::{
-    CheckedClass, CheckedInstance, CheckedSynonym, ForallBinder, ForallBinderId, Role, RowType,
-    RowTypeId, Synonym, SynonymId, Type, TypeId,
+    CheckedClass, CheckedInstance, CheckedSynonym, ForallBinder, ForallBinderId, Name, Role,
+    RowType, RowTypeId, SmolStrId, Synonym, SynonymId, Type, TypeId,
 };
 use crate::error::CheckError;
 
@@ -64,6 +64,7 @@ pub struct CheckedModule {
     pub roles: FxHashMap<TypeItemId, Arc<[Role]>>,
     pub nodes: CheckedNodes,
     pub errors: Vec<CheckError>,
+    pub names: FxHashMap<Name, SmolStrId>,
 }
 
 #[derive(Debug, Default, PartialEq, Eq)]
@@ -109,6 +110,10 @@ impl CheckedModule {
     pub fn lookup_roles(&self, id: TypeItemId) -> Option<Arc<[Role]>> {
         self.roles.get(&id).cloned()
     }
+
+    pub fn lookup_name(&self, name: Name) -> Option<SmolStrId> {
+        self.names.get(&name).copied()
+    }
 }
 
 impl CheckedNodes {

compiler-core/checking2/src/source/type_items.rs

@@ -339,9 +339,10 @@ where
         };
 
         let text = context.queries.intern_smol_str(text);
+        state.checked.names.insert(name, text);
         let visible = equation_binding.visible;
 
-        binders.push(ForallBinder { visible, name, text, kind });
+        binders.push(ForallBinder { visible, name, kind });
     }
 
     Ok(binders)

compiler-core/checking2/src/source/types.rs

@@ -366,8 +366,9 @@ where
     let text = if let Some(name) = &binding.name { SmolStr::clone(name) } else { name.as_text() };
     let text = context.queries.intern_smol_str(text);
 
+    state.checked.names.insert(name, text);
     state.bindings.bind_forall(binding.id, name, kind);
-    Ok(ForallBinder { visible, name, text, kind })
+    Ok(ForallBinder { visible, name, kind })
 }
 
 pub fn infer_application_kind<Q>(

compiler-core/checking2/src/state.rs

@@ -204,7 +204,7 @@ impl CheckState {
         Q: ExternalQueries,
     {
         let id = zonk::zonk(self, context, id)?;
-        let pretty = pretty::Pretty::new(context.queries).render(id);
+        let pretty = pretty::Pretty::new(context.queries, &self.checked).render(id);
         Ok(context.queries.intern_smol_str(pretty))
     }