gigatoken 0.1.0 → 0.2.0

data/lib/gigatoken/encodings.rb

@@ -0,0 +1,103 @@
+# frozen_string_literal: true
+
+module Gigatoken
+  # The tiktoken encodings gigatoken vendors ranks for, and the pieces a
+  # .tiktoken file doesn't carry: its pretokenizer scheme and special-token
+  # table (see Tokenizer.from_tiktoken and lib/gigatoken/encodings/
+  # PROVENANCE.md, the source of truth this is transcribed from).
+  module Encodings
+    DATA_DIR = File.expand_path("encodings", __dir__)
+    private_constant :DATA_DIR
+
+    # The non-contiguous head of openai_public.py's o200k_harmony()
+    # special-token table, transcribed verbatim (see PROVENANCE.md): its ten
+    # named control tokens, plus the six reserved slots sitting in the gaps
+    # between them (200000, 200001, 200004, 200009, 200010, 200011). The
+    # reserved range only goes contiguous at 200013.
+    HARMONY_HEAD_TOKENS = {
+      "<|startoftext|>" => 199998,
+      "<|endoftext|>" => 199999,
+      "<|reserved_200000|>" => 200000,
+      "<|reserved_200001|>" => 200001,
+      "<|return|>" => 200002,
+      "<|constrain|>" => 200003,
+      "<|reserved_200004|>" => 200004,
+      "<|channel|>" => 200005,
+      "<|start|>" => 200006,
+      "<|end|>" => 200007,
+      "<|message|>" => 200008,
+      "<|reserved_200009|>" => 200009,
+      "<|reserved_200010|>" => 200010,
+      "<|reserved_200011|>" => 200011,
+      "<|call|>" => 200012,
+      "<|endofprompt|>" => 200018
+    }.freeze
+    private_constant :HARMONY_HEAD_TOKENS
+
+    # The contiguous tail of that reserved range: 200013..201087. With the
+    # head above, 1091 entries total — 10 named, 1081 reserved (see
+    # PROVENANCE.md).
+    HARMONY_RESERVED_TAIL = (200013..201087).to_h { |id| ["<|reserved_#{id}|>", id] }.freeze
+    private_constant :HARMONY_RESERVED_TAIL
+
+    REGISTRY = {
+      "r50k_base" => {
+        rank_file: File.join(DATA_DIR, "r50k_base.tiktoken"),
+        pretokenizer: "gpt2",
+        special_tokens: {"<|endoftext|>" => 50256}
+      },
+      "cl100k_base" => {
+        rank_file: File.join(DATA_DIR, "cl100k_base.tiktoken"),
+        pretokenizer: "gpt4",
+        special_tokens: {
+          "<|endoftext|>" => 100257,
+          "<|fim_prefix|>" => 100258,
+          "<|fim_middle|>" => 100259,
+          "<|fim_suffix|>" => 100260,
+          "<|endofprompt|>" => 100276
+        }
+      },
+      "o200k_base" => {
+        rank_file: File.join(DATA_DIR, "o200k_base.tiktoken"),
+        pretokenizer: "o200k",
+        special_tokens: {"<|endoftext|>" => 199999, "<|endofprompt|>" => 200018}
+      },
+      "o200k_harmony" => {
+        rank_file: File.join(DATA_DIR, "o200k_base.tiktoken"),
+        pretokenizer: "o200k",
+        special_tokens: HARMONY_HEAD_TOKENS.merge(HARMONY_RESERVED_TAIL).freeze
+      }
+    }.freeze
+    private_constant :REGISTRY
+
+    # The packaged encoding names — the single source error messages naming
+    # what's available are built from (the same discipline
+    # Native.pretokenizer_names / PretokenizerType::NAMES applies to
+    # pretokenizer scheme names).
+    NAMES = REGISTRY.keys.freeze
+
+    # Encodings known by name but deliberately not packaged, keyed to the
+    # reason a caller asking for one by name deserves to hear.
+    UNPACKABLE_REASONS = {
+      "p50k_base" => "its ranks are not dense (50256 is left free for <|endoftext|>), " \
+        "and the rank loader rejects non-dense ranks",
+      "p50k_edit" => "it loads the same p50k_base.tiktoken ranks, which are not dense " \
+        "(50256 is left free for <|endoftext|>), and the rank loader rejects non-dense ranks"
+    }.freeze
+    private_constant :UNPACKABLE_REASONS
+
+    class << self
+      # The {rank_file:, pretokenizer:, special_tokens:} registered for a
+      # packaged encoding name, or nil.
+      def [](name)
+        REGISTRY[name]
+      end
+
+      # Why `name` can't be packaged, or nil when there's no reason on
+      # record (it's either packaged, or simply not one gigatoken knows of).
+      def unpackable_reason(name)
+        UNPACKABLE_REASONS[name]
+      end
+    end
+  end
+end

data/lib/gigatoken/tokenizer.rb

@@ -7,9 +7,6 @@ module Gigatoken
   # over a native `Gigatoken::Native::BPETokenizer` or
   # `Gigatoken::Native::SentencePieceTokenizer`.
   class Tokenizer
-    TIKTOKEN_ENDOFTEXT = "<|endoftext|>"
-    private_constant :TIKTOKEN_ENDOFTEXT
-
     FILE_SOURCE_CLASSES = [Native::TextFileSource, Native::JsonlFileSource, Native::ParquetFileSource].freeze
     private_constant :FILE_SOURCE_CLASSES
 
@@ -27,10 +24,27 @@ module Gigatoken
       from_json(File.binread(path))
     end
 
-    # Load from a .tiktoken mergeable-ranks file.
-    def self.from_tiktoken(path)
-      native = Native::BPETokenizer.from_tiktoken(path.to_s)
-      new(native, special_tokens: {TIKTOKEN_ENDOFTEXT => native.vocab_size - 1})
+    # Load from a .tiktoken mergeable-ranks file. The file carries neither a
+    # pretokenization scheme nor special tokens — nothing is guessed here,
+    # so `pretokenizer:` is required: one of the schemes gigatoken ships
+    # (see Native.pretokenizer_names, e.g. "gpt2"/"r50k", "gpt4"/"cl100k",
+    # "o200k", "qwen2", "qwen35", "olmo3", "deepseek_v3", "nemotron", "kimi").
+    # `special_tokens:` maps token content to id (none by default).
+    def self.from_tiktoken(path, pretokenizer:, special_tokens: {})
+      native = Native::BPETokenizer.from_tiktoken(path.to_s, pretokenizer, special_tokens)
+      new(native, special_tokens: special_tokens)
+    end
+
+    # Load one of the tiktoken encodings gigatoken vendors ranks for, by
+    # name — see Gigatoken::Encodings::NAMES — entirely from the vendored
+    # files: no network, no writable cache.
+    def self.from_encoding(name)
+      encoding = Encodings[name]
+      return from_tiktoken(encoding[:rank_file], pretokenizer: encoding[:pretokenizer], special_tokens: encoding[:special_tokens]) if encoding
+
+      reason = Encodings.unpackable_reason(name)
+      detail = reason ? " — #{reason}" : ""
+      raise Error, "#{name.inspect}: not a packaged encoding#{detail} (packaged encodings: #{Encodings::NAMES.join(", ")})"
     end
 
     # Load tokenizer.json from HuggingFace Hub repo `repo_id` at `revision`
@@ -41,12 +55,29 @@ module Gigatoken
 
     # Load from any of the supported source shapes: an existing file or
     # directory path (a tokenizer.json, or a directory containing one), a
-    # .tiktoken vocabulary file, or a HuggingFace Hub repo id like
-    # "openai-community/gpt2".
-    def self.load(source, revision: "main", hub: Hub.new)
+    # .tiktoken vocabulary file, a packaged encoding name (see
+    # Gigatoken::Encodings::NAMES, e.g. "cl100k_base"), or a HuggingFace Hub
+    # repo id like "openai-community/gpt2". A .tiktoken file carries no
+    # pretokenizer scheme of its own, so one must be named explicitly via
+    # `pretokenizer:` — nothing here is guessed. Packaged encoding names are
+    # checked before the Hub-repo-id shape: a bare name like "o200k_base" is
+    # also shaped like a legacy repo id, and must resolve locally rather
+    # than reach the network. Names the registry knows but doesn't package
+    # (see Encodings.unpackable_reason, e.g. "p50k_base") are intercepted
+    # here too, raising the same explanation from_encoding gives rather than
+    # reaching the Hub — but only those; an unrecognized bare name like
+    # "gpt2" still dispatches to the Hub.
+    def self.load(source, pretokenizer: nil, special_tokens: {}, revision: "main", hub: Hub.new)
       source = source.to_s
-      return from_tiktoken(source) if source.end_with?(".tiktoken")
+      if source.end_with?(".tiktoken")
+        unless pretokenizer
+          raise Error, "#{source.inspect}: a .tiktoken file carries no pretokenizer scheme of its own — " \
+            "pass pretokenizer: (one of #{Native.pretokenizer_names.join(", ")})"
+        end
+        return from_tiktoken(source, pretokenizer: pretokenizer, special_tokens: special_tokens)
+      end
       return from_file(source) if File.exist?(source)
+      return from_encoding(source) if Encodings::NAMES.include?(source) || Encodings.unpackable_reason(source)
       return from_hub(source, revision: revision, hub: hub) if Hub.looks_like_repo_id?(source)
 
       raise Error, "#{source.inspect}: no such file or directory, not a .tiktoken path, and doesn't look like a HuggingFace Hub repo id"
@@ -112,6 +143,12 @@ module Gigatoken
       @native.merges
     end
 
+    # Cached pretoken/unit entries on this tokenizer's single-document
+    # encode path — see Gigatoken.max_cache_bytes.
+    def cache_entries
+      @native.cache_entries
+    end
+
     attr_reader :special_tokens
   end
 end

data/lib/gigatoken/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module Gigatoken
-  VERSION = "0.1.0"
+  VERSION = "0.2.0"
 end

data/lib/gigatoken.rb

@@ -7,6 +7,20 @@ module Gigatoken
   # extension — never a raw Rust panic across the Ruby boundary.
   class Error < StandardError; end
 
+  class << self
+    # The process-global encode-cache budget in bytes per worker (a parallel
+    # batch encode may use up to workers x budget), applied to tokenizers of
+    # either backend constructed *afterward* — changing it never affects an
+    # already-built Tokenizer. nil means unbounded. Default: 512 MiB.
+    def max_cache_bytes
+      Native.get_max_cache_bytes
+    end
+
+    def max_cache_bytes=(bytes)
+      Native.set_max_cache_bytes(bytes)
+    end
+  end
+
   NATIVE_EXTENSIONS = %w[.bundle .so .rb].freeze
 
   # Precompiled native gems ship per-ABI subdirs (`gigatoken/4.0/...`),
@@ -26,4 +40,5 @@ require native
 
 require_relative "gigatoken/hub"
 require_relative "gigatoken/packed_result"
+require_relative "gigatoken/encodings"
 require_relative "gigatoken/tokenizer"

data/src/batch.rs

@@ -600,13 +600,8 @@ impl Committer {
         unsafe {
             flat.set_len(total);
         }
-        // Return the unused reservation. Large allocations trim in place on
-        // the mainstream allocators (macOS libmalloc large entries, glibc
-        // mmap'd chunks via mremap): pointer-stable, no copy, and the
-        // untouched tail pages were never faulted so there is nothing to
-        // tear down. An allocator that copies instead only costs one
-        // memcpy; correctness is unaffected.
-        flat.shrink_to_fit();
+        // The unused reservation tail is unfaulted virtual address space —
+        // freed with the Vec on drop, no trim needed to reclaim it.
         Some(flat)
     }
 

data/src/bindings/cache.rs

@@ -0,0 +1,38 @@
+//! The process-global cache-budget knob, read once per tokenizer
+//! construction. Semantics in the .pyi stubs, mechanism in
+//! `Tokenizer::set_max_cache_bytes`.
+
+use crate::bpe::{SentencePieceBPE, Tokenizer};
+use pyo3::prelude::*;
+use std::sync::Mutex;
+
+/// The budget applied to tokenizers constructed after the last
+/// `set_max_cache_bytes` call; `None` = unbounded.
+static MAX_CACHE_BYTES: Mutex<Option<usize>> = Mutex::new(Some(Tokenizer::DEFAULT_MAX_CACHE_BYTES));
+
+/// Apply the global setting to a freshly constructed tokenizer (which
+/// already carries the built-in default, hence the `!=` skip).
+pub(crate) fn apply_max_cache_bytes(mut tokenizer: Tokenizer) -> Tokenizer {
+    let configured = *MAX_CACHE_BYTES.lock().unwrap();
+    if configured != tokenizer.max_cache_bytes() {
+        tokenizer.set_max_cache_bytes(configured);
+    }
+    tokenizer
+}
+
+/// SentencePiece analog of [`apply_max_cache_bytes`].
+pub(crate) fn apply_max_cache_bytes_sp(mut model: SentencePieceBPE) -> SentencePieceBPE {
+    model.set_max_cache_bytes(*MAX_CACHE_BYTES.lock().unwrap());
+    model
+}
+
+#[pyfunction]
+#[pyo3(signature = (max_bytes))]
+pub(crate) fn set_max_cache_bytes(max_bytes: Option<usize>) {
+    *MAX_CACHE_BYTES.lock().unwrap() = max_bytes;
+}
+
+#[pyfunction]
+pub(crate) fn get_max_cache_bytes() -> Option<usize> {
+    *MAX_CACHE_BYTES.lock().unwrap()
+}

data/src/bindings/mod.rs

@@ -3,9 +3,11 @@
 //! submodule per feature: Python<->Rust bridging shared by the bindings
 //! (bridge), the FileSource/BytesSource classes (sources), BPE training (train), padded
 //! compat-API encoding (padding), the compat layers' special-token scanner
-//! (matcher), and the pretokenizer helpers (pretokenize).
+//! (matcher), the pretokenizer helpers (pretokenize), and the
+//! process-global cache-budget knob (cache).
 
 pub(crate) mod bridge;
+pub(crate) mod cache;
 pub(crate) mod hub;
 pub(crate) mod matcher;
 pub(crate) mod padding;

data/src/bindings/pretokenize.rs

@@ -29,6 +29,17 @@ impl PretokenizerIter {
     }
 }
 
+/// The pretokenization scheme a loader was given by name, for the sources
+/// whose format carries no split regex (tiktoken rank files).
+pub(crate) fn pretokenizer_scheme(name: &str) -> PyResult<pretokenize::PretokenizerType> {
+    pretokenize::PretokenizerType::from_name(name).ok_or_else(|| {
+        pyo3::exceptions::PyValueError::new_err(format!(
+            "unknown pretokenizer scheme {name:?}; expected one of {}",
+            pretokenize::PretokenizerType::NAMES.join(", ")
+        ))
+    })
+}
+
 #[pyfunction]
 pub(crate) fn pretokenizer<'py>(text: Bound<'py, PyBytes>) -> PyResult<PretokenizerIter> {
     Ok(PretokenizerIter {

data/src/bpe/pretoken_cache.rs

@@ -178,11 +178,42 @@ impl ShortPretokenCache {
     /// 3/4 load. Either way the table is constructed exactly once, at the
     /// max of the two requirements.
     pub(crate) fn with_at_least(n: usize, min_slots: usize) -> Self {
+        Self::with_pow2_capacity(Self::required_capacity(n, min_slots))
+    }
+
+    /// The slot count [`Self::with_at_least`] would allocate for `n`
+    /// entries starting from `min_slots`, without allocating anything.
+    pub(crate) fn required_capacity(n: usize, min_slots: usize) -> usize {
         let mut cap = min_slots.max(1 << 16).next_power_of_two();
         while (n + 1) * 4 > cap * 3 {
             cap *= 2;
         }
-        Self::with_pow2_capacity(cap)
+        cap
+    }
+
+    /// Zero every slot in place (all-zero == all-empty), keeping the
+    /// allocation — the generational-wipe path. Like [`Self::insert`]'s
+    /// grow, this logically invalidates any outstanding [`ProbeView`].
+    pub(crate) fn clear(&mut self) {
+        // SAFETY: the allocation holds exactly `cap` entries.
+        unsafe { std::ptr::write_bytes(self.slots.ptr.as_ptr(), 0, self.slots.cap) };
+        self.len = 0;
+    }
+
+    /// Reset to an empty table of exactly `cap` slots (power of two,
+    /// >= 2), reusing the current allocation when the capacity matches.
+    pub(crate) fn reset_to_capacity(&mut self, cap: usize) {
+        debug_assert!(cap.is_power_of_two() && cap >= 2);
+        if cap == self.slots.cap {
+            self.clear();
+            return;
+        }
+        // Drop the old table (via a 2-slot placeholder) before building
+        // the new one, so two full-size tables never coexist.
+        self.slots = Slots::new_zeroed(2);
+        self.slots = Slots::new_zeroed(cap);
+        self.mask = cap - 1;
+        self.len = 0;
     }
 
     /// Address of `h`'s home pair; both slots share the addressed line.

data/src/bpe/sentencepiece.rs

@@ -296,6 +296,8 @@ pub struct SentencePieceBPE {
     /// Bitset over the byte just before a candidate boundary (the last byte
     /// of some `cross_pieces` pre) for a one-load rejection in the scanner.
     pub(crate) cross_prev: [u64; 4],
+    /// Cache budget snapshot by each [`Self::encoder`]; `None` = unbounded.
+    pub(crate) max_cache_bytes: Option<usize>,
 }
 
 /// How many distinct punctuation bytes the unit splitter checks for.
@@ -764,10 +766,22 @@ impl SentencePieceBPE {
     pub fn encoder(&self) -> Encoder<'_> {
         Encoder {
             model: self,
-            state: EncodeState::new(),
+            state: EncodeState::with_budget(self.max_cache_bytes),
         }
     }
 
+    /// SentencePiece analog of [`super::Tokenizer::set_max_cache_bytes`],
+    /// same default: bounds the caches of [`EncodeState`]s created
+    /// afterward, which wipe back to empty (there is no vocab seed here).
+    pub fn set_max_cache_bytes(&mut self, budget: Option<usize>) {
+        self.max_cache_bytes = budget;
+    }
+
+    /// The configured cache budget in bytes; `None` when unbounded.
+    pub fn max_cache_bytes(&self) -> Option<usize> {
+        self.max_cache_bytes
+    }
+
     /// Size of the vocabulary: one greater than the largest token ID,
     /// including added tokens (IDs with no assigned content count too).
     pub fn vocab_size(&self) -> usize {
@@ -854,6 +868,14 @@ fn collapse_space_runs(input: &str, content: &str) -> String {
 /// Zipf-hot working set resident despite eviction churn from tail units.
 const FRONT_BITS: u32 = 20;
 
+/// Fixed footprint of the front cache (keys + vals), subtracted from the
+/// budget before bounding the growing caches.
+const FRONT_BYTES: usize = (1 << FRONT_BITS) * (16 + 8);
+/// Estimated bytes per map entry, payload + SwissTable slack (`long` adds
+/// its key bytes on top). Length-based, because the wipe keeps allocations.
+const SHORT_ENTRY_BYTES: usize = 48;
+const LONG_ENTRY_BYTES: usize = 48;
+
 /// Index of `key` in the front cache: multiplicative hash, top bits.
 #[inline(always)]
 fn front_index(key: u128) -> usize {
@@ -882,10 +904,26 @@ pub struct EncodeState {
     symbols: Vec<TokenId>,
     /// Scratch for composing keys with a virtual leading space.
     key_buf: Vec<u8>,
+    /// Byte budget for the growing caches ([`Self::with_budget`]);
+    /// `usize::MAX` = unbounded.
+    budget: usize,
+    /// Estimated `long` bytes: key bytes + [`LONG_ENTRY_BYTES`] per entry.
+    long_bytes_used: usize,
+    /// Largest single encoding, as wipe-trigger slack: one recurring giant
+    /// unit must not force a wipe per occurrence. Kept across wipes.
+    max_encoding: usize,
 }
 
 impl EncodeState {
     pub fn new() -> Self {
+        Self::with_budget(Some(super::Tokenizer::DEFAULT_MAX_CACHE_BYTES))
+    }
+
+    /// A state whose growing caches (short/long maps + arena; the front
+    /// cache is fixed-size) wipe back to empty when their estimated
+    /// footprint exceeds `max_bytes` minus the front cache, floored at
+    /// 1 MiB so tiny budgets stay functional.
+    pub fn with_budget(max_bytes: Option<usize>) -> Self {
         EncodeState {
             arena: Vec::new(),
             front_keys: vec![0u128; 1 << FRONT_BITS],
@@ -894,6 +932,10 @@ impl EncodeState {
             long: HashMap::with_hasher(FxBuildHasher),
             symbols: Vec::new(),
             key_buf: Vec::new(),
+            budget: max_bytes
+                .map_or(usize::MAX, |t| t.saturating_sub(FRONT_BYTES).max(1 << 20)),
+            long_bytes_used: 0,
+            max_encoding: 0,
         }
     }
 
@@ -901,6 +943,24 @@ impl EncodeState {
     pub fn cache_size(&self) -> usize {
         self.short.len() + self.long.len()
     }
+
+    #[inline]
+    fn over_budget(&self) -> bool {
+        self.short.len() * SHORT_ENTRY_BYTES + self.long_bytes_used + self.arena.len() * 4
+            > self.budget.saturating_add(self.max_encoding * 4)
+    }
+
+    /// Wipe every cache back to empty, keeping allocations. The front keys
+    /// must go too: their values are slices of the truncated arena.
+    #[cold]
+    #[inline(never)]
+    fn wipe(&mut self) {
+        self.arena.clear();
+        self.front_keys.fill(0);
+        self.short.clear();
+        self.long.clear();
+        self.long_bytes_used = 0;
+    }
 }
 
 impl Default for EncodeState {
@@ -1330,7 +1390,8 @@ impl SentencePieceBPE {
                 let (offset, len) = unsafe { *state.front_vals.get_unchecked(front_idx) };
                 let start = offset as usize;
                 // SAFETY: entries are recorded right after appending `len`
-                // tokens at `offset`; `arena` never shrinks.
+                // tokens at `offset`; the arena only clears in a budget
+                // wipe, which also zeroes every front key.
                 f(unsafe { state.arena.get_unchecked(start..start + len as usize) });
                 return;
             }
@@ -1347,13 +1408,18 @@ impl SentencePieceBPE {
             }
             let start = offset as usize;
             // SAFETY: every cached (offset, len) was recorded right after
-            // appending those `len` tokens at `offset`, and `arena` never
-            // shrinks, so the range is always in bounds.
+            // appending those `len` tokens at `offset`; the arena only
+            // clears in a budget wipe, which also clears both maps.
             f(unsafe { state.arena.get_unchecked(start..start + len as usize) });
             return;
         }
 
-        // Miss: character init → ranked merge, then record in the arena.
+        // Miss: budget check first (the lookups above already missed, and
+        // stay misses against the freshly wiped caches), then character
+        // init → ranked merge, then record in the arena.
+        if state.over_budget() {
+            state.wipe();
+        }
         state.symbols.clear();
         if virtual_prefix {
             state.symbols.extend_from_slice(&self.space_init);
@@ -1367,6 +1433,7 @@ impl SentencePieceBPE {
         let offset = state.arena.len() as u32;
         let len = state.symbols.len() as u32;
         state.arena.extend_from_slice(&state.symbols);
+        state.max_encoding = state.max_encoding.max(len as usize);
         match packed {
             Some(key) => {
                 state.short.insert(key, (offset, len));
@@ -1379,6 +1446,7 @@ impl SentencePieceBPE {
                 } else {
                     unit.into()
                 };
+                state.long_bytes_used += key.len() + LONG_ENTRY_BYTES;
                 state.long.insert(key, (offset, len));
             }
         }
@@ -1482,4 +1550,68 @@ mod tests {
         assert_eq!(collapse_space_runs("  ", "▁"), "▁");
         assert_eq!(collapse_space_runs("no runs", "▁"), "no runs");
     }
+
+    /// PARITY + BOUNDS: a budgeted EncodeState that wipes mid-corpus must
+    /// produce the exact token stream of an unbounded one, while the
+    /// estimated footprint of the growing caches stays within the budget
+    /// (plus the documented giant-encoding slack).
+    #[test]
+    fn budgeted_wipe_matches_unbounded() {
+        let Some(path) = crate::test_hub::hf_tokenizer_json("TinyLlama/TinyLlama-1.1B-Chat-v1.0")
+        else {
+            eprintln!("Skipping: TinyLlama tokenizer.json not in the HF cache");
+            return;
+        };
+        let model = crate::load_tokenizer::hf::load_hf_sentencepiece(&path).unwrap();
+        assert_eq!(
+            model.max_cache_bytes(),
+            Some(super::super::Tokenizer::DEFAULT_MAX_CACHE_BYTES)
+        );
+
+        // ~100k distinct short words plus interleaved > 15-byte words (the
+        // long-map path) and repeats of a common head.
+        let mut text = String::new();
+        let mut x = 0x1234_5678_9ABC_DEF0u64;
+        let mut rand = |m: u64| {
+            x ^= x << 13;
+            x ^= x >> 7;
+            x ^= x << 17;
+            x % m
+        };
+        for i in 0..100_000u32 {
+            let len = if i % 17 == 0 { 16 + rand(24) } else { 3 + rand(8) };
+            for _ in 0..len {
+                text.push((b'a' + rand(26) as u8) as char);
+            }
+            text.push(' ');
+            if i % 5 == 0 {
+                text.push_str("the quick brown fox ");
+            }
+        }
+
+        let mut unbounded = EncodeState::with_budget(None);
+        let mut expected = Vec::new();
+        model.encode_raw_with(&mut unbounded, &text, &mut expected);
+
+        // 25 MiB total = ~1 MiB effective past the 24 MiB front cache:
+        // several wipes over ~100k distinct units.
+        let mut budgeted = EncodeState::with_budget(Some(25 << 20));
+        let mut actual = Vec::new();
+        model.encode_raw_with(&mut budgeted, &text, &mut actual);
+        assert_eq!(actual, expected, "budgeted output diverged");
+        assert!(budgeted.long_bytes_used > 0, "corpus never hit the long map");
+        // Both saw ~100k distinct units, so a small survivor count proves
+        // wipes happened, and the estimate must end up within budget.
+        assert!(
+            budgeted.cache_size() < unbounded.cache_size() / 3,
+            "survivors {} vs unbounded {}",
+            budgeted.cache_size(),
+            unbounded.cache_size()
+        );
+        let used = budgeted.short.len() * SHORT_ENTRY_BYTES
+            + budgeted.long_bytes_used
+            + budgeted.arena.len() * 4;
+        let limit = budgeted.budget + budgeted.max_encoding * 4;
+        assert!(used <= limit + 4096, "estimate {used} exceeds budget {limit}");
+    }
 }