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

flash_map/
lib.rs
