diff --git a/.github/workflows/rust.yml b/.github/workflows/rust.yml index 6eb3b703e6..42f3ecaaf1 100644 --- a/.github/workflows/rust.yml +++ b/.github/workflows/rust.yml @@ -42,10 +42,12 @@ jobs: - name: program-libs-fast packages: aligned-sized light-hasher light-compressed-account light-account-checks \ - light-verifier light-merkle-tree-metadata light-zero-copy light-hash-set light-concurrent-merkle-tree + light-verifier light-merkle-tree-metadata light-zero-copy light-hash-set light-concurrent-merkle-tree \ + light-array-map test_cmd: | cargo test -p light-macros cargo test -p aligned-sized + cargo test -p light-array-map --all-features cargo test -p light-hasher --all-features cargo test -p light-compressed-account --all-features cargo test -p light-compressed-account --features new-unique,poseidon diff --git a/Cargo.lock b/Cargo.lock index 02dd0a583e..4d7bf0835a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3375,6 +3375,13 @@ dependencies = [ "thiserror 2.0.17", ] +[[package]] +name = "light-array-map" +version = "0.1.0" +dependencies = [ + "tinyvec", +] + [[package]] name = "light-batched-merkle-tree" version = "0.5.0" diff --git a/Cargo.toml b/Cargo.toml index 3641c22599..97dbdb0882 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,7 @@ [workspace] members = [ "program-libs/account-checks", + "program-libs/array-map", "program-libs/compressed-account", "program-libs/aligned-sized", "program-libs/batched-merkle-tree", @@ -200,11 +201,13 @@ light-bounded-vec = { version = "2.0.0" } light-poseidon = { version = "0.3.0" } light-test-utils = { path = "program-tests/utils", version = "1.2.1" } light-indexed-array = { path = "program-libs/indexed-array", version = "0.2.0" } +light-array-map = { path = "program-libs/array-map", version = "0.1.0" } light-program-profiler = { version = "0.1.0" } create-address-program-test = { path = "program-tests/create-address-test-program", version = "1.0.0" } groth16-solana = { version = "0.2.0" } bytemuck = { version = "1.19.0" } arrayvec = "0.7" +tinyvec = "1.10.0" # Math and crypto num-bigint = "0.4.6" diff --git a/program-libs/array-map/Cargo.toml b/program-libs/array-map/Cargo.toml new file mode 100644 index 0000000000..6035eb42e7 --- /dev/null +++ b/program-libs/array-map/Cargo.toml @@ -0,0 +1,11 @@ +[package] +name = "light-array-map" +version = "0.1.0" +description = "Generic array-backed map with O(n) lookup for small collections" +repository = "https://github.com/Lightprotocol/light-protocol" +license = "Apache-2.0" +edition = "2021" + + +[dependencies] +tinyvec = { workspace = true } diff --git a/program-libs/array-map/src/lib.rs b/program-libs/array-map/src/lib.rs new file mode 100644 index 0000000000..560d6ed45b --- /dev/null +++ b/program-libs/array-map/src/lib.rs @@ -0,0 +1,181 @@ +#![no_std] + +use core::ptr::read_unaligned; + +use tinyvec::ArrayVec; + +/// A generic tinyvec::ArrayVec backed map with O(n) lookup. +/// Maintains insertion order and tracks the last updated entry index. +pub struct ArrayMap +where + K: PartialEq + Default, + V: Default, +{ + entries: ArrayVec<[(K, V); N]>, + last_updated_index: Option, +} + +impl ArrayMap +where + K: PartialEq + Default, + V: Default, +{ + pub fn new() -> Self { + Self { + entries: ArrayVec::new(), + last_updated_index: None, + } + } + + pub fn len(&self) -> usize { + self.entries.len() + } + + pub fn is_empty(&self) -> bool { + self.entries.is_empty() + } + + pub fn last_updated_index(&self) -> Option { + self.last_updated_index + } + + pub fn get(&self, index: usize) -> Option<&(K, V)> { + self.entries.get(index) + } + + pub fn get_mut(&mut self, index: usize) -> Option<&mut (K, V)> { + self.entries.get_mut(index) + } + + pub fn get_u8(&self, index: u8) -> Option<&(K, V)> { + self.get(index as usize) + } + + pub fn get_mut_u8(&mut self, index: u8) -> Option<&mut (K, V)> { + self.get_mut(index as usize) + } + + pub fn get_by_key(&self, key: &K) -> Option<&V> { + self.entries.iter().find(|(k, _)| k == key).map(|(_, v)| v) + } + + pub fn get_mut_by_key(&mut self, key: &K) -> Option<&mut V> { + self.entries + .iter_mut() + .find(|(k, _)| k == key) + .map(|(_, v)| v) + } + + pub fn find(&self, key: &K) -> Option<(usize, &(K, V))> { + self.entries.iter().enumerate().find(|(_, (k, _))| k == key) + } + + pub fn find_mut(&mut self, key: &K) -> Option<(usize, &mut (K, V))> { + self.entries + .iter_mut() + .enumerate() + .find(|(_, (k, _))| k == key) + } + + pub fn find_index(&self, key: &K) -> Option { + self.find(key).map(|(idx, _)| idx) + } + + pub fn set_last_updated_index(&mut self, index: usize) -> Result<(), E> + where + E: From, + { + if index < self.entries.len() { + self.last_updated_index = Some(index); + Ok(()) + } else { + Err(ArrayMapError::IndexOutOfBounds.into()) + } + } + + pub fn insert(&mut self, key: K, value: V, error: E) -> Result { + let new_idx = self.entries.len(); + // tinyvec's try_push returns Some(item) on failure, None on success + if self.entries.try_push((key, value)).is_some() { + return Err(error); + } + self.last_updated_index = Some(new_idx); + Ok(new_idx) + } +} + +impl Default for ArrayMap +where + K: PartialEq + Default, + V: Default, +{ + fn default() -> Self { + Self::new() + } +} + +// Optimized [u8; 32] key methods (4x u64 comparison instead of 32x u8). +impl ArrayMap<[u8; 32], V, N> +where + V: Default, +{ + pub fn get_by_pubkey(&self, key: &[u8; 32]) -> Option<&V> { + self.entries + .iter() + .find(|(k, _)| pubkey_eq(k, key)) + .map(|(_, v)| v) + } + + pub fn get_mut_by_pubkey(&mut self, key: &[u8; 32]) -> Option<&mut V> { + self.entries + .iter_mut() + .find(|(k, _)| pubkey_eq(k, key)) + .map(|(_, v)| v) + } + + pub fn find_by_pubkey(&self, key: &[u8; 32]) -> Option<(usize, &([u8; 32], V))> { + self.entries + .iter() + .enumerate() + .find(|(_, (k, _))| pubkey_eq(k, key)) + } + + pub fn find_mut_by_pubkey(&mut self, key: &[u8; 32]) -> Option<(usize, &mut ([u8; 32], V))> { + self.entries + .iter_mut() + .enumerate() + .find(|(_, (k, _))| pubkey_eq(k, key)) + } + + pub fn find_pubkey_index(&self, key: &[u8; 32]) -> Option { + self.find_by_pubkey(key).map(|(idx, _)| idx) + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ArrayMapError { + CapacityExceeded, + IndexOutOfBounds, +} + +impl core::fmt::Display for ArrayMapError { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + match self { + ArrayMapError::CapacityExceeded => write!(f, "ArrayMap capacity exceeded"), + ArrayMapError::IndexOutOfBounds => write!(f, "ArrayMap index out of bounds"), + } + } +} + +#[inline(always)] +pub const fn pubkey_eq(p1: &[u8; 32], p2: &[u8; 32]) -> bool { + let p1_ptr = p1.as_ptr() as *const u64; + let p2_ptr = p2.as_ptr() as *const u64; + + unsafe { + read_unaligned(p1_ptr) == read_unaligned(p2_ptr) + && read_unaligned(p1_ptr.add(1)) == read_unaligned(p2_ptr.add(1)) + && read_unaligned(p1_ptr.add(2)) == read_unaligned(p2_ptr.add(2)) + && read_unaligned(p1_ptr.add(3)) == read_unaligned(p2_ptr.add(3)) + } +} diff --git a/program-libs/array-map/tests/array_map_tests.rs b/program-libs/array-map/tests/array_map_tests.rs new file mode 100644 index 0000000000..765c9bf595 --- /dev/null +++ b/program-libs/array-map/tests/array_map_tests.rs @@ -0,0 +1,198 @@ +use light_array_map::{ArrayMap, ArrayMapError}; + +// Test error type for testing +#[derive(Debug, PartialEq)] +enum TestError { + ArrayMap(ArrayMapError), + Custom, +} + +impl From for TestError { + fn from(e: ArrayMapError) -> Self { + TestError::ArrayMap(e) + } +} + +#[test] +fn test_new_map() { + let map = ArrayMap::::new(); + assert_eq!(map.len(), 0); + assert!(map.is_empty()); + assert!(map.last_updated_index().is_none()); +} + +#[test] +fn test_insert() { + let mut map = ArrayMap::::new(); + + let idx = map.insert(1, "one".to_string(), TestError::Custom).unwrap(); + + assert_eq!(idx, 0); + assert_eq!(map.len(), 1); + assert_eq!(map.last_updated_index(), Some(0)); + assert_eq!(map.get(0).unwrap().1, "one"); +} + +#[test] +fn test_get_by_key() { + let mut map = ArrayMap::::new(); + + map.insert(1, "one".to_string(), TestError::Custom).unwrap(); + map.insert(2, "two".to_string(), TestError::Custom).unwrap(); + + assert_eq!(map.get_by_key(&1), Some(&"one".to_string())); + assert_eq!(map.get_by_key(&2), Some(&"two".to_string())); + assert_eq!(map.get_by_key(&3), None); +} + +#[test] +fn test_get_mut_by_key() { + let mut map = ArrayMap::::new(); + + map.insert(1, "one".to_string(), TestError::Custom).unwrap(); + + if let Some(val) = map.get_mut_by_key(&1) { + *val = "ONE".to_string(); + } + + assert_eq!(map.get_by_key(&1), Some(&"ONE".to_string())); +} + +#[test] +fn test_find_index() { + let mut map = ArrayMap::::new(); + + map.insert(10, "ten".to_string(), TestError::Custom) + .unwrap(); + map.insert(20, "twenty".to_string(), TestError::Custom) + .unwrap(); + + assert_eq!(map.find_index(&10), Some(0)); + assert_eq!(map.find_index(&20), Some(1)); + assert_eq!(map.find_index(&30), None); +} + +#[test] +fn test_set_last_updated_index() { + let mut map = ArrayMap::::new(); + + map.insert(1, "one".to_string(), TestError::Custom).unwrap(); + map.insert(2, "two".to_string(), TestError::Custom).unwrap(); + + // Should be at index 1 after last insert + assert_eq!(map.last_updated_index(), Some(1)); + + // Set to 0 + map.set_last_updated_index::(0).unwrap(); + assert_eq!(map.last_updated_index(), Some(0)); + + // Out of bounds should fail + let result = map.set_last_updated_index::(10); + assert!(result.is_err()); + assert_eq!( + result.unwrap_err(), + TestError::ArrayMap(ArrayMapError::IndexOutOfBounds) + ); +} + +#[test] +fn test_capacity_limit() { + let mut map = ArrayMap::::new(); + + // Fill to capacity + for i in 0..5 { + map.insert(i, format!("val{}", i), TestError::Custom) + .unwrap(); + } + + assert_eq!(map.len(), 5); + + // 6th entry should fail + let result = map.insert(5, "val5".to_string(), TestError::Custom); + assert!(result.is_err()); +} + +#[test] +fn test_get_mut_direct() { + let mut map = ArrayMap::::new(); + + map.insert(1, 100, TestError::Custom).unwrap(); + + if let Some(entry) = map.get_mut(0) { + entry.1 += 50; + } + + assert_eq!(map.get(0).unwrap().1, 150); +} + +#[test] +fn test_last_updated_index_updates() { + let mut map = ArrayMap::::new(); + + // Insert first entry + map.insert(1, 100, TestError::Custom).unwrap(); + assert_eq!(map.last_updated_index(), Some(0)); + + // Insert second entry + map.insert(2, 200, TestError::Custom).unwrap(); + assert_eq!(map.last_updated_index(), Some(1)); +} + +#[test] +fn test_capacity_overflow_without_alloc() { + // Demonstrate that ArrayVec has fixed capacity regardless of alloc feature + let mut map = ArrayMap::::new(); + + // Fill to capacity + map.insert(1, 100, TestError::Custom).unwrap(); + map.insert(2, 200, TestError::Custom).unwrap(); + map.insert(3, 300, TestError::Custom).unwrap(); + + assert_eq!(map.len(), 3); + + // 4th insert should fail - fixed capacity + let result = map.insert(4, 400, TestError::Custom); + assert!(result.is_err(), "ArrayVec has fixed capacity"); +} + +#[test] +fn test_get_u8() { + let mut map = ArrayMap::::new(); + + map.insert(1, "one".to_string(), TestError::Custom).unwrap(); + map.insert(2, "two".to_string(), TestError::Custom).unwrap(); + map.insert(3, "three".to_string(), TestError::Custom) + .unwrap(); + + // Test valid indices + assert_eq!(map.get_u8(0).unwrap().1, "one"); + assert_eq!(map.get_u8(1).unwrap().1, "two"); + assert_eq!(map.get_u8(2).unwrap().1, "three"); + + // Test out of bounds + assert!(map.get_u8(3).is_none()); + assert!(map.get_u8(255).is_none()); +} + +#[test] +fn test_get_mut_u8() { + let mut map = ArrayMap::::new(); + + map.insert(1, 100, TestError::Custom).unwrap(); + map.insert(2, 200, TestError::Custom).unwrap(); + map.insert(3, 300, TestError::Custom).unwrap(); + + // Modify via get_mut_u8 + if let Some(entry) = map.get_mut_u8(1) { + entry.1 += 50; + } + + // Verify modification + assert_eq!(map.get_u8(1).unwrap().1, 250); + assert_eq!(map.get_u8(0).unwrap().1, 100); + assert_eq!(map.get_u8(2).unwrap().1, 300); + + // Test out of bounds + assert!(map.get_mut_u8(3).is_none()); + assert!(map.get_mut_u8(255).is_none()); +} diff --git a/program-libs/compressed-account/src/pubkey.rs b/program-libs/compressed-account/src/pubkey.rs index 9325f96988..7a6d3cc99b 100644 --- a/program-libs/compressed-account/src/pubkey.rs +++ b/program-libs/compressed-account/src/pubkey.rs @@ -74,6 +74,10 @@ impl Pubkey { array.copy_from_slice(slice); Self(array) } + + pub fn array_ref(&self) -> &[u8; 32] { + &self.0 + } } impl AsRef for Pubkey {