Merge commit 'd803bfe2b1fe7f5e219e50ac20d6801a0a58ac75' as 'vendor/ruvector'

This commit is contained in:
ruv
2026-02-28 14:39:40 -05:00
7854 changed files with 3522914 additions and 0 deletions

View File

@@ -0,0 +1,18 @@
[package]
name = "ruvector-temporal-tensor"
version.workspace = true
edition.workspace = true
rust-version.workspace = true
license.workspace = true
authors.workspace = true
repository.workspace = true
description = "Temporal tensor compression with tiered quantization for RuVector"
[features]
default = []
ffi = [] # Enable WASM/C FFI exports
simd = [] # Enable SIMD-accelerated quantization (future)
persistence = [] # Enable disk-backed BlockIO and MetaLog (uses std::fs)
[lib]
crate-type = ["lib"]

View File

@@ -0,0 +1,279 @@
# ruvector-temporal-tensor
[![Crates.io](https://img.shields.io/crates/v/ruvector-temporal-tensor.svg)](https://crates.io/crates/ruvector-temporal-tensor)
[![Documentation](https://docs.rs/ruvector-temporal-tensor/badge.svg)](https://docs.rs/ruvector-temporal-tensor)
[![License: MIT](https://img.shields.io/badge/License-MIT-blue.svg)](https://opensource.org/licenses/MIT)
[![Rust](https://img.shields.io/badge/rust-1.77%2B-orange.svg)](https://www.rust-lang.org)
**Shrink your vector data 4-10x without losing the signal.**
`ruvector-temporal-tensor` compresses streams of floating-point tensors by exploiting two properties that most vector workloads share:
1. **Values within a group are similar** — so a single scale factor per group captures the range, and a small integer code captures the value. This is *groupwise symmetric quantization*.
2. **Consecutive frames barely change** — so the same scale factors can be reused across many frames until the data drifts. This is *temporal segment reuse*.
The crate automatically picks the right bit-width based on how "hot" (frequently accessed) the tensor is, giving you aggressive compression on cold data while preserving accuracy on hot data.
Zero external dependencies. Compiles to WASM. Ships with a C FFI.
## How It Works
```
f32 frame ──► tier policy ──► quantizer ──► bitpack ──► segment blob
"How hot is this tensor?"
Hot → 8-bit (lossless-ish)
Warm → 7 or 5-bit
Cold → 3-bit (10x smaller)
```
Each frame of `f32` values is divided into fixed-size groups (default 64). Per group, the compressor computes a single scale factor (`max_abs / qmax`) and maps every value to a signed integer code. Codes are packed into a tight bitstream with no byte-alignment waste.
When the next frame arrives, the compressor checks whether the existing scale factors still cover the new data (within a configurable drift tolerance). If they do, the frame is appended to the current **segment** — reusing the same scales. If they don't, the segment is finalized and a new one starts.
Segments are self-contained binary blobs with a 22-byte header, the f16-encoded scales, and the packed data. They can be decoded independently, or you can random-access a single frame by index.
## Compression Ratios
| Tier | Bits | Ratio vs f32 | Typical Error | When Used |
|------|------|-------------|---------------|-----------|
| Hot | 8 | **~4x** | < 0.5% | Frequently accessed tensors |
| Warm | 7 | **~4.6x** | < 1% | Moderate access patterns |
| Warm | 5 | **~6.4x** | < 3% | Aggressively compressed warm data |
| Cold | 3 | **~10.7x** | < 15% | Rarely accessed / archival |
Ratios improve further with temporal reuse — the scale overhead is amortized across all frames in a segment.
## Quick Start
Add to your `Cargo.toml`:
```toml
[dependencies]
ruvector-temporal-tensor = "2.0"
```
### Compress and decompress
```rust
use ruvector_temporal_tensor::{TemporalTensorCompressor, TierPolicy};
// 1. Create a compressor for 128-element tensors
let mut comp = TemporalTensorCompressor::new(TierPolicy::default(), 128, 0);
comp.set_access(100, 0); // mark as hot → 8-bit quantization
let frame = vec![1.0f32; 128];
let mut segment = Vec::new();
// 2. Push frames — segment stays empty until a boundary is crossed
comp.push_frame(&frame, 1, &mut segment);
// 3. Force-emit the current segment
comp.flush(&mut segment);
// 4. Decode back to f32
let mut decoded = Vec::new();
ruvector_temporal_tensor::segment::decode(&segment, &mut decoded);
assert_eq!(decoded.len(), 128);
```
### Stream many frames
```rust
use ruvector_temporal_tensor::{TemporalTensorCompressor, TierPolicy};
let mut comp = TemporalTensorCompressor::new(TierPolicy::default(), 512, 0);
comp.set_access(100, 0);
let mut segments: Vec<Vec<u8>> = Vec::new();
let mut seg = Vec::new();
for t in 0..1000 {
let frame: Vec<f32> = (0..512).map(|i| ((i + t) as f32 * 0.01).sin()).collect();
comp.push_frame(&frame, t as u32, &mut seg);
if !seg.is_empty() {
segments.push(seg.clone());
}
}
comp.flush(&mut seg);
if !seg.is_empty() {
segments.push(seg);
}
```
### Random-access a single frame
```rust
use ruvector_temporal_tensor::segment;
# use ruvector_temporal_tensor::{TemporalTensorCompressor, TierPolicy};
# let mut comp = TemporalTensorCompressor::new(TierPolicy::default(), 64, 0);
# let mut seg = Vec::new();
# comp.push_frame(&vec![1.0f32; 64], 0, &mut seg);
# comp.flush(&mut seg);
// Decode only frame 0 — skips all other frames in the segment
let values = segment::decode_single_frame(&seg, 0).unwrap();
assert_eq!(values.len(), 64);
// Check compression ratio
let ratio = segment::compression_ratio(&seg);
assert!(ratio > 1.0);
```
### Custom tier policy
```rust
use ruvector_temporal_tensor::{TemporalTensorCompressor, TierPolicy};
let policy = TierPolicy {
hot_min_score: 512, // score threshold for 8-bit
warm_min_score: 64, // score threshold for warm tier
warm_bits: 5, // use 5-bit instead of default 7 for warm
drift_pct_q8: 26, // ~10% drift tolerance (Q8 fixed-point)
group_len: 32, // smaller groups = more scales, tighter fit
};
let mut comp = TemporalTensorCompressor::new(policy, 256, 0);
```
## Feature Flags
```toml
[dependencies]
ruvector-temporal-tensor = { version = "2.0", features = ["ffi"] }
```
| Feature | Default | Description |
|---------|---------|-------------|
| `ffi` | off | Enable `extern "C"` exports for WASM and C interop |
| `simd` | off | Reserved for future SIMD-accelerated quantization |
## API Reference
### Core Types
| Type | Description |
|------|-------------|
| `TemporalTensorCompressor` | Main entry point — push frames, get segments |
| `TierPolicy` | Controls bit-width selection and drift tolerance |
### Compressor Methods
| Method | Description |
|--------|-------------|
| `new(policy, len, now_ts)` | Create a compressor for tensors of `len` elements |
| `push_frame(frame, now_ts, out)` | Compress a frame; emits a segment on boundary crossings |
| `flush(out)` | Force-emit the current segment |
| `touch(now_ts)` | Record an access event (increments count + updates timestamp) |
| `set_access(count, ts)` | Set access stats directly (for restoring state) |
| `active_bits()` | Current quantization bit-width |
| `active_frame_count()` | Frames buffered in the current segment |
| `len()` / `is_empty()` | Tensor length |
### Segment Functions
| Function | Description |
|----------|-------------|
| `segment::decode(bytes, out)` | Decode all frames from a segment |
| `segment::decode_single_frame(bytes, idx)` | Decode one frame by index |
| `segment::parse_header(bytes)` | Read segment metadata without decoding |
| `segment::compression_ratio(bytes)` | Compute raw-to-compressed ratio |
| `segment::encode(...)` | Low-level segment encoder (used internally) |
### Low-Level Modules
| Module | Description |
|--------|-------------|
| `quantizer` | Groupwise symmetric quantization and dequantization |
| `bitpack` | Arbitrary-width bitstream packer and unpacker |
| `f16` | Software IEEE 754 half-precision conversion |
| `tier_policy` | Access-pattern scoring and bit-width selection |
## Segment Binary Format
Segments are self-contained, portable, and version-tagged:
```
Offset Size Field
────── ──── ─────────────────
0 4 Magic: 0x43545154 ("TQTC")
4 1 Version (currently 1)
5 1 Bits per code (3, 5, 7, or 8)
6 4 Group length
10 4 Tensor length (elements per frame)
14 4 Frame count
18 4 Scale count (S)
22 2*S Scales (f16, little-endian)
22+2S 4 Data length (D)
26+2S D Packed quantization codes
```
## FFI / WASM Usage
Enable the `ffi` feature and compile with `--target wasm32-unknown-unknown`:
```bash
cargo build --release --target wasm32-unknown-unknown --features ffi
```
Exported C functions:
| Function | Description |
|----------|-------------|
| `ttc_create(len, now_ts, out_handle)` | Create compressor, get handle |
| `ttc_create_with_policy(...)` | Create with custom tier policy |
| `ttc_free(handle)` | Free a compressor |
| `ttc_touch(handle, now_ts)` | Record access |
| `ttc_set_access(handle, count, ts)` | Set access stats |
| `ttc_push_frame(handle, ts, in, len, out, cap, written)` | Compress a frame |
| `ttc_flush(handle, out, cap, written)` | Flush current segment |
| `ttc_decode_segment(seg, len, out, cap, written)` | Decode a segment |
| `ttc_alloc(size, out_ptr)` | Allocate WASM linear memory |
| `ttc_dealloc(ptr, cap)` | Free allocated memory |
## Design Decisions
See **[ADR-017](../../docs/adr/ADR-017-temporal-tensor-compression.md)** for the full architecture decision record, including SOTA survey, compression math, safety analysis, and integration guidance.
Key decisions:
- **Groupwise symmetric** (no zero-point) — simpler, faster, well-suited for normally-distributed embeddings
- **f16 scales** — 2 bytes per group vs 4 for f32, with negligible accuracy loss
- **64-bit bitstream accumulator** — handles any sub-byte width without byte-alignment waste
- **Score-based tiering** — `access_count * 1024 / age` balances recency and frequency
- **~10% drift tolerance** — Q8 fixed-point configurable, default 26/256
## Building and Testing
```bash
# Build
cargo build -p ruvector-temporal-tensor --release
# Run all tests (41 unit + 3 doc-tests)
cargo test -p ruvector-temporal-tensor
# Clippy
cargo clippy -p ruvector-temporal-tensor -- -W clippy::all
# Build WASM target
cargo build -p ruvector-temporal-tensor --release --target wasm32-unknown-unknown --features ffi
```
## Related Crates
| Crate | Relationship |
|-------|-------------|
| [ruvector-core](../ruvector-core/) | Parent vector database engine; temporal tensors integrate as a storage backend |
| [ruvector-temporal-tensor-wasm](../ruvector-temporal-tensor-wasm/) | Thin WASM re-export wrapper |
## License
MIT License — see [LICENSE](../../LICENSE) for details.
---
<div align="center">
**Part of [Ruvector](https://github.com/ruvnet/ruvector)**
</div>

View File

@@ -0,0 +1,831 @@
//! AgentDB adapter for pattern-aware tiering.
//!
//! Provides a bridge between the TieredStore and an external HNSW
//! vector index. When connected, tiering decisions can be influenced
//! by semantic similarity to frequently-accessed patterns.
//!
//! # Overview
//!
//! Block metadata is converted into a compact 4-dimensional embedding
//! via [`pattern_from_meta`], then stored in a [`PatternIndex`]. The
//! [`AdaptiveTiering`] struct combines the index with a
//! [`TierConfig`](crate::tiering::TierConfig) to produce tier
//! suggestions based on weighted neighbor voting.
//!
//! The default [`InMemoryPatternIndex`] uses brute-force linear scan
//! with cosine similarity, suitable for up to ~10K blocks. A real
//! deployment would swap in an HNSW-backed implementation.
use crate::store::{BlockKey, BlockMeta, Tier};
use crate::tiering::TierConfig;
use std::collections::HashMap;
// ---------------------------------------------------------------------------
// PatternVector
// ---------------------------------------------------------------------------
/// A block's access-pattern embedding for similarity search.
#[derive(Clone, Debug)]
pub struct PatternVector {
/// The block this vector represents.
pub key: BlockKey,
/// Access-pattern embedding (typically 4 dimensions).
pub embedding: Vec<f32>,
/// Tiering score at the time of insertion.
pub score: f32,
}
// ---------------------------------------------------------------------------
// PatternIndex trait
// ---------------------------------------------------------------------------
/// Trait for a vector index over access-pattern embeddings.
///
/// Implementations range from a simple brute-force scan
/// ([`InMemoryPatternIndex`]) to an HNSW-backed production index.
pub trait PatternIndex {
/// Insert (or replace) a pattern vector.
fn insert(&mut self, vec: &PatternVector);
/// Return the `k` nearest neighbors to `query`, sorted by
/// descending cosine similarity. Each result is `(key, similarity)`.
fn search_nearest(&self, query: &[f32], k: usize) -> Vec<(BlockKey, f32)>;
/// Remove the pattern for `key`, if present.
fn remove(&mut self, key: BlockKey);
/// Number of pattern vectors currently stored.
fn len(&self) -> usize;
/// Returns `true` if the index contains no vectors.
fn is_empty(&self) -> bool {
self.len() == 0
}
}
// ---------------------------------------------------------------------------
// Cosine similarity
// ---------------------------------------------------------------------------
/// Compute the cosine similarity between two vectors.
///
/// Returns 0.0 if either vector has zero magnitude or they differ in length.
fn cosine_similarity(a: &[f32], b: &[f32]) -> f32 {
if a.len() != b.len() || a.is_empty() {
return 0.0;
}
let mut dot = 0.0f32;
let mut norm_a_sq = 0.0f32;
let mut norm_b_sq = 0.0f32;
for (&x, &y) in a.iter().zip(b.iter()) {
dot += x * y;
norm_a_sq += x * x;
norm_b_sq += y * y;
}
let denom = norm_a_sq.sqrt() * norm_b_sq.sqrt();
if denom == 0.0 {
0.0
} else {
dot / denom
}
}
// ---------------------------------------------------------------------------
// InMemoryPatternIndex
// ---------------------------------------------------------------------------
/// Brute-force in-memory implementation of [`PatternIndex`].
///
/// Uses a `Vec<PatternVector>` with linear-scan cosine similarity.
/// Adequate for small collections (<10K blocks); a real AgentDB
/// deployment would use HNSW for sub-linear search.
pub struct InMemoryPatternIndex {
vectors: Vec<PatternVector>,
}
impl InMemoryPatternIndex {
/// Create a new empty index.
pub fn new() -> Self {
Self {
vectors: Vec::new(),
}
}
}
impl Default for InMemoryPatternIndex {
fn default() -> Self {
Self::new()
}
}
impl PatternIndex for InMemoryPatternIndex {
fn insert(&mut self, vec: &PatternVector) {
// Remove any existing entry for the same key, then append.
self.vectors.retain(|v| v.key != vec.key);
self.vectors.push(vec.clone());
}
fn search_nearest(&self, query: &[f32], k: usize) -> Vec<(BlockKey, f32)> {
if k == 0 || self.vectors.is_empty() {
return Vec::new();
}
let mut scored: Vec<(BlockKey, f32)> = self
.vectors
.iter()
.map(|v| (v.key, cosine_similarity(query, &v.embedding)))
.collect();
// Sort by descending similarity.
scored.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(core::cmp::Ordering::Equal));
scored.truncate(k);
scored
}
fn remove(&mut self, key: BlockKey) {
self.vectors.retain(|v| v.key != key);
}
fn len(&self) -> usize {
self.vectors.len()
}
}
// ---------------------------------------------------------------------------
// pattern_from_meta
// ---------------------------------------------------------------------------
/// Convert block metadata into a 4-dimensional pattern vector.
///
/// The dimensions encode access-pattern features that are useful for
/// clustering blocks with similar tiering behaviour:
///
/// | Index | Feature | Range | Description |
/// |-------|------------------|---------|------------------------------------------|
/// | 0 | `ema_rate` | [0, 1] | Exponential moving average of access rate|
/// | 1 | `popcount/64` | [0, 1] | Fraction of recent ticks with access |
/// | 2 | `recency_decay` | (0, 1] | `1 / (1 + tier_age)` -- inverse staleness|
/// | 3 | `access_count_log` | [0, 1] | `log2(1 + count) / 32` -- normalized log |
pub fn pattern_from_meta(meta: &BlockMeta) -> Vec<f32> {
let ema = meta.ema_rate.clamp(0.0, 1.0);
let pop = meta.window.count_ones() as f32 / 64.0;
let recency = 1.0 / (1.0 + meta.tier_age as f32);
let count_log = ((1.0 + meta.access_count as f32).log2() / 32.0).clamp(0.0, 1.0);
vec![ema, pop, recency, count_log]
}
// ---------------------------------------------------------------------------
// AdaptiveTiering
// ---------------------------------------------------------------------------
/// Pattern-aware tiering advisor.
///
/// Combines a [`PatternIndex`] with a [`TierConfig`] to suggest tier
/// assignments based on the tiers of semantically similar blocks.
///
/// # Algorithm
///
/// Given a block's metadata and a set of nearest neighbors (from the
/// pattern index), each neighbor's known tier contributes a weighted
/// vote proportional to its cosine similarity. The tier with the
/// highest cumulative vote is suggested, unless it matches the block's
/// current tier (in which case `None` is returned).
pub struct AdaptiveTiering<I: PatternIndex> {
/// The underlying pattern vector index.
pub index: I,
/// Tiering configuration (thresholds, hysteresis, etc.).
pub config: TierConfig,
/// Known tier for each block, updated via [`register_block`].
block_tiers: HashMap<BlockKey, Tier>,
}
impl<I: PatternIndex> AdaptiveTiering<I> {
/// Create a new `AdaptiveTiering` with the given index and config.
pub fn new(index: I, config: TierConfig) -> Self {
Self {
index,
config,
block_tiers: HashMap::new(),
}
}
/// Register (or update) the known tier for a block.
///
/// This must be called whenever a block changes tier so that
/// [`suggest_tier`](Self::suggest_tier) can use accurate neighbor
/// tier information for voting.
pub fn register_block(&mut self, key: BlockKey, tier: Tier) {
self.block_tiers.insert(key, tier);
}
/// Remove a block from the tier registry and the pattern index.
pub fn remove_block(&mut self, key: BlockKey) {
self.block_tiers.remove(&key);
self.index.remove(key);
}
/// Number of blocks registered in the tier map.
pub fn registered_count(&self) -> usize {
self.block_tiers.len()
}
/// Suggest a tier for `meta` based on its nearest neighbors.
///
/// `neighbors` should be the output of
/// [`PatternIndex::search_nearest`]: a list of `(BlockKey, similarity)`
/// pairs. Each neighbor whose tier is known contributes a weighted
/// vote. The tier with the highest total vote is returned, unless it
/// matches the block's current tier.
///
/// Returns `None` if:
/// - `neighbors` is empty,
/// - no neighbors have known tiers, or
/// - the consensus tier matches the block's current tier.
pub fn suggest_tier(&self, meta: &BlockMeta, neighbors: &[(BlockKey, f32)]) -> Option<Tier> {
if neighbors.is_empty() {
return None;
}
// Accumulate weighted votes per tier.
// Index 0 = Tier0, 1 = Tier1, 2 = Tier2, 3 = Tier3.
let mut votes = [0.0f32; 4];
let mut total_weight = 0.0f32;
for &(key, similarity) in neighbors {
if let Some(&tier) = self.block_tiers.get(&key) {
let weight = similarity.max(0.0);
votes[tier as u8 as usize] += weight;
total_weight += weight;
}
}
if total_weight == 0.0 {
return None;
}
// Find the tier with the highest vote. On ties, prefer the
// hotter tier (lower index) since it was found first.
let mut best_idx = 0usize;
let mut best_vote = votes[0];
for i in 1..4 {
if votes[i] > best_vote {
best_vote = votes[i];
best_idx = i;
}
}
let suggested = match best_idx {
0 => Tier::Tier0,
1 => Tier::Tier1,
2 => Tier::Tier2,
3 => Tier::Tier3,
_ => unreachable!(),
};
if suggested == meta.tier {
None
} else {
Some(suggested)
}
}
}
// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------
#[cfg(test)]
mod tests {
use super::*;
use crate::store::{DType, ReconstructPolicy};
fn make_key(tid: u128, idx: u32) -> BlockKey {
BlockKey {
tensor_id: tid,
block_index: idx,
}
}
fn make_store_meta(
key: BlockKey,
tier: Tier,
ema_rate: f32,
window: u64,
access_count: u32,
tier_age: u32,
) -> BlockMeta {
BlockMeta {
key,
dtype: DType::F32,
tier,
bits: 8,
scale: 1.0,
zero_point: 0,
created_at: 0,
last_access_at: 100,
access_count,
ema_rate,
window,
checksum: 0,
reconstruct: ReconstructPolicy::None,
tier_age,
lineage_parent: None,
block_bytes: 1024,
}
}
// -- cosine_similarity -------------------------------------------------
#[test]
fn cosine_identical_vectors() {
let v = vec![1.0, 2.0, 3.0, 4.0];
let sim = cosine_similarity(&v, &v);
assert!((sim - 1.0).abs() < 1e-6, "sim={sim}");
}
#[test]
fn cosine_orthogonal_vectors() {
let a = vec![1.0, 0.0];
let b = vec![0.0, 1.0];
let sim = cosine_similarity(&a, &b);
assert!(sim.abs() < 1e-6, "sim={sim}");
}
#[test]
fn cosine_opposite_vectors() {
let a = vec![1.0, 0.0, 0.0];
let b = vec![-1.0, 0.0, 0.0];
let sim = cosine_similarity(&a, &b);
assert!((sim - (-1.0)).abs() < 1e-6, "sim={sim}");
}
#[test]
fn cosine_zero_vector() {
let a = vec![1.0, 2.0];
let b = vec![0.0, 0.0];
assert_eq!(cosine_similarity(&a, &b), 0.0);
}
#[test]
fn cosine_different_lengths() {
let a = vec![1.0, 2.0];
let b = vec![1.0, 2.0, 3.0];
assert_eq!(cosine_similarity(&a, &b), 0.0);
}
#[test]
fn cosine_empty() {
let a: Vec<f32> = vec![];
let b: Vec<f32> = vec![];
assert_eq!(cosine_similarity(&a, &b), 0.0);
}
#[test]
fn cosine_known_value() {
// cos([1,1], [1,0]) = 1/sqrt(2) ~ 0.7071
let a = vec![1.0, 1.0];
let b = vec![1.0, 0.0];
let sim = cosine_similarity(&a, &b);
let expected = 1.0 / 2.0f32.sqrt();
assert!(
(sim - expected).abs() < 1e-6,
"sim={sim}, expected={expected}"
);
}
// -- InMemoryPatternIndex ----------------------------------------------
#[test]
fn index_insert_and_len() {
let mut idx = InMemoryPatternIndex::new();
assert!(idx.is_empty());
idx.insert(&PatternVector {
key: make_key(1, 0),
embedding: vec![1.0, 0.0, 0.0, 0.0],
score: 0.5,
});
assert_eq!(idx.len(), 1);
assert!(!idx.is_empty());
}
#[test]
fn index_insert_replaces_duplicate_key() {
let mut idx = InMemoryPatternIndex::new();
let key = make_key(1, 0);
idx.insert(&PatternVector {
key,
embedding: vec![1.0, 0.0, 0.0, 0.0],
score: 0.5,
});
idx.insert(&PatternVector {
key,
embedding: vec![0.0, 1.0, 0.0, 0.0],
score: 0.8,
});
assert_eq!(idx.len(), 1);
// The search should find the updated embedding.
let results = idx.search_nearest(&[0.0, 1.0, 0.0, 0.0], 1);
assert_eq!(results.len(), 1);
assert_eq!(results[0].0, key);
// Similarity should be ~1.0 since embeddings match.
assert!((results[0].1 - 1.0).abs() < 1e-6);
}
#[test]
fn index_remove() {
let mut idx = InMemoryPatternIndex::new();
let key = make_key(1, 0);
idx.insert(&PatternVector {
key,
embedding: vec![1.0, 0.0, 0.0, 0.0],
score: 0.5,
});
assert_eq!(idx.len(), 1);
idx.remove(key);
assert_eq!(idx.len(), 0);
}
#[test]
fn index_remove_nonexistent() {
let mut idx = InMemoryPatternIndex::new();
idx.remove(make_key(99, 0)); // should not panic
assert_eq!(idx.len(), 0);
}
#[test]
fn index_search_nearest_ordering() {
let mut idx = InMemoryPatternIndex::new();
// Insert three vectors with known geometry.
idx.insert(&PatternVector {
key: make_key(1, 0),
embedding: vec![1.0, 0.0, 0.0, 0.0],
score: 0.0,
});
idx.insert(&PatternVector {
key: make_key(2, 0),
embedding: vec![0.7, 0.7, 0.0, 0.0],
score: 0.0,
});
idx.insert(&PatternVector {
key: make_key(3, 0),
embedding: vec![0.0, 1.0, 0.0, 0.0],
score: 0.0,
});
// Query close to [1, 0, 0, 0].
let results = idx.search_nearest(&[1.0, 0.1, 0.0, 0.0], 3);
assert_eq!(results.len(), 3);
// Closest should be key 1 (nearly identical direction).
assert_eq!(results[0].0, make_key(1, 0));
// Second should be key 2 (partial overlap).
assert_eq!(results[1].0, make_key(2, 0));
// Third should be key 3 (mostly orthogonal).
assert_eq!(results[2].0, make_key(3, 0));
// Similarities should be descending.
assert!(results[0].1 >= results[1].1);
assert!(results[1].1 >= results[2].1);
}
#[test]
fn index_search_nearest_k_larger_than_size() {
let mut idx = InMemoryPatternIndex::new();
idx.insert(&PatternVector {
key: make_key(1, 0),
embedding: vec![1.0, 0.0],
score: 0.0,
});
let results = idx.search_nearest(&[1.0, 0.0], 10);
assert_eq!(results.len(), 1);
}
#[test]
fn index_search_nearest_k_zero() {
let mut idx = InMemoryPatternIndex::new();
idx.insert(&PatternVector {
key: make_key(1, 0),
embedding: vec![1.0],
score: 0.0,
});
let results = idx.search_nearest(&[1.0], 0);
assert!(results.is_empty());
}
#[test]
fn index_search_nearest_empty() {
let idx = InMemoryPatternIndex::new();
let results = idx.search_nearest(&[1.0, 0.0], 5);
assert!(results.is_empty());
}
// -- pattern_from_meta -------------------------------------------------
#[test]
fn pattern_from_meta_dimensions() {
let meta = make_store_meta(make_key(1, 0), Tier::Tier1, 0.5, 0xFFFF, 100, 10);
let pat = pattern_from_meta(&meta);
assert_eq!(pat.len(), 4);
}
#[test]
fn pattern_from_meta_ema_component() {
let meta = make_store_meta(make_key(1, 0), Tier::Tier1, 0.8, 0, 0, 0);
let pat = pattern_from_meta(&meta);
assert!((pat[0] - 0.8).abs() < 1e-6, "ema={}", pat[0]);
}
#[test]
fn pattern_from_meta_popcount_component() {
// All 64 bits set.
let meta = make_store_meta(make_key(1, 0), Tier::Tier1, 0.0, u64::MAX, 0, 0);
let pat = pattern_from_meta(&meta);
assert!((pat[1] - 1.0).abs() < 1e-6, "pop={}", pat[1]);
// No bits set.
let meta2 = make_store_meta(make_key(1, 0), Tier::Tier1, 0.0, 0, 0, 0);
let pat2 = pattern_from_meta(&meta2);
assert!((pat2[1]).abs() < 1e-6, "pop={}", pat2[1]);
// 32 bits set.
let meta3 = make_store_meta(make_key(1, 0), Tier::Tier1, 0.0, 0xFFFF_FFFF, 0, 0);
let pat3 = pattern_from_meta(&meta3);
assert!((pat3[1] - 0.5).abs() < 1e-6, "pop={}", pat3[1]);
}
#[test]
fn pattern_from_meta_recency_component() {
// tier_age = 0 => recency = 1.0 / (1.0 + 0) = 1.0
let meta = make_store_meta(make_key(1, 0), Tier::Tier1, 0.0, 0, 0, 0);
let pat = pattern_from_meta(&meta);
assert!((pat[2] - 1.0).abs() < 1e-6, "recency={}", pat[2]);
// tier_age = 9 => recency = 1.0 / 10.0 = 0.1
let meta2 = make_store_meta(make_key(1, 0), Tier::Tier1, 0.0, 0, 0, 9);
let pat2 = pattern_from_meta(&meta2);
assert!((pat2[2] - 0.1).abs() < 1e-6, "recency={}", pat2[2]);
}
#[test]
fn pattern_from_meta_access_count_log_component() {
// access_count = 0 => log2(1) / 32 = 0
let meta = make_store_meta(make_key(1, 0), Tier::Tier1, 0.0, 0, 0, 0);
let pat = pattern_from_meta(&meta);
assert!(pat[3].abs() < 1e-6, "count_log={}", pat[3]);
// access_count = 1 => log2(2) / 32 = 1/32 ~ 0.03125
let meta2 = make_store_meta(make_key(1, 0), Tier::Tier1, 0.0, 0, 1, 0);
let pat2 = pattern_from_meta(&meta2);
assert!((pat2[3] - 1.0 / 32.0).abs() < 1e-4, "count_log={}", pat2[3]);
}
#[test]
fn pattern_from_meta_values_in_unit_range() {
// Use extreme values to verify clamping.
let meta = make_store_meta(
make_key(1, 0),
Tier::Tier1,
2.0, // ema > 1, should be clamped
u64::MAX, // all bits set
u32::MAX, // max access count
u32::MAX, // max tier age
);
let pat = pattern_from_meta(&meta);
for (i, &v) in pat.iter().enumerate() {
assert!(v >= 0.0 && v <= 1.0, "dim {i} out of [0,1]: {v}");
}
}
// -- AdaptiveTiering ---------------------------------------------------
#[test]
fn adaptive_new_and_register() {
let idx = InMemoryPatternIndex::new();
let config = TierConfig::default();
let mut at = AdaptiveTiering::new(idx, config);
assert_eq!(at.registered_count(), 0);
at.register_block(make_key(1, 0), Tier::Tier1);
assert_eq!(at.registered_count(), 1);
at.register_block(make_key(1, 0), Tier::Tier2);
assert_eq!(at.registered_count(), 1); // same key, updated
}
#[test]
fn adaptive_remove_block() {
let mut idx = InMemoryPatternIndex::new();
let key = make_key(1, 0);
idx.insert(&PatternVector {
key,
embedding: vec![1.0, 0.0, 0.0, 0.0],
score: 0.5,
});
let config = TierConfig::default();
let mut at = AdaptiveTiering::new(idx, config);
at.register_block(key, Tier::Tier1);
assert_eq!(at.registered_count(), 1);
assert_eq!(at.index.len(), 1);
at.remove_block(key);
assert_eq!(at.registered_count(), 0);
assert_eq!(at.index.len(), 0);
}
#[test]
fn suggest_tier_empty_neighbors() {
let idx = InMemoryPatternIndex::new();
let config = TierConfig::default();
let at = AdaptiveTiering::new(idx, config);
let meta = make_store_meta(make_key(1, 0), Tier::Tier1, 0.5, 0, 10, 5);
let result = at.suggest_tier(&meta, &[]);
assert_eq!(result, None);
}
#[test]
fn suggest_tier_no_known_neighbors() {
let idx = InMemoryPatternIndex::new();
let config = TierConfig::default();
let at = AdaptiveTiering::new(idx, config);
let meta = make_store_meta(make_key(1, 0), Tier::Tier1, 0.5, 0, 10, 5);
// Neighbors exist but their tiers are not registered.
let neighbors = vec![(make_key(2, 0), 0.9), (make_key(3, 0), 0.8)];
let result = at.suggest_tier(&meta, &neighbors);
assert_eq!(result, None);
}
#[test]
fn suggest_tier_unanimous_vote() {
let idx = InMemoryPatternIndex::new();
let config = TierConfig::default();
let mut at = AdaptiveTiering::new(idx, config);
// Register three neighbors all in Tier3.
at.register_block(make_key(2, 0), Tier::Tier3);
at.register_block(make_key(3, 0), Tier::Tier3);
at.register_block(make_key(4, 0), Tier::Tier3);
let meta = make_store_meta(make_key(1, 0), Tier::Tier1, 0.5, 0, 10, 5);
let neighbors = vec![
(make_key(2, 0), 0.9),
(make_key(3, 0), 0.8),
(make_key(4, 0), 0.7),
];
let result = at.suggest_tier(&meta, &neighbors);
assert_eq!(result, Some(Tier::Tier3));
}
#[test]
fn suggest_tier_same_as_current_returns_none() {
let idx = InMemoryPatternIndex::new();
let config = TierConfig::default();
let mut at = AdaptiveTiering::new(idx, config);
// Neighbors all in Tier1, same as the block.
at.register_block(make_key(2, 0), Tier::Tier1);
at.register_block(make_key(3, 0), Tier::Tier1);
let meta = make_store_meta(make_key(1, 0), Tier::Tier1, 0.5, 0, 10, 5);
let neighbors = vec![(make_key(2, 0), 0.9), (make_key(3, 0), 0.8)];
let result = at.suggest_tier(&meta, &neighbors);
assert_eq!(result, None);
}
#[test]
fn suggest_tier_weighted_majority() {
let idx = InMemoryPatternIndex::new();
let config = TierConfig::default();
let mut at = AdaptiveTiering::new(idx, config);
// Two neighbors in Tier1 with moderate similarity.
at.register_block(make_key(2, 0), Tier::Tier1);
at.register_block(make_key(3, 0), Tier::Tier1);
// One neighbor in Tier3 with very high similarity.
at.register_block(make_key(4, 0), Tier::Tier3);
let meta = make_store_meta(make_key(1, 0), Tier::Tier2, 0.5, 0, 10, 5);
let neighbors = vec![
(make_key(2, 0), 0.3), // votes Tier1 with weight 0.3
(make_key(3, 0), 0.3), // votes Tier1 with weight 0.3
(make_key(4, 0), 0.9), // votes Tier3 with weight 0.9
];
// Tier1 total = 0.6, Tier3 total = 0.9. Tier3 wins.
let result = at.suggest_tier(&meta, &neighbors);
assert_eq!(result, Some(Tier::Tier3));
}
#[test]
fn suggest_tier_negative_similarity_ignored() {
let idx = InMemoryPatternIndex::new();
let config = TierConfig::default();
let mut at = AdaptiveTiering::new(idx, config);
at.register_block(make_key(2, 0), Tier::Tier3);
at.register_block(make_key(3, 0), Tier::Tier1);
let meta = make_store_meta(make_key(1, 0), Tier::Tier2, 0.5, 0, 10, 5);
let neighbors = vec![
(make_key(2, 0), -0.5), // negative similarity, weight clamped to 0
(make_key(3, 0), 0.5), // positive similarity, votes Tier1
];
// Tier3 gets 0 weight (clamped), Tier1 gets 0.5. Tier1 wins.
let result = at.suggest_tier(&meta, &neighbors);
assert_eq!(result, Some(Tier::Tier1));
}
#[test]
fn suggest_tier_zero_similarity_all() {
let idx = InMemoryPatternIndex::new();
let config = TierConfig::default();
let mut at = AdaptiveTiering::new(idx, config);
at.register_block(make_key(2, 0), Tier::Tier3);
let meta = make_store_meta(make_key(1, 0), Tier::Tier1, 0.5, 0, 10, 5);
let neighbors = vec![(make_key(2, 0), 0.0)];
// Zero similarity means zero weight => total_weight == 0 => None.
let result = at.suggest_tier(&meta, &neighbors);
assert_eq!(result, None);
}
// -- Integration: pattern_from_meta + index + adaptive -----------------
#[test]
fn integration_end_to_end() {
let mut idx = InMemoryPatternIndex::new();
let config = TierConfig::default();
// Create several blocks with different access patterns.
let hot_key = make_key(1, 0);
let warm_key = make_key(2, 0);
let cold_key = make_key(3, 0);
let hot_meta = make_store_meta(hot_key, Tier::Tier1, 0.9, u64::MAX, 1000, 2);
let warm_meta = make_store_meta(warm_key, Tier::Tier2, 0.5, 0xFFFF_FFFF, 100, 10);
let cold_meta = make_store_meta(cold_key, Tier::Tier3, 0.05, 0x0F, 5, 100);
// Build embeddings and insert into index.
let hot_emb = pattern_from_meta(&hot_meta);
let warm_emb = pattern_from_meta(&warm_meta);
let cold_emb = pattern_from_meta(&cold_meta);
idx.insert(&PatternVector {
key: hot_key,
embedding: hot_emb.clone(),
score: 0.9,
});
idx.insert(&PatternVector {
key: warm_key,
embedding: warm_emb.clone(),
score: 0.5,
});
idx.insert(&PatternVector {
key: cold_key,
embedding: cold_emb.clone(),
score: 0.1,
});
let mut at = AdaptiveTiering::new(idx, config);
at.register_block(hot_key, Tier::Tier1);
at.register_block(warm_key, Tier::Tier2);
at.register_block(cold_key, Tier::Tier3);
// Query: a new block with a hot-like pattern.
let new_key = make_key(4, 0);
let new_meta = make_store_meta(new_key, Tier::Tier3, 0.85, u64::MAX, 800, 3);
let new_emb = pattern_from_meta(&new_meta);
let neighbors = at.index.search_nearest(&new_emb, 3);
assert!(!neighbors.is_empty());
let suggestion = at.suggest_tier(&new_meta, &neighbors);
// The new block's pattern is closest to the hot block, so
// the suggestion should be to promote it (away from Tier3).
assert!(
suggestion.is_some(),
"expected a tier suggestion for a hot-like pattern in Tier3"
);
let suggested = suggestion.unwrap();
assert_ne!(suggested, Tier::Tier3, "should not stay cold");
}
}

View File

@@ -0,0 +1,158 @@
//! Bitstream packer/unpacker for arbitrary bit widths (1-8).
//!
//! Uses a 64-bit accumulator for sub-byte codes with no alignment padding.
/// Pack unsigned codes of `bits` width into a byte stream.
///
/// Each code occupies exactly `bits` bits in the output with no alignment
/// padding between codes. A trailing partial byte is emitted if needed.
///
/// For 8-bit codes, writes bytes directly without bit accumulation.
#[inline]
pub fn pack(codes: &[u32], bits: u32, out: &mut Vec<u8>) {
// Fast path: 8-bit codes map 1:1 to bytes.
if bits == 8 {
out.extend(codes.iter().map(|&c| c as u8));
return;
}
let mut acc: u64 = 0;
let mut acc_bits: u32 = 0;
for &code in codes {
acc |= (code as u64) << acc_bits;
acc_bits += bits;
while acc_bits >= 8 {
out.push((acc & 0xFF) as u8);
acc >>= 8;
acc_bits -= 8;
}
}
if acc_bits > 0 {
out.push((acc & 0xFF) as u8);
}
}
/// Unpack `count` unsigned codes of `bits` width from a byte stream.
///
/// Stops early if the data is exhausted before `count` codes are extracted.
///
/// For 8-bit codes, reads bytes directly without bit accumulation.
#[inline]
pub fn unpack(data: &[u8], bits: u32, count: usize, out: &mut Vec<u32>) {
// Fast path: 8-bit codes map 1:1 from bytes.
if bits == 8 {
let n = count.min(data.len());
out.extend(data[..n].iter().map(|&b| b as u32));
return;
}
let mask = (1u64 << bits) - 1;
let mut acc: u64 = 0;
let mut acc_bits: u32 = 0;
let mut byte_idx = 0usize;
let mut decoded = 0usize;
while decoded < count {
while acc_bits < bits && byte_idx < data.len() {
acc |= (data[byte_idx] as u64) << acc_bits;
acc_bits += 8;
byte_idx += 1;
}
if acc_bits < bits {
break;
}
out.push((acc & mask) as u32);
acc >>= bits;
acc_bits -= bits;
decoded += 1;
}
}
/// Compute qmax for a given bit width: `2^(bits-1) - 1`.
///
/// Returns 0 for invalid bit widths (0 or >8).
///
/// | bits | qmax |
/// |------|------|
/// | 8 | 127 |
/// | 7 | 63 |
/// | 5 | 15 |
/// | 3 | 3 |
#[inline]
pub fn qmax_from_bits(bits: u8) -> i32 {
if bits == 0 || bits > 8 {
return 0;
}
(1i32 << (bits - 1)) - 1
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_roundtrip_8bit() {
let codes: Vec<u32> = (0..256).collect();
let mut packed = Vec::new();
pack(&codes, 8, &mut packed);
assert_eq!(packed.len(), 256);
let mut unpacked = Vec::new();
unpack(&packed, 8, 256, &mut unpacked);
assert_eq!(codes, unpacked);
}
#[test]
fn test_roundtrip_3bit() {
let codes: Vec<u32> = (0..7).collect();
let mut packed = Vec::new();
pack(&codes, 3, &mut packed);
let mut unpacked = Vec::new();
unpack(&packed, 3, 7, &mut unpacked);
assert_eq!(codes, unpacked);
}
#[test]
fn test_roundtrip_5bit() {
let codes: Vec<u32> = (0..31).collect();
let mut packed = Vec::new();
pack(&codes, 5, &mut packed);
let mut unpacked = Vec::new();
unpack(&packed, 5, 31, &mut unpacked);
assert_eq!(codes, unpacked);
}
#[test]
fn test_roundtrip_7bit() {
let codes: Vec<u32> = (0..127).collect();
let mut packed = Vec::new();
pack(&codes, 7, &mut packed);
let mut unpacked = Vec::new();
unpack(&packed, 7, 127, &mut unpacked);
assert_eq!(codes, unpacked);
}
#[test]
fn test_packing_density() {
let codes = vec![5u32; 100];
let mut packed = Vec::new();
pack(&codes, 3, &mut packed);
assert_eq!(packed.len(), 38); // ceil(300/8) = 38
}
#[test]
fn test_qmax() {
assert_eq!(qmax_from_bits(8), 127);
assert_eq!(qmax_from_bits(7), 63);
assert_eq!(qmax_from_bits(5), 15);
assert_eq!(qmax_from_bits(3), 3);
assert_eq!(qmax_from_bits(1), 0);
assert_eq!(qmax_from_bits(0), 0);
}
}

View File

@@ -0,0 +1,552 @@
//! Coherence gate: read-after-write validation for the temporal tensor store.
//!
//! Ensures data integrity by verifying that a `get()` immediately after `put()`
//! returns data within the expected quantization error bounds for the tier.
//!
//! # Overview
//!
//! Quantization is lossy -- the error introduced depends on the tier's bit
//! width (8-bit for Tier1, 7-bit for Tier2, 3-bit for Tier3). The coherence
//! gate validates that the round-trip error stays within configurable
//! per-tier bounds, catching silent corruption or encoding bugs.
//!
//! # Epoch Tracking
//!
//! [`EpochTracker`] provides a lightweight write-epoch mechanism so that
//! readers can detect stale data (i.e. data that was overwritten between
//! the time it was read and the time it was consumed).
use std::collections::HashMap;
use crate::store::{BlockKey, StoreError, Tier, TieredStore};
// ---------------------------------------------------------------------------
// CoherenceResult
// ---------------------------------------------------------------------------
/// Outcome of a coherence check.
#[derive(Clone, Debug, PartialEq)]
pub struct CoherenceResult {
/// Maximum relative error observed across all elements.
pub max_error: f32,
/// The tier at which the block is stored.
pub tier: Tier,
/// Whether the observed error is within the configured bound for this tier.
pub passed: bool,
}
// ---------------------------------------------------------------------------
// CoherenceCheck
// ---------------------------------------------------------------------------
/// Per-tier maximum relative error bounds for read-after-write validation.
///
/// After a `put()`, the block is immediately read back and the maximum
/// relative error (per-element `|orig - decoded| / |orig|`) is compared
/// against the bound for the block's current tier.
#[derive(Clone, Debug)]
pub struct CoherenceCheck {
/// Maximum acceptable relative error for each tier, indexed by
/// `Tier as usize`: `[Tier0, Tier1, Tier2, Tier3]`.
///
/// Tier0 (evicted) has no payload, so any read will fail before the
/// error comparison is reached. The bound is set to `f32::MAX` as a
/// sentinel.
pub max_relative_errors: [f32; 4],
}
impl Default for CoherenceCheck {
fn default() -> Self {
Self {
// Tier0: evicted, reads always fail (sentinel value).
// Tier1: 8-bit, very tight bound.
// Tier2: 7-bit, slightly looser.
// Tier3: 3-bit, aggressive quantization allows up to 35% error.
max_relative_errors: [f32::MAX, 0.01, 0.02, 0.35],
}
}
}
impl CoherenceCheck {
/// Create a `CoherenceCheck` with custom per-tier error bounds.
pub fn new(max_relative_errors: [f32; 4]) -> Self {
Self {
max_relative_errors,
}
}
/// Validate read-after-write coherence for a block that was just written.
///
/// Reads the block back from `store`, computes the maximum relative
/// error against `original_data`, and checks whether it falls within
/// the configured bound for the block's tier.
///
/// # Errors
///
/// Returns [`StoreError::BlockNotFound`] if the key does not exist,
/// [`StoreError::TensorEvicted`] if the block is in Tier0, or any
/// other `StoreError` from the underlying read.
pub fn check_coherence(
&self,
store: &mut TieredStore,
key: BlockKey,
original_data: &[f32],
now: u64,
) -> Result<CoherenceResult, StoreError> {
// Look up the tier before reading (needed for the error bound).
let tier = store.meta(key).ok_or(StoreError::BlockNotFound)?.tier;
// Read back the block.
let mut buf = vec![0.0f32; original_data.len()];
let n = store.get(key, &mut buf, now)?;
// Compute the maximum relative error.
let max_error = compute_max_relative_error(original_data, &buf[..n]);
let tier_idx = tier as usize;
let bound = if tier_idx < self.max_relative_errors.len() {
self.max_relative_errors[tier_idx]
} else {
f32::MAX
};
Ok(CoherenceResult {
max_error,
tier,
passed: max_error <= bound,
})
}
/// Convenience: `put` followed by `check_coherence` in one call.
///
/// Stores the data at the given tier, then immediately reads it back
/// and validates the round-trip error. Returns the coherence result
/// so the caller can decide whether to retry at a higher-fidelity tier.
///
/// # Errors
///
/// Propagates errors from both `put` and the subsequent `get`.
pub fn verify_put(
&self,
store: &mut TieredStore,
key: BlockKey,
data: &[f32],
tier: Tier,
now: u64,
) -> Result<CoherenceResult, StoreError> {
store.put(key, data, tier, now)?;
self.check_coherence(store, key, data, now)
}
}
// ---------------------------------------------------------------------------
// Helper: relative error computation
// ---------------------------------------------------------------------------
/// Compute the maximum element-wise relative error between `original` and
/// `decoded`.
///
/// For elements where `|original| < epsilon` (near-zero), the absolute
/// error is used directly to avoid division-by-zero amplification.
fn compute_max_relative_error(original: &[f32], decoded: &[f32]) -> f32 {
const EPSILON: f32 = 1e-6;
let len = original.len().min(decoded.len());
let mut max_err: f32 = 0.0;
for i in 0..len {
let orig = original[i];
let dec = decoded[i];
let abs_err = (orig - dec).abs();
let rel_err = if orig.abs() > EPSILON {
abs_err / orig.abs()
} else {
abs_err
};
if rel_err > max_err {
max_err = rel_err;
}
}
max_err
}
// ---------------------------------------------------------------------------
// EpochTracker
// ---------------------------------------------------------------------------
/// Monotonic write-epoch tracker keyed by [`BlockKey`].
///
/// Each call to [`record_write`](EpochTracker::record_write) increments a
/// global counter and associates the new epoch with the given key. Readers
/// can later check whether their snapshot is stale via
/// [`is_stale`](EpochTracker::is_stale).
#[derive(Clone, Debug)]
pub struct EpochTracker {
/// Global monotonically increasing write counter.
next_epoch: u64,
/// Per-key latest write epoch.
epochs: HashMap<BlockKey, u64>,
}
impl EpochTracker {
/// Create a new tracker with epoch starting at 1.
pub fn new() -> Self {
Self {
next_epoch: 1,
epochs: HashMap::new(),
}
}
/// Record a write for `key`, returning the new epoch number.
///
/// The epoch is strictly monotonically increasing across all keys.
pub fn record_write(&mut self, key: BlockKey) -> u64 {
let epoch = self.next_epoch;
self.next_epoch += 1;
self.epochs.insert(key, epoch);
epoch
}
/// Return the latest write epoch for `key`, if any write has been recorded.
pub fn check_epoch(&self, key: BlockKey) -> Option<u64> {
self.epochs.get(&key).copied()
}
/// Returns `true` if the block identified by `key` has been written
/// after `read_epoch`, meaning the reader's snapshot is stale.
///
/// Returns `false` if no write has been recorded for `key` (the key
/// does not exist in the tracker).
pub fn is_stale(&self, key: BlockKey, read_epoch: u64) -> bool {
match self.epochs.get(&key) {
Some(&write_epoch) => write_epoch > read_epoch,
None => false,
}
}
}
impl Default for EpochTracker {
fn default() -> Self {
Self::new()
}
}
// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------
#[cfg(test)]
mod tests {
use super::*;
use crate::store::{BlockKey, Tier, TieredStore};
fn make_key(tid: u128, idx: u32) -> BlockKey {
BlockKey {
tensor_id: tid,
block_index: idx,
}
}
// -- CoherenceCheck -----------------------------------------------------
#[test]
fn test_coherence_check_default_bounds() {
let cc = CoherenceCheck::default();
assert_eq!(cc.max_relative_errors[0], f32::MAX);
assert!((cc.max_relative_errors[1] - 0.01).abs() < 1e-9);
assert!((cc.max_relative_errors[2] - 0.02).abs() < 1e-9);
assert!((cc.max_relative_errors[3] - 0.35).abs() < 1e-9);
}
#[test]
fn test_coherence_check_custom_bounds() {
let bounds = [0.0, 0.05, 0.10, 0.50];
let cc = CoherenceCheck::new(bounds);
assert_eq!(cc.max_relative_errors, bounds);
}
#[test]
fn test_check_coherence_tier1_passes() {
let mut store = TieredStore::new(4096);
let key = make_key(1, 0);
let data: Vec<f32> = (0..64).map(|i| (i as f32 + 1.0) * 0.25).collect();
store.put(key, &data, Tier::Tier1, 0).unwrap();
let cc = CoherenceCheck::default();
let result = cc.check_coherence(&mut store, key, &data, 1).unwrap();
assert_eq!(result.tier, Tier::Tier1);
assert!(
result.passed,
"Tier1 coherence should pass; max_error={}, bound={}",
result.max_error, cc.max_relative_errors[1],
);
assert!(
result.max_error < cc.max_relative_errors[1],
"max_error {} should be < bound {}",
result.max_error,
cc.max_relative_errors[1],
);
}
#[test]
fn test_check_coherence_tier3_passes() {
let mut store = TieredStore::new(4096);
let key = make_key(2, 0);
// Use values with large magnitude to keep relative error low under
// 3-bit quantization (only 7 levels). Avoid near-zero values where
// even small absolute error produces large relative error.
let data: Vec<f32> = (0..32).map(|i| 10.0 + (i as f32) * 0.1).collect();
store.put(key, &data, Tier::Tier3, 0).unwrap();
let cc = CoherenceCheck::default();
let result = cc.check_coherence(&mut store, key, &data, 1).unwrap();
assert_eq!(result.tier, Tier::Tier3);
assert!(
result.passed,
"Tier3 coherence should pass with default 0.35 bound; max_error={}",
result.max_error,
);
}
#[test]
fn test_check_coherence_missing_block() {
let mut store = TieredStore::new(4096);
let key = make_key(99, 0);
let data = vec![1.0f32; 8];
let cc = CoherenceCheck::default();
let err = cc.check_coherence(&mut store, key, &data, 0);
assert_eq!(err, Err(StoreError::BlockNotFound));
}
#[test]
fn test_check_coherence_evicted_block() {
use crate::store::ReconstructPolicy;
let mut store = TieredStore::new(4096);
let key = make_key(3, 0);
let data = vec![1.0f32; 16];
store.put(key, &data, Tier::Tier1, 0).unwrap();
store.evict(key, ReconstructPolicy::None).unwrap();
let cc = CoherenceCheck::default();
let err = cc.check_coherence(&mut store, key, &data, 1);
assert_eq!(err, Err(StoreError::TensorEvicted));
}
#[test]
fn test_check_coherence_tight_bound_fails() {
let mut store = TieredStore::new(4096);
let key = make_key(4, 0);
// Data with large dynamic range to maximize quantization error.
let data: Vec<f32> = (0..64).map(|i| (i as f32 - 32.0) * 10.0).collect();
// Store at Tier3 (3-bit) for maximum quantization error.
store.put(key, &data, Tier::Tier3, 0).unwrap();
// Use an extremely tight bound that 3-bit quantization cannot meet.
let cc = CoherenceCheck::new([f32::MAX, 0.001, 0.001, 0.001]);
let result = cc.check_coherence(&mut store, key, &data, 1).unwrap();
assert_eq!(result.tier, Tier::Tier3);
assert!(
!result.passed,
"Tier3 with 0.001 bound should fail; max_error={}",
result.max_error,
);
}
// -- verify_put ---------------------------------------------------------
#[test]
fn test_verify_put_tier1() {
let mut store = TieredStore::new(4096);
let key = make_key(10, 0);
let data: Vec<f32> = (0..64).map(|i| (i as f32 + 1.0) * 0.1).collect();
let cc = CoherenceCheck::default();
let result = cc
.verify_put(&mut store, key, &data, Tier::Tier1, 0)
.unwrap();
assert_eq!(result.tier, Tier::Tier1);
assert!(result.passed, "verify_put Tier1 should pass");
assert_eq!(store.block_count(), 1);
}
#[test]
fn test_verify_put_tier0_rejected() {
let mut store = TieredStore::new(4096);
let key = make_key(11, 0);
let data = vec![1.0f32; 16];
let cc = CoherenceCheck::default();
let err = cc.verify_put(&mut store, key, &data, Tier::Tier0, 0);
assert_eq!(err, Err(StoreError::InvalidBlock));
}
#[test]
fn test_verify_put_tier2() {
let mut store = TieredStore::new(4096);
let key = make_key(12, 0);
let data: Vec<f32> = (0..64).map(|i| (i as f32 + 1.0) * 0.3).collect();
let cc = CoherenceCheck::default();
let result = cc
.verify_put(&mut store, key, &data, Tier::Tier2, 0)
.unwrap();
assert_eq!(result.tier, Tier::Tier2);
assert!(
result.passed,
"verify_put Tier2 should pass; max_error={}",
result.max_error
);
}
// -- compute_max_relative_error -----------------------------------------
#[test]
fn test_relative_error_identical() {
let a = vec![1.0, 2.0, 3.0];
let b = vec![1.0, 2.0, 3.0];
assert_eq!(compute_max_relative_error(&a, &b), 0.0);
}
#[test]
fn test_relative_error_known() {
let original = vec![10.0, 20.0, 50.0];
let decoded = vec![10.5, 20.0, 48.0];
let err = compute_max_relative_error(&original, &decoded);
// Element 0: |0.5| / 10.0 = 0.05
// Element 1: 0.0
// Element 2: |2.0| / 50.0 = 0.04
assert!((err - 0.05).abs() < 1e-6, "expected 0.05, got {err}");
}
#[test]
fn test_relative_error_near_zero() {
// Near-zero original values should use absolute error.
let original = vec![0.0, 1e-8, 1.0];
let decoded = vec![0.001, 0.0, 1.0];
let err = compute_max_relative_error(&original, &decoded);
// Element 0: |0.001| (absolute, since orig < epsilon)
// Element 1: |1e-8| (absolute, since orig < epsilon)
// Element 2: 0.0
assert!((err - 0.001).abs() < 1e-6, "expected ~0.001, got {err}");
}
#[test]
fn test_relative_error_empty() {
assert_eq!(compute_max_relative_error(&[], &[]), 0.0);
}
#[test]
fn test_relative_error_mismatched_lengths() {
let a = vec![1.0, 2.0, 3.0];
let b = vec![1.0, 2.0];
// Should only compare up to min(len(a), len(b)) = 2 elements.
let err = compute_max_relative_error(&a, &b);
assert_eq!(err, 0.0);
}
// -- EpochTracker -------------------------------------------------------
#[test]
fn test_epoch_tracker_new() {
let tracker = EpochTracker::new();
let key = make_key(1, 0);
assert_eq!(tracker.check_epoch(key), None);
assert!(!tracker.is_stale(key, 0));
}
#[test]
fn test_epoch_tracker_record_write() {
let mut tracker = EpochTracker::new();
let key = make_key(1, 0);
let e1 = tracker.record_write(key);
assert_eq!(e1, 1);
assert_eq!(tracker.check_epoch(key), Some(1));
let e2 = tracker.record_write(key);
assert_eq!(e2, 2);
assert_eq!(tracker.check_epoch(key), Some(2));
}
#[test]
fn test_epoch_tracker_monotonic_across_keys() {
let mut tracker = EpochTracker::new();
let key_a = make_key(1, 0);
let key_b = make_key(2, 0);
let e1 = tracker.record_write(key_a);
let e2 = tracker.record_write(key_b);
let e3 = tracker.record_write(key_a);
assert_eq!(e1, 1);
assert_eq!(e2, 2);
assert_eq!(e3, 3);
assert_eq!(tracker.check_epoch(key_a), Some(3));
assert_eq!(tracker.check_epoch(key_b), Some(2));
}
#[test]
fn test_epoch_tracker_is_stale() {
let mut tracker = EpochTracker::new();
let key = make_key(1, 0);
let epoch = tracker.record_write(key);
assert!(
!tracker.is_stale(key, epoch),
"same epoch should not be stale"
);
assert!(
!tracker.is_stale(key, epoch + 1),
"future epoch should not be stale"
);
// Write again -> epoch advances.
let _e2 = tracker.record_write(key);
assert!(
tracker.is_stale(key, epoch),
"old epoch should now be stale after a new write"
);
}
#[test]
fn test_epoch_tracker_unknown_key_not_stale() {
let tracker = EpochTracker::new();
let key = make_key(99, 0);
assert!(!tracker.is_stale(key, 0));
assert!(!tracker.is_stale(key, u64::MAX));
}
#[test]
fn test_epoch_tracker_multiple_keys_independent() {
let mut tracker = EpochTracker::new();
let key_a = make_key(1, 0);
let key_b = make_key(2, 0);
let ea = tracker.record_write(key_a);
let _eb = tracker.record_write(key_b);
// Writing key_b should not make key_a stale at its own epoch.
assert!(!tracker.is_stale(key_a, ea));
}
#[test]
fn test_epoch_tracker_default_trait() {
let tracker = EpochTracker::default();
assert_eq!(tracker.check_epoch(make_key(1, 0)), None);
}
}

View File

@@ -0,0 +1,342 @@
//! TemporalTensorCompressor: the main entry point.
//!
//! Manages temporal segments, drift detection, and tier transitions.
//! Caches f32-converted scales to avoid repeated f16 conversion in hot paths.
use crate::quantizer;
use crate::segment;
use crate::tier_policy::TierPolicy;
pub struct TemporalTensorCompressor {
policy: TierPolicy,
len: u32,
access_count: u32,
last_access_ts: u32,
active_bits: u8,
active_group_len: usize,
active_scales_f16: Vec<u16>,
active_scales_f32: Vec<f32>, // Cached f32 conversion of scales
active_frames: u32,
active_data: Vec<u8>,
}
impl TemporalTensorCompressor {
/// Create a new compressor for tensors of the given length.
pub fn new(policy: TierPolicy, len: u32, now_ts: u32) -> Self {
let bits = policy.select_bits(0, now_ts, now_ts);
Self {
policy,
len,
access_count: 0,
last_access_ts: now_ts,
active_bits: bits,
active_group_len: policy.group_len.max(1) as usize,
active_scales_f16: Vec::new(),
active_scales_f32: Vec::new(),
active_frames: 0,
active_data: Vec::new(),
}
}
/// Record an access (increments count, updates timestamp).
pub fn touch(&mut self, now_ts: u32) {
self.access_count = self.access_count.wrapping_add(1);
self.last_access_ts = now_ts;
}
/// Set access stats directly (for restoring state).
pub fn set_access(&mut self, access_count: u32, last_access_ts: u32) {
self.access_count = access_count;
self.last_access_ts = last_access_ts;
}
/// Current tier bits.
pub fn active_bits(&self) -> u8 {
self.active_bits
}
/// Number of frames in the current segment.
pub fn active_frame_count(&self) -> u32 {
self.active_frames
}
/// Current policy.
pub fn policy(&self) -> &TierPolicy {
&self.policy
}
/// Tensor length.
pub fn len(&self) -> u32 {
self.len
}
/// Returns `true` if the tensor length is zero.
pub fn is_empty(&self) -> bool {
self.len == 0
}
/// Bytes currently buffered in the active segment data.
pub fn active_data_bytes(&self) -> usize {
self.active_data.len()
}
/// Push a frame. If a segment boundary is crossed, the completed segment
/// bytes are written to `out_segment`. Otherwise `out_segment` is cleared.
pub fn push_frame(&mut self, frame: &[f32], now_ts: u32, out_segment: &mut Vec<u8>) {
out_segment.clear();
if frame.len() != self.len as usize {
return;
}
let desired_bits = self
.policy
.select_bits(self.access_count, self.last_access_ts, now_ts);
let drift_factor = self.policy.drift_factor();
// Use cached f32 scales for drift check (avoids f16 conversion per group)
let need_new_segment = self.active_frames == 0
|| desired_bits != self.active_bits
|| !quantizer::frame_fits_scales_f32(
frame,
&self.active_scales_f32,
self.active_group_len,
self.active_bits,
drift_factor,
);
if need_new_segment {
self.flush(out_segment);
self.active_bits = desired_bits;
self.active_group_len = self.policy.group_len.max(1) as usize;
self.active_scales_f16 =
quantizer::compute_scales(frame, self.active_group_len, self.active_bits);
self.active_scales_f32 = quantizer::scales_to_f32(&self.active_scales_f16);
}
// Use cached f32 scales for quantization (avoids f16 conversion per group)
quantizer::quantize_and_pack_f32(
frame,
&self.active_scales_f32,
self.active_group_len,
self.active_bits,
&mut self.active_data,
);
self.active_frames = self.active_frames.wrapping_add(1);
}
/// Flush the current segment. Writes segment bytes to `out_segment`.
/// Resets internal state for the next segment.
pub fn flush(&mut self, out_segment: &mut Vec<u8>) {
if self.active_frames == 0 {
return;
}
segment::encode(
self.active_bits,
self.active_group_len as u32,
self.len,
self.active_frames,
&self.active_scales_f16,
&self.active_data,
out_segment,
);
self.active_frames = 0;
self.active_scales_f16.clear();
self.active_scales_f32.clear();
self.active_data.clear();
}
}
#[cfg(test)]
mod tests {
use super::*;
fn default_policy() -> TierPolicy {
TierPolicy::default()
}
#[test]
fn test_create_and_push() {
let mut comp = TemporalTensorCompressor::new(default_policy(), 64, 0);
let frame = vec![1.0f32; 64];
let mut seg = Vec::new();
comp.push_frame(&frame, 0, &mut seg);
assert!(seg.is_empty()); // First frame, no completed segment
assert_eq!(comp.active_frame_count(), 1);
}
#[test]
fn test_flush_produces_segment() {
let mut comp = TemporalTensorCompressor::new(default_policy(), 64, 0);
let frame = vec![1.0f32; 64];
let mut seg = Vec::new();
comp.push_frame(&frame, 0, &mut seg);
comp.flush(&mut seg);
assert!(!seg.is_empty());
let mut decoded = Vec::new();
segment::decode(&seg, &mut decoded);
assert_eq!(decoded.len(), 64);
}
#[test]
fn test_tier_transition_flushes() {
let policy = TierPolicy {
hot_min_score: 512,
warm_min_score: 64,
warm_bits: 7,
drift_pct_q8: 26,
group_len: 64,
};
let mut comp = TemporalTensorCompressor::new(policy, 64, 0);
comp.set_access(100, 0); // Hot
let frame = vec![1.0f32; 64];
let mut seg = Vec::new();
comp.push_frame(&frame, 1, &mut seg);
assert_eq!(comp.active_bits(), 8);
// Make it cold
comp.set_access(1, 0);
comp.push_frame(&frame, 10000, &mut seg);
assert!(!seg.is_empty());
assert_eq!(comp.active_bits(), 3);
}
#[test]
fn test_drift_triggers_new_segment() {
let mut comp = TemporalTensorCompressor::new(default_policy(), 64, 0);
let mut seg = Vec::new();
let frame1 = vec![1.0f32; 64];
comp.push_frame(&frame1, 0, &mut seg);
let frame2 = vec![5.0f32; 64];
comp.push_frame(&frame2, 0, &mut seg);
assert!(!seg.is_empty());
}
#[test]
fn test_multi_frame_same_segment() {
let mut comp = TemporalTensorCompressor::new(default_policy(), 64, 0);
let mut seg = Vec::new();
let frame = vec![1.0f32; 64];
comp.push_frame(&frame, 0, &mut seg);
assert!(seg.is_empty());
let frame2 = vec![1.05f32; 64];
comp.push_frame(&frame2, 0, &mut seg);
assert!(seg.is_empty());
assert_eq!(comp.active_frame_count(), 2);
}
#[test]
fn test_full_roundtrip_hot() {
let mut comp = TemporalTensorCompressor::new(default_policy(), 128, 0);
comp.set_access(100, 0);
let frame: Vec<f32> = (0..128).map(|i| (i as f32 - 64.0) * 0.01).collect();
let mut seg = Vec::new();
for _ in 0..10 {
comp.push_frame(&frame, 1, &mut seg);
}
comp.flush(&mut seg);
let mut decoded = Vec::new();
segment::decode(&seg, &mut decoded);
assert_eq!(decoded.len(), 128 * 10);
let max_abs = frame.iter().map(|v| v.abs()).fold(0.0f32, f32::max);
for i in 0..128 {
let err = (decoded[i] - frame[i]).abs();
assert!(
err < max_abs * 0.02,
"i={i} orig={} dec={} err={err}",
frame[i],
decoded[i]
);
}
}
#[test]
fn test_full_roundtrip_cold() {
let mut comp = TemporalTensorCompressor::new(default_policy(), 64, 0);
// Default: access_count=0, cold -> 3-bit
let frame: Vec<f32> = (0..64).map(|i| (i as f32 - 32.0) * 0.1).collect();
let mut seg = Vec::new();
comp.push_frame(&frame, 0, &mut seg);
comp.flush(&mut seg);
let header = segment::parse_header(&seg).unwrap();
assert_eq!(header.bits, 3);
let mut decoded = Vec::new();
segment::decode(&seg, &mut decoded);
assert_eq!(decoded.len(), 64);
let max_abs = frame.iter().map(|v| v.abs()).fold(0.0f32, f32::max);
for (i, (&orig, &dec)) in frame.iter().zip(decoded.iter()).enumerate() {
let err = (orig - dec).abs();
// 3-bit: qmax=3, max relative error ~33%
assert!(err < max_abs * 0.4, "i={i} orig={orig} dec={dec} err={err}");
}
}
#[test]
fn test_wrong_length_frame_rejected() {
let mut comp = TemporalTensorCompressor::new(default_policy(), 64, 0);
let frame = vec![1.0f32; 32];
let mut seg = Vec::new();
comp.push_frame(&frame, 0, &mut seg);
assert_eq!(comp.active_frame_count(), 0);
}
#[test]
fn test_accessor_methods() {
let policy = TierPolicy::default();
let comp = TemporalTensorCompressor::new(policy, 256, 42);
assert_eq!(comp.len(), 256);
assert_eq!(comp.active_frame_count(), 0);
assert_eq!(comp.active_data_bytes(), 0);
assert_eq!(comp.policy().group_len, 64);
}
#[test]
fn test_large_tensor_multi_group() {
let mut comp = TemporalTensorCompressor::new(default_policy(), 512, 0);
comp.set_access(100, 0); // hot -> 8-bit
let frame: Vec<f32> = (0..512).map(|i| ((i as f32) * 0.731).sin()).collect();
let mut seg = Vec::new();
for _ in 0..50 {
comp.push_frame(&frame, 1, &mut seg);
}
comp.flush(&mut seg);
let header = segment::parse_header(&seg).unwrap();
assert_eq!(header.bits, 8);
assert_eq!(header.tensor_len, 512);
assert_eq!(header.frame_count, 50);
assert_eq!(header.scale_count, 8); // 512/64 = 8 groups
let mut decoded = Vec::new();
segment::decode(&seg, &mut decoded);
assert_eq!(decoded.len(), 512 * 50);
// Verify compression ratio
let raw = 512 * 4 * 50;
let compressed = seg.len();
let ratio = raw as f32 / compressed as f32;
assert!(ratio > 3.5, "ratio={ratio:.2}x, expected >3.5x");
}
}

View File

@@ -0,0 +1,531 @@
//! Abstract trait interface for tensor block storage.
//!
//! Defines [`TensorStore`] so that other crates can depend on a thin
//! abstraction rather than the concrete [`crate::store::TieredStore`].
//! An extension trait [`TensorStoreExt`] provides convenience helpers
//! via a blanket implementation for all `TensorStore` implementors.
#![allow(dead_code)]
use crate::store::{BlockKey, BlockMeta, ReconstructPolicy, StoreError, Tier, TieredStore};
// ---------------------------------------------------------------------------
// TensorStore trait
// ---------------------------------------------------------------------------
/// Abstract interface for a tiered tensor block store.
///
/// All methods mirror the public API of [`TieredStore`] so that higher-level
/// crates can interact with the store without depending on the concrete type.
pub trait TensorStore {
/// Quantize `data` at the bit width for `tier` and store the block.
///
/// Replaces any existing block with the same `key`.
fn put(&mut self, key: BlockKey, data: &[f32], tier: Tier, now: u64) -> Result<(), StoreError>;
/// Dequantize the block identified by `key` into `out`.
///
/// Returns the number of f32 elements written.
fn get(&mut self, key: BlockKey, out: &mut [f32], now: u64) -> Result<usize, StoreError>;
/// Update access statistics for `key` at tick `now`.
fn touch(&mut self, key: BlockKey, now: u64);
/// Evict a block to Tier0, preserving metadata with the given policy.
fn evict(&mut self, key: BlockKey, policy: ReconstructPolicy) -> Result<(), StoreError>;
/// Return a reference to the metadata for `key`, if it exists.
fn meta(&self, key: BlockKey) -> Option<&BlockMeta>;
/// Total number of blocks tracked (including Tier0 evicted blocks).
fn block_count(&self) -> usize;
/// Number of blocks currently in the given tier.
fn tier_count(&self, tier: Tier) -> usize;
/// Total bytes of quantized data stored across all active tiers.
fn total_bytes(&self) -> usize;
/// Whether a block with the given key exists in the store.
fn contains(&self, key: BlockKey) -> bool;
/// Capture a read-only snapshot of the store's current state.
fn snapshot(&self) -> TensorStoreSnapshot;
}
// ---------------------------------------------------------------------------
// TensorStore impl for TieredStore
// ---------------------------------------------------------------------------
impl TensorStore for TieredStore {
fn put(&mut self, key: BlockKey, data: &[f32], tier: Tier, now: u64) -> Result<(), StoreError> {
TieredStore::put(self, key, data, tier, now)
}
fn get(&mut self, key: BlockKey, out: &mut [f32], now: u64) -> Result<usize, StoreError> {
TieredStore::get(self, key, out, now)
}
fn touch(&mut self, key: BlockKey, now: u64) {
TieredStore::touch(self, key, now);
}
fn evict(&mut self, key: BlockKey, policy: ReconstructPolicy) -> Result<(), StoreError> {
TieredStore::evict(self, key, policy)
}
fn meta(&self, key: BlockKey) -> Option<&BlockMeta> {
TieredStore::meta(self, key)
}
fn block_count(&self) -> usize {
TieredStore::block_count(self)
}
fn tier_count(&self, tier: Tier) -> usize {
TieredStore::tier_count(self, tier)
}
fn total_bytes(&self) -> usize {
TieredStore::total_bytes(self)
}
fn contains(&self, key: BlockKey) -> bool {
TieredStore::meta(self, key).is_some()
}
fn snapshot(&self) -> TensorStoreSnapshot {
let tier_counts = [
TieredStore::tier_count(self, Tier::Tier0),
TieredStore::tier_count(self, Tier::Tier1),
TieredStore::tier_count(self, Tier::Tier2),
TieredStore::tier_count(self, Tier::Tier3),
];
// Compute per-tier byte totals from the store metrics.
let metrics = TieredStore::metrics(self);
let tier_bytes = [
0, // Tier0 holds no payload data
metrics.tier1_bytes as usize,
metrics.tier2_bytes as usize,
metrics.tier3_bytes as usize,
];
TensorStoreSnapshot {
block_count: TieredStore::block_count(self),
tier_counts,
total_bytes: TieredStore::total_bytes(self),
tier_bytes,
}
}
}
// ---------------------------------------------------------------------------
// TensorStoreSnapshot
// ---------------------------------------------------------------------------
/// Read-only snapshot of the store's current state.
///
/// Captures block counts, byte totals, and per-tier breakdowns at a single
/// point in time. Useful for monitoring, dashboards, and tiering decisions
/// that need a consistent view without holding a borrow on the store.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct TensorStoreSnapshot {
/// Total number of blocks tracked (including evicted Tier0 blocks).
pub block_count: usize,
/// Number of blocks in each tier, indexed as `[Tier0, Tier1, Tier2, Tier3]`.
pub tier_counts: [usize; 4],
/// Total bytes of quantized data across all active tiers.
pub total_bytes: usize,
/// Bytes of quantized data per tier, indexed as `[Tier0, Tier1, Tier2, Tier3]`.
pub tier_bytes: [usize; 4],
}
impl TensorStoreSnapshot {
/// Fraction of total blocks that reside in the given tier.
///
/// Returns 0.0 if the store is empty.
pub fn tier_fraction(&self, tier: Tier) -> f64 {
if self.block_count == 0 {
return 0.0;
}
self.tier_counts[tier as usize] as f64 / self.block_count as f64
}
/// Fraction of total bytes stored in the given tier.
///
/// Returns 0.0 if the store holds no data.
pub fn byte_fraction(&self, tier: Tier) -> f64 {
if self.total_bytes == 0 {
return 0.0;
}
self.tier_bytes[tier as usize] as f64 / self.total_bytes as f64
}
}
// ---------------------------------------------------------------------------
// TensorStoreExt extension trait
// ---------------------------------------------------------------------------
/// Convenience methods available on every [`TensorStore`] implementor.
pub trait TensorStoreExt: TensorStore {
/// Allocate a `Vec<f32>` of length `len` and read the block into it.
///
/// This is a convenience wrapper around [`TensorStore::get`] for callers
/// that do not want to manage the output buffer themselves.
fn get_vec(&mut self, key: BlockKey, len: usize, now: u64) -> Result<Vec<f32>, StoreError>;
/// Store a block in Tier1 (hot, 8-bit quantization).
///
/// Shorthand for `put(key, data, Tier::Tier1, now)`.
fn put_tier1(&mut self, key: BlockKey, data: &[f32], now: u64) -> Result<(), StoreError>;
/// Check whether a block has been evicted to Tier0.
///
/// Returns `false` if the block does not exist.
fn is_evicted(&self, key: BlockKey) -> bool;
}
/// Blanket implementation of [`TensorStoreExt`] for all `TensorStore` types.
impl<T: TensorStore> TensorStoreExt for T {
fn get_vec(&mut self, key: BlockKey, len: usize, now: u64) -> Result<Vec<f32>, StoreError> {
let mut buf = vec![0.0f32; len];
let n = self.get(key, &mut buf, now)?;
buf.truncate(n);
Ok(buf)
}
fn put_tier1(&mut self, key: BlockKey, data: &[f32], now: u64) -> Result<(), StoreError> {
self.put(key, data, Tier::Tier1, now)
}
fn is_evicted(&self, key: BlockKey) -> bool {
self.meta(key)
.map(|m| m.tier == Tier::Tier0)
.unwrap_or(false)
}
}
// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------
#[cfg(test)]
mod tests {
use super::*;
use crate::store::{BlockKey, Tier, TieredStore};
fn make_key(tid: u128, idx: u32) -> BlockKey {
BlockKey {
tensor_id: tid,
block_index: idx,
}
}
// -- TensorStore trait delegation ----------------------------------------
#[test]
fn test_trait_put_get_roundtrip() {
let mut store = TieredStore::new(4096);
let key = make_key(1, 0);
let data: Vec<f32> = (0..64).map(|i| i as f32 * 0.25).collect();
// Use trait method
TensorStore::put(&mut store, key, &data, Tier::Tier1, 0).unwrap();
assert_eq!(TensorStore::block_count(&store), 1);
assert!(TensorStore::contains(&store, key));
let mut out = vec![0.0f32; 64];
let n = TensorStore::get(&mut store, key, &mut out, 1).unwrap();
assert_eq!(n, 64);
for (i, (&orig, &dec)) in data.iter().zip(out.iter()).enumerate() {
let err = (orig - dec).abs();
let tol = if orig.abs() > 0.01 {
orig.abs() * 0.02
} else {
0.15
};
assert!(err < tol, "i={i} orig={orig} dec={dec} err={err}");
}
}
#[test]
fn test_trait_touch_updates_access() {
let mut store = TieredStore::new(4096);
let key = make_key(1, 0);
TensorStore::put(&mut store, key, &[1.0; 16], Tier::Tier1, 0).unwrap();
let meta = TensorStore::meta(&store, key).unwrap();
assert_eq!(meta.access_count, 1);
TensorStore::touch(&mut store, key, 10);
let meta = TensorStore::meta(&store, key).unwrap();
assert_eq!(meta.access_count, 2);
assert_eq!(meta.last_access_at, 10);
}
#[test]
fn test_trait_evict() {
let mut store = TieredStore::new(4096);
let key = make_key(1, 0);
TensorStore::put(&mut store, key, &[1.0; 32], Tier::Tier1, 0).unwrap();
assert_eq!(TensorStore::tier_count(&store, Tier::Tier1), 1);
TensorStore::evict(&mut store, key, ReconstructPolicy::Delta).unwrap();
let meta = TensorStore::meta(&store, key).unwrap();
assert_eq!(meta.tier, Tier::Tier0);
assert_eq!(meta.reconstruct, ReconstructPolicy::Delta);
assert_eq!(TensorStore::tier_count(&store, Tier::Tier0), 1);
assert_eq!(TensorStore::tier_count(&store, Tier::Tier1), 0);
}
#[test]
fn test_trait_contains_false_for_missing() {
let store = TieredStore::new(4096);
assert!(!TensorStore::contains(&store, make_key(99, 0)));
}
#[test]
fn test_trait_total_bytes() {
let mut store = TieredStore::new(4096);
assert_eq!(TensorStore::total_bytes(&store), 0);
TensorStore::put(&mut store, make_key(1, 0), &[1.0; 64], Tier::Tier1, 0).unwrap();
assert!(TensorStore::total_bytes(&store) > 0);
}
// -- TensorStoreSnapshot -------------------------------------------------
#[test]
fn test_snapshot_empty_store() {
let store = TieredStore::new(4096);
let snap = TensorStore::snapshot(&store);
assert_eq!(snap.block_count, 0);
assert_eq!(snap.tier_counts, [0, 0, 0, 0]);
assert_eq!(snap.total_bytes, 0);
assert_eq!(snap.tier_bytes, [0, 0, 0, 0]);
}
#[test]
fn test_snapshot_populated_store() {
let mut store = TieredStore::new(4096);
let data = vec![1.0f32; 32];
TensorStore::put(&mut store, make_key(1, 0), &data, Tier::Tier1, 0).unwrap();
TensorStore::put(&mut store, make_key(2, 0), &data, Tier::Tier1, 0).unwrap();
TensorStore::put(&mut store, make_key(3, 0), &data, Tier::Tier2, 0).unwrap();
TensorStore::put(&mut store, make_key(4, 0), &data, Tier::Tier3, 0).unwrap();
let snap = TensorStore::snapshot(&store);
assert_eq!(snap.block_count, 4);
assert_eq!(snap.tier_counts[0], 0); // Tier0
assert_eq!(snap.tier_counts[1], 2); // Tier1
assert_eq!(snap.tier_counts[2], 1); // Tier2
assert_eq!(snap.tier_counts[3], 1); // Tier3
assert!(snap.total_bytes > 0);
assert!(snap.tier_bytes[1] > 0); // Tier1 bytes
assert!(snap.tier_bytes[2] > 0); // Tier2 bytes
assert!(snap.tier_bytes[3] > 0); // Tier3 bytes
assert_eq!(snap.tier_bytes[0], 0); // Tier0 holds no data
}
#[test]
fn test_snapshot_tier_fraction() {
let mut store = TieredStore::new(4096);
let data = vec![1.0f32; 16];
TensorStore::put(&mut store, make_key(1, 0), &data, Tier::Tier1, 0).unwrap();
TensorStore::put(&mut store, make_key(2, 0), &data, Tier::Tier1, 0).unwrap();
TensorStore::put(&mut store, make_key(3, 0), &data, Tier::Tier2, 0).unwrap();
TensorStore::put(&mut store, make_key(4, 0), &data, Tier::Tier3, 0).unwrap();
let snap = TensorStore::snapshot(&store);
assert!((snap.tier_fraction(Tier::Tier1) - 0.5).abs() < 1e-10);
assert!((snap.tier_fraction(Tier::Tier2) - 0.25).abs() < 1e-10);
assert!((snap.tier_fraction(Tier::Tier3) - 0.25).abs() < 1e-10);
assert!((snap.tier_fraction(Tier::Tier0) - 0.0).abs() < 1e-10);
}
#[test]
fn test_snapshot_tier_fraction_empty() {
let snap = TensorStoreSnapshot {
block_count: 0,
tier_counts: [0; 4],
total_bytes: 0,
tier_bytes: [0; 4],
};
assert_eq!(snap.tier_fraction(Tier::Tier1), 0.0);
}
#[test]
fn test_snapshot_byte_fraction_empty() {
let snap = TensorStoreSnapshot {
block_count: 0,
tier_counts: [0; 4],
total_bytes: 0,
tier_bytes: [0; 4],
};
assert_eq!(snap.byte_fraction(Tier::Tier1), 0.0);
}
#[test]
fn test_snapshot_after_eviction() {
let mut store = TieredStore::new(4096);
let data = vec![1.0f32; 32];
TensorStore::put(&mut store, make_key(1, 0), &data, Tier::Tier1, 0).unwrap();
TensorStore::put(&mut store, make_key(2, 0), &data, Tier::Tier2, 0).unwrap();
TensorStore::evict(&mut store, make_key(1, 0), ReconstructPolicy::None).unwrap();
let snap = TensorStore::snapshot(&store);
assert_eq!(snap.block_count, 2); // metadata preserved
assert_eq!(snap.tier_counts[0], 1); // one evicted
assert_eq!(snap.tier_counts[1], 0); // tier1 now empty
assert_eq!(snap.tier_counts[2], 1); // tier2 still has one
assert_eq!(snap.tier_bytes[0], 0); // evicted holds no data
assert_eq!(snap.tier_bytes[1], 0); // tier1 bytes gone
assert!(snap.tier_bytes[2] > 0); // tier2 bytes remain
}
// -- TensorStoreExt convenience methods ----------------------------------
#[test]
fn test_ext_get_vec() {
let mut store = TieredStore::new(4096);
let key = make_key(1, 0);
let data: Vec<f32> = (0..32).map(|i| i as f32 * 0.5).collect();
TensorStore::put(&mut store, key, &data, Tier::Tier1, 0).unwrap();
let result = TensorStoreExt::get_vec(&mut store, key, 32, 1).unwrap();
assert_eq!(result.len(), 32);
for (i, (&orig, &dec)) in data.iter().zip(result.iter()).enumerate() {
let err = (orig - dec).abs();
let tol = if orig.abs() > 0.01 {
orig.abs() * 0.05
} else {
0.15
};
assert!(err < tol, "i={i} orig={orig} dec={dec} err={err}");
}
}
#[test]
fn test_ext_get_vec_truncates_to_actual() {
let mut store = TieredStore::new(4096);
let key = make_key(1, 0);
TensorStore::put(&mut store, key, &[1.0; 16], Tier::Tier1, 0).unwrap();
// Request a larger buffer than the block contains; vec should be truncated.
let result = TensorStoreExt::get_vec(&mut store, key, 64, 1).unwrap();
assert_eq!(result.len(), 16);
}
#[test]
fn test_ext_get_vec_not_found() {
let mut store = TieredStore::new(4096);
let result = TensorStoreExt::get_vec(&mut store, make_key(99, 0), 16, 0);
assert_eq!(result, Err(StoreError::BlockNotFound));
}
#[test]
fn test_ext_put_tier1() {
let mut store = TieredStore::new(4096);
let key = make_key(1, 0);
let data = vec![2.0f32; 16];
TensorStoreExt::put_tier1(&mut store, key, &data, 0).unwrap();
let meta = TensorStore::meta(&store, key).unwrap();
assert_eq!(meta.tier, Tier::Tier1);
assert_eq!(meta.bits, 8);
}
#[test]
fn test_ext_is_evicted_false_when_active() {
let mut store = TieredStore::new(4096);
let key = make_key(1, 0);
TensorStore::put(&mut store, key, &[1.0; 8], Tier::Tier1, 0).unwrap();
assert!(!TensorStoreExt::is_evicted(&store, key));
}
#[test]
fn test_ext_is_evicted_true_after_evict() {
let mut store = TieredStore::new(4096);
let key = make_key(1, 0);
TensorStore::put(&mut store, key, &[1.0; 8], Tier::Tier1, 0).unwrap();
TensorStore::evict(&mut store, key, ReconstructPolicy::None).unwrap();
assert!(TensorStoreExt::is_evicted(&store, key));
}
#[test]
fn test_ext_is_evicted_false_when_missing() {
let store = TieredStore::new(4096);
assert!(!TensorStoreExt::is_evicted(&store, make_key(99, 0)));
}
// -- Trait object safety check -------------------------------------------
#[test]
fn test_trait_object_usable() {
let mut store = TieredStore::new(4096);
let key = make_key(1, 0);
// Ensure TensorStore can be used as a trait object for the subset
// of methods that are object-safe. Since &BlockMeta borrows prevent
// full dyn dispatch for meta(), we verify the non-borrowing methods.
fn use_store(s: &mut dyn TensorStore) -> usize {
s.block_count()
}
TensorStore::put(&mut store, key, &[1.0; 8], Tier::Tier1, 0).unwrap();
assert_eq!(use_store(&mut store), 1);
}
// -- Integration: mixed trait + ext usage --------------------------------
#[test]
fn test_integration_mixed_usage() {
let mut store = TieredStore::new(4096);
let k1 = make_key(1, 0);
let k2 = make_key(2, 0);
let k3 = make_key(3, 0);
// Insert via ext shorthand and trait method.
TensorStoreExt::put_tier1(&mut store, k1, &[1.0; 32], 0).unwrap();
TensorStore::put(&mut store, k2, &[2.0; 32], Tier::Tier2, 0).unwrap();
TensorStore::put(&mut store, k3, &[3.0; 32], Tier::Tier3, 0).unwrap();
assert_eq!(TensorStore::block_count(&store), 3);
assert!(TensorStore::contains(&store, k1));
assert!(TensorStore::contains(&store, k2));
assert!(TensorStore::contains(&store, k3));
// Evict k3 and verify via ext method.
TensorStore::evict(&mut store, k3, ReconstructPolicy::Delta).unwrap();
assert!(TensorStoreExt::is_evicted(&store, k3));
assert!(!TensorStoreExt::is_evicted(&store, k1));
// Read back via ext.
let v1 = TensorStoreExt::get_vec(&mut store, k1, 32, 10).unwrap();
assert_eq!(v1.len(), 32);
// Snapshot should reflect the current state.
let snap = TensorStore::snapshot(&store);
assert_eq!(snap.block_count, 3);
assert_eq!(snap.tier_counts[0], 1); // k3 evicted
assert_eq!(snap.tier_counts[1], 1); // k1
assert_eq!(snap.tier_counts[2], 1); // k2
assert_eq!(snap.tier_counts[3], 0); // k3 was here but evicted
}
}

View File

@@ -0,0 +1,824 @@
//! Delta compression, delta chains, and reconstruction policies (ADR-021).
//!
//! Sparse delta encoding for incremental tensor updates, bounded-depth delta
//! chain management with automatic compaction, and SVD-based low-rank factor
//! reconstruction. All structures are WASM-safe (no `f64` in hot paths).
use crate::store::StoreError;
#[allow(unused_imports)]
use crate::store::{BlockKey, ReconstructPolicy};
/// Size of the fixed portion of a serialized delta (header + scale).
const DELTA_HEADER_BYTES: usize = 34;
/// Size of a single serialized sparse entry (index: u16 + value: i16).
const DELTA_ENTRY_BYTES: usize = 4;
/// Maximum power-iteration steps per singular component.
const POWER_ITER_MAX: usize = 30;
/// Convergence threshold for power iteration.
const POWER_ITER_EPS: f32 = 1e-10;
/// Header for a delta record.
#[derive(Clone, Debug)]
pub struct DeltaHeader {
pub tensor_id: u128,
pub block_index: u32,
pub base_epoch: u64,
pub nnz: u16,
}
/// A single sparse delta entry: index + quantized value.
#[derive(Clone, Copy, Debug)]
pub struct SparseEntry {
pub index: u16,
pub value: i16,
}
/// Complete delta record: header + sparse entries + scale.
///
/// Actual diff = `entry.value as f32 * delta_scale`.
#[derive(Clone, Debug)]
pub struct DeltaRecord {
pub header: DeltaHeader,
pub delta_scale: f32,
pub entries: Vec<SparseEntry>,
}
/// Compute a sparse delta between `old` and `new` data.
///
/// Keeps entries whose absolute change exceeds `threshold`. Returns `None`
/// if the changed fraction meets or exceeds `max_change_fraction`.
///
/// # Panics
///
/// Panics if `old.len() != new.len()`.
pub fn compute_delta(
old: &[f32],
new: &[f32],
tensor_id: u128,
block_index: u32,
base_epoch: u64,
threshold: f32,
max_change_fraction: f32,
) -> Option<DeltaRecord> {
assert_eq!(old.len(), new.len(), "old and new must have equal length");
let n = old.len();
if n == 0 {
return Some(DeltaRecord {
header: DeltaHeader {
tensor_id,
block_index,
base_epoch,
nnz: 0,
},
delta_scale: 0.0,
entries: Vec::new(),
});
}
let mut changed: Vec<(u16, f32)> = Vec::new();
let mut max_abs = 0.0f32;
for i in 0..n {
let diff = new[i] - old[i];
if diff.abs() >= threshold {
changed.push((i as u16, diff));
if diff.abs() > max_abs {
max_abs = diff.abs();
}
}
}
if changed.len() as f32 / n as f32 >= max_change_fraction {
return None;
}
let delta_scale = if max_abs == 0.0 {
1.0
} else {
max_abs / i16::MAX as f32
};
let inv_scale = 1.0 / delta_scale;
let entries: Vec<SparseEntry> = changed
.iter()
.map(|&(idx, diff)| {
let q = (diff * inv_scale).round() as i32;
SparseEntry {
index: idx,
value: q.clamp(i16::MIN as i32, i16::MAX as i32) as i16,
}
})
.collect();
Some(DeltaRecord {
header: DeltaHeader {
tensor_id,
block_index,
base_epoch,
nnz: entries.len() as u16,
},
delta_scale,
entries,
})
}
/// Apply a delta to a base data vector in-place.
///
/// Entries whose indices exceed the base length are silently skipped.
pub fn apply_delta(base: &mut [f32], delta: &DeltaRecord) {
let scale = delta.delta_scale;
for entry in &delta.entries {
let idx = entry.index as usize;
if idx < base.len() {
base[idx] += entry.value as f32 * scale;
}
}
}
/// A chain of deltas applied to a base block.
/// Invariant: `deltas.len() <= max_chain_len`.
#[derive(Clone, Debug)]
pub struct DeltaChain {
base_data: Vec<f32>,
deltas: Vec<DeltaRecord>,
max_chain_len: u8,
}
impl DeltaChain {
/// Create a new chain with a base block.
pub fn new(base_data: Vec<f32>, max_chain_len: u8) -> Self {
Self {
base_data,
deltas: Vec::new(),
max_chain_len,
}
}
/// Append a delta. Returns `Err(StoreError::DeltaChainTooLong)` at max length.
pub fn append(&mut self, delta: DeltaRecord) -> Result<(), StoreError> {
if self.deltas.len() >= self.max_chain_len as usize {
return Err(StoreError::DeltaChainTooLong);
}
self.deltas.push(delta);
Ok(())
}
/// Reconstruct the current state by applying all deltas to the base.
pub fn reconstruct(&self) -> Vec<f32> {
let mut result = self.base_data.clone();
for delta in &self.deltas {
apply_delta(&mut result, delta);
}
result
}
/// Compact the chain: apply all deltas to base, clear delta list.
pub fn compact(&mut self) {
if self.deltas.is_empty() {
return;
}
for delta in &self.deltas {
apply_delta(&mut self.base_data, delta);
}
self.deltas.clear();
}
/// Number of deltas in the chain.
#[inline]
pub fn chain_len(&self) -> usize {
self.deltas.len()
}
/// Whether the chain needs compaction (at max length).
#[inline]
pub fn needs_compaction(&self) -> bool {
self.deltas.len() >= self.max_chain_len as usize
}
/// Total storage bytes: base + serialized size of all deltas.
pub fn total_bytes(&self) -> usize {
let base_bytes = self.base_data.len() * 4;
let delta_bytes: usize = self
.deltas
.iter()
.map(|d| DELTA_HEADER_BYTES + d.entries.len() * DELTA_ENTRY_BYTES)
.sum();
base_bytes + delta_bytes
}
}
/// Low-rank factor representation for reconstruction.
///
/// Stores U (m x k), S (k), V (k x n) such that data ~ U * diag(S) * V.
/// All matrices are row-major.
#[derive(Clone, Debug)]
pub struct FactorSet {
pub m: usize,
pub n: usize,
pub k: usize,
pub u_data: Vec<f32>, // m * k elements
pub s_data: Vec<f32>, // k elements
pub v_data: Vec<f32>, // k * n elements
}
impl FactorSet {
/// Reconstruct the full data from factors: U * diag(S) * V.
pub fn reconstruct(&self) -> Vec<f32> {
let mut out = vec![0.0f32; self.m * self.n];
for r in 0..self.k {
let s_r = self.s_data[r];
for i in 0..self.m {
let u_s = self.u_data[i * self.k + r] * s_r;
let row = i * self.n;
let v_off = r * self.n;
for j in 0..self.n {
out[row + j] += u_s * self.v_data[v_off + j];
}
}
}
out
}
/// Compute storage size in bytes: (m*k + k + k*n) * 4.
pub fn storage_bytes(&self) -> usize {
(self.m * self.k + self.k + self.k * self.n) * 4
}
/// Create from a flat data vector using truncated SVD via power iteration.
///
/// Simplified implementation suitable for moderate-sized matrices.
/// Extracts top-`rank` singular triplets with successive deflation.
///
/// # Panics
///
/// Panics if `data.len() != rows * cols`.
pub fn from_data(data: &[f32], rows: usize, cols: usize, rank: usize) -> Self {
assert_eq!(
data.len(),
rows * cols,
"data length must equal rows * cols"
);
let (m, n) = (rows, cols);
let k = rank.min(m).min(n);
let mut work = data.to_vec();
let mut u_data = vec![0.0f32; m * k];
let mut s_data = vec![0.0f32; k];
let mut v_data = vec![0.0f32; k * n];
for r in 0..k {
// Deterministic initial vector: Fibonacci-hash sign pattern.
let inv_sqrt_n = 1.0 / (n as f32).sqrt();
let mut v = vec![0.0f32; n];
for j in 0..n {
let seed = (j as u32)
.wrapping_mul(2_654_435_761)
.wrapping_add((r as u32).wrapping_mul(0x9E37_79B9));
v[j] = if seed & 1 == 0 {
inv_sqrt_n
} else {
-inv_sqrt_n
};
}
let mut u = vec![0.0f32; m];
let mut sigma = 0.0f32;
for _ in 0..POWER_ITER_MAX {
// u = work * v
for i in 0..m {
let mut acc = 0.0f32;
let row = i * n;
for j in 0..n {
acc += work[row + j] * v[j];
}
u[i] = acc;
}
let su: f32 = u.iter().map(|x| x * x).sum::<f32>().sqrt();
if su < POWER_ITER_EPS {
sigma = 0.0;
break;
}
let inv = 1.0 / su;
for x in u.iter_mut() {
*x *= inv;
}
// v = work^T * u
for j in 0..n {
let mut acc = 0.0f32;
for i in 0..m {
acc += work[i * n + j] * u[i];
}
v[j] = acc;
}
let sv: f32 = v.iter().map(|x| x * x).sum::<f32>().sqrt();
if sv < POWER_ITER_EPS {
sigma = su;
break;
}
sigma = sv;
let inv = 1.0 / sv;
for x in v.iter_mut() {
*x *= inv;
}
}
s_data[r] = sigma;
for i in 0..m {
u_data[i * k + r] = u[i];
}
for j in 0..n {
v_data[r * n + j] = v[j];
}
// Deflate: work -= sigma * u * v^T
if sigma > POWER_ITER_EPS {
for i in 0..m {
let us = u[i] * sigma;
let row = i * n;
for j in 0..n {
work[row + j] -= us * v[j];
}
}
}
}
Self {
m,
n,
k,
u_data,
s_data,
v_data,
}
}
/// Compute the relative reconstruction error (Frobenius norm).
///
/// Returns `||original - reconstructed|| / ||original||`.
/// Returns 0.0 if the original has zero norm.
pub fn reconstruction_error(&self, original: &[f32]) -> f32 {
let reconstructed = self.reconstruct();
let mut diff_sq = 0.0f32;
let mut orig_sq = 0.0f32;
for (i, &o) in original.iter().enumerate() {
let r = if i < reconstructed.len() {
reconstructed[i]
} else {
0.0
};
diff_sq += (o - r) * (o - r);
orig_sq += o * o;
}
if orig_sq < 1e-30 {
return 0.0;
}
(diff_sq / orig_sq).sqrt()
}
/// Estimate the fraction of total energy (Frobenius norm) captured by factors.
///
/// Uses `sum(s_i^2)` as captured energy. Requires the original data to compute
/// total energy as `||data||_F^2`. Returns 1.0 if total energy is near zero.
pub fn energy_captured(&self, original: &[f32]) -> f32 {
let total_energy: f32 = original.iter().map(|x| x * x).sum();
if total_energy < 1e-30 {
return 1.0;
}
let captured: f32 = self.s_data.iter().map(|s| s * s).sum();
(captured / total_energy).min(1.0)
}
/// Compression ratio: original_elements * 4 bytes / storage_bytes.
///
/// Returns 0.0 if storage_bytes is zero.
pub fn compression_ratio(&self, original_elements: usize) -> f32 {
let raw = original_elements * 4;
let stored = self.storage_bytes();
if stored == 0 {
return 0.0;
}
raw as f32 / stored as f32
}
/// Create factors with adaptive rank selection.
///
/// Starts with rank 1 and increases until either `max_rank` is reached or
/// the reconstruction error falls below `target_error`.
pub fn from_data_adaptive(
data: &[f32],
rows: usize,
cols: usize,
max_rank: usize,
target_error: f32,
) -> Self {
let max_k = max_rank.min(rows).min(cols);
let mut best = Self::from_data(data, rows, cols, 1);
for rank in 2..=max_k {
let err = best.reconstruction_error(data);
if err <= target_error {
break;
}
best = Self::from_data(data, rows, cols, rank);
}
best
}
}
/// Encode a [`DeltaRecord`] to bytes (little-endian, ADR-021 section 4.1).
pub fn encode_delta(delta: &DeltaRecord) -> Vec<u8> {
let mut buf = Vec::with_capacity(DELTA_HEADER_BYTES + delta.entries.len() * DELTA_ENTRY_BYTES);
buf.extend_from_slice(&delta.header.tensor_id.to_le_bytes());
buf.extend_from_slice(&delta.header.block_index.to_le_bytes());
buf.extend_from_slice(&delta.header.base_epoch.to_le_bytes());
buf.extend_from_slice(&delta.header.nnz.to_le_bytes());
buf.extend_from_slice(&delta.delta_scale.to_le_bytes());
for entry in &delta.entries {
buf.extend_from_slice(&entry.index.to_le_bytes());
buf.extend_from_slice(&entry.value.to_le_bytes());
}
buf
}
/// Decode a [`DeltaRecord`] from bytes.
///
/// Returns `Err(StoreError::InvalidBlock)` on truncated or malformed input.
pub fn decode_delta(data: &[u8]) -> Result<DeltaRecord, StoreError> {
if data.len() < DELTA_HEADER_BYTES {
return Err(StoreError::InvalidBlock);
}
let tensor_id = u128::from_le_bytes(
data[0..16]
.try_into()
.map_err(|_| StoreError::InvalidBlock)?,
);
let block_index = u32::from_le_bytes(
data[16..20]
.try_into()
.map_err(|_| StoreError::InvalidBlock)?,
);
let base_epoch = u64::from_le_bytes(
data[20..28]
.try_into()
.map_err(|_| StoreError::InvalidBlock)?,
);
let nnz = u16::from_le_bytes(
data[28..30]
.try_into()
.map_err(|_| StoreError::InvalidBlock)?,
);
let delta_scale = f32::from_le_bytes(
data[30..34]
.try_into()
.map_err(|_| StoreError::InvalidBlock)?,
);
if data.len() < DELTA_HEADER_BYTES + (nnz as usize) * DELTA_ENTRY_BYTES {
return Err(StoreError::InvalidBlock);
}
let mut entries = Vec::with_capacity(nnz as usize);
let mut off = DELTA_HEADER_BYTES;
for _ in 0..nnz {
let index = u16::from_le_bytes(
data[off..off + 2]
.try_into()
.map_err(|_| StoreError::InvalidBlock)?,
);
let value = i16::from_le_bytes(
data[off + 2..off + 4]
.try_into()
.map_err(|_| StoreError::InvalidBlock)?,
);
entries.push(SparseEntry { index, value });
off += DELTA_ENTRY_BYTES;
}
Ok(DeltaRecord {
header: DeltaHeader {
tensor_id,
block_index,
base_epoch,
nnz,
},
delta_scale,
entries,
})
}
#[cfg(test)]
mod tests {
use super::*;
fn make_delta(entries: Vec<(u16, i16)>, scale: f32) -> DeltaRecord {
let sparse: Vec<SparseEntry> = entries
.iter()
.map(|&(i, v)| SparseEntry { index: i, value: v })
.collect();
DeltaRecord {
header: DeltaHeader {
tensor_id: 42,
block_index: 0,
base_epoch: 1,
nnz: sparse.len() as u16,
},
delta_scale: scale,
entries: sparse,
}
}
#[test]
fn test_compute_delta_small_change() {
let old = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0];
let mut new = old.clone();
new[2] = 3.5;
let d = compute_delta(&old, &new, 1, 0, 0, 0.01, 0.5).unwrap();
assert_eq!(d.entries.len(), 1);
assert_eq!(d.entries[0].index, 2);
assert!(d.delta_scale > 0.0);
}
#[test]
fn test_compute_delta_large_change_returns_none() {
let old = vec![1.0; 10];
let new = vec![5.0; 10];
assert!(compute_delta(&old, &new, 1, 0, 0, 0.01, 0.5).is_none());
}
#[test]
fn test_apply_delta_modifies_base() {
let mut base = vec![1.0, 2.0, 3.0, 4.0];
apply_delta(&mut base, &make_delta(vec![(1, 100), (3, -50)], 0.01));
assert!((base[0] - 1.0).abs() < 1e-6);
assert!((base[1] - 3.0).abs() < 1e-6); // 2.0 + 100*0.01
assert!((base[2] - 3.0).abs() < 1e-6);
assert!((base[3] - 3.5).abs() < 1e-6); // 4.0 - 50*0.01
}
#[test]
fn test_chain_append_and_reconstruct() {
let mut chain = DeltaChain::new(vec![1.0, 2.0, 3.0, 4.0], 4);
chain.append(make_delta(vec![(0, 1000)], 0.001)).unwrap(); // +1.0
assert_eq!(chain.chain_len(), 1);
let r = chain.reconstruct();
assert!((r[0] - 2.0).abs() < 1e-3);
assert!((r[1] - 2.0).abs() < 1e-6);
}
#[test]
fn test_chain_compact_preserves_state() {
let mut chain = DeltaChain::new(vec![0.0; 4], 8);
chain.append(make_delta(vec![(0, 100)], 0.1)).unwrap(); // +10.0
chain.append(make_delta(vec![(1, 200)], 0.1)).unwrap(); // +20.0
let before = chain.reconstruct();
chain.compact();
assert_eq!(chain.chain_len(), 0);
let after = chain.reconstruct();
for (a, b) in before.iter().zip(after.iter()) {
assert!((a - b).abs() < 1e-6);
}
}
#[test]
fn test_chain_max_length_enforcement() {
let mut chain = DeltaChain::new(vec![1.0; 4], 2);
assert!(chain.append(make_delta(vec![(0, 1)], 0.1)).is_ok());
assert!(chain.append(make_delta(vec![(1, 1)], 0.1)).is_ok());
assert!(chain.append(make_delta(vec![(2, 1)], 0.1)).is_err());
}
#[test]
fn test_chain_needs_compaction() {
let mut chain = DeltaChain::new(vec![1.0; 4], 2);
assert!(!chain.needs_compaction());
chain.append(make_delta(vec![(0, 1)], 0.1)).unwrap();
assert!(!chain.needs_compaction());
chain.append(make_delta(vec![(1, 1)], 0.1)).unwrap();
assert!(chain.needs_compaction());
}
#[test]
fn test_factor_reconstruct() {
let (u, v, s) = (vec![1.0, 2.0, 3.0], vec![4.0, 5.0], 2.0);
let f = FactorSet {
m: 3,
n: 2,
k: 1,
u_data: u.clone(),
s_data: vec![s],
v_data: v.clone(),
};
let r = f.reconstruct();
assert_eq!(r.len(), 6);
for i in 0..3 {
for j in 0..2 {
assert!((r[i * 2 + j] - u[i] * s * v[j]).abs() < 1e-6);
}
}
}
#[test]
fn test_factor_from_data_approximation() {
let (m, n) = (8, 6);
let data: Vec<f32> = (0..m * n)
.map(|idx| {
let (i, j) = (idx / n, idx % n);
(i as f32 + 1.0) * (j as f32 + 1.0)
})
.collect();
let reconstructed = FactorSet::from_data(&data, m, n, 1).reconstruct();
let max_err = data
.iter()
.zip(reconstructed.iter())
.map(|(a, b)| (a - b).abs())
.fold(0.0f32, f32::max);
assert!(
max_err < 0.5,
"max error {max_err} too large for rank-1 input"
);
}
#[test]
fn test_encode_decode_roundtrip() {
let orig = DeltaRecord {
header: DeltaHeader {
tensor_id: 0xDEADBEEFCAFEBABE,
block_index: 42,
base_epoch: 100,
nnz: 3,
},
delta_scale: 0.001,
entries: vec![
SparseEntry {
index: 10,
value: 500,
},
SparseEntry {
index: 20,
value: -300,
},
SparseEntry {
index: 30,
value: 1,
},
],
};
let bytes = encode_delta(&orig);
assert_eq!(bytes.len(), DELTA_HEADER_BYTES + 3 * DELTA_ENTRY_BYTES);
let dec = decode_delta(&bytes).unwrap();
assert_eq!(dec.header.tensor_id, orig.header.tensor_id);
assert_eq!(dec.header.block_index, orig.header.block_index);
assert_eq!(dec.header.nnz, orig.header.nnz);
assert!((dec.delta_scale - orig.delta_scale).abs() < 1e-10);
for (a, b) in dec.entries.iter().zip(orig.entries.iter()) {
assert_eq!(a.index, b.index);
assert_eq!(a.value, b.value);
}
}
#[test]
fn test_decode_truncated_header() {
assert!(decode_delta(&vec![0u8; 20]).is_err());
}
#[test]
fn test_decode_truncated_entries() {
let mut bytes = encode_delta(&make_delta(vec![(0, 1), (1, 2)], 1.0));
bytes[28] = 5;
bytes[29] = 0; // claim 5 entries, only 2 present
assert!(decode_delta(&bytes).is_err());
}
#[test]
fn test_empty_delta_roundtrip() {
let d = DeltaRecord {
header: DeltaHeader {
tensor_id: 99,
block_index: 7,
base_epoch: 50,
nnz: 0,
},
delta_scale: 0.0,
entries: Vec::new(),
};
let dec = decode_delta(&encode_delta(&d)).unwrap();
assert_eq!(dec.entries.len(), 0);
}
#[test]
fn test_single_entry_delta() {
let old = vec![1.0; 100];
let mut new = old.clone();
new[50] = 2.0;
let d = compute_delta(&old, &new, 1, 0, 0, 0.01, 0.5).unwrap();
assert_eq!(d.entries.len(), 1);
assert_eq!(d.entries[0].index, 50);
let mut base = old.clone();
apply_delta(&mut base, &d);
assert!((base[50] - 2.0).abs() < 0.01);
}
#[test]
fn test_full_density_delta() {
let old = vec![0.0; 4];
let new = vec![0.1, 0.2, 0.3, 0.4];
let d = compute_delta(&old, &new, 1, 0, 0, 0.001, 1.1).unwrap();
assert_eq!(d.entries.len(), 4);
let mut base = old.clone();
apply_delta(&mut base, &d);
for i in 0..4 {
assert!((base[i] - new[i]).abs() < 0.01, "index {i}");
}
}
#[test]
fn test_compute_apply_roundtrip_64() {
let old: Vec<f32> = (0..64).map(|i| i as f32 * 0.1).collect();
let mut new = old.clone();
new[5] += 0.5;
new[10] -= 0.3;
new[60] += 1.0;
let d = compute_delta(&old, &new, 1, 0, 0, 0.01, 0.5).unwrap();
let mut recon = old.clone();
apply_delta(&mut recon, &d);
for i in 0..64 {
assert!((recon[i] - new[i]).abs() < 0.01, "index {i}");
}
}
#[test]
fn test_reconstruction_error_zero_for_exact() {
// Rank-1 data should be exactly reconstructed with rank-1 factors
let (m, n) = (4, 3);
let data: Vec<f32> = (0..m * n)
.map(|idx| {
let (i, j) = (idx / n, idx % n);
(i as f32 + 1.0) * (j as f32 + 1.0)
})
.collect();
let factors = FactorSet::from_data(&data, m, n, 1);
let err = factors.reconstruction_error(&data);
assert!(err < 0.01, "err={err} too large for rank-1 data");
}
#[test]
fn test_reconstruction_error_decreases_with_rank() {
let (m, n) = (8, 6);
let data: Vec<f32> = (0..m * n).map(|i| (i as f32 * 0.7).sin()).collect();
let err1 = FactorSet::from_data(&data, m, n, 1).reconstruction_error(&data);
let err3 = FactorSet::from_data(&data, m, n, 3).reconstruction_error(&data);
assert!(err3 <= err1 + 1e-6, "err3={err3} > err1={err1}");
}
#[test]
fn test_energy_captured_rank1_data() {
let (m, n) = (4, 3);
let data: Vec<f32> = (0..m * n)
.map(|idx| {
let (i, j) = (idx / n, idx % n);
(i as f32 + 1.0) * (j as f32 + 1.0)
})
.collect();
let factors = FactorSet::from_data(&data, m, n, 1);
let energy = factors.energy_captured(&data);
assert!(energy > 0.95, "energy={energy} too low for rank-1 data");
}
#[test]
fn test_compression_ratio_meaningful() {
let (m, n) = (16, 16);
let data: Vec<f32> = (0..m * n).map(|i| i as f32).collect();
let factors = FactorSet::from_data(&data, m, n, 2);
let ratio = factors.compression_ratio(m * n);
// rank-2 storage: (16*2 + 2 + 2*16) * 4 = 264 bytes vs 16*16*4 = 1024 bytes
assert!(ratio > 1.0, "ratio={ratio} should be > 1");
}
#[test]
fn test_from_data_adaptive_stops_early() {
let (m, n) = (4, 3);
// Rank-1 data: adaptive should stop at rank 1
let data: Vec<f32> = (0..m * n)
.map(|idx| {
let (i, j) = (idx / n, idx % n);
(i as f32 + 1.0) * (j as f32 + 1.0)
})
.collect();
let factors = FactorSet::from_data_adaptive(&data, m, n, 5, 0.05);
// Should use rank 1 since data is rank 1
assert!(
factors.k <= 2,
"k={} should be small for rank-1 data",
factors.k
);
}
#[test]
fn test_from_data_adaptive_increases_rank() {
let (m, n) = (8, 6);
// Multi-rank data
let data: Vec<f32> = (0..m * n)
.map(|i| (i as f32 * 0.3).sin() + (i as f32 * 0.7).cos())
.collect();
let factors = FactorSet::from_data_adaptive(&data, m, n, 6, 0.01);
let err = factors.reconstruction_error(&data);
// Should achieve close to target error or use max rank
assert!(err < 0.1 || factors.k == 6, "err={err}, k={}", factors.k);
}
}

View File

@@ -0,0 +1,150 @@
//! Software IEEE 754 half-precision (f16) conversion.
//!
//! No external crate dependencies. Handles normals, denormals, infinity, and NaN.
//! Round-to-nearest with ties-to-even for normal values.
/// Convert f32 to f16 bit representation.
///
/// Handles all IEEE 754 special cases: infinity, NaN, denormals, and zero (both signs).
/// Values outside f16 range saturate to infinity. Values too small for f16 denormals
/// flush to zero.
#[inline]
pub fn f32_to_f16_bits(x: f32) -> u16 {
let b = x.to_bits();
let sign = ((b >> 16) & 0x8000) as u16;
let exp = ((b >> 23) & 0xFF) as i32;
let mant = b & 0x7F_FFFF;
// Infinity or NaN
if exp == 255 {
if mant == 0 {
return sign | 0x7C00;
}
let nan_m = (mant >> 13) as u16;
return sign | 0x7C00 | nan_m | 1;
}
let exp16 = exp - 127 + 15;
// Overflow -> Infinity
if exp16 >= 31 {
return sign | 0x7C00;
}
// Underflow -> denormal or zero
if exp16 <= 0 {
if exp16 < -10 {
return sign;
}
let shift = (14 - exp16) as u32;
let mut mant32 = mant | 0x80_0000;
let round_bit = 1u32.wrapping_shl(shift.wrapping_sub(1));
mant32 = mant32.wrapping_add(round_bit);
let sub = (mant32 >> shift) as u16;
return sign | sub;
}
// Normal case
let mant16 = (mant >> 13) as u16;
let round = (mant >> 12) & 1;
let mut res = sign | ((exp16 as u16) << 10) | mant16;
if round != 0 {
res = res.wrapping_add(1);
}
res
}
/// Convert f16 bit representation to f32.
///
/// Exactly reconstructs the f32 value represented by the f16 bit pattern.
/// Handles denormals by normalizing the mantissa before constructing the f32 bits.
#[inline]
pub fn f16_bits_to_f32(h: u16) -> f32 {
let sign = ((h & 0x8000) as u32) << 16;
let exp = ((h >> 10) & 0x1F) as i32;
let mant = (h & 0x03FF) as u32;
// Zero or denormal
if exp == 0 {
if mant == 0 {
return f32::from_bits(sign);
}
let mut e = 1i32;
let mut m = mant;
while (m & 0x0400) == 0 {
m <<= 1;
e += 1;
}
m &= 0x03FF;
let exp32 = 127 - 15 - e + 1;
let mant32 = m << 13;
return f32::from_bits(sign | ((exp32 as u32) << 23) | mant32);
}
// Infinity or NaN
if exp == 31 {
return f32::from_bits(sign | 0x7F80_0000 | (mant << 13));
}
// Normal
let exp32 = exp - 15 + 127;
let mant32 = mant << 13;
f32::from_bits(sign | ((exp32 as u32) << 23) | mant32)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_roundtrip_normal() {
for &v in &[0.0f32, 1.0, -1.0, 0.5, 65504.0, -65504.0, 0.0001] {
let h = f32_to_f16_bits(v);
let back = f16_bits_to_f32(h);
if v == 0.0 {
assert_eq!(back, 0.0);
} else {
let rel_err = ((back - v) / v).abs();
assert!(rel_err < 0.01, "v={v}, back={back}, rel_err={rel_err}");
}
}
}
#[test]
fn test_infinity() {
let h = f32_to_f16_bits(f32::INFINITY);
assert_eq!(h, 0x7C00);
assert!(f16_bits_to_f32(h).is_infinite());
}
#[test]
fn test_neg_infinity() {
let h = f32_to_f16_bits(f32::NEG_INFINITY);
assert_eq!(h, 0xFC00);
let back = f16_bits_to_f32(h);
assert!(back.is_infinite() && back < 0.0);
}
#[test]
fn test_nan() {
let h = f32_to_f16_bits(f32::NAN);
assert!(f16_bits_to_f32(h).is_nan());
}
#[test]
fn test_zero_signs() {
assert_eq!(f32_to_f16_bits(0.0f32), 0x0000);
assert_eq!(f32_to_f16_bits(-0.0f32), 0x8000);
}
#[test]
fn test_scale_range_accuracy() {
for exp in -4..=4i32 {
let v = 10.0f32.powi(exp);
let h = f32_to_f16_bits(v);
let back = f16_bits_to_f32(h);
let rel_err = ((back - v) / v).abs();
assert!(rel_err < 0.002, "v={v}, back={back}, rel_err={rel_err}");
}
}
}

View File

@@ -0,0 +1,250 @@
//! WASM/C FFI interface with handle-based resource management.
//!
//! Exports `extern "C"` functions for:
//! - Compressor lifecycle (`ttc_create`, `ttc_free`, `ttc_touch`, `ttc_set_access`)
//! - Frame compression (`ttc_push_frame`, `ttc_flush`)
//! - Segment decoding (`ttc_decode_segment`)
//! - Memory management (`ttc_alloc`, `ttc_dealloc`)
use crate::compressor::TemporalTensorCompressor;
use crate::segment;
use crate::tier_policy::TierPolicy;
static mut STORE: Option<Vec<Option<TemporalTensorCompressor>>> = None;
fn store_init() {
unsafe {
if STORE.is_none() {
STORE = Some(Vec::new());
}
}
}
fn with_store<F, R>(f: F) -> R
where
F: FnOnce(&mut Vec<Option<TemporalTensorCompressor>>) -> R,
{
store_init();
unsafe { f(STORE.as_mut().unwrap()) }
}
fn with_compressor<F>(handle: u32, f: F)
where
F: FnOnce(&mut TemporalTensorCompressor),
{
with_store(|store| {
let idx = handle as usize;
if idx < store.len() {
if let Some(comp) = store[idx].as_mut() {
f(comp);
}
}
});
}
/// Create a new compressor. Returns handle via out_handle.
#[no_mangle]
pub extern "C" fn ttc_create(len: u32, now_ts: u32, out_handle: *mut u32) {
let policy = TierPolicy::default();
let comp = TemporalTensorCompressor::new(policy, len, now_ts);
with_store(|store| {
// Find a free slot
for (i, slot) in store.iter_mut().enumerate() {
if slot.is_none() {
*slot = Some(comp);
if !out_handle.is_null() {
unsafe { *out_handle = i as u32 };
}
return;
}
}
// No free slot, push
let idx = store.len();
store.push(Some(comp));
if !out_handle.is_null() {
unsafe { *out_handle = idx as u32 };
}
});
}
/// Create a compressor with custom policy parameters.
#[no_mangle]
pub extern "C" fn ttc_create_with_policy(
len: u32,
now_ts: u32,
hot_min_score: u32,
warm_min_score: u32,
warm_bits: u8,
drift_pct_q8: u32,
group_len: u32,
out_handle: *mut u32,
) {
let policy = TierPolicy {
hot_min_score,
warm_min_score,
warm_bits,
drift_pct_q8,
group_len,
};
let comp = TemporalTensorCompressor::new(policy, len, now_ts);
with_store(|store| {
for (i, slot) in store.iter_mut().enumerate() {
if slot.is_none() {
*slot = Some(comp);
if !out_handle.is_null() {
unsafe { *out_handle = i as u32 };
}
return;
}
}
let idx = store.len();
store.push(Some(comp));
if !out_handle.is_null() {
unsafe { *out_handle = idx as u32 };
}
});
}
/// Free a compressor.
#[no_mangle]
pub extern "C" fn ttc_free(handle: u32) {
with_store(|store| {
let idx = handle as usize;
if idx < store.len() {
store[idx] = None;
}
});
}
/// Record an access event.
#[no_mangle]
pub extern "C" fn ttc_touch(handle: u32, now_ts: u32) {
with_compressor(handle, |comp| comp.touch(now_ts));
}
/// Set access stats directly.
#[no_mangle]
pub extern "C" fn ttc_set_access(handle: u32, access_count: u32, last_access_ts: u32) {
with_compressor(handle, |comp| comp.set_access(access_count, last_access_ts));
}
/// Push a frame. If a segment boundary is crossed, the completed segment
/// is written to out_ptr/out_cap, and out_written is set to the byte count.
#[no_mangle]
pub extern "C" fn ttc_push_frame(
handle: u32,
now_ts: u32,
in_ptr: *const f32,
len: u32,
out_ptr: *mut u8,
out_cap: u32,
out_written: *mut u32,
) {
if out_written.is_null() {
return;
}
unsafe { *out_written = 0 };
if in_ptr.is_null() || out_ptr.is_null() {
return;
}
let frame = unsafe { std::slice::from_raw_parts(in_ptr, len as usize) };
let mut seg = Vec::new();
with_compressor(handle, |comp| {
comp.push_frame(frame, now_ts, &mut seg);
});
if seg.is_empty() || (seg.len() as u32) > out_cap {
return;
}
unsafe {
let out = std::slice::from_raw_parts_mut(out_ptr, out_cap as usize);
out[..seg.len()].copy_from_slice(&seg);
*out_written = seg.len() as u32;
}
}
/// Flush the current segment.
#[no_mangle]
pub extern "C" fn ttc_flush(handle: u32, out_ptr: *mut u8, out_cap: u32, out_written: *mut u32) {
if out_written.is_null() {
return;
}
unsafe { *out_written = 0 };
let mut seg = Vec::new();
with_compressor(handle, |comp| {
comp.flush(&mut seg);
});
if seg.is_empty() || out_ptr.is_null() || (seg.len() as u32) > out_cap {
return;
}
unsafe {
let out = std::slice::from_raw_parts_mut(out_ptr, out_cap as usize);
out[..seg.len()].copy_from_slice(&seg);
*out_written = seg.len() as u32;
}
}
/// Decode a segment into f32 values.
#[no_mangle]
pub extern "C" fn ttc_decode_segment(
seg_ptr: *const u8,
seg_len: u32,
out_ptr: *mut f32,
out_cap_f32: u32,
out_written_f32: *mut u32,
) {
if out_written_f32.is_null() {
return;
}
unsafe { *out_written_f32 = 0 };
if seg_ptr.is_null() || out_ptr.is_null() {
return;
}
let seg = unsafe { std::slice::from_raw_parts(seg_ptr, seg_len as usize) };
let mut values = Vec::new();
segment::decode(seg, &mut values);
if values.is_empty() || (values.len() as u32) > out_cap_f32 {
return;
}
unsafe {
let out = std::slice::from_raw_parts_mut(out_ptr, out_cap_f32 as usize);
out[..values.len()].copy_from_slice(&values);
*out_written_f32 = values.len() as u32;
}
}
/// Allocate a buffer in WASM linear memory.
#[no_mangle]
pub extern "C" fn ttc_alloc(size: u32, out_ptr: *mut u32) {
if out_ptr.is_null() {
return;
}
let mut v: Vec<u8> = Vec::with_capacity(size as usize);
let p = v.as_mut_ptr();
std::mem::forget(v);
unsafe {
*out_ptr = p as u32;
}
}
/// Free a buffer previously allocated with ttc_alloc.
#[no_mangle]
pub extern "C" fn ttc_dealloc(ptr: u32, cap: u32) {
if ptr == 0 || cap == 0 {
return;
}
unsafe {
let _ = Vec::<u8>::from_raw_parts(ptr as *mut u8, 0, cap as usize);
}
}

View File

@@ -0,0 +1,99 @@
//! Temporal Tensor Compression with Tiered Quantization
//!
//! Implements ADR-017: groupwise symmetric quantization with temporal segment
//! reuse and access-pattern-driven tier selection (8/7/5/3 bit).
//!
//! # Architecture
//!
//! ```text
//! f32 frame → tier_policy → quantizer → bitpack → segment
//! segment → bitpack → quantizer → f32 output
//! ```
//!
//! # Compression Ratios
//!
//! | Tier | Bits | Ratio vs f32 | Use Case |
//! |------|------|-------------|----------|
//! | Hot | 8 | ~4.0x | Frequently accessed tensors |
//! | Warm | 7 | ~4.57x | Moderately accessed |
//! | Warm | 5 | ~6.4x | Aggressively compressed warm |
//! | Cold | 3 | ~10.67x | Rarely accessed |
//!
//! # Zero Dependencies
//!
//! This crate has no external dependencies, making it fully WASM-compatible.
//!
//! # Quick Start
//!
//! ```rust
//! use ruvector_temporal_tensor::{TemporalTensorCompressor, TierPolicy};
//!
//! // Create a compressor for 128-element tensors
//! let mut comp = TemporalTensorCompressor::new(TierPolicy::default(), 128, 0);
//! comp.set_access(100, 0); // hot tensor -> 8-bit quantization
//!
//! let frame = vec![1.0f32; 128];
//! let mut segment = Vec::new();
//!
//! // Push frames; segment is populated when a boundary is crossed
//! comp.push_frame(&frame, 1, &mut segment);
//! comp.flush(&mut segment); // force-emit the current segment
//!
//! // Decode the segment back to f32
//! let mut decoded = Vec::new();
//! ruvector_temporal_tensor::segment::decode(&segment, &mut decoded);
//! assert_eq!(decoded.len(), 128);
//! ```
//!
//! # Random-Access Decode
//!
//! ```rust
//! # use ruvector_temporal_tensor::{TemporalTensorCompressor, TierPolicy};
//! # let mut comp = TemporalTensorCompressor::new(TierPolicy::default(), 64, 0);
//! # let frame = vec![1.0f32; 64];
//! # let mut seg = Vec::new();
//! # comp.push_frame(&frame, 0, &mut seg);
//! # comp.flush(&mut seg);
//! // Decode only frame 0 without decoding the entire segment
//! let single = ruvector_temporal_tensor::segment::decode_single_frame(&seg, 0);
//! assert!(single.is_some());
//! ```
//!
//! # Compression Ratio Inspection
//!
//! ```rust
//! # use ruvector_temporal_tensor::{TemporalTensorCompressor, TierPolicy};
//! # let mut comp = TemporalTensorCompressor::new(TierPolicy::default(), 64, 0);
//! # let frame = vec![1.0f32; 64];
//! # let mut seg = Vec::new();
//! # comp.push_frame(&frame, 0, &mut seg);
//! # comp.flush(&mut seg);
//! let ratio = ruvector_temporal_tensor::segment::compression_ratio(&seg);
//! assert!(ratio > 1.0);
//! ```
pub mod bitpack;
pub mod compressor;
pub mod delta;
pub mod f16;
pub mod metrics;
pub mod quantizer;
pub mod segment;
pub mod store;
pub mod tier_policy;
pub mod tiering;
pub mod agentdb;
pub mod coherence;
pub mod core_trait;
#[cfg(feature = "persistence")]
pub mod persistence;
#[cfg(feature = "ffi")]
pub mod ffi;
#[cfg(feature = "ffi")]
pub mod store_ffi;
pub use compressor::TemporalTensorCompressor;
pub use tier_policy::TierPolicy;

File diff suppressed because it is too large Load Diff

View File

@@ -0,0 +1,859 @@
//! Disk-backed BlockIO and MetaLog implementations.
//!
//! Gated behind the `persistence` feature flag. Uses raw file I/O
//! with a simple binary format. No external dependencies.
#![cfg(feature = "persistence")]
use crate::store::{
BlockIO, BlockKey, BlockMeta, DType, MetaLog, ReconstructPolicy, StoreError, Tier,
};
use std::collections::HashMap;
use std::fs;
use std::io::Write;
use std::path::{Path, PathBuf};
/// Fixed size of a single encoded [`BlockMeta`] record in bytes.
///
/// Layout (all little-endian):
///
/// | Offset | Size | Field |
/// |--------|------|-----------------|
/// | 0 | 16 | tensor_id |
/// | 16 | 4 | block_index |
/// | 20 | 1 | dtype |
/// | 21 | 1 | tier |
/// | 22 | 1 | bits |
/// | 23 | 4 | scale |
/// | 27 | 2 | zero_point |
/// | 29 | 8 | created_at |
/// | 37 | 8 | last_access_at |
/// | 45 | 4 | access_count |
/// | 49 | 4 | ema_rate |
/// | 53 | 8 | window |
/// | 61 | 4 | checksum |
/// | 65 | 1 | reconstruct |
/// | 66 | 4 | tier_age |
/// | 70 | 1 | has_lineage |
/// | 71 | 16 | lineage_parent |
/// | 87 | 4 | block_bytes |
const RECORD_SIZE: usize = 91;
// ---------------------------------------------------------------------------
// Serialization helpers
// ---------------------------------------------------------------------------
/// Serialize a [`BlockMeta`] into a fixed-size byte vector.
///
/// The encoding uses little-endian byte order for all multi-byte fields
/// and occupies exactly [`RECORD_SIZE`] bytes.
pub fn encode_meta(meta: &BlockMeta) -> Vec<u8> {
let mut buf = Vec::with_capacity(RECORD_SIZE);
// key
buf.extend_from_slice(&meta.key.tensor_id.to_le_bytes());
buf.extend_from_slice(&meta.key.block_index.to_le_bytes());
// scalar metadata
buf.push(meta.dtype as u8);
buf.push(meta.tier as u8);
buf.push(meta.bits);
buf.extend_from_slice(&meta.scale.to_le_bytes());
buf.extend_from_slice(&meta.zero_point.to_le_bytes());
// timestamps and counters
buf.extend_from_slice(&meta.created_at.to_le_bytes());
buf.extend_from_slice(&meta.last_access_at.to_le_bytes());
buf.extend_from_slice(&meta.access_count.to_le_bytes());
buf.extend_from_slice(&meta.ema_rate.to_le_bytes());
buf.extend_from_slice(&meta.window.to_le_bytes());
buf.extend_from_slice(&meta.checksum.to_le_bytes());
// policy and age
buf.push(meta.reconstruct as u8);
buf.extend_from_slice(&meta.tier_age.to_le_bytes());
// optional lineage parent
match meta.lineage_parent {
Some(parent) => {
buf.push(1);
buf.extend_from_slice(&parent.to_le_bytes());
}
None => {
buf.push(0);
buf.extend_from_slice(&0u128.to_le_bytes());
}
}
// payload size
buf.extend_from_slice(&meta.block_bytes.to_le_bytes());
debug_assert_eq!(buf.len(), RECORD_SIZE);
buf
}
/// Deserialize a [`BlockMeta`] from a byte slice of at least [`RECORD_SIZE`] bytes.
///
/// Returns [`StoreError::InvalidData`] if the slice is too short or
/// contains invalid enum discriminants.
pub fn decode_meta(bytes: &[u8]) -> Result<BlockMeta, StoreError> {
if bytes.len() < RECORD_SIZE {
return Err(StoreError::InvalidData);
}
let tensor_id = u128::from_le_bytes(
bytes[0..16]
.try_into()
.map_err(|_| StoreError::InvalidData)?,
);
let block_index = u32::from_le_bytes(
bytes[16..20]
.try_into()
.map_err(|_| StoreError::InvalidData)?,
);
let dtype = match bytes[20] {
0 => DType::F32,
1 => DType::F16,
2 => DType::BF16,
_ => return Err(StoreError::InvalidData),
};
let tier = match bytes[21] {
0 => Tier::Tier0,
1 => Tier::Tier1,
2 => Tier::Tier2,
3 => Tier::Tier3,
_ => return Err(StoreError::InvalidData),
};
let bits = bytes[22];
let scale = f32::from_le_bytes(
bytes[23..27]
.try_into()
.map_err(|_| StoreError::InvalidData)?,
);
let zero_point = i16::from_le_bytes(
bytes[27..29]
.try_into()
.map_err(|_| StoreError::InvalidData)?,
);
let created_at = u64::from_le_bytes(
bytes[29..37]
.try_into()
.map_err(|_| StoreError::InvalidData)?,
);
let last_access_at = u64::from_le_bytes(
bytes[37..45]
.try_into()
.map_err(|_| StoreError::InvalidData)?,
);
let access_count = u32::from_le_bytes(
bytes[45..49]
.try_into()
.map_err(|_| StoreError::InvalidData)?,
);
let ema_rate = f32::from_le_bytes(
bytes[49..53]
.try_into()
.map_err(|_| StoreError::InvalidData)?,
);
let window = u64::from_le_bytes(
bytes[53..61]
.try_into()
.map_err(|_| StoreError::InvalidData)?,
);
let checksum = u32::from_le_bytes(
bytes[61..65]
.try_into()
.map_err(|_| StoreError::InvalidData)?,
);
let reconstruct = match bytes[65] {
0 => ReconstructPolicy::None,
1 => ReconstructPolicy::Delta,
2 => ReconstructPolicy::Factor,
_ => return Err(StoreError::InvalidData),
};
let tier_age = u32::from_le_bytes(
bytes[66..70]
.try_into()
.map_err(|_| StoreError::InvalidData)?,
);
let has_lineage = bytes[70];
let lineage_value = u128::from_le_bytes(
bytes[71..87]
.try_into()
.map_err(|_| StoreError::InvalidData)?,
);
let lineage_parent = if has_lineage != 0 {
Some(lineage_value)
} else {
None
};
let block_bytes = u32::from_le_bytes(
bytes[87..91]
.try_into()
.map_err(|_| StoreError::InvalidData)?,
);
Ok(BlockMeta {
key: BlockKey {
tensor_id,
block_index,
},
dtype,
tier,
bits,
scale,
zero_point,
created_at,
last_access_at,
access_count,
ema_rate,
window,
checksum,
reconstruct,
tier_age,
lineage_parent,
block_bytes,
})
}
// ---------------------------------------------------------------------------
// FileBlockIO
// ---------------------------------------------------------------------------
/// Disk-backed [`BlockIO`] that stores each block as a separate file.
///
/// Directory layout:
/// ```text
/// {base_dir}/
/// tier0/
/// tier1/
/// tier2/
/// tier3/
/// ```
///
/// Each block file is named `{tensor_id_hex}_{block_index}.bin`.
pub struct FileBlockIO {
base_dir: PathBuf,
}
impl FileBlockIO {
/// Create a new `FileBlockIO` rooted at `base_dir`.
///
/// Creates the tier subdirectories if they do not already exist.
pub fn new(base_dir: impl Into<PathBuf>) -> Result<Self, StoreError> {
let base_dir = base_dir.into();
for tier_num in 0..=3u8 {
let tier_dir = base_dir.join(format!("tier{}", tier_num));
fs::create_dir_all(&tier_dir).map_err(|_| StoreError::IOError)?;
}
Ok(Self { base_dir })
}
/// Return the filesystem path for a given block.
fn block_path(&self, tier: Tier, key: BlockKey) -> PathBuf {
self.base_dir
.join(format!("tier{}", tier as u8))
.join(format!("{:032x}_{}.bin", key.tensor_id, key.block_index))
}
/// Return the base directory.
pub fn base_dir(&self) -> &Path {
&self.base_dir
}
}
impl BlockIO for FileBlockIO {
fn read_block(&self, tier: Tier, key: BlockKey, dst: &mut [u8]) -> Result<usize, StoreError> {
let path = self.block_path(tier, key);
let data = fs::read(&path).map_err(|_| StoreError::BlockNotFound)?;
let n = data.len().min(dst.len());
dst[..n].copy_from_slice(&data[..n]);
Ok(n)
}
fn write_block(&mut self, tier: Tier, key: BlockKey, src: &[u8]) -> Result<(), StoreError> {
if tier == Tier::Tier0 {
return Err(StoreError::InvalidBlock);
}
let path = self.block_path(tier, key);
fs::write(&path, src).map_err(|_| StoreError::IOError)
}
fn delete_block(&mut self, tier: Tier, key: BlockKey) -> Result<(), StoreError> {
let path = self.block_path(tier, key);
fs::remove_file(&path).map_err(|_| StoreError::BlockNotFound)
}
}
// ---------------------------------------------------------------------------
// FileMetaLog
// ---------------------------------------------------------------------------
/// Append-only file-backed [`MetaLog`].
///
/// Each [`append`](MetaLog::append) call writes a fixed-size binary record
/// to `{base_dir}/meta.log`. On construction the log is replayed into an
/// in-memory [`HashMap`] so that [`get`](MetaLog::get) is a simple lookup.
///
/// Because the log is append-only, multiple records for the same key may
/// exist on disk. The last record wins when the log is replayed.
pub struct FileMetaLog {
log_path: PathBuf,
index: HashMap<BlockKey, BlockMeta>,
}
impl FileMetaLog {
/// Open (or create) a `FileMetaLog` rooted at `base_dir`.
///
/// If `{base_dir}/meta.log` already exists it is replayed to populate
/// the in-memory index.
pub fn new(base_dir: impl Into<PathBuf>) -> Result<Self, StoreError> {
let base_dir = base_dir.into();
fs::create_dir_all(&base_dir).map_err(|_| StoreError::IOError)?;
let log_path = base_dir.join("meta.log");
let mut index = HashMap::new();
if log_path.exists() {
let data = fs::read(&log_path).map_err(|_| StoreError::IOError)?;
let mut offset = 0;
while offset + RECORD_SIZE <= data.len() {
if let Ok(meta) = decode_meta(&data[offset..offset + RECORD_SIZE]) {
index.insert(meta.key, meta);
}
offset += RECORD_SIZE;
}
}
Ok(Self { log_path, index })
}
/// Return the path to the underlying log file.
pub fn log_path(&self) -> &Path {
&self.log_path
}
/// Number of unique blocks tracked in the in-memory index.
pub fn len(&self) -> usize {
self.index.len()
}
/// Returns `true` if no metadata records are tracked.
pub fn is_empty(&self) -> bool {
self.index.is_empty()
}
}
impl MetaLog for FileMetaLog {
fn append(&mut self, rec: &BlockMeta) -> Result<(), StoreError> {
let encoded = encode_meta(rec);
let mut file = fs::OpenOptions::new()
.create(true)
.append(true)
.open(&self.log_path)
.map_err(|_| StoreError::IOError)?;
file.write_all(&encoded).map_err(|_| StoreError::IOError)?;
file.flush().map_err(|_| StoreError::IOError)?;
self.index.insert(rec.key, rec.clone());
Ok(())
}
fn get(&self, key: BlockKey) -> Option<&BlockMeta> {
self.index.get(&key)
}
fn iter(&self) -> Box<dyn Iterator<Item = &BlockMeta> + '_> {
Box::new(self.index.values())
}
}
// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------
#[cfg(test)]
mod tests {
use super::*;
use std::sync::atomic::{AtomicU32, Ordering};
/// Monotonic counter for unique test directory names.
static TEST_ID: AtomicU32 = AtomicU32::new(0);
/// Create a unique temporary directory for a test.
fn test_dir(prefix: &str) -> PathBuf {
let id = TEST_ID.fetch_add(1, Ordering::SeqCst);
let pid = std::process::id();
let dir =
std::env::temp_dir().join(format!("ruvector_persistence_{}_{}_{}", prefix, pid, id));
let _ = fs::remove_dir_all(&dir);
fs::create_dir_all(&dir).unwrap();
dir
}
/// Clean up a test directory (best-effort).
fn cleanup(dir: &Path) {
let _ = fs::remove_dir_all(dir);
}
fn make_key(tid: u128, idx: u32) -> BlockKey {
BlockKey {
tensor_id: tid,
block_index: idx,
}
}
fn sample_meta(key: BlockKey) -> BlockMeta {
BlockMeta {
key,
dtype: DType::F32,
tier: Tier::Tier1,
bits: 8,
scale: 0.03125,
zero_point: 0,
created_at: 1000,
last_access_at: 2000,
access_count: 42,
ema_rate: 0.75,
window: 0xAAAA_BBBB_CCCC_DDDD,
checksum: 0xDEAD_BEEF,
reconstruct: ReconstructPolicy::None,
tier_age: 15,
lineage_parent: None,
block_bytes: 512,
}
}
// -- encode / decode roundtrip -----------------------------------------
#[test]
fn encode_decode_roundtrip_basic() {
let key = make_key(0x0123_4567_89AB_CDEF_FEDC_BA98_7654_3210, 7);
let meta = sample_meta(key);
let encoded = encode_meta(&meta);
assert_eq!(encoded.len(), RECORD_SIZE);
let decoded = decode_meta(&encoded).unwrap();
assert_eq!(decoded.key, meta.key);
assert_eq!(decoded.dtype, meta.dtype);
assert_eq!(decoded.tier, meta.tier);
assert_eq!(decoded.bits, meta.bits);
assert!((decoded.scale - meta.scale).abs() < 1e-10);
assert_eq!(decoded.zero_point, meta.zero_point);
assert_eq!(decoded.created_at, meta.created_at);
assert_eq!(decoded.last_access_at, meta.last_access_at);
assert_eq!(decoded.access_count, meta.access_count);
assert!((decoded.ema_rate - meta.ema_rate).abs() < 1e-6);
assert_eq!(decoded.window, meta.window);
assert_eq!(decoded.checksum, meta.checksum);
assert_eq!(decoded.reconstruct, meta.reconstruct);
assert_eq!(decoded.tier_age, meta.tier_age);
assert_eq!(decoded.lineage_parent, meta.lineage_parent);
assert_eq!(decoded.block_bytes, meta.block_bytes);
}
#[test]
fn encode_decode_with_lineage() {
let key = make_key(1, 0);
let mut meta = sample_meta(key);
meta.lineage_parent = Some(0xFFFF_FFFF_FFFF_FFFF_0000_0000_0000_0001);
let encoded = encode_meta(&meta);
let decoded = decode_meta(&encoded).unwrap();
assert_eq!(
decoded.lineage_parent,
Some(0xFFFF_FFFF_FFFF_FFFF_0000_0000_0000_0001)
);
}
#[test]
fn encode_decode_all_dtypes() {
for (dtype_val, expected) in [(0u8, DType::F32), (1, DType::F16), (2, DType::BF16)] {
let key = make_key(dtype_val as u128, 0);
let mut meta = sample_meta(key);
meta.dtype = expected;
let decoded = decode_meta(&encode_meta(&meta)).unwrap();
assert_eq!(decoded.dtype, expected);
}
}
#[test]
fn encode_decode_all_tiers() {
for (tier_val, expected) in [
(0u8, Tier::Tier0),
(1, Tier::Tier1),
(2, Tier::Tier2),
(3, Tier::Tier3),
] {
let key = make_key(tier_val as u128, 0);
let mut meta = sample_meta(key);
meta.tier = expected;
let decoded = decode_meta(&encode_meta(&meta)).unwrap();
assert_eq!(decoded.tier, expected);
}
}
#[test]
fn encode_decode_all_reconstruct_policies() {
for (_, expected) in [
(0u8, ReconstructPolicy::None),
(1, ReconstructPolicy::Delta),
(2, ReconstructPolicy::Factor),
] {
let key = make_key(1, 0);
let mut meta = sample_meta(key);
meta.reconstruct = expected;
let decoded = decode_meta(&encode_meta(&meta)).unwrap();
assert_eq!(decoded.reconstruct, expected);
}
}
#[test]
fn decode_too_short() {
let result = decode_meta(&[0u8; RECORD_SIZE - 1]);
assert!(
matches!(result, Err(StoreError::InvalidData)),
"expected InvalidData, got {:?}",
result.err()
);
}
#[test]
fn decode_invalid_dtype() {
let key = make_key(1, 0);
let mut encoded = encode_meta(&sample_meta(key));
encoded[20] = 255; // invalid dtype
assert!(
matches!(decode_meta(&encoded), Err(StoreError::InvalidData)),
"expected InvalidData for bad dtype"
);
}
#[test]
fn decode_invalid_tier() {
let key = make_key(1, 0);
let mut encoded = encode_meta(&sample_meta(key));
encoded[21] = 99; // invalid tier
assert!(
matches!(decode_meta(&encoded), Err(StoreError::InvalidData)),
"expected InvalidData for bad tier"
);
}
#[test]
fn decode_invalid_reconstruct() {
let key = make_key(1, 0);
let mut encoded = encode_meta(&sample_meta(key));
encoded[65] = 77; // invalid reconstruct policy
assert!(
matches!(decode_meta(&encoded), Err(StoreError::InvalidData)),
"expected InvalidData for bad reconstruct"
);
}
// -- FileBlockIO -------------------------------------------------------
#[test]
fn file_block_io_write_read() {
let dir = test_dir("bio_wr");
let mut io = FileBlockIO::new(&dir).unwrap();
let key = make_key(0xABCD, 3);
let data = vec![1u8, 2, 3, 4, 5, 6, 7, 8];
io.write_block(Tier::Tier1, key, &data).unwrap();
let mut dst = vec![0u8; 16];
let n = io.read_block(Tier::Tier1, key, &mut dst).unwrap();
assert_eq!(n, 8);
assert_eq!(&dst[..8], &data);
cleanup(&dir);
}
#[test]
fn file_block_io_write_tier0_rejected() {
let dir = test_dir("bio_t0");
let mut io = FileBlockIO::new(&dir).unwrap();
let key = make_key(1, 0);
assert_eq!(
io.write_block(Tier::Tier0, key, &[1]),
Err(StoreError::InvalidBlock)
);
cleanup(&dir);
}
#[test]
fn file_block_io_read_not_found() {
let dir = test_dir("bio_nf");
let io = FileBlockIO::new(&dir).unwrap();
let key = make_key(99, 99);
let mut dst = vec![0u8; 4];
assert_eq!(
io.read_block(Tier::Tier2, key, &mut dst),
Err(StoreError::BlockNotFound)
);
cleanup(&dir);
}
#[test]
fn file_block_io_delete() {
let dir = test_dir("bio_del");
let mut io = FileBlockIO::new(&dir).unwrap();
let key = make_key(5, 0);
io.write_block(Tier::Tier2, key, &[10, 20, 30]).unwrap();
io.delete_block(Tier::Tier2, key).unwrap();
let mut dst = vec![0u8; 4];
assert_eq!(
io.read_block(Tier::Tier2, key, &mut dst),
Err(StoreError::BlockNotFound)
);
cleanup(&dir);
}
#[test]
fn file_block_io_delete_not_found() {
let dir = test_dir("bio_del_nf");
let mut io = FileBlockIO::new(&dir).unwrap();
let key = make_key(1, 0);
assert_eq!(
io.delete_block(Tier::Tier1, key),
Err(StoreError::BlockNotFound)
);
cleanup(&dir);
}
#[test]
fn file_block_io_overwrite() {
let dir = test_dir("bio_ow");
let mut io = FileBlockIO::new(&dir).unwrap();
let key = make_key(1, 0);
io.write_block(Tier::Tier1, key, &[1, 2, 3]).unwrap();
io.write_block(Tier::Tier1, key, &[4, 5, 6, 7]).unwrap();
let mut dst = vec![0u8; 8];
let n = io.read_block(Tier::Tier1, key, &mut dst).unwrap();
assert_eq!(n, 4);
assert_eq!(&dst[..4], &[4, 5, 6, 7]);
cleanup(&dir);
}
#[test]
fn file_block_io_multiple_tiers() {
let dir = test_dir("bio_mt");
let mut io = FileBlockIO::new(&dir).unwrap();
let key = make_key(1, 0);
io.write_block(Tier::Tier1, key, &[1]).unwrap();
io.write_block(Tier::Tier2, key, &[2]).unwrap();
io.write_block(Tier::Tier3, key, &[3]).unwrap();
let mut dst = [0u8; 1];
let n = io.read_block(Tier::Tier1, key, &mut dst).unwrap();
assert_eq!(n, 1);
assert_eq!(dst[0], 1);
let n = io.read_block(Tier::Tier2, key, &mut dst).unwrap();
assert_eq!(n, 1);
assert_eq!(dst[0], 2);
let n = io.read_block(Tier::Tier3, key, &mut dst).unwrap();
assert_eq!(n, 1);
assert_eq!(dst[0], 3);
cleanup(&dir);
}
#[test]
fn file_block_io_path_format() {
let dir = test_dir("bio_path");
let io = FileBlockIO::new(&dir).unwrap();
let key = make_key(0xFF, 42);
let path = io.block_path(Tier::Tier1, key);
let expected = dir
.join("tier1")
.join("000000000000000000000000000000ff_42.bin");
assert_eq!(path, expected);
cleanup(&dir);
}
// -- FileMetaLog -------------------------------------------------------
#[test]
fn file_meta_log_append_get() {
let dir = test_dir("ml_ag");
let mut log = FileMetaLog::new(&dir).unwrap();
let key = make_key(1, 0);
let meta = sample_meta(key);
log.append(&meta).unwrap();
let retrieved = log.get(key).unwrap();
assert_eq!(retrieved.key, key);
assert_eq!(retrieved.created_at, 1000);
assert_eq!(log.len(), 1);
cleanup(&dir);
}
#[test]
fn file_meta_log_get_missing() {
let dir = test_dir("ml_miss");
let log = FileMetaLog::new(&dir).unwrap();
assert!(log.get(make_key(99, 0)).is_none());
cleanup(&dir);
}
#[test]
fn file_meta_log_upsert() {
let dir = test_dir("ml_ups");
let mut log = FileMetaLog::new(&dir).unwrap();
let key = make_key(1, 0);
let mut meta = sample_meta(key);
meta.access_count = 10;
log.append(&meta).unwrap();
meta.access_count = 20;
log.append(&meta).unwrap();
// In-memory should reflect the latest write.
let retrieved = log.get(key).unwrap();
assert_eq!(retrieved.access_count, 20);
assert_eq!(log.len(), 1);
cleanup(&dir);
}
#[test]
fn file_meta_log_iter() {
let dir = test_dir("ml_iter");
let mut log = FileMetaLog::new(&dir).unwrap();
for i in 0..5u32 {
let key = make_key(i as u128, 0);
log.append(&sample_meta(key)).unwrap();
}
let entries: Vec<_> = log.iter().collect();
assert_eq!(entries.len(), 5);
cleanup(&dir);
}
#[test]
fn file_meta_log_persistence_across_opens() {
let dir = test_dir("ml_persist");
let key1 = make_key(1, 0);
let key2 = make_key(2, 5);
// First open: write two records.
{
let mut log = FileMetaLog::new(&dir).unwrap();
log.append(&sample_meta(key1)).unwrap();
let mut meta2 = sample_meta(key2);
meta2.tier = Tier::Tier3;
meta2.bits = 3;
meta2.lineage_parent = Some(0x42);
log.append(&meta2).unwrap();
assert_eq!(log.len(), 2);
}
// Second open: records should be recovered from disk.
{
let log = FileMetaLog::new(&dir).unwrap();
assert_eq!(log.len(), 2);
let r1 = log.get(key1).unwrap();
assert_eq!(r1.tier, Tier::Tier1);
let r2 = log.get(key2).unwrap();
assert_eq!(r2.tier, Tier::Tier3);
assert_eq!(r2.lineage_parent, Some(0x42));
}
cleanup(&dir);
}
#[test]
fn file_meta_log_replay_last_wins() {
let dir = test_dir("ml_lw");
let key = make_key(1, 0);
// Write two versions of the same key.
{
let mut log = FileMetaLog::new(&dir).unwrap();
let mut meta = sample_meta(key);
meta.access_count = 100;
log.append(&meta).unwrap();
meta.access_count = 200;
log.append(&meta).unwrap();
}
// Reopen: last record should win during replay.
{
let log = FileMetaLog::new(&dir).unwrap();
assert_eq!(log.len(), 1);
let retrieved = log.get(key).unwrap();
assert_eq!(retrieved.access_count, 200);
}
cleanup(&dir);
}
#[test]
fn file_meta_log_empty_on_fresh_dir() {
let dir = test_dir("ml_empty");
let log = FileMetaLog::new(&dir).unwrap();
assert!(log.is_empty());
assert_eq!(log.len(), 0);
assert_eq!(log.iter().count(), 0);
cleanup(&dir);
}
// -- Integration: FileBlockIO + FileMetaLog ----------------------------
#[test]
fn integration_block_io_and_meta_log() {
let dir = test_dir("integ");
let mut io = FileBlockIO::new(&dir).unwrap();
let mut log = FileMetaLog::new(&dir).unwrap();
let key = make_key(0x1234, 0);
let block_data = vec![0xFFu8; 256];
// Write block and metadata.
io.write_block(Tier::Tier1, key, &block_data).unwrap();
let mut meta = sample_meta(key);
meta.block_bytes = 256;
log.append(&meta).unwrap();
// Read back and verify.
let mut dst = vec![0u8; 512];
let n = io.read_block(Tier::Tier1, key, &mut dst).unwrap();
assert_eq!(n, 256);
assert!(dst[..256].iter().all(|&b| b == 0xFF));
let retrieved = log.get(key).unwrap();
assert_eq!(retrieved.block_bytes, 256);
cleanup(&dir);
}
#[test]
fn record_size_constant_matches() {
// Verify that RECORD_SIZE matches the actual encoded size.
let meta = sample_meta(make_key(0, 0));
let encoded = encode_meta(&meta);
assert_eq!(encoded.len(), RECORD_SIZE);
}
}

View File

@@ -0,0 +1,834 @@
//! Groupwise symmetric quantization with f16 scales.
//!
//! For each group of `group_len` values:
//! - `scale = max(|v_i|) / qmax`
//! - `q_i = round(v_i / scale)`, clamped to `[-qmax, +qmax]`
//! - `u_i = q_i + qmax` (bias to unsigned for packing)
use crate::bitpack::qmax_from_bits;
use crate::f16;
/// Compute f16 group scales for a frame.
///
/// Returns one f16-encoded scale per group of `group_len` elements.
/// Each scale is `max(|v|) / qmax` for that group, stored as IEEE 754 half-precision.
#[inline]
pub fn compute_scales(frame: &[f32], group_len: usize, bits: u8) -> Vec<u16> {
let qmax = qmax_from_bits(bits);
if qmax == 0 {
return Vec::new();
}
let qmax_f = qmax as f32;
let num_groups = frame.len().div_ceil(group_len);
let mut scales = Vec::with_capacity(num_groups);
for chunk in frame.chunks(group_len) {
let mut max_abs = 0.0f32;
for &v in chunk {
if v.is_finite() {
let a = v.abs();
if a > max_abs {
max_abs = a;
}
}
}
let scale = if max_abs == 0.0 {
0.0
} else {
max_abs / qmax_f
};
scales.push(f16::f32_to_f16_bits(scale));
}
scales
}
/// Pre-convert f16 scales to f32 for hot-path use.
#[inline]
pub fn scales_to_f32(scales_f16: &[u16]) -> Vec<f32> {
scales_f16
.iter()
.map(|&s| f16::f16_bits_to_f32(s))
.collect()
}
/// Check if a frame fits within existing scales (within drift tolerance).
///
/// Uses pre-converted f32 scales to avoid repeated f16 conversion.
/// Returns `false` if any group's max absolute value exceeds
/// `scale * qmax * drift_factor`.
pub fn frame_fits_scales_f32(
frame: &[f32],
scales_f32: &[f32],
group_len: usize,
bits: u8,
drift_factor: f32,
) -> bool {
let qmax = qmax_from_bits(bits);
if qmax == 0 || scales_f32.is_empty() {
return false;
}
let qmax_f = qmax as f32;
for (group_idx, chunk) in frame.chunks(group_len).enumerate() {
if group_idx >= scales_f32.len() {
return false;
}
let allowed = scales_f32[group_idx] * qmax_f * drift_factor;
for &v in chunk {
if v.is_finite() && v.abs() > allowed {
return false;
}
}
}
true
}
/// Quantize a frame using pre-computed f32 scales and pack into bitstream.
///
/// Appends packed bytes to `out`. Pre-reserves the expected output size
/// to avoid reallocations.
///
/// For 8-bit quantization, writes bytes directly without bit accumulation
/// since each quantized value maps 1:1 to a u8.
#[inline]
pub fn quantize_and_pack_f32(
frame: &[f32],
scales_f32: &[f32],
group_len: usize,
bits: u8,
out: &mut Vec<u8>,
) {
let qmax = qmax_from_bits(bits);
if qmax == 0 {
return;
}
// Fast path: 8-bit quantization writes bytes directly, no bit accumulator.
if bits == 8 {
out.reserve(frame.len());
for (group_idx, chunk) in frame.chunks(group_len).enumerate() {
let scale = if group_idx < scales_f32.len() {
scales_f32[group_idx]
} else {
0.0
};
let inv_scale = if scale == 0.0 { 0.0 } else { 1.0 / scale };
for &v in chunk {
let mut q: i32 = 0;
if v.is_finite() {
let scaled = v * inv_scale;
q = if scaled >= 0.0 {
(scaled + 0.5) as i32
} else {
(scaled - 0.5) as i32
};
q = q.clamp(-127, 127);
}
out.push((q + 127) as u8);
}
}
return;
}
// Fast path: 5-bit quantization packs 8 values into 5 bytes.
// 8 values * 5 bits = 40 bits = 5 bytes exactly, avoiding the bit accumulator.
// LSB-first packing layout for 8 values in 5 bytes:
// byte0 = v0 | (v1 << 5)
// byte1 = (v1 >> 3) | (v2 << 2) | (v3 << 7)
// byte2 = (v3 >> 1) | (v4 << 4)
// byte3 = (v4 >> 4) | (v5 << 1) | (v6 << 6)
// byte4 = (v6 >> 2) | (v7 << 3)
#[inline]
fn pack_5bit_group(chunk: &[f32], inv_scale: f32, out: &mut Vec<u8>) {
let quantize = |v: f32| -> u32 {
let mut q: i32 = 0;
if v.is_finite() {
let scaled = v * inv_scale;
q = if scaled >= 0.0 {
(scaled + 0.5) as i32
} else {
(scaled - 0.5) as i32
};
q = q.clamp(-15, 15);
}
(q + 15) as u32
};
let v0 = quantize(chunk[0]);
let v1 = quantize(chunk[1]);
let v2 = quantize(chunk[2]);
let v3 = quantize(chunk[3]);
let v4 = quantize(chunk[4]);
let v5 = quantize(chunk[5]);
let v6 = quantize(chunk[6]);
let v7 = quantize(chunk[7]);
out.push((v0 | (v1 << 5)) as u8);
out.push(((v1 >> 3) | (v2 << 2) | (v3 << 7)) as u8);
out.push(((v3 >> 1) | (v4 << 4)) as u8);
out.push(((v4 >> 4) | (v5 << 1) | (v6 << 6)) as u8);
out.push(((v6 >> 2) | (v7 << 3)) as u8);
}
if bits == 5 {
let needed_bytes = (frame.len() * 5).div_ceil(8);
out.reserve(needed_bytes);
let mut acc: u64 = 0;
let mut acc_bits: u32 = 0;
for (group_idx, chunk) in frame.chunks(group_len).enumerate() {
let scale = if group_idx < scales_f32.len() {
scales_f32[group_idx]
} else {
0.0
};
let inv_scale = if scale == 0.0 { 0.0 } else { 1.0 / scale };
let mut i = 0;
// Process 8 values at a time into 5 bytes when byte-aligned
while acc_bits == 0 && i + 8 <= chunk.len() {
pack_5bit_group(&chunk[i..i + 8], inv_scale, out);
i += 8;
}
// Remainder (or misaligned) with bit accumulator
while i < chunk.len() {
let mut q: i32 = 0;
if chunk[i].is_finite() {
let scaled = chunk[i] * inv_scale;
q = if scaled >= 0.0 {
(scaled + 0.5) as i32
} else {
(scaled - 0.5) as i32
};
q = q.clamp(-15, 15);
}
let u = (q + 15) as u32;
acc |= (u as u64) << acc_bits;
acc_bits += 5;
while acc_bits >= 8 {
out.push((acc & 0xFF) as u8);
acc >>= 8;
acc_bits -= 8;
}
i += 1;
}
}
if acc_bits > 0 {
out.push((acc & 0xFF) as u8);
}
return;
}
// Generic path for sub-byte bit widths.
let qmax_i = qmax;
let bias = qmax;
let bits_u32 = bits as u32;
let needed_bytes = (frame.len() * bits as usize).div_ceil(8);
out.reserve(needed_bytes);
let mut acc: u64 = 0;
let mut acc_bits: u32 = 0;
for (group_idx, chunk) in frame.chunks(group_len).enumerate() {
let scale = if group_idx < scales_f32.len() {
scales_f32[group_idx]
} else {
0.0
};
let inv_scale = if scale == 0.0 { 0.0 } else { 1.0 / scale };
for &v in chunk {
let mut q: i32 = 0;
if v.is_finite() {
let scaled = v * inv_scale;
q = if scaled >= 0.0 {
(scaled + 0.5) as i32
} else {
(scaled - 0.5) as i32
};
q = q.clamp(-qmax_i, qmax_i);
}
let u = (q + bias) as u32;
acc |= (u as u64) << acc_bits;
acc_bits += bits_u32;
while acc_bits >= 8 {
out.push((acc & 0xFF) as u8);
acc >>= 8;
acc_bits -= 8;
}
}
}
if acc_bits > 0 {
out.push((acc & 0xFF) as u8);
}
}
/// Dequantize packed codes using f32 scales, writing f32 values.
///
/// Iterates by frame then by group to avoid per-value modulo/division
/// and caches the f32 scale per group.
///
/// For 8-bit data, reads bytes directly without bit accumulation.
#[inline]
pub fn dequantize_f32(
data: &[u8],
scales_f32: &[f32],
group_len: usize,
bits: u8,
tensor_len: usize,
frame_count: usize,
out: &mut Vec<f32>,
) {
let qmax = qmax_from_bits(bits);
if qmax == 0 {
return;
}
let total = tensor_len * frame_count;
out.resize(total, 0.0);
// Fast path: 8-bit dequantization reads bytes directly, no bit accumulator.
if bits == 8 {
let mut out_idx = 0usize;
let mut byte_idx = 0usize;
for _frame in 0..frame_count {
let mut pos = 0usize;
let mut group_idx = 0usize;
while pos < tensor_len {
let group_end = (pos + group_len).min(tensor_len);
let scale = if group_idx < scales_f32.len() {
scales_f32[group_idx]
} else {
0.0
};
while pos < group_end && byte_idx < data.len() {
let u = data[byte_idx] as i32;
let q = u - 127;
out[out_idx] = (q as f32) * scale;
out_idx += 1;
byte_idx += 1;
pos += 1;
}
group_idx += 1;
}
}
return;
}
// Fast path: 3-bit dequantization processes 8 values from 3 bytes.
// 8 values * 3 bits = 24 bits = 3 bytes exactly, avoiding the bit accumulator.
// LSB-first packing layout for 8 values in 3 bytes:
// byte0 = v0 | (v1 << 3) | ((v2 & 0x3) << 6)
// byte1 = (v2 >> 2) | (v3 << 1) | (v4 << 4) | ((v5 & 0x1) << 7)
// byte2 = (v5 >> 1) | (v6 << 2) | (v7 << 5)
if bits == 3 {
let bias = 3i32; // qmax for 3-bit
let mut out_idx = 0usize;
let mut byte_idx = 0usize;
for _frame in 0..frame_count {
let mut pos = 0usize;
let mut group_idx = 0usize;
while pos < tensor_len {
let group_end = (pos + group_len).min(tensor_len);
let scale = if group_idx < scales_f32.len() {
scales_f32[group_idx]
} else {
0.0
};
// Process 8 values at a time from 3 bytes
while pos + 8 <= group_end && byte_idx + 3 <= data.len() {
let b0 = data[byte_idx] as u32;
let b1 = data[byte_idx + 1] as u32;
let b2 = data[byte_idx + 2] as u32;
byte_idx += 3;
out[out_idx] = ((b0 & 0x7) as i32 - bias) as f32 * scale;
out[out_idx + 1] = (((b0 >> 3) & 0x7) as i32 - bias) as f32 * scale;
out[out_idx + 2] =
((((b0 >> 6) | (b1 << 2)) & 0x7) as i32 - bias) as f32 * scale;
out[out_idx + 3] = (((b1 >> 1) & 0x7) as i32 - bias) as f32 * scale;
out[out_idx + 4] = (((b1 >> 4) & 0x7) as i32 - bias) as f32 * scale;
out[out_idx + 5] =
((((b1 >> 7) | (b2 << 1)) & 0x7) as i32 - bias) as f32 * scale;
out[out_idx + 6] = (((b2 >> 2) & 0x7) as i32 - bias) as f32 * scale;
out[out_idx + 7] = (((b2 >> 5) & 0x7) as i32 - bias) as f32 * scale;
out_idx += 8;
pos += 8;
}
// Handle remaining values (< 8) with a local bit accumulator
if pos < group_end {
let remaining = group_end - pos;
let mut acc: u64 = 0;
let mut acc_bits: u32 = 0;
while acc_bits < (remaining as u32) * 3 && byte_idx < data.len() {
acc |= (data[byte_idx] as u64) << acc_bits;
acc_bits += 8;
byte_idx += 1;
}
for _ in 0..remaining {
if acc_bits < 3 {
break;
}
let u = (acc & 0x7) as i32;
acc >>= 3;
acc_bits -= 3;
out[out_idx] = (u - bias) as f32 * scale;
out_idx += 1;
pos += 1;
}
}
group_idx += 1;
}
}
return;
}
// Fast path: 7-bit dequantization processes 8 values from 7 bytes.
// 8 values * 7 bits = 56 bits = 7 bytes exactly, avoiding the bit accumulator.
// LSB-first packing layout for 8 values in 7 bytes:
// v0 = b0 & 0x7F
// v1 = ((b0 >> 7) | (b1 << 1)) & 0x7F
// v2 = ((b1 >> 6) | (b2 << 2)) & 0x7F
// v3 = ((b2 >> 5) | (b3 << 3)) & 0x7F
// v4 = ((b3 >> 4) | (b4 << 4)) & 0x7F
// v5 = ((b4 >> 3) | (b5 << 5)) & 0x7F
// v6 = ((b5 >> 2) | (b6 << 6)) & 0x7F
// v7 = (b6 >> 1) & 0x7F
if bits == 7 {
let bias = 63i32; // qmax for 7-bit
let mut out_idx = 0usize;
let mut byte_idx = 0usize;
for _frame in 0..frame_count {
let mut pos = 0usize;
let mut group_idx = 0usize;
while pos < tensor_len {
let group_end = (pos + group_len).min(tensor_len);
let scale = if group_idx < scales_f32.len() {
scales_f32[group_idx]
} else {
0.0
};
// Process 8 values at a time from 7 bytes
#[inline]
fn unpack_7bit(
out: &mut [f32],
out_idx: usize,
data: &[u8],
byte_idx: usize,
bias: i32,
scale: f32,
) {
let b0 = data[byte_idx] as u32;
let b1 = data[byte_idx + 1] as u32;
let b2 = data[byte_idx + 2] as u32;
let b3 = data[byte_idx + 3] as u32;
let b4 = data[byte_idx + 4] as u32;
let b5 = data[byte_idx + 5] as u32;
let b6 = data[byte_idx + 6] as u32;
out[out_idx] = ((b0 & 0x7F) as i32 - bias) as f32 * scale;
out[out_idx + 1] =
((((b0 >> 7) | (b1 << 1)) & 0x7F) as i32 - bias) as f32 * scale;
out[out_idx + 2] =
((((b1 >> 6) | (b2 << 2)) & 0x7F) as i32 - bias) as f32 * scale;
out[out_idx + 3] =
((((b2 >> 5) | (b3 << 3)) & 0x7F) as i32 - bias) as f32 * scale;
out[out_idx + 4] =
((((b3 >> 4) | (b4 << 4)) & 0x7F) as i32 - bias) as f32 * scale;
out[out_idx + 5] =
((((b4 >> 3) | (b5 << 5)) & 0x7F) as i32 - bias) as f32 * scale;
out[out_idx + 6] =
((((b5 >> 2) | (b6 << 6)) & 0x7F) as i32 - bias) as f32 * scale;
out[out_idx + 7] = (((b6 >> 1) & 0x7F) as i32 - bias) as f32 * scale;
}
while pos + 8 <= group_end && byte_idx + 7 <= data.len() {
unpack_7bit(out, out_idx, data, byte_idx, bias, scale);
byte_idx += 7;
out_idx += 8;
pos += 8;
}
// Handle remaining values (< 8) with a local bit accumulator
if pos < group_end {
let remaining = group_end - pos;
let mut acc: u64 = 0;
let mut acc_bits: u32 = 0;
while acc_bits < (remaining as u32) * 7 && byte_idx < data.len() {
acc |= (data[byte_idx] as u64) << acc_bits;
acc_bits += 8;
byte_idx += 1;
}
for _ in 0..remaining {
if acc_bits < 7 {
break;
}
let u = (acc & 0x7F) as i32;
acc >>= 7;
acc_bits -= 7;
out[out_idx] = (u - bias) as f32 * scale;
out_idx += 1;
pos += 1;
}
}
group_idx += 1;
}
}
return;
}
// Fast path: 5-bit dequantization processes 8 values from 5 bytes.
// 8 values * 5 bits = 40 bits = 5 bytes exactly, avoiding the bit accumulator.
// LSB-first packing layout for 8 values in 5 bytes:
// v0 = b0 & 0x1F
// v1 = ((b0 >> 5) | (b1 << 3)) & 0x1F
// v2 = (b1 >> 2) & 0x1F
// v3 = ((b1 >> 7) | (b2 << 1)) & 0x1F
// v4 = ((b2 >> 4) | (b3 << 4)) & 0x1F
// v5 = (b3 >> 1) & 0x1F
// v6 = ((b3 >> 6) | (b4 << 2)) & 0x1F
// v7 = (b4 >> 3) & 0x1F
if bits == 5 {
let bias = 15i32; // qmax for 5-bit
let mut out_idx = 0usize;
let mut byte_idx = 0usize;
for _frame in 0..frame_count {
let mut pos = 0usize;
let mut group_idx = 0usize;
while pos < tensor_len {
let group_end = (pos + group_len).min(tensor_len);
let scale = if group_idx < scales_f32.len() {
scales_f32[group_idx]
} else {
0.0
};
// Process 8 values at a time from 5 bytes
#[inline]
fn unpack_5bit(
out: &mut [f32],
out_idx: usize,
data: &[u8],
byte_idx: usize,
bias: i32,
scale: f32,
) {
let b0 = data[byte_idx] as u32;
let b1 = data[byte_idx + 1] as u32;
let b2 = data[byte_idx + 2] as u32;
let b3 = data[byte_idx + 3] as u32;
let b4 = data[byte_idx + 4] as u32;
out[out_idx] = ((b0 & 0x1F) as i32 - bias) as f32 * scale;
out[out_idx + 1] =
((((b0 >> 5) | (b1 << 3)) & 0x1F) as i32 - bias) as f32 * scale;
out[out_idx + 2] = (((b1 >> 2) & 0x1F) as i32 - bias) as f32 * scale;
out[out_idx + 3] =
((((b1 >> 7) | (b2 << 1)) & 0x1F) as i32 - bias) as f32 * scale;
out[out_idx + 4] =
((((b2 >> 4) | (b3 << 4)) & 0x1F) as i32 - bias) as f32 * scale;
out[out_idx + 5] = (((b3 >> 1) & 0x1F) as i32 - bias) as f32 * scale;
out[out_idx + 6] =
((((b3 >> 6) | (b4 << 2)) & 0x1F) as i32 - bias) as f32 * scale;
out[out_idx + 7] = (((b4 >> 3) & 0x1F) as i32 - bias) as f32 * scale;
}
while pos + 8 <= group_end && byte_idx + 5 <= data.len() {
unpack_5bit(out, out_idx, data, byte_idx, bias, scale);
byte_idx += 5;
out_idx += 8;
pos += 8;
}
// Handle remaining values (< 8) with a local bit accumulator
if pos < group_end {
let remaining = group_end - pos;
let mut acc: u64 = 0;
let mut acc_bits: u32 = 0;
while acc_bits < (remaining as u32) * 5 && byte_idx < data.len() {
acc |= (data[byte_idx] as u64) << acc_bits;
acc_bits += 8;
byte_idx += 1;
}
for _ in 0..remaining {
if acc_bits < 5 {
break;
}
let u = (acc & 0x1F) as i32;
acc >>= 5;
acc_bits -= 5;
out[out_idx] = (u - bias) as f32 * scale;
out_idx += 1;
pos += 1;
}
}
group_idx += 1;
}
}
return;
}
// Generic path for sub-byte bit widths.
let bias = qmax;
let bits_u32 = bits as u32;
let mask = (1u64 << bits_u32) - 1;
let mut acc: u64 = 0;
let mut acc_bits: u32 = 0;
let mut byte_idx = 0usize;
let mut out_idx = 0usize;
for _frame in 0..frame_count {
let mut pos = 0usize;
let mut group_idx = 0usize;
while pos < tensor_len {
let group_end = (pos + group_len).min(tensor_len);
let scale = if group_idx < scales_f32.len() {
scales_f32[group_idx]
} else {
0.0
};
while pos < group_end {
while acc_bits < bits_u32 && byte_idx < data.len() {
acc |= (data[byte_idx] as u64) << acc_bits;
acc_bits += 8;
byte_idx += 1;
}
if acc_bits < bits_u32 {
return;
}
let u = (acc & mask) as u32;
acc >>= bits_u32;
acc_bits -= bits_u32;
let q = (u as i32) - bias;
out[out_idx] = (q as f32) * scale;
out_idx += 1;
pos += 1;
}
group_idx += 1;
}
}
}
// --- Legacy API (delegates to f32 variants) ---
/// Check if a frame fits within existing f16 scales (within drift tolerance).
pub fn frame_fits_scales(
frame: &[f32],
scales: &[u16],
group_len: usize,
bits: u8,
drift_factor: f32,
) -> bool {
let scales_f32 = scales_to_f32(scales);
frame_fits_scales_f32(frame, &scales_f32, group_len, bits, drift_factor)
}
/// Quantize a frame using pre-computed f16 scales and pack into bitstream.
pub fn quantize_and_pack(
frame: &[f32],
scales: &[u16],
group_len: usize,
bits: u8,
out: &mut Vec<u8>,
) {
let scales_f32 = scales_to_f32(scales);
quantize_and_pack_f32(frame, &scales_f32, group_len, bits, out)
}
/// Dequantize packed codes using f16 scales, writing f32 values.
pub fn dequantize(
data: &[u8],
scales: &[u16],
group_len: usize,
bits: u8,
tensor_len: usize,
frame_count: usize,
out: &mut Vec<f32>,
) {
let scales_f32 = scales_to_f32(scales);
dequantize_f32(
data,
&scales_f32,
group_len,
bits,
tensor_len,
frame_count,
out,
)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_quantize_roundtrip_8bit() {
let frame: Vec<f32> = (0..128).map(|i| (i as f32 - 64.0) * 0.1).collect();
let scales = compute_scales(&frame, 64, 8);
let mut packed = Vec::new();
quantize_and_pack(&frame, &scales, 64, 8, &mut packed);
let mut decoded = Vec::new();
dequantize(&packed, &scales, 64, 8, frame.len(), 1, &mut decoded);
assert_eq!(decoded.len(), frame.len());
for (i, (&orig, &dec)) in frame.iter().zip(decoded.iter()).enumerate() {
let err = (orig - dec).abs();
let max_err = if orig.abs() > 0.01 {
orig.abs() * 0.02
} else {
0.1
};
assert!(err < max_err, "i={i}, orig={orig}, dec={dec}, err={err}");
}
}
#[test]
fn test_quantize_roundtrip_3bit() {
let frame: Vec<f32> = (0..64).map(|i| (i as f32 - 32.0) * 0.5).collect();
let scales = compute_scales(&frame, 64, 3);
let mut packed = Vec::new();
quantize_and_pack(&frame, &scales, 64, 3, &mut packed);
let mut decoded = Vec::new();
dequantize(&packed, &scales, 64, 3, frame.len(), 1, &mut decoded);
let max_val = frame.iter().map(|v| v.abs()).fold(0.0f32, f32::max);
for (&orig, &dec) in frame.iter().zip(decoded.iter()) {
let err = (orig - dec).abs();
assert!(err < max_val * 0.35, "orig={orig}, dec={dec}, err={err}");
}
}
#[test]
fn test_quantize_roundtrip_5bit() {
let frame: Vec<f32> = (0..256).map(|i| (i as f32 - 128.0) * 0.05).collect();
let scales = compute_scales(&frame, 64, 5);
let mut packed = Vec::new();
quantize_and_pack(&frame, &scales, 64, 5, &mut packed);
let mut decoded = Vec::new();
dequantize(&packed, &scales, 64, 5, frame.len(), 1, &mut decoded);
let max_val = frame.iter().map(|v| v.abs()).fold(0.0f32, f32::max);
for (&orig, &dec) in frame.iter().zip(decoded.iter()) {
let err = (orig - dec).abs();
assert!(err < max_val * 0.08, "orig={orig}, dec={dec}, err={err}");
}
}
#[test]
fn test_quantize_roundtrip_7bit() {
let frame: Vec<f32> = (0..256).map(|i| (i as f32 - 128.0) * 0.05).collect();
let scales = compute_scales(&frame, 64, 7);
let mut packed = Vec::new();
quantize_and_pack(&frame, &scales, 64, 7, &mut packed);
let mut decoded = Vec::new();
dequantize(&packed, &scales, 64, 7, frame.len(), 1, &mut decoded);
for (i, (&orig, &dec)) in frame.iter().zip(decoded.iter()).enumerate() {
let err = (orig - dec).abs();
let max_err = if orig.abs() > 0.01 {
orig.abs() * 0.02
} else {
0.1
};
assert!(err < max_err, "i={i}, orig={orig}, dec={dec}, err={err}");
}
}
#[test]
fn test_drift_detection() {
let frame1: Vec<f32> = vec![1.0; 64];
let frame2: Vec<f32> = vec![1.05; 64];
let frame3: Vec<f32> = vec![2.0; 64];
let scales = compute_scales(&frame1, 64, 8);
let drift_factor = 1.0 + 26.0 / 256.0;
assert!(frame_fits_scales(&frame2, &scales, 64, 8, drift_factor));
assert!(!frame_fits_scales(&frame3, &scales, 64, 8, drift_factor));
}
#[test]
fn test_zero_frame() {
let frame = vec![0.0f32; 128];
let scales = compute_scales(&frame, 64, 8);
let mut packed = Vec::new();
quantize_and_pack(&frame, &scales, 64, 8, &mut packed);
let mut decoded = Vec::new();
dequantize(&packed, &scales, 64, 8, 128, 1, &mut decoded);
for &v in &decoded {
assert_eq!(v, 0.0);
}
}
#[test]
fn test_non_finite_values() {
let mut frame = vec![1.0f32; 64];
frame[10] = f32::NAN;
frame[20] = f32::INFINITY;
frame[30] = f32::NEG_INFINITY;
let scales = compute_scales(&frame, 64, 8);
let mut packed = Vec::new();
quantize_and_pack(&frame, &scales, 64, 8, &mut packed);
let mut decoded = Vec::new();
dequantize(&packed, &scales, 64, 8, 64, 1, &mut decoded);
assert_eq!(decoded[10], 0.0);
assert_eq!(decoded[20], 0.0);
assert_eq!(decoded[30], 0.0);
assert!((decoded[0] - 1.0).abs() < 0.02);
}
#[test]
fn test_single_element_group() {
let frame = vec![3.14f32; 16];
let scales = compute_scales(&frame, 1, 8);
assert_eq!(scales.len(), 16);
let mut packed = Vec::new();
quantize_and_pack(&frame, &scales, 1, 8, &mut packed);
let mut decoded = Vec::new();
dequantize(&packed, &scales, 1, 8, 16, 1, &mut decoded);
for (i, &v) in decoded.iter().enumerate() {
let err = (v - 3.14).abs();
assert!(err < 0.03, "i={i} v={v} err={err}");
}
}
#[test]
fn test_compression_ratio() {
let frame = vec![1.0f32; 512];
for &(bits, min_ratio) in &[(8u8, 3.5f32), (7, 4.0), (5, 5.5), (3, 8.5)] {
let scales = compute_scales(&frame, 64, bits);
let mut packed = Vec::new();
quantize_and_pack(&frame, &scales, 64, bits, &mut packed);
let raw_bytes = frame.len() * 4;
let compressed = packed.len() + scales.len() * 2;
let ratio = raw_bytes as f32 / compressed as f32;
assert!(
ratio >= min_ratio,
"bits={bits}: ratio {ratio:.2}x < expected {min_ratio}x"
);
}
}
}

View File

@@ -0,0 +1,335 @@
//! Segment binary format: encode and decode.
//!
//! Format (little-endian):
//!
//! ```text
//! [magic:4][version:1][bits:1][group_len:4][tensor_len:4][frames:4]
//! [scale_count:4][scales:2*S][data_len:4][data:D]
//! ```
//!
//! Magic: `0x43545154` ("TQTC" in LE). Header is 26 bytes before scales.
use crate::quantizer;
/// Segment magic number: `"TQTC"` in little-endian.
pub const MAGIC: u32 = 0x4354_5154;
/// Current segment format version.
pub const VERSION: u8 = 1;
/// Minimum valid segment size in bytes (header fields + data_len, no scales/data).
pub const HEADER_SIZE: usize = 26;
/// Encode a segment from metadata, scales, and packed data.
pub fn encode(
bits: u8,
group_len: u32,
tensor_len: u32,
frame_count: u32,
scales: &[u16],
data: &[u8],
out: &mut Vec<u8>,
) {
out.clear();
let estimated = HEADER_SIZE + scales.len() * 2 + data.len();
out.reserve(estimated);
// Header
out.extend_from_slice(&MAGIC.to_le_bytes());
out.push(VERSION);
out.push(bits);
out.extend_from_slice(&group_len.to_le_bytes());
out.extend_from_slice(&tensor_len.to_le_bytes());
out.extend_from_slice(&frame_count.to_le_bytes());
// Scales
let scale_count = scales.len() as u32;
out.extend_from_slice(&scale_count.to_le_bytes());
for &s in scales {
out.extend_from_slice(&s.to_le_bytes());
}
// Data
let data_len = data.len() as u32;
out.extend_from_slice(&data_len.to_le_bytes());
out.extend_from_slice(data);
}
/// Decoded segment header.
#[derive(Debug, Clone)]
pub struct SegmentHeader {
pub bits: u8,
pub group_len: u32,
pub tensor_len: u32,
pub frame_count: u32,
pub scale_count: u32,
}
/// Decode a segment, returning all frames as f32 values.
pub fn decode(segment: &[u8], out: &mut Vec<f32>) {
out.clear();
if segment.len() < HEADER_SIZE {
return;
}
let mut off = 0;
let magic = read_u32_le(segment, &mut off);
if magic != MAGIC {
return;
}
let version = segment[off];
off += 1;
if version != VERSION {
return;
}
let bits = segment[off];
off += 1;
let group_len = read_u32_le(segment, &mut off);
let tensor_len = read_u32_le(segment, &mut off);
let frame_count = read_u32_le(segment, &mut off);
let scale_count = read_u32_le(segment, &mut off);
// Read scales
let scales_end = off + (scale_count as usize) * 2;
if scales_end > segment.len() {
return;
}
let mut scales = Vec::with_capacity(scale_count as usize);
for _ in 0..scale_count {
scales.push(read_u16_le(segment, &mut off));
}
// Read data
if off + 4 > segment.len() {
return;
}
let data_len = read_u32_le(segment, &mut off) as usize;
if off + data_len > segment.len() {
return;
}
let data = &segment[off..off + data_len];
// Convert scales to f32 once, then dequantize via the optimized path
let scales_f32 = quantizer::scales_to_f32(&scales);
quantizer::dequantize_f32(
data,
&scales_f32,
group_len as usize,
bits,
tensor_len as usize,
frame_count as usize,
out,
);
}
/// Parse only the segment header (no data decoding).
pub fn parse_header(segment: &[u8]) -> Option<SegmentHeader> {
if segment.len() < HEADER_SIZE {
return None;
}
let mut off = 0;
let magic = read_u32_le(segment, &mut off);
if magic != MAGIC {
return None;
}
let version = segment[off];
off += 1;
if version != VERSION {
return None;
}
let bits = segment[off];
off += 1;
let group_len = read_u32_le(segment, &mut off);
let tensor_len = read_u32_le(segment, &mut off);
let frame_count = read_u32_le(segment, &mut off);
let scale_count = read_u32_le(segment, &mut off);
Some(SegmentHeader {
bits,
group_len,
tensor_len,
frame_count,
scale_count,
})
}
/// Compute the compression ratio for a segment: raw f32 bytes / segment bytes.
///
/// Returns `0.0` if the segment is empty or has no frames.
pub fn compression_ratio(segment: &[u8]) -> f32 {
match parse_header(segment) {
Some(h) if h.frame_count > 0 => {
let raw = h.tensor_len as usize * h.frame_count as usize * 4;
raw as f32 / segment.len() as f32
}
_ => 0.0,
}
}
/// Decode a single frame by index from a segment.
///
/// Returns `None` if the segment is invalid or `frame_idx` is out of range.
pub fn decode_single_frame(segment: &[u8], frame_idx: usize) -> Option<Vec<f32>> {
let header = parse_header(segment)?;
if frame_idx >= header.frame_count as usize {
return None;
}
// Skip past the fixed header fields (magic + version + bits + group_len +
// tensor_len + frame_count + scale_count = 4+1+1+4+4+4+4 = 22 bytes).
let mut off = 22usize;
let scale_count = header.scale_count as usize;
// Read scales
let scales_end = off + scale_count * 2;
if scales_end > segment.len() {
return None;
}
let mut scales_f16 = Vec::with_capacity(scale_count);
for _ in 0..scale_count {
scales_f16.push(read_u16_le(segment, &mut off));
}
let scales_f32 = quantizer::scales_to_f32(&scales_f16);
// Read data section
if off + 4 > segment.len() {
return None;
}
let data_len = read_u32_le(segment, &mut off) as usize;
if off + data_len > segment.len() {
return None;
}
let data = &segment[off..off + data_len];
// Compute byte offset for the requested frame
let tensor_len = header.tensor_len as usize;
let bits = header.bits;
let bits_per_frame = tensor_len * bits as usize;
let bytes_per_frame = bits_per_frame.div_ceil(8);
let frame_start = frame_idx * bytes_per_frame;
if frame_start + bytes_per_frame > data.len() {
return None;
}
let frame_data = &data[frame_start..frame_start + bytes_per_frame];
let mut out = Vec::new();
quantizer::dequantize_f32(
frame_data,
&scales_f32,
header.group_len as usize,
bits,
tensor_len,
1,
&mut out,
);
Some(out)
}
#[inline]
fn read_u32_le(bytes: &[u8], offset: &mut usize) -> u32 {
let o = *offset;
let arr = [bytes[o], bytes[o + 1], bytes[o + 2], bytes[o + 3]];
*offset = o + 4;
u32::from_le_bytes(arr)
}
fn read_u16_le(bytes: &[u8], offset: &mut usize) -> u16 {
let o = *offset;
let arr = [bytes[o], bytes[o + 1]];
*offset = o + 2;
u16::from_le_bytes(arr)
}
#[cfg(test)]
mod tests {
use super::*;
use crate::quantizer;
#[test]
fn test_encode_decode_roundtrip() {
let frame: Vec<f32> = (0..128).map(|i| (i as f32 - 64.0) * 0.1).collect();
let group_len = 64usize;
let bits = 8u8;
let scales = quantizer::compute_scales(&frame, group_len, bits);
let mut packed = Vec::new();
quantizer::quantize_and_pack(&frame, &scales, group_len, bits, &mut packed);
let mut seg = Vec::new();
encode(
bits,
group_len as u32,
frame.len() as u32,
1,
&scales,
&packed,
&mut seg,
);
let mut decoded = Vec::new();
decode(&seg, &mut decoded);
assert_eq!(decoded.len(), frame.len());
for (i, (&orig, &dec)) in frame.iter().zip(decoded.iter()).enumerate() {
let err = (orig - dec).abs();
assert!(err < 0.1, "i={i} orig={orig} dec={dec} err={err}");
}
}
#[test]
fn test_magic_validation() {
let mut decoded = Vec::new();
decode(&[0, 0, 0, 0], &mut decoded);
assert!(decoded.is_empty()); // Wrong magic
}
#[test]
fn test_parse_header() {
let frame = vec![1.0f32; 64];
let scales = quantizer::compute_scales(&frame, 64, 7);
let mut packed = Vec::new();
quantizer::quantize_and_pack(&frame, &scales, 64, 7, &mut packed);
let mut seg = Vec::new();
encode(7, 64, 64, 1, &scales, &packed, &mut seg);
let header = parse_header(&seg).unwrap();
assert_eq!(header.bits, 7);
assert_eq!(header.group_len, 64);
assert_eq!(header.tensor_len, 64);
assert_eq!(header.frame_count, 1);
}
#[test]
fn test_multi_frame_roundtrip() {
let group_len = 32usize;
let bits = 5u8;
let tensor_len = 64;
let frame1: Vec<f32> = (0..tensor_len).map(|i| (i as f32) * 0.1).collect();
let frame2: Vec<f32> = (0..tensor_len).map(|i| (i as f32) * 0.09).collect();
let scales = quantizer::compute_scales(&frame1, group_len, bits);
let mut packed = Vec::new();
quantizer::quantize_and_pack(&frame1, &scales, group_len, bits, &mut packed);
quantizer::quantize_and_pack(&frame2, &scales, group_len, bits, &mut packed);
let mut seg = Vec::new();
encode(
bits,
group_len as u32,
tensor_len as u32,
2,
&scales,
&packed,
&mut seg,
);
let mut decoded = Vec::new();
decode(&seg, &mut decoded);
assert_eq!(decoded.len(), tensor_len * 2);
}
}

File diff suppressed because it is too large Load Diff

View File

@@ -0,0 +1,877 @@
//! WASM/C FFI for the block-based temporal tensor store (ADR-022).
//!
//! Exports `extern "C"` functions prefixed with `tts_` for:
//! - Store lifecycle (`tts_init`)
//! - Block ingest and read (`tts_put`, `tts_get`)
//! - Access tracking (`tts_touch`)
//! - Maintenance (`tts_tick`, `tts_evict`)
//! - Statistics (`tts_stats`, `tts_block_count`, `tts_tier_count`)
//!
//! Coexists with `ffi.rs` which exports `ttc_*` functions for the
//! frame-based compressor.
use std::collections::HashMap;
use crate::quantizer;
use crate::segment;
// ── Error codes ──────────────────────────────────────────────────────
#[allow(dead_code)]
const ERR_NOT_INITIALIZED: i32 = -1;
const ERR_NULL_POINTER: i32 = -2;
const ERR_INVALID_CONFIG: i32 = -3;
const ERR_BLOCK_NOT_FOUND: i32 = -4;
const ERR_BUFFER_TOO_SMALL: i32 = -5;
const ERR_EMPTY_DATA: i32 = -6;
// ── Types ────────────────────────────────────────────────────────────
// These mirror the types defined in store.rs and tiering.rs which are
// being written in parallel. Once those modules land, these can be
// replaced with `use crate::store::*` / `use crate::tiering::*`.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
struct BlockKey {
tensor_id: u128,
block_index: u32,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
#[repr(u8)]
enum Tier {
Hot = 0,
Warm = 1,
Cool = 2,
Cold = 3,
}
impl Tier {
fn from_u8(v: u8) -> Option<Self> {
match v {
0 => Some(Tier::Hot),
1 => Some(Tier::Warm),
2 => Some(Tier::Cool),
3 => Some(Tier::Cold),
_ => None,
}
}
/// Quantization bit-width for this tier.
fn bits(self) -> u8 {
match self {
Tier::Hot => 8,
Tier::Warm => 7,
Tier::Cool => 5,
Tier::Cold => 3,
}
}
}
#[derive(Clone, Debug)]
struct BlockMeta {
tier: Tier,
access_count: u32,
last_access_ts: u64,
ema_score: f32,
/// Original f32 count; used when re-tiering to size the decode buffer.
#[allow(dead_code)]
element_count: usize,
}
/// Binary config layout (little-endian, 45 bytes):
/// ```text
/// [block_bytes:u32][alpha:f32][tau:f32][w_ema:f32][w_pop:f32][w_rec:f32]
/// [t1:f32][t2:f32][t3:f32][hysteresis:f32][min_residency:u32][max_delta_chain:u8]
/// ```
#[derive(Clone, Debug)]
struct TierConfig {
block_bytes: u32,
alpha: f32,
tau: f32,
w_ema: f32,
w_pop: f32,
w_rec: f32,
t1: f32,
t2: f32,
t3: f32,
hysteresis: f32,
min_residency: u32,
max_delta_chain: u8,
}
const CONFIG_BINARY_LEN: usize = 45;
impl Default for TierConfig {
fn default() -> Self {
Self {
block_bytes: 4096,
alpha: 0.3,
tau: 100.0,
w_ema: 0.5,
w_pop: 0.3,
w_rec: 0.2,
t1: 0.8,
t2: 0.5,
t3: 0.2,
hysteresis: 0.05,
min_residency: 10,
max_delta_chain: 4,
}
}
}
impl TierConfig {
fn from_bytes(bytes: &[u8]) -> Option<Self> {
if bytes.len() < CONFIG_BINARY_LEN {
return None;
}
let mut off = 0usize;
let block_bytes = read_u32_le(bytes, &mut off);
let alpha = read_f32_le(bytes, &mut off);
let tau = read_f32_le(bytes, &mut off);
let w_ema = read_f32_le(bytes, &mut off);
let w_pop = read_f32_le(bytes, &mut off);
let w_rec = read_f32_le(bytes, &mut off);
let t1 = read_f32_le(bytes, &mut off);
let t2 = read_f32_le(bytes, &mut off);
let t3 = read_f32_le(bytes, &mut off);
let hysteresis = read_f32_le(bytes, &mut off);
let min_residency = read_u32_le(bytes, &mut off);
let max_delta_chain = bytes[off];
if ![alpha, tau, w_ema, w_pop, w_rec, t1, t2, t3, hysteresis]
.iter()
.all(|v| v.is_finite())
{
return None;
}
Some(Self {
block_bytes,
alpha,
tau,
w_ema,
w_pop,
w_rec,
t1,
t2,
t3,
hysteresis,
min_residency,
max_delta_chain,
})
}
}
// ── Store ────────────────────────────────────────────────────────────
struct TieredStore {
blocks: HashMap<BlockKey, (BlockMeta, Vec<u8>)>,
}
impl TieredStore {
fn new() -> Self {
Self {
blocks: HashMap::new(),
}
}
fn block_count(&self) -> usize {
self.blocks.len()
}
fn tier_count(&self, tier: Tier) -> usize {
self.blocks.values().filter(|(m, _)| m.tier == tier).count()
}
fn total_bytes(&self) -> usize {
self.blocks.values().map(|(_, d)| d.len()).sum()
}
}
// ── Global state ─────────────────────────────────────────────────────
struct StoreState {
store: TieredStore,
config: TierConfig,
tick_count: u64,
}
static mut STORE_STATE: Option<StoreState> = None;
// ── Helpers ──────────────────────────────────────────────────────────
/// Combine hi/lo u64 into u128 tensor_id.
#[inline]
fn make_tensor_id(hi: u64, lo: u64) -> u128 {
((hi as u128) << 64) | (lo as u128)
}
/// Access the global store state, initializing with defaults if needed.
fn with_state<F, R>(f: F) -> R
where
F: FnOnce(&mut StoreState) -> R,
{
unsafe {
if STORE_STATE.is_none() {
STORE_STATE = Some(StoreState {
store: TieredStore::new(),
config: TierConfig::default(),
tick_count: 0,
});
}
f(STORE_STATE.as_mut().unwrap())
}
}
const DEFAULT_GROUP_LEN: usize = 64;
/// Composite access score used for tier selection.
fn compute_score(config: &TierConfig, meta: &BlockMeta, tick: u64) -> f32 {
let recency = if tick > meta.last_access_ts {
(-((tick - meta.last_access_ts) as f32) / config.tau).exp()
} else {
1.0
};
let popularity = (meta.access_count as f32).ln_1p();
config.w_ema * meta.ema_score + config.w_pop * popularity + config.w_rec * recency
}
/// Map a score to a tier using the config thresholds.
fn choose_tier(config: &TierConfig, score: f32) -> Tier {
if score >= config.t1 {
Tier::Hot
} else if score >= config.t2 {
Tier::Warm
} else if score >= config.t3 {
Tier::Cool
} else {
Tier::Cold
}
}
/// Quantize f32 data and encode into a compressed segment.
fn encode_block(data: &[f32], tier: Tier) -> Vec<u8> {
let bits = tier.bits();
let group_len = DEFAULT_GROUP_LEN;
let scales = quantizer::compute_scales(data, group_len, bits);
let mut packed = Vec::new();
quantizer::quantize_and_pack(data, &scales, group_len, bits, &mut packed);
let mut seg = Vec::new();
segment::encode(
bits,
group_len as u32,
data.len() as u32,
1,
&scales,
&packed,
&mut seg,
);
seg
}
/// Decode a compressed segment back to f32.
fn decode_block(seg: &[u8]) -> Vec<f32> {
let mut out = Vec::new();
segment::decode(seg, &mut out);
out
}
#[inline]
fn read_u32_le(bytes: &[u8], off: &mut usize) -> u32 {
let o = *off;
let arr = [bytes[o], bytes[o + 1], bytes[o + 2], bytes[o + 3]];
*off = o + 4;
u32::from_le_bytes(arr)
}
#[inline]
fn read_f32_le(bytes: &[u8], off: &mut usize) -> f32 {
f32::from_bits(read_u32_le(bytes, off))
}
#[inline]
fn write_u32_le(buf: &mut [u8], off: &mut usize, v: u32) {
buf[*off..*off + 4].copy_from_slice(&v.to_le_bytes());
*off += 4;
}
#[inline]
fn write_u64_le(buf: &mut [u8], off: &mut usize, v: u64) {
buf[*off..*off + 8].copy_from_slice(&v.to_le_bytes());
*off += 8;
}
/// Stats binary layout (36 bytes, little-endian):
/// ```text
/// [block_count:u32][hot:u32][warm:u32][cool:u32][cold:u32]
/// [total_bytes:u64][tick_count:u64]
/// ```
const STATS_SIZE: usize = 5 * 4 + 2 * 8;
// ── FFI exports ──────────────────────────────────────────────────────
/// Initialize the temporal tensor store with a serialized config.
/// If `policy_ptr` is null or `policy_len` is 0, uses `TierConfig::default()`.
/// Returns 0 on success, negative on error.
#[no_mangle]
pub extern "C" fn tts_init(policy_ptr: *const u8, policy_len: usize) -> i32 {
let config = if policy_ptr.is_null() || policy_len == 0 {
TierConfig::default()
} else {
let bytes = unsafe { std::slice::from_raw_parts(policy_ptr, policy_len) };
match TierConfig::from_bytes(bytes) {
Some(c) => c,
None => return ERR_INVALID_CONFIG,
}
};
unsafe {
STORE_STATE = Some(StoreState {
store: TieredStore::new(),
config,
tick_count: 0,
});
}
0
}
/// Store a tensor block. Quantizes according to the block's current tier
/// (or Hot for new blocks). `tensor_id` is split into hi/lo because WASM
/// does not support u128.
/// Returns 0 on success, negative on error.
#[no_mangle]
pub extern "C" fn tts_put(
tensor_id_hi: u64,
tensor_id_lo: u64,
block_index: u32,
data_ptr: *const f32,
data_len: usize,
) -> i32 {
if data_ptr.is_null() {
return ERR_NULL_POINTER;
}
if data_len == 0 {
return ERR_EMPTY_DATA;
}
let data = unsafe { std::slice::from_raw_parts(data_ptr, data_len) };
let key = BlockKey {
tensor_id: make_tensor_id(tensor_id_hi, tensor_id_lo),
block_index,
};
with_state(|state| {
let tier = state
.store
.blocks
.get(&key)
.map(|(m, _)| m.tier)
.unwrap_or(Tier::Hot);
let seg = encode_block(data, tier);
let meta = BlockMeta {
tier,
access_count: 1,
last_access_ts: state.tick_count,
ema_score: 1.0,
element_count: data_len,
};
state.store.blocks.insert(key, (meta, seg));
0
})
}
/// Read a tensor block, dequantized to f32.
/// Returns the number of f32 elements written, or negative on error.
#[no_mangle]
pub extern "C" fn tts_get(
tensor_id_hi: u64,
tensor_id_lo: u64,
block_index: u32,
out_ptr: *mut f32,
out_len: usize,
) -> i32 {
if out_ptr.is_null() {
return ERR_NULL_POINTER;
}
let key = BlockKey {
tensor_id: make_tensor_id(tensor_id_hi, tensor_id_lo),
block_index,
};
with_state(|state| match state.store.blocks.get(&key) {
None => ERR_BLOCK_NOT_FOUND,
Some((_meta, seg)) => {
let decoded = decode_block(seg);
if decoded.len() > out_len {
return ERR_BUFFER_TOO_SMALL;
}
let out = unsafe { std::slice::from_raw_parts_mut(out_ptr, out_len) };
out[..decoded.len()].copy_from_slice(&decoded);
decoded.len() as i32
}
})
}
/// Run a maintenance tick with byte and operation budgets.
/// Re-scores every block and migrates those whose tier has changed,
/// subject to hysteresis.
/// Returns number of migration operations performed, or negative on error.
#[no_mangle]
pub extern "C" fn tts_tick(budget_bytes: u32, budget_ops: u32) -> i32 {
with_state(|state| {
state.tick_count += 1;
let tick = state.tick_count;
// Snapshot keys and scores so we can mutate blocks afterwards.
let entries: Vec<(BlockKey, f32)> = state
.store
.blocks
.iter()
.map(|(k, (m, _))| (*k, compute_score(&state.config, m, tick)))
.collect();
let mut ops = 0u32;
let mut bytes_used = 0u32;
for (key, score) in entries {
if ops >= budget_ops || bytes_used >= budget_bytes {
break;
}
if let Some((meta, seg)) = state.store.blocks.get_mut(&key) {
let new_tier = choose_tier(&state.config, score);
let current_threshold = match meta.tier {
Tier::Hot => state.config.t1,
Tier::Warm => state.config.t2,
Tier::Cool => state.config.t3,
Tier::Cold => 0.0,
};
let needs_change = new_tier != meta.tier
&& (score - current_threshold).abs() > state.config.hysteresis;
if needs_change {
let decoded = decode_block(seg);
if !decoded.is_empty() {
let new_seg = encode_block(&decoded, new_tier);
bytes_used = bytes_used.saturating_add(new_seg.len() as u32);
*seg = new_seg;
meta.tier = new_tier;
ops += 1;
}
}
// Update EMA for every block regardless of migration.
meta.ema_score =
state.config.alpha * score + (1.0 - state.config.alpha) * meta.ema_score;
}
}
ops as i32
})
}
/// Write a statistics snapshot to `out_ptr`.
/// Returns number of bytes written, or negative on error.
#[no_mangle]
pub extern "C" fn tts_stats(out_ptr: *mut u8, out_len: usize) -> i32 {
if out_ptr.is_null() {
return ERR_NULL_POINTER;
}
if out_len < STATS_SIZE {
return ERR_BUFFER_TOO_SMALL;
}
with_state(|state| {
let out = unsafe { std::slice::from_raw_parts_mut(out_ptr, out_len) };
let mut off = 0usize;
write_u32_le(out, &mut off, state.store.block_count() as u32);
write_u32_le(out, &mut off, state.store.tier_count(Tier::Hot) as u32);
write_u32_le(out, &mut off, state.store.tier_count(Tier::Warm) as u32);
write_u32_le(out, &mut off, state.store.tier_count(Tier::Cool) as u32);
write_u32_le(out, &mut off, state.store.tier_count(Tier::Cold) as u32);
write_u64_le(out, &mut off, state.store.total_bytes() as u64);
write_u64_le(out, &mut off, state.tick_count);
STATS_SIZE as i32
})
}
/// Record an access event for a block (increments count, updates timestamp).
/// Returns 0 on success, negative on error.
#[no_mangle]
pub extern "C" fn tts_touch(tensor_id_hi: u64, tensor_id_lo: u64, block_index: u32) -> i32 {
let key = BlockKey {
tensor_id: make_tensor_id(tensor_id_hi, tensor_id_lo),
block_index,
};
with_state(|state| match state.store.blocks.get_mut(&key) {
None => ERR_BLOCK_NOT_FOUND,
Some((meta, _)) => {
meta.access_count = meta.access_count.saturating_add(1);
meta.last_access_ts = state.tick_count;
0
}
})
}
/// Evict a block, removing it from the store entirely.
/// Returns 0 on success, negative on error.
#[no_mangle]
pub extern "C" fn tts_evict(tensor_id_hi: u64, tensor_id_lo: u64, block_index: u32) -> i32 {
let key = BlockKey {
tensor_id: make_tensor_id(tensor_id_hi, tensor_id_lo),
block_index,
};
with_state(|state| match state.store.blocks.remove(&key) {
None => ERR_BLOCK_NOT_FOUND,
Some(_) => 0,
})
}
/// Get total number of blocks in the store.
#[no_mangle]
pub extern "C" fn tts_block_count() -> i32 {
with_state(|state| state.store.block_count() as i32)
}
/// Get number of blocks in a specific tier (0=Hot, 1=Warm, 2=Cool, 3=Cold).
#[no_mangle]
pub extern "C" fn tts_tier_count(tier: u8) -> i32 {
match Tier::from_u8(tier) {
Some(t) => with_state(|state| state.store.tier_count(t) as i32),
None => ERR_INVALID_CONFIG,
}
}
// ── Tests ────────────────────────────────────────────────────────────
#[cfg(test)]
mod tests {
use super::*;
/// Reset global state before each test.
fn reset() {
unsafe {
STORE_STATE = None;
}
}
/// Build a binary config buffer from the default TierConfig.
fn default_config_bytes() -> Vec<u8> {
let c = TierConfig::default();
let mut buf = Vec::with_capacity(CONFIG_BINARY_LEN);
buf.extend_from_slice(&c.block_bytes.to_le_bytes());
buf.extend_from_slice(&c.alpha.to_bits().to_le_bytes());
buf.extend_from_slice(&c.tau.to_bits().to_le_bytes());
buf.extend_from_slice(&c.w_ema.to_bits().to_le_bytes());
buf.extend_from_slice(&c.w_pop.to_bits().to_le_bytes());
buf.extend_from_slice(&c.w_rec.to_bits().to_le_bytes());
buf.extend_from_slice(&c.t1.to_bits().to_le_bytes());
buf.extend_from_slice(&c.t2.to_bits().to_le_bytes());
buf.extend_from_slice(&c.t3.to_bits().to_le_bytes());
buf.extend_from_slice(&c.hysteresis.to_bits().to_le_bytes());
buf.extend_from_slice(&c.min_residency.to_le_bytes());
buf.push(c.max_delta_chain);
buf
}
#[test]
fn test_init_default() {
reset();
let rc = tts_init(std::ptr::null(), 0);
assert_eq!(rc, 0);
assert_eq!(tts_block_count(), 0);
}
#[test]
fn test_init_with_config() {
reset();
let cfg = default_config_bytes();
let rc = tts_init(cfg.as_ptr(), cfg.len());
assert_eq!(rc, 0);
assert_eq!(tts_block_count(), 0);
}
#[test]
fn test_init_invalid_config_too_short() {
reset();
let buf = [0u8; 10];
let rc = tts_init(buf.as_ptr(), buf.len());
assert_eq!(rc, ERR_INVALID_CONFIG);
}
#[test]
fn test_put_get_roundtrip() {
reset();
tts_init(std::ptr::null(), 0);
let data: Vec<f32> = (0..64).map(|i| (i as f32 - 32.0) * 0.1).collect();
let rc = tts_put(0, 1, 0, data.as_ptr(), data.len());
assert_eq!(rc, 0);
let mut out = vec![0.0f32; 64];
let n = tts_get(0, 1, 0, out.as_mut_ptr(), out.len());
assert_eq!(n, 64);
// 8-bit quantization: expect low error.
let max_abs = data.iter().map(|v| v.abs()).fold(0.0f32, f32::max);
for (i, (&orig, &dec)) in data.iter().zip(out.iter()).enumerate() {
let err = (orig - dec).abs();
assert!(
err < max_abs * 0.05,
"i={i} orig={orig} dec={dec} err={err}"
);
}
}
#[test]
fn test_put_null_pointer() {
reset();
tts_init(std::ptr::null(), 0);
let rc = tts_put(0, 1, 0, std::ptr::null(), 64);
assert_eq!(rc, ERR_NULL_POINTER);
}
#[test]
fn test_put_empty_data() {
reset();
tts_init(std::ptr::null(), 0);
let data = [1.0f32; 1];
let rc = tts_put(0, 1, 0, data.as_ptr(), 0);
assert_eq!(rc, ERR_EMPTY_DATA);
}
#[test]
fn test_get_not_found() {
reset();
tts_init(std::ptr::null(), 0);
let mut out = vec![0.0f32; 64];
let rc = tts_get(0, 99, 0, out.as_mut_ptr(), out.len());
assert_eq!(rc, ERR_BLOCK_NOT_FOUND);
}
#[test]
fn test_get_null_pointer() {
reset();
tts_init(std::ptr::null(), 0);
let rc = tts_get(0, 1, 0, std::ptr::null_mut(), 64);
assert_eq!(rc, ERR_NULL_POINTER);
}
#[test]
fn test_get_buffer_too_small() {
reset();
tts_init(std::ptr::null(), 0);
let data = vec![1.0f32; 64];
tts_put(0, 1, 0, data.as_ptr(), data.len());
let mut out = vec![0.0f32; 2]; // too small
let rc = tts_get(0, 1, 0, out.as_mut_ptr(), out.len());
assert_eq!(rc, ERR_BUFFER_TOO_SMALL);
}
#[test]
fn test_block_count_after_puts() {
reset();
tts_init(std::ptr::null(), 0);
let data = vec![1.0f32; 64];
tts_put(0, 1, 0, data.as_ptr(), data.len());
tts_put(0, 1, 1, data.as_ptr(), data.len());
tts_put(0, 2, 0, data.as_ptr(), data.len());
assert_eq!(tts_block_count(), 3);
}
#[test]
fn test_tier_count_initial() {
reset();
tts_init(std::ptr::null(), 0);
let data = vec![1.0f32; 64];
tts_put(0, 1, 0, data.as_ptr(), data.len());
tts_put(0, 1, 1, data.as_ptr(), data.len());
// New blocks default to Hot.
assert_eq!(tts_tier_count(0), 2); // Hot
assert_eq!(tts_tier_count(1), 0); // Warm
assert_eq!(tts_tier_count(2), 0); // Cool
assert_eq!(tts_tier_count(3), 0); // Cold
}
#[test]
fn test_tier_count_invalid_tier() {
reset();
tts_init(std::ptr::null(), 0);
assert_eq!(tts_tier_count(99), ERR_INVALID_CONFIG);
}
#[test]
fn test_touch() {
reset();
tts_init(std::ptr::null(), 0);
let data = vec![1.0f32; 64];
tts_put(0, 1, 0, data.as_ptr(), data.len());
let rc = tts_touch(0, 1, 0);
assert_eq!(rc, 0);
// Touch a non-existent block.
let rc = tts_touch(0, 99, 0);
assert_eq!(rc, ERR_BLOCK_NOT_FOUND);
}
#[test]
fn test_evict() {
reset();
tts_init(std::ptr::null(), 0);
let data = vec![1.0f32; 64];
tts_put(0, 1, 0, data.as_ptr(), data.len());
assert_eq!(tts_block_count(), 1);
let rc = tts_evict(0, 1, 0);
assert_eq!(rc, 0);
assert_eq!(tts_block_count(), 0);
// Evict again should fail.
let rc = tts_evict(0, 1, 0);
assert_eq!(rc, ERR_BLOCK_NOT_FOUND);
}
#[test]
fn test_tick_does_not_crash() {
reset();
tts_init(std::ptr::null(), 0);
let data = vec![1.0f32; 64];
tts_put(0, 1, 0, data.as_ptr(), data.len());
tts_put(0, 1, 1, data.as_ptr(), data.len());
// Run several ticks with generous budgets.
for _ in 0..10 {
let ops = tts_tick(1_000_000, 1000);
assert!(ops >= 0);
}
// Blocks should still be readable.
let mut out = vec![0.0f32; 64];
let n = tts_get(0, 1, 0, out.as_mut_ptr(), out.len());
assert!(n > 0);
}
#[test]
fn test_tick_with_zero_budget() {
reset();
tts_init(std::ptr::null(), 0);
let data = vec![1.0f32; 64];
tts_put(0, 1, 0, data.as_ptr(), data.len());
let ops = tts_tick(0, 0);
assert_eq!(ops, 0);
}
#[test]
fn test_stats_returns_valid_data() {
reset();
tts_init(std::ptr::null(), 0);
let data = vec![1.0f32; 64];
tts_put(0, 1, 0, data.as_ptr(), data.len());
tts_put(0, 1, 1, data.as_ptr(), data.len());
let mut buf = vec![0u8; STATS_SIZE];
let written = tts_stats(buf.as_mut_ptr(), buf.len());
assert_eq!(written, STATS_SIZE as i32);
// Parse the stats back.
let mut off = 0usize;
let block_count = read_u32_le(&buf, &mut off);
let hot = read_u32_le(&buf, &mut off);
let warm = read_u32_le(&buf, &mut off);
let cool = read_u32_le(&buf, &mut off);
let cold = read_u32_le(&buf, &mut off);
assert_eq!(block_count, 2);
assert_eq!(hot, 2);
assert_eq!(warm, 0);
assert_eq!(cool, 0);
assert_eq!(cold, 0);
}
#[test]
fn test_stats_null_pointer() {
reset();
tts_init(std::ptr::null(), 0);
let rc = tts_stats(std::ptr::null_mut(), 64);
assert_eq!(rc, ERR_NULL_POINTER);
}
#[test]
fn test_stats_buffer_too_small() {
reset();
tts_init(std::ptr::null(), 0);
let mut buf = vec![0u8; 4]; // too small
let rc = tts_stats(buf.as_mut_ptr(), buf.len());
assert_eq!(rc, ERR_BUFFER_TOO_SMALL);
}
#[test]
fn test_make_tensor_id() {
assert_eq!(make_tensor_id(0, 0), 0u128);
assert_eq!(make_tensor_id(0, 1), 1u128);
assert_eq!(make_tensor_id(1, 0), 1u128 << 64);
assert_eq!(make_tensor_id(u64::MAX, u64::MAX), u128::MAX,);
}
#[test]
fn test_multiple_tensor_ids() {
reset();
tts_init(std::ptr::null(), 0);
let data = vec![1.0f32; 64];
tts_put(0, 1, 0, data.as_ptr(), data.len());
tts_put(0, 2, 0, data.as_ptr(), data.len());
tts_put(1, 0, 0, data.as_ptr(), data.len());
assert_eq!(tts_block_count(), 3);
// Each should be independently readable.
let mut out = vec![0.0f32; 64];
assert!(tts_get(0, 1, 0, out.as_mut_ptr(), out.len()) > 0);
assert!(tts_get(0, 2, 0, out.as_mut_ptr(), out.len()) > 0);
assert!(tts_get(1, 0, 0, out.as_mut_ptr(), out.len()) > 0);
}
#[test]
fn test_overwrite_block() {
reset();
tts_init(std::ptr::null(), 0);
let data1 = vec![1.0f32; 64];
tts_put(0, 1, 0, data1.as_ptr(), data1.len());
let data2 = vec![2.0f32; 64];
tts_put(0, 1, 0, data2.as_ptr(), data2.len());
assert_eq!(tts_block_count(), 1);
// Should read back the second write.
let mut out = vec![0.0f32; 64];
let n = tts_get(0, 1, 0, out.as_mut_ptr(), out.len());
assert_eq!(n, 64);
for &v in &out {
assert!((v - 2.0).abs() < 0.1);
}
}
}

View File

@@ -0,0 +1,104 @@
//! Tier policy for access-pattern-driven bit-width selection.
//!
//! Score = `access_count * 1024 / (now_ts - last_access_ts + 1)`
//!
//! | Tier | Condition | Bits |
//! |------|-----------|------|
//! | Hot | score >= hot_min_score | 8 |
//! | Warm | score >= warm_min_score | warm_bits (7 or 5) |
//! | Cold | otherwise | 3 |
#[derive(Clone, Copy, Debug)]
pub struct TierPolicy {
pub hot_min_score: u32,
pub warm_min_score: u32,
pub warm_bits: u8,
/// Drift tolerance as Q8 fixed-point. 26 means ~10.2% (26/256).
pub drift_pct_q8: u32,
pub group_len: u32,
}
impl Default for TierPolicy {
fn default() -> Self {
Self {
hot_min_score: 512,
warm_min_score: 64,
warm_bits: 7,
drift_pct_q8: 26,
group_len: 64,
}
}
}
impl TierPolicy {
/// Select bit width based on access pattern.
pub fn select_bits(&self, access_count: u32, last_access_ts: u32, now_ts: u32) -> u8 {
let age = now_ts.wrapping_sub(last_access_ts).wrapping_add(1);
let score = access_count.saturating_mul(1024).wrapping_div(age);
if score >= self.hot_min_score {
8
} else if score >= self.warm_min_score {
self.warm_bits
} else {
3
}
}
/// Compute the drift factor as 1.0 + drift_pct_q8/256.
pub fn drift_factor(&self) -> f32 {
1.0 + (self.drift_pct_q8 as f32) / 256.0
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_default_policy() {
let p = TierPolicy::default();
assert_eq!(p.hot_min_score, 512);
assert_eq!(p.warm_min_score, 64);
assert_eq!(p.warm_bits, 7);
assert_eq!(p.drift_pct_q8, 26);
assert_eq!(p.group_len, 64);
}
#[test]
fn test_tier_selection_hot() {
let p = TierPolicy::default();
// 100 accesses, age=10 -> score = 100*1024/10 = 10240 >= 512
assert_eq!(p.select_bits(100, 0, 9), 8);
}
#[test]
fn test_tier_selection_warm() {
let p = TierPolicy::default();
// 10 accesses, age=100 -> score = 10*1024/100 = 102 >= 64, < 512
assert_eq!(p.select_bits(10, 0, 99), 7);
}
#[test]
fn test_tier_selection_cold() {
let p = TierPolicy::default();
// 1 access, age=1000 -> score = 1024/1000 = 1 < 64
assert_eq!(p.select_bits(1, 0, 999), 3);
}
#[test]
fn test_drift_factor() {
let p = TierPolicy::default();
let df = p.drift_factor();
assert!((df - 1.1015625).abs() < 1e-6);
}
#[test]
fn test_warm_bits_5() {
let p = TierPolicy {
warm_bits: 5,
..Default::default()
};
assert_eq!(p.select_bits(10, 0, 99), 5);
}
}

File diff suppressed because it is too large Load Diff

File diff suppressed because it is too large Load Diff

View File

@@ -0,0 +1,701 @@
//! End-to-end integration tests for the temporal tensor store.
//!
//! Exercises the full lifecycle: put, get, tier migration, delta compression,
//! quantization quality, eviction, checksums, witness logging, and factor
//! reconstruction.
//!
//! Run via: `cargo test -p ruvector-temporal-tensor --test integration`
use ruvector_temporal_tensor::delta::{
compute_delta, decode_delta, encode_delta, DeltaChain, FactorSet,
};
use ruvector_temporal_tensor::metrics::{TierChangeReason, WitnessEvent, WitnessLog};
use ruvector_temporal_tensor::quantizer;
use ruvector_temporal_tensor::segment;
use ruvector_temporal_tensor::store::{BlockKey, ReconstructPolicy, StoreError, Tier, TieredStore};
use ruvector_temporal_tensor::tiering::{self, TierConfig};
use ruvector_temporal_tensor::{TemporalTensorCompressor, TierPolicy};
// ---------------------------------------------------------------------------
// Deterministic PRNG (LCG) -- no external deps
// ---------------------------------------------------------------------------
/// Simple linear congruential generator. Constants from Knuth MMIX.
struct SimpleRng {
state: u64,
}
impl SimpleRng {
fn new(seed: u64) -> Self {
Self { state: seed }
}
fn next_u64(&mut self) -> u64 {
self.state = self
.state
.wrapping_mul(6_364_136_223_846_793_005)
.wrapping_add(1_442_695_040_888_963_407);
self.state
}
fn next_f64(&mut self) -> f64 {
(self.next_u64() >> 11) as f64 / (1u64 << 53) as f64
}
fn next_f32(&mut self) -> f32 {
self.next_f64() as f32
}
}
// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------
fn make_key(tid: u128, idx: u32) -> BlockKey {
BlockKey {
tensor_id: tid,
block_index: idx,
}
}
/// Map tiering module Tier to store module Tier.
fn tiering_to_store_tier(t: tiering::Tier) -> Tier {
match t {
tiering::Tier::Tier0 => Tier::Tier0,
tiering::Tier::Tier1 => Tier::Tier1,
tiering::Tier::Tier2 => Tier::Tier2,
tiering::Tier::Tier3 => Tier::Tier3,
}
}
// ===========================================================================
// 1. Full Lifecycle Test
// ===========================================================================
/// Put 100 blocks as hot, simulate 1000 ticks touching only 10, then verify
/// that the 90 untouched blocks migrate to colder tiers.
#[test]
fn test_full_lifecycle() {
let mut store = TieredStore::new(4096);
let tier_config = TierConfig::default();
let n_elems = 64;
let mut rng = SimpleRng::new(42);
let block_data: Vec<Vec<f32>> = (0..100)
.map(|_| (0..n_elems).map(|_| rng.next_f32() * 2.0 - 1.0).collect())
.collect();
// Put 100 blocks as Tier1 (hot).
for i in 0..100u32 {
store
.put(make_key(1, i), &block_data[i as usize], Tier::Tier1, 0)
.unwrap();
}
assert_eq!(store.tier_count(Tier::Tier1), 100);
assert_eq!(store.block_count(), 100);
// Parallel tiering metadata for migration scoring.
let mut tiering_metas: Vec<tiering::BlockMeta> =
(0..100).map(|_| tiering::BlockMeta::new(0)).collect();
// Simulate 1000 ticks -- only blocks 0..10 are accessed.
for tick in 1..=1000u64 {
for i in 0..10 {
store.touch(make_key(1, i as u32), tick);
tiering::touch(&tier_config, tick, &mut tiering_metas[i]);
}
for i in 10..100 {
tiering::tick_decay(&tier_config, &mut tiering_metas[i]);
}
}
// Apply tier migration decisions.
let mut migrated = 0u32;
for i in 0..100u32 {
if let Some(target) = tiering::choose_tier(&tier_config, 1000, &tiering_metas[i as usize]) {
let st = tiering_to_store_tier(target);
if st != Tier::Tier0 {
store
.put(make_key(1, i), &block_data[i as usize], st, 1000)
.unwrap();
migrated += 1;
}
}
}
let tier1 = store.tier_count(Tier::Tier1);
let tier2 = store.tier_count(Tier::Tier2);
let tier3 = store.tier_count(Tier::Tier3);
assert!(migrated > 0, "expected migrations, got none");
assert!(
tier1 < 100,
"expected fewer Tier1 blocks after migration, got {}",
tier1
);
assert!(tier1 <= 20, "hot blocks should be ~10, got {}", tier1);
assert!(
tier2 + tier3 >= 80,
"expected >=80 in lower tiers, got {} + {}",
tier2,
tier3
);
assert_eq!(store.block_count(), 100);
}
// ===========================================================================
// 2. Delta Chain Lifecycle Test
// ===========================================================================
/// Build a delta chain with 5 incremental deltas, reconstruct, compact,
/// verify encode/decode roundtrip.
#[test]
fn test_delta_chain_lifecycle() {
let n = 256;
let mut rng = SimpleRng::new(99);
let base: Vec<f32> = (0..n).map(|_| rng.next_f32() * 2.0 - 1.0).collect();
let mut chain = DeltaChain::new(base.clone(), 8);
// Build 5 incremental deltas (~10% change each).
let mut current = base.clone();
for epoch in 0..5u64 {
let mut next = current.clone();
for i in 0..n {
if (rng.next_u64() % 10) == 0 {
next[i] += (rng.next_f32() - 0.5) * 0.1;
}
}
let delta = compute_delta(&current, &next, 1, 0, epoch, 0.001, 0.5)
.expect("delta should be computable for ~10% change");
chain.append(delta).unwrap();
current = next;
}
assert_eq!(chain.chain_len(), 5);
// Reconstruct and verify accuracy against the final state.
let reconstructed = chain.reconstruct();
assert_eq!(reconstructed.len(), n);
for i in 0..n {
let err = (reconstructed[i] - current[i]).abs();
assert!(
err < 0.01,
"recon err at {}: {} vs {} (err={})",
i,
reconstructed[i],
current[i],
err
);
}
// Encode/decode the last delta and verify roundtrip.
let last_delta = compute_delta(&base, &current, 1, 0, 99, 0.001, 1.1).unwrap();
let encoded = encode_delta(&last_delta);
let decoded = decode_delta(&encoded).unwrap();
assert_eq!(decoded.header.tensor_id, 1);
assert_eq!(decoded.entries.len(), last_delta.entries.len());
// Compact the chain; delta list drops to 0 but state is preserved.
let before_compact = reconstructed.clone();
chain.compact();
assert_eq!(chain.chain_len(), 0);
let after_compact = chain.reconstruct();
for i in 0..n {
let err = (after_compact[i] - before_compact[i]).abs();
assert!(
err < 1e-6,
"compact mismatch at {}: {} vs {}",
i,
after_compact[i],
before_compact[i]
);
}
}
// ===========================================================================
// 3. Quantization Quality Sweep
// ===========================================================================
/// For each bit width (8, 7, 5, 3) verify MSE and max relative error
/// stay within ADR-023 bounds.
#[test]
fn test_quality_sweep_all_tiers() {
let n_elems = 256;
let mut rng = SimpleRng::new(7777);
// Sinusoidal + noise with guaranteed minimum magnitude.
let data: Vec<f32> = (0..n_elems)
.map(|i| {
let base = (i as f32 * 0.05).sin();
let noise = (rng.next_f32() - 0.5) * 0.1;
let val = base + noise;
if val.abs() < 0.05 {
if val >= 0.0 {
0.05 + rng.next_f32() * 0.1
} else {
-0.05 - rng.next_f32() * 0.1
}
} else {
val
}
})
.collect();
let max_abs: f32 = data.iter().map(|v| v.abs()).fold(0.0f32, f32::max);
// Store-backed tiers: (tier, bound_vs_max, label).
let store_configs: &[(Tier, f64, &str)] = &[
(Tier::Tier1, 0.01, "8-bit/Tier1"),
(Tier::Tier2, 0.02, "7-bit/Tier2"),
(Tier::Tier3, 0.35, "3-bit/Tier3"),
];
let mut store = TieredStore::new(4096);
for &(tier, bound, label) in store_configs {
let key = make_key(tier as u128 + 100, 0);
store.put(key, &data, tier, 0).unwrap();
let mut out = vec![0.0f32; n_elems];
let n = store.get(key, &mut out, 0).unwrap();
assert_eq!(n, n_elems);
let mut max_rel = 0.0f64;
let mut mse = 0.0f64;
for i in 0..n_elems {
let err = (data[i] - out[i]) as f64;
mse += err * err;
let rel = err.abs() / max_abs as f64;
if rel > max_rel {
max_rel = rel;
}
}
mse /= n_elems as f64;
assert!(
max_rel < bound,
"{}: max_rel {:.4} >= bound {:.4} (MSE={:.8})",
label,
max_rel,
bound,
mse
);
}
// 5-bit via groupwise quantizer directly (no store tier for 5-bit).
{
let scales = quantizer::compute_scales(&data, 64, 5);
let mut packed = Vec::new();
quantizer::quantize_and_pack(&data, &scales, 64, 5, &mut packed);
let mut decoded = Vec::new();
quantizer::dequantize(&packed, &scales, 64, 5, n_elems, 1, &mut decoded);
let mut max_rel = 0.0f64;
for i in 0..n_elems {
let err = (data[i] - decoded[i]) as f64;
let rel = err.abs() / max_abs as f64;
if rel > max_rel {
max_rel = rel;
}
}
assert!(max_rel < 0.07, "5-bit: max_rel {:.4} >= 0.07", max_rel);
}
}
// ===========================================================================
// 4. Store Persistence Roundtrip
// ===========================================================================
/// Put 50 blocks with varied data and tiers, get each back and verify data
/// and metadata.
#[test]
fn test_store_put_get_roundtrip() {
let mut store = TieredStore::new(4096);
let mut rng = SimpleRng::new(1234);
let n_elems = 64;
let tiers = [Tier::Tier1, Tier::Tier2, Tier::Tier3];
let mut block_data: Vec<Vec<f32>> = Vec::new();
let mut block_tiers: Vec<Tier> = Vec::new();
for i in 0..50u32 {
let d: Vec<f32> = (0..n_elems).map(|_| rng.next_f32() * 2.0 - 1.0).collect();
let tier = tiers[(i % 3) as usize];
store.put(make_key(42, i), &d, tier, i as u64).unwrap();
block_data.push(d);
block_tiers.push(tier);
}
assert_eq!(store.block_count(), 50);
for i in 0..50u32 {
let key = make_key(42, i);
let mut out = vec![0.0f32; n_elems];
let n = store.get(key, &mut out, i as u64).unwrap();
assert_eq!(n, n_elems);
let meta = store.meta(key).unwrap();
assert_eq!(meta.tier, block_tiers[i as usize]);
assert_eq!(meta.created_at, i as u64);
let max_abs: f32 = block_data[i as usize]
.iter()
.map(|v| v.abs())
.fold(0.0f32, f32::max);
let tol = match block_tiers[i as usize] {
Tier::Tier1 => max_abs * 0.01,
Tier::Tier2 => max_abs * 0.02,
Tier::Tier3 => max_abs * 0.35,
Tier::Tier0 => unreachable!(),
}
.max(1e-6);
for j in 0..n_elems {
let err = (block_data[i as usize][j] - out[j]).abs();
assert!(err < tol, "block {} elem {}: err={} tol={}", i, j, err, tol);
}
}
}
// ===========================================================================
// 5. Eviction and Tier0
// ===========================================================================
/// Put a block at Tier1, evict it, verify reads fail and metadata reflects
/// eviction state.
#[test]
fn test_eviction_to_tier0() {
let mut store = TieredStore::new(4096);
let key = make_key(1, 0);
let data = vec![1.0f32; 64];
store.put(key, &data, Tier::Tier1, 0).unwrap();
assert_eq!(store.tier_count(Tier::Tier1), 1);
assert!(store.total_bytes() > 0);
store.evict(key, ReconstructPolicy::None).unwrap();
// Read should fail.
let mut out = vec![0.0f32; 64];
assert_eq!(store.get(key, &mut out, 1), Err(StoreError::TensorEvicted));
// Metadata should reflect Tier0.
let meta = store.meta(key).unwrap();
assert_eq!(meta.tier, Tier::Tier0);
assert_eq!(meta.bits, 0);
assert_eq!(meta.block_bytes, 0);
assert_eq!(meta.reconstruct, ReconstructPolicy::None);
assert_eq!(store.tier_count(Tier::Tier1), 0);
assert_eq!(store.tier_count(Tier::Tier0), 1);
assert_eq!(store.block_count(), 1);
assert_eq!(store.total_bytes(), 0);
}
// ===========================================================================
// 6. Checksum Integrity
// ===========================================================================
/// Verify that checksums are non-zero and deterministic for the same data.
#[test]
fn test_checksum_integrity() {
let mut store = TieredStore::new(4096);
let data: Vec<f32> = (0..128).map(|i| (i as f32) * 0.1).collect();
let key1 = make_key(1, 0);
store.put(key1, &data, Tier::Tier1, 0).unwrap();
let cksum1 = store.meta(key1).unwrap().checksum;
assert_ne!(
cksum1, 0,
"checksum should be non-zero for non-trivial data"
);
// Same data under a different key produces the same checksum.
let key2 = make_key(1, 1);
store.put(key2, &data, Tier::Tier1, 0).unwrap();
assert_eq!(store.meta(key2).unwrap().checksum, cksum1);
// Different data produces a different checksum.
let other: Vec<f32> = (0..128).map(|i| (i as f32) * 0.2).collect();
let key3 = make_key(1, 2);
store.put(key3, &other, Tier::Tier1, 0).unwrap();
assert_ne!(store.meta(key3).unwrap().checksum, cksum1);
}
// ===========================================================================
// 7. Multi-Tensor Store
// ===========================================================================
/// Blocks from 3 different tensor_ids are stored and retrieved independently.
#[test]
fn test_multiple_tensors() {
let mut store = TieredStore::new(4096);
let n_elems = 32;
let mut rng = SimpleRng::new(555);
let tensor_ids: [u128; 3] = [100, 200, 300];
let mut all_data: Vec<Vec<Vec<f32>>> = Vec::new();
for &tid in &tensor_ids {
let mut tensor_blocks = Vec::new();
for blk in 0..5u32 {
let d: Vec<f32> = (0..n_elems).map(|_| rng.next_f32() * 2.0 - 1.0).collect();
store.put(make_key(tid, blk), &d, Tier::Tier1, 0).unwrap();
tensor_blocks.push(d);
}
all_data.push(tensor_blocks);
}
assert_eq!(store.block_count(), 15);
for (t_idx, &tid) in tensor_ids.iter().enumerate() {
for blk in 0..5u32 {
let key = make_key(tid, blk);
let mut out = vec![0.0f32; n_elems];
let n = store.get(key, &mut out, 0).unwrap();
assert_eq!(n, n_elems);
let meta = store.meta(key).unwrap();
assert_eq!(meta.key.tensor_id, tid);
assert_eq!(meta.key.block_index, blk);
let orig = &all_data[t_idx][blk as usize];
let max_abs: f32 = orig.iter().map(|v| v.abs()).fold(0.0f32, f32::max);
let tol = (max_abs * 0.01).max(1e-6);
for j in 0..n_elems {
let err = (orig[j] - out[j]).abs();
assert!(err < tol, "tid={} blk={} j={}: err={}", tid, blk, j, err);
}
}
}
}
// ===========================================================================
// 8. Stress Test
// ===========================================================================
/// Put 1000 blocks with random tiers, touch random blocks 10000 times,
/// verify no panics and all blocks remain readable.
#[test]
fn test_stress_1000_blocks() {
let mut store = TieredStore::new(4096);
let mut rng = SimpleRng::new(0xDEADBEEF);
let n_elems = 32;
let tiers = [Tier::Tier1, Tier::Tier2, Tier::Tier3];
for i in 0..1000u32 {
let d: Vec<f32> = (0..n_elems).map(|_| rng.next_f32() * 2.0 - 1.0).collect();
let tier = tiers[(rng.next_u64() % 3) as usize];
store.put(make_key(1, i), &d, tier, i as u64).unwrap();
}
assert_eq!(store.block_count(), 1000);
assert!(store.total_bytes() > 0);
for t in 0..10_000u64 {
let idx = (rng.next_u64() % 1000) as u32;
store.touch(make_key(1, idx), 1000 + t);
}
for i in 0..1000u32 {
let mut out = vec![0.0f32; n_elems];
let n = store.get(make_key(1, i), &mut out, 20_000).unwrap();
assert_eq!(n, n_elems);
for j in 0..n_elems {
assert!(out[j].is_finite(), "block {} elem {} not finite", i, j);
}
}
assert!(store.total_bytes() > 0);
}
// ===========================================================================
// 9. Compressor + Store Integration
// ===========================================================================
/// Compress frames via TemporalTensorCompressor, decode the segment, store
/// each decoded frame as a block, and verify roundtrip.
#[test]
fn test_compressor_to_store() {
let tensor_len = 128u32;
let policy = TierPolicy::default();
let mut comp = TemporalTensorCompressor::new(policy, tensor_len, 0);
comp.set_access(100, 0); // hot -> 8-bit
let mut rng = SimpleRng::new(0xCAFE);
let n_frames = 10usize;
let frames: Vec<Vec<f32>> = (0..n_frames)
.map(|_| {
(0..tensor_len as usize)
.map(|_| rng.next_f32() * 2.0 - 1.0)
.collect()
})
.collect();
let mut seg = Vec::new();
for (i, frame) in frames.iter().enumerate() {
comp.push_frame(frame, (i + 1) as u32, &mut seg);
}
comp.flush(&mut seg);
assert!(!seg.is_empty(), "compressor should produce a segment");
let mut decoded = Vec::new();
segment::decode(&seg, &mut decoded);
assert_eq!(decoded.len(), tensor_len as usize * n_frames);
// Store each decoded frame as a block.
let mut store = TieredStore::new(4096);
for i in 0..n_frames {
let start = i * tensor_len as usize;
let end = start + tensor_len as usize;
store
.put(
make_key(50, i as u32),
&decoded[start..end],
Tier::Tier1,
i as u64,
)
.unwrap();
}
assert_eq!(store.block_count(), n_frames);
// Read back and verify against the decoded data (double quantization).
for i in 0..n_frames {
let mut out = vec![0.0f32; tensor_len as usize];
let n = store
.get(make_key(50, i as u32), &mut out, n_frames as u64)
.unwrap();
assert_eq!(n, tensor_len as usize);
let start = i * tensor_len as usize;
for j in 0..tensor_len as usize {
let expected = decoded[start + j];
let err = (expected - out[j]).abs();
// Double quantization (compressor + store) compounds error.
let tol = if expected.abs() > 0.01 {
expected.abs() * 0.04
} else {
0.05
};
assert!(
err < tol,
"frame {} elem {}: exp={} got={} err={}",
i,
j,
expected,
out[j],
err
);
}
}
}
// ===========================================================================
// 10. Factor Reconstruction Quality
// ===========================================================================
/// Create a low-rank matrix, factor it, reconstruct, and verify error is low.
#[test]
fn test_factor_reconstruction_quality() {
let m = 16;
let n = 16;
// Rank-1 matrix: data[i][j] = (i+1)*(j+1) / (m*n).
let data: Vec<f32> = (0..m * n)
.map(|idx| {
let (i, j) = (idx / n, idx % n);
(i as f32 + 1.0) * (j as f32 + 1.0) / (m * n) as f32
})
.collect();
let factors = FactorSet::from_data(&data, m, n, 1);
assert_eq!(factors.m, m);
assert_eq!(factors.n, n);
assert_eq!(factors.k, 1);
let reconstructed = factors.reconstruct();
assert_eq!(reconstructed.len(), m * n);
let max_abs: f32 = data.iter().map(|v| v.abs()).fold(0.0f32, f32::max);
let mut max_err = 0.0f32;
for i in 0..m * n {
let err = (data[i] - reconstructed[i]).abs();
if err > max_err {
max_err = err;
}
}
assert!(
max_err < max_abs * 0.01,
"factor reconstruction error too high: max_err={} (max_abs={})",
max_err,
max_abs
);
// Factor storage should be smaller than the full matrix.
assert!(factors.storage_bytes() > 0);
assert!(
factors.storage_bytes() < m * n * 4,
"factor storage {} should be < original {}",
factors.storage_bytes(),
m * n * 4
);
}
// ===========================================================================
// 11. Witness Logging Integration
// ===========================================================================
/// Record access, tier-change, and eviction events; verify counters and
/// flip-rate calculation.
#[test]
fn test_witness_logging() {
let mut log = WitnessLog::new(256);
let mut store = TieredStore::new(4096);
let key = make_key(1, 0);
store.put(key, &vec![1.0f32; 64], Tier::Tier1, 0).unwrap();
log.record(
0,
WitnessEvent::Access {
key,
score: 0.95,
tier: Tier::Tier1,
},
);
log.record(
100,
WitnessEvent::TierChange {
key,
from_tier: Tier::Tier1,
to_tier: Tier::Tier2,
score: 0.45,
reason: TierChangeReason::ScoreDowngrade,
},
);
store.evict(key, ReconstructPolicy::None).unwrap();
log.record(
200,
WitnessEvent::Eviction {
key,
score: 0.05,
bytes_freed: 64,
},
);
assert_eq!(log.len(), 3);
assert_eq!(log.count_tier_changes(), 1);
assert_eq!(log.count_evictions(), 1);
assert_eq!(log.count_checksum_failures(), 0);
let recent = log.recent(2);
assert_eq!(recent.len(), 2);
assert_eq!(recent[0].timestamp, 100);
assert_eq!(recent[1].timestamp, 200);
// One tier change across 1 block in the window = flip rate 1.0.
let rate = log.tier_flip_rate(300, 1);
assert!(
(rate - 1.0).abs() < 1e-6,
"expected flip rate 1.0, got {}",
rate
);
}

View File

@@ -0,0 +1,225 @@
#![cfg(feature = "persistence")]
use ruvector_temporal_tensor::persistence::{FileBlockIO, FileMetaLog};
use ruvector_temporal_tensor::store::{
BlockIO, BlockKey, BlockMeta, DType, MetaLog, ReconstructPolicy, Tier,
};
use std::path::PathBuf;
fn test_dir(name: &str) -> PathBuf {
let dir = std::env::temp_dir().join(format!("ruvector_test_{}", name));
let _ = std::fs::remove_dir_all(&dir);
std::fs::create_dir_all(&dir).unwrap();
dir
}
fn cleanup(dir: &PathBuf) {
let _ = std::fs::remove_dir_all(dir);
}
fn make_key(id: u128, idx: u32) -> BlockKey {
BlockKey {
tensor_id: id,
block_index: idx,
}
}
fn make_meta(key: BlockKey, tier: Tier) -> BlockMeta {
BlockMeta {
key,
dtype: DType::F32,
tier,
bits: 8,
scale: 0.5,
zero_point: 0,
created_at: 100,
last_access_at: 200,
access_count: 5,
ema_rate: 0.1,
window: 0xFF,
checksum: 0xDEADBEEF,
reconstruct: ReconstructPolicy::None,
tier_age: 10,
lineage_parent: None,
block_bytes: 64,
}
}
// -----------------------------------------------------------------------
// FileBlockIO tests
// -----------------------------------------------------------------------
#[test]
fn test_file_block_io_write_read() {
let dir = test_dir("block_io_write_read");
let mut bio = FileBlockIO::new(&dir).unwrap();
let key = make_key(1, 0);
let data = vec![0xAB, 0xCD, 0xEF, 0x01, 0x23, 0x45, 0x67, 0x89];
bio.write_block(Tier::Tier1, key, &data).unwrap();
let mut dst = vec![0u8; 32];
let n = bio.read_block(Tier::Tier1, key, &mut dst).unwrap();
assert_eq!(n, data.len());
assert_eq!(&dst[..n], &data[..]);
cleanup(&dir);
}
#[test]
fn test_file_block_io_different_tiers() {
let dir = test_dir("block_io_tiers");
let mut bio = FileBlockIO::new(&dir).unwrap();
let key = make_key(1, 0);
let data1 = vec![1u8; 16];
let data2 = vec![2u8; 8];
let data3 = vec![3u8; 4];
bio.write_block(Tier::Tier1, key, &data1).unwrap();
bio.write_block(Tier::Tier2, key, &data2).unwrap();
bio.write_block(Tier::Tier3, key, &data3).unwrap();
let mut buf = vec![0u8; 32];
let n1 = bio.read_block(Tier::Tier1, key, &mut buf).unwrap();
assert_eq!(&buf[..n1], &data1[..]);
let n2 = bio.read_block(Tier::Tier2, key, &mut buf).unwrap();
assert_eq!(&buf[..n2], &data2[..]);
let n3 = bio.read_block(Tier::Tier3, key, &mut buf).unwrap();
assert_eq!(&buf[..n3], &data3[..]);
cleanup(&dir);
}
#[test]
fn test_file_block_io_delete() {
let dir = test_dir("block_io_delete");
let mut bio = FileBlockIO::new(&dir).unwrap();
let key = make_key(1, 0);
bio.write_block(Tier::Tier1, key, &[1, 2, 3]).unwrap();
bio.delete_block(Tier::Tier1, key).unwrap();
let mut buf = vec![0u8; 32];
let result = bio.read_block(Tier::Tier1, key, &mut buf);
assert!(result.is_err() || result.unwrap() == 0);
cleanup(&dir);
}
#[test]
fn test_file_block_io_overwrite() {
let dir = test_dir("block_io_overwrite");
let mut bio = FileBlockIO::new(&dir).unwrap();
let key = make_key(1, 0);
bio.write_block(Tier::Tier1, key, &[1, 2, 3]).unwrap();
bio.write_block(Tier::Tier1, key, &[4, 5, 6, 7]).unwrap();
let mut buf = vec![0u8; 32];
let n = bio.read_block(Tier::Tier1, key, &mut buf).unwrap();
assert_eq!(&buf[..n], &[4, 5, 6, 7]);
cleanup(&dir);
}
#[test]
fn test_file_block_io_missing_key() {
let dir = test_dir("block_io_missing");
let bio = FileBlockIO::new(&dir).unwrap();
let mut buf = vec![0u8; 32];
let result = bio.read_block(Tier::Tier1, make_key(99, 0), &mut buf);
assert!(result.is_err() || result.unwrap() == 0);
cleanup(&dir);
}
// -----------------------------------------------------------------------
// FileMetaLog tests
// -----------------------------------------------------------------------
#[test]
fn test_file_meta_log_append_get() {
let dir = test_dir("meta_log_append");
let mut log = FileMetaLog::new(&dir).unwrap();
let key = make_key(1, 0);
let meta = make_meta(key, Tier::Tier1);
log.append(&meta).unwrap();
let retrieved = log.get(key).unwrap();
assert_eq!(retrieved.key, key);
assert_eq!(retrieved.tier, Tier::Tier1);
assert_eq!(retrieved.bits, 8);
assert!((retrieved.scale - 0.5).abs() < 1e-6);
assert_eq!(retrieved.checksum, 0xDEADBEEF);
cleanup(&dir);
}
#[test]
fn test_file_meta_log_upsert() {
let dir = test_dir("meta_log_upsert");
let mut log = FileMetaLog::new(&dir).unwrap();
let key = make_key(1, 0);
let meta1 = make_meta(key, Tier::Tier1);
log.append(&meta1).unwrap();
let mut meta2 = make_meta(key, Tier::Tier2);
meta2.bits = 7;
log.append(&meta2).unwrap();
let retrieved = log.get(key).unwrap();
assert_eq!(retrieved.tier, Tier::Tier2);
assert_eq!(retrieved.bits, 7);
cleanup(&dir);
}
#[test]
fn test_file_meta_log_iter() {
let dir = test_dir("meta_log_iter");
let mut log = FileMetaLog::new(&dir).unwrap();
for i in 0..5u128 {
let key = make_key(i, 0);
log.append(&make_meta(key, Tier::Tier1)).unwrap();
}
let count = log.iter().count();
assert_eq!(count, 5);
cleanup(&dir);
}
#[test]
fn test_file_meta_log_missing_key() {
let dir = test_dir("meta_log_missing");
let log = FileMetaLog::new(&dir).unwrap();
assert!(log.get(make_key(99, 0)).is_none());
cleanup(&dir);
}
#[test]
fn test_file_meta_log_multiple_blocks_same_tensor() {
let dir = test_dir("meta_log_multi_block");
let mut log = FileMetaLog::new(&dir).unwrap();
for idx in 0..3u32 {
let key = make_key(1, idx);
log.append(&make_meta(key, Tier::Tier1)).unwrap();
}
assert!(log.get(make_key(1, 0)).is_some());
assert!(log.get(make_key(1, 1)).is_some());
assert!(log.get(make_key(1, 2)).is_some());
assert!(log.get(make_key(1, 3)).is_none());
cleanup(&dir);
}

View File

@@ -0,0 +1,821 @@
//! Property-based roundtrip tests for temporal tensor compression.
//!
//! Verifies quantization roundtrip correctness across many random inputs
//! using a deterministic PRNG. No external dependencies.
//!
//! Run with:
//! ```sh
//! cargo test --release -p ruvector-temporal-tensor --test property_tests -- --nocapture
//! ```
use ruvector_temporal_tensor::bitpack;
use ruvector_temporal_tensor::delta;
use ruvector_temporal_tensor::f16;
use ruvector_temporal_tensor::quantizer;
use ruvector_temporal_tensor::segment;
use ruvector_temporal_tensor::tiering::{self, BlockMeta, TierConfig};
// ---------------------------------------------------------------------------
// Deterministic PRNG (LCG) -- no external deps
// ---------------------------------------------------------------------------
/// Simple linear congruential generator. Constants from Knuth MMIX.
struct SimpleRng {
state: u64,
}
impl SimpleRng {
fn new(seed: u64) -> Self {
Self { state: seed }
}
fn next_u64(&mut self) -> u64 {
self.state = self
.state
.wrapping_mul(6364136223846793005)
.wrapping_add(1442695040888963407);
self.state
}
fn next_f32(&mut self) -> f32 {
(self.next_u64() >> 40) as f32 / (1u64 << 24) as f32
}
fn next_f32_range(&mut self, lo: f32, hi: f32) -> f32 {
lo + self.next_f32() * (hi - lo)
}
fn next_usize_range(&mut self, lo: usize, hi: usize) -> usize {
let range = (hi - lo) as u64;
if range == 0 {
return lo;
}
lo + (self.next_u64() % range) as usize
}
}
// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------
const GROUP_LEN: usize = 64;
/// Generate a random f32 vector of the given length with values in [lo, hi].
fn random_vec(rng: &mut SimpleRng, len: usize, lo: f32, hi: f32) -> Vec<f32> {
(0..len).map(|_| rng.next_f32_range(lo, hi)).collect()
}
/// Compute group-level maximum absolute values for error bounding.
fn group_max_abs(frame: &[f32], group_len: usize) -> Vec<f32> {
frame
.chunks(group_len)
.map(|chunk| {
chunk
.iter()
.filter(|v| v.is_finite())
.map(|v| v.abs())
.fold(0.0f32, f32::max)
})
.collect()
}
// ---------------------------------------------------------------------------
// 1. Quantize/Dequant Roundtrip Property
// ---------------------------------------------------------------------------
#[test]
fn prop_roundtrip_error_bounded() {
let mut rng = SimpleRng::new(0xDEAD_BEEF_CAFE_BABE);
// Error bounds as fraction of each group's max absolute value.
// The absolute error per element is bounded by:
// scale * 1 (one quantization step) + f16 rounding (~0.1% of scale)
// where scale = group_max_abs / qmax. So the error fraction of group_max is
// approximately 1/qmax + small f16 term.
// 8-bit: qmax=127, ~0.8% + margin -> 1%
// 7-bit: qmax=63, ~1.6% + margin -> 2%
// 5-bit: qmax=15, ~6.7% + margin -> 7%
// 3-bit: qmax=3, ~33% + margin -> 35%
let bit_configs: &[(u8, f32)] = &[
(8, 0.01), // 8-bit: < 1% of group max
(7, 0.02), // 7-bit: < 2% of group max
(5, 0.07), // 5-bit: < 7% of group max
(3, 0.35), // 3-bit: < 35% of group max
];
for trial in 0..1000 {
let len = rng.next_usize_range(64, 513); // 64..512 inclusive
let frame = random_vec(&mut rng, len, -10.0, 10.0);
for &(bits, max_err_frac) in bit_configs {
let scales = quantizer::compute_scales(&frame, GROUP_LEN, bits);
let scales_f32 = quantizer::scales_to_f32(&scales);
let mut packed = Vec::new();
quantizer::quantize_and_pack_f32(&frame, &scales_f32, GROUP_LEN, bits, &mut packed);
let mut decoded = Vec::new();
quantizer::dequantize_f32(
&packed,
&scales_f32,
GROUP_LEN,
bits,
frame.len(),
1,
&mut decoded,
);
assert_eq!(
decoded.len(),
frame.len(),
"trial={trial}, bits={bits}: length mismatch"
);
// Compute per-group max absolute value for error bounding.
let gmax = group_max_abs(&frame, GROUP_LEN);
for (i, (&orig, &dec)) in frame.iter().zip(decoded.iter()).enumerate() {
let abs_err = (orig - dec).abs();
let group_idx = i / GROUP_LEN;
let group_m = if group_idx < gmax.len() {
gmax[group_idx]
} else {
1.0
};
// Bound: max_err_frac * group_max + small absolute floor for near-zero groups.
let bound = max_err_frac * group_m + 1e-6;
assert!(
abs_err <= bound,
"trial={trial}, bits={bits}, i={i}: orig={orig}, dec={dec}, \
abs_err={abs_err}, bound={bound}, group_max={group_m}"
);
}
}
}
}
// ---------------------------------------------------------------------------
// 2. Bit Packing Roundtrip Property
// ---------------------------------------------------------------------------
#[test]
fn prop_bitpack_roundtrip() {
let mut rng = SimpleRng::new(0x1234_5678_9ABC_DEF0);
let bit_widths: &[u32] = &[3, 5, 7, 8];
for _trial in 0..1000 {
let count = rng.next_usize_range(1, 513);
for &bits in bit_widths {
let max_val = (1u32 << bits) - 1;
let codes: Vec<u32> = (0..count)
.map(|_| (rng.next_u64() as u32) % (max_val + 1))
.collect();
let mut packed = Vec::new();
bitpack::pack(&codes, bits, &mut packed);
let mut unpacked = Vec::new();
bitpack::unpack(&packed, bits, count, &mut unpacked);
assert_eq!(
codes, unpacked,
"bits={bits}, count={count}: pack/unpack mismatch"
);
}
}
}
// ---------------------------------------------------------------------------
// 3. Segment Encode/Decode Property
// ---------------------------------------------------------------------------
#[test]
fn prop_segment_roundtrip() {
let mut rng = SimpleRng::new(0xFEED_FACE_DEAD_C0DE);
let tensor_lens: &[usize] = &[32, 64, 128, 256, 512];
let frame_counts: &[usize] = &[1, 2, 5, 10, 20];
let bit_widths: &[u8] = &[3, 5, 7, 8];
for _trial in 0..200 {
let tensor_len = tensor_lens[rng.next_usize_range(0, tensor_lens.len())];
let frame_count = frame_counts[rng.next_usize_range(0, frame_counts.len())];
let bits = bit_widths[rng.next_usize_range(0, bit_widths.len())];
// Generate the first frame and compute scales from it (shared across frames).
let first_frame = random_vec(&mut rng, tensor_len, -5.0, 5.0);
let scales = quantizer::compute_scales(&first_frame, GROUP_LEN, bits);
let scales_f32 = quantizer::scales_to_f32(&scales);
// Quantize all frames with the same scales.
let mut packed = Vec::new();
quantizer::quantize_and_pack_f32(&first_frame, &scales_f32, GROUP_LEN, bits, &mut packed);
for _ in 1..frame_count {
// Subsequent frames use values within the first frame's range to fit scales.
let frame = random_vec(&mut rng, tensor_len, -4.0, 4.0);
quantizer::quantize_and_pack_f32(&frame, &scales_f32, GROUP_LEN, bits, &mut packed);
}
// Encode into segment format.
let mut seg = Vec::new();
segment::encode(
bits,
GROUP_LEN as u32,
tensor_len as u32,
frame_count as u32,
&scales,
&packed,
&mut seg,
);
// Decode the segment.
let mut decoded = Vec::new();
segment::decode(&seg, &mut decoded);
assert_eq!(
decoded.len(),
tensor_len * frame_count,
"trial={_trial}, bits={bits}, tensor_len={tensor_len}, frames={frame_count}: \
decoded length mismatch"
);
// Parse the header and verify metadata.
let header = segment::parse_header(&seg).expect("header should parse");
assert_eq!(header.bits, bits);
assert_eq!(header.tensor_len, tensor_len as u32);
assert_eq!(header.frame_count, frame_count as u32);
assert_eq!(header.group_len, GROUP_LEN as u32);
}
}
// ---------------------------------------------------------------------------
// 4. f16 Roundtrip Property
// ---------------------------------------------------------------------------
#[test]
fn prop_f16_roundtrip() {
let mut rng = SimpleRng::new(0xAAAA_BBBB_CCCC_DDDD);
for _trial in 0..10_000 {
// Generate value in scale-relevant range [1e-4, 1e4].
let v = rng.next_f32_range(1e-4, 1e4);
// Randomly negate half the values.
let v = if rng.next_u64() & 1 == 0 { v } else { -v };
let h = f16::f32_to_f16_bits(v);
let back = f16::f16_bits_to_f32(h);
// f16 has ~0.1% relative error for normal values in this range.
let rel_err = ((back - v) / v).abs();
assert!(
rel_err < 0.002,
"trial={_trial}: v={v}, back={back}, rel_err={rel_err}"
);
}
}
// ---------------------------------------------------------------------------
// 5. Delta Compute/Apply Property
// ---------------------------------------------------------------------------
#[test]
fn prop_delta_apply_recovers_new() {
let mut rng = SimpleRng::new(0x0123_4567_89AB_CDEF);
for trial in 0..500 {
let len = rng.next_usize_range(8, 257);
let old = random_vec(&mut rng, len, -5.0, 5.0);
// Create "new" as old with a small number of perturbations.
let mut new = old.clone();
let num_changes = rng.next_usize_range(1, (len / 4).max(2));
for _ in 0..num_changes {
let idx = rng.next_usize_range(0, len);
new[idx] += rng.next_f32_range(-1.0, 1.0);
}
let threshold = 0.001;
let max_change_frac = 0.8;
let result =
delta::compute_delta(&old, &new, trial as u128, 0, 0, threshold, max_change_frac);
match result {
Some(d) => {
// Apply delta to old, verify it approximates new.
let mut reconstructed = old.clone();
delta::apply_delta(&mut reconstructed, &d);
for i in 0..len {
let err = (reconstructed[i] - new[i]).abs();
// Two sources of error:
// 1. Entries below threshold are not captured in the delta,
// so the reconstruction error for those is up to `threshold`.
// 2. Captured entries have i16 quantization error of at most
// delta_scale / 2 (half a quantization step).
let tolerance = threshold + d.delta_scale * 1.5 + 1e-6;
assert!(
err <= tolerance,
"trial={trial}, i={i}: recon={}, new={}, err={err}, tol={tolerance}",
reconstructed[i],
new[i]
);
}
}
None => {
// Delta was too large (>= max_change_fraction).
// Verify that indeed many values changed.
let changed = old
.iter()
.zip(new.iter())
.filter(|(&o, &n)| (o - n).abs() >= threshold)
.count();
let fraction = changed as f32 / len as f32;
assert!(
fraction >= max_change_frac,
"trial={trial}: delta was None but change fraction={fraction} < {max_change_frac}"
);
}
}
}
}
// ---------------------------------------------------------------------------
// 6. Compression Ratio Property
// ---------------------------------------------------------------------------
#[test]
fn prop_compression_ratio_matches_theory() {
let mut rng = SimpleRng::new(0xCAFE_D00D_BEEF_FEED);
let expected: &[(u8, f32)] = &[(8, 3.5), (7, 4.0), (5, 5.5), (3, 8.5)];
for &(bits, min_ratio) in expected {
// Use a 512-element tensor with group_len=64 for consistent measurement.
let frame = random_vec(&mut rng, 512, -1.0, 1.0);
let scales = quantizer::compute_scales(&frame, GROUP_LEN, bits);
let mut packed = Vec::new();
quantizer::quantize_and_pack(&frame, &scales, GROUP_LEN, bits, &mut packed);
let raw_bytes = frame.len() * 4; // f32 = 4 bytes
let compressed = packed.len() + scales.len() * 2; // packed data + f16 scales
let ratio = raw_bytes as f32 / compressed as f32;
assert!(
ratio >= min_ratio,
"bits={bits}: ratio={ratio:.2}x < expected={min_ratio}x \
(raw={raw_bytes}, compressed={compressed})"
);
}
}
// ---------------------------------------------------------------------------
// 7. Score Monotonicity Property
// ---------------------------------------------------------------------------
#[test]
fn prop_score_monotonic_with_access() {
let mut rng = SimpleRng::new(0x7777_8888_9999_AAAA);
let config = TierConfig::default();
for _trial in 0..100 {
let start_tick = rng.next_u64() % 1000;
let mut meta = BlockMeta::new(start_tick);
// Score before any touch.
let score_before = tiering::compute_score(&config, start_tick, &meta);
// Touch the block.
tiering::touch(&config, start_tick + 1, &mut meta);
let score_after_touch = tiering::compute_score(&config, start_tick + 1, &meta);
// Touching should increase (or at minimum maintain) the score.
assert!(
score_after_touch >= score_before - 1e-6,
"trial={_trial}: score decreased after touch: \
before={score_before}, after={score_after_touch}"
);
// Now let time pass without access -- score should decrease.
let score_at_touch = tiering::compute_score(&config, start_tick + 1, &meta);
let score_later = tiering::compute_score(&config, start_tick + 1000, &meta);
assert!(
score_later <= score_at_touch + 1e-6,
"trial={_trial}: score increased without access: \
at_touch={score_at_touch}, later={score_later}"
);
}
}
// ---------------------------------------------------------------------------
// 8. Zero Vector Property
// ---------------------------------------------------------------------------
#[test]
fn prop_zero_vector_roundtrip() {
let bit_widths: &[u8] = &[3, 5, 7, 8];
for &len in &[64, 128, 256, 512] {
let frame = vec![0.0f32; len];
for &bits in bit_widths {
let scales = quantizer::compute_scales(&frame, GROUP_LEN, bits);
let scales_f32 = quantizer::scales_to_f32(&scales);
// All scales should be zero for a zero vector.
for (i, &s) in scales_f32.iter().enumerate() {
assert_eq!(
s, 0.0,
"len={len}, bits={bits}, group={i}: scale should be 0.0, got {s}"
);
}
let mut packed = Vec::new();
quantizer::quantize_and_pack_f32(&frame, &scales_f32, GROUP_LEN, bits, &mut packed);
let mut decoded = Vec::new();
quantizer::dequantize_f32(&packed, &scales_f32, GROUP_LEN, bits, len, 1, &mut decoded);
assert_eq!(decoded.len(), len);
for (i, &v) in decoded.iter().enumerate() {
assert_eq!(
v, 0.0,
"len={len}, bits={bits}, i={i}: expected 0.0, got {v}"
);
}
}
}
}
// ---------------------------------------------------------------------------
// 9. Single-Value (Uniform) Vector Property
// ---------------------------------------------------------------------------
#[test]
fn prop_uniform_vector_roundtrip() {
let mut rng = SimpleRng::new(0xBBBB_CCCC_DDDD_EEEE);
let bit_widths: &[u8] = &[3, 5, 7, 8];
for _trial in 0..200 {
let len = rng.next_usize_range(64, 513);
let value = rng.next_f32_range(-10.0, 10.0);
let frame = vec![value; len];
for &bits in bit_widths {
let qmax = bitpack::qmax_from_bits(bits);
if qmax == 0 {
continue;
}
let scales = quantizer::compute_scales(&frame, GROUP_LEN, bits);
let scales_f32 = quantizer::scales_to_f32(&scales);
let mut packed = Vec::new();
quantizer::quantize_and_pack_f32(&frame, &scales_f32, GROUP_LEN, bits, &mut packed);
let mut decoded = Vec::new();
quantizer::dequantize_f32(&packed, &scales_f32, GROUP_LEN, bits, len, 1, &mut decoded);
assert_eq!(decoded.len(), len);
// For a uniform vector, the quantization step is value.abs() / qmax.
// Max error should be at most half a step (rounding) plus f16 scale error.
let step = if value.abs() > 0.0 {
value.abs() / qmax as f32
} else {
0.0
};
// Allow step/2 plus a small f16 rounding margin.
let max_err = step * 0.5 + value.abs() * 0.002 + 1e-6;
for (i, &dec) in decoded.iter().enumerate() {
let err = (dec - value).abs();
assert!(
err <= max_err,
"trial={_trial}, bits={bits}, i={i}: value={value}, dec={dec}, \
err={err}, max_err={max_err}, step={step}"
);
}
}
}
}
// ---------------------------------------------------------------------------
// 10. Extreme Value Property
// ---------------------------------------------------------------------------
#[test]
fn prop_extreme_values_dont_panic() {
let bit_widths: &[u8] = &[3, 5, 7, 8];
// Frames where scales stay within f16 representable range -- decoded values
// must be finite.
let finite_frames: Vec<Vec<f32>> = vec![
// Very small positive values
vec![f32::MIN_POSITIVE; 128],
// Contains infinities and NaN (quantizer maps non-finite to 0)
{
let mut v = vec![1.0f32; 128];
v[0] = f32::INFINITY;
v[1] = f32::NEG_INFINITY;
v[2] = f32::NAN;
v[3] = -0.0;
v
},
// All subnormal
vec![1e-40f32; 128],
// Alternating zero and large (within f16 scale range)
(0..128)
.map(|i| if i % 2 == 0 { 0.0 } else { 1e4 })
.collect(),
];
// Frames with magnitudes that overflow f16 scales -- we only assert
// no panics and correct output length. The decoded values may be NaN/Inf
// because scale overflows to f16 infinity.
let overflow_frames: Vec<Vec<f32>> = vec![
// All f32::MAX
vec![f32::MAX; 128],
// All f32::MIN (most negative finite)
vec![f32::MIN; 128],
// Mixed signs of large magnitude
(0..128)
.map(|i| if i % 2 == 0 { f32::MAX } else { f32::MIN })
.collect(),
// Mix of tiny and huge
(0..128)
.map(|i| {
if i % 3 == 0 {
f32::MIN_POSITIVE
} else if i % 3 == 1 {
1e30
} else {
-1e30
}
})
.collect(),
];
// Test finite-output frames: no panics, correct length, all decoded finite.
for (frame_idx, frame) in finite_frames.iter().enumerate() {
for &bits in bit_widths {
let scales = quantizer::compute_scales(frame, GROUP_LEN, bits);
let scales_f32 = quantizer::scales_to_f32(&scales);
let mut packed = Vec::new();
quantizer::quantize_and_pack_f32(frame, &scales_f32, GROUP_LEN, bits, &mut packed);
let mut decoded = Vec::new();
quantizer::dequantize_f32(
&packed,
&scales_f32,
GROUP_LEN,
bits,
frame.len(),
1,
&mut decoded,
);
assert_eq!(
decoded.len(),
frame.len(),
"finite frame_idx={frame_idx}, bits={bits}: length mismatch"
);
for (i, &d) in decoded.iter().enumerate() {
assert!(
d.is_finite(),
"finite frame_idx={frame_idx}, bits={bits}, i={i}: \
decoded value is not finite: {d}"
);
}
}
}
// Test overflow frames: no panics, correct length (decoded may contain NaN/Inf).
for (frame_idx, frame) in overflow_frames.iter().enumerate() {
for &bits in bit_widths {
let scales = quantizer::compute_scales(frame, GROUP_LEN, bits);
let scales_f32 = quantizer::scales_to_f32(&scales);
let mut packed = Vec::new();
quantizer::quantize_and_pack_f32(frame, &scales_f32, GROUP_LEN, bits, &mut packed);
let mut decoded = Vec::new();
quantizer::dequantize_f32(
&packed,
&scales_f32,
GROUP_LEN,
bits,
frame.len(),
1,
&mut decoded,
);
assert_eq!(
decoded.len(),
frame.len(),
"overflow frame_idx={frame_idx}, bits={bits}: length mismatch"
);
}
}
// Bitpack roundtrip with boundary codes -- must not panic and must be exact.
for &bits in bit_widths {
let qmax = bitpack::qmax_from_bits(bits) as u32;
if qmax > 0 {
let max_code = qmax * 2;
let codes: Vec<u32> = (0..128).map(|i| i as u32 % (max_code + 1)).collect();
let mut bp = Vec::new();
bitpack::pack(&codes, bits as u32, &mut bp);
let mut unpacked = Vec::new();
bitpack::unpack(&bp, bits as u32, codes.len(), &mut unpacked);
assert_eq!(codes, unpacked);
}
}
}
// ---------------------------------------------------------------------------
// 11. Segment Compression Ratio is Positive
// ---------------------------------------------------------------------------
#[test]
fn prop_segment_compression_ratio_positive() {
let mut rng = SimpleRng::new(0x1111_2222_3333_4444);
for _trial in 0..100 {
let tensor_len = 128;
let bits = [3u8, 5, 7, 8][rng.next_usize_range(0, 4)];
let frame = random_vec(&mut rng, tensor_len, -1.0, 1.0);
let scales = quantizer::compute_scales(&frame, GROUP_LEN, bits);
let mut packed = Vec::new();
quantizer::quantize_and_pack(&frame, &scales, GROUP_LEN, bits, &mut packed);
let mut seg = Vec::new();
segment::encode(
bits,
GROUP_LEN as u32,
tensor_len as u32,
1,
&scales,
&packed,
&mut seg,
);
let ratio = segment::compression_ratio(&seg);
assert!(
ratio > 1.0,
"trial={_trial}, bits={bits}: compression ratio {ratio} should be > 1.0"
);
}
}
// ---------------------------------------------------------------------------
// 12. Single-Frame Decode Matches Full Decode
// ---------------------------------------------------------------------------
#[test]
fn prop_single_frame_decode_consistency() {
let mut rng = SimpleRng::new(0x5555_6666_7777_8888);
for _trial in 0..100 {
let tensor_len = 64;
let frame_count = rng.next_usize_range(1, 6);
let bits = [3u8, 5, 7, 8][rng.next_usize_range(0, 4)];
let first_frame = random_vec(&mut rng, tensor_len, -3.0, 3.0);
let scales = quantizer::compute_scales(&first_frame, GROUP_LEN, bits);
let scales_f32 = quantizer::scales_to_f32(&scales);
let mut packed = Vec::new();
quantizer::quantize_and_pack_f32(&first_frame, &scales_f32, GROUP_LEN, bits, &mut packed);
for _ in 1..frame_count {
let frame = random_vec(&mut rng, tensor_len, -2.5, 2.5);
quantizer::quantize_and_pack_f32(&frame, &scales_f32, GROUP_LEN, bits, &mut packed);
}
let mut seg = Vec::new();
segment::encode(
bits,
GROUP_LEN as u32,
tensor_len as u32,
frame_count as u32,
&scales,
&packed,
&mut seg,
);
// Full decode.
let mut all_decoded = Vec::new();
segment::decode(&seg, &mut all_decoded);
assert_eq!(all_decoded.len(), tensor_len * frame_count);
// Single-frame decode should match the corresponding slice.
for f in 0..frame_count {
let single = segment::decode_single_frame(&seg, f);
assert!(
single.is_some(),
"trial={_trial}, frame={f}: single-frame decode returned None"
);
let single = single.unwrap();
let expected = &all_decoded[f * tensor_len..(f + 1) * tensor_len];
assert_eq!(
single.len(),
expected.len(),
"trial={_trial}, frame={f}: length mismatch"
);
for (i, (&s, &e)) in single.iter().zip(expected.iter()).enumerate() {
assert!(
(s - e).abs() < 1e-6,
"trial={_trial}, frame={f}, i={i}: single={s}, full={e}"
);
}
}
}
}
// ---------------------------------------------------------------------------
// 13. Delta Encode/Decode Binary Roundtrip
// ---------------------------------------------------------------------------
#[test]
fn prop_delta_encode_decode_binary() {
let mut rng = SimpleRng::new(0x9999_0000_1111_2222);
for trial in 0..500 {
let nnz = rng.next_usize_range(0, 100);
let entries: Vec<delta::SparseEntry> = (0..nnz)
.map(|_| delta::SparseEntry {
index: (rng.next_u64() % 65536) as u16,
value: (rng.next_u64() % 65536) as i16,
})
.collect();
let scale = rng.next_f32_range(1e-6, 100.0);
let record = delta::DeltaRecord {
header: delta::DeltaHeader {
tensor_id: rng.next_u64() as u128 | ((rng.next_u64() as u128) << 64),
block_index: rng.next_u64() as u32,
base_epoch: rng.next_u64(),
nnz: nnz as u16,
},
delta_scale: scale,
entries,
};
let bytes = delta::encode_delta(&record);
let decoded = delta::decode_delta(&bytes)
.unwrap_or_else(|e| panic!("trial={trial}: decode failed: {e:?}"));
assert_eq!(decoded.header.tensor_id, record.header.tensor_id);
assert_eq!(decoded.header.block_index, record.header.block_index);
assert_eq!(decoded.header.base_epoch, record.header.base_epoch);
assert_eq!(decoded.header.nnz, record.header.nnz);
assert!(
(decoded.delta_scale - record.delta_scale).abs() < 1e-10,
"trial={trial}: scale mismatch"
);
assert_eq!(decoded.entries.len(), record.entries.len());
for (i, (a, b)) in decoded
.entries
.iter()
.zip(record.entries.iter())
.enumerate()
{
assert_eq!(a.index, b.index, "trial={trial}, entry={i}: index mismatch");
assert_eq!(a.value, b.value, "trial={trial}, entry={i}: value mismatch");
}
}
}
// ---------------------------------------------------------------------------
// 14. Quantization is Deterministic
// ---------------------------------------------------------------------------
#[test]
fn prop_quantization_deterministic() {
let mut rng = SimpleRng::new(0xABCD_EF01_2345_6789);
for _trial in 0..200 {
let len = rng.next_usize_range(64, 257);
let frame = random_vec(&mut rng, len, -5.0, 5.0);
let bits = [3u8, 5, 7, 8][rng.next_usize_range(0, 4)];
let scales = quantizer::compute_scales(&frame, GROUP_LEN, bits);
let scales_f32 = quantizer::scales_to_f32(&scales);
let mut packed1 = Vec::new();
quantizer::quantize_and_pack_f32(&frame, &scales_f32, GROUP_LEN, bits, &mut packed1);
let mut packed2 = Vec::new();
quantizer::quantize_and_pack_f32(&frame, &scales_f32, GROUP_LEN, bits, &mut packed2);
assert_eq!(
packed1, packed2,
"trial={_trial}, bits={bits}: quantization is not deterministic"
);
}
}

View File

@@ -0,0 +1,920 @@
//! Stress and fuzz-like tests for temporal tensor compression.
//!
//! Exercises the storage engine, delta chains, and checksum integrity under
//! heavy random workloads using a deterministic PRNG. No external dependencies.
//!
//! Run with:
//! ```sh
//! cargo test --release -p ruvector-temporal-tensor --test stress_tests -- --nocapture
//! ```
use ruvector_temporal_tensor::delta::{compute_delta, DeltaChain};
use ruvector_temporal_tensor::store::{BlockKey, ReconstructPolicy, StoreError, Tier, TieredStore};
// ---------------------------------------------------------------------------
// Deterministic PRNG (LCG) -- same as other test files, no external deps
// ---------------------------------------------------------------------------
/// Simple linear congruential generator. Constants from Knuth MMIX.
struct SimpleRng {
state: u64,
}
impl SimpleRng {
fn new(seed: u64) -> Self {
Self { state: seed }
}
fn next_u64(&mut self) -> u64 {
self.state = self
.state
.wrapping_mul(6_364_136_223_846_793_005)
.wrapping_add(1_442_695_040_888_963_407);
self.state
}
fn next_f64(&mut self) -> f64 {
(self.next_u64() >> 11) as f64 / (1u64 << 53) as f64
}
fn next_f32(&mut self) -> f32 {
self.next_f64() as f32
}
fn next_f32_range(&mut self, lo: f32, hi: f32) -> f32 {
lo + self.next_f32() * (hi - lo)
}
fn next_usize_range(&mut self, lo: usize, hi: usize) -> usize {
let range = (hi - lo) as u64;
if range == 0 {
return lo;
}
lo + (self.next_u64() % range) as usize
}
}
// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------
fn make_key(tid: u128, idx: u32) -> BlockKey {
BlockKey {
tensor_id: tid,
block_index: idx,
}
}
fn random_tier(rng: &mut SimpleRng) -> Tier {
match rng.next_usize_range(0, 3) {
0 => Tier::Tier1,
1 => Tier::Tier2,
_ => Tier::Tier3,
}
}
fn random_data(rng: &mut SimpleRng, len: usize) -> Vec<f32> {
(0..len).map(|_| rng.next_f32_range(-1.0, 1.0)).collect()
}
// ===========================================================================
// 1. Random put/get/evict cycle
// ===========================================================================
/// Exercises the store with 5000 random operations (put 40%, get 30%,
/// touch 20%, evict 10%) on a pool of 200 block keys. After all
/// iterations the block count must equal `inserted - evicted`.
#[test]
fn test_random_put_get_evict_cycle() {
let mut store = TieredStore::new(4096);
let mut rng = SimpleRng::new(0xDEAD_BEEF);
const NUM_KEYS: usize = 200;
const NUM_ITERS: usize = 5_000;
const ELEM_COUNT: usize = 64;
// Track which keys have been inserted and not yet evicted.
let mut inserted: std::collections::HashSet<u32> = std::collections::HashSet::new();
let mut evicted: std::collections::HashSet<u32> = std::collections::HashSet::new();
for iter in 0..NUM_ITERS {
let roll = rng.next_usize_range(0, 100);
let key_idx = rng.next_usize_range(0, NUM_KEYS) as u32;
let key = make_key(1, key_idx);
let tick = iter as u64;
if roll < 40 {
// PUT (40%)
let data = random_data(&mut rng, ELEM_COUNT);
let tier = random_tier(&mut rng);
store.put(key, &data, tier, tick).unwrap();
inserted.insert(key_idx);
evicted.remove(&key_idx);
} else if roll < 70 {
// GET (30%)
let mut out = vec![0.0f32; ELEM_COUNT];
match store.get(key, &mut out, tick) {
Ok(n) => {
assert!(n > 0, "get returned 0 elements for an existing block");
assert!(n <= ELEM_COUNT);
}
Err(StoreError::BlockNotFound) => {
// Key was never inserted or was evicted -- valid.
}
Err(StoreError::TensorEvicted) => {
// Block was evicted to Tier0 -- valid.
assert!(
evicted.contains(&key_idx),
"TensorEvicted for key not in evicted set"
);
}
Err(e) => {
panic!("unexpected error on get at iter {}: {:?}", iter, e);
}
}
} else if roll < 90 {
// TOUCH (20%)
store.touch(key, tick);
} else {
// EVICT (10%)
match store.evict(key, ReconstructPolicy::None) {
Ok(()) => {
if inserted.contains(&key_idx) {
evicted.insert(key_idx);
}
}
Err(StoreError::BlockNotFound) => {
// Key never existed -- valid.
}
Err(e) => {
panic!("unexpected error on evict at iter {}: {:?}", iter, e);
}
}
}
}
// Final invariant: block_count = all unique keys ever put (including evicted ones,
// since eviction keeps metadata).
let all_known: std::collections::HashSet<u32> = inserted.union(&evicted).copied().collect();
assert_eq!(
store.block_count(),
all_known.len(),
"block_count mismatch after random cycle"
);
// Verify: non-evicted blocks are readable.
let live_keys: Vec<u32> = inserted.difference(&evicted).copied().collect();
for &kid in &live_keys {
let mut out = vec![0.0f32; ELEM_COUNT];
let key = make_key(1, kid);
let result = store.get(key, &mut out, NUM_ITERS as u64);
assert!(
result.is_ok(),
"live block {} should be readable, got {:?}",
kid,
result
);
}
println!(
"random_put_get_evict_cycle: {} iters, {} live blocks, {} evicted",
NUM_ITERS,
live_keys.len(),
evicted.len()
);
}
// ===========================================================================
// 2. Rapid tier oscillation (stress hysteresis)
// ===========================================================================
/// Puts 50 blocks at Tier1, then alternately touches 25 blocks intensively
/// (50 touches/tick) and ignores them for 500 ticks. Verifies that all
/// blocks remain readable and no panics occur during rapid access-pattern
/// changes.
#[test]
fn test_rapid_tier_oscillation() {
let mut store = TieredStore::new(4096);
let mut rng = SimpleRng::new(0xCAFE_BABE);
const NUM_BLOCKS: usize = 50;
const ELEM_COUNT: usize = 64;
const TOTAL_TICKS: u64 = 500;
const HOT_COUNT: usize = 25;
const TOUCHES_PER_TICK: usize = 50;
// Insert all blocks at Tier1.
let block_data: Vec<Vec<f32>> = (0..NUM_BLOCKS)
.map(|_| random_data(&mut rng, ELEM_COUNT))
.collect();
for i in 0..NUM_BLOCKS {
store
.put(make_key(2, i as u32), &block_data[i], Tier::Tier1, 0)
.unwrap();
}
assert_eq!(store.block_count(), NUM_BLOCKS);
// Oscillate: even ticks -> heavy touching of first HOT_COUNT blocks,
// odd ticks -> no touching (cold period).
for tick in 1..=TOTAL_TICKS {
if tick % 2 == 0 {
// Hot phase: touch first HOT_COUNT blocks repeatedly.
for _ in 0..TOUCHES_PER_TICK {
let idx = rng.next_usize_range(0, HOT_COUNT) as u32;
store.touch(make_key(2, idx), tick);
}
}
// Odd ticks: silence (no touches).
}
// All blocks must remain readable.
for i in 0..NUM_BLOCKS {
let key = make_key(2, i as u32);
let mut out = vec![0.0f32; ELEM_COUNT];
let n = store
.get(key, &mut out, TOTAL_TICKS + 1)
.unwrap_or_else(|e| panic!("block {} unreadable after oscillation: {:?}", i, e));
assert_eq!(n, ELEM_COUNT);
// Values must be finite.
for (j, &v) in out.iter().enumerate() {
assert!(v.is_finite(), "block {} elem {} is non-finite: {}", i, j, v);
}
}
// Verify metadata is intact for all blocks.
for i in 0..NUM_BLOCKS {
let m = store.meta(make_key(2, i as u32)).expect("meta missing");
assert!(
m.tier == Tier::Tier1 || m.tier == Tier::Tier2 || m.tier == Tier::Tier3,
"block {} has unexpected tier {:?}",
i,
m.tier
);
}
println!(
"rapid_tier_oscillation: {} ticks, {} blocks, no panics",
TOTAL_TICKS, NUM_BLOCKS
);
}
// ===========================================================================
// 3. Large block stress (memory pressure)
// ===========================================================================
/// Puts 500 blocks of 4096 elements each (total ~8MB at 8-bit), touches
/// them randomly, reads them all back verifying finite values, evicts half,
/// and verifies the other half is still readable and total_bytes decreased.
#[test]
fn test_large_block_stress() {
let mut store = TieredStore::new(4096);
let mut rng = SimpleRng::new(0x1234_5678);
const NUM_BLOCKS: usize = 500;
const ELEM_COUNT: usize = 4096;
// Insert all blocks at Tier1 (8-bit = 1 byte/elem = 4096 bytes/block).
for i in 0..NUM_BLOCKS {
let data = random_data(&mut rng, ELEM_COUNT);
store
.put(make_key(3, i as u32), &data, Tier::Tier1, i as u64)
.unwrap();
}
assert_eq!(store.block_count(), NUM_BLOCKS);
let bytes_before = store.total_bytes();
assert!(
bytes_before > 0,
"total_bytes should be positive after inserting {} blocks",
NUM_BLOCKS
);
println!(
"large_block_stress: {} blocks inserted, total_bytes = {}",
NUM_BLOCKS, bytes_before
);
// Touch all blocks randomly.
for _ in 0..NUM_BLOCKS {
let idx = rng.next_usize_range(0, NUM_BLOCKS) as u32;
store.touch(make_key(3, idx), NUM_BLOCKS as u64 + 1);
}
// Read all blocks back and verify finite values.
for i in 0..NUM_BLOCKS {
let key = make_key(3, i as u32);
let mut out = vec![0.0f32; ELEM_COUNT];
let n = store
.get(key, &mut out, NUM_BLOCKS as u64 + 2)
.unwrap_or_else(|e| panic!("block {} unreadable: {:?}", i, e));
assert_eq!(n, ELEM_COUNT);
for (j, &v) in out.iter().enumerate() {
assert!(v.is_finite(), "block {} elem {} is non-finite: {}", i, j, v);
}
}
// Evict the first half.
for i in 0..(NUM_BLOCKS / 2) {
store
.evict(make_key(3, i as u32), ReconstructPolicy::None)
.unwrap();
}
let bytes_after = store.total_bytes();
assert!(
bytes_after < bytes_before,
"total_bytes should decrease after evicting half: before={}, after={}",
bytes_before,
bytes_after
);
// Verify the second half is still readable.
for i in (NUM_BLOCKS / 2)..NUM_BLOCKS {
let key = make_key(3, i as u32);
let mut out = vec![0.0f32; ELEM_COUNT];
let n = store
.get(key, &mut out, NUM_BLOCKS as u64 + 3)
.unwrap_or_else(|e| {
panic!(
"block {} should still be readable after evicting first half: {:?}",
i, e
)
});
assert_eq!(n, ELEM_COUNT);
}
// Verify evicted blocks return TensorEvicted.
for i in 0..(NUM_BLOCKS / 2) {
let key = make_key(3, i as u32);
let mut out = vec![0.0f32; ELEM_COUNT];
let result = store.get(key, &mut out, NUM_BLOCKS as u64 + 4);
assert_eq!(
result,
Err(StoreError::TensorEvicted),
"evicted block {} should return TensorEvicted",
i
);
}
println!(
"large_block_stress: bytes before={}, after={}, reduction={}%",
bytes_before,
bytes_after,
((bytes_before - bytes_after) as f64 / bytes_before as f64 * 100.0) as u32
);
}
// ===========================================================================
// 4. Delta chain stress
// ===========================================================================
/// Creates a 1024-element base vector, builds a DeltaChain with max_depth=8,
/// appends 8 deltas each modifying ~5% of values, reconstructs and verifies
/// error < 1%, compacts, rebuilds to max, and checks that an extra append
/// yields DeltaChainTooLong.
#[test]
fn test_delta_chain_stress() {
let mut rng = SimpleRng::new(0xABCD_EF01);
const DIM: usize = 1024;
const MAX_DEPTH: u8 = 8;
const CHANGE_FRACTION: f32 = 0.05; // ~5% of values per delta
// Create a base vector with random values in [-1, 1].
let base: Vec<f32> = (0..DIM).map(|_| rng.next_f32_range(-1.0, 1.0)).collect();
let mut chain = DeltaChain::new(base.clone(), MAX_DEPTH);
// Build the expected ground-truth by applying modifications cumulatively.
let mut truth = base.clone();
// Append MAX_DEPTH deltas, each modifying ~5% of elements.
for epoch in 0..MAX_DEPTH {
let mut modified = truth.clone();
let num_changes = (DIM as f32 * CHANGE_FRACTION) as usize;
for _ in 0..num_changes {
let idx = rng.next_usize_range(0, DIM);
let perturbation = rng.next_f32_range(-0.1, 0.1);
modified[idx] += perturbation;
}
let delta = compute_delta(
&truth,
&modified,
42, // tensor_id
0, // block_index
epoch as u64, // base_epoch
1e-8, // threshold (very small to capture all changes)
1.0, // max_change_fraction (allow up to 100%)
)
.expect("compute_delta should succeed for small changes");
chain
.append(delta)
.unwrap_or_else(|e| panic!("append should succeed at depth {}: {:?}", epoch, e));
truth = modified;
}
assert_eq!(chain.chain_len(), MAX_DEPTH as usize);
// Reconstruct and verify error < 1%.
let reconstructed = chain.reconstruct();
assert_eq!(reconstructed.len(), DIM);
let mut max_err: f32 = 0.0;
for i in 0..DIM {
let err = (reconstructed[i] - truth[i]).abs();
if err > max_err {
max_err = err;
}
}
// The error comes from i16 quantization of deltas; for small perturbations
// the relative error should be well under 1% of the value range.
let value_range = truth.iter().fold(0.0f32, |acc, &v| acc.max(v.abs()));
let relative_max_err = if value_range > 0.0 {
max_err / value_range
} else {
0.0
};
assert!(
relative_max_err < 0.01,
"reconstruction error {:.6} ({:.4}%) exceeds 1% of value range {:.4}",
max_err,
relative_max_err * 100.0,
value_range
);
println!(
"delta_chain_stress: max reconstruction error = {:.6} ({:.4}% of range {:.4})",
max_err,
relative_max_err * 100.0,
value_range
);
// Compact: apply all deltas to base, chain_len should become 0.
chain.compact();
assert_eq!(
chain.chain_len(),
0,
"chain_len should be 0 after compaction"
);
// Verify reconstruction after compaction still yields correct data.
let after_compact = chain.reconstruct();
for i in 0..DIM {
let err = (after_compact[i] - truth[i]).abs();
assert!(
err < 0.01,
"post-compaction error at elem {}: {:.6}",
i,
err
);
}
// Rebuild chain to max depth.
let compacted_base = after_compact.clone();
let mut chain2 = DeltaChain::new(compacted_base.clone(), MAX_DEPTH);
let mut truth2 = compacted_base.clone();
for epoch in 0..MAX_DEPTH {
let mut modified = truth2.clone();
let num_changes = (DIM as f32 * CHANGE_FRACTION) as usize;
for _ in 0..num_changes {
let idx = rng.next_usize_range(0, DIM);
modified[idx] += rng.next_f32_range(-0.05, 0.05);
}
let delta = compute_delta(&truth2, &modified, 42, 0, epoch as u64, 1e-8, 1.0)
.expect("compute_delta should succeed");
chain2.append(delta).unwrap();
truth2 = modified;
}
assert_eq!(chain2.chain_len(), MAX_DEPTH as usize);
// One more append should fail with DeltaChainTooLong.
let mut overflow_modified = truth2.clone();
overflow_modified[0] += 0.01;
let overflow_delta = compute_delta(
&truth2,
&overflow_modified,
42,
0,
MAX_DEPTH as u64,
1e-8,
1.0,
)
.expect("compute_delta for overflow");
let result = chain2.append(overflow_delta);
assert_eq!(
result,
Err(StoreError::DeltaChainTooLong),
"appending beyond max_depth should return DeltaChainTooLong"
);
// Reconstruct should still work after the failed append.
let after_fail = chain2.reconstruct();
assert_eq!(after_fail.len(), DIM);
for i in 0..DIM {
let err = (after_fail[i] - truth2[i]).abs();
assert!(
err < 0.01,
"reconstruction after failed append: elem {} error {:.6}",
i,
err
);
}
println!("delta_chain_stress: all chain operations verified");
}
// ===========================================================================
// 5. Checksum sensitivity
// ===========================================================================
/// Verifies that the checksum stored in block metadata is deterministic
/// and sensitive to even tiny changes in input data.
#[test]
fn test_checksum_sensitivity() {
let mut store = TieredStore::new(4096);
let mut rng = SimpleRng::new(0xFEED_FACE);
const ELEM_COUNT: usize = 128;
let data: Vec<f32> = (0..ELEM_COUNT)
.map(|_| rng.next_f32_range(-1.0, 1.0))
.collect();
let key = make_key(5, 0);
// Put and record the checksum.
store.put(key, &data, Tier::Tier1, 0).unwrap();
let checksum1 = store.meta(key).unwrap().checksum;
// Put the same data again with the same key -> same checksum.
store.put(key, &data, Tier::Tier1, 1).unwrap();
let checksum2 = store.meta(key).unwrap().checksum;
assert_eq!(
checksum1, checksum2,
"identical data should produce identical checksums"
);
// Modify one element by a tiny amount (1e-6), put again.
let mut data_tiny = data.clone();
data_tiny[ELEM_COUNT / 2] += 1e-6;
store.put(key, &data_tiny, Tier::Tier1, 2).unwrap();
let checksum3 = store.meta(key).unwrap().checksum;
// Note: due to 8-bit quantization, a 1e-6 change on values in [-1,1]
// might not change the quantized representation. If it does, checksums
// differ; if not, they are the same. We test a larger perturbation below
// to guarantee a difference.
// Modify one element by a larger amount that will definitely change quantized value.
let mut data_modified = data.clone();
data_modified[ELEM_COUNT / 2] += 0.1;
store.put(key, &data_modified, Tier::Tier1, 3).unwrap();
let checksum4 = store.meta(key).unwrap().checksum;
assert_ne!(
checksum1, checksum4,
"modifying one element by 0.1 should change the checksum"
);
// Put very different data -> very different checksum.
let data_different: Vec<f32> = (0..ELEM_COUNT)
.map(|_| rng.next_f32_range(-10.0, 10.0))
.collect();
store.put(key, &data_different, Tier::Tier1, 4).unwrap();
let checksum5 = store.meta(key).unwrap().checksum;
assert_ne!(
checksum1, checksum5,
"very different data should produce a different checksum"
);
// Also verify it differs from the slightly-modified version.
assert_ne!(
checksum4, checksum5,
"two different datasets should have different checksums"
);
println!(
"checksum_sensitivity: c1={:#010X} c2={:#010X} c3={:#010X} c4={:#010X} c5={:#010X}",
checksum1, checksum2, checksum3, checksum4, checksum5
);
}
// ===========================================================================
// 6. Concurrent simulation (simulated multi-reader)
// ===========================================================================
/// Puts 100 blocks, then runs 10 simulated "reader threads" (sequential
/// loops) each performing 100 iterations of random touches and reads.
/// Verifies all reads succeed and return finite data, and metadata remains
/// consistent.
#[test]
fn test_concurrent_simulation() {
let mut store = TieredStore::new(4096);
let mut rng = SimpleRng::new(0xC0DE_C0DE);
const NUM_BLOCKS: usize = 100;
const NUM_READERS: usize = 10;
const ITERS_PER_READER: usize = 100;
const ELEM_COUNT: usize = 64;
// Insert all blocks.
for i in 0..NUM_BLOCKS {
let data = random_data(&mut rng, ELEM_COUNT);
store
.put(make_key(6, i as u32), &data, Tier::Tier1, 0)
.unwrap();
}
assert_eq!(store.block_count(), NUM_BLOCKS);
let mut total_reads: usize = 0;
let mut total_touches: usize = 0;
// Simulate NUM_READERS concurrent readers.
for reader_id in 0..NUM_READERS {
let base_tick = (reader_id as u64 + 1) * 1000;
for iter in 0..ITERS_PER_READER {
let key_idx = rng.next_usize_range(0, NUM_BLOCKS) as u32;
let key = make_key(6, key_idx);
let tick = base_tick + iter as u64;
// Touch the block.
store.touch(key, tick);
total_touches += 1;
// Read the block.
let mut out = vec![0.0f32; ELEM_COUNT];
let n = store.get(key, &mut out, tick).unwrap_or_else(|e| {
panic!(
"reader {} iter {} key {} failed: {:?}",
reader_id, iter, key_idx, e
)
});
assert_eq!(n, ELEM_COUNT);
total_reads += 1;
// Verify finite values.
for (j, &v) in out.iter().enumerate() {
assert!(
v.is_finite(),
"reader {} iter {} block {} elem {} non-finite: {}",
reader_id,
iter,
key_idx,
j,
v
);
}
}
}
// Verify metadata integrity for all blocks.
for i in 0..NUM_BLOCKS {
let key = make_key(6, i as u32);
let m = store.meta(key).expect("meta should exist");
assert!(
m.tier == Tier::Tier1 || m.tier == Tier::Tier2 || m.tier == Tier::Tier3,
"block {} has invalid tier {:?}",
i,
m.tier
);
assert!(
m.access_count > 0,
"block {} should have been accessed at least once",
i
);
}
println!(
"concurrent_simulation: {} readers x {} iters = {} reads, {} touches",
NUM_READERS, ITERS_PER_READER, total_reads, total_touches
);
}
// ===========================================================================
// 7. Extreme tick values
// ===========================================================================
/// Tests behavior at tick value boundaries: 0, u64::MAX-1, and u64::MAX.
/// Verifies no overflow or underflow panics in access-pattern tracking.
#[test]
fn test_extreme_tick_values() {
let mut store = TieredStore::new(4096);
const ELEM_COUNT: usize = 32;
let data = vec![0.5f32; ELEM_COUNT];
// -- Test 1: Put at tick=0, touch at tick=u64::MAX-1 --
let key_a = make_key(7, 0);
store.put(key_a, &data, Tier::Tier1, 0).unwrap();
store.touch(key_a, u64::MAX - 1);
let meta_a = store.meta(key_a).unwrap();
assert_eq!(meta_a.last_access_at, u64::MAX - 1);
assert!(
meta_a.access_count >= 2,
"access_count should reflect put + touch"
);
// Read should still work.
let mut out = vec![0.0f32; ELEM_COUNT];
let n = store.get(key_a, &mut out, u64::MAX - 1).unwrap();
assert_eq!(n, ELEM_COUNT);
// -- Test 2: Put at tick=u64::MAX --
let key_b = make_key(7, 1);
store.put(key_b, &data, Tier::Tier1, u64::MAX).unwrap();
let meta_b = store.meta(key_b).unwrap();
assert_eq!(meta_b.created_at, u64::MAX);
assert_eq!(meta_b.last_access_at, u64::MAX);
// Read at u64::MAX.
let mut out2 = vec![0.0f32; ELEM_COUNT];
let n2 = store.get(key_b, &mut out2, u64::MAX).unwrap();
assert_eq!(n2, ELEM_COUNT);
// -- Test 3: Touch at tick=0 when last_access=u64::MAX --
// This tests that saturating_sub prevents underflow.
store.touch(key_b, 0);
let meta_b2 = store.meta(key_b).unwrap();
// last_access should update to 0 (the tick we passed).
// The delta computation uses saturating_sub, so 0 - u64::MAX saturates to 0,
// meaning delta=0 and the window/ema are handled without panic.
assert_eq!(meta_b2.last_access_at, 0);
// -- Test 4: Touch at tick=u64::MAX after last_access=0 --
store.touch(key_b, u64::MAX);
let meta_b3 = store.meta(key_b).unwrap();
assert_eq!(meta_b3.last_access_at, u64::MAX);
// The delta is u64::MAX, which is >= 64, so window resets to 1.
assert_eq!(meta_b3.window, 1);
// Verify all blocks still readable after extreme tick gymnastics.
for i in 0..2u32 {
let key = make_key(7, i);
let mut out = vec![0.0f32; ELEM_COUNT];
let result = store.get(key, &mut out, u64::MAX);
assert!(
result.is_ok(),
"block {} should be readable after extreme ticks: {:?}",
i,
result
);
}
println!("extreme_tick_values: all boundary conditions passed without panic");
}
// ===========================================================================
// 8. All tiers coexist
// ===========================================================================
/// Puts 100 blocks in each of Tier1, Tier2, Tier3 (300 total), verifies
/// tier counts, reads all blocks verifying accuracy matches tier expectations
/// (higher tiers = less quantization error), evicts all Tier3 blocks, and
/// verifies Tier1 and Tier2 are still readable.
#[test]
fn test_all_tiers_coexist() {
let mut store = TieredStore::new(4096);
let mut rng = SimpleRng::new(0xBAAD_F00D);
const BLOCKS_PER_TIER: usize = 100;
const ELEM_COUNT: usize = 128;
// Store original data for roundtrip error comparison.
let mut originals: Vec<Vec<f32>> = Vec::new();
// Insert 100 blocks at Tier1 (tensor_id=81).
for i in 0..BLOCKS_PER_TIER {
let data = random_data(&mut rng, ELEM_COUNT);
store
.put(make_key(81, i as u32), &data, Tier::Tier1, 0)
.unwrap();
originals.push(data);
}
// Insert 100 blocks at Tier2 (tensor_id=82).
for i in 0..BLOCKS_PER_TIER {
let data = random_data(&mut rng, ELEM_COUNT);
store
.put(make_key(82, i as u32), &data, Tier::Tier2, 0)
.unwrap();
originals.push(data);
}
// Insert 100 blocks at Tier3 (tensor_id=83).
for i in 0..BLOCKS_PER_TIER {
let data = random_data(&mut rng, ELEM_COUNT);
store
.put(make_key(83, i as u32), &data, Tier::Tier3, 0)
.unwrap();
originals.push(data);
}
// Verify tier counts.
assert_eq!(store.tier_count(Tier::Tier1), BLOCKS_PER_TIER);
assert_eq!(store.tier_count(Tier::Tier2), BLOCKS_PER_TIER);
assert_eq!(store.tier_count(Tier::Tier3), BLOCKS_PER_TIER);
assert_eq!(store.block_count(), 3 * BLOCKS_PER_TIER);
// Read all blocks and compute per-tier max roundtrip error.
let mut tier1_max_err: f32 = 0.0;
let mut tier2_max_err: f32 = 0.0;
let mut tier3_max_err: f32 = 0.0;
for i in 0..BLOCKS_PER_TIER {
// Tier1
let key = make_key(81, i as u32);
let mut out = vec![0.0f32; ELEM_COUNT];
store.get(key, &mut out, 1).unwrap();
let orig = &originals[i];
for j in 0..ELEM_COUNT {
let err = (out[j] - orig[j]).abs();
if err > tier1_max_err {
tier1_max_err = err;
}
}
// Tier2
let key = make_key(82, i as u32);
store.get(key, &mut out, 1).unwrap();
let orig = &originals[BLOCKS_PER_TIER + i];
for j in 0..ELEM_COUNT {
let err = (out[j] - orig[j]).abs();
if err > tier2_max_err {
tier2_max_err = err;
}
}
// Tier3
let key = make_key(83, i as u32);
store.get(key, &mut out, 1).unwrap();
let orig = &originals[2 * BLOCKS_PER_TIER + i];
for j in 0..ELEM_COUNT {
let err = (out[j] - orig[j]).abs();
if err > tier3_max_err {
tier3_max_err = err;
}
}
}
// Tier1 (8-bit) should have the lowest error, Tier3 (3-bit) the highest.
// Values are in [-1, 1], so 8-bit qmax=127 -> step ~0.0079, 3-bit qmax=3 -> step ~0.33.
assert!(
tier1_max_err <= tier3_max_err,
"Tier1 error ({:.6}) should not exceed Tier3 error ({:.6})",
tier1_max_err,
tier3_max_err
);
// Tier3 with 3-bit quantization has significant error for [-1,1] data.
assert!(
tier3_max_err > 0.0,
"Tier3 (3-bit) should have nonzero quantization error"
);
println!(
"all_tiers_coexist: tier1_err={:.6}, tier2_err={:.6}, tier3_err={:.6}",
tier1_max_err, tier2_max_err, tier3_max_err
);
// Evict all Tier3 blocks.
for i in 0..BLOCKS_PER_TIER {
store
.evict(make_key(83, i as u32), ReconstructPolicy::None)
.unwrap();
}
assert_eq!(store.tier_count(Tier::Tier3), 0);
assert_eq!(store.tier_count(Tier::Tier0), BLOCKS_PER_TIER);
// Total blocks unchanged (eviction preserves metadata).
assert_eq!(store.block_count(), 3 * BLOCKS_PER_TIER);
// Tier1 and Tier2 must still be readable.
for i in 0..BLOCKS_PER_TIER {
let mut out = vec![0.0f32; ELEM_COUNT];
let key1 = make_key(81, i as u32);
store.get(key1, &mut out, 2).unwrap_or_else(|e| {
panic!("Tier1 block {} unreadable after Tier3 eviction: {:?}", i, e)
});
let key2 = make_key(82, i as u32);
store.get(key2, &mut out, 2).unwrap_or_else(|e| {
panic!("Tier2 block {} unreadable after Tier3 eviction: {:?}", i, e)
});
}
// Evicted Tier3 blocks should return TensorEvicted.
for i in 0..BLOCKS_PER_TIER {
let key = make_key(83, i as u32);
let mut out = vec![0.0f32; ELEM_COUNT];
let result = store.get(key, &mut out, 2);
assert_eq!(
result,
Err(StoreError::TensorEvicted),
"evicted Tier3 block {} should return TensorEvicted",
i
);
}
println!(
"all_tiers_coexist: evicted Tier3, Tier1 ({}) and Tier2 ({}) still intact",
store.tier_count(Tier::Tier1),
store.tier_count(Tier::Tier2)
);
}

View File

@@ -0,0 +1,353 @@
//! FFI interface tests for the temporal tensor store.
//!
//! These tests exercise the `tts_*` extern "C" functions exposed by
//! `store_ffi.rs` through their public API. Because the FFI layer uses
//! a single global `STORE_STATE`, tests **must** run sequentially:
//!
//! ```bash
//! cargo test -p ruvector-temporal-tensor --test wasm_ffi_test --features ffi -- --test-threads=1
//! ```
#![cfg(feature = "ffi")]
use ruvector_temporal_tensor::store_ffi::{
tts_block_count, tts_evict, tts_get, tts_init, tts_put, tts_stats, tts_tier_count, tts_touch,
};
// ── Constants mirrored from store_ffi.rs ────────────────────────────────
const ERR_BLOCK_NOT_FOUND: i32 = -4;
const ERR_BUFFER_TOO_SMALL: i32 = -5;
/// Binary stats size: 5 * u32 + 2 * u64 = 36 bytes.
const STATS_SIZE: usize = 5 * 4 + 2 * 8;
// ── Helpers ─────────────────────────────────────────────────────────────
/// Re-initialize the global store with default config before each test.
/// This replaces whatever state was left by a previous test.
fn reset() {
let rc = tts_init(std::ptr::null(), 0);
assert_eq!(rc, 0, "tts_init with default config must succeed");
}
/// Read a little-endian u32 from `buf` at the given byte offset.
fn read_u32_le(buf: &[u8], off: usize) -> u32 {
u32::from_le_bytes([buf[off], buf[off + 1], buf[off + 2], buf[off + 3]])
}
/// Read a little-endian u64 from `buf` at the given byte offset.
fn read_u64_le(buf: &[u8], off: usize) -> u64 {
let mut arr = [0u8; 8];
arr.copy_from_slice(&buf[off..off + 8]);
u64::from_le_bytes(arr)
}
// ── Tests ───────────────────────────────────────────────────────────────
#[test]
fn test_ffi_init_and_destroy() {
// Calling tts_init with a null pointer and zero length should use
// the default TierConfig and return success (0).
let rc = tts_init(std::ptr::null(), 0);
assert_eq!(rc, 0, "tts_init should return 0 on success");
// The freshly initialized store must contain zero blocks.
assert_eq!(tts_block_count(), 0, "new store should have 0 blocks");
// Re-initializing must also succeed (replaces old state).
let rc2 = tts_init(std::ptr::null(), 0);
assert_eq!(rc2, 0, "re-init should succeed");
assert_eq!(tts_block_count(), 0, "re-init should reset block count");
}
#[test]
fn test_ffi_put_get_roundtrip() {
reset();
// Create 64 f32 values with a clear pattern.
let data: Vec<f32> = (0..64).map(|i| (i as f32 - 32.0) * 0.1).collect();
let rc = tts_put(0, 1, 0, data.as_ptr(), data.len());
assert_eq!(rc, 0, "tts_put should return 0 on success");
let mut out = vec![0.0f32; 64];
let n = tts_get(0, 1, 0, out.as_mut_ptr(), out.len());
assert_eq!(n, 64, "tts_get should return 64 elements");
// Verify accuracy. New blocks default to Hot (8-bit quantization)
// so the error should be small.
let max_abs = data.iter().map(|v| v.abs()).fold(0.0f32, f32::max);
for (i, (&orig, &dec)) in data.iter().zip(out.iter()).enumerate() {
let err = (orig - dec).abs();
assert!(
err < max_abs * 0.05,
"element {i}: orig={orig}, decoded={dec}, err={err}, tolerance={}",
max_abs * 0.05,
);
}
}
#[test]
fn test_ffi_multi_tensor() {
reset();
let data_a: Vec<f32> = (0..64).map(|i| i as f32 * 0.5).collect();
let data_b: Vec<f32> = (0..64).map(|i| -(i as f32) * 0.3).collect();
let data_c: Vec<f32> = (0..64).map(|i| (i as f32).sin()).collect();
// Three different tensor IDs using hi/lo split for u128:
// tensor A: hi=0, lo=1 -> tensor_id = 1
// tensor B: hi=0, lo=2 -> tensor_id = 2
// tensor C: hi=1, lo=0 -> tensor_id = 1 << 64
assert_eq!(tts_put(0, 1, 0, data_a.as_ptr(), data_a.len()), 0);
assert_eq!(tts_put(0, 2, 0, data_b.as_ptr(), data_b.len()), 0);
assert_eq!(tts_put(1, 0, 0, data_c.as_ptr(), data_c.len()), 0);
assert_eq!(tts_block_count(), 3, "should have 3 blocks total");
// Read back each tensor independently.
let mut out = vec![0.0f32; 64];
let n_a = tts_get(0, 1, 0, out.as_mut_ptr(), out.len());
assert_eq!(n_a, 64);
// Spot-check first element of tensor A.
assert!(
(out[0] - data_a[0]).abs() < 0.5,
"tensor A readback mismatch"
);
let n_b = tts_get(0, 2, 0, out.as_mut_ptr(), out.len());
assert_eq!(n_b, 64);
assert!(
(out[0] - data_b[0]).abs() < 0.5,
"tensor B readback mismatch"
);
let n_c = tts_get(1, 0, 0, out.as_mut_ptr(), out.len());
assert_eq!(n_c, 64);
assert!(
(out[0] - data_c[0]).abs() < 0.5,
"tensor C readback mismatch"
);
}
#[test]
fn test_ffi_eviction() {
reset();
let data = vec![1.0f32; 64];
assert_eq!(tts_put(0, 42, 0, data.as_ptr(), data.len()), 0);
assert_eq!(tts_block_count(), 1);
// Evict the block.
let rc = tts_evict(0, 42, 0);
assert_eq!(rc, 0, "tts_evict should return 0 on success");
assert_eq!(tts_block_count(), 0, "evicted block should be gone");
// A subsequent get should return ERR_BLOCK_NOT_FOUND.
let mut out = vec![0.0f32; 64];
let rc_get = tts_get(0, 42, 0, out.as_mut_ptr(), out.len());
assert_eq!(
rc_get, ERR_BLOCK_NOT_FOUND,
"get after evict should return block-not-found"
);
// Evicting again should also return block-not-found.
let rc2 = tts_evict(0, 42, 0);
assert_eq!(rc2, ERR_BLOCK_NOT_FOUND);
}
#[test]
fn test_ffi_touch_updates_access() {
reset();
let data = vec![1.0f32; 64];
assert_eq!(tts_put(0, 7, 3, data.as_ptr(), data.len()), 0);
assert_eq!(tts_block_count(), 1);
// Touch the block multiple times.
for _ in 0..5 {
let rc = tts_touch(0, 7, 3);
assert_eq!(rc, 0, "tts_touch should return 0 on success");
}
// Block count should remain unchanged (touch does not add/remove blocks).
assert_eq!(tts_block_count(), 1, "touch should not change block count");
// The block should still be readable.
let mut out = vec![0.0f32; 64];
let n = tts_get(0, 7, 3, out.as_mut_ptr(), out.len());
assert_eq!(n, 64, "block should still be readable after touches");
// Touching a non-existent block should fail.
let rc_missing = tts_touch(0, 99, 0);
assert_eq!(rc_missing, ERR_BLOCK_NOT_FOUND);
}
#[test]
fn test_ffi_tier_counts() {
reset();
// All new blocks are placed in Hot (tier 0) by default.
let data = vec![1.0f32; 64];
assert_eq!(tts_put(0, 1, 0, data.as_ptr(), data.len()), 0);
assert_eq!(tts_put(0, 1, 1, data.as_ptr(), data.len()), 0);
assert_eq!(tts_put(0, 2, 0, data.as_ptr(), data.len()), 0);
assert_eq!(tts_block_count(), 3);
assert_eq!(tts_tier_count(0), 3, "all blocks should be Hot");
assert_eq!(tts_tier_count(1), 0, "no Warm blocks");
assert_eq!(tts_tier_count(2), 0, "no Cool blocks");
assert_eq!(tts_tier_count(3), 0, "no Cold blocks");
// Invalid tier should return an error.
assert!(tts_tier_count(99) < 0, "invalid tier should return error");
}
#[test]
fn test_ffi_stats_output() {
reset();
let data = vec![1.0f32; 64];
assert_eq!(tts_put(0, 1, 0, data.as_ptr(), data.len()), 0);
assert_eq!(tts_put(0, 1, 1, data.as_ptr(), data.len()), 0);
assert_eq!(tts_put(0, 2, 0, data.as_ptr(), data.len()), 0);
let mut buf = vec![0u8; STATS_SIZE];
let written = tts_stats(buf.as_mut_ptr(), buf.len());
assert_eq!(
written, STATS_SIZE as i32,
"tts_stats should write exactly {STATS_SIZE} bytes"
);
// Parse the binary stats layout:
// [block_count:u32][hot:u32][warm:u32][cool:u32][cold:u32]
// [total_bytes:u64][tick_count:u64]
let block_count = read_u32_le(&buf, 0);
let hot = read_u32_le(&buf, 4);
let warm = read_u32_le(&buf, 8);
let cool = read_u32_le(&buf, 12);
let cold = read_u32_le(&buf, 16);
let total_bytes = read_u64_le(&buf, 20);
let _tick_count = read_u64_le(&buf, 28);
assert_eq!(block_count, 3, "block_count mismatch");
assert_eq!(hot, 3, "hot count mismatch");
assert_eq!(warm, 0, "warm count mismatch");
assert_eq!(cool, 0, "cool count mismatch");
assert_eq!(cold, 0, "cold count mismatch");
assert!(total_bytes > 0, "total_bytes should be > 0 after puts");
// Verify stats rejects a too-small buffer.
let mut small_buf = vec![0u8; 4];
let rc = tts_stats(small_buf.as_mut_ptr(), small_buf.len());
assert_eq!(rc, ERR_BUFFER_TOO_SMALL);
}
#[test]
fn test_ffi_put_multiple_blocks_same_tensor() {
reset();
let data = vec![2.5f32; 64];
// Put 5 blocks for the same tensor (different block indices).
for idx in 0..5u32 {
let rc = tts_put(0, 10, idx, data.as_ptr(), data.len());
assert_eq!(rc, 0, "put block_index={idx} should succeed");
}
assert_eq!(tts_block_count(), 5);
// Each block should be independently readable.
let mut out = vec![0.0f32; 64];
for idx in 0..5u32 {
let n = tts_get(0, 10, idx, out.as_mut_ptr(), out.len());
assert_eq!(n, 64, "block_index={idx} should return 64 elements");
}
}
#[test]
fn test_ffi_overwrite_block() {
reset();
let data1 = vec![1.0f32; 64];
assert_eq!(tts_put(0, 5, 0, data1.as_ptr(), data1.len()), 0);
let data2 = vec![9.0f32; 64];
assert_eq!(tts_put(0, 5, 0, data2.as_ptr(), data2.len()), 0);
// Block count should still be 1 (overwrite, not insert).
assert_eq!(tts_block_count(), 1);
// Should read back the second write.
let mut out = vec![0.0f32; 64];
let n = tts_get(0, 5, 0, out.as_mut_ptr(), out.len());
assert_eq!(n, 64);
for &v in &out {
assert!(
(v - 9.0).abs() < 0.5,
"expected ~9.0 after overwrite, got {v}"
);
}
}
#[test]
fn test_ffi_get_buffer_too_small() {
reset();
let data = vec![1.0f32; 64];
assert_eq!(tts_put(0, 1, 0, data.as_ptr(), data.len()), 0);
let mut small_out = vec![0.0f32; 2];
let rc = tts_get(0, 1, 0, small_out.as_mut_ptr(), small_out.len());
assert_eq!(
rc, ERR_BUFFER_TOO_SMALL,
"get with undersized buffer should return buffer-too-small"
);
}
#[test]
fn test_ffi_evict_then_reinsert() {
reset();
let data = vec![3.0f32; 64];
assert_eq!(tts_put(0, 1, 0, data.as_ptr(), data.len()), 0);
assert_eq!(tts_block_count(), 1);
// Evict.
assert_eq!(tts_evict(0, 1, 0), 0);
assert_eq!(tts_block_count(), 0);
// Re-insert at the same key.
let data2 = vec![7.0f32; 64];
assert_eq!(tts_put(0, 1, 0, data2.as_ptr(), data2.len()), 0);
assert_eq!(tts_block_count(), 1);
// Should read back the new data.
let mut out = vec![0.0f32; 64];
let n = tts_get(0, 1, 0, out.as_mut_ptr(), out.len());
assert_eq!(n, 64);
for &v in &out {
assert!(
(v - 7.0).abs() < 0.5,
"expected ~7.0 after re-insert, got {v}"
);
}
}
#[test]
fn test_ffi_large_tensor_id() {
reset();
// Use the full u128 range: hi=u64::MAX, lo=u64::MAX -> tensor_id = u128::MAX.
let data = vec![0.5f32; 64];
assert_eq!(
tts_put(u64::MAX, u64::MAX, 0, data.as_ptr(), data.len()),
0,
"put with max tensor_id should succeed"
);
let mut out = vec![0.0f32; 64];
let n = tts_get(u64::MAX, u64::MAX, 0, out.as_mut_ptr(), out.len());
assert_eq!(n, 64, "get with max tensor_id should succeed");
}