Skip to main content

uqa_analysis/token/allocation/
cloning.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! Reserved token copies preserve attributes and reuse retained source allocations.
8
9use uqa_core::memory::{Budgeted, MemoryBudget};
10
11use super::{AnalysisToken, AnalyzedText, TokenBuffer};
12use crate::term::TermBoundary;
13use crate::{AnalysisResult, SourceOffsets, TokenTerm};
14
15impl AnalysisToken {
16    /// Copy the term and morphology into independent reservations, retaining all position and source attributes.
17    pub fn clone_budgeted(
18        &self,
19        budget: &MemoryBudget,
20        mut poll: impl FnMut() -> AnalysisResult<()>,
21    ) -> AnalysisResult<Budgeted<Self>> {
22        let term = self.term.clone_budgeted(budget, &mut poll)?;
23        self.copy_attributes(term, self.verbatim, budget, &mut poll)
24    }
25
26    pub(crate) fn rewrite_budgeted(
27        &self,
28        term: Budgeted<TokenTerm>,
29        budget: &MemoryBudget,
30        poll: &mut dyn FnMut() -> AnalysisResult<()>,
31    ) -> AnalysisResult<Budgeted<Self>> {
32        let verbatim = self.verbatim && self.term.eq_with_control(&term, poll)?;
33        self.copy_attributes(term, verbatim, budget, poll)
34    }
35
36    fn copy_attributes(
37        &self,
38        term: Budgeted<TokenTerm>,
39        verbatim: bool,
40        budget: &MemoryBudget,
41        poll: &mut dyn FnMut() -> AnalysisResult<()>,
42    ) -> AnalysisResult<Budgeted<Self>> {
43        poll()?;
44        #[cfg(feature = "nori")]
45        let morphology = self
46            .korean_morphology
47            .as_ref()
48            .map(|value| value.clone_budgeted(budget, poll))
49            .transpose()?;
50        #[cfg(not(feature = "nori"))]
51        let _ = budget;
52        let (term, memory) = term.into_parts();
53        #[cfg(feature = "nori")]
54        let (korean_morphology, memory) = match morphology {
55            Some(morphology) => {
56                let (morphology, allocation) = morphology.into_parts();
57                let mut memory = memory;
58                memory.absorb(allocation);
59                (Some(morphology), memory)
60            }
61            None => (None, memory),
62        };
63        Ok(Budgeted::new(
64            Self {
65                term,
66                offsets: self.offsets.clone(),
67                position_increment: self.position_increment,
68                position_length: self.position_length,
69                keyword: self.keyword,
70                filtered_utf16: self.filtered_utf16.clone(),
71                #[cfg(feature = "nori")]
72                korean_morphology,
73                verbatim,
74            },
75            memory,
76        ))
77    }
78
79    pub(crate) fn substring_budgeted(
80        &self,
81        start: TermBoundary,
82        end: TermBoundary,
83        complete: TermBoundary,
84        budget: &MemoryBudget,
85        poll: &mut dyn FnMut() -> AnalysisResult<()>,
86    ) -> AnalysisResult<Budgeted<Self>> {
87        let term = self
88            .term
89            .substring_budgeted(start.offset..end.offset, budget, poll)?;
90        let (mut token, memory) = self
91            .copy_attributes(term, self.verbatim, budget, poll)?
92            .into_parts();
93        if self.verbatim {
94            if let Some(offsets) = &self.offsets {
95                token.offsets = Some(SourceOffsets {
96                    utf8: offsets.utf8.start + start.offset..offsets.utf8.start + end.offset,
97                    utf16: offsets.utf16.start + start.utf16..offsets.utf16.start + end.utf16,
98                });
99                if let Some(filtered) = &self.filtered_utf16 {
100                    if filtered.len() == complete.utf16 {
101                        token.filtered_utf16 =
102                            Some(filtered.start + start.utf16..filtered.start + end.utf16);
103                    }
104                }
105            }
106        }
107        Ok(Budgeted::new(token, memory))
108    }
109}
110
111impl AnalyzedText {
112    /// Copy owned token buffers into this allowance while sharing retained source allocations with the original.
113    ///
114    /// Every term, morphology value, token buffer and hidden terminal token receives its own reservation. Existing source projection and mapping leases remain shared with the input. Failure leaves the borrowed result unchanged.
115    ///
116    /// ```
117    /// use uqa_analysis::Tokenizer;
118    /// use uqa_core::memory::MemoryBudget;
119    ///
120    /// let budget = MemoryBudget::new(64 * 1024);
121    /// let input = Tokenizer::Whitespace.tokenize_with_offsets_budgeted(
122    ///     "韓🙂 UQA", &budget, || Ok(()),
123    /// )?;
124    /// let copied = input.clone_budgeted(&budget, || Ok(()))?;
125    /// assert_eq!(*copied, *input);
126    /// drop(input);
127    /// assert!(budget.used() >= copied.reserved_bytes());
128    /// drop(copied);
129    /// assert_eq!(budget.used(), 0);
130    /// # Ok::<(), uqa_analysis::AnalysisError>(())
131    /// ```
132    pub fn clone_budgeted(
133        &self,
134        budget: &MemoryBudget,
135        mut poll: impl FnMut() -> AnalysisResult<()>,
136    ) -> AnalysisResult<Budgeted<Self>> {
137        poll()?;
138        let mut output = TokenBuffer::new(budget);
139        output.tokens.reserve(self.batch.tokens.len())?;
140        for token in &self.batch.tokens {
141            output.push(token.clone_budgeted(budget, &mut poll)?)?;
142        }
143        if let Some(terminal) = &self.batch.terminal {
144            let terminal = terminal.clone_budgeted(budget, &mut poll)?;
145            output.memory.grow(std::mem::size_of::<AnalysisToken>())?;
146            let (terminal, memory) = terminal.into_parts();
147            output.memory.absorb(memory);
148            output.terminal = Some(Box::new(terminal));
149        }
150        output.finish_retained(
151            self.final_offsets.clone(),
152            self.batch.final_position_increment,
153            #[cfg(feature = "nori")]
154            self.projection.clone(),
155            &mut poll,
156        )
157    }
158}