uqa_storage/term_key/
allocation.rs1use super::TokenTermKey;
10use uqa_analysis::TokenTerm;
11use uqa_core::memory::{Budgeted, BudgetedVec, MemoryBudget, MemoryError};
12
13impl TokenTermKey {
14 pub fn from_text_budgeted<E: From<MemoryError>>(
16 text: &str,
17 budget: &MemoryBudget,
18 mut poll: impl FnMut() -> Result<(), E>,
19 ) -> Result<Budgeted<Self>, E> {
20 poll()?;
21 let mut bytes = BudgetedVec::new(budget);
22 bytes.reserve(text.len().checked_add(1).ok_or(MemoryError::SizeOverflow)?)?;
23 bytes.push(0)?;
24 for (index, byte) in text.bytes().enumerate() {
25 if index % 1024 == 0 {
26 poll()?;
27 }
28 bytes.push(byte)?;
29 }
30 poll()?;
31 let (bytes, memory) = bytes.into_parts();
32 Ok(Budgeted::new(Self(bytes), memory))
33 }
34
35 pub fn from_term_budgeted<E: From<MemoryError>>(
37 term: &TokenTerm,
38 budget: &MemoryBudget,
39 mut poll: impl FnMut() -> Result<(), E>,
40 ) -> Result<Budgeted<Self>, E> {
41 if let Some(text) = term.as_str() {
42 return Self::from_text_budgeted(text, budget, poll);
43 }
44 poll()?;
45 let units = term.utf16();
46 let mut bytes = BudgetedVec::new(budget);
47 bytes.reserve(
48 units
49 .len()
50 .checked_mul(2)
51 .and_then(|size| size.checked_add(1))
52 .ok_or(MemoryError::SizeOverflow)?,
53 )?;
54 bytes.push(1)?;
55 for (index, unit) in units.iter().enumerate() {
56 if index % 1024 == 0 {
57 poll()?;
58 }
59 for byte in unit.to_be_bytes() {
60 bytes.push(byte)?;
61 }
62 }
63 poll()?;
64 let (bytes, memory) = bytes.into_parts();
65 Ok(Budgeted::new(Self(bytes), memory))
66 }
67
68 pub fn clone_budgeted<E: From<MemoryError>>(
70 &self,
71 budget: &MemoryBudget,
72 mut poll: impl FnMut() -> Result<(), E>,
73 ) -> Result<Budgeted<Self>, E> {
74 poll()?;
75 let mut bytes = BudgetedVec::new(budget);
76 bytes.reserve(self.0.len())?;
77 for (index, byte) in self.0.iter().copied().enumerate() {
78 if index % 1024 == 0 {
79 poll()?;
80 }
81 bytes.push(byte)?;
82 }
83 poll()?;
84 let (bytes, memory) = bytes.into_parts();
85 Ok(Budgeted::new(Self(bytes), memory))
86 }
87
88 pub fn cmp_with_control<E>(
90 &self,
91 other: &Self,
92 poll: &mut dyn FnMut() -> Result<(), E>,
93 ) -> Result<std::cmp::Ordering, E> {
94 poll()?;
95 for (left, right) in self.0.chunks(1024).zip(other.0.chunks(1024)) {
96 poll()?;
97 let order = left.cmp(right);
98 if !order.is_eq() {
99 return Ok(order);
100 }
101 }
102 Ok(self.0.len().cmp(&other.0.len()))
103 }
104}
105
106#[cfg(test)]
107mod tests;