From 5e27c5b4631124999d56de7825e9caa42ed15b0f Mon Sep 17 00:00:00 2001 From: Esteban Ordano Date: Sun, 1 Mar 2026 16:48:17 -0300 Subject: [PATCH] feat: compile-time trie tokenizer, mmap file reads MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit gen_vocab.py builds a double-array tree from the vocabulary at code-gen time. At runtime there is nothing to initialize — just two static uint32 arrays and a greedy longest-match loop. File reads use mmap instead of ifstream. Benchmark (1MB of War and Peace, hyperfine -N, 500 runs): optimized upstream 1 B 260 µs 91 ms (350×) 1 MB 5.6 ms 258 ms (46×) Co-Authored-By: Claude Opus 4.6 --- .bazelrc | 1 + .gitignore | 2 + BUILD.bazel | 8 +-- ctoc.cc | 155 +++++++++++++++++++---------------------- ctoc_smoke_test.cc | 40 ++++++++++- gen_vocab.py | 170 +++++++++++++++++++++++++++++++++------------ 6 files changed, 237 insertions(+), 139 deletions(-) diff --git a/.bazelrc b/.bazelrc index 2574bf3..104940e 100644 --- a/.bazelrc +++ b/.bazelrc @@ -6,6 +6,7 @@ build:macos --sandbox_add_mount_pair=/var/tmp build --action_env=BAZEL_DO_NOT_DETECT_CPP_TOOLCHAIN=1 build --cxxopt=-std=c++17 build --host_cxxopt=-std=c++17 +build -c opt # Cross-compilation build:linux_amd64 --platforms=@zig_sdk//platform:linux_amd64 diff --git a/.gitignore b/.gitignore index cb8d71a..44188d5 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,5 @@ bazel-* +vocab_data.cc +vocab_data.h *.jsonl *.pt diff --git a/BUILD.bazel b/BUILD.bazel index 2921030..a457f01 100644 --- a/BUILD.bazel +++ b/BUILD.bazel @@ -13,15 +13,11 @@ genrule( cc_binary( name = "ctoc", - srcs = [ - "ctoc.cc", - "vocab_data.cc", - "vocab_data.h", - ], + srcs = ["ctoc.cc", ":gen_vocab"], ) cc_test( name = "ctoc_smoke_test", size = "small", - srcs = ["ctoc_smoke_test.cc"], + srcs = ["ctoc_smoke_test.cc", ":gen_vocab"], ) diff --git a/ctoc.cc b/ctoc.cc index a2f2c28..4f7b246 100644 --- a/ctoc.cc +++ b/ctoc.cc @@ -1,22 +1,26 @@ // ctoc — Count Tokens of Code // Like cloc, but for Claude tokens. // -// Uses a greedy longest-match tokenizer built from a reverse-engineered -// vocabulary of 36,495 verified Claude tokens (95-96% accuracy). +// Uses a greedy longest-match tokenizer built into a double-array tree, +// from a reverse-engineered vocabulary of 38,360 verified Claude tokens +// (95-96% accuracy). #include #include #include #include -#include #include #include -#include #include #include #include #include +#include +#include +#include +#include + #include "vocab_data.h" namespace fs = std::filesystem; @@ -38,72 +42,45 @@ static const std::unordered_set DEFAULT_EXCLUDED_DIRS = { static constexpr size_t MAX_FILE_SIZE = 1 * 1024 * 1024; // 1 MB static constexpr size_t BINARY_CHECK_SIZE = 8192; -// ─── Trie ──────────────────────────────────────────────────────────── +// ─── Tokenizer ────────────────────────────────────────────────────── -struct TrieNode { - std::unordered_map children; - bool is_terminal = false; +static constexpr uint32_t TERM_BIT = 0x80000000u; +static constexpr uint32_t IDX_MASK = 0x7FFFFFFFu; - ~TrieNode() { - for (auto& [_, child] : children) - delete child; - } -}; - -class Trie { -public: - Trie() : root_(new TrieNode()) {} - ~Trie() { delete root_; } - - Trie(const Trie&) = delete; - Trie& operator=(const Trie&) = delete; - - void insert(const std::string& token) { - TrieNode* node = root_; - for (unsigned char c : token) { - auto it = node->children.find(c); - if (it == node->children.end()) { - node->children[c] = new TrieNode(); - node = node->children[c]; - } else { - node = it->second; - } - } - node->is_terminal = true; - } +// Returns the longest match length starting at data[pos], minimum 1. +static inline size_t match_len(const char* data, size_t len, size_t pos) { + unsigned char byte = static_cast(data[pos]); + uint32_t t = DA_BASE[0] + byte; + if (t >= DA_TRIE_SIZE) + return 1; + uint32_t c = DA_CHECK[t]; + if (c == 0xFFFFFFFFu || (c & IDX_MASK) != 0) + return 1; - // Returns the length of the longest match starting at data[pos], or 0. - size_t longest_match(const std::string& data, size_t pos) const { - TrieNode* node = root_; - size_t best = 0; - for (size_t i = pos; i < data.size(); ++i) { - auto it = node->children.find(static_cast(data[i])); - if (it == node->children.end()) - break; - node = it->second; - if (node->is_terminal) - best = i - pos + 1; - } - return best; + size_t best = (c & TERM_BIT) ? 1 : 0; + uint32_t state = t; + + for (size_t i = pos + 1; i < len; ++i) { + byte = static_cast(data[i]); + t = DA_BASE[state] + byte; + if (t >= DA_TRIE_SIZE) + break; + c = DA_CHECK[t]; + if (c == 0xFFFFFFFFu || (c & IDX_MASK) != state) + break; + state = t; + if (c & TERM_BIT) + best = i - pos + 1; } -private: - TrieNode* root_; -}; - -// ─── Tokenizer ─────────────────────────────────────────────────────── + return best ? best : 1; +} -static size_t count_tokens(const std::string& text, const Trie& trie) { +static size_t count_tokens(const char* data, size_t len) { size_t count = 0; size_t pos = 0; - while (pos < text.size()) { - size_t match_len = trie.longest_match(text, pos); - if (match_len == 0) { - // Unknown byte — count as 1 token (single-byte fallback) - ++pos; - } else { - pos += match_len; - } + while (pos < len) { + pos += match_len(data, len, pos); ++count; } return count; @@ -111,8 +88,8 @@ static size_t count_tokens(const std::string& text, const Trie& trie) { // ─── File discovery ────────────────────────────────────────────────── -static bool is_binary(const std::string& data) { - size_t check_len = std::min(data.size(), BINARY_CHECK_SIZE); +static bool is_binary(const char* data, size_t len) { + size_t check_len = std::min(len, BINARY_CHECK_SIZE); for (size_t i = 0; i < check_len; ++i) { if (data[i] == '\0') return true; @@ -120,13 +97,27 @@ static bool is_binary(const std::string& data) { return false; } -static std::string read_file(const fs::path& path) { - std::ifstream f(path, std::ios::binary); - if (!f) - return {}; - std::ostringstream ss; - ss << f.rdbuf(); - return ss.str(); +struct MappedFile { + const char* data = nullptr; + size_t size = 0; + int fd = -1; + void close() { + if (data) munmap(const_cast(data), size); + if (fd >= 0) ::close(fd); + data = nullptr; fd = -1; + } + ~MappedFile() { close(); } + explicit operator bool() const { return data != nullptr; } +}; + +static MappedFile map_file(const fs::path& path) { + int fd = open(path.c_str(), O_RDONLY); + if (fd < 0) return {}; + struct stat st; + if (fstat(fd, &st) < 0 || st.st_size == 0) { ::close(fd); return {}; } + auto* p = mmap(nullptr, st.st_size, PROT_READ, MAP_PRIVATE, fd, 0); + if (p == MAP_FAILED) { ::close(fd); return {}; } + return {static_cast(p), static_cast(st.st_size), fd}; } struct FileEntry { @@ -159,8 +150,7 @@ static bool is_bazel_dir(const fs::path& path) { static std::vector discover_files( const std::vector& paths, const std::unordered_set& excluded_dirs, - const std::unordered_set& include_exts, - const Trie& trie) + const std::unordered_set& include_exts) { std::vector files; @@ -181,12 +171,12 @@ static std::vector discover_files( std::string ext = get_ext(p); if (!include_exts.empty() && include_exts.find(ext) == include_exts.end()) continue; - std::string content = read_file(p); - if (content.empty() || is_binary(content)) + auto mf = map_file(p); + if (!mf || is_binary(mf.data, mf.size)) continue; if (ext.empty()) ext = "(none)"; - files.push_back({p, ext, count_tokens(content, trie)}); + files.push_back({p, ext, count_tokens(mf.data, mf.size)}); continue; } @@ -226,11 +216,11 @@ static std::vector discover_files( if (!include_exts.empty() && include_exts.find(ext) == include_exts.end()) continue; - std::string content = read_file(it->path()); - if (content.empty() || is_binary(content)) + auto mf = map_file(it->path()); + if (!mf || is_binary(mf.data, mf.size)) continue; - files.push_back({it->path(), ext, count_tokens(content, trie)}); + files.push_back({it->path(), ext, count_tokens(mf.data, mf.size)}); } } @@ -432,13 +422,8 @@ int main(int argc, char* argv[]) { auto excluded_dirs = DEFAULT_EXCLUDED_DIRS; excluded_dirs.insert(extra_excluded_dirs.begin(), extra_excluded_dirs.end()); - // Build trie from embedded vocabulary - Trie trie; - for (size_t i = 0; i < VOCAB_COUNT; ++i) - trie.insert(VOCAB_TOKENS[i]); - // Discover and tokenize files - auto files = discover_files(input_paths, excluded_dirs, include_exts, trie); + auto files = discover_files(input_paths, excluded_dirs, include_exts); if (files.empty()) { std::cerr << "ctoc: no files found\n"; diff --git a/ctoc_smoke_test.cc b/ctoc_smoke_test.cc index 9463b97..3669045 100644 --- a/ctoc_smoke_test.cc +++ b/ctoc_smoke_test.cc @@ -1,5 +1,41 @@ -#include +#include +#include +#include +#include "vocab_data.h" + +static constexpr uint32_t TERM_BIT = 0x80000000u; +static constexpr uint32_t IDX_MASK = 0x7FFFFFFFu; + +static size_t match_len(const char* data, size_t len, size_t pos) { + unsigned char byte = static_cast(data[pos]); + uint32_t t = DA_BASE[0] + byte; + if (t >= DA_TRIE_SIZE) return 1; + uint32_t c = DA_CHECK[t]; + if (c == 0xFFFFFFFFu || (c & IDX_MASK) != 0) return 1; + size_t best = (c & TERM_BIT) ? 1 : 0; + uint32_t state = t; + for (size_t i = pos + 1; i < len; ++i) { + byte = static_cast(data[i]); + t = DA_BASE[state] + byte; + if (t >= DA_TRIE_SIZE) break; + c = DA_CHECK[t]; + if (c == 0xFFFFFFFFu || (c & IDX_MASK) != state) break; + state = t; + if (c & TERM_BIT) best = i - pos + 1; + } + return best ? best : 1; +} + +static size_t count(const char* s) { + size_t len = strlen(s), n = 0, pos = 0; + while (pos < len) { pos += match_len(s, len, pos); ++n; } + return n; +} int main() { - return EXIT_SUCCESS; + assert(DA_TRIE_SIZE > 0); + assert(count("hello") == 1); + assert(count("\x01") == 1); + assert(count("") == 0); + assert(count("Hello, world!") >= 1 && count("Hello, world!") <= 10); } diff --git a/gen_vocab.py b/gen_vocab.py index 3ebd496..cd4e97a 100644 --- a/gen_vocab.py +++ b/gen_vocab.py @@ -1,52 +1,120 @@ #!/usr/bin/env python3 """Generate vocab_data.cc and vocab_data.h from vocab.json. -Reads the "verified" array from the vocab JSON and emits a C++ source file -containing the token strings as a compile-time array, so the binary is -self-contained and doesn't need the JSON at runtime. +Builds a double-array tree at code-gen time and emits it as static C++ arrays. +At runtime, the tree is ready to use with zero initialization — just array +lookups, no heap allocation, no hash maps. """ import json import sys +from collections import deque -def c_escape(s: str) -> str: - """Escape a Python string for embedding as a C string literal. +TERM_BIT = 0x80000000 +IDX_MASK = 0x7FFFFFFF - Uses octal escapes for non-ASCII bytes to avoid the C++ problem where - \\xNN followed by a hex digit gets parsed as a longer hex literal. + +def build_da_trie(tokens): + """Build a double-array tree from a list of token strings. + + Returns (base, check, array_size) where base and check are lists of uint32. """ - out = [] - for ch in s: - o = ord(ch) - if ch == '\\': - out.append('\\\\') - elif ch == '"': - out.append('\\"') - elif ch == '\n': - out.append('\\n') - elif ch == '\r': - out.append('\\r') - elif ch == '\t': - out.append('\\t') - elif ch == '\0': - out.append('\\0') - elif ch == '\a': - out.append('\\a') - elif ch == '\b': - out.append('\\b') - elif ch == '\f': - out.append('\\f') - elif ch == '\v': - out.append('\\v') - elif 0x20 <= o <= 0x7E: - out.append(ch) - else: - # Emit raw UTF-8 bytes as octal escapes (max 3 digits, so no - # ambiguity with following characters unlike hex escapes). - for b in ch.encode('utf-8'): - out.append(f'\\{b:03o}') - return ''.join(out) + # Phase 1: Build simple trie + node_children = [[]] # list of (byte, child_index) per node + node_terminal = [False] + + for token in tokens: + cur = 0 + for byte in token.encode('utf-8'): + existing = None + for k, idx in node_children[cur]: + if k == byte: + existing = idx + break + if existing is not None: + cur = existing + else: + idx = len(node_children) + node_children.append([]) + node_terminal.append(False) + node_children[cur].append((byte, idx)) + cur = idx + node_terminal[cur] = True + + # Sort children by byte value + for children in node_children: + children.sort(key=lambda x: x[0]) + + num_nodes = len(node_children) + + # Phase 2: Convert to double-array via BFS + initial_size = num_nodes + 512 + base = [0] * initial_size + check = [0xFFFFFFFF] * initial_size + occupied = [False] * initial_size + + da_pos = [0] * num_nodes + da_pos[0] = 0 # root at position 0 + occupied[0] = True + + def find_base(keys): + n = len(occupied) + first_key = keys[0] + b = 0 + while True: + fpos = b + first_key + if fpos < n and occupied[fpos]: + b += 1 + continue + ok = True + for k in keys[1:]: + pos = b + k + if pos < n and occupied[pos]: + ok = False + break + if ok: + return b + b += 1 + + queue = deque([0]) + + while queue: + trie_node = queue.popleft() + s = da_pos[trie_node] + ch = node_children[trie_node] + + if not ch: + continue + + keys = [k for k, _ in ch] + b = find_base(keys) + + max_pos = b + 256 + if max_pos >= len(base): + new_size = max_pos + 512 + base.extend([0] * (new_size - len(base))) + check.extend([0xFFFFFFFF] * (new_size - len(check))) + occupied.extend([False] * (new_size - len(occupied))) + + base[s] = b + + for key, child_trie_idx in ch: + t = b + key + term = TERM_BIT if node_terminal[child_trie_idx] else 0 + check[t] = s | term + occupied[t] = True + da_pos[child_trie_idx] = t + queue.append(child_trie_idx) + + # Trim to actual size + actual_size = 0 + for i in range(len(occupied) - 1, -1, -1): + if occupied[i]: + actual_size = i + 1 + break + + return base[:actual_size], check[:actual_size], actual_size def main(): @@ -61,26 +129,36 @@ def main(): tokens = data['verified'] + base, check, array_size = build_da_trie(tokens) + # Write header with open(h_path, 'w', encoding='utf-8') as f: f.write('#ifndef VOCAB_DATA_H_\n') f.write('#define VOCAB_DATA_H_\n\n') + f.write('#include \n') f.write('#include \n\n') - f.write(f'extern const char* const VOCAB_TOKENS[{len(tokens)}];\n') - f.write(f'extern const size_t VOCAB_COUNT;\n\n') + f.write(f'static constexpr size_t DA_TRIE_SIZE = {array_size};\n') + f.write(f'extern const uint32_t DA_BASE[{array_size}];\n') + f.write(f'extern const uint32_t DA_CHECK[{array_size}];\n\n') f.write('#endif // VOCAB_DATA_H_\n') # Write source with open(cc_path, 'w', encoding='utf-8') as f: f.write('#include "vocab_data.h"\n\n') - f.write(f'const char* const VOCAB_TOKENS[{len(tokens)}] = {{\n') - for i, token in enumerate(tokens): - escaped = c_escape(token) - f.write(f' "{escaped}",\n') + + f.write(f'const uint32_t DA_BASE[{array_size}] = {{\n') + for i in range(0, array_size, 16): + chunk = base[i:i+16] + f.write(' ' + ','.join(str(v) for v in chunk) + ',\n') f.write('};\n\n') - f.write(f'const size_t VOCAB_COUNT = {len(tokens)};\n') - print(f"Generated {cc_path} and {h_path} with {len(tokens)} tokens", file=sys.stderr) + f.write(f'const uint32_t DA_CHECK[{array_size}] = {{\n') + for i in range(0, array_size, 16): + chunk = check[i:i+16] + f.write(' ' + ','.join(str(v) for v in chunk) + ',\n') + f.write('};\n') + + print(f"Generated {len(tokens)} tokens, {array_size} slots", file=sys.stderr) if __name__ == '__main__':