//! # Arms Engine //! //! The main ARMS orchestrator. //! //! This struct wires together: //! - Storage (Place port) //! - Index (Near port) //! - Configuration //! //! And exposes a unified API for storing and retrieving points. use crate::core::{Blob, Id, PlacedPoint, Point}; use crate::core::config::ArmsConfig; use crate::ports::{Near, NearResult, Place, PlaceResult, SearchResult}; use crate::adapters::storage::MemoryStorage; use crate::adapters::index::FlatIndex; /// The main ARMS engine /// /// Orchestrates storage and indexing with a unified API. pub struct Arms { /// Configuration config: ArmsConfig, /// Storage backend (Place port) storage: Box, /// Index backend (Near port) index: Box, } impl Arms { /// Create a new ARMS instance with default adapters /// /// Uses MemoryStorage and FlatIndex. /// For production, use `Arms::with_adapters` with appropriate backends. pub fn new(config: ArmsConfig) -> Self { let storage = Box::new(MemoryStorage::new(config.dimensionality)); let index = Box::new(FlatIndex::new( config.dimensionality, config.proximity.clone(), true, // Assuming cosine-like similarity by default )); Self { config, storage, index, } } /// Create with custom adapters pub fn with_adapters( config: ArmsConfig, storage: Box, index: Box, ) -> Self { Self { config, storage, index, } } /// Get the configuration pub fn config(&self) -> &ArmsConfig { &self.config } /// Get the dimensionality of this space pub fn dimensionality(&self) -> usize { self.config.dimensionality } // ======================================================================== // PLACE OPERATIONS // ======================================================================== /// Place a point in the space /// /// The point will be normalized if configured to do so. /// Returns the assigned ID. pub fn place(&mut self, point: Point, blob: Blob) -> PlaceResult { // Normalize if configured let point = if self.config.normalize_on_insert { point.normalize() } else { point }; // Store in storage let id = self.storage.place(point.clone(), blob)?; // Add to index if let Err(e) = self.index.add(id, &point) { // Rollback storage if index fails self.storage.remove(id); return Err(crate::ports::PlaceError::StorageError(format!( "Index error: {:?}", e ))); } Ok(id) } /// Place multiple points at once pub fn place_batch(&mut self, items: Vec<(Point, Blob)>) -> Vec> { items .into_iter() .map(|(point, blob)| self.place(point, blob)) .collect() } /// Remove a point from the space pub fn remove(&mut self, id: Id) -> Option { // Remove from index first let _ = self.index.remove(id); // Then from storage self.storage.remove(id) } /// Get a point by ID pub fn get(&self, id: Id) -> Option<&PlacedPoint> { self.storage.get(id) } /// Check if a point exists pub fn contains(&self, id: Id) -> bool { self.storage.contains(id) } /// Get the number of stored points pub fn len(&self) -> usize { self.storage.len() } /// Check if the space is empty pub fn is_empty(&self) -> bool { self.storage.is_empty() } /// Clear all points pub fn clear(&mut self) { self.storage.clear(); let _ = self.index.rebuild(); // Reset index } // ======================================================================== // NEAR OPERATIONS // ======================================================================== /// Find k nearest points to query pub fn near(&self, query: &Point, k: usize) -> NearResult> { // Normalize query if configured let query = if self.config.normalize_on_insert { query.normalize() } else { query.clone() }; self.index.near(&query, k) } /// Find all points within threshold pub fn within(&self, query: &Point, threshold: f32) -> NearResult> { let query = if self.config.normalize_on_insert { query.normalize() } else { query.clone() }; self.index.within(&query, threshold) } /// Find and retrieve k nearest points (with full data) pub fn near_with_data(&self, query: &Point, k: usize) -> NearResult> { let results = self.near(query, k)?; Ok(results .into_iter() .filter_map(|r| self.storage.get(r.id).map(|p| (p, r.score))) .collect()) } // ======================================================================== // MERGE OPERATIONS // ======================================================================== /// Merge multiple points into one using the configured merge function pub fn merge(&self, points: &[Point]) -> Point { self.config.merge.merge(points) } /// Compute proximity between two points pub fn proximity(&self, a: &Point, b: &Point) -> f32 { self.config.proximity.proximity(a, b) } // ======================================================================== // STATS // ======================================================================== /// Get storage size in bytes pub fn size_bytes(&self) -> usize { self.storage.size_bytes() } /// Get index stats pub fn index_len(&self) -> usize { self.index.len() } /// Check if index is ready pub fn is_ready(&self) -> bool { self.index.is_ready() } } #[cfg(test)] mod tests { use super::*; fn create_test_arms() -> Arms { Arms::new(ArmsConfig::new(3)) } #[test] fn test_arms_place_and_get() { let mut arms = create_test_arms(); let point = Point::new(vec![1.0, 0.0, 0.0]); let blob = Blob::from_str("test data"); let id = arms.place(point, blob).unwrap(); let retrieved = arms.get(id).unwrap(); assert_eq!(retrieved.blob.as_str(), Some("test data")); } #[test] fn test_arms_near() { let mut arms = create_test_arms(); // Add some points arms.place(Point::new(vec![1.0, 0.0, 0.0]), Blob::from_str("x")).unwrap(); arms.place(Point::new(vec![0.0, 1.0, 0.0]), Blob::from_str("y")).unwrap(); arms.place(Point::new(vec![0.0, 0.0, 1.0]), Blob::from_str("z")).unwrap(); // Query let query = Point::new(vec![1.0, 0.0, 0.0]); let results = arms.near(&query, 2).unwrap(); assert_eq!(results.len(), 2); // First result should have highest similarity assert!(results[0].score > results[1].score); } #[test] fn test_arms_near_with_data() { let mut arms = create_test_arms(); arms.place(Point::new(vec![1.0, 0.0, 0.0]), Blob::from_str("x")).unwrap(); arms.place(Point::new(vec![0.0, 1.0, 0.0]), Blob::from_str("y")).unwrap(); let query = Point::new(vec![1.0, 0.0, 0.0]); let results = arms.near_with_data(&query, 1).unwrap(); assert_eq!(results.len(), 1); assert_eq!(results[0].0.blob.as_str(), Some("x")); } #[test] fn test_arms_remove() { let mut arms = create_test_arms(); let id = arms.place(Point::new(vec![1.0, 0.0, 0.0]), Blob::empty()).unwrap(); assert!(arms.contains(id)); assert_eq!(arms.len(), 1); arms.remove(id); assert!(!arms.contains(id)); assert_eq!(arms.len(), 0); } #[test] fn test_arms_merge() { let arms = create_test_arms(); let points = vec![ Point::new(vec![1.0, 0.0, 0.0]), Point::new(vec![0.0, 1.0, 0.0]), ]; let merged = arms.merge(&points); // Mean of [1,0,0] and [0,1,0] = [0.5, 0.5, 0] assert!((merged.dims()[0] - 0.5).abs() < 0.0001); assert!((merged.dims()[1] - 0.5).abs() < 0.0001); assert!((merged.dims()[2] - 0.0).abs() < 0.0001); } #[test] fn test_arms_clear() { let mut arms = create_test_arms(); for i in 0..10 { arms.place(Point::new(vec![i as f32, 0.0, 0.0]), Blob::empty()).unwrap(); } assert_eq!(arms.len(), 10); arms.clear(); assert_eq!(arms.len(), 0); assert!(arms.is_empty()); } #[test] fn test_arms_normalizes_on_insert() { let mut arms = create_test_arms(); // Insert a non-normalized point let point = Point::new(vec![3.0, 4.0, 0.0]); // magnitude = 5 let id = arms.place(point, Blob::empty()).unwrap(); let retrieved = arms.get(id).unwrap(); // Should be normalized assert!(retrieved.point.is_normalized()); } }