Skip to main content

logicaffeine_compile/model/tokenizer/
cxx.rs

1//! The libstdc++ behaviors llama.cpp's vocabulary inherits, ported step for step.
2//!
3//! [`sort`] is `std::sort`: llama.cpp orders its special-token cache by length alone, so tokens of equal length end
4//! in the order this introsort leaves them, and that order decides which of two overlapping special tokens of one
5//! length a text splits on first. [`unordered_map_order`] is the iteration order of a
6//! `std::unordered_map<std::string, int>` filled one key at a time: llama.cpp takes the first token its map yields
7//! that matches a list of texts as the end-of-turn and fill-in-the-middle tokens, so with two candidates (Qwen's
8//! `<|endoftext|>` and `<|im_end|>`) the map's order picks one. `model_tokenizer_lock` holds both equal to the
9//! library on this box.
10
11const THRESHOLD: usize = 16;
12
13/// `std::sort(v.begin(), v.end(), less)`.
14pub fn sort<T: Copy>(v: &mut [T], less: &impl Fn(&T, &T) -> bool) {
15    let n = v.len();
16    if n == 0 {
17        return;
18    }
19    let lg = (usize::BITS - 1 - n.leading_zeros()) as usize;
20    introsort_loop(v, 0, n, 2 * lg, less);
21    final_insertion_sort(v, 0, n, less);
22}
23
24fn introsort_loop<T: Copy>(v: &mut [T], first: usize, mut last: usize, mut depth: usize, less: &impl Fn(&T, &T) -> bool) {
25    while last - first > THRESHOLD {
26        if depth == 0 {
27            heap_select(v, first, last, less);
28            sort_heap(v, first, last, less);
29            return;
30        }
31        depth -= 1;
32        let mid = first + (last - first) / 2;
33        move_median_to_first(v, first, first + 1, mid, last - 1, less);
34        let cut = unguarded_partition(v, first + 1, last, first, less);
35        introsort_loop(v, cut, last, depth, less);
36        last = cut;
37    }
38}
39
40fn move_median_to_first<T: Copy>(v: &mut [T], result: usize, a: usize, b: usize, c: usize, less: &impl Fn(&T, &T) -> bool) {
41    if less(&v[a], &v[b]) {
42        if less(&v[b], &v[c]) {
43            v.swap(result, b);
44        } else if less(&v[a], &v[c]) {
45            v.swap(result, c);
46        } else {
47            v.swap(result, a);
48        }
49    } else if less(&v[a], &v[c]) {
50        v.swap(result, a);
51    } else if less(&v[b], &v[c]) {
52        v.swap(result, c);
53    } else {
54        v.swap(result, b);
55    }
56}
57
58fn unguarded_partition<T: Copy>(v: &mut [T], mut first: usize, mut last: usize, pivot: usize, less: &impl Fn(&T, &T) -> bool) -> usize {
59    loop {
60        while less(&v[first], &v[pivot]) {
61            first += 1;
62        }
63        last -= 1;
64        while less(&v[pivot], &v[last]) {
65            last -= 1;
66        }
67        if first >= last {
68            return first;
69        }
70        v.swap(first, last);
71        first += 1;
72    }
73}
74
75fn adjust_heap<T: Copy>(v: &mut [T], first: usize, mut hole: usize, len: usize, value: T, less: &impl Fn(&T, &T) -> bool) {
76    let top = hole;
77    let mut second = hole;
78    while second < (len - 1) / 2 {
79        second = 2 * (second + 1);
80        if less(&v[first + second], &v[first + second - 1]) {
81            second -= 1;
82        }
83        v[first + hole] = v[first + second];
84        hole = second;
85    }
86    if len & 1 == 0 && second == (len - 2) / 2 {
87        second = 2 * (second + 1);
88        v[first + hole] = v[first + second - 1];
89        hole = second - 1;
90    }
91    let mut parent = hole.wrapping_sub(1) / 2;
92    while hole > top && less(&v[first + parent], &value) {
93        v[first + hole] = v[first + parent];
94        hole = parent;
95        parent = hole.wrapping_sub(1) / 2;
96    }
97    v[first + hole] = value;
98}
99
100fn make_heap<T: Copy>(v: &mut [T], first: usize, last: usize, less: &impl Fn(&T, &T) -> bool) {
101    let len = last - first;
102    if len < 2 {
103        return;
104    }
105    let mut parent = (len - 2) / 2;
106    loop {
107        let value = v[first + parent];
108        adjust_heap(v, first, parent, len, value, less);
109        if parent == 0 {
110            return;
111        }
112        parent -= 1;
113    }
114}
115
116fn pop_heap<T: Copy>(v: &mut [T], first: usize, last: usize, result: usize, less: &impl Fn(&T, &T) -> bool) {
117    let value = v[result];
118    v[result] = v[first];
119    adjust_heap(v, first, 0, last - first, value, less);
120}
121
122fn heap_select<T: Copy>(v: &mut [T], first: usize, middle: usize, less: &impl Fn(&T, &T) -> bool) {
123    make_heap(v, first, middle, less);
124}
125
126fn sort_heap<T: Copy>(v: &mut [T], first: usize, mut last: usize, less: &impl Fn(&T, &T) -> bool) {
127    while last - first > 1 {
128        last -= 1;
129        pop_heap(v, first, last, last, less);
130    }
131}
132
133fn final_insertion_sort<T: Copy>(v: &mut [T], first: usize, last: usize, less: &impl Fn(&T, &T) -> bool) {
134    if last - first > THRESHOLD {
135        insertion_sort(v, first, first + THRESHOLD, less);
136        for i in first + THRESHOLD..last {
137            unguarded_linear_insert(v, i, less);
138        }
139    } else {
140        insertion_sort(v, first, last, less);
141    }
142}
143
144fn insertion_sort<T: Copy>(v: &mut [T], first: usize, last: usize, less: &impl Fn(&T, &T) -> bool) {
145    if first == last {
146        return;
147    }
148    for i in first + 1..last {
149        if less(&v[i], &v[first]) {
150            let val = v[i];
151            v.copy_within(first..i, first + 1);
152            v[first] = val;
153        } else {
154            unguarded_linear_insert(v, i, less);
155        }
156    }
157}
158
159fn unguarded_linear_insert<T: Copy>(v: &mut [T], mut last: usize, less: &impl Fn(&T, &T) -> bool) {
160    let val = v[last];
161    let mut next = last - 1;
162    while less(&val, &v[next]) {
163        v[last] = v[next];
164        last = next;
165        next -= 1;
166    }
167    v[last] = val;
168}
169
170/// `std::_Hash_bytes` for 64-bit `size_t` (MurmurHash64A over unaligned little-endian words), the hash
171/// `std::hash<std::string>` runs with seed `0xc70f6907`.
172pub fn hash_bytes(b: &[u8], seed: u64) -> u64 {
173    const MUL: u64 = (0xc6a4a793u64 << 32) + 0x5bd1e995;
174    let shift_mix = |v: u64| v ^ (v >> 47);
175    let aligned = b.len() & !7;
176    let mut hash = seed ^ (b.len() as u64).wrapping_mul(MUL);
177    for w in b[..aligned].chunks_exact(8) {
178        let data = shift_mix(u64::from_le_bytes(w.try_into().expect("an 8-byte word")).wrapping_mul(MUL)).wrapping_mul(MUL);
179        hash ^= data;
180        hash = hash.wrapping_mul(MUL);
181    }
182    if b.len() & 7 != 0 {
183        let data = b[aligned..].iter().rev().fold(0u64, |acc, &c| (acc << 8) + c as u64);
184        hash ^= data;
185        hash = hash.wrapping_mul(MUL);
186    }
187    hash = shift_mix(hash).wrapping_mul(MUL);
188    shift_mix(hash)
189}
190
191/// `std::hash<std::string>`.
192pub fn hash_string(b: &[u8]) -> u64 {
193    hash_bytes(b, 0xc70f6907)
194}
195
196/// `__prime_list` past its first six entries, the bucket counts `_Prime_rehash_policy::_M_next_bkt` chooses among;
197/// the last is the sentinel past which no table grows.
198const PRIMES: [u64; 298] = [
199    17, 19, 23, 29, 31, 37, 41, 43, 47, 53, 59, 61, 67, 71, 73, 79, 83, 89, 97, 103, 109, 113, 127, 137, 139, 149,
200    157, 167, 179, 193, 199, 211, 227, 241, 257, 277, 293, 313, 337, 359, 383, 409, 439, 467, 503, 541, 577, 619,
201    661, 709, 761, 823, 887, 953, 1031, 1109, 1193, 1289, 1381, 1493, 1613, 1741, 1879, 2029, 2179, 2357, 2549,
202    2753, 2971, 3209, 3469, 3739, 4027, 4349, 4703, 5087, 5503, 5953, 6427, 6949, 7517, 8123, 8783, 9497, 10273,
203    11113, 12011, 12983, 14033, 15173, 16411, 17749, 19183, 20753, 22447, 24281, 26267, 28411, 30727, 33223, 35933,
204    38873, 42043, 45481, 49201, 53201, 57557, 62233, 67307, 72817, 78779, 85229, 92203, 99733, 107897, 116731,
205    126271, 136607, 147793, 159871, 172933, 187091, 202409, 218971, 236897, 256279, 277261, 299951, 324503, 351061,
206    379787, 410857, 444487, 480881, 520241, 562841, 608903, 658753, 712697, 771049, 834181, 902483, 976369,
207    1056323, 1142821, 1236397, 1337629, 1447153, 1565659, 1693859, 1832561, 1982627, 2144977, 2320627, 2510653,
208    2716249, 2938679, 3179303, 3439651, 3721303, 4026031, 4355707, 4712381, 5098259, 5515729, 5967347, 6456007,
209    6984629, 7556579, 8175383, 8844859, 9569143, 10352717, 11200489, 12117689, 13109983, 14183539, 15345007,
210    16601593, 17961079, 19431899, 21023161, 22744717, 24607243, 26622317, 28802401, 31160981, 33712729, 36473443,
211    39460231, 42691603, 46187573, 49969847, 54061849, 58488943, 63278561, 68460391, 74066549, 80131819, 86693767,
212    93793069, 101473717, 109783337, 118773397, 128499677, 139022417, 150406843, 162723577, 176048909, 190465427,
213    206062531, 222936881, 241193053, 260944219, 282312799, 305431229, 330442829, 357502601, 386778277, 418451333,
214    452718089, 489790921, 529899637, 573292817, 620239453, 671030513, 725980837, 785430967, 849749479, 919334987,
215    994618837, 1076067617, 1164186217, 1259520799, 1362662261, 1474249943, 1594975441, 1725587117, 1866894511,
216    2019773507, 2185171673, 2364114217, 2557710269, 2767159799, 2993761039, 3238918481, 3504151727, 3791104843,
217    4101556399, 4294967291, 6442450933, 8589934583, 12884901857, 17179869143, 25769803693, 34359738337,
218    51539607367, 68719476731, 103079215087, 137438953447, 206158430123, 274877906899, 412316860387, 549755813881,
219    824633720731, 1099511627689, 1649267441579, 2199023255531, 3298534883309, 4398046511093, 6597069766607,
220    8796093022151, 13194139533241, 17592186044399, 26388279066581, 35184372088777, 52776558133177, 70368744177643,
221    105553116266399, 140737488355213, 211106232532861, 281474976710597, 562949953421231, 1125899906842597,
222    2251799813685119, 4503599627370449, 9007199254740881, 18014398509481951, 36028797018963913, 72057594037927931,
223    144115188075855859, 288230376151711717, 576460752303423433, 1152921504606846883, 2305843009213693951,
224    4611686018427387847, 9223372036854775783, 18446744073709551557,
225];
226
227/// `_Prime_rehash_policy` at its default maximum load factor of 1.
228struct RehashPolicy {
229    next_resize: u64,
230}
231
232impl RehashPolicy {
233    fn next_bkt(&mut self, n: u64) -> u64 {
234        const FAST: [u64; 14] = [2, 2, 2, 3, 5, 5, 7, 7, 11, 11, 11, 11, 13, 13];
235        if n < FAST.len() as u64 {
236            if n == 0 {
237                return 1;
238            }
239            self.next_resize = FAST[n as usize];
240            return FAST[n as usize];
241        }
242        let last = PRIMES.len() - 1;
243        let i = PRIMES[..last].partition_point(|&p| p < n);
244        self.next_resize = if i == last { u64::MAX } else { PRIMES[i] };
245        PRIMES[i]
246    }
247
248    fn need_rehash(&mut self, n_bkt: u64, n_elt: u64, n_ins: u64) -> Option<u64> {
249        if n_elt + n_ins <= self.next_resize {
250            return None;
251        }
252        let min_bkts = (n_elt + n_ins).max(if self.next_resize == 0 { 11 } else { 0 }) as f64;
253        if min_bkts >= n_bkt as f64 {
254            return Some(self.next_bkt((min_bkts.floor() as u64 + 1).max(n_bkt * 2)));
255        }
256        self.next_resize = n_bkt;
257        None
258    }
259}
260
261/// The order a `std::unordered_map` iterates keys inserted in the order of `hashes` (each key distinct): libstdc++'s
262/// `_Hashtable` with cached hash codes, a new node at the front of its bucket (or of the whole list when its bucket
263/// is empty), and the unique-key rehash that relinks nodes front-first. The result lists insertion indices.
264pub fn unordered_map_order(hashes: &[u64]) -> Vec<u32> {
265    const NIL: u32 = u32::MAX;
266    const EMPTY: i64 = -1;
267    const BEFORE_BEGIN: i64 = -2;
268    let mut next = vec![NIL; hashes.len()];
269    let mut begin = NIL;
270    let mut policy = RehashPolicy { next_resize: 0 };
271    let mut bucket_count: u64 = 1;
272    let mut buckets: Vec<i64> = vec![EMPTY];
273    let get = |b: i64, begin: u32, next: &[u32]| if b == BEFORE_BEGIN { begin } else { next[b as usize] };
274    for (node, &h) in hashes.iter().enumerate() {
275        let node = node as u32;
276        if let Some(n) = policy.need_rehash(bucket_count, node as u64, 1) {
277            let mut fresh = vec![EMPTY; n as usize];
278            let mut p = begin;
279            begin = NIL;
280            let mut bbegin = 0usize;
281            while p != NIL {
282                let after = next[p as usize];
283                let bkt = (hashes[p as usize] % n) as usize;
284                if fresh[bkt] == EMPTY {
285                    next[p as usize] = begin;
286                    begin = p;
287                    fresh[bkt] = BEFORE_BEGIN;
288                    if next[p as usize] != NIL {
289                        fresh[bbegin] = p as i64;
290                    }
291                    bbegin = bkt;
292                } else {
293                    let b = fresh[bkt];
294                    next[p as usize] = get(b, begin, &next);
295                    if b == BEFORE_BEGIN {
296                        begin = p;
297                    } else {
298                        next[b as usize] = p;
299                    }
300                }
301                p = after;
302            }
303            buckets = fresh;
304            bucket_count = n;
305        }
306        let bkt = (h % bucket_count) as usize;
307        if buckets[bkt] != EMPTY {
308            let b = buckets[bkt];
309            next[node as usize] = get(b, begin, &next);
310            if b == BEFORE_BEGIN {
311                begin = node;
312            } else {
313                next[b as usize] = node;
314            }
315        } else {
316            next[node as usize] = begin;
317            begin = node;
318            let after = next[node as usize];
319            if after != NIL {
320                buckets[(hashes[after as usize] % bucket_count) as usize] = node as i64;
321            }
322            buckets[bkt] = BEFORE_BEGIN;
323        }
324    }
325    let mut order = Vec::with_capacity(hashes.len());
326    let mut p = begin;
327    while p != NIL {
328        order.push(p);
329        p = next[p as usize];
330    }
331    order
332}