diff --git a/compiler/noirc_evaluator/src/ssa/opt/mem2reg/alias_set.rs b/compiler/noirc_evaluator/src/ssa/opt/mem2reg/alias_set.rs index e32eaa70186..50a170a0fa2 100644 --- a/compiler/noirc_evaluator/src/ssa/opt/mem2reg/alias_set.rs +++ b/compiler/noirc_evaluator/src/ssa/opt/mem2reg/alias_set.rs @@ -3,77 +3,127 @@ use std::collections::BTreeSet; use crate::ssa::ir::value::ValueId; /// A set of possible aliases. Each ValueId in this set represents one possible value the reference -/// holding this AliasSet may be aliased to. This struct wrapper is provided so that when we take -/// the union of multiple alias sets, the result should be empty if any individual set is empty. -/// -/// Note that we distinguish between "definitely has no aliases" - `Some(BTreeSet::new())`, and -/// "unknown which aliases this may refer to" - `None`. +/// holding this AliasSet may be aliased to. This enum wrapper is provided so that when we take +/// the union of multiple alias sets, the result should be unknown if any individual set is unknown. #[derive(Debug, Default, Clone)] -pub(super) struct AliasSet { - aliases: Option>, +pub(super) enum AliasSet { + /// Unknown which aliases this may refer to. + #[default] + Unknown, + /// Definitely has no aliases. + Empty, + Known(ValueId), + KnownMultiple(BTreeSet), } impl AliasSet { pub(super) fn unknown() -> AliasSet { - Self { aliases: None } + AliasSet::Unknown } pub(super) fn known(value: ValueId) -> AliasSet { - let mut aliases = BTreeSet::new(); - aliases.insert(value); - Self { aliases: Some(aliases) } + AliasSet::Known(value) } pub(super) fn known_multiple(values: BTreeSet) -> AliasSet { - Self { aliases: Some(values) } + if values.is_empty() { + AliasSet::Empty + } else if values.len() == 1 { + AliasSet::Known(*values.first().unwrap()) + } else { + AliasSet::KnownMultiple(values) + } } /// In rare cases, such as when creating an empty array of references, the set of aliases for a /// particular value will be known to be zero, which is distinct from being unknown and /// possibly referring to any alias. pub(super) fn known_empty() -> AliasSet { - Self { aliases: Some(BTreeSet::new()) } + AliasSet::Empty } pub(super) fn is_unknown(&self) -> bool { - self.aliases.is_none() + matches!(self, AliasSet::Unknown) } /// Return the single known alias if there is exactly one. /// Otherwise, return None. pub(super) fn single_alias(&self) -> Option { - self.aliases - .as_ref() - .and_then(|aliases| (aliases.len() == 1).then(|| *aliases.first().unwrap())) + if let AliasSet::Known(alias) = self { + Some(*alias) + } else { + None + } } - /// Unify this alias set with another. The result of this set is empty if either set is empty. + /// Unify this alias set with another. The result of this set is unknown if either set is unknown. /// Otherwise, it is the union of both alias sets. pub(super) fn unify(&mut self, other: &Self) { - if let (Some(self_aliases), Some(other_aliases)) = (&mut self.aliases, &other.aliases) { - self_aliases.extend(other_aliases); - } else { - self.aliases = None; + match self { + AliasSet::Unknown => (), + AliasSet::Empty => *self = other.clone(), + AliasSet::Known(id) => match other { + AliasSet::Unknown => *self = AliasSet::Unknown, + AliasSet::Empty => (), + AliasSet::Known(other_id) => { + if id != other_id { + *self = AliasSet::KnownMultiple([*id, *other_id].iter().copied().collect()); + } + } + AliasSet::KnownMultiple(other_values) => { + let mut values = other_values.clone(); + values.insert(*id); + *self = AliasSet::KnownMultiple(values); + } + }, + AliasSet::KnownMultiple(values) => match other { + AliasSet::Unknown => *self = AliasSet::Unknown, + AliasSet::Empty => (), + AliasSet::Known(other_id) => { + values.insert(*other_id); + } + AliasSet::KnownMultiple(other_values) => { + values.extend(other_values); + } + }, } } /// Inserts a new alias into this set if it is not unknown pub(super) fn insert(&mut self, new_alias: ValueId) { - if let Some(aliases) = &mut self.aliases { - aliases.insert(new_alias); + match self { + AliasSet::Unknown => (), + AliasSet::Empty => *self = AliasSet::Known(new_alias), + AliasSet::Known(id) => { + if *id != new_alias { + *self = AliasSet::KnownMultiple([*id, new_alias].iter().copied().collect()); + } + } + AliasSet::KnownMultiple(values) => { + values.insert(new_alias); + } } } /// Returns `Some(true)` if `f` returns true for any known alias in this set. /// If this alias set is unknown, None is returned. - pub(super) fn any(&self, f: impl FnMut(ValueId) -> bool) -> Option { - self.aliases.as_ref().map(|aliases| aliases.iter().copied().any(f)) + pub(super) fn any(&self, mut f: impl FnMut(ValueId) -> bool) -> Option { + match self { + AliasSet::Unknown => None, + AliasSet::Empty => Some(false), + AliasSet::Known(id) => Some(f(*id)), + AliasSet::KnownMultiple(values) => Some(values.iter().copied().any(f)), + } } pub(super) fn for_each(&self, mut f: impl FnMut(ValueId)) { - if let Some(aliases) = &self.aliases { - for alias in aliases { - f(*alias); + match self { + AliasSet::Unknown | AliasSet::Empty => (), + AliasSet::Known(id) => f(*id), + AliasSet::KnownMultiple(values) => { + for value in values { + f(*value); + } } } } @@ -82,6 +132,10 @@ impl AliasSet { /// The ordering is arbitrary (by lowest ValueId) so this method should only be /// used when you need an arbitrary ValueId from the alias set. pub(super) fn first(&self) -> Option { - self.aliases.as_ref().and_then(|aliases| aliases.first().copied()) + match self { + AliasSet::Unknown | AliasSet::Empty => None, + AliasSet::Known(id) => Some(*id), + AliasSet::KnownMultiple(values) => values.first().copied(), + } } }