1use crate::store::{BlockKey, BlockMeta, Tier};
20use crate::tiering::TierConfig;
21use std::collections::HashMap;
22
23#[derive(Clone, Debug)]
29pub struct PatternVector {
30 pub key: BlockKey,
32 pub embedding: Vec<f32>,
34 pub score: f32,
36}
37
38pub trait PatternIndex {
47 fn insert(&mut self, vec: &PatternVector);
49
50 fn search_nearest(&self, query: &[f32], k: usize) -> Vec<(BlockKey, f32)>;
53
54 fn remove(&mut self, key: BlockKey);
56
57 fn len(&self) -> usize;
59
60 fn is_empty(&self) -> bool {
62 self.len() == 0
63 }
64}
65
66fn cosine_similarity(a: &[f32], b: &[f32]) -> f32 {
74 if a.len() != b.len() || a.is_empty() {
75 return 0.0;
76 }
77
78 let mut dot = 0.0f32;
79 let mut norm_a_sq = 0.0f32;
80 let mut norm_b_sq = 0.0f32;
81
82 for (&x, &y) in a.iter().zip(b.iter()) {
83 dot += x * y;
84 norm_a_sq += x * x;
85 norm_b_sq += y * y;
86 }
87
88 let denom = norm_a_sq.sqrt() * norm_b_sq.sqrt();
89 if denom == 0.0 {
90 0.0
91 } else {
92 dot / denom
93 }
94}
95
96pub struct InMemoryPatternIndex {
106 vectors: Vec<PatternVector>,
107}
108
109impl InMemoryPatternIndex {
110 pub fn new() -> Self {
112 Self {
113 vectors: Vec::new(),
114 }
115 }
116}
117
118impl Default for InMemoryPatternIndex {
119 fn default() -> Self {
120 Self::new()
121 }
122}
123
124impl PatternIndex for InMemoryPatternIndex {
125 fn insert(&mut self, vec: &PatternVector) {
126 self.vectors.retain(|v| v.key != vec.key);
128 self.vectors.push(vec.clone());
129 }
130
131 fn search_nearest(&self, query: &[f32], k: usize) -> Vec<(BlockKey, f32)> {
132 if k == 0 || self.vectors.is_empty() {
133 return Vec::new();
134 }
135
136 let mut scored: Vec<(BlockKey, f32)> = self
137 .vectors
138 .iter()
139 .map(|v| (v.key, cosine_similarity(query, &v.embedding)))
140 .collect();
141
142 scored.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(core::cmp::Ordering::Equal));
144 scored.truncate(k);
145 scored
146 }
147
148 fn remove(&mut self, key: BlockKey) {
149 self.vectors.retain(|v| v.key != key);
150 }
151
152 fn len(&self) -> usize {
153 self.vectors.len()
154 }
155}
156
157pub fn pattern_from_meta(meta: &BlockMeta) -> Vec<f32> {
173 let ema = meta.ema_rate.clamp(0.0, 1.0);
174 let pop = meta.window.count_ones() as f32 / 64.0;
175 let recency = 1.0 / (1.0 + meta.tier_age as f32);
176 let count_log = ((1.0 + meta.access_count as f32).log2() / 32.0).clamp(0.0, 1.0);
177
178 vec![ema, pop, recency, count_log]
179}
180
181pub struct AdaptiveTiering<I: PatternIndex> {
198 pub index: I,
200 pub config: TierConfig,
202 block_tiers: HashMap<BlockKey, Tier>,
204}
205
206impl<I: PatternIndex> AdaptiveTiering<I> {
207 pub fn new(index: I, config: TierConfig) -> Self {
209 Self {
210 index,
211 config,
212 block_tiers: HashMap::new(),
213 }
214 }
215
216 pub fn register_block(&mut self, key: BlockKey, tier: Tier) {
222 self.block_tiers.insert(key, tier);
223 }
224
225 pub fn remove_block(&mut self, key: BlockKey) {
227 self.block_tiers.remove(&key);
228 self.index.remove(key);
229 }
230
231 pub fn registered_count(&self) -> usize {
233 self.block_tiers.len()
234 }
235
236 pub fn suggest_tier(&self, meta: &BlockMeta, neighbors: &[(BlockKey, f32)]) -> Option<Tier> {
249 if neighbors.is_empty() {
250 return None;
251 }
252
253 let mut votes = [0.0f32; 4];
256 let mut total_weight = 0.0f32;
257
258 for &(key, similarity) in neighbors {
259 if let Some(&tier) = self.block_tiers.get(&key) {
260 let weight = similarity.max(0.0);
261 votes[tier as u8 as usize] += weight;
262 total_weight += weight;
263 }
264 }
265
266 if total_weight == 0.0 {
267 return None;
268 }
269
270 let mut best_idx = 0usize;
273 let mut best_vote = votes[0];
274 for i in 1..4 {
275 if votes[i] > best_vote {
276 best_vote = votes[i];
277 best_idx = i;
278 }
279 }
280
281 let suggested = match best_idx {
282 0 => Tier::Tier0,
283 1 => Tier::Tier1,
284 2 => Tier::Tier2,
285 3 => Tier::Tier3,
286 _ => unreachable!(),
287 };
288
289 if suggested == meta.tier {
290 None
291 } else {
292 Some(suggested)
293 }
294 }
295}
296
297#[cfg(test)]
302mod tests {
303 use super::*;
304 use crate::store::{DType, ReconstructPolicy};
305
306 fn make_key(tid: u128, idx: u32) -> BlockKey {
307 BlockKey {
308 tensor_id: tid,
309 block_index: idx,
310 }
311 }
312
313 fn make_store_meta(
314 key: BlockKey,
315 tier: Tier,
316 ema_rate: f32,
317 window: u64,
318 access_count: u32,
319 tier_age: u32,
320 ) -> BlockMeta {
321 BlockMeta {
322 key,
323 dtype: DType::F32,
324 tier,
325 bits: 8,
326 scale: 1.0,
327 zero_point: 0,
328 created_at: 0,
329 last_access_at: 100,
330 access_count,
331 ema_rate,
332 window,
333 checksum: 0,
334 reconstruct: ReconstructPolicy::None,
335 tier_age,
336 lineage_parent: None,
337 block_bytes: 1024,
338 }
339 }
340
341 #[test]
344 fn cosine_identical_vectors() {
345 let v = vec![1.0, 2.0, 3.0, 4.0];
346 let sim = cosine_similarity(&v, &v);
347 assert!((sim - 1.0).abs() < 1e-6, "sim={sim}");
348 }
349
350 #[test]
351 fn cosine_orthogonal_vectors() {
352 let a = vec![1.0, 0.0];
353 let b = vec![0.0, 1.0];
354 let sim = cosine_similarity(&a, &b);
355 assert!(sim.abs() < 1e-6, "sim={sim}");
356 }
357
358 #[test]
359 fn cosine_opposite_vectors() {
360 let a = vec![1.0, 0.0, 0.0];
361 let b = vec![-1.0, 0.0, 0.0];
362 let sim = cosine_similarity(&a, &b);
363 assert!((sim - (-1.0)).abs() < 1e-6, "sim={sim}");
364 }
365
366 #[test]
367 fn cosine_zero_vector() {
368 let a = vec![1.0, 2.0];
369 let b = vec![0.0, 0.0];
370 assert_eq!(cosine_similarity(&a, &b), 0.0);
371 }
372
373 #[test]
374 fn cosine_different_lengths() {
375 let a = vec![1.0, 2.0];
376 let b = vec![1.0, 2.0, 3.0];
377 assert_eq!(cosine_similarity(&a, &b), 0.0);
378 }
379
380 #[test]
381 fn cosine_empty() {
382 let a: Vec<f32> = vec![];
383 let b: Vec<f32> = vec![];
384 assert_eq!(cosine_similarity(&a, &b), 0.0);
385 }
386
387 #[test]
388 fn cosine_known_value() {
389 let a = vec![1.0, 1.0];
391 let b = vec![1.0, 0.0];
392 let sim = cosine_similarity(&a, &b);
393 let expected = 1.0 / 2.0f32.sqrt();
394 assert!(
395 (sim - expected).abs() < 1e-6,
396 "sim={sim}, expected={expected}"
397 );
398 }
399
400 #[test]
403 fn index_insert_and_len() {
404 let mut idx = InMemoryPatternIndex::new();
405 assert!(idx.is_empty());
406
407 idx.insert(&PatternVector {
408 key: make_key(1, 0),
409 embedding: vec![1.0, 0.0, 0.0, 0.0],
410 score: 0.5,
411 });
412 assert_eq!(idx.len(), 1);
413 assert!(!idx.is_empty());
414 }
415
416 #[test]
417 fn index_insert_replaces_duplicate_key() {
418 let mut idx = InMemoryPatternIndex::new();
419 let key = make_key(1, 0);
420
421 idx.insert(&PatternVector {
422 key,
423 embedding: vec![1.0, 0.0, 0.0, 0.0],
424 score: 0.5,
425 });
426 idx.insert(&PatternVector {
427 key,
428 embedding: vec![0.0, 1.0, 0.0, 0.0],
429 score: 0.8,
430 });
431
432 assert_eq!(idx.len(), 1);
433
434 let results = idx.search_nearest(&[0.0, 1.0, 0.0, 0.0], 1);
436 assert_eq!(results.len(), 1);
437 assert_eq!(results[0].0, key);
438 assert!((results[0].1 - 1.0).abs() < 1e-6);
440 }
441
442 #[test]
443 fn index_remove() {
444 let mut idx = InMemoryPatternIndex::new();
445 let key = make_key(1, 0);
446
447 idx.insert(&PatternVector {
448 key,
449 embedding: vec![1.0, 0.0, 0.0, 0.0],
450 score: 0.5,
451 });
452 assert_eq!(idx.len(), 1);
453
454 idx.remove(key);
455 assert_eq!(idx.len(), 0);
456 }
457
458 #[test]
459 fn index_remove_nonexistent() {
460 let mut idx = InMemoryPatternIndex::new();
461 idx.remove(make_key(99, 0)); assert_eq!(idx.len(), 0);
463 }
464
465 #[test]
466 fn index_search_nearest_ordering() {
467 let mut idx = InMemoryPatternIndex::new();
468
469 idx.insert(&PatternVector {
471 key: make_key(1, 0),
472 embedding: vec![1.0, 0.0, 0.0, 0.0],
473 score: 0.0,
474 });
475 idx.insert(&PatternVector {
476 key: make_key(2, 0),
477 embedding: vec![0.7, 0.7, 0.0, 0.0],
478 score: 0.0,
479 });
480 idx.insert(&PatternVector {
481 key: make_key(3, 0),
482 embedding: vec![0.0, 1.0, 0.0, 0.0],
483 score: 0.0,
484 });
485
486 let results = idx.search_nearest(&[1.0, 0.1, 0.0, 0.0], 3);
488 assert_eq!(results.len(), 3);
489
490 assert_eq!(results[0].0, make_key(1, 0));
492 assert_eq!(results[1].0, make_key(2, 0));
494 assert_eq!(results[2].0, make_key(3, 0));
496
497 assert!(results[0].1 >= results[1].1);
499 assert!(results[1].1 >= results[2].1);
500 }
501
502 #[test]
503 fn index_search_nearest_k_larger_than_size() {
504 let mut idx = InMemoryPatternIndex::new();
505 idx.insert(&PatternVector {
506 key: make_key(1, 0),
507 embedding: vec![1.0, 0.0],
508 score: 0.0,
509 });
510
511 let results = idx.search_nearest(&[1.0, 0.0], 10);
512 assert_eq!(results.len(), 1);
513 }
514
515 #[test]
516 fn index_search_nearest_k_zero() {
517 let mut idx = InMemoryPatternIndex::new();
518 idx.insert(&PatternVector {
519 key: make_key(1, 0),
520 embedding: vec![1.0],
521 score: 0.0,
522 });
523
524 let results = idx.search_nearest(&[1.0], 0);
525 assert!(results.is_empty());
526 }
527
528 #[test]
529 fn index_search_nearest_empty() {
530 let idx = InMemoryPatternIndex::new();
531 let results = idx.search_nearest(&[1.0, 0.0], 5);
532 assert!(results.is_empty());
533 }
534
535 #[test]
538 fn pattern_from_meta_dimensions() {
539 let meta = make_store_meta(make_key(1, 0), Tier::Tier1, 0.5, 0xFFFF, 100, 10);
540 let pat = pattern_from_meta(&meta);
541 assert_eq!(pat.len(), 4);
542 }
543
544 #[test]
545 fn pattern_from_meta_ema_component() {
546 let meta = make_store_meta(make_key(1, 0), Tier::Tier1, 0.8, 0, 0, 0);
547 let pat = pattern_from_meta(&meta);
548 assert!((pat[0] - 0.8).abs() < 1e-6, "ema={}", pat[0]);
549 }
550
551 #[test]
552 fn pattern_from_meta_popcount_component() {
553 let meta = make_store_meta(make_key(1, 0), Tier::Tier1, 0.0, u64::MAX, 0, 0);
555 let pat = pattern_from_meta(&meta);
556 assert!((pat[1] - 1.0).abs() < 1e-6, "pop={}", pat[1]);
557
558 let meta2 = make_store_meta(make_key(1, 0), Tier::Tier1, 0.0, 0, 0, 0);
560 let pat2 = pattern_from_meta(&meta2);
561 assert!((pat2[1]).abs() < 1e-6, "pop={}", pat2[1]);
562
563 let meta3 = make_store_meta(make_key(1, 0), Tier::Tier1, 0.0, 0xFFFF_FFFF, 0, 0);
565 let pat3 = pattern_from_meta(&meta3);
566 assert!((pat3[1] - 0.5).abs() < 1e-6, "pop={}", pat3[1]);
567 }
568
569 #[test]
570 fn pattern_from_meta_recency_component() {
571 let meta = make_store_meta(make_key(1, 0), Tier::Tier1, 0.0, 0, 0, 0);
573 let pat = pattern_from_meta(&meta);
574 assert!((pat[2] - 1.0).abs() < 1e-6, "recency={}", pat[2]);
575
576 let meta2 = make_store_meta(make_key(1, 0), Tier::Tier1, 0.0, 0, 0, 9);
578 let pat2 = pattern_from_meta(&meta2);
579 assert!((pat2[2] - 0.1).abs() < 1e-6, "recency={}", pat2[2]);
580 }
581
582 #[test]
583 fn pattern_from_meta_access_count_log_component() {
584 let meta = make_store_meta(make_key(1, 0), Tier::Tier1, 0.0, 0, 0, 0);
586 let pat = pattern_from_meta(&meta);
587 assert!(pat[3].abs() < 1e-6, "count_log={}", pat[3]);
588
589 let meta2 = make_store_meta(make_key(1, 0), Tier::Tier1, 0.0, 0, 1, 0);
591 let pat2 = pattern_from_meta(&meta2);
592 assert!((pat2[3] - 1.0 / 32.0).abs() < 1e-4, "count_log={}", pat2[3]);
593 }
594
595 #[test]
596 fn pattern_from_meta_values_in_unit_range() {
597 let meta = make_store_meta(
599 make_key(1, 0),
600 Tier::Tier1,
601 2.0, u64::MAX, u32::MAX, u32::MAX, );
606 let pat = pattern_from_meta(&meta);
607 for (i, &v) in pat.iter().enumerate() {
608 assert!(v >= 0.0 && v <= 1.0, "dim {i} out of [0,1]: {v}");
609 }
610 }
611
612 #[test]
615 fn adaptive_new_and_register() {
616 let idx = InMemoryPatternIndex::new();
617 let config = TierConfig::default();
618 let mut at = AdaptiveTiering::new(idx, config);
619
620 assert_eq!(at.registered_count(), 0);
621
622 at.register_block(make_key(1, 0), Tier::Tier1);
623 assert_eq!(at.registered_count(), 1);
624
625 at.register_block(make_key(1, 0), Tier::Tier2);
626 assert_eq!(at.registered_count(), 1); }
628
629 #[test]
630 fn adaptive_remove_block() {
631 let mut idx = InMemoryPatternIndex::new();
632 let key = make_key(1, 0);
633 idx.insert(&PatternVector {
634 key,
635 embedding: vec![1.0, 0.0, 0.0, 0.0],
636 score: 0.5,
637 });
638
639 let config = TierConfig::default();
640 let mut at = AdaptiveTiering::new(idx, config);
641 at.register_block(key, Tier::Tier1);
642 assert_eq!(at.registered_count(), 1);
643 assert_eq!(at.index.len(), 1);
644
645 at.remove_block(key);
646 assert_eq!(at.registered_count(), 0);
647 assert_eq!(at.index.len(), 0);
648 }
649
650 #[test]
651 fn suggest_tier_empty_neighbors() {
652 let idx = InMemoryPatternIndex::new();
653 let config = TierConfig::default();
654 let at = AdaptiveTiering::new(idx, config);
655
656 let meta = make_store_meta(make_key(1, 0), Tier::Tier1, 0.5, 0, 10, 5);
657 let result = at.suggest_tier(&meta, &[]);
658 assert_eq!(result, None);
659 }
660
661 #[test]
662 fn suggest_tier_no_known_neighbors() {
663 let idx = InMemoryPatternIndex::new();
664 let config = TierConfig::default();
665 let at = AdaptiveTiering::new(idx, config);
666
667 let meta = make_store_meta(make_key(1, 0), Tier::Tier1, 0.5, 0, 10, 5);
668 let neighbors = vec![(make_key(2, 0), 0.9), (make_key(3, 0), 0.8)];
670 let result = at.suggest_tier(&meta, &neighbors);
671 assert_eq!(result, None);
672 }
673
674 #[test]
675 fn suggest_tier_unanimous_vote() {
676 let idx = InMemoryPatternIndex::new();
677 let config = TierConfig::default();
678 let mut at = AdaptiveTiering::new(idx, config);
679
680 at.register_block(make_key(2, 0), Tier::Tier3);
682 at.register_block(make_key(3, 0), Tier::Tier3);
683 at.register_block(make_key(4, 0), Tier::Tier3);
684
685 let meta = make_store_meta(make_key(1, 0), Tier::Tier1, 0.5, 0, 10, 5);
686 let neighbors = vec![
687 (make_key(2, 0), 0.9),
688 (make_key(3, 0), 0.8),
689 (make_key(4, 0), 0.7),
690 ];
691
692 let result = at.suggest_tier(&meta, &neighbors);
693 assert_eq!(result, Some(Tier::Tier3));
694 }
695
696 #[test]
697 fn suggest_tier_same_as_current_returns_none() {
698 let idx = InMemoryPatternIndex::new();
699 let config = TierConfig::default();
700 let mut at = AdaptiveTiering::new(idx, config);
701
702 at.register_block(make_key(2, 0), Tier::Tier1);
704 at.register_block(make_key(3, 0), Tier::Tier1);
705
706 let meta = make_store_meta(make_key(1, 0), Tier::Tier1, 0.5, 0, 10, 5);
707 let neighbors = vec![(make_key(2, 0), 0.9), (make_key(3, 0), 0.8)];
708
709 let result = at.suggest_tier(&meta, &neighbors);
710 assert_eq!(result, None);
711 }
712
713 #[test]
714 fn suggest_tier_weighted_majority() {
715 let idx = InMemoryPatternIndex::new();
716 let config = TierConfig::default();
717 let mut at = AdaptiveTiering::new(idx, config);
718
719 at.register_block(make_key(2, 0), Tier::Tier1);
721 at.register_block(make_key(3, 0), Tier::Tier1);
722 at.register_block(make_key(4, 0), Tier::Tier3);
724
725 let meta = make_store_meta(make_key(1, 0), Tier::Tier2, 0.5, 0, 10, 5);
726 let neighbors = vec![
727 (make_key(2, 0), 0.3), (make_key(3, 0), 0.3), (make_key(4, 0), 0.9), ];
731 let result = at.suggest_tier(&meta, &neighbors);
733 assert_eq!(result, Some(Tier::Tier3));
734 }
735
736 #[test]
737 fn suggest_tier_negative_similarity_ignored() {
738 let idx = InMemoryPatternIndex::new();
739 let config = TierConfig::default();
740 let mut at = AdaptiveTiering::new(idx, config);
741
742 at.register_block(make_key(2, 0), Tier::Tier3);
743 at.register_block(make_key(3, 0), Tier::Tier1);
744
745 let meta = make_store_meta(make_key(1, 0), Tier::Tier2, 0.5, 0, 10, 5);
746 let neighbors = vec![
747 (make_key(2, 0), -0.5), (make_key(3, 0), 0.5), ];
750 let result = at.suggest_tier(&meta, &neighbors);
752 assert_eq!(result, Some(Tier::Tier1));
753 }
754
755 #[test]
756 fn suggest_tier_zero_similarity_all() {
757 let idx = InMemoryPatternIndex::new();
758 let config = TierConfig::default();
759 let mut at = AdaptiveTiering::new(idx, config);
760
761 at.register_block(make_key(2, 0), Tier::Tier3);
762
763 let meta = make_store_meta(make_key(1, 0), Tier::Tier1, 0.5, 0, 10, 5);
764 let neighbors = vec![(make_key(2, 0), 0.0)];
765
766 let result = at.suggest_tier(&meta, &neighbors);
768 assert_eq!(result, None);
769 }
770
771 #[test]
774 fn integration_end_to_end() {
775 let mut idx = InMemoryPatternIndex::new();
776 let config = TierConfig::default();
777
778 let hot_key = make_key(1, 0);
780 let warm_key = make_key(2, 0);
781 let cold_key = make_key(3, 0);
782
783 let hot_meta = make_store_meta(hot_key, Tier::Tier1, 0.9, u64::MAX, 1000, 2);
784 let warm_meta = make_store_meta(warm_key, Tier::Tier2, 0.5, 0xFFFF_FFFF, 100, 10);
785 let cold_meta = make_store_meta(cold_key, Tier::Tier3, 0.05, 0x0F, 5, 100);
786
787 let hot_emb = pattern_from_meta(&hot_meta);
789 let warm_emb = pattern_from_meta(&warm_meta);
790 let cold_emb = pattern_from_meta(&cold_meta);
791
792 idx.insert(&PatternVector {
793 key: hot_key,
794 embedding: hot_emb.clone(),
795 score: 0.9,
796 });
797 idx.insert(&PatternVector {
798 key: warm_key,
799 embedding: warm_emb.clone(),
800 score: 0.5,
801 });
802 idx.insert(&PatternVector {
803 key: cold_key,
804 embedding: cold_emb.clone(),
805 score: 0.1,
806 });
807
808 let mut at = AdaptiveTiering::new(idx, config);
809 at.register_block(hot_key, Tier::Tier1);
810 at.register_block(warm_key, Tier::Tier2);
811 at.register_block(cold_key, Tier::Tier3);
812
813 let new_key = make_key(4, 0);
815 let new_meta = make_store_meta(new_key, Tier::Tier3, 0.85, u64::MAX, 800, 3);
816 let new_emb = pattern_from_meta(&new_meta);
817
818 let neighbors = at.index.search_nearest(&new_emb, 3);
819 assert!(!neighbors.is_empty());
820
821 let suggestion = at.suggest_tier(&new_meta, &neighbors);
822 assert!(
825 suggestion.is_some(),
826 "expected a tier suggestion for a hot-like pattern in Tier3"
827 );
828 let suggested = suggestion.unwrap();
829 assert_ne!(suggested, Tier::Tier3, "should not stay cold");
830 }
831}