zerotrie/builder/nonconst/
builder.rs

1// This file is part of ICU4X. For terms of use, please see the file
2// called LICENSE at the top level of the ICU4X source tree
3// (online at: https://github.com/unicode-org/icu4x/blob/main/LICENSE ).
4
5use core::cmp::Ordering;
6
7use super::super::branch_meta::BranchMeta;
8use super::store::NonConstLengthsStack;
9use super::store::TrieBuilderStore;
10use crate::builder::slice_indices::ByteSliceWithIndices;
11use crate::byte_phf::PerfectByteHashMapCacheOwned;
12use crate::error::ZeroTrieBuildError;
13use crate::options::*;
14use crate::varint;
15use alloc::vec::Vec;
16
17/// A low-level builder for [`ZeroTrie`](crate::ZeroTrie). Supports all options.
18pub(crate) struct ZeroTrieBuilder<S> {
19    data: S,
20    phf_cache: PerfectByteHashMapCacheOwned,
21    options: ZeroTrieBuilderOptions,
22}
23
24impl<S: TrieBuilderStore> ZeroTrieBuilder<S> {
25    /// Returns the trie data as a `Vec<u8>`.
26    pub fn to_bytes(&self) -> Vec<u8> {
27        self.data.atbs_to_bytes()
28    }
29
30    /// Prepends a byte value to the front of the builder. If it is ASCII, an ASCII
31    /// node is prepended. If it is non-ASCII, if there is already a span node at
32    /// the front, we modify the span node to add the new byte; otherwise, we create
33    /// a new span node. Returns the delta in length, which is either 1 or 2.
34    fn prepend_ascii(&mut self, ascii: u8) -> Result<usize, ZeroTrieBuildError> {
35        if ascii <= 127 {
36            self.data.atbs_push_front(ascii);
37            Ok(1)
38        } else if matches!(self.options.ascii_mode, AsciiMode::BinarySpans) {
39            if let Some(old_front) = self.data.atbs_pop_front() {
40                let old_byte_len = self.data.atbs_len() + 1;
41                if old_front & 0b11100000 == 0b10100000 {
42                    // Extend an existing span
43                    // Unwrap OK: there is a varint at this location in the buffer
44                    #[expect(clippy::unwrap_used)]
45                    let old_span_size =
46                        varint::try_read_varint_meta3_from_tstore(old_front, &mut self.data)
47                            .unwrap();
48                    self.data.atbs_push_front(ascii);
49                    let varint_array = varint::write_varint_meta3(old_span_size + 1);
50                    self.data.atbs_extend_front(varint_array.as_slice());
51                    self.data.atbs_bitor_assign(0, 0b10100000);
52                    let new_byte_len = self.data.atbs_len();
53                    return Ok(new_byte_len - old_byte_len);
54                } else {
55                    self.data.atbs_push_front(old_front);
56                }
57            }
58            // Create a new span
59            self.data.atbs_push_front(ascii);
60            self.data.atbs_push_front(0b10100001);
61            Ok(2)
62        } else {
63            Err(ZeroTrieBuildError::NonAsciiError)
64        }
65    }
66
67    /// Prepends a value node to the front of the builder. Returns the
68    /// delta in length, which depends on the size of the varint.
69    #[must_use]
70    fn prepend_value(&mut self, value: usize) -> usize {
71        let varint_array = varint::write_varint_meta3(value);
72        self.data.atbs_extend_front(varint_array.as_slice());
73        self.data.atbs_bitor_assign(0, 0b10000000);
74        varint_array.len()
75    }
76
77    /// Prepends a branch node to the front of the builder. Returns the
78    /// delta in length, which depends on the size of the varint.
79    #[must_use]
80    fn prepend_branch(&mut self, value: usize) -> usize {
81        let varint_array = varint::write_varint_meta2(value);
82        self.data.atbs_extend_front(varint_array.as_slice());
83        self.data.atbs_bitor_assign(0, 0b11000000);
84        varint_array.len()
85    }
86
87    /// Prepends multiple arbitrary bytes to the front of the builder. Returns the
88    /// delta in length, which is the length of the slice.
89    #[must_use]
90    fn prepend_slice(&mut self, s: &[u8]) -> usize {
91        self.data.atbs_extend_front(s);
92        s.len()
93    }
94
95    /// Builds a [`ZeroTrie`](crate::ZeroTrie) from an iterator of bytes. It first collects and sorts the iterator.
96    pub fn from_bytes_iter<K: AsRef<[u8]>, I: IntoIterator<Item = (K, usize)>>(
97        iter: I,
98        options: ZeroTrieBuilderOptions,
99    ) -> Result<Self, ZeroTrieBuildError> {
100        let items = Vec::<(K, usize)>::from_iter(iter);
101        let mut items = items
102            .iter()
103            .map(|(k, v)| (k.as_ref(), *v))
104            .collect::<Vec<(&[u8], usize)>>();
105        items.sort_by(|a, b| cmp_keys_values(options, *a, *b));
106        let ascii_str_slice = items.as_slice();
107        let byte_str_slice = ByteSliceWithIndices::from_byte_slice(ascii_str_slice);
108        Self::from_sorted_tuple_slice_impl(byte_str_slice, options)
109    }
110
111    /// Builds a [`ZeroTrie`](crate::ZeroTrie) with the given items and options. Assumes that the items are sorted,
112    /// except for a case-insensitive trie where the items are re-sorted.
113    ///
114    /// # Panics
115    ///
116    /// May panic if the items are not sorted.
117    pub fn from_sorted_tuple_slice(
118        items: ByteSliceWithIndices,
119        options: ZeroTrieBuilderOptions,
120    ) -> Result<Self, ZeroTrieBuildError> {
121        if matches!(options.case_sensitivity, CaseSensitivity::IgnoreCase) {
122            // We need to re-sort the items with our custom comparator.
123            let mut items_vec = items.to_vec_u8();
124            items_vec.sort_by(|a, b| cmp_keys_values(options, *a, *b));
125            Self::from_sorted_tuple_slice_impl(
126                ByteSliceWithIndices::from_byte_slice(&items_vec),
127                options,
128            )
129        } else {
130            Self::from_sorted_tuple_slice_impl(items, options)
131        }
132    }
133
134    /// Internal constructor that does not re-sort the items.
135    fn from_sorted_tuple_slice_impl(
136        items: ByteSliceWithIndices,
137        options: ZeroTrieBuilderOptions,
138    ) -> Result<Self, ZeroTrieBuildError> {
139        #[allow(clippy::indexing_slicing)] // a debug assertion only
140        let mut i = 0;
141        while i + 1 < items.len() {
142            let ab0 = items.get_or_panic(i);
143            let ab1 = items.get_or_panic(i + 1);
144            debug_assert!(cmp_keys_values(options, (ab0.0, ab0.1), (ab1.0, ab1.1)).is_lt());
145            i += 1;
146        }
147
148        let mut result = Self {
149            data: S::atbs_new_empty(),
150            phf_cache: PerfectByteHashMapCacheOwned::new_empty(),
151            options,
152        };
153        let total_size = result.create(items)?;
154        debug_assert!(total_size == result.data.atbs_len());
155        Ok(result)
156    }
157
158    /// The actual builder algorithm. For an explanation, see [`crate::builder`].
159    #[expect(clippy::unwrap_used, clippy::indexing_slicing)] // lots of indexing, but all indexes should be in range
160    fn create(&mut self, all_items: ByteSliceWithIndices) -> Result<usize, ZeroTrieBuildError> {
161        let mut prefix_len = match all_items.last() {
162            Some(x) => x.0.len(),
163            // Empty slice:
164            None => return Ok(0),
165        };
166        // Initialize the main loop to point at the last string.
167        let mut lengths_stack = NonConstLengthsStack::new();
168        let mut i = all_items.len() - 1;
169        let mut j = all_items.len();
170        let mut current_len = 0;
171        // Start the main loop.
172        loop {
173            let item_i = all_items.get_or_panic(i);
174            let item_j = all_items.get_or_panic(j - 1);
175            debug_assert!(crate::builder::slice_indices::prefix_eq_or_panic(
176                item_i.0, item_j.0, prefix_len
177            ));
178            // Check if we need to add a value node here.
179            if item_i.0.len() == prefix_len {
180                let len = self.prepend_value(item_i.1);
181                current_len += len;
182            }
183            if prefix_len == 0 {
184                // All done! Leave the main loop.
185                break;
186            }
187            // Reduce the prefix length by 1 and recalculate i and j.
188            prefix_len -= 1;
189            let mut new_i = i;
190            let mut new_j = j;
191            let mut ascii_i = item_i.0[prefix_len];
192            let mut ascii_j = item_j.0[prefix_len];
193            debug_assert_eq!(ascii_i, ascii_j);
194            let key_ascii = ascii_i;
195            loop {
196                if new_i == 0 {
197                    break;
198                }
199                let candidate = all_items.get_or_panic(new_i - 1).0;
200                if candidate.len() < prefix_len {
201                    // Too short
202                    break;
203                }
204                if crate::builder::slice_indices::prefix_eq_or_panic(
205                    item_i.0, candidate, prefix_len,
206                ) {
207                    new_i -= 1;
208                } else {
209                    break;
210                }
211                if candidate.len() == prefix_len {
212                    // A string that equals the prefix does not take part in the branch node.
213                    break;
214                }
215                let candidate = candidate[prefix_len];
216                if candidate != ascii_i {
217                    ascii_i = candidate;
218                }
219            }
220            loop {
221                if new_j == all_items.len() {
222                    break;
223                }
224                let candidate = all_items.get_or_panic(new_j).0;
225                if candidate.len() < prefix_len {
226                    // Too short
227                    break;
228                }
229                if crate::builder::slice_indices::prefix_eq_or_panic(
230                    item_j.0, candidate, prefix_len,
231                ) {
232                    new_j += 1;
233                } else {
234                    break;
235                }
236                if candidate.len() == prefix_len {
237                    unreachable!("A shorter string should be earlier in the sequence");
238                }
239                let candidate = candidate[prefix_len];
240                if candidate != ascii_j {
241                    ascii_j = candidate;
242                }
243            }
244            // If there are no different bytes at this prefix level, we can add an ASCII or Span
245            // node and then continue to the next iteration of the main loop.
246            if ascii_i == key_ascii && ascii_j == key_ascii {
247                let len = self.prepend_ascii(key_ascii)?;
248                current_len += len;
249                if matches!(self.options.case_sensitivity, CaseSensitivity::IgnoreCase)
250                    && i == new_i + 2
251                {
252                    // This can happen if two strings were picked up, each with a different case
253                    return Err(ZeroTrieBuildError::MixedCase);
254                }
255                debug_assert!(
256                    i == new_i || i == new_i + 1,
257                    "only the exact prefix string can be picked up at this level: {key_ascii}"
258                );
259                i = new_i;
260                debug_assert_eq!(j, new_j);
261                continue;
262            }
263            // If i and j changed, we are a target of a branch node.
264            if ascii_j == key_ascii {
265                // We are the _last_ target of a branch node.
266                lengths_stack.push(BranchMeta {
267                    ascii: key_ascii,
268                    cumulative_length: current_len,
269                    local_length: current_len,
270                    count: 1,
271                });
272            } else {
273                // We are the _not the last_ target of a branch node.
274                let BranchMeta {
275                    cumulative_length,
276                    count,
277                    ..
278                } = lengths_stack.peek_or_panic();
279                lengths_stack.push(BranchMeta {
280                    ascii: key_ascii,
281                    cumulative_length: cumulative_length + current_len,
282                    local_length: current_len,
283                    count: count + 1,
284                });
285            }
286            if ascii_i != key_ascii {
287                // We are _not the first_ target of a branch node.
288                // Set the cursor to the previous string and continue the loop.
289                j = i;
290                i -= 1;
291                prefix_len = all_items.get_or_panic(i).0.len();
292                current_len = 0;
293                continue;
294            }
295            // Branch (first)
296            // std::println!("lengths_stack: {lengths_stack:?}");
297            let (total_length, total_count) = {
298                let BranchMeta {
299                    cumulative_length,
300                    count,
301                    ..
302                } = lengths_stack.peek_or_panic();
303                (cumulative_length, count)
304            };
305            let mut branch_metas = lengths_stack.pop_many_or_panic(total_count);
306            let original_keys = branch_metas.map_to_ascii_bytes();
307            if matches!(self.options.case_sensitivity, CaseSensitivity::IgnoreCase) {
308                // Check to see if we have the same letter in two different cases
309                let mut seen_ascii_alpha = [false; 26];
310                for c in original_keys.as_const_slice().as_slice() {
311                    if c.is_ascii_alphabetic() {
312                        let i = (c.to_ascii_lowercase() - b'a') as usize;
313                        #[allow(clippy::indexing_slicing)] // 26 letters
314                        if seen_ascii_alpha[i] {
315                            return Err(ZeroTrieBuildError::MixedCase);
316                        } else {
317                            seen_ascii_alpha[i] = true;
318                        }
319                    }
320                }
321            }
322            let use_phf = matches!(self.options.phf_mode, PhfMode::UsePhf);
323            let opt_phf_vec = if total_count > 15 && use_phf {
324                let phf_vec = self
325                    .phf_cache
326                    .try_get_or_insert(original_keys.as_const_slice().as_slice().to_vec())?;
327                // Put everything in order via bubble sort
328                // Note: branch_metas is stored in reverse order (0 = last element)
329                loop {
330                    let mut l = total_count - 1;
331                    let mut changes = 0;
332                    let mut start = 0;
333                    while l > 0 {
334                        let a = *branch_metas.as_const_slice().get_or_panic(l);
335                        let b = *branch_metas.as_const_slice().get_or_panic(l - 1);
336                        let a_idx = phf_vec.keys().iter().position(|x| x == &a.ascii).unwrap();
337                        let b_idx = phf_vec.keys().iter().position(|x| x == &b.ascii).unwrap();
338                        if a_idx > b_idx {
339                            // std::println!("{a:?} <=> {b:?} ({phf_vec:?})");
340                            // This method call won't panic because the ranges are valid.
341                            self.data.atbs_swap_ranges(
342                                start,
343                                start + a.local_length,
344                                start + a.local_length + b.local_length,
345                            );
346                            branch_metas.swap_or_panic(l - 1, l);
347                            start += b.local_length;
348                            changes += 1;
349                            // FIXME: fix the `length` field
350                        } else {
351                            start += a.local_length;
352                        }
353                        l -= 1;
354                    }
355                    if changes == 0 {
356                        break;
357                    }
358                }
359                Some(phf_vec)
360            } else {
361                None
362            };
363            // Write out the offset table
364            current_len = total_length;
365            let w = (usize::BITS as usize - (total_length.leading_zeros() as usize) - 1) / 8;
366            if w > 3 && matches!(self.options.capacity_mode, CapacityMode::Normal) {
367                return Err(ZeroTrieBuildError::CapacityExceeded);
368            }
369            let mut k = 0;
370            while k <= w {
371                self.data.atbs_prepend_n_zeros(total_count - 1);
372                current_len += total_count - 1;
373                let mut l = 0;
374                let mut length_to_write = 0;
375                while l < total_count {
376                    let BranchMeta { local_length, .. } = *branch_metas
377                        .as_const_slice()
378                        .get_or_panic(total_count - l - 1);
379                    let mut adjusted_length = length_to_write;
380                    let mut m = 0;
381                    while m < k {
382                        adjusted_length >>= 8;
383                        m += 1;
384                    }
385                    if l > 0 {
386                        self.data.atbs_bitor_assign(l - 1, adjusted_length as u8);
387                    }
388                    l += 1;
389                    length_to_write += local_length;
390                }
391                k += 1;
392            }
393            // Write out the lookup table
394            assert!(0 < total_count && total_count <= 256);
395            let branch_value = (w << 8) + (total_count & 0xff);
396            if let Some(phf_vec) = opt_phf_vec {
397                self.data.atbs_extend_front(phf_vec.as_bytes());
398                let phf_len = phf_vec.as_bytes().len();
399                let branch_len = self.prepend_branch(branch_value);
400                current_len += phf_len + branch_len;
401            } else {
402                let search_len = self.prepend_slice(original_keys.as_slice());
403                let branch_len = self.prepend_branch(branch_value);
404                current_len += search_len + branch_len;
405            }
406            i = new_i;
407            j = new_j;
408        }
409        assert!(lengths_stack.is_empty());
410        Ok(current_len)
411    }
412}
413
414fn cmp_keys_values(
415    options: ZeroTrieBuilderOptions,
416    a: (&[u8], usize),
417    b: (&[u8], usize),
418) -> Ordering {
419    if matches!(options.case_sensitivity, CaseSensitivity::Sensitive) {
420        a.0.cmp(b.0)
421    } else {
422        let a_iter = a.0.iter().map(|x| x.to_ascii_lowercase());
423        let b_iter = b.0.iter().map(|x| x.to_ascii_lowercase());
424        Iterator::cmp(a_iter, b_iter)
425    }
426    .then_with(|| a.1.cmp(&b.1))
427}
zerotrie/builder/nonconst/builder.rs

zerotrie/builder/nonconst/
builder.rs
