flash_map/
lib.rs

1//! FlashMap — GPU-native concurrent hash map.
2//!
3//! Bulk-only API designed for maximum GPU throughput:
4//! - `bulk_get`: parallel key lookup
5//! - `bulk_insert`: parallel key-value insertion (updates existing keys)
6//! - `bulk_remove`: parallel key removal (tombstone-based)
7//!
8//! SoA (Struct of Arrays) memory layout on GPU for coalesced access.
9//! Robin Hood hashing with probe distance early exit.
10//!
11//! # Features
12//!
13//! - `cuda` — GPU backend via CUDA (requires NVIDIA GPU + CUDA toolkit)
14//! - `rayon` — multi-threaded CPU backend (default)
15//!
16//! # Example
17//!
18//! ```rust,no_run
19//! use flash_map::{FlashMap, HashStrategy};
20//!
21//! let mut map: FlashMap<[u8; 32], [u8; 128]> =
22//!     FlashMap::with_capacity(1_000_000).unwrap();
23//!
24//! // Insert 1M key-value pairs in one GPU kernel launch
25//! let pairs: Vec<([u8; 32], [u8; 128])> = generate_pairs();
26//! map.bulk_insert(&pairs).unwrap();
27//!
28//! // Lookup
29//! let keys: Vec<[u8; 32]> = pairs.iter().map(|(k, _)| *k).collect();
30//! let results: Vec<Option<[u8; 128]>> = map.bulk_get(&keys).unwrap();
31//! # fn generate_pairs() -> Vec<([u8; 32], [u8; 128])> { vec![] }
32//! ```
33
34#[cfg(not(any(feature = "cuda", feature = "rayon")))]
35compile_error!(
36    "flash-map: enable at least one of 'cuda' or 'rayon' features"
37);
38
39mod error;
40mod hash;
41
42#[cfg(feature = "cuda")]
43mod gpu;
44
45#[cfg(feature = "rayon")]
46mod rayon_cpu;
47
48#[cfg(feature = "tokio")]
49mod async_map;
50
51pub use bytemuck::Pod;
52pub use error::FlashMapError;
53pub use hash::HashStrategy;
54
55#[cfg(feature = "tokio")]
56pub use async_map::AsyncFlashMap;
57
58use bytemuck::Pod as PodBound;
59
60// ---------------------------------------------------------------------------
61// FlashMap — public API
62// ---------------------------------------------------------------------------
63
64/// GPU-native concurrent hash map with bulk-only operations.
65///
66/// Generic over fixed-size key `K` and value `V` types that implement
67/// [`bytemuck::Pod`] (plain old data — `Copy`, fixed layout, any bit
68/// pattern valid).
69///
70/// Common type combinations:
71/// - `FlashMap<[u8; 32], [u8; 128]>` — blockchain state (pubkey → account)
72/// - `FlashMap<u64, u64>` — numeric keys and values
73/// - `FlashMap<[u8; 32], [u8; 32]>` — hash → hash mappings
74pub struct FlashMap<K: PodBound, V: PodBound> {
75    inner: FlashMapBackend<K, V>,
76}
77
78enum FlashMapBackend<K: PodBound, V: PodBound> {
79    #[cfg(feature = "cuda")]
80    Gpu(gpu::GpuFlashMap<K, V>),
81    #[cfg(feature = "rayon")]
82    Rayon(rayon_cpu::RayonFlashMap<K, V>),
83}
84
85impl<K: PodBound + Send + Sync, V: PodBound + Send + Sync> FlashMap<K, V> {
86    /// Create a FlashMap with the given capacity using default settings.
87    ///
88    /// Tries GPU first (if `cuda` feature enabled), falls back to Rayon.
89    /// Capacity is rounded up to the next power of 2.
90    pub fn with_capacity(capacity: usize) -> Result<Self, FlashMapError> {
91        FlashMapBuilder::new(capacity).build()
92    }
93
94    /// Create a builder for fine-grained configuration.
95    pub fn builder(capacity: usize) -> FlashMapBuilder {
96        FlashMapBuilder::new(capacity)
97    }
98
99    /// Look up multiple keys in parallel. Returns `None` for missing keys.
100    pub fn bulk_get(&self, keys: &[K]) -> Result<Vec<Option<V>>, FlashMapError> {
101        match &self.inner {
102            #[cfg(feature = "cuda")]
103            FlashMapBackend::Gpu(m) => m.bulk_get(keys),
104            #[cfg(feature = "rayon")]
105            FlashMapBackend::Rayon(m) => m.bulk_get(keys),
106        }
107    }
108
109    /// Insert multiple key-value pairs in parallel.
110    ///
111    /// Returns the number of **new** insertions (updates don't count).
112    /// If a key already exists, its value is updated in place.
113    ///
114    /// # Invariant
115    ///
116    /// No duplicate keys within a single batch. If the same key appears
117    /// multiple times, behavior is undefined (one will win, but which
118    /// one is non-deterministic on GPU).
119    pub fn bulk_insert(&mut self, pairs: &[(K, V)]) -> Result<usize, FlashMapError> {
120        match &mut self.inner {
121            #[cfg(feature = "cuda")]
122            FlashMapBackend::Gpu(m) => m.bulk_insert(pairs),
123            #[cfg(feature = "rayon")]
124            FlashMapBackend::Rayon(m) => m.bulk_insert(pairs),
125        }
126    }
127
128    /// Remove multiple keys in parallel (tombstone-based).
129    ///
130    /// Returns the number of keys actually removed.
131    pub fn bulk_remove(&mut self, keys: &[K]) -> Result<usize, FlashMapError> {
132        match &mut self.inner {
133            #[cfg(feature = "cuda")]
134            FlashMapBackend::Gpu(m) => m.bulk_remove(keys),
135            #[cfg(feature = "rayon")]
136            FlashMapBackend::Rayon(m) => m.bulk_remove(keys),
137        }
138    }
139
140    /// Number of occupied entries.
141    pub fn len(&self) -> usize {
142        match &self.inner {
143            #[cfg(feature = "cuda")]
144            FlashMapBackend::Gpu(m) => m.len(),
145            #[cfg(feature = "rayon")]
146            FlashMapBackend::Rayon(m) => m.len(),
147        }
148    }
149
150    /// Whether the map is empty.
151    pub fn is_empty(&self) -> bool {
152        self.len() == 0
153    }
154
155    /// Total slot capacity (always a power of 2).
156    pub fn capacity(&self) -> usize {
157        match &self.inner {
158            #[cfg(feature = "cuda")]
159            FlashMapBackend::Gpu(m) => m.capacity(),
160            #[cfg(feature = "rayon")]
161            FlashMapBackend::Rayon(m) => m.capacity(),
162        }
163    }
164
165    /// Current load factor (0.0 to 1.0).
166    pub fn load_factor(&self) -> f64 {
167        match &self.inner {
168            #[cfg(feature = "cuda")]
169            FlashMapBackend::Gpu(m) => m.load_factor(),
170            #[cfg(feature = "rayon")]
171            FlashMapBackend::Rayon(m) => m.load_factor(),
172        }
173    }
174
175    /// Remove all entries, resetting to empty.
176    pub fn clear(&mut self) -> Result<(), FlashMapError> {
177        match &mut self.inner {
178            #[cfg(feature = "cuda")]
179            FlashMapBackend::Gpu(m) => m.clear(),
180            #[cfg(feature = "rayon")]
181            FlashMapBackend::Rayon(m) => m.clear(),
182        }
183    }
184}
185
186impl<K: PodBound + Send + Sync, V: PodBound + Send + Sync> std::fmt::Debug
187    for FlashMap<K, V>
188{
189    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
190        let backend = match &self.inner {
191            #[cfg(feature = "cuda")]
192            FlashMapBackend::Gpu(_) => "GPU",
193            #[cfg(feature = "rayon")]
194            FlashMapBackend::Rayon(_) => "Rayon",
195        };
196        f.debug_struct("FlashMap")
197            .field("backend", &backend)
198            .field("len", &self.len())
199            .field("capacity", &self.capacity())
200            .field("load_factor", &format!("{:.1}%", self.load_factor() * 100.0))
201            .finish()
202    }
203}
204
205// ---------------------------------------------------------------------------
206// Builder
207// ---------------------------------------------------------------------------
208
209/// Builder for configuring a [`FlashMap`].
210pub struct FlashMapBuilder {
211    capacity: usize,
212    hash_strategy: HashStrategy,
213    device_id: usize,
214    force_rayon: bool,
215}
216
217impl FlashMapBuilder {
218    /// Create a builder targeting the given capacity.
219    pub fn new(capacity: usize) -> Self {
220        Self {
221            capacity,
222            hash_strategy: HashStrategy::Identity,
223            device_id: 0,
224            force_rayon: false,
225        }
226    }
227
228    /// Set the hash strategy (default: Identity).
229    pub fn hash_strategy(mut self, strategy: HashStrategy) -> Self {
230        self.hash_strategy = strategy;
231        self
232    }
233
234    /// Set the CUDA device ordinal (default: 0).
235    pub fn device_id(mut self, id: usize) -> Self {
236        self.device_id = id;
237        self
238    }
239
240    /// Force Rayon backend even if CUDA is available.
241    pub fn force_cpu(mut self) -> Self {
242        self.force_rayon = true;
243        self
244    }
245
246    /// Build the FlashMap. Tries GPU first, falls back to Rayon.
247    pub fn build<K: PodBound + Send + Sync, V: PodBound + Send + Sync>(
248        self,
249    ) -> Result<FlashMap<K, V>, FlashMapError> {
250        let mut _gpu_err: Option<FlashMapError> = None;
251
252        #[cfg(feature = "cuda")]
253        if !self.force_rayon {
254            match gpu::GpuFlashMap::<K, V>::new(
255                self.capacity,
256                self.hash_strategy,
257                self.device_id,
258            ) {
259                Ok(m) => return Ok(FlashMap { inner: FlashMapBackend::Gpu(m) }),
260                Err(e) => _gpu_err = Some(e),
261            }
262        }
263
264        #[cfg(feature = "rayon")]
265        {
266            if let Some(ref e) = _gpu_err {
267                eprintln!("[flash-map] GPU unavailable ({e}), using Rayon backend");
268            }
269            return Ok(FlashMap {
270                inner: FlashMapBackend::Rayon(rayon_cpu::RayonFlashMap::new(
271                    self.capacity,
272                    self.hash_strategy,
273                )),
274            });
275        }
276
277        #[allow(unreachable_code)]
278        match _gpu_err {
279            Some(e) => Err(e),
280            None => Err(FlashMapError::NoBackend),
281        }
282    }
283}
flash_map/lib.rs

flash_map/
lib.rs
