chore(perf): try using vec-collections's VecSet in AliasSet (#7058)

noir-lang · Jan 14, 2025 · b6097a0 · b6097a0
1 parent 31dbfe3
commit b6097a0
Show file tree

Hide file tree

Showing 4 changed files with 50 additions and 12 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/compiler/noirc_evaluator/Cargo.toml b/compiler/noirc_evaluator/Cargo.toml
@@ -29,6 +29,7 @@ chrono = "0.4.37"
 rayon.workspace = true
 cfg-if.workspace = true
 smallvec = { version = "1.13.2", features = ["serde"] }
+vec-collections = "0.4.3"
 
 [dev-dependencies]
 proptest.workspace = true

diff --git a/compiler/noirc_evaluator/src/ssa/opt/mem2reg.rs b/compiler/noirc_evaluator/src/ssa/opt/mem2reg.rs
@@ -79,6 +79,7 @@ mod block;
 use std::collections::{BTreeMap, BTreeSet};
 
 use fxhash::{FxHashMap as HashMap, FxHashSet as HashSet};
+use vec_collections::VecSet;
 
 use crate::ssa::{
     ir::{
@@ -619,7 +620,7 @@ impl<'f> PerFunctionContext<'f> {
                 // then those parameters also alias each other.
                 // We save parameters with repeat arguments to later mark those
                 // parameters as aliasing one another.
-                let mut arg_set: HashMap<ValueId, BTreeSet<ValueId>> = HashMap::default();
+                let mut arg_set = HashMap::default();
 
                 // Add an alias for each reference parameter
                 for (parameter, argument) in destination_parameters.iter().zip(arguments) {
@@ -632,7 +633,8 @@ impl<'f> PerFunctionContext<'f> {
                                 aliases.insert(*parameter);
 
                                 // Check if we have seen the same argument
-                                let seen_parameters = arg_set.entry(argument).or_default();
+                                let seen_parameters =
+                                    arg_set.entry(argument).or_insert_with(VecSet::empty);
                                 // Add the current parameter to the parameters we have seen for this argument.
                                 // The previous parameters and the current one alias one another.
                                 seen_parameters.insert(*parameter);

diff --git a/compiler/noirc_evaluator/src/ssa/opt/mem2reg/alias_set.rs b/compiler/noirc_evaluator/src/ssa/opt/mem2reg/alias_set.rs
@@ -1,4 +1,4 @@
-use std::collections::BTreeSet;
+use vec_collections::VecSet;
 
 use crate::ssa::ir::value::ValueId;
 
@@ -10,7 +10,7 @@ use crate::ssa::ir::value::ValueId;
 /// "unknown which aliases this may refer to" - `None`.
 #[derive(Debug, Default, Clone)]
 pub(super) struct AliasSet {
-    aliases: Option<BTreeSet<ValueId>>,
+    aliases: Option<VecSet<[ValueId; 1]>>,
 }
 
 impl AliasSet {
@@ -19,20 +19,18 @@ impl AliasSet {
     }
 
     pub(super) fn known(value: ValueId) -> AliasSet {
-        let mut aliases = BTreeSet::new();
-        aliases.insert(value);
-        Self { aliases: Some(aliases) }
+        Self { aliases: Some(VecSet::single(value)) }
     }
 
-    pub(super) fn known_multiple(values: BTreeSet<ValueId>) -> AliasSet {
+    pub(super) fn known_multiple(values: VecSet<[ValueId; 1]>) -> AliasSet {
         Self { aliases: Some(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()) }
+        Self { aliases: Some(VecSet::empty()) }
     }
 
     pub(super) fn is_unknown(&self) -> bool {
@@ -44,14 +42,14 @@ impl AliasSet {
     pub(super) fn single_alias(&self) -> Option<ValueId> {
         self.aliases
             .as_ref()
-            .and_then(|aliases| (aliases.len() == 1).then(|| *aliases.first().unwrap()))
+            .and_then(|aliases| (aliases.len() == 1).then(|| *aliases.iter().next().unwrap()))
     }
 
     /// Unify this alias set with another. The result of this set is empty if either set is empty.
     /// 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);
+            self_aliases.extend(other_aliases.iter().cloned());
         } else {
             self.aliases = None;
         }
@@ -82,6 +80,6 @@ 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<ValueId> {
-        self.aliases.as_ref().and_then(|aliases| aliases.first().copied())
+        self.aliases.as_ref().and_then(|aliases| aliases.iter().next().copied())
     }
 }