Merge commit 'd803bfe2b1fe7f5e219e50ac20d6801a0a58ac75' as 'vendor/ruvector'

vendor/ruvector/crates/sona/src/trajectory.rs

@@ -0,0 +1,362 @@
+//! Lock-free trajectory buffer for SONA
+//!
+//! Provides efficient, non-blocking trajectory recording during inference.
+
+use crate::time_compat::Instant;
+use crate::types::{QueryTrajectory, TrajectoryStep};
+use crossbeam::queue::ArrayQueue;
+use std::sync::atomic::{AtomicU64, Ordering};
+
+/// Lock-free trajectory buffer using crossbeam ArrayQueue
+pub struct TrajectoryBuffer {
+    /// Internal queue
+    buffer: ArrayQueue<QueryTrajectory>,
+    /// Capacity
+    capacity: usize,
+    /// Count of dropped trajectories
+    dropped: AtomicU64,
+    /// Total trajectories seen
+    total_seen: AtomicU64,
+}
+
+impl TrajectoryBuffer {
+    /// Create new buffer with capacity
+    pub fn new(capacity: usize) -> Self {
+        Self {
+            buffer: ArrayQueue::new(capacity),
+            capacity,
+            dropped: AtomicU64::new(0),
+            total_seen: AtomicU64::new(0),
+        }
+    }
+
+    /// Record trajectory (non-blocking)
+    ///
+    /// Returns true if recorded, false if buffer full
+    pub fn record(&self, trajectory: QueryTrajectory) -> bool {
+        self.total_seen.fetch_add(1, Ordering::Relaxed);
+
+        match self.buffer.push(trajectory) {
+            Ok(()) => true,
+            Err(_) => {
+                self.dropped.fetch_add(1, Ordering::Relaxed);
+                false
+            }
+        }
+    }
+
+    /// Try to pop single trajectory
+    pub fn pop(&self) -> Option<QueryTrajectory> {
+        self.buffer.pop()
+    }
+
+    /// Drain all trajectories
+    pub fn drain(&self) -> Vec<QueryTrajectory> {
+        let mut result = Vec::with_capacity(self.len());
+        while let Some(t) = self.buffer.pop() {
+            result.push(t);
+        }
+        result
+    }
+
+    /// Drain up to n trajectories
+    pub fn drain_n(&self, n: usize) -> Vec<QueryTrajectory> {
+        let mut result = Vec::with_capacity(n.min(self.len()));
+        for _ in 0..n {
+            match self.buffer.pop() {
+                Some(t) => result.push(t),
+                None => break,
+            }
+        }
+        result
+    }
+
+    /// Get current length
+    pub fn len(&self) -> usize {
+        self.buffer.len()
+    }
+
+    /// Check if empty
+    pub fn is_empty(&self) -> bool {
+        self.buffer.is_empty()
+    }
+
+    /// Check if full
+    pub fn is_full(&self) -> bool {
+        self.buffer.is_full()
+    }
+
+    /// Get capacity
+    pub fn capacity(&self) -> usize {
+        self.capacity
+    }
+
+    /// Get dropped count
+    pub fn dropped_count(&self) -> u64 {
+        self.dropped.load(Ordering::Relaxed)
+    }
+
+    /// Get total seen count
+    pub fn total_seen(&self) -> u64 {
+        self.total_seen.load(Ordering::Relaxed)
+    }
+
+    /// Get success rate
+    pub fn success_rate(&self) -> f64 {
+        let total = self.total_seen.load(Ordering::Relaxed);
+        let dropped = self.dropped.load(Ordering::Relaxed);
+        if total == 0 {
+            1.0
+        } else {
+            (total - dropped) as f64 / total as f64
+        }
+    }
+
+    /// Reset statistics (not the buffer contents)
+    pub fn reset_stats(&self) {
+        self.dropped.store(0, Ordering::Relaxed);
+        self.total_seen.store(0, Ordering::Relaxed);
+    }
+}
+
+/// Builder for constructing trajectories during inference
+pub struct TrajectoryBuilder {
+    /// Trajectory ID
+    id: u64,
+    /// Query embedding
+    query_embedding: Vec<f32>,
+    /// Steps collected
+    steps: Vec<TrajectoryStep>,
+    /// Start time
+    start_time: Instant,
+    /// Model route
+    model_route: Option<String>,
+    /// Context IDs
+    context_ids: Vec<String>,
+}
+
+impl TrajectoryBuilder {
+    /// Start new trajectory
+    pub fn new(id: u64, query_embedding: Vec<f32>) -> Self {
+        Self {
+            id,
+            query_embedding,
+            steps: Vec::with_capacity(16),
+            start_time: Instant::now(),
+            model_route: None,
+            context_ids: Vec::new(),
+        }
+    }
+
+    /// Add execution step
+    pub fn add_step(&mut self, activations: Vec<f32>, attention_weights: Vec<f32>, reward: f32) {
+        let step_idx = self.steps.len();
+        self.steps.push(TrajectoryStep::new(
+            activations,
+            attention_weights,
+            reward,
+            step_idx,
+        ));
+    }
+
+    /// Add step with layer name
+    pub fn add_named_step(
+        &mut self,
+        name: &str,
+        activations: Vec<f32>,
+        attention_weights: Vec<f32>,
+        reward: f32,
+    ) {
+        let step_idx = self.steps.len();
+        self.steps.push(
+            TrajectoryStep::new(activations, attention_weights, reward, step_idx).with_layer(name),
+        );
+    }
+
+    /// Set model route
+    pub fn set_model_route(&mut self, route: &str) {
+        self.model_route = Some(route.to_string());
+    }
+
+    /// Add context ID
+    pub fn add_context(&mut self, context_id: &str) {
+        self.context_ids.push(context_id.to_string());
+    }
+
+    /// Get current step count
+    pub fn step_count(&self) -> usize {
+        self.steps.len()
+    }
+
+    /// Get elapsed time
+    pub fn elapsed(&self) -> std::time::Duration {
+        self.start_time.elapsed()
+    }
+
+    /// Finalize and build trajectory
+    pub fn build(self, final_quality: f32) -> QueryTrajectory {
+        let latency_us = self.start_time.elapsed().as_micros() as u64;
+
+        QueryTrajectory {
+            id: self.id,
+            query_embedding: self.query_embedding,
+            steps: self.steps,
+            final_quality,
+            latency_us,
+            model_route: self.model_route,
+            context_ids: self.context_ids,
+        }
+    }
+
+    /// Build with explicit latency
+    pub fn build_with_latency(self, final_quality: f32, latency_us: u64) -> QueryTrajectory {
+        QueryTrajectory {
+            id: self.id,
+            query_embedding: self.query_embedding,
+            steps: self.steps,
+            final_quality,
+            latency_us,
+            model_route: self.model_route,
+            context_ids: self.context_ids,
+        }
+    }
+}
+
+/// Trajectory ID generator
+pub struct TrajectoryIdGen {
+    counter: AtomicU64,
+}
+
+impl TrajectoryIdGen {
+    /// Create new generator
+    pub fn new() -> Self {
+        Self {
+            counter: AtomicU64::new(0),
+        }
+    }
+
+    /// Create with starting ID
+    pub fn with_start(start: u64) -> Self {
+        Self {
+            counter: AtomicU64::new(start),
+        }
+    }
+
+    /// Generate next ID
+    pub fn next(&self) -> u64 {
+        self.counter.fetch_add(1, Ordering::Relaxed)
+    }
+
+    /// Get current value without incrementing
+    pub fn current(&self) -> u64 {
+        self.counter.load(Ordering::Relaxed)
+    }
+}
+
+impl Default for TrajectoryIdGen {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_buffer_basic_ops() {
+        let buffer = TrajectoryBuffer::new(10);
+
+        assert!(buffer.is_empty());
+        assert_eq!(buffer.capacity(), 10);
+
+        let trajectory = QueryTrajectory::new(1, vec![0.1, 0.2]);
+        assert!(buffer.record(trajectory));
+
+        assert_eq!(buffer.len(), 1);
+        assert!(!buffer.is_empty());
+    }
+
+    #[test]
+    fn test_buffer_overflow() {
+        let buffer = TrajectoryBuffer::new(3);
+
+        for i in 0..5 {
+            let trajectory = QueryTrajectory::new(i, vec![0.1]);
+            buffer.record(trajectory);
+        }
+
+        assert_eq!(buffer.len(), 3);
+        assert_eq!(buffer.dropped_count(), 2);
+        assert_eq!(buffer.total_seen(), 5);
+    }
+
+    #[test]
+    fn test_buffer_drain() {
+        let buffer = TrajectoryBuffer::new(10);
+
+        for i in 0..5 {
+            let trajectory = QueryTrajectory::new(i, vec![0.1]);
+            buffer.record(trajectory);
+        }
+
+        let drained = buffer.drain();
+        assert_eq!(drained.len(), 5);
+        assert!(buffer.is_empty());
+    }
+
+    #[test]
+    fn test_buffer_drain_n() {
+        let buffer = TrajectoryBuffer::new(10);
+
+        for i in 0..5 {
+            let trajectory = QueryTrajectory::new(i, vec![0.1]);
+            buffer.record(trajectory);
+        }
+
+        let partial = buffer.drain_n(3);
+        assert_eq!(partial.len(), 3);
+        assert_eq!(buffer.len(), 2);
+    }
+
+    #[test]
+    fn test_builder() {
+        let mut builder = TrajectoryBuilder::new(42, vec![0.1, 0.2, 0.3]);
+
+        builder.add_step(vec![0.5], vec![0.4, 0.6], 0.7);
+        builder.add_step(vec![0.6], vec![0.3, 0.7], 0.8);
+        builder.set_model_route("llama-7b");
+        builder.add_context("ctx-123");
+
+        assert_eq!(builder.step_count(), 2);
+
+        let trajectory = builder.build(0.85);
+
+        assert_eq!(trajectory.id, 42);
+        assert_eq!(trajectory.steps.len(), 2);
+        assert_eq!(trajectory.final_quality, 0.85);
+        assert_eq!(trajectory.model_route, Some("llama-7b".to_string()));
+        assert!(trajectory.latency_us > 0);
+    }
+
+    #[test]
+    fn test_id_generator() {
+        let gen = TrajectoryIdGen::new();
+
+        assert_eq!(gen.next(), 0);
+        assert_eq!(gen.next(), 1);
+        assert_eq!(gen.next(), 2);
+        assert_eq!(gen.current(), 3);
+    }
+
+    #[test]
+    fn test_success_rate() {
+        let buffer = TrajectoryBuffer::new(2);
+
+        for i in 0..4 {
+            buffer.record(QueryTrajectory::new(i, vec![]));
+        }
+
+        assert!((buffer.success_rate() - 0.5).abs() < 1e-6);
+    }
+}