Merge pull request #3197 from rtfeldman/subs-cleanup

subs cleanup

compiler/types/src/subs.rs

@@ -1747,8 +1747,8 @@ impl Subs {
 
     /// Unions two keys without the possibility of failure.
     pub fn union(&mut self, left: Variable, right: Variable, desc: Descriptor) {
-        let l_root = self.utable.inlined_get_root_key(left);
-        let r_root = self.utable.inlined_get_root_key(right);
+        let l_root = self.utable.root_key(left);
+        let r_root = self.utable.root_key(right);
 
         // NOTE this swapping is intentional! most of our unifying commands are based on the elm
         // source, but unify_roots is from `ena`, not the elm source. Turns out that they have
@@ -1803,7 +1803,7 @@ impl Subs {
 
     #[inline(always)]
     pub fn get_root_key(&mut self, key: Variable) -> Variable {
-        self.utable.inlined_get_root_key(key)
+        self.utable.root_key(key)
     }
 
     #[inline(always)]
@@ -1813,7 +1813,7 @@ impl Subs {
 
     #[inline(always)]
     pub fn set(&mut self, key: Variable, r_value: Descriptor) {
-        let l_key = self.utable.inlined_get_root_key(key);
+        let l_key = self.utable.root_key(key);
 
         // self.utable.update_value(l_key, |node| node.value = r_value);
         self.utable.set_descriptor(l_key, r_value)

compiler/types/src/unification_table.rs

@@ -1,25 +1,46 @@
+use std::hint::unreachable_unchecked;
+
 use crate::subs::{Content, Descriptor, Mark, OptVariable, Rank, Variable, VariableSubsSlice};
 
 #[derive(Clone, Default)]
 pub struct UnificationTable {
     contents: Vec<Content>,
-    ranks: Vec<Rank>,
-    marks: Vec<Mark>,
-    copies: Vec<OptVariable>,
-    redirects: Vec<OptVariable>,
+    metadata: Vec<Combine>,
 }
 
 pub(crate) struct Snapshot(UnificationTable);
 
+#[derive(Debug, Clone, Copy)]
+enum Combine {
+    Redirect(Variable),
+    Root(Root),
+}
+
+#[derive(Debug, Clone, Copy)]
+pub struct Root {
+    rank: Rank,
+    mark: Mark,
+    copy: OptVariable,
+}
+
+impl Root {
+    const NONE: Self = Self {
+        rank: Rank::NONE,
+        mark: Mark::NONE,
+        copy: OptVariable::NONE,
+    };
+}
+
+impl Combine {
+    const NONE: Self = Self::Root(Root::NONE);
+}
+
 impl UnificationTable {
     #[allow(unused)]
     pub fn with_capacity(cap: usize) -> Self {
         Self {
-            contents: Vec::with_capacity(cap),  // vec![Content::Error; cap],
-            ranks: Vec::with_capacity(cap),     // vec![Rank::NONE; cap],
-            marks: Vec::with_capacity(cap),     // vec![Mark::NONE; cap],
-            copies: Vec::with_capacity(cap),    // vec![OptVariable::NONE; cap],
-            redirects: Vec::with_capacity(cap), // vec![OptVariable::NONE; cap],
+            contents: Vec::with_capacity(cap),
+            metadata: Vec::with_capacity(cap),
         }
     }
 
@@ -38,12 +59,8 @@ impl UnificationTable {
 
         self.contents
             .extend(repeat(Content::FlexVar(None)).take(extra_length));
-        self.ranks.extend(repeat(Rank::NONE).take(extra_length));
-        self.marks.extend(repeat(Mark::NONE).take(extra_length));
-        self.copies
-            .extend(repeat(OptVariable::NONE).take(extra_length));
-        self.redirects
-            .extend(repeat(OptVariable::NONE).take(extra_length));
+        self.metadata
+            .extend(repeat(Combine::NONE).take(extra_length));
 
         VariableSubsSlice::new(start as _, extra_length as _)
     }
@@ -58,10 +75,10 @@ impl UnificationTable {
         let variable = unsafe { Variable::from_index(self.len() as _) };
 
         self.contents.push(content);
-        self.ranks.push(rank);
-        self.marks.push(mark);
-        self.copies.push(copy);
-        self.redirects.push(OptVariable::NONE);
+
+        let combine = Combine::Root(Root { rank, mark, copy });
+
+        self.metadata.push(combine);
 
         variable
     }
@@ -75,53 +92,78 @@ impl UnificationTable {
         mark: Mark,
         copy: OptVariable,
     ) {
-        let index = self.root_key(key).index() as usize;
+        let root = self.root_key(key);
+        self.set_unchecked(root, content, rank, mark, copy)
+    }
+
+    pub fn set_unchecked(
+        &mut self,
+        key: Variable,
+        content: Content,
+        rank: Rank,
+        mark: Mark,
+        copy: OptVariable,
+    ) {
+        let index = key.index() as usize;
 
         self.contents[index] = content;
-        self.ranks[index] = rank;
-        self.marks[index] = mark;
-        self.copies[index] = copy;
+
+        self.metadata[index] = Combine::Root(Root { rank, mark, copy });
     }
 
     pub fn modify<F, T>(&mut self, key: Variable, mapper: F) -> T
     where
         F: FnOnce(&mut Descriptor) -> T,
     {
-        let index = self.root_key(key).index() as usize;
+        let root_key = self.root_key(key);
+        let index = root_key.index() as usize;
+
+        let root = self.get_root_unchecked(root_key);
 
         let mut desc = Descriptor {
             content: self.contents[index],
-            rank: self.ranks[index],
-            mark: self.marks[index],
-            copy: self.copies[index],
+            rank: root.rank,
+            mark: root.mark,
+            copy: root.copy,
         };
 
         let result = mapper(&mut desc);
 
-        self.contents[index] = desc.content;
-        self.ranks[index] = desc.rank;
-        self.marks[index] = desc.mark;
-        self.copies[index] = desc.copy;
+        self.set_unchecked(root_key, desc.content, desc.rank, desc.mark, desc.copy);
 
         result
     }
 
     // GET UNCHECKED
 
+    #[inline(always)]
+    pub fn get_root_unchecked(&self, key: Variable) -> Root {
+        match self.metadata[key.index() as usize] {
+            Combine::Root(root) => root,
+            Combine::Redirect(_) => {
+                if cfg!(debug_assertions) {
+                    unreachable!("unchecked access on non-root variable {:?}", key);
+                } else {
+                    unsafe { unreachable_unchecked() }
+                }
+            }
+        }
+    }
+
     #[inline(always)]
     pub fn get_rank_unchecked(&self, key: Variable) -> Rank {
-        self.ranks[key.index() as usize]
+        self.get_root_unchecked(key).rank
     }
 
     #[inline(always)]
     pub fn get_mark_unchecked(&self, key: Variable) -> Mark {
-        self.marks[key.index() as usize]
+        self.get_root_unchecked(key).mark
     }
 
     #[allow(unused)]
     #[inline(always)]
     pub fn get_copy_unchecked(&self, key: Variable) -> OptVariable {
-        self.copies[key.index() as usize]
+        self.get_root_unchecked(key).copy
     }
 
     #[inline(always)]
@@ -131,20 +173,34 @@ impl UnificationTable {
 
     // GET CHECKED
 
+    #[inline(always)]
+    pub fn get_root(&self, key: Variable) -> Root {
+        let root_key = self.root_key_without_compacting(key);
+        match self.metadata[root_key.index() as usize] {
+            Combine::Root(root) => root,
+            Combine::Redirect(_) => {
+                if cfg!(debug_assertions) {
+                    unreachable!("the root key {:?} is not actually a root key", root_key);
+                } else {
+                    unsafe { unreachable_unchecked() }
+                }
+            }
+        }
+    }
+
     #[inline(always)]
     pub fn get_rank(&self, key: Variable) -> Rank {
-        self.ranks[self.root_key_without_compacting(key).index() as usize]
+        self.get_root(key).rank
     }
 
     #[inline(always)]
     pub fn get_mark(&self, key: Variable) -> Mark {
-        self.marks[self.root_key_without_compacting(key).index() as usize]
+        self.get_root(key).mark
     }
 
     #[inline(always)]
     pub fn get_copy(&self, key: Variable) -> OptVariable {
-        let index = self.root_key_without_compacting(key).index() as usize;
-        self.copies[index]
+        self.get_root(key).copy
     }
 
     #[inline(always)]
@@ -154,20 +210,37 @@ impl UnificationTable {
 
     // SET UNCHECKED
 
+    #[inline(always)]
+    pub fn modify_root_unchecked<F, T>(&mut self, key: Variable, f: F) -> T
+    where
+        F: Fn(&mut Root) -> T,
+    {
+        match &mut self.metadata[key.index() as usize] {
+            Combine::Root(root) => f(root),
+            Combine::Redirect(_) => {
+                if cfg!(debug_assertions) {
+                    unreachable!("unchecked access on non-root variable {:?}", key);
+                } else {
+                    unsafe { unreachable_unchecked() }
+                }
+            }
+        }
+    }
+
     #[inline(always)]
     pub fn set_rank_unchecked(&mut self, key: Variable, value: Rank) {
-        self.ranks[key.index() as usize] = value;
+        self.modify_root_unchecked(key, |root| root.rank = value)
     }
 
     #[inline(always)]
     pub fn set_mark_unchecked(&mut self, key: Variable, value: Mark) {
-        self.marks[key.index() as usize] = value;
+        self.modify_root_unchecked(key, |root| root.mark = value)
     }
 
     #[allow(unused)]
     #[inline(always)]
     pub fn set_copy_unchecked(&mut self, key: Variable, value: OptVariable) {
-        self.copies[key.index() as usize] = value;
+        self.modify_root_unchecked(key, |root| root.copy = value)
     }
 
     #[allow(unused)]
@@ -180,20 +253,20 @@ impl UnificationTable {
 
     #[inline(always)]
     pub fn set_rank(&mut self, key: Variable, value: Rank) {
-        let index = self.root_key(key).index() as usize;
-        self.ranks[index] = value;
+        let root_key = self.root_key(key);
+        self.modify_root_unchecked(root_key, |root| root.rank = value)
     }
 
     #[inline(always)]
     pub fn set_mark(&mut self, key: Variable, value: Mark) {
-        let index = self.root_key(key).index() as usize;
-        self.marks[index] = value;
+        let root_key = self.root_key(key);
+        self.modify_root_unchecked(root_key, |root| root.mark = value)
     }
 
     #[inline(always)]
     pub fn set_copy(&mut self, key: Variable, value: OptVariable) {
-        let index = self.root_key(key).index() as usize;
-        self.copies[index] = value;
+        let root_key = self.root_key(key);
+        self.modify_root_unchecked(root_key, |root| root.copy = value)
     }
 
     #[inline(always)]
@@ -206,22 +279,19 @@ impl UnificationTable {
 
     #[inline(always)]
     pub fn root_key(&mut self, mut key: Variable) -> Variable {
-        let index = key.index() as usize;
+        let root = self.root_key_without_compacting(key);
 
-        while let Some(redirect) = self.redirects[key.index() as usize].into_variable() {
-            key = redirect;
+        while let Combine::Redirect(redirect) = &mut self.metadata[key.index() as usize] {
+            key = *redirect;
+            *redirect = root;
         }
 
-        if index != key.index() as usize {
-            self.redirects[index] = OptVariable::from(key);
-        }
-
-        key
+        root
     }
 
     #[inline(always)]
     pub fn root_key_without_compacting(&self, mut key: Variable) -> Variable {
-        while let Some(redirect) = self.redirects[key.index() as usize].into_variable() {
+        while let Combine::Redirect(redirect) = self.metadata[key.index() as usize] {
             key = redirect;
         }
 
@@ -246,7 +316,7 @@ impl UnificationTable {
     }
 
     pub fn is_redirect(&self, key: Variable) -> bool {
-        self.redirects[key.index() as usize].is_some()
+        matches!(self.metadata[key.index() as usize], Combine::Redirect(_))
     }
 
     pub fn unioned(&mut self, a: Variable, b: Variable) -> bool {
@@ -256,17 +326,13 @@ impl UnificationTable {
     // custom very specific helpers
     #[inline(always)]
     pub fn get_rank_set_mark(&mut self, key: Variable, mark: Mark) -> Rank {
-        let index = self.root_key(key).index() as usize;
+        let root_key = self.root_key(key);
 
-        self.marks[index] = mark;
+        self.modify_root_unchecked(root_key, |root| {
+            root.mark = mark;
 
-        self.ranks[index]
-    }
-
-    // TODO remove
-    #[inline(always)]
-    pub fn inlined_get_root_key(&mut self, key: Variable) -> Variable {
-        self.root_key(key)
+            root.rank
+        })
     }
 
     /// NOTE: assumes variables are root
@@ -276,34 +342,29 @@ impl UnificationTable {
 
         // redirect from -> to
         if from_index != to_index {
-            self.redirects[from_index] = OptVariable::from(to);
+            self.metadata[from_index] = Combine::Redirect(to)
         }
 
         // update to's Descriptor
-        self.contents[to_index] = desc.content;
-        self.ranks[to_index] = desc.rank;
-        self.marks[to_index] = desc.mark;
-        self.copies[to_index] = desc.copy;
+        self.set_unchecked(to, desc.content, desc.rank, desc.mark, desc.copy);
     }
 
     pub fn get_descriptor(&self, key: Variable) -> Descriptor {
-        let index = self.root_key_without_compacting(key).index() as usize;
+        let root_key = self.root_key_without_compacting(key);
+        let index = root_key.index() as usize;
+
+        let root = self.get_root_unchecked(root_key);
 
         Descriptor {
             content: self.contents[index],
-            rank: self.ranks[index],
-            mark: self.marks[index],
-            copy: self.copies[index],
+            rank: root.rank,
+            mark: root.mark,
+            copy: root.copy,
         }
     }
 
     pub fn set_descriptor(&mut self, key: Variable, desc: Descriptor) {
-        let index = self.root_key(key).index() as usize;
-
-        self.contents[index] = desc.content;
-        self.ranks[index] = desc.rank;
-        self.marks[index] = desc.mark;
-        self.copies[index] = desc.copy;
+        self.set(key, desc.content, desc.rank, desc.mark, desc.copy);
     }
 
     pub(crate) fn serialize(
@@ -314,10 +375,37 @@ impl UnificationTable {
         use crate::subs::Subs;
 
         written = Subs::serialize_slice(&self.contents, writer, written)?;
-        written = Subs::serialize_slice(&self.ranks, writer, written)?;
-        written = Subs::serialize_slice(&self.marks, writer, written)?;
-        written = Subs::serialize_slice(&self.copies, writer, written)?;
-        written = Subs::serialize_slice(&self.redirects, writer, written)?;
+
+        let mut ranks = Vec::new();
+        let mut marks = Vec::new();
+        let mut copies = Vec::new();
+        let mut redirects = Vec::new();
+
+        for c in self.metadata.iter() {
+            match c {
+                Combine::Redirect(redirect) => {
+                    let root = Root::NONE;
+
+                    ranks.push(root.rank);
+                    marks.push(root.mark);
+                    copies.push(root.copy);
+
+                    redirects.push(OptVariable::from(*redirect));
+                }
+                Combine::Root(root) => {
+                    ranks.push(root.rank);
+                    marks.push(root.mark);
+                    copies.push(root.copy);
+
+                    redirects.push(OptVariable::NONE);
+                }
+            }
+        }
+
+        written = Subs::serialize_slice(&ranks, writer, written)?;
+        written = Subs::serialize_slice(&marks, writer, written)?;
+        written = Subs::serialize_slice(&copies, writer, written)?;
+        written = Subs::serialize_slice(&redirects, writer, written)?;
 
         Ok(written)
     }
@@ -331,12 +419,29 @@ impl UnificationTable {
         let (copies, offset) = Subs::deserialize_slice::<OptVariable>(bytes, length, offset);
         let (redirects, offset) = Subs::deserialize_slice::<OptVariable>(bytes, length, offset);
 
+        let mut metadata = Vec::with_capacity(ranks.len());
+
+        let it = ranks
+            .iter()
+            .zip(marks.iter())
+            .zip(copies.iter())
+            .zip(redirects.iter());
+        for (((rank, mark), copy), redirect) in it {
+            match redirect.into_variable() {
+                Some(redirect) => metadata.push(Combine::Redirect(redirect)),
+                None => {
+                    metadata.push(Combine::Root(Root {
+                        rank: *rank,
+                        mark: *mark,
+                        copy: *copy,
+                    }));
+                }
+            }
+        }
+
         let this = Self {
             contents: contents.to_vec(),
-            ranks: ranks.to_vec(),
-            marks: marks.to_vec(),
-            copies: copies.to_vec(),
-            redirects: redirects.to_vec(),
+            metadata,
         };
 
         (this, offset)