Skip to main content

dynamo_tokens/
lib.rs

1// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4#![deny(missing_docs)]
5
6//! Types and utilities for handling sequences of tokens, including block creation and hashing.
7
8use bytemuck::cast_slice;
9use derive_getters::Dissolve;
10use std::ops::Range;
11
12pub mod blocks;
13mod radix;
14pub use radix::PositionalRadixTree;
15
16/// Trait for hashes that include position information.
17pub trait PositionalHash {
18    /// Returns the position associated with the hash.
19    fn position(&self) -> u64;
20}
21
22/// A token is represented as a 32-bit unsigned integer.
23pub type Token = u32;
24
25/// A salt used for hashing, represented as a vector of bytes.
26/// This might encode model architecture, weights, PEFT info, etc.
27pub type Salt = Vec<u8>;
28
29/// A 64-bit hash of the salt. Computed once per request and used as the seed for
30/// every block-hash computation in that request.
31///
32/// The canonical construction path is [`compute_salt_hash_from_bytes`] (or
33/// `dynamo_kv_hashing::Request::salt_hash` at the application layer).
34pub type SaltHash = u64;
35
36/// A 64-bit hash computed from the tokens within a single block (with optional MM
37/// frames), seeded by the request's [`SaltHash`].
38///
39/// The canonical construction path is [`compute_block_hash`].
40pub type BlockHash = u64;
41
42/// A 64-bit sequence-aware hash. Equals the [`BlockHash`] at position 0 and
43/// [`compute_next_sequence_hash(prev_seq, block_hash)`](compute_next_sequence_hash)
44/// at every subsequent position. Salt propagates through `seq_hash[0]` since
45/// `block_hash[0]` already encodes it.
46pub type SequenceHash = u64;
47
48/// Computes a hash of the data using the given seed (raw u64).
49///
50/// Prefer [`compute_block_hash`] / [`compute_salt_hash_from_bytes`] for typed
51/// construction; this raw-u64 form is kept for low-level callers.
52pub fn compute_hash_v2(data: &[u8], seed: u64) -> u64 {
53    xxhash_rust::xxh3::xxh3_64_with_seed(data, seed)
54}
55
56/// Canonical XXH3 seed used by every chain-step `(parent_seq, child_block_hash) → next_seq`
57/// in this codebase. Must match `dynamo_kv_router::protocols::XXH3_SEED` — the router's
58/// `compute_seq_hash_for_block` and the `PositionalIndexer`'s chain re-computation both
59/// route through [`compute_next_sequence_hash`] (or equivalently seeded helpers in
60/// `kv-router`), so changing this constant requires all of them to flip in lockstep.
61pub const CHAIN_XXH3_SEED: u64 = 1337;
62
63/// Chain-step for sequence hashing: returns the [`SequenceHash`] at `position + 1`
64/// given the parent's `SequenceHash` and the child block's [`BlockHash`].
65///
66/// This is the single source of truth for the chain recurrence used by:
67/// - [`PositionalLineageHash::extend`]
68/// - [`TokenBlock::from_chunk`]
69/// - `dynamo_kv_router::protocols::compute_next_seq_hash` (request side)
70/// - `dynamo_kv_router::indexer::PositionalIndexer` (chain re-validation)
71///
72/// Salt is already mixed into `block_hash[0]` and propagates through every parent, so
73/// the per-step seed is a constant: re-feeding salt at each step would be redundant.
74/// Any constant seed preserves PLH composability — the value is set to
75/// [`CHAIN_XXH3_SEED`] to match the router's existing wire format.
76#[inline]
77pub fn compute_next_sequence_hash(
78    parent_sequence_hash: SequenceHash,
79    child_block_hash: BlockHash,
80) -> SequenceHash {
81    let combined = [parent_sequence_hash, child_block_hash];
82    compute_hash_v2(cast_slice(&combined), CHAIN_XXH3_SEED)
83}
84
85/// Custom serde codec that encodes a `u128` as a 16-byte big-endian byte sequence.
86///
87/// MessagePack (`rmp-serde`) has no native 128-bit integer type, so the default
88/// `u128` derive does not roundtrip reliably. Encoding as raw bytes is supported
89/// uniformly across msgpack, JSON, CBOR, etc.
90mod serde_bytes_u128 {
91    use serde::{Deserializer, Serializer};
92
93    pub fn serialize<S>(val: &u128, serializer: S) -> Result<S::Ok, S::Error>
94    where
95        S: Serializer,
96    {
97        serializer.serialize_bytes(&val.to_be_bytes())
98    }
99
100    pub fn deserialize<'de, D>(deserializer: D) -> Result<u128, D::Error>
101    where
102        D: Deserializer<'de>,
103    {
104        use serde::de::{self, SeqAccess, Visitor};
105        use std::fmt;
106
107        struct V;
108        impl<'de> Visitor<'de> for V {
109            type Value = [u8; 16];
110
111            fn expecting(&self, f: &mut fmt::Formatter) -> fmt::Result {
112                f.write_str("16 bytes (msgpack bin) or a sequence of 16 u8 values")
113            }
114
115            fn visit_bytes<E: de::Error>(self, v: &[u8]) -> Result<[u8; 16], E> {
116                v.try_into()
117                    .map_err(|_| E::invalid_length(v.len(), &"16 bytes"))
118            }
119
120            fn visit_borrowed_bytes<E: de::Error>(self, v: &'de [u8]) -> Result<[u8; 16], E> {
121                self.visit_bytes(v)
122            }
123
124            fn visit_byte_buf<E: de::Error>(self, v: Vec<u8>) -> Result<[u8; 16], E> {
125                self.visit_bytes(&v)
126            }
127
128            fn visit_seq<A: SeqAccess<'de>>(self, mut seq: A) -> Result<[u8; 16], A::Error> {
129                let mut arr = [0u8; 16];
130                for (i, slot) in arr.iter_mut().enumerate() {
131                    *slot = seq
132                        .next_element()?
133                        .ok_or_else(|| de::Error::invalid_length(i, &"16 u8 elements"))?;
134                }
135                Ok(arr)
136            }
137        }
138
139        let arr = deserializer.deserialize_bytes(V)?;
140        Ok(u128::from_be_bytes(arr))
141    }
142}
143
144/// Canonical [`BlockHash`] construction: XXH3 over the per-block byte buffer
145/// (already encoded by [`compute_block_bytes_with_mm`] or `cast_slice` for the
146/// no-MM path), seeded by [`SaltHash`].
147#[inline]
148pub fn compute_block_hash(block_bytes: &[u8], salt: SaltHash) -> BlockHash {
149    compute_hash_v2(block_bytes, salt)
150}
151
152/// Canonical [`BlockHash`] construction for a token-only block.
153///
154/// Multimodal token sequences must first route through
155/// [`compute_block_bytes_with_mm`] so placeholder slots use their multimodal
156/// identities instead of their token IDs.
157#[inline]
158pub fn compute_block_hash_for_tokens(tokens: &[Token], salt: SaltHash) -> BlockHash {
159    compute_block_hash(cast_slice(tokens), salt)
160}
161
162/// Canonical [`SaltHash`] construction from a pre-canonicalized salt-payload byte
163/// buffer. Application-layer callers should use `dynamo_kv_hashing::Request::salt_hash`
164/// which canonicalizes `(salt, lora_name)` first; this function is the low-level path.
165#[inline]
166pub fn compute_salt_hash_from_bytes(payload: &[u8]) -> SaltHash {
167    compute_hash_v2(payload, 0)
168}
169
170/// Metadata describing a single multimodal placeholder run within a token sequence.
171///
172/// A run occupies `length` consecutive slots starting at `offset`. The token IDs at
173/// those slot positions are opaque for hashing — the `(mm_hash, run_offset)` pair drives
174/// the per-slot bytes during block formation.
175///
176/// See [`compute_block_bytes_with_mm`] for the byte-encoding rule.
177#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, serde::Serialize, serde::Deserialize)]
178pub struct TokenBlockMmInfo {
179    /// Hash identifying the multimodal object (image / audio / etc.).
180    pub mm_hash: u64,
181    /// Start position of the placeholder run in the full token sequence (zero-based).
182    pub offset: usize,
183    /// Number of placeholder slots in the run.
184    pub length: usize,
185}
186
187/// Slot-tag byte distinguishing real-token slots from multimodal placeholder slots in
188/// the per-block byte buffer. See [`compute_block_bytes_with_mm`].
189pub const MM_SLOT_TAG_TOKEN: u8 = 0x00;
190/// Slot-tag byte for multimodal placeholder slots. See [`compute_block_bytes_with_mm`].
191pub const MM_SLOT_TAG_PLACEHOLDER: u8 = 0x01;
192
193impl TokenBlockMmInfo {
194    /// Returns the exclusive end position of this run, or `None` on `usize` overflow.
195    #[inline]
196    pub fn checked_end(&self) -> Option<usize> {
197        self.offset.checked_add(self.length)
198    }
199
200    /// Returns the exclusive end position of this run.
201    ///
202    /// # Panics
203    /// Panics if `offset + length` overflows `usize`. Prefer [`Self::checked_end`] in
204    /// validation paths; this helper is for already-validated runs.
205    #[inline]
206    pub fn end(&self) -> usize {
207        self.checked_end()
208            .expect("TokenBlockMmInfo::end overflowed usize; run was not validated")
209    }
210
211    /// Returns `true` if the given absolute position falls inside this run.
212    /// Returns `false` if the run's end overflows `usize` (such a run is invalid; use
213    /// [`validate_and_sort_mm_info`] before relying on this method).
214    #[inline]
215    pub fn covers(&self, position: usize) -> bool {
216        match self.checked_end() {
217            Some(end) => position >= self.offset && position < end,
218            None => false,
219        }
220    }
221}
222
223/// Errors raised while validating [`TokenBlockMmInfo`] inputs.
224#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)]
225pub enum MmInfoError {
226    /// The run extends past the end of the token sequence.
227    #[error(
228        "mm_info range starting at {offset} (length {length}) exceeds tokens length {tokens_len}"
229    )]
230    OutOfBounds {
231        /// Run start.
232        offset: usize,
233        /// Run length.
234        length: usize,
235        /// Length of the token sequence the run was validated against.
236        tokens_len: usize,
237    },
238    /// `offset + length` overflows `usize`.
239    #[error("mm_info range starting at {offset} (length {length}) overflows usize")]
240    OffsetOverflow {
241        /// Run start.
242        offset: usize,
243        /// Run length.
244        length: usize,
245    },
246    /// Two runs overlap.
247    #[error("mm_info ranges overlap at position {position}")]
248    Overlapping {
249        /// Position where the overlap begins.
250        position: usize,
251    },
252    /// A run has zero length.
253    #[error("mm_info length must be greater than zero")]
254    EmptyRun,
255}
256
257/// Validates `mm_info` against `tokens_len` and returns a copy sorted by `offset`.
258///
259/// Validation rules:
260/// - Every run must have `length > 0`.
261/// - `offset + length` must not overflow `usize`.
262/// - Every run's end (`offset + length`) must be `<= tokens_len`.
263/// - No two runs may overlap.
264pub fn validate_and_sort_mm_info(
265    mm_info: &[TokenBlockMmInfo],
266    tokens_len: usize,
267) -> Result<Vec<TokenBlockMmInfo>, MmInfoError> {
268    let mut sorted: Vec<TokenBlockMmInfo> = mm_info.to_vec();
269    sorted.sort_by_key(|m| m.offset);
270    let mut prev_end = 0usize;
271    for m in &sorted {
272        if m.length == 0 {
273            return Err(MmInfoError::EmptyRun);
274        }
275        let end = m
276            .offset
277            .checked_add(m.length)
278            .ok_or(MmInfoError::OffsetOverflow {
279                offset: m.offset,
280                length: m.length,
281            })?;
282        if end > tokens_len {
283            return Err(MmInfoError::OutOfBounds {
284                offset: m.offset,
285                length: m.length,
286                tokens_len,
287            });
288        }
289        if m.offset < prev_end {
290            return Err(MmInfoError::Overlapping { position: m.offset });
291        }
292        prev_end = end;
293    }
294    Ok(sorted)
295}
296
297/// Returns `true` if any run in `mm_runs` overlaps the block `[block_offset, block_offset + len)`.
298/// `mm_runs` must be validated and sorted.
299fn block_has_mm(block_offset: usize, len: usize, mm_runs: &[TokenBlockMmInfo]) -> bool {
300    let block_end = block_offset.saturating_add(len);
301    mm_runs
302        .iter()
303        .any(|m| m.offset < block_end && m.end() > block_offset)
304}
305
306/// Builds the byte buffer used to compute a block's [`BlockHash`].
307///
308/// **Two encodings, picked per-block:**
309///
310/// 1. **Legacy / zero-MM** — when no run in `mm_runs` overlaps this block, the buffer is
311///    `bytemuck::cast_slice(tokens)` (4 bytes per slot, LE u32). This matches the existing
312///    `compute_hash_v2(cast_slice(&tokens), salt_hash)` path used by every existing
313///    [`TokenBlock`] in `dynamo_tokens` and kvbm — preserving cache identity for any block
314///    that is not itself MM-affected.
315///
316/// 2. **Tagged / MM-affected** — when at least one run overlaps this block, every slot
317///    emits a fixed 13-byte frame:
318///    - Real-token slot: `[MM_SLOT_TAG_TOKEN | token_id u32 LE | 0u64 LE]`
319///      The trailing `0u64` is **frame padding only** — it has no semantic meaning,
320///      it is there so the real-token frame matches the placeholder frame's width.
321///      Token IDs at placeholder positions are *ignored* by this encoder; whether
322///      a slot is a placeholder is determined solely by `mm_runs[run_idx].covers(g)`.
323///    - Placeholder slot: `[MM_SLOT_TAG_PLACEHOLDER | run_offset u32 LE | mm_hash u64 LE]`,
324///      where `run_offset = (block_offset + s) - run.offset`.
325///
326///    The 1-byte tag plus the fixed-width frame make the encoding self-delimiting and
327///    slot-position-preserving: two MM-affected byte buffers compare equal iff they describe
328///    the same `(slot_kind, slot_payload)` sequence at the same slot positions.
329///
330/// The two encodings have different per-slot widths (4 vs 13), so an all-tokens block can
331/// never produce the same byte buffer as an MM-affected block — eliminating cross-encoding
332/// collisions.
333///
334/// `mm_runs` must be validated and sorted (typically the output of
335/// [`validate_and_sort_mm_info`]).
336pub fn compute_block_bytes_with_mm(
337    tokens: &[Token],
338    block_offset: usize,
339    mm_runs: &[TokenBlockMmInfo],
340) -> Vec<u8> {
341    // Defense-in-depth: routing of each slot to the placeholder vs real-token branch
342    // depends on `mm_runs` being sorted by offset and non-overlapping. The public
343    // entry points (`Request::new`, `TokenBlockSequence::new_with_mm`,
344    // `split_tokens_with_mm`) enforce this via `validate_and_sort_mm_info`, but this
345    // function is also pub so we re-check in debug to catch direct misuse.
346    debug_assert!(
347        mm_runs.windows(2).all(|w| w[0].end() <= w[1].offset),
348        "compute_block_bytes_with_mm: mm_runs must be sorted by offset and non-overlapping (use validate_and_sort_mm_info)",
349    );
350    debug_assert!(
351        mm_runs.iter().all(|r| r.length > 0),
352        "compute_block_bytes_with_mm: mm_runs must have non-zero length",
353    );
354
355    if !block_has_mm(block_offset, tokens.len(), mm_runs) {
356        // Zero-MM-affecting-this-block: use the legacy encoding so the resulting block_hash
357        // matches what TokenBlockSequence::new would produce.
358        return cast_slice::<Token, u8>(tokens).to_vec();
359    }
360
361    const FRAME: usize = 13;
362    let mut out: Vec<u8> = Vec::with_capacity(tokens.len() * FRAME);
363    let mut run_idx = 0usize;
364    // Skip runs that ended at or before this block starts.
365    while run_idx < mm_runs.len() && mm_runs[run_idx].end() <= block_offset {
366        run_idx += 1;
367    }
368    for (s, &tok) in tokens.iter().enumerate() {
369        let g = block_offset + s;
370        // Advance past runs that end at or before g (validated non-overlapping => monotonic).
371        while run_idx < mm_runs.len() && mm_runs[run_idx].end() <= g {
372            run_idx += 1;
373        }
374        if run_idx < mm_runs.len() && mm_runs[run_idx].covers(g) {
375            let run = &mm_runs[run_idx];
376            let run_offset = (g - run.offset) as u32;
377            out.push(MM_SLOT_TAG_PLACEHOLDER);
378            out.extend_from_slice(&run_offset.to_le_bytes());
379            out.extend_from_slice(&run.mm_hash.to_le_bytes());
380        } else {
381            out.push(MM_SLOT_TAG_TOKEN);
382            out.extend_from_slice(&tok.to_le_bytes());
383            out.extend_from_slice(&0u64.to_le_bytes());
384        }
385    }
386    out
387}
388
389/// A 128-bit positional sequence hash combining traditional sequence hash with positional information.
390///
391/// Layout:
392/// - Lower 64 bits: Traditional SequenceHash
393/// - Upper 64 bits: 2-bit mode + position + LocalBlockHash (BlockHash)
394///
395/// Modes (automatically selected based on position):
396/// - Mode 00: 8-bit position (max 255) + 54-bit LBH
397/// - Mode 01: 16-bit position (max 65,535) + 46-bit LBH
398/// - Mode 10: 24-bit position (max 16,777,215) + 38-bit LBH
399/// - Mode 11: 31-bit position (max 2,147,483,647) + 31-bit LBH
400#[derive(Clone, Copy, PartialEq, Eq, Hash, Default, serde::Serialize, serde::Deserialize)]
401#[serde(transparent)]
402pub struct PositionalSequenceHash(#[serde(with = "serde_bytes_u128")] u128);
403
404impl PositionalSequenceHash {
405    /// Creates a new PositionalSequenceHash from components.
406    ///
407    /// The mode is automatically selected based on the position value to use the minimal
408    /// representation that can fit the position.
409    pub fn new(sequence_hash: SequenceHash, position: u64, local_block_hash: BlockHash) -> Self {
410        let mode = Self::select_mode(position);
411        let upper = Self::encode_upper(mode, position, local_block_hash);
412        let value = ((upper as u128) << 64) | (sequence_hash as u128);
413        PositionalSequenceHash(value)
414    }
415
416    /// Returns the sequence hash component (lower 64 bits).
417    pub fn sequence_hash(&self) -> SequenceHash {
418        (self.0 & 0xFFFF_FFFF_FFFF_FFFF) as u64
419    }
420
421    /// Returns the block position.
422    pub fn position(&self) -> u64 {
423        let (_, position, _) = self.decode_upper();
424        position
425    }
426
427    /// Returns the local block hash (BlockHash) component.
428    pub fn local_block_hash(&self) -> BlockHash {
429        let (_, _, lbh) = self.decode_upper();
430        lbh
431    }
432
433    /// Returns the mode used for encoding (0, 1, 2, or 3).
434    pub fn mode(&self) -> u8 {
435        let (mode, _, _) = self.decode_upper();
436        mode
437    }
438
439    /// Returns the inner 128-bit value.
440    #[inline(always)]
441    pub fn as_u128(&self) -> u128 {
442        self.0
443    }
444
445    /// Selects the minimal mode that can represent the given position.
446    fn select_mode(position: u64) -> u8 {
447        if position < (1u64 << 8) {
448            0 // Mode 00: 8-bit position
449        } else if position < (1u64 << 16) {
450            1 // Mode 01: 16-bit position
451        } else if position < (1u64 << 24) {
452            2 // Mode 10: 24-bit position
453        } else if position < (1u64 << 31) {
454            3 // Mode 11: 31-bit position
455        } else {
456            panic!(
457                "Position {} exceeds maximum supported value (2^31 - 1)",
458                position
459            );
460        }
461    }
462
463    /// Encodes the upper 64 bits from mode, position, and local block hash.
464    fn encode_upper(mode: u8, position: u64, local_block_hash: u64) -> u64 {
465        let (position_bits, lbh_bits) = match mode {
466            0 => (8, 54),  // 2 + 8 + 54 = 64
467            1 => (16, 46), // 2 + 16 + 46 = 64
468            2 => (24, 38), // 2 + 24 + 38 = 64
469            3 => (31, 31), // 2 + 31 + 31 = 64
470            _ => unreachable!(
471                "Invalid mode {} when encoding PositionalSequenceHash; mode must be 0, 1, 2, or 3",
472                mode
473            ),
474        };
475
476        // Create masks for extracting the relevant bits
477        let position_mask = (1u64 << position_bits) - 1;
478        let lbh_mask = (1u64 << lbh_bits) - 1;
479
480        // Extract and position components
481        let position_part = position & position_mask;
482        let lbh_part = local_block_hash & lbh_mask;
483
484        // Combine: [mode (2 bits)][position (X bits)][lbh (R bits)]
485        ((mode as u64) << 62) | (position_part << lbh_bits) | lbh_part
486    }
487
488    /// Decodes the upper 64 bits into (mode, position, local_block_hash).
489    fn decode_upper(&self) -> (u8, u64, u64) {
490        let upper = (self.0 >> 64) as u64;
491
492        // Extract mode from top 2 bits
493        let mode = (upper >> 62) as u8;
494
495        let (position_bits, lbh_bits) = match mode {
496            0 => (8, 54),
497            1 => (16, 46),
498            2 => (24, 38),
499            3 => (31, 31),
500            _ => unreachable!(
501                "Invalid mode {} in PositionalSequenceHash - value may be corrupted",
502                mode
503            ),
504        };
505
506        // Create masks
507        let lbh_mask = (1u64 << lbh_bits) - 1;
508        let position_mask = (1u64 << position_bits) - 1;
509
510        // Extract components
511        let lbh = upper & lbh_mask;
512        let position = (upper >> lbh_bits) & position_mask;
513
514        (mode, position, lbh)
515    }
516}
517
518impl std::fmt::Debug for PositionalSequenceHash {
519    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
520        f.debug_struct("PositionalSequenceHash")
521            .field("sequence_hash", &self.sequence_hash())
522            .field("local_block_hash", &self.local_block_hash())
523            .field("position", &self.position())
524            .finish()
525    }
526}
527
528/// A 128-bit positional lineage hash encoding parental lineage for tree traversal.
529///
530/// Layout (using full 128 bits):
531/// - Mode (2 bits): Determines position field size
532/// - Position (8/16/24 bits): Block position in sequence
533/// - Current Sequence Hash (64 bits): Full u64 sequence hash for this block
534/// - Parent Fragment (variable bits): Truncated lower bits of the parent's sequence hash
535///
536/// Modes (automatically selected based on position):
537/// - Mode 00: 8-bit position (max 255) + 64-bit current + 54-bit parent fragment
538/// - Mode 01: 16-bit position (max 65,535) + 64-bit current + 46-bit parent fragment
539/// - Mode 10: 24-bit position (max 16,777,215) + 64-bit current + 38-bit parent fragment
540///
541/// Carrying the full 64-bit current sequence hash inline makes PLH self-contained for
542/// chain extension: a child PLH can be derived from a parent PLH plus the child's
543/// `BlockHash` alone, with no out-of-band state. The parent fragment shrinks as
544/// position grows, but radix backward-traversal always matches on
545/// `(position, parent_fragment)`, and the probability of a shared parent at high
546/// positions decreases faster than the fragment-collision bound rises.
547#[derive(Clone, Copy, PartialEq, Eq, Hash, Default, serde::Serialize, serde::Deserialize)]
548#[serde(transparent)]
549pub struct PositionalLineageHash(#[serde(with = "serde_bytes_u128")] u128);
550
551impl PositionalLineageHash {
552    /// Creates a new PositionalLineageHash from components.
553    ///
554    /// The mode is automatically selected based on the position value to use the minimal
555    /// representation that can fit the position.
556    ///
557    /// # Arguments
558    ///
559    /// * `current_seq_hash` - The full u64 sequence hash of the current block
560    /// * `parent_seq_hash` - The full u64 sequence hash of the parent block (None for root).
561    ///   Will be truncated to this position's parent-fragment width.
562    /// * `position` - The block position in the sequence
563    ///
564    /// # Panics
565    ///
566    /// Panics if position >= 2^24 (16,777,216).
567    pub fn new(
568        current_seq_hash: SequenceHash,
569        parent_seq_hash: Option<SequenceHash>,
570        position: u64,
571    ) -> Self {
572        if position >= (1u64 << 24) {
573            panic!(
574                "Position {} exceeds maximum supported value (2^24 - 1 = 16,777,215)",
575                position
576            );
577        }
578
579        let mode = Self::select_mode(position);
580        let (position_bits, parent_bits) = Self::bit_layout(mode);
581
582        let position_mask = (1u128 << position_bits) - 1;
583        let parent_mask = (1u128 << parent_bits) - 1;
584
585        let position_part = (position as u128) & position_mask;
586        let current_part = current_seq_hash as u128;
587        let parent_part = (parent_seq_hash.unwrap_or(0) as u128) & parent_mask;
588
589        // Pack: [mode (2)][position (P)][current_u64 (64)][parent_fragment (R)]
590        let value = ((mode as u128) << 126)
591            | (position_part << (64 + parent_bits))
592            | (current_part << parent_bits)
593            | parent_part;
594
595        PositionalLineageHash(value)
596    }
597
598    /// Creates a root [`PositionalLineageHash`] (position 0).
599    ///
600    /// At the root, the sequence hash equals the block hash and there is no parent.
601    pub fn root(block_hash: BlockHash) -> Self {
602        Self::new(block_hash, None, 0)
603    }
604
605    /// Extends this lineage by one block, producing the child PLH.
606    ///
607    /// The chain recurrence is [`compute_next_sequence_hash`]. Salt does not seed the
608    /// per-step xxh3: `BlockHash` is already `xxh3(tokens, salt_hash)`, so salt is mixed
609    /// into `seq_hash[0]` and propagates through every parent. Re-feeding it at each step
610    /// would be redundant.
611    ///
612    /// # Panics
613    ///
614    /// Panics if `self.position() + 1 >= 2^24`.
615    pub fn extend(&self, child_block_hash: BlockHash) -> Self {
616        let parent_seq = self.current_sequence_hash();
617        let child_seq = compute_next_sequence_hash(parent_seq, child_block_hash);
618        Self::new(child_seq, Some(parent_seq), self.position() + 1)
619    }
620
621    /// Returns the block position.
622    pub fn position(&self) -> u64 {
623        let mode = self.mode();
624        let (position_bits, parent_bits) = Self::bit_layout(mode);
625        let position_mask = (1u128 << position_bits) - 1;
626        ((self.0 >> (64 + parent_bits)) & position_mask) as u64
627    }
628
629    /// Returns the full 64-bit sequence hash of the current block.
630    ///
631    /// Unlike the legacy `current_hash_fragment` (now removed), this is the complete
632    /// `SequenceHash` and is what [`extend`](Self::extend) feeds into the chain.
633    pub fn current_sequence_hash(&self) -> SequenceHash {
634        let mode = self.mode();
635        let (_, parent_bits) = Self::bit_layout(mode);
636        ((self.0 >> parent_bits) & 0xFFFF_FFFF_FFFF_FFFFu128) as u64
637    }
638
639    /// Returns the parent sequence hash fragment as stored in this PLH.
640    ///
641    /// Width depends on this PLH's mode (54/46/38 bits). Backward radix lookup must
642    /// always combine `(position, parent_fragment)` — the fragment alone is not unique.
643    pub fn parent_hash_fragment(&self) -> u64 {
644        let mode = self.mode();
645        let (_, parent_bits) = Self::bit_layout(mode);
646        let parent_mask = (1u128 << parent_bits) - 1;
647        (self.0 & parent_mask) as u64
648    }
649
650    /// Truncates this PLH's `current_sequence_hash` to the parent-fragment width that
651    /// a child at `child_position` would store.
652    ///
653    /// Useful when an external builder wants to construct a child PLH or verify a
654    /// parent/child edge: `child.parent_hash_fragment() ==
655    /// parent.parent_fragment_for_child_position(child.position())`.
656    pub fn parent_fragment_for_child_position(&self, child_position: u64) -> u64 {
657        let child_mode = Self::select_mode(child_position);
658        let (_, child_parent_bits) = Self::bit_layout(child_mode);
659        let mask = (1u64 << child_parent_bits).wrapping_sub(1);
660        self.current_sequence_hash() & mask
661    }
662
663    /// Returns the mode used for encoding (0, 1, or 2).
664    pub fn mode(&self) -> u8 {
665        (self.0 >> 126) as u8
666    }
667
668    /// Returns the inner 128-bit value.
669    #[inline(always)]
670    pub fn as_u128(&self) -> u128 {
671        self.0
672    }
673
674    /// Selects the minimal mode that can represent the given position.
675    fn select_mode(position: u64) -> u8 {
676        if position < (1u64 << 8) {
677            0 // Mode 00: 8-bit position
678        } else if position < (1u64 << 16) {
679            1 // Mode 01: 16-bit position
680        } else {
681            2 // Mode 10: 24-bit position
682        }
683    }
684
685    /// Returns the bit layout for a given mode: (position_bits, parent_fragment_bits).
686    /// Current is always 64 bits.
687    fn bit_layout(mode: u8) -> (u32, u32) {
688        match mode {
689            0 => (8, 54),  // 2 + 8 + 64 + 54 = 128
690            1 => (16, 46), // 2 + 16 + 64 + 46 = 128
691            2 => (24, 38), // 2 + 24 + 64 + 38 = 128
692            _ => unreachable!(
693                "Invalid mode {} in PositionalLineageHash; mode must be 0, 1, or 2",
694                mode
695            ),
696        }
697    }
698}
699
700impl PositionalLineageHash {
701    fn format_impl(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
702        let position = self.position();
703        let current_hash = self.current_sequence_hash();
704        let current_hash_b58 = bs58::encode(current_hash.to_be_bytes()).into_string();
705
706        if position == 0 {
707            write!(f, "{}:{}", position, current_hash_b58)
708        } else {
709            let parent_hash = self.parent_hash_fragment();
710            let parent_hash_b58 = bs58::encode(parent_hash.to_be_bytes()).into_string();
711            write!(f, "{}:{}:{}", position, current_hash_b58, parent_hash_b58)
712        }
713    }
714}
715
716impl std::fmt::Debug for PositionalLineageHash {
717    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
718        self.format_impl(f)
719    }
720}
721
722impl std::fmt::Display for PositionalLineageHash {
723    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
724        self.format_impl(f)
725    }
726}
727
728impl std::cmp::PartialOrd for PositionalLineageHash {
729    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
730        Some(self.cmp(other))
731    }
732}
733
734impl std::cmp::Ord for PositionalLineageHash {
735    /// Lexicographic order: [`Self::position`], then [`Self::current_sequence_hash`],
736    /// then the full packed [`Self::as_u128`] so the order is total and consistent with [`Eq`].
737    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
738        self.position()
739            .cmp(&other.position())
740            .then_with(|| {
741                self.current_sequence_hash()
742                    .cmp(&other.current_sequence_hash())
743            })
744            .then_with(|| self.0.cmp(&other.0))
745    }
746}
747
748/// A collection of tokens, represented as a `Vec<Token>`.
749///
750/// Provides convenience methods for conversion and manipulation.
751#[derive(Debug, Clone, Dissolve, Default, Eq)]
752pub struct Tokens(Vec<Token>);
753
754impl AsRef<[Token]> for Tokens {
755    fn as_ref(&self) -> &[Token] {
756        &self.0
757    }
758}
759
760impl std::ops::Deref for Tokens {
761    type Target = [Token];
762
763    fn deref(&self) -> &Self::Target {
764        &self.0
765    }
766}
767
768impl std::borrow::Borrow<[Token]> for Tokens {
769    fn borrow(&self) -> &[Token] {
770        &self.0
771    }
772}
773
774impl From<Vec<Token>> for Tokens {
775    fn from(tokens: Vec<Token>) -> Self {
776        Tokens(tokens)
777    }
778}
779
780impl From<&[Token]> for Tokens {
781    fn from(tokens: &[Token]) -> Self {
782        Tokens(tokens.to_vec())
783    }
784}
785
786impl From<Vec<usize>> for Tokens {
787    fn from(tokens: Vec<usize>) -> Self {
788        Tokens(
789            tokens
790                .into_iter()
791                .map(|t| t.try_into().expect("Token ID exceeds u32::MAX"))
792                .collect(),
793        )
794    }
795}
796
797impl From<Vec<i32>> for Tokens {
798    /// Converts `Vec<i32>` to `Tokens`, casting each `i32` to `u32`.
799    fn from(tokens: Vec<i32>) -> Self {
800        Tokens(tokens.into_iter().map(|t| t as u32).collect())
801    }
802}
803
804impl From<&[i32]> for Tokens {
805    /// Converts `&[i32]` to `Tokens`, casting each `i32` to `u32`.
806    fn from(tokens: &[i32]) -> Self {
807        Tokens(tokens.iter().map(|&t| t as u32).collect())
808    }
809}
810
811impl From<Tokens> for Vec<Token> {
812    fn from(tokens: Tokens) -> Self {
813        tokens.0
814    }
815}
816
817// PartialEq implementations for comparing Tokens with Vec<Token> and &[Token]
818// (Generated implementations are usually sufficient, but explicit ones can be clearer)
819impl PartialEq<Vec<Token>> for Tokens {
820    fn eq(&self, other: &Vec<Token>) -> bool {
821        self.0 == *other
822    }
823}
824
825impl PartialEq<Tokens> for Vec<Token> {
826    fn eq(&self, other: &Tokens) -> bool {
827        *self == other.0
828    }
829}
830
831impl PartialEq<[Token]> for Tokens {
832    fn eq(&self, other: &[Token]) -> bool {
833        self.0.as_slice() == other
834    }
835}
836
837impl PartialEq<Tokens> for &[Token] {
838    fn eq(&self, other: &Tokens) -> bool {
839        *self == other.0.as_slice()
840    }
841}
842
843impl PartialEq for Tokens {
844    fn eq(&self, other: &Self) -> bool {
845        self.0 == other.0
846    }
847}
848
849// Add PartialEq<&[T]> where T: Into<Token> + Copy could be more general,
850// but specifically implementing for &[Token] is sufficient for the tests.
851impl PartialEq<&[Token]> for Tokens {
852    fn eq(&self, other: &&[Token]) -> bool {
853        self.0.as_slice() == *other
854    }
855}
856
857impl Tokens {
858    fn with_capacity(capacity: usize) -> Self {
859        Tokens(Vec::with_capacity(capacity))
860    }
861
862    /// Consumes the [`Tokens`] object and creates a [`TokenBlockSequence`].
863    ///
864    /// The sequence is initialized with the provided tokens, splitting them into blocks
865    /// of the specified `block_size` using the given `salt_hash` (or 0 if `None`).
866    ///
867    /// # Arguments
868    ///
869    /// * `block_size` - The fixed size for each [`TokenBlock`].
870    /// * `salt_hash` - An optional [`SaltHash`] used as the base seed for hashing. Defaults to 0.
871    pub fn into_sequence(self, block_size: u32, salt_hash: Option<SaltHash>) -> TokenBlockSequence {
872        TokenBlockSequence::new(self, block_size, salt_hash)
873    }
874}
875
876/// Errors that can occur during [`PartialTokenBlock`] operations.
877#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)]
878pub enum TokenBlockError {
879    /// The operation could not be completed because the block is full.
880    #[error("TokenBlock is full")]
881    Full,
882
883    /// The operation requires a full block, but the block is incomplete.
884    #[error("TokenBlock is incomplete")]
885    Incomplete,
886
887    /// The operation could not be completed because the block is empty.
888    #[error("TokenBlock is empty")]
889    Empty,
890
891    /// The operation requires more tokens than are currently in the block.
892    #[error("TokenBlock has insufficient tokens")]
893    InsufficientTokens,
894
895    /// Multimodal info validation failed.
896    #[error(transparent)]
897    MmInfo(#[from] MmInfoError),
898
899    /// A mutating operation is not supported on a sequence with multimodal runs.
900    #[error("operation is not supported on a TokenBlockSequence with multimodal runs")]
901    MmRunsPresent,
902}
903
904/// Represents a partially filled block of tokens within a sequence.
905///
906/// This structure accumulates tokens until it reaches the specified `block_size`,
907/// at which point it can be [`commit`](PartialTokenBlock::commit)ted into a full [`TokenBlock`].
908#[derive(Debug, PartialEq)] // No Clone: intended to be unique within a sequence
909pub struct PartialTokenBlock {
910    tokens: Tokens,
911    block_size: u32,
912    salt_hash: SaltHash,
913    parent_sequence_hash: Option<SequenceHash>,
914    position: usize, // The position this block will have when committed
915}
916
917impl PartialTokenBlock {
918    /// Creates the first partial block (root) for a new sequence.
919    ///
920    /// # Arguments
921    ///
922    /// * `block_size` - The fixed size for blocks in this sequence.
923    /// * `salt_hash` - The [`SaltHash`] for the sequence.
924    pub(crate) fn create_sequence_root(block_size: u32, salt_hash: SaltHash) -> Self {
925        Self {
926            tokens: Tokens::with_capacity(block_size as usize),
927            block_size,
928            salt_hash,
929            parent_sequence_hash: None, // Root has no parent
930            position: 0,                // First block is at position 0
931        }
932    }
933
934    /// Attempts to push multiple tokens onto the block from a [`Tokens`] object.
935    ///
936    /// Tokens are added until the block is full or all input tokens are consumed.
937    ///
938    /// # Arguments
939    ///
940    /// * `tokens` - The [`Tokens`] to push.
941    ///
942    /// # Returns
943    ///
944    /// A new [`Tokens`] object containing any tokens that did not fit,
945    /// if all tokens were added, the returned object will be empty.
946    pub(crate) fn push_tokens(&mut self, tokens: Tokens) -> Tokens {
947        let remaining_space = self.remaining();
948
949        if remaining_space == 0 {
950            return tokens; // Block is already full
951        }
952
953        if tokens.0.len() <= remaining_space {
954            // All tokens fit
955            self.tokens.0.extend(tokens.0);
956            Tokens::default() // No remaining tokens
957        } else {
958            // Only some tokens fit
959            let (to_add, remaining) = tokens.0.split_at(remaining_space);
960            self.tokens.0.extend_from_slice(to_add);
961            Tokens(remaining.to_vec()) // Return the leftover tokens
962        }
963    }
964
965    /// Attempts to push a single token onto the block.
966    ///
967    /// # Returns
968    ///
969    /// * `Ok(())` - If the token was successfully added.
970    /// * `Err(TokenBlockError::Full)` - If the block already contains `block_size` tokens.
971    #[cfg(test)]
972    pub(crate) fn push_token(&mut self, token: Token) -> Result<(), TokenBlockError> {
973        if self.tokens.0.len() >= self.block_size as usize {
974            return Err(TokenBlockError::Full);
975        }
976        self.tokens.0.push(token);
977        Ok(())
978    }
979
980    /// Attempts to remove the last `count` tokens from the block.
981    ///
982    /// # Arguments
983    ///
984    /// * `count` - The number of tokens to remove.
985    ///
986    /// # Returns
987    ///
988    /// * `Ok(())` - If the specified number of tokens were successfully removed.
989    /// * `Err(TokenBlockError::InsufficientTokens)` - If `count` is greater than the number of tokens in the block.
990    pub(crate) fn pop_tokens(&mut self, count: usize) -> Result<(), TokenBlockError> {
991        if self.tokens.0.len() < count {
992            return Err(TokenBlockError::InsufficientTokens);
993        }
994        self.tokens.0.truncate(self.tokens.0.len() - count);
995        Ok(())
996    }
997
998    /// Attempts to commit the current partial block into a full [`TokenBlock`].
999    ///
1000    /// This operation consumes the tokens within the partial block.
1001    /// After a successful commit, this `PartialTokenBlock` instance is reset
1002    /// to represent the *next* partial block in the sequence, inheriting the
1003    /// sequence hash from the block just committed.
1004    ///
1005    /// # Returns
1006    ///
1007    /// * `Ok(TokenBlock)` - The newly created full [`TokenBlock`].
1008    /// * `Err(TokenBlockError::Incomplete)` - If the block does not contain exactly `block_size` tokens.
1009    pub fn commit(&mut self) -> Result<TokenBlock, TokenBlockError> {
1010        if self.tokens.0.len() != self.block_size as usize {
1011            // Check for exact size match for committing
1012            return Err(TokenBlockError::Incomplete);
1013        }
1014
1015        // Take ownership of the tokens, leaving the internal tokens empty
1016        let tokens = std::mem::replace(
1017            &mut self.tokens,
1018            Tokens::with_capacity(self.block_size as usize),
1019        );
1020
1021        let chunk = TokenBlockChunk::new(tokens, self.salt_hash);
1022        let block = TokenBlock::from_chunk(chunk, self.parent_sequence_hash, self.position);
1023
1024        // Reset self to be the next block in the sequence
1025        self.parent_sequence_hash = Some(block.sequence_hash());
1026        self.position += 1; // Increment position for the next block
1027        // self.block_size and self.salt_hash remain the same
1028
1029        Ok(block)
1030    }
1031
1032    /// Returns the number of additional tokens required to fill the block.
1033    pub fn remaining(&self) -> usize {
1034        // Use saturating_sub to prevent underflow if len somehow exceeds block_size
1035        (self.block_size as usize).saturating_sub(self.tokens.0.len())
1036    }
1037
1038    /// Returns the number of tokens currently in the block.
1039    pub fn len(&self) -> usize {
1040        self.tokens.0.len()
1041    }
1042
1043    /// Returns `true` if the block contains no tokens.
1044    pub fn is_empty(&self) -> bool {
1045        self.tokens.0.is_empty()
1046    }
1047
1048    /// Returns a reference to the tokens currently in the block.
1049    pub fn tokens(&self) -> &Tokens {
1050        &self.tokens
1051    }
1052}
1053
1054// Deref allows treating &PartialTokenBlock like &Tokens for read-only access.
1055impl std::ops::Deref for PartialTokenBlock {
1056    type Target = Tokens;
1057
1058    fn deref(&self) -> &Self::Target {
1059        &self.tokens
1060    }
1061}
1062
1063/// An intermediate structure holding a chunk of tokens destined to become a [`TokenBlock`].
1064///
1065/// This calculates the [`BlockHash`] but does not compute the final [`SequenceHash`],
1066/// allowing chunks to be processed independently (e.g., in parallel).
1067#[derive(Debug)] // No Clone: temporary intermediate value
1068struct TokenBlockChunk {
1069    tokens: Tokens,
1070    salt_hash: SaltHash,
1071    block_hash: BlockHash,
1072}
1073
1074impl TokenBlockChunk {
1075    /// Creates a new chunk from [`Tokens`], calculating the [`BlockHash`].
1076    fn new(tokens: Tokens, salt_hash: SaltHash) -> Self {
1077        let block_hash = compute_block_hash_for_tokens(&tokens, salt_hash);
1078        Self {
1079            tokens,
1080            salt_hash,
1081            block_hash,
1082        }
1083    }
1084
1085    /// Creates a new chunk from a slice of `&[Token]`, calculating the [`BlockHash`].
1086    fn from_tokens(tokens: &[Token], salt_hash: SaltHash) -> Self {
1087        let block_hash = compute_block_hash_for_tokens(tokens, salt_hash);
1088        Self {
1089            tokens: tokens.into(), // Converts slice to owned Tokens
1090            salt_hash,
1091            block_hash,
1092        }
1093    }
1094}
1095
1096/// Represents a completed, immutable block of tokens with associated hashes.
1097///
1098/// Contains exactly `block_size` tokens and includes the [`SaltHash`], [`BlockHash`],
1099/// [`SequenceHash`], [`PositionalSequenceHash`], [`PositionalLineageHash`], and optionally the parent's [`SequenceHash`].
1100#[derive(Debug, Clone, Default, PartialEq)] // Add PartialEq for tests
1101pub struct TokenBlock {
1102    tokens: Tokens,
1103    salt_hash: SaltHash,
1104    block_hash: BlockHash,
1105    sequence_hash: SequenceHash,
1106    parent_sequence_hash: Option<SequenceHash>,
1107    positional_sequence_hash: PositionalSequenceHash,
1108    positional_lineage_hash: PositionalLineageHash,
1109}
1110
1111impl TokenBlock {
1112    /// Creates a new [`PartialTokenBlock`] representing the block immediately following this one.
1113    ///
1114    /// The new partial block will have the correct `parent_sequence_hash` and `position` set.
1115    pub fn next_block(&self) -> PartialTokenBlock {
1116        PartialTokenBlock {
1117            tokens: Tokens::with_capacity(self.tokens.len()),
1118            block_size: self.tokens.len() as u32, // Should be == self.block_size
1119            salt_hash: self.salt_hash,
1120            parent_sequence_hash: Some(self.sequence_hash), // Link to this block
1121            position: self.position() as usize + 1,         // Next position
1122        }
1123    }
1124
1125    /// Finalizes a [`TokenBlock`] from a [`TokenBlockChunk`], parent's sequence hash, and position.
1126    ///
1127    /// This computes the final [`SequenceHash`], [`PositionalSequenceHash`], and [`PositionalLineageHash`] for the block.
1128    fn from_chunk(
1129        chunk: TokenBlockChunk,
1130        parent_sequence_hash: Option<SequenceHash>,
1131        position: usize,
1132    ) -> Self {
1133        let sequence_hash = match parent_sequence_hash {
1134            Some(parent) => compute_next_sequence_hash(parent, chunk.block_hash),
1135            None => {
1136                // First block: sequence hash is just the block hash
1137                chunk.block_hash
1138            }
1139        };
1140
1141        let positional_sequence_hash = PositionalSequenceHash::new(
1142            sequence_hash,
1143            position as u64,
1144            chunk.block_hash, // LocalBlockHash is the same as BlockHash
1145        );
1146
1147        let positional_lineage_hash =
1148            PositionalLineageHash::new(sequence_hash, parent_sequence_hash, position as u64);
1149
1150        Self {
1151            tokens: chunk.tokens,
1152            salt_hash: chunk.salt_hash,
1153            block_hash: chunk.block_hash,
1154            sequence_hash,
1155            parent_sequence_hash,
1156            positional_sequence_hash,
1157            positional_lineage_hash,
1158        }
1159    }
1160
1161    /// Returns a reference to the tokens in this block.
1162    pub fn tokens(&self) -> &Tokens {
1163        &self.tokens
1164    }
1165
1166    /// Returns the salt hash used for this block's hashing.
1167    pub fn salt_hash(&self) -> SaltHash {
1168        self.salt_hash
1169    }
1170
1171    /// Returns the hash of only the tokens within this block.
1172    pub fn block_hash(&self) -> BlockHash {
1173        self.block_hash
1174    }
1175
1176    /// Returns the sequence-aware hash for this block.
1177    pub fn sequence_hash(&self) -> SequenceHash {
1178        self.sequence_hash
1179    }
1180
1181    /// Returns the sequence hash of the preceding block, if any.
1182    pub fn parent_sequence_hash(&self) -> Option<SequenceHash> {
1183        self.parent_sequence_hash
1184    }
1185
1186    /// Returns the number of tokens in the block.
1187    pub fn block_size(&self) -> usize {
1188        self.tokens.0.len()
1189    }
1190
1191    /// Returns the positional sequence hash for this block.
1192    pub fn positional_sequence_hash(&self) -> PositionalSequenceHash {
1193        self.positional_sequence_hash
1194    }
1195
1196    /// Returns the positional lineage hash for this block.
1197    pub fn positional_lineage_hash(&self) -> PositionalLineageHash {
1198        self.positional_lineage_hash
1199    }
1200
1201    /// Returns the position of this block in the sequence.
1202    pub fn position(&self) -> u64 {
1203        self.positional_sequence_hash.position()
1204    }
1205}
1206
1207impl PositionalHash for PositionalSequenceHash {
1208    fn position(&self) -> u64 {
1209        self.position()
1210    }
1211}
1212
1213impl PositionalHash for PositionalLineageHash {
1214    fn position(&self) -> u64 {
1215        self.position()
1216    }
1217}
1218
1219/// Represents a sequence of tokens, segmented into fixed-size, hashed blocks.
1220///
1221/// This structure manages a series of completed [`TokenBlock`]s and one
1222/// [`PartialTokenBlock`] for accumulating incoming tokens.
1223/// It provides methods for appending tokens (`append`, `extend`), removing tokens
1224/// (`pop`, `truncate`, `unwind`), and accessing sequence information.
1225///
1226/// Hashing incorporates an initial [`SaltHash`] to ensure uniqueness across different
1227/// contexts (e.g., different models, PEFTs).
1228///
1229/// Key Hashes:
1230/// - [`BlockHash`]: Hash of tokens within a single block (seeded by [`SaltHash`]).
1231/// - [`SequenceHash`]: Hash combining the previous block's [`SequenceHash`] and the current
1232///   block's [`BlockHash`] (also seeded by [`SaltHash`]).
1233#[derive(Debug, PartialEq)]
1234pub struct TokenBlockSequence {
1235    blocks: Vec<TokenBlock>,
1236    current_block: PartialTokenBlock,
1237    salt_hash: SaltHash,
1238    block_size: usize,
1239    /// Validated, sorted multimodal runs covering committed and partial slots.
1240    /// Empty for sequences built via the zero-MM constructors; populated by
1241    /// [`TokenBlockSequence::new_with_mm`] and the streaming MM helpers.
1242    mm_runs: Vec<TokenBlockMmInfo>,
1243}
1244
1245impl TokenBlockSequence {
1246    /// Creates a new [`TokenBlockSequence`] from an initial set of tokens.
1247    ///
1248    /// The tokens are split into blocks of `block_size`. Any remaining tokens
1249    /// form the initial `current_block`.
1250    ///
1251    /// # Arguments
1252    ///
1253    /// * `tokens` - The initial [`Tokens`] for the sequence.
1254    /// * `block_size` - The fixed size for each [`TokenBlock`]. Must be greater than 0.
1255    /// * `salt_hash` - An optional [`SaltHash`]. Defaults to 0 if `None`.
1256    ///
1257    /// # Panics
1258    ///
1259    /// Panics if `block_size` is 0.
1260    pub fn new(tokens: Tokens, block_size: u32, salt_hash: Option<SaltHash>) -> Self {
1261        assert!(block_size > 0, "block_size must be greater than 0");
1262        let salt_hash = salt_hash.unwrap_or_default();
1263        let (blocks, current_block) = Self::split_tokens(&tokens, block_size, salt_hash);
1264
1265        Self {
1266            blocks,
1267            current_block,
1268            salt_hash,
1269            block_size: block_size as usize,
1270            mm_runs: Vec::new(),
1271        }
1272    }
1273
1274    /// Extends the sequence with the given tokens, potentially completing multiple blocks.
1275    ///
1276    /// This method processes all tokens from the input [`Tokens`] object.
1277    /// If adding tokens causes one or more blocks to become full, they are committed
1278    /// and added to the internal list of completed blocks.
1279    ///
1280    /// # Arguments
1281    ///
1282    /// * `tokens` - The [`Tokens`] object containing the tokens to extend the sequence with.
1283    ///
1284    /// # Returns
1285    ///
1286    /// * `Ok(Some(Range<usize>))` - The range of indices in the `blocks` vector corresponding
1287    ///   to the blocks completed during this `extend` operation.
1288    /// * `Ok(None)` - If no blocks were completed.
1289    /// * `Err(TokenBlockError)` - If an internal error occurs during commit.
1290    pub fn extend(&mut self, tokens: Tokens) -> Result<Option<Range<usize>>, TokenBlockError> {
1291        let start_block_index = self.blocks.len();
1292        let mut tokens_to_append = tokens;
1293
1294        while !tokens_to_append.is_empty() {
1295            let remaining_in_current = self.current_block.remaining();
1296
1297            if remaining_in_current == 0 {
1298                // Current block is full, commit it first.
1299                let new_block = self.commit_current()?;
1300                self.blocks.push(new_block);
1301                // Continue loop to add tokens to the *new* current_block.
1302            }
1303
1304            // Push as many tokens as possible into the current (potentially new) block.
1305            let available_tokens = tokens_to_append;
1306            tokens_to_append = self.current_block.push_tokens(available_tokens);
1307
1308            // Check if the current block *became* full after pushing tokens.
1309            if self.current_block.remaining() == 0 {
1310                // If it became full AND there are still more tokens to append,
1311                // commit it now so the next loop iteration starts with a fresh block.
1312                let new_block = self.commit_current()?;
1313                self.blocks.push(new_block);
1314            }
1315        }
1316
1317        let end_block_index = self.blocks.len();
1318        if start_block_index == end_block_index {
1319            Ok(None) // No blocks were completed.
1320        } else {
1321            Ok(Some(start_block_index..end_block_index))
1322        }
1323    }
1324
1325    /// Commits the current partial block.
1326    ///
1327    /// Routes through the MM-aware byte encoding when [`Self::mm_runs`] is non-empty;
1328    /// otherwise behaves identically to [`PartialTokenBlock::commit`].
1329    fn commit_current(&mut self) -> Result<TokenBlock, TokenBlockError> {
1330        if self.mm_runs.is_empty() {
1331            return self.current_block.commit();
1332        }
1333        // MM-aware path: compute block_hash from the substituted byte buffer.
1334        if self.current_block.tokens.0.len() != self.current_block.block_size as usize {
1335            return Err(TokenBlockError::Incomplete);
1336        }
1337        let block_offset = self.blocks.len() * (self.current_block.block_size as usize);
1338        let tokens = std::mem::take(&mut self.current_block.tokens);
1339        let block_bytes = compute_block_bytes_with_mm(&tokens, block_offset, &self.mm_runs);
1340        let block_hash = compute_block_hash(&block_bytes, self.current_block.salt_hash);
1341        let chunk = TokenBlockChunk {
1342            tokens,
1343            salt_hash: self.current_block.salt_hash,
1344            block_hash,
1345        };
1346        let block = TokenBlock::from_chunk(
1347            chunk,
1348            self.current_block.parent_sequence_hash,
1349            self.current_block.position,
1350        );
1351        self.current_block.parent_sequence_hash = Some(block.sequence_hash());
1352        self.current_block.position += 1;
1353        Ok(block)
1354    }
1355
1356    /// Appends a single token to the sequence.
1357    ///
1358    /// If adding this token completes the current partial block, the block is committed,
1359    /// and the index of the newly completed block is returned.
1360    ///
1361    /// This method is equivalent to calling [`extend`] with a single-token [`Tokens`] object.
1362    ///
1363    /// # Arguments
1364    ///
1365    /// * `token` - The [`Token`] to append.
1366    ///
1367    /// # Returns
1368    ///
1369    /// * `Ok(Some(usize))` - The index of the block that was just completed.
1370    /// * `Ok(None)` - No block was completed by adding this token.
1371    /// * `Err(TokenBlockError)` - If an internal error occurs during processing.
1372    pub fn append(&mut self, token: Token) -> Result<Option<usize>, TokenBlockError> {
1373        let before = self.blocks.len();
1374        self.extend(Tokens::from(vec![token]))?;
1375        Ok(if self.blocks.len() > before {
1376            Some(before)
1377        } else {
1378            None
1379        })
1380    }
1381
1382    /// Shortens the sequence, keeping the first `len` tokens and removing the rest.
1383    ///
1384    /// If `len` is greater than the sequence's current length, this has no effect.
1385    ///
1386    /// This operation is analogous to `Vec::truncate`.
1387    /// It may involve removing tokens from the current partial block, removing entire
1388    /// completed blocks, and adjusting the current partial block
1389    /// to reflect the new end of the sequence.
1390    ///
1391    /// # Arguments
1392    ///
1393    /// * `len` - The number of tokens to keep.
1394    ///
1395    /// # Returns
1396    ///
1397    /// * `Ok(())` - If the sequence was successfully truncated.
1398    /// * `Err(TokenBlockError::InsufficientTokens)` - This error should ideally not occur if `len`
1399    ///   is correctly checked against `total_tokens`, but the underlying `pop_tokens` might return it.
1400    pub fn truncate(&mut self, len: usize) -> Result<(), TokenBlockError> {
1401        if !self.mm_runs.is_empty() {
1402            return Err(TokenBlockError::MmRunsPresent);
1403        }
1404        let current_total_len = self.total_tokens();
1405        if len >= current_total_len {
1406            return Ok(()); // Nothing to truncate
1407        }
1408
1409        let n = current_total_len - len; // Number of tokens to remove
1410
1411        // This inner block handles the actual removal logic based on `n` tokens to remove.
1412        {
1413            let current_len = self.current_block.len();
1414            // Avoid division by zero if block_size is somehow 0 (though asserted in new)
1415            let block_size = self.current_block.block_size.max(1);
1416
1417            if n <= current_len {
1418                // Only need to pop from the current partial block
1419                self.current_block.pop_tokens(n)?;
1420            } else {
1421                // Need to pop from full blocks as well
1422                let tokens_to_pop_from_blocks = n - current_len;
1423
1424                // Calculate how many blocks are affected (including the one partially popped)
1425                let num_blocks_to_affect = tokens_to_pop_from_blocks.div_ceil(block_size as usize);
1426
1427                // Check if we need to pop more blocks than available (should be prevented by initial len check)
1428                if num_blocks_to_affect > self.blocks.len() {
1429                    // This indicates an inconsistency between total_tokens() and internal state.
1430                    debug_assert!(
1431                        false,
1432                        "Truncate calculation error: trying to pop too many blocks."
1433                    );
1434                    return Err(TokenBlockError::InsufficientTokens);
1435                }
1436
1437                // Determine the index of the block that will be the source for the new partial block
1438                let source_block_index = self.blocks.len() - num_blocks_to_affect;
1439
1440                // Calculate how many tokens to keep from that source block
1441                let num_full_blocks_completely_popped = num_blocks_to_affect - 1;
1442                let num_tokens_to_pop_from_source_block = tokens_to_pop_from_blocks
1443                    - num_full_blocks_completely_popped * block_size as usize;
1444                let num_tokens_to_keep_in_new_partial =
1445                    (block_size as usize).saturating_sub(num_tokens_to_pop_from_source_block);
1446
1447                // Get the tokens for the new partial block
1448                let new_partial_tokens = if num_tokens_to_keep_in_new_partial > 0 {
1449                    self.blocks[source_block_index].tokens().as_ref()
1450                        [..num_tokens_to_keep_in_new_partial]
1451                        .to_vec()
1452                } else {
1453                    Vec::new()
1454                };
1455
1456                // Truncate the blocks vector to remove popped blocks
1457                self.blocks.truncate(source_block_index);
1458
1459                // Update the current_block state
1460                self.current_block.tokens = Tokens(new_partial_tokens);
1461                // Correctly set the parent hash based on the *new* last block
1462                self.current_block.parent_sequence_hash =
1463                    self.blocks.last().map(|b| b.sequence_hash());
1464                // Update position to match the number of complete blocks
1465                self.current_block.position = self.blocks.len();
1466                // salt_hash and block_size remain the same for current_block
1467            }
1468        }
1469        Ok(())
1470    }
1471
1472    /// Removes the last `count` tokens from the sequence.
1473    ///
1474    /// This is a convenience method that calculates the required length and calls [`truncate`].
1475    ///
1476    /// # Arguments
1477    ///
1478    /// * `count` - The number of tokens to remove from the end.
1479    ///
1480    /// # Returns
1481    ///
1482    /// * `Ok(())` - If the tokens were successfully removed.
1483    /// * `Err(TokenBlockError::InsufficientTokens)` - If `count` is greater than or equal to
1484    ///   the total number of tokens in the sequence.
1485    pub fn unwind(&mut self, count: usize) -> Result<(), TokenBlockError> {
1486        let current_total_len = self.total_tokens();
1487        if count > current_total_len {
1488            // Allow count == current_total_len, which truncates to 0.
1489            return Err(TokenBlockError::InsufficientTokens);
1490        }
1491
1492        // number of tokens remaining in the sequence after undoing the given count
1493        let len = current_total_len - count;
1494        self.truncate(len)
1495    }
1496
1497    /// Resets the sequence to the initial state.
1498    ///
1499    /// Clears any accumulated multimodal runs; after `reset` the sequence behaves
1500    /// identically to a freshly-constructed zero-MM sequence with the same `salt_hash`
1501    /// and `block_size`.
1502    pub fn reset(&mut self) {
1503        self.blocks.clear();
1504        self.current_block =
1505            PartialTokenBlock::create_sequence_root(self.block_size as u32, self.salt_hash);
1506        self.mm_runs.clear();
1507    }
1508
1509    /// Removes the last token from the sequence and returns it, or [`None`] if it is empty.
1510    ///
1511    /// This operation is analogous to `Vec::pop`.
1512    ///
1513    /// # Returns
1514    ///
1515    /// * `Some(Token)` - The last token, if the sequence was not empty.
1516    /// * `None` - If the sequence was empty.
1517    ///
1518    /// # Panics
1519    ///
1520    /// Panics if the sequence has accumulated multimodal runs (see [`Self::try_pop`] for a
1521    /// non-panicking variant). `pop` returns `Option<Token>` and cannot signal
1522    /// "operation unsupported on MM sequence" through its return type without breaking the
1523    /// `Vec::pop` analogy and silently lying to callers.
1524    pub fn pop(&mut self) -> Option<Token> {
1525        if !self.mm_runs.is_empty() {
1526            panic!(
1527                "TokenBlockSequence::pop is not supported on a sequence with multimodal runs; \
1528                 use try_pop or reset before pop"
1529            );
1530        }
1531        let current_total_len = self.total_tokens();
1532        if current_total_len == 0 {
1533            return None;
1534        }
1535
1536        // Determine the last token. It must be in the current_block if current_block is not empty.
1537        // If current_block is empty, it must be the last token of the last full block.
1538        let last_token = if !self.current_block.tokens.is_empty() {
1539            // Last token is in the partial block
1540            *self
1541                .current_block
1542                .tokens
1543                .last()
1544                .expect("Current block checked for non-empty")
1545        } else {
1546            // Current block is empty, sequence is not. Must be in the last full block.
1547            let last_block = self
1548                .blocks
1549                .last()
1550                .expect("Sequence is not empty but has no blocks and empty current block?");
1551            *last_block
1552                .tokens()
1553                .last()
1554                .expect("Last block cannot be empty")
1555        };
1556
1557        // Truncate the sequence by one element.
1558        // We expect this to succeed since we know the length > 0.
1559        match self.truncate(current_total_len - 1) {
1560            Ok(_) => Some(last_token),
1561            Err(_) => {
1562                // This should be logically impossible if total_tokens() and truncate() are correct.
1563                // Panic in debug, return None in release as a fallback, though it indicates a bug.
1564                debug_assert!(
1565                    false,
1566                    "truncate failed unexpectedly after checking length in pop"
1567                );
1568                None
1569            }
1570        }
1571    }
1572
1573    /// Non-panicking variant of [`Self::pop`].
1574    ///
1575    /// Returns:
1576    /// - `Ok(Some(token))` when the sequence had at least one token and pop succeeded.
1577    /// - `Ok(None)` when the sequence was empty.
1578    /// - `Err(TokenBlockError::MmRunsPresent)` when the sequence has multimodal runs.
1579    pub fn try_pop(&mut self) -> Result<Option<Token>, TokenBlockError> {
1580        if !self.mm_runs.is_empty() {
1581            return Err(TokenBlockError::MmRunsPresent);
1582        }
1583        Ok(self.pop())
1584    }
1585
1586    /// Returns a slice containing all the completed [`TokenBlock`]s in the sequence.
1587    pub fn blocks(&self) -> &[TokenBlock] {
1588        &self.blocks
1589    }
1590
1591    /// Returns a reference to the last completed [`TokenBlock`] in the sequence, if any.
1592    pub fn last_complete_block(&self) -> Option<&TokenBlock> {
1593        self.blocks.last()
1594    }
1595
1596    /// Returns a reference to the current [`PartialTokenBlock`] where new tokens are added.
1597    pub fn current_block(&self) -> &PartialTokenBlock {
1598        &self.current_block
1599    }
1600
1601    /// Consumes the sequence and returns its parts: a `Vec` of completed blocks and the final partial block.
1602    pub fn into_parts(self) -> (Vec<TokenBlock>, PartialTokenBlock) {
1603        (self.blocks, self.current_block)
1604    }
1605
1606    /// Returns the block size used for this sequence.
1607    pub fn block_size(&self) -> usize {
1608        self.block_size
1609    }
1610
1611    /// Returns the [`SaltHash`] used for this sequence.
1612    pub fn salt_hash(&self) -> SaltHash {
1613        self.salt_hash
1614    }
1615
1616    /// Returns the total number of tokens in the sequence (sum of tokens in all completed blocks
1617    /// plus tokens in the current partial block).
1618    pub fn total_tokens(&self) -> usize {
1619        let block_size = self.current_block.block_size as usize;
1620        (self.blocks.len() * block_size) + self.current_block.len()
1621    }
1622
1623    /// Extract the token with the range
1624    pub fn tokens_at(&self, range: Range<usize>) -> Tokens {
1625        let total = self.total_tokens();
1626
1627        // Validate range - return empty tokens for invalid ranges
1628        if range.start > range.end || range.end > total {
1629            return Tokens::default();
1630        }
1631
1632        // Handle empty range
1633        if range.is_empty() {
1634            return Tokens::default();
1635        }
1636
1637        let mut result = Vec::with_capacity(range.len());
1638
1639        for i in range {
1640            if i < self.blocks.len() * self.block_size {
1641                // Token is in a completed block
1642                let block_index = i / self.block_size;
1643                let token_index = i % self.block_size;
1644                result.push(self.blocks[block_index].tokens()[token_index]);
1645            } else {
1646                // Token is in the current partial block
1647                let current_block_index = i - (self.blocks.len() * self.block_size);
1648                result.push(self.current_block.tokens()[current_block_index]);
1649            }
1650        }
1651
1652        Tokens::from(result)
1653    }
1654
1655    /// Splits a [`Tokens`] object into a vector of completed blocks and a final partial block.
1656    ///
1657    /// This is primarily used internally by [`TokenBlockSequence::new`] but can be used externally.
1658    ///
1659    /// # Arguments
1660    ///
1661    /// * `tokens` - The [`Tokens`] to split.
1662    /// * `block_size` - The size of each block.
1663    /// * `salt_hash` - The [`SaltHash`] to use for hashing.
1664    ///
1665    /// # Returns
1666    ///
1667    /// A tuple containing `(Vec<TokenBlock>, PartialTokenBlock)`.
1668    ///
1669    /// # Panics
1670    ///
1671    /// Panics if `block_size` is 0.
1672    pub fn split_tokens(
1673        tokens: &[Token],
1674        block_size: u32,
1675        salt_hash: SaltHash,
1676    ) -> (Vec<TokenBlock>, PartialTokenBlock) {
1677        assert!(block_size > 0, "block_size must be greater than 0");
1678        let chunks: Vec<TokenBlockChunk> = tokens
1679            .as_ref()
1680            .chunks_exact(block_size as usize)
1681            .map(|chunk| TokenBlockChunk::from_tokens(chunk, salt_hash))
1682            .collect();
1683
1684        let mut result_blocks = Vec::with_capacity(chunks.len());
1685        let mut last_sequence_hash: Option<SequenceHash> = None;
1686
1687        // Sequentially combine chunks to compute sequence hashes
1688        for (position, chunk) in chunks.into_iter().enumerate() {
1689            let new_block = TokenBlock::from_chunk(chunk, last_sequence_hash, position);
1690            last_sequence_hash = Some(new_block.sequence_hash());
1691            result_blocks.push(new_block);
1692        }
1693
1694        // Handle any remaining tokens
1695        let remainder = tokens
1696            .as_ref()
1697            .chunks_exact(block_size as usize)
1698            .remainder();
1699
1700        let next_position = result_blocks.len(); // Position for the next block to be committed
1701
1702        let mut partial_tokens = Tokens::with_capacity(block_size as usize);
1703        partial_tokens.0.extend_from_slice(remainder);
1704
1705        let current_block = PartialTokenBlock {
1706            tokens: partial_tokens,
1707            block_size,
1708            salt_hash,
1709            // Parent hash is the sequence hash of the last *full* block computed
1710            parent_sequence_hash: last_sequence_hash,
1711            position: next_position,
1712        };
1713
1714        (result_blocks, current_block)
1715    }
1716
1717    /// Creates a new [`TokenBlockSequence`] from a slice of tokens.
1718    ///
1719    /// The tokens are split into blocks of `block_size`. Any remaining tokens
1720    /// form the initial `current_block`.
1721    ///
1722    /// # Arguments
1723    ///
1724    /// * `tokens` - The slice of tokens to create the sequence from.
1725    /// * `block_size` - The size of each block.
1726    /// * `salt_hash` - The [`SaltHash`] to use for hashing.
1727    pub fn from_slice(tokens: &[Token], block_size: u32, salt_hash: Option<SaltHash>) -> Self {
1728        assert!(block_size > 0, "block_size must be greater than 0");
1729        let salt_hash = salt_hash.unwrap_or_default();
1730        let (blocks, current_block) = Self::split_tokens(tokens, block_size, salt_hash);
1731
1732        Self {
1733            blocks,
1734            current_block,
1735            salt_hash,
1736            block_size: block_size as usize,
1737            mm_runs: Vec::new(),
1738        }
1739    }
1740
1741    /// Creates a [`TokenBlockSequence`] with multimodal placeholder runs.
1742    ///
1743    /// `mm_info` is validated and sorted via [`validate_and_sort_mm_info`]. Each block's
1744    /// [`BlockHash`] is computed using the per-block byte encoding documented on
1745    /// [`compute_block_bytes_with_mm`], which selects one of two encodings:
1746    ///
1747    /// - **MM-affected block** (at least one run overlaps the block): every slot emits
1748    ///   13 bytes — a 1-byte tag ([`MM_SLOT_TAG_TOKEN`] or [`MM_SLOT_TAG_PLACEHOLDER`])
1749    ///   followed by a 12-byte payload. Real-token slots carry `token_id u32 LE` plus
1750    ///   8 bytes of padding; placeholder slots carry `run_offset u32 LE | mm_hash u64 LE`.
1751    /// - **Non-MM block** (no run overlaps): the legacy `bytemuck::cast_slice(tokens)`
1752    ///   form is used (4 bytes per slot, LE u32), preserving cache identity with blocks
1753    ///   produced by [`Self::from_slice`].
1754    ///
1755    /// Returns an error if `mm_info` is invalid (overlap, out of bounds, zero-length run).
1756    pub fn new_with_mm(
1757        tokens: Tokens,
1758        mm_info: &[TokenBlockMmInfo],
1759        block_size: u32,
1760        salt_hash: Option<SaltHash>,
1761    ) -> Result<Self, TokenBlockError> {
1762        assert!(block_size > 0, "block_size must be greater than 0");
1763        let salt_hash = salt_hash.unwrap_or_default();
1764        let validated =
1765            validate_and_sort_mm_info(mm_info, tokens.len()).map_err(TokenBlockError::MmInfo)?;
1766        let (blocks, current_block) =
1767            Self::split_tokens_with_mm(&tokens, &validated, block_size, salt_hash);
1768        Ok(Self {
1769            blocks,
1770            current_block,
1771            salt_hash,
1772            block_size: block_size as usize,
1773            mm_runs: validated,
1774        })
1775    }
1776
1777    /// MM-aware variant of [`Self::split_tokens`].
1778    ///
1779    /// `mm_runs` must be pre-validated and sorted (e.g., via [`validate_and_sort_mm_info`]).
1780    pub fn split_tokens_with_mm(
1781        tokens: &[Token],
1782        mm_runs: &[TokenBlockMmInfo],
1783        block_size: u32,
1784        salt_hash: SaltHash,
1785    ) -> (Vec<TokenBlock>, PartialTokenBlock) {
1786        assert!(block_size > 0, "block_size must be greater than 0");
1787        let bs = block_size as usize;
1788        let n_complete = tokens.len() / bs;
1789        let mut result_blocks = Vec::with_capacity(n_complete);
1790        let mut last_seq_hash: Option<SequenceHash> = None;
1791        for i in 0..n_complete {
1792            let block_offset = i * bs;
1793            let block_tokens = &tokens[block_offset..block_offset + bs];
1794            let block_bytes = compute_block_bytes_with_mm(block_tokens, block_offset, mm_runs);
1795            let block_hash = compute_block_hash(&block_bytes, salt_hash);
1796            let chunk = TokenBlockChunk {
1797                tokens: block_tokens.into(),
1798                salt_hash,
1799                block_hash,
1800            };
1801            let new_block = TokenBlock::from_chunk(chunk, last_seq_hash, i);
1802            last_seq_hash = Some(new_block.sequence_hash());
1803            result_blocks.push(new_block);
1804        }
1805        let remainder = &tokens[n_complete * bs..];
1806        let current_block = PartialTokenBlock {
1807            tokens: remainder.into(),
1808            block_size,
1809            salt_hash,
1810            parent_sequence_hash: last_seq_hash,
1811            position: n_complete,
1812        };
1813        (result_blocks, current_block)
1814    }
1815
1816    /// Returns the validated, sorted multimodal runs accumulated by this sequence.
1817    pub fn mm_runs(&self) -> &[TokenBlockMmInfo] {
1818        &self.mm_runs
1819    }
1820
1821    /// Appends a single real token to the sequence.
1822    ///
1823    /// Equivalent to [`Self::append`] but named for symmetry with [`Self::push_mm_run`].
1824    pub fn push_token(&mut self, token: Token) -> Result<Option<usize>, TokenBlockError> {
1825        self.append(token)
1826    }
1827
1828    /// Appends a multimodal placeholder run of `length` slots all tagged with `mm_hash`.
1829    ///
1830    /// The placeholder run starts at the current end of the sequence (`total_tokens()` before
1831    /// the call). The token IDs at placeholder slot positions are filled with zero sentinels;
1832    /// hashing uses the `(mm_hash, run_offset)` pair instead of those token bytes.
1833    ///
1834    /// Returns the range of fully-committed block indices completed during the call (if any).
1835    pub fn push_mm_run(
1836        &mut self,
1837        mm_hash: u64,
1838        length: usize,
1839    ) -> Result<Option<Range<usize>>, TokenBlockError> {
1840        if length == 0 {
1841            return Err(TokenBlockError::MmInfo(MmInfoError::EmptyRun));
1842        }
1843        let offset = self.total_tokens();
1844        self.mm_runs.push(TokenBlockMmInfo {
1845            mm_hash,
1846            offset,
1847            length,
1848        });
1849        // The token values at placeholder slots are opaque for hashing; use 0 sentinels.
1850        let placeholders = Tokens::from(vec![0u32; length]);
1851        self.extend(placeholders)
1852    }
1853
1854    /// Batch-validated extension: appends `tokens` (with embedded multimodal runs in `mm_info`)
1855    /// to the sequence in a single validated step.
1856    ///
1857    /// `mm_info` offsets are **relative to the start of `tokens`** (not the existing sequence).
1858    /// The function validates the chunk, then translates to absolute offsets and applies the
1859    /// updates atomically: it errors before mutating any state if `mm_info` is invalid.
1860    ///
1861    /// Real-token regions and placeholder runs are interleaved per the chunk's layout.
1862    pub fn extend_with_mm(
1863        &mut self,
1864        tokens: &[Token],
1865        mm_info: &[TokenBlockMmInfo],
1866    ) -> Result<Option<Range<usize>>, TokenBlockError> {
1867        let validated =
1868            validate_and_sort_mm_info(mm_info, tokens.len()).map_err(TokenBlockError::MmInfo)?;
1869        let start_block = self.blocks.len();
1870        let mut cursor = 0usize;
1871        for run in &validated {
1872            if run.offset > cursor {
1873                let real = Tokens::from(tokens[cursor..run.offset].to_vec());
1874                self.extend(real)?;
1875            }
1876            self.push_mm_run(run.mm_hash, run.length)?;
1877            cursor = run.offset + run.length;
1878        }
1879        if cursor < tokens.len() {
1880            let real = Tokens::from(tokens[cursor..].to_vec());
1881            self.extend(real)?;
1882        }
1883        let end_block = self.blocks.len();
1884        if start_block == end_block {
1885            Ok(None)
1886        } else {
1887            Ok(Some(start_block..end_block))
1888        }
1889    }
1890}
1891
1892#[cfg(test)]
1893mod tests {
1894    use super::*;
1895    use bytemuck::cast_slice;
1896
1897    // Helper to create a sequence for testing
1898    fn create_test_sequence(
1899        initial_tokens: &[Token],
1900        block_size: u32,
1901        salt_hash: Option<SaltHash>,
1902    ) -> TokenBlockSequence {
1903        TokenBlockSequence::new(Tokens::from(initial_tokens), block_size, salt_hash)
1904    }
1905
1906    // Helper to get expected hashes (replace with actual calculated values if needed)
1907    const TEST_SALT_HASH: SaltHash = 1337;
1908    const HASH_1_4: BlockHash = 14643705804678351452; // hash([1,2,3,4], 1337)
1909    const SEQ_HASH_1_4: SequenceHash = HASH_1_4;
1910    const HASH_5_8: BlockHash = 16777012769546811212; // hash([5,6,7,8], 1337)
1911    const SEQ_HASH_5_8: SequenceHash = 4945711292740353085; // hash([SEQ_HASH_1_4, HASH_5_8], CHAIN_XXH3_SEED)
1912    const HASH_9_12: BlockHash = 483935686894639516; // hash([9,10,11,12], 1337)
1913    const SEQ_HASH_9_12: SequenceHash = 12583592247330656132; // hash([SEQ_HASH_5_8, HASH_9_12], CHAIN_XXH3_SEED)
1914
1915    #[test]
1916    fn token_hash_helper_matches_canonical_byte_encoding() {
1917        let tokens = [1u32, 2, 3, 4];
1918        assert_eq!(
1919            compute_block_hash_for_tokens(&tokens, TEST_SALT_HASH),
1920            compute_block_hash(cast_slice(&tokens), TEST_SALT_HASH)
1921        );
1922        assert_eq!(
1923            compute_block_hash_for_tokens(&tokens, TEST_SALT_HASH),
1924            HASH_1_4
1925        );
1926    }
1927
1928    impl PartialTokenBlock {
1929        /// Attempts to remove the last token from the block.
1930        ///
1931        /// # Returns
1932        ///
1933        /// * `Ok(())` - If a token was successfully removed.
1934        /// * `Err(TokenBlockError::Empty)` - If the block was already empty.
1935        pub fn pop_token(&mut self) -> Result<(), TokenBlockError> {
1936            if self.tokens.0.is_empty() {
1937                return Err(TokenBlockError::Empty);
1938            }
1939            self.tokens.0.pop();
1940            Ok(())
1941        }
1942    }
1943
1944    #[test]
1945    fn test_validate_hash_constants() {
1946        let salt = TEST_SALT_HASH;
1947
1948        // Block 1: [1, 2, 3, 4]
1949        let tokens_1_4 = &[1u32, 2, 3, 4];
1950        let computed_hash_1_4 = compute_block_hash(cast_slice(tokens_1_4), salt);
1951        assert_eq!(computed_hash_1_4, HASH_1_4, "Mismatch for HASH_1_4");
1952        // First block's sequence hash is its block hash
1953        assert_eq!(computed_hash_1_4, SEQ_HASH_1_4, "Mismatch for SEQ_HASH_1_4");
1954
1955        // Block 2: [5, 6, 7, 8]
1956        let tokens_5_8 = &[5u32, 6, 7, 8];
1957        let computed_hash_5_8 = compute_block_hash(cast_slice(tokens_5_8), salt);
1958        assert_eq!(computed_hash_5_8, HASH_5_8, "Mismatch for HASH_5_8");
1959        // Chain step uses the shared CHAIN_XXH3_SEED; salt is already mixed into block_hash.
1960        let computed_seq_hash_5_8 = compute_next_sequence_hash(SEQ_HASH_1_4, HASH_5_8);
1961        assert_eq!(
1962            computed_seq_hash_5_8, SEQ_HASH_5_8,
1963            "Mismatch for SEQ_HASH_5_8"
1964        );
1965
1966        // Block 3: [9, 10, 11, 12]
1967        let tokens_9_12 = &[9u32, 10, 11, 12];
1968        let computed_hash_9_12 = compute_block_hash(cast_slice(tokens_9_12), salt);
1969        assert_eq!(computed_hash_9_12, HASH_9_12, "Mismatch for HASH_9_12");
1970        let computed_seq_hash_9_12 = compute_next_sequence_hash(SEQ_HASH_5_8, HASH_9_12);
1971        assert_eq!(
1972            computed_seq_hash_9_12, SEQ_HASH_9_12,
1973            "Mismatch for SEQ_HASH_9_12"
1974        );
1975    }
1976
1977    #[test]
1978    fn test_positional_sequence_hash_encoding_decoding() {
1979        // Test Mode 0: position fits in 8 bits (< 256)
1980        let seq_hash_0 = 0x1234567890ABCDEF;
1981        let position_0 = 100;
1982        let lbh_0 = 0xFEDCBA9876543210;
1983        let psh_0 = PositionalSequenceHash::new(seq_hash_0, position_0, lbh_0);
1984
1985        assert_eq!(psh_0.mode(), 0, "Position 100 should use mode 0");
1986        assert_eq!(psh_0.sequence_hash(), seq_hash_0);
1987        assert_eq!(psh_0.position(), position_0);
1988        // LBH is truncated to 54 bits in mode 0
1989        assert_eq!(
1990            psh_0.local_block_hash(),
1991            lbh_0 & ((1u64 << 54) - 1),
1992            "LBH should be truncated to 54 bits"
1993        );
1994
1995        // Test Mode 1: position fits in 16 bits (256 <= pos < 65536)
1996        let position_1 = 1000;
1997        let psh_1 = PositionalSequenceHash::new(seq_hash_0, position_1, lbh_0);
1998
1999        assert_eq!(psh_1.mode(), 1, "Position 1000 should use mode 1");
2000        assert_eq!(psh_1.sequence_hash(), seq_hash_0);
2001        assert_eq!(psh_1.position(), position_1);
2002        // LBH is truncated to 46 bits in mode 1
2003        assert_eq!(
2004            psh_1.local_block_hash(),
2005            lbh_0 & ((1u64 << 46) - 1),
2006            "LBH should be truncated to 46 bits"
2007        );
2008
2009        // Test Mode 2: position fits in 24 bits (65536 <= pos < 16777216)
2010        let position_2 = 100_000;
2011        let psh_2 = PositionalSequenceHash::new(seq_hash_0, position_2, lbh_0);
2012
2013        assert_eq!(psh_2.mode(), 2, "Position 100,000 should use mode 2");
2014        assert_eq!(psh_2.sequence_hash(), seq_hash_0);
2015        assert_eq!(psh_2.position(), position_2);
2016        // LBH is truncated to 38 bits in mode 2
2017        assert_eq!(
2018            psh_2.local_block_hash(),
2019            lbh_0 & ((1u64 << 38) - 1),
2020            "LBH should be truncated to 38 bits"
2021        );
2022
2023        // Test Mode 3: position fits in 31 bits (16777216 <= pos < 2^31)
2024        let position_3 = 20_000_000;
2025        let psh_3 = PositionalSequenceHash::new(seq_hash_0, position_3, lbh_0);
2026
2027        assert_eq!(psh_3.mode(), 3, "Position 20,000,000 should use mode 3");
2028        assert_eq!(psh_3.sequence_hash(), seq_hash_0);
2029        assert_eq!(psh_3.position(), position_3);
2030        // LBH is truncated to 31 bits in mode 3
2031        assert_eq!(
2032            psh_3.local_block_hash(),
2033            lbh_0 & ((1u64 << 31) - 1),
2034            "LBH should be truncated to 31 bits"
2035        );
2036
2037        // Test edge case: position at boundary
2038        let position_255 = 255;
2039        let psh_255 = PositionalSequenceHash::new(seq_hash_0, position_255, lbh_0);
2040        assert_eq!(psh_255.mode(), 0, "Position 255 should use mode 0");
2041        assert_eq!(psh_255.position(), position_255);
2042
2043        let position_256 = 256;
2044        let psh_256 = PositionalSequenceHash::new(seq_hash_0, position_256, lbh_0);
2045        assert_eq!(psh_256.mode(), 1, "Position 256 should use mode 1");
2046        assert_eq!(psh_256.position(), position_256);
2047    }
2048
2049    #[test]
2050    fn test_positional_lineage_hash() {
2051        // Test Mode 0: position fits in 8 bits (< 256)
2052        let current_hash_0 = 0x1234567890ABCDEF;
2053        let parent_hash_0 = 0xFEDCBA9876543210;
2054        let position_0 = 100;
2055        let plh_0 = PositionalLineageHash::new(current_hash_0, Some(parent_hash_0), position_0);
2056
2057        assert_eq!(plh_0.mode(), 0, "Position 100 should use mode 0");
2058        assert_eq!(plh_0.position(), position_0);
2059        // Current carries the full u64; parent fragment is 54 bits in mode 0
2060        assert_eq!(
2061            plh_0.current_sequence_hash(),
2062            current_hash_0,
2063            "Current sequence hash should be stored in full"
2064        );
2065        assert_eq!(
2066            plh_0.parent_hash_fragment(),
2067            parent_hash_0 & ((1u64 << 54) - 1),
2068            "Parent fragment should be truncated to 54 bits in mode 0"
2069        );
2070
2071        // Test Mode 1: position fits in 16 bits (256 <= pos < 65536)
2072        let position_1 = 1000;
2073        let plh_1 = PositionalLineageHash::new(current_hash_0, Some(parent_hash_0), position_1);
2074
2075        assert_eq!(plh_1.mode(), 1, "Position 1000 should use mode 1");
2076        assert_eq!(plh_1.position(), position_1);
2077        assert_eq!(plh_1.current_sequence_hash(), current_hash_0);
2078        assert_eq!(
2079            plh_1.parent_hash_fragment(),
2080            parent_hash_0 & ((1u64 << 46) - 1),
2081            "Parent fragment should be truncated to 46 bits in mode 1"
2082        );
2083
2084        // Test Mode 2: position fits in 24 bits (65536 <= pos < 16777216)
2085        let position_2 = 100_000;
2086        let plh_2 = PositionalLineageHash::new(current_hash_0, Some(parent_hash_0), position_2);
2087
2088        assert_eq!(plh_2.mode(), 2, "Position 100,000 should use mode 2");
2089        assert_eq!(plh_2.position(), position_2);
2090        assert_eq!(plh_2.current_sequence_hash(), current_hash_0);
2091        assert_eq!(
2092            plh_2.parent_hash_fragment(),
2093            parent_hash_0 & ((1u64 << 38) - 1),
2094            "Parent fragment should be truncated to 38 bits in mode 2"
2095        );
2096
2097        // Test edge cases: position at boundaries
2098        let position_255 = 255;
2099        let plh_255 = PositionalLineageHash::new(current_hash_0, Some(parent_hash_0), position_255);
2100        assert_eq!(plh_255.mode(), 0, "Position 255 should use mode 0");
2101        assert_eq!(plh_255.position(), position_255);
2102
2103        let position_256 = 256;
2104        let plh_256 = PositionalLineageHash::new(current_hash_0, Some(parent_hash_0), position_256);
2105        assert_eq!(plh_256.mode(), 1, "Position 256 should use mode 1");
2106        assert_eq!(plh_256.position(), position_256);
2107
2108        let position_65535 = 65535;
2109        let plh_65535 =
2110            PositionalLineageHash::new(current_hash_0, Some(parent_hash_0), position_65535);
2111        assert_eq!(plh_65535.mode(), 1, "Position 65535 should use mode 1");
2112        assert_eq!(plh_65535.position(), position_65535);
2113
2114        let position_65536 = 65536;
2115        let plh_65536 =
2116            PositionalLineageHash::new(current_hash_0, Some(parent_hash_0), position_65536);
2117        assert_eq!(plh_65536.mode(), 2, "Position 65536 should use mode 2");
2118        assert_eq!(plh_65536.position(), position_65536);
2119
2120        // Test with None parent (root block)
2121        let plh_root = PositionalLineageHash::new(current_hash_0, None, 0);
2122        assert_eq!(plh_root.mode(), 0);
2123        assert_eq!(plh_root.position(), 0);
2124        assert_eq!(
2125            plh_root.parent_hash_fragment(),
2126            0,
2127            "Root should have zero parent fragment"
2128        );
2129        assert_eq!(plh_root.current_sequence_hash(), current_hash_0);
2130    }
2131
2132    #[test]
2133    #[should_panic(expected = "Position 16777216 exceeds maximum supported value")]
2134    fn test_positional_lineage_hash_panic_on_large_position() {
2135        let current_hash = 0x1234567890ABCDEF;
2136        let parent_hash = 0xFEDCBA9876543210;
2137        let position = 1u64 << 24; // 2^24 = 16,777,216
2138        let _ = PositionalLineageHash::new(current_hash, Some(parent_hash), position);
2139    }
2140
2141    #[test]
2142    fn test_positional_lineage_hash_mode_boundary_alignment() {
2143        // Test that backward radix traversal still works across mode boundaries.
2144        // Under the asymmetric layout, the parent fragment is sized to the *child's*
2145        // mode and pulled from the parent's full u64 via
2146        // `parent_fragment_for_child_position`.
2147
2148        let parent_hash = 0xFEDCBA9876543210;
2149        let current_hash_255 = 0x1234567890ABCDEF;
2150        let current_hash_256 = 0xABCDEF0123456789;
2151
2152        // Position 255: Mode 0 (last position before boundary)
2153        let plh_255 = PositionalLineageHash::new(current_hash_255, Some(parent_hash), 255);
2154        assert_eq!(plh_255.mode(), 0);
2155        assert_eq!(plh_255.current_sequence_hash(), current_hash_255);
2156
2157        // Position 256: Mode 1 (first position after boundary)
2158        let plh_256 = PositionalLineageHash::new(current_hash_256, Some(current_hash_255), 256);
2159        assert_eq!(plh_256.mode(), 1);
2160
2161        // CRITICAL: child's parent_fragment must equal parent's current truncated to
2162        // the child's mode width (46 bits at mode 1).
2163        let mask_46 = (1u64 << 46) - 1;
2164        assert_eq!(
2165            plh_256.parent_hash_fragment(),
2166            current_hash_255 & mask_46,
2167            "Mode boundary: position 256's parent fragment matches position 255's current truncated to 46 bits"
2168        );
2169        assert_eq!(
2170            plh_256.parent_hash_fragment(),
2171            plh_255.parent_fragment_for_child_position(256),
2172            "parent_fragment_for_child_position helper should match"
2173        );
2174
2175        // Test the other boundary: 65535 -> 65536 (Mode 1 -> Mode 2)
2176        let current_hash_65535 = 0x1111222233334444;
2177        let current_hash_65536 = 0x5555666677778888;
2178
2179        let plh_65535 = PositionalLineageHash::new(current_hash_65535, Some(parent_hash), 65535);
2180        assert_eq!(plh_65535.mode(), 1);
2181
2182        let plh_65536 =
2183            PositionalLineageHash::new(current_hash_65536, Some(current_hash_65535), 65536);
2184        assert_eq!(plh_65536.mode(), 2);
2185
2186        // Mode 2 parent fragment is 38 bits.
2187        let mask_38 = (1u64 << 38) - 1;
2188        assert_eq!(
2189            plh_65536.parent_hash_fragment(),
2190            current_hash_65535 & mask_38,
2191            "Mode boundary: position 65536's parent fragment matches position 65535's current truncated to 38 bits"
2192        );
2193        assert_eq!(
2194            plh_65536.parent_hash_fragment(),
2195            plh_65535.parent_fragment_for_child_position(65536),
2196        );
2197    }
2198
2199    #[test]
2200    fn test_positional_lineage_hash_extend() {
2201        // PLH must be self-extending: a chain built from PLH::root + extend should be
2202        // bitwise identical to the chain produced by full TokenBlock construction.
2203        let salt: SaltHash = 1337;
2204        let bh: [BlockHash; 3] = [
2205            compute_block_hash(cast_slice(&[1u32, 2, 3, 4]), salt),
2206            compute_block_hash(cast_slice(&[5u32, 6, 7, 8]), salt),
2207            compute_block_hash(cast_slice(&[9u32, 10, 11, 12]), salt),
2208        ];
2209
2210        // Direct construction via TokenBlock
2211        let blk0 =
2212            TokenBlock::from_chunk(TokenBlockChunk::from_tokens(&[1, 2, 3, 4], salt), None, 0);
2213        let blk1 = TokenBlock::from_chunk(
2214            TokenBlockChunk::from_tokens(&[5, 6, 7, 8], salt),
2215            Some(blk0.sequence_hash()),
2216            1,
2217        );
2218        let blk2 = TokenBlock::from_chunk(
2219            TokenBlockChunk::from_tokens(&[9, 10, 11, 12], salt),
2220            Some(blk1.sequence_hash()),
2221            2,
2222        );
2223
2224        // Self-extending construction via PLH::root + extend
2225        let plh0 = PositionalLineageHash::root(bh[0]);
2226        let plh1 = plh0.extend(bh[1]);
2227        let plh2 = plh1.extend(bh[2]);
2228
2229        assert_eq!(plh0.as_u128(), blk0.positional_lineage_hash().as_u128());
2230        assert_eq!(plh1.as_u128(), blk1.positional_lineage_hash().as_u128());
2231        assert_eq!(plh2.as_u128(), blk2.positional_lineage_hash().as_u128());
2232
2233        // Chain definition: extended.current_sequence_hash == compute_next_sequence_hash(parent, child)
2234        assert_eq!(
2235            plh1.current_sequence_hash(),
2236            compute_next_sequence_hash(plh0.current_sequence_hash(), bh[1]),
2237        );
2238
2239        // Salt propagation: changing salt changes block_hash[0] and therefore every
2240        // subsequent PLH, even though salt no longer seeds the per-step chain.
2241        let alt_salt: SaltHash = 4242;
2242        let alt_bh0 = compute_block_hash(cast_slice(&[1u32, 2, 3, 4]), alt_salt);
2243        assert_ne!(alt_bh0, bh[0]);
2244        let alt_plh0 = PositionalLineageHash::root(alt_bh0);
2245        let alt_plh1 = alt_plh0.extend(compute_block_hash(cast_slice(&[5u32, 6, 7, 8]), alt_salt));
2246        assert_ne!(alt_plh0.as_u128(), plh0.as_u128());
2247        assert_ne!(alt_plh1.as_u128(), plh1.as_u128());
2248    }
2249
2250    #[test]
2251    fn test_tokens_from() {
2252        let vec_u32: Vec<u32> = vec![1, 2, 3];
2253        let tokens_u32: Tokens = vec_u32.clone().into();
2254        assert_eq!(tokens_u32.0, vec_u32);
2255
2256        let slice_u32: &[u32] = &[4, 5];
2257        let tokens_slice_u32: Tokens = slice_u32.into();
2258        assert_eq!(tokens_slice_u32.0, vec![4, 5]);
2259
2260        let vec_i32: Vec<i32> = vec![-1, 0, 1]; // Note: -1 becomes large u32
2261        let tokens_i32: Tokens = vec_i32.into();
2262        assert_eq!(tokens_i32.0, vec![u32::MAX, 0, 1]);
2263
2264        let slice_i32: &[i32] = &[100, 200];
2265        let tokens_slice_i32: Tokens = slice_i32.into();
2266        assert_eq!(tokens_slice_i32.0, vec![100, 200]);
2267
2268        let into_vec: Vec<u32> = tokens_slice_i32.into();
2269        assert_eq!(into_vec, vec![100, 200]);
2270    }
2271
2272    #[test]
2273    fn test_tokens_equality() {
2274        let tokens = Tokens::from(vec![1, 2, 3]);
2275        assert_eq!(tokens, vec![1, 2, 3]);
2276        assert_eq!(vec![1, 2, 3], tokens);
2277        assert_eq!(tokens, &[1, 2, 3][..]);
2278        assert_eq!(&[1, 2, 3][..], tokens);
2279        assert_eq!(tokens, Tokens::from(vec![1, 2, 3]));
2280        assert_ne!(tokens, Tokens::from(vec![1, 2, 4]));
2281    }
2282
2283    #[test]
2284    fn test_tokens_deref_asref() {
2285        let tokens = Tokens::from(vec![10, 20, 30]);
2286
2287        // Deref to &[Token]
2288        assert_eq!(tokens.len(), 3);
2289        assert_eq!(tokens[1], 20);
2290        let slice: &[Token] = &tokens;
2291        assert_eq!(slice, &[10, 20, 30]);
2292
2293        // AsRef<[Token]>
2294        let as_ref_slice: &[Token] = tokens.as_ref();
2295        assert_eq!(as_ref_slice, &[10, 20, 30]);
2296
2297        // Borrow<[Token]>
2298        let borrowed_slice: &[Token] = std::borrow::Borrow::borrow(&tokens);
2299        assert_eq!(borrowed_slice, &[10, 20, 30]);
2300    }
2301
2302    #[test]
2303    fn test_tokens_into_sequence() {
2304        let tokens = Tokens::from(vec![1, 2, 3, 4, 5]);
2305        let seq = tokens.into_sequence(3, Some(TEST_SALT_HASH));
2306        assert_eq!(seq.blocks().len(), 1);
2307        assert_eq!(seq.blocks[0].tokens().as_ref(), &[1, 2, 3]);
2308        assert_eq!(seq.current_block().tokens().as_ref(), &[4, 5]);
2309        assert_eq!(seq.salt_hash(), TEST_SALT_HASH);
2310    }
2311
2312    #[test]
2313    fn test_partial_block_ops() {
2314        let mut partial = PartialTokenBlock::create_sequence_root(3, TEST_SALT_HASH);
2315        assert_eq!(partial.len(), 0);
2316        assert_eq!(partial.remaining(), 3);
2317        assert!(partial.is_empty());
2318
2319        // Push tokens
2320        assert!(partial.push_token(1).is_ok());
2321        assert_eq!(partial.len(), 1);
2322        assert_eq!(partial.remaining(), 2);
2323        let remaining = partial.push_tokens(Tokens::from(vec![2, 3, 4]));
2324        assert_eq!(partial.len(), 3);
2325        assert_eq!(partial.remaining(), 0);
2326        assert_eq!(remaining.as_ref(), &[4]); // Token 4 didn't fit
2327        assert_eq!(partial.tokens().as_ref(), &[1, 2, 3]);
2328
2329        // Push when full
2330        assert_eq!(partial.push_token(5), Err(TokenBlockError::Full));
2331        let remaining_full = partial.push_tokens(Tokens::from(vec![5]));
2332        assert_eq!(remaining_full.as_ref(), &[5]);
2333
2334        // Pop tokens
2335        assert!(partial.pop_token().is_ok());
2336        assert_eq!(partial.len(), 2);
2337        assert_eq!(partial.tokens().as_ref(), &[1, 2]);
2338        assert!(partial.pop_tokens(2).is_ok());
2339        assert!(partial.is_empty());
2340
2341        // Pop when empty
2342        assert_eq!(partial.pop_token(), Err(TokenBlockError::Empty));
2343        assert_eq!(
2344            partial.pop_tokens(1),
2345            Err(TokenBlockError::InsufficientTokens)
2346        );
2347
2348        // Commit incomplete
2349        assert!(partial.push_token(10).is_ok());
2350        assert_eq!(partial.commit(), Err(TokenBlockError::Incomplete));
2351
2352        // Commit complete
2353        assert!(partial.push_token(11).is_ok());
2354        assert!(partial.push_token(12).is_ok());
2355        assert_eq!(partial.len(), 3);
2356        let commit_result = partial.commit();
2357        assert!(commit_result.is_ok());
2358        let committed_block = commit_result.unwrap();
2359        assert_eq!(committed_block.tokens().as_ref(), &[10, 11, 12]);
2360
2361        // Check state after commit (partial block is now the next one)
2362        assert!(partial.is_empty());
2363        assert_eq!(
2364            partial.parent_sequence_hash,
2365            Some(committed_block.sequence_hash())
2366        );
2367        assert_eq!(partial.block_size, 3);
2368    }
2369
2370    #[test]
2371    fn test_token_block_creation_and_hashes() {
2372        let salt = TEST_SALT_HASH;
2373        let tokens1 = Tokens::from(vec![1, 2, 3, 4]);
2374        let chunk1 = TokenBlockChunk::new(tokens1.clone(), salt);
2375        let block1 = TokenBlock::from_chunk(chunk1, None, 0);
2376
2377        assert_eq!(block1.tokens(), &tokens1);
2378        assert_eq!(block1.salt_hash(), salt);
2379        assert_eq!(block1.parent_sequence_hash(), None);
2380        assert_eq!(block1.block_hash(), HASH_1_4);
2381        assert_eq!(block1.sequence_hash(), SEQ_HASH_1_4); // First block seq_hash == block_hash
2382        assert_eq!(block1.position(), 0); // First block is at position 0
2383
2384        // Verify positional lineage hash for block 1
2385        let plh1 = block1.positional_lineage_hash();
2386        assert_eq!(plh1.position(), 0);
2387        assert_eq!(plh1.parent_hash_fragment(), 0); // Root has no parent
2388        assert_eq!(plh1.current_sequence_hash(), SEQ_HASH_1_4);
2389
2390        let tokens2 = Tokens::from(vec![5, 6, 7, 8]);
2391        let chunk2 = TokenBlockChunk::new(tokens2.clone(), salt);
2392        let block2 = TokenBlock::from_chunk(chunk2, block1.parent_sequence_hash(), 1); // Incorrect parent
2393        // Sequence hash should differ if parent is wrong
2394        assert_ne!(block2.sequence_hash(), SEQ_HASH_5_8);
2395
2396        let chunk2_correct = TokenBlockChunk::new(tokens2.clone(), salt);
2397        let block2_correct =
2398            TokenBlock::from_chunk(chunk2_correct, Some(block1.sequence_hash()), 1);
2399
2400        assert_eq!(block2_correct.tokens(), &tokens2);
2401        assert_eq!(block2_correct.salt_hash(), salt);
2402        assert_eq!(
2403            block2_correct.parent_sequence_hash(),
2404            Some(block1.sequence_hash())
2405        );
2406        assert_eq!(block2_correct.block_hash(), HASH_5_8);
2407        assert_eq!(block2_correct.sequence_hash(), SEQ_HASH_5_8);
2408        assert_eq!(block2_correct.position(), 1); // Second block is at position 1
2409
2410        // Verify positional lineage hash for block 2
2411        let plh2 = block2_correct.positional_lineage_hash();
2412        assert_eq!(plh2.position(), 1);
2413        assert_eq!(
2414            plh2.parent_hash_fragment(),
2415            SEQ_HASH_1_4 & ((1u64 << 54) - 1)
2416        ); // Parent fragment is the parent's u64 truncated to 54 bits in mode 0
2417        assert_eq!(plh2.current_sequence_hash(), SEQ_HASH_5_8);
2418    }
2419
2420    #[test]
2421    fn test_new_sequence() {
2422        // Empty initial tokens
2423        let seq_empty = create_test_sequence(&[], 4, Some(TEST_SALT_HASH));
2424        assert!(seq_empty.blocks().is_empty());
2425        assert!(seq_empty.current_block().is_empty());
2426        assert_eq!(seq_empty.total_tokens(), 0);
2427        assert_eq!(seq_empty.salt_hash(), TEST_SALT_HASH);
2428        assert_eq!(seq_empty.current_block().parent_sequence_hash, None);
2429
2430        // Less than one block
2431        let seq_partial = create_test_sequence(&[1, 2], 4, Some(TEST_SALT_HASH));
2432        assert!(seq_partial.blocks().is_empty());
2433        assert_eq!(seq_partial.current_block().tokens().as_ref(), &[1, 2]);
2434        assert_eq!(seq_partial.total_tokens(), 2);
2435        assert_eq!(seq_partial.current_block().parent_sequence_hash, None);
2436
2437        // Exactly one block
2438        let seq_one_block = create_test_sequence(&[1, 2, 3, 4], 4, Some(TEST_SALT_HASH));
2439        assert_eq!(seq_one_block.blocks().len(), 1);
2440        assert!(seq_one_block.current_block().is_empty());
2441        assert_eq!(seq_one_block.total_tokens(), 4);
2442        assert_eq!(seq_one_block.blocks[0].tokens().as_ref(), &[1, 2, 3, 4]);
2443        assert_eq!(seq_one_block.blocks[0].sequence_hash(), SEQ_HASH_1_4);
2444        assert_eq!(
2445            seq_one_block.current_block().parent_sequence_hash,
2446            Some(SEQ_HASH_1_4)
2447        );
2448
2449        // More than one block
2450        let seq_multi = create_test_sequence(&[1, 2, 3, 4, 5, 6, 7, 8, 9], 4, Some(TEST_SALT_HASH));
2451        assert_eq!(seq_multi.blocks().len(), 2);
2452        assert_eq!(seq_multi.current_block().tokens().as_ref(), &[9]);
2453        assert_eq!(seq_multi.total_tokens(), 9);
2454        assert_eq!(seq_multi.blocks[0].sequence_hash(), SEQ_HASH_1_4);
2455        assert_eq!(seq_multi.blocks[1].sequence_hash(), SEQ_HASH_5_8);
2456        assert_eq!(
2457            seq_multi.current_block().parent_sequence_hash,
2458            Some(SEQ_HASH_5_8)
2459        );
2460
2461        // Test tokens_at across blocks and partial block
2462        assert_eq!(seq_multi.tokens_at(0..4).as_ref(), &[1, 2, 3, 4]); // First complete block
2463        assert_eq!(seq_multi.tokens_at(4..8).as_ref(), &[5, 6, 7, 8]); // Second complete block
2464        assert_eq!(seq_multi.tokens_at(8..9).as_ref(), &[9]); // Current partial block
2465        assert_eq!(seq_multi.tokens_at(2..6).as_ref(), &[3, 4, 5, 6]); // Spanning blocks
2466        assert_eq!(seq_multi.tokens_at(6..9).as_ref(), &[7, 8, 9]); // Spanning to partial
2467        assert_eq!(seq_multi.tokens_at(5..5).as_ref(), &[0u32; 0]); // Empty range
2468        assert_eq!(seq_multi.tokens_at(10..15).as_ref(), &[0u32; 0]); // Out of bounds
2469
2470        // No salt hash
2471        let seq_no_salt = create_test_sequence(&[1, 2, 3, 4, 5], 4, None);
2472        assert_eq!(seq_no_salt.salt_hash(), 0);
2473        assert_eq!(seq_no_salt.blocks().len(), 1);
2474        assert_ne!(seq_no_salt.blocks[0].block_hash(), HASH_1_4); // Hash differs with salt 0
2475        assert_eq!(seq_no_salt.current_block().tokens().as_ref(), &[5]);
2476    }
2477
2478    #[test]
2479    #[should_panic]
2480    fn test_new_sequence_zero_block_size() {
2481        let _ = create_test_sequence(&[1], 0, None);
2482    }
2483
2484    #[test]
2485    fn test_append_single_token() {
2486        let mut sequence =
2487            create_test_sequence(&[1, 2, 3, 4, 5, 6, 7, 8, 9, 10], 4, Some(TEST_SALT_HASH));
2488        assert_eq!(sequence.blocks().len(), 2);
2489        assert_eq!(sequence.current_block().tokens.len(), 2);
2490        assert_eq!(sequence.current_block().tokens, vec![9, 10]);
2491        assert_eq!(
2492            sequence.current_block().parent_sequence_hash,
2493            Some(SEQ_HASH_5_8)
2494        );
2495
2496        // Append token 11 - should not complete a block
2497        let completed_idx = sequence.append(11).unwrap();
2498        assert_eq!(completed_idx, None);
2499        assert_eq!(sequence.blocks().len(), 2);
2500        assert_eq!(sequence.current_block().tokens.as_ref(), &[9, 10, 11]);
2501
2502        // Append token 12 - should complete block 2 (index 2)
2503        // This will also commit block 2
2504        let completed_idx = sequence.append(12).unwrap();
2505        assert_eq!(completed_idx, Some(2));
2506        assert_eq!(sequence.blocks().len(), 3);
2507        assert_eq!(sequence.current_block.tokens.as_ref(), &[0u32; 0]);
2508        assert_eq!(sequence.current_block.remaining(), 4);
2509        assert_eq!(
2510            sequence.current_block().parent_sequence_hash,
2511            Some(SEQ_HASH_9_12)
2512        ); // Still linked to block 1
2513
2514        // Append token 13 - should not complete a block
2515        let completed_idx_13 = sequence.append(13).unwrap();
2516        assert_eq!(completed_idx_13, None);
2517        assert_eq!(sequence.blocks().len(), 3);
2518        assert_eq!(sequence.blocks[2].tokens().as_ref(), &[9, 10, 11, 12]);
2519        assert_eq!(sequence.blocks[2].sequence_hash(), SEQ_HASH_9_12);
2520        assert_eq!(sequence.current_block.tokens.as_ref(), &[13]); // New current block has 13
2521        assert_eq!(sequence.current_block.remaining(), 3);
2522        assert_eq!(
2523            sequence.current_block.parent_sequence_hash,
2524            Some(SEQ_HASH_9_12)
2525        ); // Linked to new block 2
2526    }
2527
2528    #[test]
2529    fn test_extend() {
2530        let block_size = 4;
2531        let salt_hash = Some(TEST_SALT_HASH);
2532
2533        // Case 1: Extend less than block size
2534        let mut seq1 = create_test_sequence(&[], block_size, salt_hash);
2535        let tokens1 = Tokens::from(vec![1, 2]);
2536        let completed1 = seq1.extend(tokens1).unwrap();
2537        assert_eq!(completed1, None); // No blocks completed
2538        assert_eq!(seq1.blocks.len(), 0);
2539        assert_eq!(seq1.current_block.tokens.as_ref(), &[1, 2]);
2540        assert_eq!(seq1.current_block.remaining(), 2);
2541        assert_eq!(seq1.current_block.parent_sequence_hash, None); // Still the root block
2542
2543        // Case 2: Extend exactly block size
2544        let mut seq2 = create_test_sequence(&[], block_size, salt_hash);
2545        let tokens2 = Tokens::from(vec![1, 2, 3, 4]);
2546        let completed2 = seq2.extend(tokens2).unwrap();
2547        assert_eq!(completed2, Some(0..1));
2548        assert_eq!(seq2.blocks.len(), 1);
2549        assert_eq!(seq2.current_block.tokens.as_ref(), &[0u32; 0]); // Current block is empty
2550        assert_eq!(seq2.current_block.remaining(), 4);
2551        assert_eq!(seq2.current_block.parent_sequence_hash, Some(SEQ_HASH_1_4)); // Still the root block
2552
2553        // Case 3: Extend more than block size, less than two blocks
2554        let mut seq3 = create_test_sequence(&[], block_size, salt_hash);
2555        let tokens3 = Tokens::from(vec![1, 2, 3, 4, 5, 6]);
2556        let completed3 = seq3.extend(tokens3).unwrap();
2557        assert_eq!(completed3, Some(0..1)); // Block at index 0 completed
2558        assert_eq!(seq3.blocks.len(), 1);
2559        assert_eq!(seq3.current_block.tokens.as_ref(), &[5, 6]); // Partial block has remainder
2560        assert_eq!(seq3.blocks[0].tokens().as_ref(), &[1, 2, 3, 4]);
2561        assert_eq!(seq3.current_block.parent_sequence_hash, Some(SEQ_HASH_1_4));
2562        assert_eq!(seq3.current_block.remaining(), 2);
2563
2564        // Case 4: Extend exactly two blocks
2565        let mut seq4 = create_test_sequence(&[], block_size, salt_hash);
2566        let tokens4 = Tokens::from(vec![1, 2, 3, 4, 5, 6, 7, 8]);
2567        let completed4 = seq4.extend(tokens4).unwrap();
2568        assert_eq!(completed4, Some(0..2)); // Only block 0 is committed
2569        assert_eq!(seq4.blocks.len(), 2); // Only 1 block committed
2570        assert_eq!(seq4.current_block.tokens.as_ref(), &[0u32; 0]);
2571        assert_eq!(seq4.current_block.remaining(), 4);
2572        assert_eq!(seq4.blocks[0].tokens().as_ref(), &[1, 2, 3, 4]);
2573        assert_eq!(seq4.blocks[0].sequence_hash(), SEQ_HASH_1_4);
2574        assert_eq!(seq4.current_block.parent_sequence_hash, Some(SEQ_HASH_5_8)); // Parent is the first block
2575
2576        // Case 5: Extend multiple times, completing blocks across calls
2577        let mut seq5 = create_test_sequence(&[], block_size, salt_hash);
2578        let tokens5a = Tokens::from(vec![1, 2]);
2579        let completed5a = seq5.extend(tokens5a).unwrap();
2580        assert_eq!(completed5a, None);
2581        assert_eq!(seq5.blocks.len(), 0);
2582        assert_eq!(seq5.current_block.tokens.as_ref(), &[1, 2]);
2583
2584        let tokens5b = Tokens::from(vec![3, 4, 5]);
2585        let completed5b = seq5.extend(tokens5b).unwrap();
2586        assert_eq!(completed5b, Some(0..1)); // Block at index 0 completed
2587        assert_eq!(seq5.blocks.len(), 1);
2588        assert_eq!(seq5.current_block.tokens.as_ref(), &[5]);
2589        assert_eq!(seq5.blocks[0].tokens().as_ref(), &[1, 2, 3, 4]);
2590        assert_eq!(seq5.current_block.parent_sequence_hash, Some(SEQ_HASH_1_4));
2591        assert_eq!(seq5.current_block.remaining(), 3);
2592
2593        let tokens5c = Tokens::from(vec![6, 7, 8, 9, 10]);
2594        let completed5c = seq5.extend(tokens5c).unwrap();
2595        assert_eq!(completed5c, Some(1..2)); // Block at index 1 completed
2596        assert_eq!(seq5.blocks.len(), 2);
2597        assert_eq!(seq5.current_block.tokens.as_ref(), &[9, 10]);
2598        assert_eq!(seq5.blocks[1].tokens().as_ref(), &[5, 6, 7, 8]);
2599        assert_eq!(seq5.current_block.parent_sequence_hash, Some(SEQ_HASH_5_8));
2600        assert_eq!(seq5.current_block.remaining(), 2);
2601
2602        // Case 6: Extend empty tokens
2603        let mut seq6 = create_test_sequence(&[1], block_size, salt_hash);
2604        let completed6 = seq6.extend(Tokens::default()).unwrap();
2605        assert_eq!(completed6, None);
2606        assert_eq!(seq6.blocks.len(), 0);
2607        assert_eq!(seq6.current_block.tokens.as_ref(), &[1]);
2608        assert_eq!(seq6.total_tokens(), 1);
2609
2610        // Case 7: Extend fills current exactly, no remainder
2611        let mut seq7 = create_test_sequence(&[1, 2], block_size, salt_hash);
2612        let tokens7 = Tokens::from(vec![3, 4]);
2613        let completed7 = seq7.extend(tokens7).unwrap();
2614        assert_eq!(completed7, Some(0..1)); // Block is full but not committed yet
2615        assert_eq!(seq7.blocks.len(), 1);
2616        assert_eq!(seq7.current_block.tokens.as_ref(), &[0u32; 0]); // Current block is full
2617        assert_eq!(seq7.current_block.remaining(), 4);
2618        assert_eq!(seq7.total_tokens(), 4);
2619        assert_eq!(seq7.current_block.parent_sequence_hash, Some(SEQ_HASH_1_4)); // Still the root block
2620
2621        // Test tokens_at extraction
2622        assert_eq!(seq7.tokens_at(0..2).as_ref(), &[1, 2]);
2623        assert_eq!(seq7.tokens_at(1..3).as_ref(), &[2, 3]);
2624        assert_eq!(seq7.tokens_at(0..4).as_ref(), &[1, 2, 3, 4]);
2625        assert_eq!(seq7.tokens_at(2..2).as_ref(), &[0u32; 0]); // Empty range
2626    }
2627
2628    #[test]
2629    fn test_truncate() {
2630        let block_size = 4;
2631        let salt_hash = Some(TEST_SALT_HASH);
2632        let initial_tokens = &[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]; // 10 tokens
2633
2634        // Case 1: Truncate within current block (len 9)
2635        let mut seq1 = create_test_sequence(initial_tokens, block_size, salt_hash);
2636        assert!(seq1.truncate(9).is_ok());
2637        assert_eq!(seq1.total_tokens(), 9);
2638        assert_eq!(seq1.blocks().len(), 2);
2639        assert_eq!(seq1.current_block().tokens.as_ref(), &[9]);
2640        assert_eq!(
2641            seq1.current_block().parent_sequence_hash,
2642            Some(SEQ_HASH_5_8)
2643        );
2644
2645        // Case 2: Truncate to exact block boundary (len 8)
2646        let mut seq2 = create_test_sequence(initial_tokens, block_size, salt_hash);
2647        assert!(seq2.truncate(8).is_ok());
2648        assert_eq!(seq2.total_tokens(), 8);
2649        assert_eq!(seq2.blocks().len(), 2);
2650        assert!(seq2.current_block().tokens.is_empty());
2651        assert_eq!(
2652            seq2.current_block().parent_sequence_hash,
2653            Some(SEQ_HASH_5_8)
2654        );
2655
2656        // Case 3: Truncate into last full block (len 7)
2657        let mut seq3 = create_test_sequence(initial_tokens, block_size, salt_hash);
2658        assert!(seq3.truncate(7).is_ok());
2659        assert_eq!(seq3.total_tokens(), 7);
2660        assert_eq!(seq3.blocks().len(), 1); // Block [5,6,7,8] removed conceptually
2661        assert_eq!(seq3.current_block().tokens.as_ref(), &[5, 6, 7]); // Kept 3 from [5,6,7,8]
2662        assert_eq!(
2663            seq3.current_block().parent_sequence_hash,
2664            Some(SEQ_HASH_1_4)
2665        ); // Parent is hash of [1,2,3,4]
2666        assert_eq!(seq3.blocks()[0].tokens().as_ref(), &[1, 2, 3, 4]);
2667
2668        // Case 4: Truncate removing full block(s) exactly (len 4)
2669        let mut seq4 = create_test_sequence(initial_tokens, block_size, salt_hash);
2670        assert!(seq4.truncate(4).is_ok());
2671        assert_eq!(seq4.total_tokens(), 4);
2672        assert_eq!(seq4.blocks().len(), 1); // Block [5,6,7,8] removed
2673        assert!(seq4.current_block().tokens.is_empty()); // New partial based on block [1,2,3,4]
2674        assert_eq!(
2675            seq4.current_block().parent_sequence_hash,
2676            Some(SEQ_HASH_1_4)
2677        );
2678        assert_eq!(seq4.blocks()[0].tokens().as_ref(), &[1, 2, 3, 4]);
2679
2680        // Case 5: Truncate into first block (len 3)
2681        let mut seq5 = create_test_sequence(initial_tokens, block_size, salt_hash);
2682        assert!(seq5.truncate(3).is_ok());
2683        assert_eq!(seq5.total_tokens(), 3);
2684        assert!(seq5.blocks().is_empty()); // Both blocks removed conceptually
2685        assert_eq!(seq5.current_block().tokens.as_ref(), &[1, 2, 3]); // Kept 3 from [1,2,3,4]
2686        assert_eq!(seq5.current_block().parent_sequence_hash, None); // No parent
2687
2688        // Case 6: Truncate to zero length (len 0)
2689        let mut seq6 = create_test_sequence(initial_tokens, block_size, salt_hash);
2690        assert!(seq6.truncate(0).is_ok());
2691        assert_eq!(seq6.total_tokens(), 0);
2692        assert!(seq6.blocks().is_empty());
2693        assert!(seq6.current_block().tokens.is_empty());
2694        assert_eq!(seq6.current_block().parent_sequence_hash, None);
2695
2696        // Case 7: Truncate to length greater than current (len 11)
2697        let mut seq7 = create_test_sequence(initial_tokens, block_size, salt_hash);
2698        let original_state = (seq7.blocks.clone(), seq7.current_block.tokens.clone()); // Clone for state check
2699        assert!(seq7.truncate(11).is_ok()); // Should have no effect
2700        assert_eq!(seq7.total_tokens(), 10);
2701        assert_eq!(seq7.blocks, original_state.0);
2702        assert_eq!(seq7.current_block.tokens, original_state.1);
2703
2704        // Case 8: Truncate to current length (len 10)
2705        let mut seq8 = create_test_sequence(initial_tokens, block_size, salt_hash);
2706        let original_state = (seq8.blocks.clone(), seq8.current_block.tokens.clone());
2707        assert!(seq8.truncate(10).is_ok());
2708        assert_eq!(seq8.total_tokens(), 10);
2709        assert_eq!(seq8.blocks, original_state.0);
2710        assert_eq!(seq8.current_block.tokens, original_state.1);
2711
2712        // Case 9: Truncate an empty sequence to 0
2713        let mut seq9 = create_test_sequence(&[], block_size, salt_hash);
2714        assert!(seq9.truncate(0).is_ok());
2715        assert_eq!(seq9.total_tokens(), 0);
2716        assert!(seq9.blocks().is_empty());
2717        assert!(seq9.current_block().tokens.is_empty());
2718
2719        // Case 10: Truncate on exact block boundary when current is empty (len 4)
2720        let tokens10 = &[1, 2, 3, 4, 5, 6, 7, 8]; // 8 tokens
2721        let mut seq10 = create_test_sequence(tokens10, block_size, salt_hash);
2722        assert_eq!(seq10.total_tokens(), 8);
2723        assert!(seq10.current_block().is_empty());
2724        assert!(seq10.truncate(4).is_ok()); // Remove block [5, 6, 7, 8]
2725        assert_eq!(seq10.total_tokens(), 4);
2726        assert_eq!(seq10.blocks().len(), 1);
2727        assert!(seq10.current_block().tokens.is_empty());
2728        assert_eq!(
2729            seq10.current_block().parent_sequence_hash,
2730            Some(SEQ_HASH_1_4)
2731        );
2732
2733        // Case 11: Truncate into first block when current is empty (len 3)
2734        let tokens11 = &[1, 2, 3, 4, 5, 6, 7, 8]; // 8 tokens
2735        let mut seq11 = create_test_sequence(tokens11, block_size, salt_hash);
2736        assert!(seq11.truncate(3).is_ok()); // Pop block [5,6,7,8] + 1 from [1,2,3,4]
2737        assert_eq!(seq11.total_tokens(), 3);
2738        assert!(seq11.blocks().is_empty());
2739        assert_eq!(seq11.current_block().tokens.as_ref(), &[1, 2, 3]); // Kept 3 from [1,2,3,4]
2740        assert_eq!(seq11.current_block().parent_sequence_hash, None);
2741    }
2742
2743    #[test]
2744    fn test_unwind() {
2745        let block_size = 4;
2746        let salt_hash = Some(TEST_SALT_HASH);
2747        let initial_tokens = &[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]; // 10 tokens
2748
2749        // Unwind 0
2750        let mut seq = create_test_sequence(initial_tokens, block_size, salt_hash);
2751        assert!(seq.unwind(0).is_ok());
2752        assert_eq!(seq.total_tokens(), 10);
2753
2754        // Unwind 1
2755        let mut seq = create_test_sequence(initial_tokens, block_size, salt_hash);
2756        assert!(seq.unwind(1).is_ok());
2757        assert_eq!(seq.total_tokens(), 9);
2758        assert_eq!(seq.current_block.tokens.as_ref(), &[9]);
2759
2760        // Unwind 3 (crosses boundary)
2761        let mut seq = create_test_sequence(initial_tokens, block_size, salt_hash);
2762        assert!(seq.unwind(3).is_ok());
2763        assert_eq!(seq.total_tokens(), 7);
2764        assert_eq!(seq.blocks.len(), 1);
2765        assert_eq!(seq.current_block.tokens.as_ref(), &[5, 6, 7]);
2766
2767        // Unwind all (10)
2768        let mut seq = create_test_sequence(initial_tokens, block_size, salt_hash);
2769        assert!(seq.unwind(10).is_ok());
2770        assert_eq!(seq.total_tokens(), 0);
2771        assert!(seq.blocks.is_empty());
2772        assert!(seq.current_block.is_empty());
2773
2774        // Unwind more than available (11)
2775        let mut seq = create_test_sequence(initial_tokens, block_size, salt_hash);
2776        assert_eq!(seq.unwind(11), Err(TokenBlockError::InsufficientTokens));
2777        assert_eq!(seq.total_tokens(), 10); // State unchanged
2778
2779        // Unwind from empty
2780        let mut seq_empty = create_test_sequence(&[], block_size, salt_hash);
2781        assert_eq!(
2782            seq_empty.unwind(1),
2783            Err(TokenBlockError::InsufficientTokens)
2784        );
2785    }
2786
2787    #[test]
2788    fn test_pop() {
2789        let block_size = 4;
2790        let salt_hash = Some(TEST_SALT_HASH);
2791        let initial_tokens = &[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]; // 10 tokens
2792
2793        let mut seq = create_test_sequence(initial_tokens, block_size, salt_hash);
2794
2795        // Pop 10
2796        assert_eq!(seq.pop(), Some(10));
2797        assert_eq!(seq.total_tokens(), 9);
2798        assert_eq!(seq.current_block.tokens.as_ref(), &[9]);
2799        assert_eq!(seq.blocks.len(), 2);
2800
2801        // Pop 9
2802        assert_eq!(seq.pop(), Some(9));
2803        assert_eq!(seq.total_tokens(), 8);
2804        assert!(seq.current_block.is_empty());
2805        assert_eq!(seq.blocks.len(), 2);
2806        assert_eq!(seq.current_block.parent_sequence_hash, Some(SEQ_HASH_5_8));
2807
2808        // Pop 8 (crosses boundary)
2809        assert_eq!(seq.pop(), Some(8));
2810        assert_eq!(seq.total_tokens(), 7);
2811        assert_eq!(seq.current_block.tokens.as_ref(), &[5, 6, 7]);
2812        assert_eq!(seq.blocks.len(), 1);
2813        assert_eq!(seq.current_block.parent_sequence_hash, Some(SEQ_HASH_1_4));
2814
2815        // Pop remaining partial (7, 6, 5)
2816        assert_eq!(seq.pop(), Some(7));
2817        assert_eq!(seq.pop(), Some(6));
2818        assert_eq!(seq.pop(), Some(5));
2819        assert_eq!(seq.total_tokens(), 4);
2820        assert!(seq.current_block.is_empty());
2821        assert_eq!(seq.blocks.len(), 1);
2822        assert_eq!(seq.current_block.parent_sequence_hash, Some(SEQ_HASH_1_4));
2823
2824        // Pop 4 (crosses boundary)
2825        assert_eq!(seq.pop(), Some(4));
2826        assert_eq!(seq.total_tokens(), 3);
2827        assert_eq!(seq.current_block.tokens.as_ref(), &[1, 2, 3]);
2828        assert!(seq.blocks.is_empty());
2829        assert_eq!(seq.current_block.parent_sequence_hash, None);
2830
2831        // Pop 3, 2, 1
2832        assert_eq!(seq.pop(), Some(3));
2833        assert_eq!(seq.pop(), Some(2));
2834        assert_eq!(seq.pop(), Some(1));
2835        assert_eq!(seq.total_tokens(), 0);
2836        assert!(seq.current_block.is_empty());
2837        assert!(seq.blocks.is_empty());
2838
2839        // Pop from empty
2840        assert_eq!(seq.pop(), None);
2841        assert_eq!(seq.total_tokens(), 0);
2842    }
2843
2844    #[test]
2845    fn test_total_tokens() {
2846        let block_size = 3;
2847        let salt_hash = Some(TEST_SALT_HASH);
2848
2849        let mut seq = create_test_sequence(&[], block_size, salt_hash);
2850        assert_eq!(seq.total_tokens(), 0);
2851
2852        seq.extend(Tokens::from(vec![1, 2])).unwrap();
2853        assert_eq!(seq.total_tokens(), 2);
2854
2855        seq.append(3).unwrap(); // Completes block 0
2856        assert_eq!(seq.total_tokens(), 3);
2857
2858        seq.extend(Tokens::from(vec![4, 5, 6, 7])).unwrap(); // Completes block 1, partial [7]
2859        assert_eq!(seq.total_tokens(), 7);
2860
2861        seq.pop().unwrap(); // Removes 7
2862        assert_eq!(seq.total_tokens(), 6);
2863
2864        seq.truncate(4).unwrap(); // Keep [1,2,3,4]
2865        assert_eq!(seq.total_tokens(), 4);
2866
2867        seq.unwind(2).unwrap(); // Keep [1,2]
2868        assert_eq!(seq.total_tokens(), 2);
2869    }
2870
2871    #[test]
2872    fn test_push_tokens_partial_block() {
2873        let mut partial = PartialTokenBlock::create_sequence_root(4, 1337);
2874
2875        let tokens = Tokens(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
2876
2877        let remaining = partial.push_tokens(tokens);
2878        assert_eq!(partial.tokens.len(), 4);
2879        assert_eq!(remaining.len(), 6);
2880    }
2881
2882    // ========== Additional tests for coverage improvement ==========
2883
2884    // === PositionalRadixTree Tests ===
2885
2886    #[test]
2887    fn test_positional_radix_tree_basic_operations() {
2888        use crate::PositionalRadixTree;
2889
2890        // Test new() and is_empty()
2891        let tree: PositionalRadixTree<String> = PositionalRadixTree::new();
2892        assert!(tree.is_empty());
2893        assert_eq!(tree.len(), 0);
2894
2895        // Test default()
2896        let tree2: PositionalRadixTree<i32> = PositionalRadixTree::default();
2897        assert!(tree2.is_empty());
2898
2899        // Test prefix() and insertion
2900        let psh1 = PositionalSequenceHash::new(0x1234, 0, 0xABCD);
2901        let psh2 = PositionalSequenceHash::new(0x5678, 0, 0xEF01);
2902        let psh3 = PositionalSequenceHash::new(0x9ABC, 1, 0x2345);
2903
2904        tree.prefix(&psh1).insert(psh1, "value1".to_string());
2905        assert!(!tree.is_empty());
2906        assert_eq!(tree.len(), 1);
2907
2908        tree.prefix(&psh2).insert(psh2, "value2".to_string());
2909        assert_eq!(tree.len(), 2);
2910
2911        tree.prefix(&psh3).insert(psh3, "value3".to_string());
2912        assert_eq!(tree.len(), 3);
2913
2914        // Test retrieval
2915        assert_eq!(
2916            tree.prefix(&psh1).get(&psh1).cloned(),
2917            Some("value1".to_string())
2918        );
2919    }
2920
2921    #[test]
2922    fn test_positional_radix_tree_with_lineage_hash() {
2923        use crate::PositionalRadixTree;
2924
2925        // Test generic usage with PositionalLineageHash
2926        let tree: PositionalRadixTree<u32, PositionalLineageHash> = PositionalRadixTree::new();
2927        assert!(tree.is_empty());
2928
2929        let plh1 = PositionalLineageHash::new(0x1234, None, 0);
2930        let plh2 = PositionalLineageHash::new(0x5678, Some(0x1234), 1);
2931
2932        tree.prefix(&plh1).insert(plh1, 100);
2933        tree.prefix(&plh2).insert(plh2, 200);
2934
2935        assert_eq!(tree.len(), 2);
2936        assert_eq!(tree.prefix(&plh1).get(&plh1).copied(), Some(100));
2937        assert_eq!(tree.prefix(&plh2).get(&plh2).copied(), Some(200));
2938    }
2939
2940    #[test]
2941    fn test_positional_radix_tree_position_lookup() {
2942        use crate::PositionalRadixTree;
2943
2944        let tree: PositionalRadixTree<String> = PositionalRadixTree::new();
2945
2946        // Insert at different positions
2947        let psh0 = PositionalSequenceHash::new(0x1111, 0, 0xAAAA);
2948        let psh1 = PositionalSequenceHash::new(0x2222, 1, 0xBBBB);
2949        let psh2 = PositionalSequenceHash::new(0x3333, 2, 0xCCCC);
2950
2951        tree.prefix(&psh0).insert(psh0, "pos0".to_string());
2952        tree.prefix(&psh1).insert(psh1, "pos1".to_string());
2953        tree.prefix(&psh2).insert(psh2, "pos2".to_string());
2954
2955        // Test position() method
2956        assert!(tree.position(0).is_some());
2957        assert!(tree.position(1).is_some());
2958        assert!(tree.position(2).is_some());
2959        assert!(tree.position(3).is_none()); // No entries at position 3
2960
2961        // Verify position lookup returns correct submap
2962        let pos0_map = tree.position(0).unwrap();
2963        assert_eq!(pos0_map.len(), 1);
2964    }
2965
2966    #[test]
2967    fn test_positional_radix_tree_concurrent_same_position() {
2968        use crate::PositionalRadixTree;
2969        use std::sync::Arc;
2970
2971        let tree = Arc::new(PositionalRadixTree::new());
2972        let threads: Vec<_> = (0..8_u64)
2973            .map(|value| {
2974                let tree = Arc::clone(&tree);
2975                std::thread::spawn(move || {
2976                    let key = PositionalSequenceHash::new(value, 7, value);
2977                    tree.prefix(&key).insert(key, value);
2978                })
2979            })
2980            .collect();
2981
2982        for thread in threads {
2983            thread.join().unwrap();
2984        }
2985
2986        assert_eq!(tree.len(), 8);
2987        assert_eq!(tree.position(7).unwrap().len(), 8);
2988    }
2989
2990    // === PositionalSequenceHash Additional Tests ===
2991
2992    #[test]
2993    fn test_positional_sequence_hash_mode_2_and_3() {
2994        // Mode 2: position fits in 24 bits (65536 <= pos < 16777216)
2995        let position_mode2 = 100_000u64;
2996        let seq_hash = 0x1234567890ABCDEF;
2997        let block_hash = 0xFEDCBA9876543210;
2998
2999        let psh_mode2 = PositionalSequenceHash::new(seq_hash, position_mode2, block_hash);
3000        assert_eq!(psh_mode2.mode(), 2, "Position 100,000 should use mode 2");
3001        assert_eq!(psh_mode2.position(), position_mode2);
3002        assert_eq!(psh_mode2.sequence_hash(), seq_hash);
3003        // Local block hash truncated to 38 bits in mode 2
3004        assert_eq!(
3005            psh_mode2.local_block_hash(),
3006            block_hash & ((1u64 << 38) - 1)
3007        );
3008
3009        // Mode 3: position fits in 31 bits (16777216 <= pos < 2147483648)
3010        let position_mode3 = 100_000_000u64;
3011        let psh_mode3 = PositionalSequenceHash::new(seq_hash, position_mode3, block_hash);
3012        assert_eq!(
3013            psh_mode3.mode(),
3014            3,
3015            "Position 100,000,000 should use mode 3"
3016        );
3017        assert_eq!(psh_mode3.position(), position_mode3);
3018        assert_eq!(psh_mode3.sequence_hash(), seq_hash);
3019        // Local block hash truncated to 31 bits in mode 3
3020        assert_eq!(
3021            psh_mode3.local_block_hash(),
3022            block_hash & ((1u64 << 31) - 1)
3023        );
3024    }
3025
3026    #[test]
3027    fn test_positional_sequence_hash_as_u128() {
3028        let psh = PositionalSequenceHash::new(0x1234, 100, 0xABCD);
3029        let raw = psh.as_u128();
3030
3031        // Verify we can reconstruct from raw value
3032        assert_eq!(raw & 0xFFFF_FFFF_FFFF_FFFF, 0x1234);
3033        assert!(raw > 0); // Non-zero
3034
3035        // Create another and compare
3036        let psh2 = PositionalSequenceHash::new(0x1234, 100, 0xABCD);
3037        assert_eq!(psh.as_u128(), psh2.as_u128());
3038    }
3039
3040    #[test]
3041    fn test_positional_sequence_hash_debug() {
3042        let psh = PositionalSequenceHash::new(0x1234567890ABCDEF, 42, 0xFEDCBA98);
3043        let debug_str = format!("{:?}", psh);
3044
3045        // Debug should contain field names and values
3046        assert!(debug_str.contains("PositionalSequenceHash"));
3047        assert!(debug_str.contains("sequence_hash"));
3048        assert!(debug_str.contains("local_block_hash"));
3049        assert!(debug_str.contains("position"));
3050    }
3051
3052    // === PositionalLineageHash Additional Tests ===
3053
3054    #[test]
3055    fn test_positional_lineage_hash_debug_and_display() {
3056        // Test position 0 (no parent shown)
3057        let plh_root = PositionalLineageHash::new(0x123456789ABCDEF0, None, 0);
3058        let debug_root = format!("{:?}", plh_root);
3059        let display_root = format!("{}", plh_root);
3060
3061        // Debug and Display should show position 0
3062        assert!(debug_root.starts_with("0:"));
3063        assert!(display_root.starts_with("0:"));
3064        // Position 0 should not show parent
3065        assert_eq!(debug_root.matches(':').count(), 1);
3066        assert_eq!(display_root.matches(':').count(), 1);
3067
3068        // Test position > 0 (parent shown)
3069        let plh_child = PositionalLineageHash::new(0xABCDEF0123456789, Some(0x123456789ABCDEF0), 5);
3070        let debug_child = format!("{:?}", plh_child);
3071        let display_child = format!("{}", plh_child);
3072
3073        // Should show position:current:parent
3074        assert!(debug_child.starts_with("5:"));
3075        assert!(display_child.starts_with("5:"));
3076        // Position > 0 should show parent (3 parts)
3077        assert_eq!(debug_child.matches(':').count(), 2);
3078        assert_eq!(display_child.matches(':').count(), 2);
3079    }
3080
3081    #[test]
3082    fn test_positional_lineage_hash_as_u128() {
3083        let plh = PositionalLineageHash::new(0x1234, Some(0x5678), 10);
3084        let raw = plh.as_u128();
3085
3086        assert!(raw > 0);
3087
3088        // Create another with same params and compare
3089        let plh2 = PositionalLineageHash::new(0x1234, Some(0x5678), 10);
3090        assert_eq!(plh.as_u128(), plh2.as_u128());
3091
3092        // Different params should give different hash
3093        let plh3 = PositionalLineageHash::new(0x1234, Some(0x5678), 11);
3094        assert_ne!(plh.as_u128(), plh3.as_u128());
3095    }
3096
3097    #[test]
3098    fn test_positional_lineage_hash_ord_by_position_then_current_fragment() {
3099        let at_5_low = PositionalLineageHash::new(0x10, Some(0x1111), 5);
3100        let at_5_high = PositionalLineageHash::new(0x20, Some(0x1111), 5);
3101        assert!(
3102            at_5_low.current_sequence_hash() < at_5_high.current_sequence_hash(),
3103            "test assumes distinct current sequence hashes at the same position"
3104        );
3105        assert!(at_5_low < at_5_high);
3106        assert!(at_5_high > at_5_low);
3107
3108        let at_3 = PositionalLineageHash::new(0x99, Some(0x2222), 3);
3109        assert!(at_3 < at_5_low);
3110        assert!(at_5_high < PositionalLineageHash::new(0x01, Some(0x3333), 6));
3111    }
3112
3113    #[test]
3114    fn test_positional_lineage_hash_ord_tiebreak_parent_via_packed_u128() {
3115        let same_pos_same_current = PositionalLineageHash::new(0x1234, Some(0x100), 10);
3116        let same_pos_same_current_other_parent =
3117            PositionalLineageHash::new(0x1234, Some(0x200), 10);
3118        assert_eq!(same_pos_same_current.position(), 10);
3119        assert_eq!(
3120            same_pos_same_current.position(),
3121            same_pos_same_current_other_parent.position()
3122        );
3123        assert_eq!(
3124            same_pos_same_current.current_sequence_hash(),
3125            same_pos_same_current_other_parent.current_sequence_hash()
3126        );
3127        assert_ne!(same_pos_same_current, same_pos_same_current_other_parent);
3128        assert_ne!(
3129            same_pos_same_current.cmp(&same_pos_same_current_other_parent),
3130            std::cmp::Ordering::Equal
3131        );
3132    }
3133
3134    #[test]
3135    fn test_positional_lineage_hash_vec_sort_matches_ord() {
3136        let a = PositionalLineageHash::new(0x30, None, 0);
3137        let b = PositionalLineageHash::new(0x10, Some(0x30), 2);
3138        let c = PositionalLineageHash::new(0x20, Some(0x30), 2);
3139        let mut v = vec![b, a, c];
3140        v.sort();
3141        assert_eq!(v, vec![a, b, c]);
3142    }
3143
3144    #[test]
3145    fn test_positional_lineage_hash_itertools_sorted() {
3146        use itertools::Itertools;
3147
3148        let a = PositionalLineageHash::new(0x30, None, 0);
3149        let b = PositionalLineageHash::new(0x10, Some(0x30), 2);
3150        let c = PositionalLineageHash::new(0x20, Some(0x30), 2);
3151        let sorted: Vec<_> = vec![b, a, c].into_iter().sorted().collect();
3152        assert_eq!(sorted, vec![a, b, c]);
3153    }
3154
3155    // === Tokens From Impls ===
3156
3157    #[test]
3158    fn test_tokens_from_vec_usize() {
3159        let usize_vec: Vec<usize> = vec![1, 2, 3, 4, 5];
3160        let tokens = Tokens::from(usize_vec);
3161
3162        assert_eq!(tokens.as_ref(), &[1u32, 2, 3, 4, 5]);
3163        assert_eq!(tokens.len(), 5);
3164    }
3165
3166    #[test]
3167    fn test_tokens_partial_eq_slice_ref() {
3168        let tokens = Tokens::from(vec![1u32, 2, 3, 4]);
3169        let slice: &[Token] = &[1, 2, 3, 4];
3170
3171        // Test PartialEq<&[Token]> for Tokens
3172        assert!(tokens == slice);
3173
3174        let different_slice: &[Token] = &[1, 2, 3, 5];
3175        assert!(tokens != different_slice);
3176    }
3177
3178    // === TokenBlock Accessors ===
3179
3180    #[test]
3181    fn test_token_block_accessors() {
3182        let tokens = Tokens::from(vec![1u32, 2, 3, 4]);
3183        let seq = TokenBlockSequence::new(tokens, 4, Some(1337));
3184
3185        let block = &seq.blocks()[0];
3186
3187        // Test block_size()
3188        assert_eq!(block.block_size(), 4);
3189
3190        // Test positional_sequence_hash()
3191        let psh = block.positional_sequence_hash();
3192        assert_eq!(psh.position(), 0);
3193
3194        // Test positional_lineage_hash()
3195        let plh = block.positional_lineage_hash();
3196        assert_eq!(plh.position(), 0);
3197        assert_eq!(plh.parent_hash_fragment(), 0); // Root has no parent
3198    }
3199
3200    #[test]
3201    fn test_positional_hash_trait_impls() {
3202        use crate::PositionalHash;
3203
3204        // Test PositionalHash for PositionalSequenceHash
3205        let psh = PositionalSequenceHash::new(0x1234, 42, 0xABCD);
3206        assert_eq!(PositionalHash::position(&psh), 42);
3207
3208        // Test PositionalHash for PositionalLineageHash
3209        let plh = PositionalLineageHash::new(0x1234, None, 99);
3210        assert_eq!(PositionalHash::position(&plh), 99);
3211    }
3212
3213    // === TokenBlockSequence Edge Cases ===
3214
3215    #[test]
3216    fn test_sequence_pop_from_full_block() {
3217        // Test pop when current partial block is empty (must pop from full block)
3218        let tokens = Tokens::from(vec![1u32, 2, 3, 4, 5, 6, 7, 8]);
3219        let mut seq = TokenBlockSequence::new(tokens, 4, Some(TEST_SALT_HASH));
3220
3221        // Current block should be empty, all tokens in completed blocks
3222        assert!(seq.current_block().is_empty());
3223        assert_eq!(seq.blocks().len(), 2);
3224        assert_eq!(seq.total_tokens(), 8);
3225
3226        // Pop should remove from last full block
3227        let popped = seq.pop();
3228        assert_eq!(popped, Some(8));
3229        assert_eq!(seq.total_tokens(), 7);
3230        assert_eq!(seq.blocks().len(), 1);
3231        assert_eq!(seq.current_block().tokens.as_ref(), &[5, 6, 7]);
3232    }
3233
3234    #[test]
3235    #[allow(clippy::reversed_empty_ranges)] // so we can explicitly test invalid ranges
3236    fn test_sequence_tokens_at_edge_cases() {
3237        let tokens = Tokens::from(vec![1u32, 2, 3, 4, 5]);
3238        let seq = TokenBlockSequence::new(tokens, 4, Some(TEST_SALT_HASH));
3239
3240        // Start > end (invalid range
3241        assert!(seq.tokens_at(3..2).is_empty());
3242
3243        // End > total (out of bounds)
3244        assert!(seq.tokens_at(0..10).is_empty());
3245
3246        // Valid edge case: exact boundaries
3247        assert_eq!(seq.tokens_at(0..4).as_ref(), &[1, 2, 3, 4]);
3248        assert_eq!(seq.tokens_at(4..5).as_ref(), &[5]);
3249    }
3250
3251    #[test]
3252    fn test_sequence_next_block() {
3253        let tokens = Tokens::from(vec![1u32, 2, 3, 4]);
3254        let seq = TokenBlockSequence::new(tokens, 4, Some(1337));
3255
3256        let block = &seq.blocks()[0];
3257        let next_partial = block.next_block();
3258
3259        // next_block should create a partial block linked to this block
3260        assert!(next_partial.is_empty());
3261        assert_eq!(next_partial.remaining(), 4);
3262        assert_eq!(
3263            next_partial.parent_sequence_hash,
3264            Some(block.sequence_hash())
3265        );
3266        assert_eq!(next_partial.position, 1);
3267    }
3268
3269    #[test]
3270    fn test_sequence_reset() {
3271        let tokens = Tokens::from(vec![1u32, 2, 3, 4, 5, 6, 7, 8, 9]);
3272        let mut seq = TokenBlockSequence::new(tokens, 4, Some(1337));
3273
3274        assert_eq!(seq.blocks().len(), 2);
3275        assert_eq!(seq.total_tokens(), 9);
3276
3277        seq.reset();
3278
3279        assert!(seq.blocks().is_empty());
3280        assert!(seq.current_block().is_empty());
3281        assert_eq!(seq.total_tokens(), 0);
3282        assert_eq!(seq.current_block().parent_sequence_hash, None);
3283    }
3284
3285    #[test]
3286    fn test_sequence_into_parts() {
3287        let tokens = Tokens::from(vec![1u32, 2, 3, 4, 5]);
3288        let seq = TokenBlockSequence::new(tokens, 4, Some(1337));
3289
3290        let (blocks, partial) = seq.into_parts();
3291
3292        assert_eq!(blocks.len(), 1);
3293        assert_eq!(partial.tokens.as_ref(), &[5]);
3294    }
3295
3296    #[test]
3297    fn test_sequence_last_complete_block() {
3298        // Empty sequence
3299        let seq_empty = TokenBlockSequence::new(Tokens::default(), 4, None);
3300        assert!(seq_empty.last_complete_block().is_none());
3301
3302        // With blocks
3303        let tokens = Tokens::from(vec![1u32, 2, 3, 4, 5, 6, 7, 8]);
3304        let seq = TokenBlockSequence::new(tokens, 4, Some(1337));
3305        let last = seq.last_complete_block();
3306        assert!(last.is_some());
3307        assert_eq!(last.unwrap().tokens().as_ref(), &[5, 6, 7, 8]);
3308    }
3309
3310    #[test]
3311    fn test_positional_hashes_msgpack_roundtrip() {
3312        let psh = PositionalSequenceHash::new(0xDEAD_BEEF_CAFE_BABE, 12345, 0x0123_4567_89AB_CDEF);
3313        let bytes = rmp_serde::to_vec(&psh).expect("psh serialize");
3314        let decoded: PositionalSequenceHash =
3315            rmp_serde::from_slice(&bytes).expect("psh deserialize");
3316        assert_eq!(psh, decoded);
3317        assert_eq!(psh.as_u128(), decoded.as_u128());
3318
3319        let plh =
3320            PositionalLineageHash::new(0x1111_2222_3333_4444, Some(0x5555_6666_7777_8888), 256);
3321        let bytes = rmp_serde::to_vec(&plh).expect("plh serialize");
3322        let decoded: PositionalLineageHash =
3323            rmp_serde::from_slice(&bytes).expect("plh deserialize");
3324        assert_eq!(plh, decoded);
3325        assert_eq!(plh.as_u128(), decoded.as_u128());
3326
3327        // Vec roundtrip — exercises the codec inside a container.
3328        let vec = vec![psh, PositionalSequenceHash::default(), psh];
3329        let bytes = rmp_serde::to_vec(&vec).expect("vec serialize");
3330        let decoded: Vec<PositionalSequenceHash> =
3331            rmp_serde::from_slice(&bytes).expect("vec deserialize");
3332        assert_eq!(vec, decoded);
3333    }
3334
3335    #[test]
3336    fn test_positional_hashes_json_roundtrip() {
3337        // Confirm the byte-array codec also roundtrips through JSON (array of u8).
3338        let psh = PositionalSequenceHash::new(0xAAAA_BBBB_CCCC_DDDD, 7, 0xEEEE_FFFF_0000_1111);
3339        let json = serde_json::to_string(&psh).expect("psh json serialize");
3340        let decoded: PositionalSequenceHash =
3341            serde_json::from_str(&json).expect("psh json deserialize");
3342        assert_eq!(psh, decoded);
3343
3344        let plh = PositionalLineageHash::new(0x1234_5678, Some(0xABCD_EF01), 42);
3345        let json = serde_json::to_string(&plh).expect("plh json serialize");
3346        let decoded: PositionalLineageHash =
3347            serde_json::from_str(&json).expect("plh json deserialize");
3348        assert_eq!(plh, decoded);
3349    }
3350
3351    // ----------------------------------------------------------------------------------------
3352    // Multimodal block-formation tests (#10–14 in the kv-hashing plan).
3353    // ----------------------------------------------------------------------------------------
3354
3355    /// #10: a sequence built via `new_with_mm` with empty mm_info must equal one built via `new`.
3356    #[test]
3357    fn tokens_mm_zero_mm_equivalence() {
3358        let tokens = Tokens::from(vec![1u32, 2, 3, 4, 5, 6, 7, 8, 9]);
3359        let baseline = TokenBlockSequence::new(tokens.clone(), 4, Some(TEST_SALT_HASH));
3360        let mm = TokenBlockSequence::new_with_mm(tokens, &[], 4, Some(TEST_SALT_HASH))
3361            .expect("validation should pass for empty mm_info");
3362
3363        assert_eq!(mm.blocks().len(), baseline.blocks().len());
3364        for (a, b) in mm.blocks().iter().zip(baseline.blocks().iter()) {
3365            assert_eq!(a.salt_hash(), b.salt_hash());
3366            assert_eq!(a.block_hash(), b.block_hash());
3367            assert_eq!(a.sequence_hash(), b.sequence_hash());
3368            assert_eq!(a.parent_sequence_hash(), b.parent_sequence_hash());
3369            assert_eq!(a.positional_lineage_hash(), b.positional_lineage_hash());
3370        }
3371        assert!(mm.mm_runs().is_empty());
3372    }
3373
3374    /// #11: byte layout — verify the MM-aware buffer matches the documented
3375    /// 13-bytes-per-slot tagged framing, and that block_hash is XXH3 over that exact buffer.
3376    #[test]
3377    fn tokens_mm_byte_layout() {
3378        // Block 0: tokens [t0..t3], placeholder run [4..6) with mm_hash=0xAA, then t6, t7.
3379        // block_size = 8 ⇒ block_offset = 0, run covers slots 4..6. The block is MM-affected
3380        // ⇒ tagged 13-byte frames apply to *every* slot.
3381        let tokens = Tokens::from(vec![100u32, 101, 102, 103, 0, 0, 106, 107]);
3382        let mm = vec![TokenBlockMmInfo {
3383            mm_hash: 0xAAu64,
3384            offset: 4,
3385            length: 2,
3386        }];
3387        let salt = TEST_SALT_HASH;
3388
3389        // Build expected bytes manually: each slot is 13 bytes.
3390        let mut expected = Vec::new();
3391        for &t in &[100u32, 101, 102, 103] {
3392            expected.push(MM_SLOT_TAG_TOKEN);
3393            expected.extend_from_slice(&t.to_le_bytes());
3394            expected.extend_from_slice(&0u64.to_le_bytes());
3395        }
3396        for run_off in 0u32..2 {
3397            expected.push(MM_SLOT_TAG_PLACEHOLDER);
3398            expected.extend_from_slice(&run_off.to_le_bytes());
3399            expected.extend_from_slice(&0xAAu64.to_le_bytes());
3400        }
3401        for &t in &[106u32, 107] {
3402            expected.push(MM_SLOT_TAG_TOKEN);
3403            expected.extend_from_slice(&t.to_le_bytes());
3404            expected.extend_from_slice(&0u64.to_le_bytes());
3405        }
3406        assert_eq!(expected.len(), 8 * 13);
3407
3408        // Validate helper output matches.
3409        let helper_bytes = compute_block_bytes_with_mm(&tokens, 0, &mm);
3410        assert_eq!(helper_bytes, expected, "MM-aware byte buffer mismatch");
3411
3412        // Validate block_hash equals XXH3 over the expected buffer.
3413        let expected_block_hash = compute_block_hash(&expected, salt);
3414        let seq = TokenBlockSequence::new_with_mm(tokens, &mm, 8, Some(salt)).unwrap();
3415        assert_eq!(seq.blocks().len(), 1);
3416        assert_eq!(seq.blocks()[0].block_hash(), expected_block_hash);
3417    }
3418
3419    /// #11b — collision regression. Reviewer P1: with the original 4/12 mixed encoding,
3420    /// `block_size=2` blocks `[MM(slot 0), token]` and `[token, MM(slot 1)]` could produce
3421    /// identical byte streams under chosen `mm_hash`/token values. With tagged 13-byte
3422    /// framing they MUST differ.
3423    #[test]
3424    fn tokens_mm_no_position_collision() {
3425        let salt = TEST_SALT_HASH;
3426        // Layout A: block_size=2, MM at slot 0, token at slot 1.
3427        let tokens_a = Tokens::from(vec![0u32, 0xAB]);
3428        let mm_a = vec![TokenBlockMmInfo {
3429            mm_hash: 0x1122_3344_5566_7788,
3430            offset: 0,
3431            length: 1,
3432        }];
3433        // Layout B: block_size=2, token at slot 0, MM at slot 1.
3434        let tokens_b = Tokens::from(vec![0xAB, 0u32]);
3435        let mm_b = vec![TokenBlockMmInfo {
3436            mm_hash: 0x1122_3344_5566_7788,
3437            offset: 1,
3438            length: 1,
3439        }];
3440
3441        let bytes_a = compute_block_bytes_with_mm(&tokens_a, 0, &mm_a);
3442        let bytes_b = compute_block_bytes_with_mm(&tokens_b, 0, &mm_b);
3443        assert_ne!(
3444            bytes_a, bytes_b,
3445            "tagged framing must distinguish slot kinds at different positions"
3446        );
3447
3448        let seq_a = TokenBlockSequence::new_with_mm(tokens_a, &mm_a, 2, Some(salt)).unwrap();
3449        let seq_b = TokenBlockSequence::new_with_mm(tokens_b, &mm_b, 2, Some(salt)).unwrap();
3450        assert_ne!(
3451            seq_a.blocks()[0].block_hash(),
3452            seq_b.blocks()[0].block_hash()
3453        );
3454    }
3455
3456    /// #11c — per-block legacy fallback. A block with no overlapping MM run uses the legacy
3457    /// 4-byte-per-slot encoding so its `block_hash` matches the existing zero-MM path. Block 0
3458    /// of an MM-bearing sequence (run starts in block 1) must equal block 0 of a no-MM sequence.
3459    #[test]
3460    fn tokens_mm_legacy_fallback_per_block() {
3461        let block_size: u32 = 4;
3462        let salt = Some(TEST_SALT_HASH);
3463        let raw = vec![1u32, 2, 3, 4, 5, 6, 7, 8];
3464        // MM run covers only block 1 (positions [4..7)).
3465        let mm = vec![TokenBlockMmInfo {
3466            mm_hash: 0xAB,
3467            offset: 4,
3468            length: 3,
3469        }];
3470        let seq_mm =
3471            TokenBlockSequence::new_with_mm(Tokens::from(raw.clone()), &mm, block_size, salt)
3472                .unwrap();
3473        let seq_plain = TokenBlockSequence::new(Tokens::from(raw), block_size, salt);
3474
3475        // Block 0 untouched by MM ⇒ identical hashes.
3476        assert_eq!(
3477            seq_mm.blocks()[0].block_hash(),
3478            seq_plain.blocks()[0].block_hash()
3479        );
3480        assert_eq!(
3481            seq_mm.blocks()[0].sequence_hash(),
3482            seq_plain.blocks()[0].sequence_hash()
3483        );
3484        // Block 1 IS MM-affected ⇒ hashes diverge.
3485        assert_ne!(
3486            seq_mm.blocks()[1].block_hash(),
3487            seq_plain.blocks()[1].block_hash()
3488        );
3489    }
3490
3491    /// #11d — `offset + length` overflow is rejected as a dedicated error variant rather
3492    /// than panicking or silently wrapping.
3493    #[test]
3494    fn tokens_mm_validation_overflow() {
3495        let bad = vec![TokenBlockMmInfo {
3496            mm_hash: 1,
3497            offset: usize::MAX - 2,
3498            length: 10,
3499        }];
3500        let err = validate_and_sort_mm_info(&bad, usize::MAX).expect_err("must reject overflow");
3501        assert!(matches!(err, MmInfoError::OffsetOverflow { .. }));
3502    }
3503
3504    /// #12: building incrementally via push_token / push_mm_run yields the same sequence
3505    /// as the batch `new_with_mm` constructor.
3506    #[test]
3507    fn tokens_mm_streaming_equals_batch() {
3508        // Layout: [t,t,t,(MM=0xAA len=4),t,t] over block_size=4 ⇒ 2 blocks + partial.
3509        let tokens = Tokens::from(vec![1u32, 2, 3, 0, 0, 0, 0, 6, 7]);
3510        let mm = vec![TokenBlockMmInfo {
3511            mm_hash: 0xAAu64,
3512            offset: 3,
3513            length: 4,
3514        }];
3515        let salt = Some(TEST_SALT_HASH);
3516        let batch = TokenBlockSequence::new_with_mm(tokens, &mm, 4, salt).unwrap();
3517
3518        let mut streamed = TokenBlockSequence::new(Tokens::default(), 4, salt);
3519        streamed.push_token(1).unwrap();
3520        streamed.push_token(2).unwrap();
3521        streamed.push_token(3).unwrap();
3522        streamed.push_mm_run(0xAAu64, 4).unwrap();
3523        streamed.push_token(6).unwrap();
3524        streamed.push_token(7).unwrap();
3525
3526        assert_eq!(streamed.blocks().len(), batch.blocks().len());
3527        for (a, b) in streamed.blocks().iter().zip(batch.blocks().iter()) {
3528            assert_eq!(a.block_hash(), b.block_hash(), "block_hash mismatch");
3529            assert_eq!(a.sequence_hash(), b.sequence_hash(), "seq_hash mismatch");
3530            assert_eq!(
3531                a.positional_lineage_hash(),
3532                b.positional_lineage_hash(),
3533                "PLH mismatch"
3534            );
3535        }
3536        assert_eq!(streamed.mm_runs(), batch.mm_runs());
3537    }
3538
3539    /// #13: a multi-block MM run produces distinct block_hashes for blocks fully covered by
3540    /// the run (run_offset increases monotonically), and shares prefix hashes with another
3541    /// request whose run starts at the same global offset.
3542    #[test]
3543    fn tokens_mm_multi_block_run() {
3544        let block_size: u32 = 8;
3545        let bs = block_size as usize;
3546        // Run of length 2*bs + k = 20 starting at the boundary of block 0 ⇒ spans blocks 0,1,2.
3547        // Blocks 0 and 1 are fully placeholders; block 2 starts as 4 placeholders + 4 reals.
3548        let mut tokens_a: Vec<Token> = vec![0u32; 2 * bs]; // blocks 0, 1 (placeholders)
3549        tokens_a.extend_from_slice(&[0u32, 0, 0, 0, 100, 101, 102, 103]); // block 2
3550        let tokens_a = Tokens::from(tokens_a);
3551        let mm = vec![TokenBlockMmInfo {
3552            mm_hash: 0xCAFEBABEu64,
3553            offset: 0,
3554            length: 20,
3555        }];
3556        let seq_a = TokenBlockSequence::new_with_mm(
3557            tokens_a.clone(),
3558            &mm,
3559            block_size,
3560            Some(TEST_SALT_HASH),
3561        )
3562        .unwrap();
3563        assert_eq!(seq_a.blocks().len(), 3);
3564
3565        // Block 0 and block 1 are *both* fully placeholder, but with different run_offsets
3566        // (0..7 vs 8..15) → block_hashes must differ.
3567        let bh0 = seq_a.blocks()[0].block_hash();
3568        let bh1 = seq_a.blocks()[1].block_hash();
3569        assert_ne!(
3570            bh0, bh1,
3571            "fully-placeholder blocks at different run_offsets must hash differently"
3572        );
3573
3574        // Same image at the same global starting position in another request must share blocks.
3575        let seq_b =
3576            TokenBlockSequence::new_with_mm(tokens_a, &mm, block_size, Some(TEST_SALT_HASH))
3577                .unwrap();
3578        assert_eq!(
3579            seq_a.blocks()[0].block_hash(),
3580            seq_b.blocks()[0].block_hash()
3581        );
3582        assert_eq!(
3583            seq_a.blocks()[1].block_hash(),
3584            seq_b.blocks()[1].block_hash()
3585        );
3586        assert_eq!(
3587            seq_a.blocks()[2].block_hash(),
3588            seq_b.blocks()[2].block_hash()
3589        );
3590
3591        // A different mm_hash at the same position must diverge starting at block 0.
3592        let mm_diff = vec![TokenBlockMmInfo {
3593            mm_hash: 0xDEADBEEFu64,
3594            offset: 0,
3595            length: 20,
3596        }];
3597        let mut tokens_c: Vec<Token> = vec![0u32; 2 * bs];
3598        tokens_c.extend_from_slice(&[0u32, 0, 0, 0, 100, 101, 102, 103]);
3599        let seq_c = TokenBlockSequence::new_with_mm(
3600            Tokens::from(tokens_c),
3601            &mm_diff,
3602            block_size,
3603            Some(TEST_SALT_HASH),
3604        )
3605        .unwrap();
3606        assert_ne!(
3607            seq_a.blocks()[0].block_hash(),
3608            seq_c.blocks()[0].block_hash()
3609        );
3610    }
3611
3612    /// #14: mm_info validation rejects overlap, out-of-bounds, and zero-length runs.
3613    #[test]
3614    fn tokens_mm_validation() {
3615        let tokens = Tokens::from(vec![0u32; 32]);
3616        // Overlap.
3617        let overlap = vec![
3618            TokenBlockMmInfo {
3619                mm_hash: 1,
3620                offset: 0,
3621                length: 5,
3622            },
3623            TokenBlockMmInfo {
3624                mm_hash: 2,
3625                offset: 4,
3626                length: 5,
3627            },
3628        ];
3629        let err = TokenBlockSequence::new_with_mm(tokens.clone(), &overlap, 4, None).unwrap_err();
3630        assert!(matches!(
3631            err,
3632            TokenBlockError::MmInfo(MmInfoError::Overlapping { .. })
3633        ));
3634
3635        // Out-of-bounds.
3636        let oob = vec![TokenBlockMmInfo {
3637            mm_hash: 1,
3638            offset: 30,
3639            length: 10,
3640        }];
3641        let err = TokenBlockSequence::new_with_mm(tokens.clone(), &oob, 4, None).unwrap_err();
3642        assert!(matches!(
3643            err,
3644            TokenBlockError::MmInfo(MmInfoError::OutOfBounds { .. })
3645        ));
3646
3647        // Zero-length run.
3648        let empty = vec![TokenBlockMmInfo {
3649            mm_hash: 1,
3650            offset: 0,
3651            length: 0,
3652        }];
3653        let err = TokenBlockSequence::new_with_mm(tokens, &empty, 4, None).unwrap_err();
3654        assert!(matches!(
3655            err,
3656            TokenBlockError::MmInfo(MmInfoError::EmptyRun)
3657        ));
3658
3659        // push_mm_run with length 0.
3660        let mut seq = TokenBlockSequence::new(Tokens::default(), 4, None);
3661        let err = seq.push_mm_run(0xAB, 0).unwrap_err();
3662        assert!(matches!(
3663            err,
3664            TokenBlockError::MmInfo(MmInfoError::EmptyRun)
3665        ));
3666
3667        // truncate / pop / unwind blocked once mm_runs is non-empty.
3668        let mut seq = TokenBlockSequence::new(Tokens::from(vec![1u32, 2, 3]), 4, None);
3669        seq.push_mm_run(0xAB, 2).unwrap();
3670        assert!(matches!(
3671            seq.truncate(0).unwrap_err(),
3672            TokenBlockError::MmRunsPresent
3673        ));
3674        assert!(matches!(
3675            seq.unwind(1).unwrap_err(),
3676            TokenBlockError::MmRunsPresent
3677        ));
3678    }
3679}