Skip to main content

somatize_core/
codec.rs

1//! `SOMA1` framed binary codec for [`Value`].
2//!
3//! The persistent cache stored values as JSON in Phase 1 — an f64
4//! tensor round-tripped through decimal text at ~3× the size. This
5//! codec writes tensors as raw little-endian f64, keeping payloads at
6//! ~1× raw size, and hashes the encoded bytes with BLAKE3 in the same
7//! pass.
8//!
9//! Frame layout:
10//!
11//! ```text
12//! [0..6)  magic  b"SOMA1\0"
13//! [6]     variant tag: 0=Empty 1=Tensor 2=Json 3=Bytes 4=Object
14//! [7]     compression: 0=raw (other values reserved)
15//! [8..]   payload
16//! ```
17//!
18//! Tensor payload: `u32 ndim` + `ndim × u64` dims + `u64 count` +
19//! `count × f64` little-endian values. Json payload: `serde_json`
20//! bytes (deterministic — the default map is sorted). Bytes/Object:
21//! raw bytes.
22
23use crate::action::ContentHash;
24use crate::error::{Result, SomaError};
25use crate::value::Value;
26use std::sync::Arc;
27
28/// Frame magic: the first six bytes of every `SOMA1`-encoded value.
29/// Public so stores can sniff a blob's format without decoding it.
30pub const MAGIC: &[u8; 6] = b"SOMA1\0";
31
32const TAG_EMPTY: u8 = 0;
33const TAG_TENSOR: u8 = 1;
34const TAG_JSON: u8 = 2;
35const TAG_BYTES: u8 = 3;
36const TAG_OBJECT: u8 = 4;
37const TAG_TEXT: u8 = 5;
38
39const COMPRESSION_RAW: u8 = 0;
40
41/// Encode a value into a `SOMA1` frame.
42pub fn encode_value(value: &Value) -> Result<Vec<u8>> {
43    let mut buf = Vec::with_capacity(64);
44    buf.extend_from_slice(MAGIC);
45    match value {
46        Value::Empty => {
47            buf.push(TAG_EMPTY);
48            buf.push(COMPRESSION_RAW);
49        }
50        Value::Tensor { values, shape } => {
51            buf.push(TAG_TENSOR);
52            buf.push(COMPRESSION_RAW);
53            buf.reserve(4 + shape.len() * 8 + 8 + values.len() * 8);
54            buf.extend_from_slice(&(shape.len() as u32).to_le_bytes());
55            for dim in shape {
56                buf.extend_from_slice(&(*dim as u64).to_le_bytes());
57            }
58            buf.extend_from_slice(&(values.len() as u64).to_le_bytes());
59            for v in values.iter() {
60                buf.extend_from_slice(&v.to_le_bytes());
61            }
62        }
63        Value::Text(s) => {
64            buf.push(TAG_TEXT);
65            buf.push(COMPRESSION_RAW);
66            buf.extend_from_slice(s.as_bytes());
67        }
68        Value::Json(v) => {
69            buf.push(TAG_JSON);
70            buf.push(COMPRESSION_RAW);
71            let bytes = serde_json::to_vec(v.as_ref())
72                .map_err(|e| SomaError::Cache(format!("codec: json encode: {e}")))?;
73            buf.extend_from_slice(&bytes);
74        }
75        Value::Bytes(b) => {
76            buf.push(TAG_BYTES);
77            buf.push(COMPRESSION_RAW);
78            buf.extend_from_slice(b);
79        }
80        Value::Object(b) => {
81            buf.push(TAG_OBJECT);
82            buf.push(COMPRESSION_RAW);
83            buf.extend_from_slice(b);
84        }
85    }
86    Ok(buf)
87}
88
89/// Encode and content-hash in one step.
90pub fn encode_and_hash(value: &Value) -> Result<(Vec<u8>, ContentHash)> {
91    let bytes = encode_value(value)?;
92    let hash = ContentHash::blake3(&bytes);
93    Ok((bytes, hash))
94}
95
96/// Decode a `SOMA1` frame back into a value.
97pub fn decode_value(bytes: &[u8]) -> Result<Value> {
98    if bytes.len() < 8 || &bytes[..6] != MAGIC {
99        return Err(SomaError::Cache("codec: not a SOMA1 frame".into()));
100    }
101    let tag = bytes[6];
102    if bytes[7] != COMPRESSION_RAW {
103        return Err(SomaError::Cache(format!(
104            "codec: unknown compression byte {}",
105            bytes[7]
106        )));
107    }
108    let payload = &bytes[8..];
109    match tag {
110        TAG_EMPTY => Ok(Value::Empty),
111        TAG_TENSOR => decode_tensor(payload),
112        TAG_JSON => {
113            let v: serde_json::Value = serde_json::from_slice(payload)
114                .map_err(|e| SomaError::Cache(format!("codec: json decode: {e}")))?;
115            Ok(Value::json(v))
116        }
117        TAG_TEXT => {
118            let s = std::str::from_utf8(payload)
119                .map_err(|e| SomaError::Cache(format!("codec: text decode: {e}")))?;
120            Ok(Value::text(s))
121        }
122        TAG_BYTES => Ok(Value::bytes(payload.to_vec())),
123        TAG_OBJECT => Ok(Value::object(payload.to_vec())),
124        other => Err(SomaError::Cache(format!("codec: unknown tag {other}"))),
125    }
126}
127
128fn decode_tensor(payload: &[u8]) -> Result<Value> {
129    let err = || SomaError::Cache("codec: truncated tensor frame".into());
130    let mut at = 0usize;
131    let take = |at: &mut usize, n: usize| -> Result<&[u8]> {
132        let slice = payload.get(*at..*at + n).ok_or_else(err)?;
133        *at += n;
134        Ok(slice)
135    };
136
137    let ndim = u32::from_le_bytes(take(&mut at, 4)?.try_into().unwrap()) as usize;
138    if ndim > 64 {
139        return Err(SomaError::Cache(format!("codec: implausible ndim {ndim}")));
140    }
141    let mut shape = Vec::with_capacity(ndim);
142    for _ in 0..ndim {
143        shape.push(u64::from_le_bytes(take(&mut at, 8)?.try_into().unwrap()) as usize);
144    }
145    let count = u64::from_le_bytes(take(&mut at, 8)?.try_into().unwrap()) as usize;
146    let data = take(&mut at, count.checked_mul(8).ok_or_else(err)?)?;
147    let mut values = Vec::with_capacity(count);
148    for chunk in data.chunks_exact(8) {
149        values.push(f64::from_le_bytes(chunk.try_into().unwrap()));
150    }
151    Ok(Value::Tensor {
152        values: Arc::new(values),
153        shape,
154    })
155}
156
157#[cfg(test)]
158mod tests {
159    use super::*;
160    use serde_json::json;
161
162    #[test]
163    fn roundtrip_all_variants() {
164        let values = vec![
165            Value::Empty,
166            Value::tensor(vec![1.0, -2.5, f64::MAX, 0.0], vec![2, 2]),
167            Value::tensor(vec![], vec![0]),
168            Value::json(json!({"a": [1, 2.5, "x"], "b": null})),
169            Value::bytes(vec![0, 255, 128]),
170            Value::object(vec![0x80, 0x04]),
171            Value::text(""),
172            Value::text("Summarize the following in one sentence."),
173            // Non-ASCII must survive the byte-level frame intact.
174            Value::text("resumen: ¿qué pasó? — 数字 🧬"),
175        ];
176        for v in values {
177            let (bytes, hash) = encode_and_hash(&v).unwrap();
178            assert!(hash.verify(&bytes));
179            assert_eq!(decode_value(&bytes).unwrap(), v, "roundtrip failed for {v}");
180        }
181    }
182
183    /// Text and a JSON string holding the same characters are different
184    /// values, and must hash differently — otherwise a prompt and a JSON
185    /// document quoting it would share a cache line.
186    #[test]
187    fn text_and_json_string_are_distinct() {
188        let (text_bytes, text_hash) = encode_and_hash(&Value::text("hello")).unwrap();
189        let (json_bytes, json_hash) = encode_and_hash(&Value::json(json!("hello"))).unwrap();
190
191        assert_ne!(text_bytes, json_bytes);
192        assert_ne!(text_hash, json_hash);
193        assert_eq!(decode_value(&text_bytes).unwrap(), Value::text("hello"));
194        assert_eq!(
195            decode_value(&json_bytes).unwrap(),
196            Value::json(json!("hello"))
197        );
198    }
199
200    /// Invalid UTF-8 in a text frame is reported, not silently replaced.
201    #[test]
202    fn text_frame_rejects_invalid_utf8() {
203        let mut frame = Vec::from(MAGIC);
204        frame.push(TAG_TEXT);
205        frame.push(COMPRESSION_RAW);
206        frame.extend_from_slice(&[0xff, 0xfe]);
207        assert!(decode_value(&frame).is_err());
208    }
209
210    #[test]
211    fn tensor_size_is_near_raw() {
212        let n = 10_000;
213        // Full-precision mantissas — the realistic case for model
214        // weights/features (short decimals like 0.7 flatter JSON).
215        let v = Value::tensor((0..n).map(|i| (i as f64).sin()).collect(), vec![n]);
216        let encoded = encode_value(&v).unwrap();
217        let raw = n * 8;
218        assert!(
219            encoded.len() <= raw + raw / 10 + 64,
220            "encoded {} bytes vs raw {} — must stay within ~1.1×",
221            encoded.len(),
222            raw
223        );
224        // And meaningfully smaller than the JSON text form.
225        let json_len = serde_json::to_vec(&v).unwrap().len();
226        assert!(
227            encoded.len() * 2 < json_len,
228            "binary ({}) should be well under half of JSON ({})",
229            encoded.len(),
230            json_len
231        );
232    }
233
234    #[test]
235    fn identical_values_hash_identically() {
236        let a = Value::tensor(vec![1.0, 2.0], vec![2]);
237        let b = Value::tensor(vec![1.0, 2.0], vec![2]);
238        assert_eq!(
239            encode_and_hash(&a).unwrap().1,
240            encode_and_hash(&b).unwrap().1
241        );
242        let c = Value::tensor(vec![1.0, 2.0], vec![2, 1]);
243        assert_ne!(
244            encode_and_hash(&a).unwrap().1,
245            encode_and_hash(&c).unwrap().1
246        );
247    }
248
249    #[test]
250    fn garbage_rejected() {
251        assert!(decode_value(b"").is_err());
252        assert!(decode_value(b"NOTSOMA1xxxx").is_err());
253        let mut frame = encode_value(&Value::Empty).unwrap();
254        frame[7] = 9; // unknown compression
255        assert!(decode_value(&frame).is_err());
256    }
257}