From 09d8ed3527a54c21313597abb9b93c89695bc7a1 Mon Sep 17 00:00:00 2001 From: Todd Nowacki Date: Tue, 26 Nov 2024 10:22:11 -0800 Subject: [PATCH] [move-compiler] Add rayon parallel iterator API to UniqueMap (#20353) ## Description - Added ParallelIterator APIs to UniqueMap ## Test plan - Used in a later API --- ## Release notes Check each box that your changes affect. If none of the boxes relate to your changes, release notes aren't required. For each box you select, include information after the relevant heading that describes the impact of your changes that a user might notice and any actions they must take to implement updates. - [ ] Protocol: - [ ] Nodes (Validators and Full nodes): - [ ] Indexer: - [ ] JSON-RPC: - [ ] GraphQL: - [ ] CLI: - [ ] Rust SDK: - [ ] REST API: --- .../move-compiler/src/shared/unique_map.rs | 143 ++++++++++++++++++ 1 file changed, 143 insertions(+) diff --git a/external-crates/move/crates/move-compiler/src/shared/unique_map.rs b/external-crates/move/crates/move-compiler/src/shared/unique_map.rs index 3d5a5ca1928e2..fe16b17b058b4 100644 --- a/external-crates/move/crates/move-compiler/src/shared/unique_map.rs +++ b/external-crates/move/crates/move-compiler/src/shared/unique_map.rs @@ -3,6 +3,7 @@ // SPDX-License-Identifier: Apache-2.0 use super::*; +use rayon::prelude::*; use std::{collections::BTreeMap, fmt::Debug, iter::IntoIterator}; //************************************************************************************************** @@ -224,6 +225,78 @@ impl UniqueMap { } } +impl UniqueMap +where + K: Sync + Send, + K::Key: Sync, + K::Loc: Send + Sync, + V: Sync, +{ + pub fn key_cloned_par_iter(&self) -> impl ParallelIterator { + self.par_iter() + .map(|(loc, k_, v)| (K::add_loc(loc, k_.clone()), v)) + } +} + +impl UniqueMap +where + K: Sync + Send, + K::Key: Sync, + K::Loc: Send + Sync, + V: Sync + Send, +{ + pub fn key_cloned_par_iter_mut(&mut self) -> impl ParallelIterator { + IntoParallelRefMutIterator::par_iter_mut(self) + .map(|(loc, k_, v)| (K::add_loc(loc, k_.clone()), v)) + } +} + +impl UniqueMap +where + K: Send, + K::Key: Send, + K::Loc: Send, + V: Send, +{ + pub fn par_map(self, f: F) -> UniqueMap + where + V2: Send, + F: Fn(K, V) -> V2 + Sync + Send, + { + UniqueMap( + self.0 + .into_par_iter() + .map(|(k_, (loc, v))| { + let v2 = f(K::add_loc(loc, k_.clone()), v); + (k_, (loc, v2)) + }) + .collect(), + ) + } + + pub fn maybe_from_par_iter( + iter: impl ParallelIterator, + ) -> Result, (K::Key, K::Loc, K::Loc)> { + iter.try_fold(Self::new, |mut m, (k, v)| { + if let Err((k, old_loc)) = m.add(k, v) { + let (loc, key_) = k.drop_loc(); + Err((key_, loc, old_loc)) + } else { + Ok(m) + } + }) + .try_reduce(Self::new, |mut m1, m2| { + for (k, v) in m2 { + if let Err((k, old_loc)) = m1.add(k, v) { + let (loc, key_) = k.drop_loc(); + return Err((key_, loc, old_loc)); + } + } + Ok(m1) + }) + } +} + impl PartialEq for UniqueMap { fn eq(&self, other: &UniqueMap) -> bool { self.iter() @@ -325,6 +398,29 @@ impl IntoIterator for UniqueMap { } } +impl IntoParallelIterator for UniqueMap +where + K: Send, + K::Key: Send, + K::Loc: Send, + V: Send, +{ + type Iter = rayon::iter::Map< + rayon::collections::btree_map::IntoIter, + fn((K::Key, (K::Loc, V))) -> (K, V), + >; + type Item = (K, V); + + fn into_par_iter(self) -> Self::Iter { + self.0.into_par_iter().map(|(k_, loc_v)| { + let loc = loc_v.0; + let v = loc_v.1; + let k = K::add_loc(loc, k_); + (k, v) + }) + } +} + //************************************************************************************************** // Iter //************************************************************************************************** @@ -366,6 +462,29 @@ impl<'a, K: TName, V> IntoIterator for &'a UniqueMap { } } +impl<'a, K: TName, V> IntoParallelRefIterator<'a> for UniqueMap +where + K: Sync, + K::Key: Sync + 'a, + K::Loc: Send + Sync + 'a, + V: Sync + 'a, +{ + type Iter = rayon::iter::Map< + rayon::collections::btree_map::Iter<'a, K::Key, (K::Loc, V)>, + fn((&'a K::Key, &'a (K::Loc, V))) -> (K::Loc, &'a K::Key, &'a V), + >; + type Item = (K::Loc, &'a K::Key, &'a V); + + fn par_iter(&'a self) -> Self::Iter { + let fix = |(k_, loc_v): (&'a K::Key, &'a (K::Loc, V))| -> (K::Loc, &'a K::Key, &'a V) { + let loc = loc_v.0; + let v = &loc_v.1; + (loc, k_, v) + }; + self.0.par_iter().map(fix) + } +} + //************************************************************************************************** // IterMut //************************************************************************************************** @@ -404,3 +523,27 @@ impl<'a, K: TName, V> IntoIterator for &'a mut UniqueMap { IterMut(self.0.iter_mut().map(fix), len) } } + +impl<'a, K: TName, V> IntoParallelRefMutIterator<'a> for UniqueMap +where + K: Sync, + K::Key: Sync + 'a, + K::Loc: Send + Sync + 'a, + V: Sync + Send + 'a, +{ + type Iter = rayon::iter::Map< + rayon::collections::btree_map::IterMut<'a, K::Key, (K::Loc, V)>, + fn((&'a K::Key, &'a mut (K::Loc, V))) -> (K::Loc, &'a K::Key, &'a mut V), + >; + type Item = (K::Loc, &'a K::Key, &'a mut V); + + fn par_iter_mut(&'a mut self) -> Self::Iter { + let fix = + |(k_, loc_v): (&'a K::Key, &'a mut (K::Loc, V))| -> (K::Loc, &'a K::Key, &'a mut V) { + let loc = loc_v.0; + let v = &mut loc_v.1; + (loc, k_, v) + }; + self.0.par_iter_mut().map(fix) + } +}