diff --git a/compiler/types/src/subs.rs b/compiler/types/src/subs.rs index efd11c23aa..925d1cd667 100644 --- a/compiler/types/src/subs.rs +++ b/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) diff --git a/compiler/types/src/unification_table.rs b/compiler/types/src/unification_table.rs index 315e92d24b..12e10f4648 100644 --- a/compiler/types/src/unification_table.rs +++ b/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, - ranks: Vec, - marks: Vec, - copies: Vec, - redirects: Vec, + metadata: Vec, } 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(&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(&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::(bytes, length, offset); let (redirects, offset) = Subs::deserialize_slice::(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)