Skip to main content

prolly/prolly/
format.rs

1//! Persisted tree-shape and node-layout descriptors.
2
3use serde::{Deserialize, Serialize};
4
5use super::cid::Cid;
6use super::encoding::{
7    Encoding, DEFAULT_CHUNKING_FACTOR, DEFAULT_HASH_SEED, DEFAULT_MAX_CHUNK_SIZE,
8    DEFAULT_MIN_CHUNK_SIZE,
9};
10use super::error::Error;
11
12const FORMAT_MAGIC: &[u8; 4] = b"CRFT";
13const FORMAT_VERSION: u8 = 1;
14
15/// Quantity used to enforce a chunk's soft size bounds.
16#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
17pub enum ChunkMeasure {
18    /// Count ordered entries.
19    EntryCount,
20    /// Count uncompressed key and value bytes.
21    LogicalBytes,
22    /// Count bytes produced by the selected node layout.
23    EncodedBytes,
24}
25
26/// Entry bytes supplied to the boundary hash.
27#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
28pub enum BoundaryInput {
29    /// Hash only the ordered key or internal separator.
30    Key,
31    /// Hash both key and value bytes.
32    KeyValue,
33}
34
35/// Stable hash algorithms available to persisted chunking policies.
36#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
37pub enum HashAlgorithm {
38    /// xxHash64 with a deterministic seed.
39    XxHash64,
40}
41
42/// Rule used to select a content-defined boundary.
43#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
44pub enum BoundaryRule {
45    /// Split when the hash falls below the factor-derived threshold.
46    HashThreshold { factor: u32 },
47    /// Split according to a bounded Weibull distribution.
48    Weibull { shape: u32 },
49    /// Split from a rolling BuzHash window.
50    RollingBuzHash { window: u16 },
51}
52
53/// Persisted content-defined chunking configuration.
54#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
55pub struct ChunkingSpec {
56    pub measure: ChunkMeasure,
57    pub input: BoundaryInput,
58    pub hash: HashAlgorithm,
59    pub rule: BoundaryRule,
60    pub min: u64,
61    pub target: u64,
62    pub max: u64,
63    pub hash_seed: u64,
64    pub level_salt: bool,
65    pub hard_max_node_bytes: u64,
66}
67
68impl Default for ChunkingSpec {
69    fn default() -> Self {
70        Self {
71            measure: ChunkMeasure::EntryCount,
72            input: BoundaryInput::Key,
73            hash: HashAlgorithm::XxHash64,
74            rule: BoundaryRule::HashThreshold {
75                factor: DEFAULT_CHUNKING_FACTOR,
76            },
77            min: DEFAULT_MIN_CHUNK_SIZE as u64,
78            target: DEFAULT_CHUNKING_FACTOR as u64,
79            max: DEFAULT_MAX_CHUNK_SIZE as u64,
80            hash_seed: DEFAULT_HASH_SEED,
81            level_salt: true,
82            hard_max_node_bytes: 16 * 1024 * 1024,
83        }
84    }
85}
86
87impl ChunkingSpec {
88    /// Validate bounds and rule parameters before the policy is persisted or run.
89    pub fn validate(&self) -> Result<(), Error> {
90        if self.min == 0 || self.min > self.target || self.target > self.max {
91            return Err(Error::InvalidFormat(
92                "chunk bounds must satisfy 0 < min <= target <= max".to_string(),
93            ));
94        }
95        if self.hard_max_node_bytes == 0 {
96            return Err(Error::InvalidFormat(
97                "hard maximum node bytes must be nonzero".to_string(),
98            ));
99        }
100        match self.rule {
101            BoundaryRule::HashThreshold { factor: 0 } => Err(Error::InvalidFormat(
102                "hash threshold factor must be nonzero".to_string(),
103            )),
104            BoundaryRule::Weibull { shape: 0 } => Err(Error::InvalidFormat(
105                "Weibull shape must be nonzero".to_string(),
106            )),
107            BoundaryRule::RollingBuzHash { window: 0 } => Err(Error::InvalidFormat(
108                "rolling hash window must be nonzero".to_string(),
109            )),
110            _ => Ok(()),
111        }
112    }
113}
114
115/// Physical encoding used for nodes in a tree.
116#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
117pub enum NodeLayoutSpec {
118    /// Prefix-compress adjacent keys.
119    #[default]
120    PrefixCompressed,
121    /// Encode complete keys and values directly.
122    Plain,
123    /// Store offsets into a validated shared payload.
124    OffsetTable,
125    /// Application-provided layout resolved by a stable identifier.
126    Custom { id: String, parameters: Vec<u8> },
127}
128
129/// Persisted settings that determine tree shape and content IDs.
130#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
131pub struct TreeFormat {
132    pub chunking: ChunkingSpec,
133    pub node_layout: NodeLayoutSpec,
134    pub value_encoding: Encoding,
135}
136
137impl Default for TreeFormat {
138    fn default() -> Self {
139        Self {
140            chunking: ChunkingSpec::default(),
141            node_layout: NodeLayoutSpec::default(),
142            value_encoding: Encoding::Raw,
143        }
144    }
145}
146
147impl TreeFormat {
148    /// Validate persisted identifiers and structural bounds known at this layer.
149    pub fn validate(&self) -> Result<(), Error> {
150        self.chunking.validate()?;
151        if let NodeLayoutSpec::Custom { id, .. } = &self.node_layout {
152            if id.is_empty() {
153                return Err(Error::InvalidFormat(
154                    "custom node layout identifier cannot be empty".to_string(),
155                ));
156            }
157        }
158        Ok(())
159    }
160
161    /// Encode the descriptor without maps or platform-sized integers.
162    pub fn canonical_bytes(&self) -> Result<Vec<u8>, Error> {
163        self.validate()?;
164        let mut out = Vec::new();
165        out.extend_from_slice(FORMAT_MAGIC);
166        out.push(FORMAT_VERSION);
167        encode_chunking(&self.chunking, &mut out);
168        encode_layout(&self.node_layout, &mut out);
169        encode_value_encoding(&self.value_encoding, &mut out);
170        Ok(out)
171    }
172
173    /// Hash the canonical persisted descriptor.
174    pub fn digest(&self) -> Result<Cid, Error> {
175        Ok(Cid::from_bytes(&self.canonical_bytes()?))
176    }
177
178    /// Decode a canonical persisted descriptor.
179    pub fn from_canonical_bytes(bytes: &[u8]) -> Result<Self, Error> {
180        let mut cursor = FormatCursor::new(bytes);
181        cursor.expect(FORMAT_MAGIC)?;
182        if cursor.read_u8()? != FORMAT_VERSION {
183            return Err(Error::InvalidFormat(
184                "unsupported tree format version".to_string(),
185            ));
186        }
187
188        let measure = match cursor.read_u8()? {
189            0 => ChunkMeasure::EntryCount,
190            1 => ChunkMeasure::LogicalBytes,
191            2 => ChunkMeasure::EncodedBytes,
192            _ => return Err(Error::InvalidFormat("invalid chunk measure".to_string())),
193        };
194        let input = match cursor.read_u8()? {
195            0 => BoundaryInput::Key,
196            1 => BoundaryInput::KeyValue,
197            _ => return Err(Error::InvalidFormat("invalid boundary input".to_string())),
198        };
199        let hash = match cursor.read_u8()? {
200            0 => HashAlgorithm::XxHash64,
201            _ => return Err(Error::InvalidFormat("invalid hash algorithm".to_string())),
202        };
203        let rule = match cursor.read_u8()? {
204            0 => BoundaryRule::HashThreshold {
205                factor: cursor.read_u32()?,
206            },
207            1 => BoundaryRule::Weibull {
208                shape: cursor.read_u32()?,
209            },
210            2 => BoundaryRule::RollingBuzHash {
211                window: cursor.read_u16()?,
212            },
213            _ => return Err(Error::InvalidFormat("invalid boundary rule".to_string())),
214        };
215        let chunking = ChunkingSpec {
216            measure,
217            input,
218            hash,
219            rule,
220            min: cursor.read_u64()?,
221            target: cursor.read_u64()?,
222            max: cursor.read_u64()?,
223            hash_seed: cursor.read_u64()?,
224            level_salt: match cursor.read_u8()? {
225                0 => false,
226                1 => true,
227                _ => return Err(Error::InvalidFormat("invalid level salt flag".to_string())),
228            },
229            hard_max_node_bytes: cursor.read_u64()?,
230        };
231        let node_layout = match cursor.read_u8()? {
232            0 => NodeLayoutSpec::PrefixCompressed,
233            1 => NodeLayoutSpec::Plain,
234            2 => NodeLayoutSpec::OffsetTable,
235            3 => NodeLayoutSpec::Custom {
236                id: cursor.read_string()?,
237                parameters: cursor.read_vec()?,
238            },
239            _ => return Err(Error::InvalidFormat("invalid node layout".to_string())),
240        };
241        let value_encoding = match cursor.read_u8()? {
242            0 => Encoding::Raw,
243            1 => Encoding::Cbor,
244            2 => Encoding::Json,
245            3 => Encoding::Custom(cursor.read_string()?),
246            _ => return Err(Error::InvalidFormat("invalid value encoding".to_string())),
247        };
248        if !cursor.is_done() {
249            return Err(Error::InvalidFormat(
250                "trailing tree format bytes".to_string(),
251            ));
252        }
253        let format = Self {
254            chunking,
255            node_layout,
256            value_encoding,
257        };
258        format.validate()?;
259        Ok(format)
260    }
261}
262
263fn encode_chunking(spec: &ChunkingSpec, out: &mut Vec<u8>) {
264    out.push(match spec.measure {
265        ChunkMeasure::EntryCount => 0,
266        ChunkMeasure::LogicalBytes => 1,
267        ChunkMeasure::EncodedBytes => 2,
268    });
269    out.push(match spec.input {
270        BoundaryInput::Key => 0,
271        BoundaryInput::KeyValue => 1,
272    });
273    out.push(match spec.hash {
274        HashAlgorithm::XxHash64 => 0,
275    });
276    match spec.rule {
277        BoundaryRule::HashThreshold { factor } => {
278            out.push(0);
279            out.extend_from_slice(&factor.to_be_bytes());
280        }
281        BoundaryRule::Weibull { shape } => {
282            out.push(1);
283            out.extend_from_slice(&shape.to_be_bytes());
284        }
285        BoundaryRule::RollingBuzHash { window } => {
286            out.push(2);
287            out.extend_from_slice(&window.to_be_bytes());
288        }
289    }
290    out.extend_from_slice(&spec.min.to_be_bytes());
291    out.extend_from_slice(&spec.target.to_be_bytes());
292    out.extend_from_slice(&spec.max.to_be_bytes());
293    out.extend_from_slice(&spec.hash_seed.to_be_bytes());
294    out.push(u8::from(spec.level_salt));
295    out.extend_from_slice(&spec.hard_max_node_bytes.to_be_bytes());
296}
297
298fn encode_layout(layout: &NodeLayoutSpec, out: &mut Vec<u8>) {
299    match layout {
300        NodeLayoutSpec::PrefixCompressed => out.push(0),
301        NodeLayoutSpec::Plain => out.push(1),
302        NodeLayoutSpec::OffsetTable => out.push(2),
303        NodeLayoutSpec::Custom { id, parameters } => {
304            out.push(3);
305            encode_bytes(id.as_bytes(), out);
306            encode_bytes(parameters, out);
307        }
308    }
309}
310
311fn encode_value_encoding(encoding: &Encoding, out: &mut Vec<u8>) {
312    match encoding {
313        Encoding::Raw => out.push(0),
314        Encoding::Cbor => out.push(1),
315        Encoding::Json => out.push(2),
316        Encoding::Custom(name) => {
317            out.push(3);
318            encode_bytes(name.as_bytes(), out);
319        }
320    }
321}
322
323fn encode_bytes(bytes: &[u8], out: &mut Vec<u8>) {
324    out.extend_from_slice(&(bytes.len() as u64).to_be_bytes());
325    out.extend_from_slice(bytes);
326}
327
328struct FormatCursor<'a> {
329    bytes: &'a [u8],
330    position: usize,
331}
332
333impl<'a> FormatCursor<'a> {
334    fn new(bytes: &'a [u8]) -> Self {
335        Self { bytes, position: 0 }
336    }
337
338    fn expect(&mut self, expected: &[u8]) -> Result<(), Error> {
339        if self.read(expected.len())? != expected {
340            return Err(Error::InvalidFormat(
341                "invalid tree format magic".to_string(),
342            ));
343        }
344        Ok(())
345    }
346
347    fn read(&mut self, len: usize) -> Result<&'a [u8], Error> {
348        let end = self
349            .position
350            .checked_add(len)
351            .ok_or_else(|| Error::InvalidFormat("tree format offset overflow".to_string()))?;
352        let value = self
353            .bytes
354            .get(self.position..end)
355            .ok_or_else(|| Error::InvalidFormat("truncated tree format".to_string()))?;
356        self.position = end;
357        Ok(value)
358    }
359
360    fn read_u8(&mut self) -> Result<u8, Error> {
361        Ok(self.read(1)?[0])
362    }
363
364    fn read_u16(&mut self) -> Result<u16, Error> {
365        Ok(u16::from_be_bytes(self.read(2)?.try_into().map_err(
366            |_| Error::InvalidFormat("invalid u16".to_string()),
367        )?))
368    }
369
370    fn read_u32(&mut self) -> Result<u32, Error> {
371        Ok(u32::from_be_bytes(self.read(4)?.try_into().map_err(
372            |_| Error::InvalidFormat("invalid u32".to_string()),
373        )?))
374    }
375
376    fn read_u64(&mut self) -> Result<u64, Error> {
377        Ok(u64::from_be_bytes(self.read(8)?.try_into().map_err(
378            |_| Error::InvalidFormat("invalid u64".to_string()),
379        )?))
380    }
381
382    fn read_vec(&mut self) -> Result<Vec<u8>, Error> {
383        let len = usize::try_from(self.read_u64()?)
384            .map_err(|_| Error::InvalidFormat("tree format length overflow".to_string()))?;
385        Ok(self.read(len)?.to_vec())
386    }
387
388    fn read_string(&mut self) -> Result<String, Error> {
389        String::from_utf8(self.read_vec()?)
390            .map_err(|_| Error::InvalidFormat("tree format string is not UTF-8".to_string()))
391    }
392
393    fn is_done(&self) -> bool {
394        self.position == self.bytes.len()
395    }
396}