Skip to main content

uqa_analysis/
source.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! Unicode source coordinates and composed character-filter edit maps.
8
9use std::cell::OnceCell;
10use std::ops::Range;
11use std::sync::Arc;
12
13use serde::Serialize;
14use uqa_core::memory::{Budgeted, BudgetedVec, MemoryBudget};
15
16use crate::{AnalysisError, AnalysisResult};
17
18mod edits;
19mod maps;
20mod runtime;
21
22pub(crate) use edits::{EditBuilder, EditedText};
23use maps::EditMaps;
24use runtime::SourceText;
25
26/// Half-open source ranges: exact UTF-16 units and the covering UTF-8 scalar range.
27///
28/// Strict projection methods require both ranges to address identical scalar boundaries. Explicit covering methods retain UTF-16 boundaries inside surrogate pairs while covering each touched scalar in UTF-8.
29#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
30pub struct SourceOffsets {
31    pub utf8: Range<usize>,
32    pub utf16: Range<usize>,
33}
34
35/// A checked conversion index containing Unicode scalar boundaries only.
36#[derive(Debug, Clone, PartialEq, Eq)]
37pub struct TextCoordinates {
38    boundaries: Vec<(usize, usize)>,
39    utf8_len: usize,
40    utf16_len: usize,
41}
42
43impl TextCoordinates {
44    pub fn new(text: &str) -> Self {
45        Self::new_budgeted(text, &MemoryBudget::new(usize::MAX), &mut || Ok(()))
46            .expect("unbounded source coordinate allocation")
47            .into_parts()
48            .0
49    }
50
51    /// Construct a scalar index with caller-owned byte limits and cancellation. ASCII needs no boundary buffer; other input reserves one entry per scalar and its end boundary.
52    pub fn new_budgeted(
53        text: &str,
54        budget: &MemoryBudget,
55        poll: &mut dyn FnMut() -> AnalysisResult<()>,
56    ) -> AnalysisResult<Budgeted<Self>> {
57        poll()?;
58        let mut ascii = true;
59        for chunk in text.as_bytes().chunks(1024) {
60            poll()?;
61            if !chunk.is_ascii() {
62                ascii = false;
63                break;
64            }
65        }
66        let mut boundaries = BudgetedVec::new(budget);
67        if ascii {
68            return Ok(Budgeted::new(
69                Self {
70                    boundaries: Vec::new(),
71                    utf8_len: text.len(),
72                    utf16_len: text.len(),
73                },
74                budget.empty_reservation(),
75            ));
76        }
77        let mut scalars = 0;
78        for _ in text.chars() {
79            if scalars % 1024 == 0 {
80                poll()?;
81            }
82            scalars += 1;
83        }
84        boundaries.reserve(scalars + 1)?;
85        let mut utf16 = 0;
86        for (index, (utf8, character)) in text.char_indices().enumerate() {
87            if index % 1024 == 0 {
88                poll()?;
89            }
90            boundaries.push((utf8, utf16))?;
91            utf16 += character.len_utf16();
92        }
93        boundaries.push((text.len(), utf16))?;
94        poll()?;
95        let (boundaries, memory) = boundaries.into_parts();
96        Ok(Budgeted::new(
97            Self {
98                boundaries,
99                utf8_len: text.len(),
100                utf16_len: utf16,
101            },
102            memory,
103        ))
104    }
105
106    pub fn utf8_len(&self) -> usize {
107        self.utf8_len
108    }
109
110    pub fn utf16_len(&self) -> usize {
111        self.utf16_len
112    }
113
114    pub fn utf8_to_utf16(&self, offset: usize) -> AnalysisResult<usize> {
115        if self.boundaries.is_empty() && offset <= self.utf8_len {
116            return Ok(offset);
117        }
118        self.boundaries
119            .binary_search_by_key(&offset, |point| point.0)
120            .map(|index| self.boundaries[index].1)
121            .map_err(|_| AnalysisError::InvalidTextOffset {
122                coordinate: "UTF-8",
123                offset,
124                length: self.utf8_len(),
125            })
126    }
127
128    pub fn utf16_to_utf8(&self, offset: usize) -> AnalysisResult<usize> {
129        if self.boundaries.is_empty() && offset <= self.utf16_len {
130            return Ok(offset);
131        }
132        self.boundaries
133            .binary_search_by_key(&offset, |point| point.1)
134            .map(|index| self.boundaries[index].0)
135            .map_err(|_| AnalysisError::InvalidTextOffset {
136                coordinate: "UTF-16",
137                offset,
138                length: self.utf16_len(),
139            })
140    }
141
142    pub fn offsets(&self, utf8: Range<usize>) -> AnalysisResult<SourceOffsets> {
143        validate_order(&utf8)?;
144        let utf16 = self.utf8_to_utf16(utf8.start)?..self.utf8_to_utf16(utf8.end)?;
145        Ok(SourceOffsets { utf8, utf16 })
146    }
147
148    /// Retain exact UTF-16 coordinates and cover split surrogate pairs in UTF-8.
149    ///
150    /// An empty range inside a pair covers that scalar; an empty range at a scalar boundary remains empty.
151    pub fn covering_offsets_utf16(&self, utf16: Range<usize>) -> AnalysisResult<SourceOffsets> {
152        self.validate_utf16_range(&utf16)?;
153        if self.boundaries.is_empty() {
154            return Ok(SourceOffsets {
155                utf8: utf16.clone(),
156                utf16,
157            });
158        }
159        let start = match self
160            .boundaries
161            .binary_search_by_key(&utf16.start, |point| point.1)
162        {
163            Ok(index) => self.boundaries[index].0,
164            Err(index) => self.boundaries[index - 1].0,
165        };
166        let end = match self
167            .boundaries
168            .binary_search_by_key(&utf16.end, |point| point.1)
169        {
170            Ok(index) | Err(index) => self.boundaries[index].0,
171        };
172        Ok(SourceOffsets {
173            utf8: start..end,
174            utf16,
175        })
176    }
177
178    fn validate_utf16_range(&self, range: &Range<usize>) -> AnalysisResult<()> {
179        validate_order(range)?;
180        if range.end > self.utf16_len {
181            return Err(AnalysisError::InvalidTextOffset {
182                coordinate: "UTF-16",
183                offset: range.end,
184                length: self.utf16_len,
185            });
186        }
187        Ok(())
188    }
189}
190
191/// Filtered text retaining its original input and the provenance of every edit.
192///
193/// Replacement output covers the replaced source range. Insertions map to an empty source range. An empty output range uses the following source boundary, including trailing deletions at the end of the input.
194///
195/// ```
196/// use uqa_analysis::CharFilter;
197/// let input = "<b>한&amp;🙂</b>";
198/// let filtered = CharFilter::HTMLStrip.filter_with_offsets(input)?;
199/// assert_eq!(filtered.as_str(), " 한&🙂 ");
200/// let entity = filtered.source_offsets(4..5)?;
201/// assert_eq!(&input[entity.utf8], "&amp;");
202/// assert_eq!(entity.utf16, 4..9);
203/// # Ok::<(), uqa_analysis::AnalysisError>(())
204/// ```
205#[derive(Debug, Clone)]
206pub struct FilteredText<'a> {
207    original: &'a str,
208    text: SourceText<'a>,
209    maps: Option<Arc<EditMaps>>,
210    original_coordinates: OnceCell<Arc<Budgeted<TextCoordinates>>>,
211    filtered_coordinates: OnceCell<Arc<Budgeted<TextCoordinates>>>,
212}
213
214impl<'a> FilteredText<'a> {
215    pub fn new(text: &'a str) -> Self {
216        Self {
217            original: text,
218            text: SourceText::Borrowed(text),
219            maps: None,
220            original_coordinates: OnceCell::new(),
221            filtered_coordinates: OnceCell::new(),
222        }
223    }
224
225    pub fn as_str(&self) -> &str {
226        self.text.as_str()
227    }
228
229    pub fn original(&self) -> &'a str {
230        self.original
231    }
232
233    pub fn into_string(self) -> String {
234        self.text.into_string()
235    }
236
237    pub(crate) fn unbounded_budget(&self) -> MemoryBudget {
238        self.maps
239            .as_ref()
240            .map(|maps| maps.budget())
241            .filter(|budget| budget.limit() == usize::MAX)
242            .cloned()
243            .unwrap_or_else(|| MemoryBudget::new(usize::MAX))
244    }
245
246    /// Project a filtered UTF-8 range into the covering original source range.
247    pub fn source_offsets(&self, mut range: Range<usize>) -> AnalysisResult<SourceOffsets> {
248        validate_utf8_range(self.as_str(), &range)?;
249        if let Some(maps) = &self.maps {
250            for map in maps.iter().rev() {
251                range = map.project(range);
252            }
253        }
254        self.original_coordinates().offsets(range)
255    }
256
257    /// Project filtered UTF-16 coordinates without accepting a split surrogate pair.
258    pub fn source_offsets_utf16(&self, range: Range<usize>) -> AnalysisResult<SourceOffsets> {
259        validate_order(&range)?;
260        let coordinates = self.filtered_coordinates();
261        let utf8 = coordinates.utf16_to_utf8(range.start)?..coordinates.utf16_to_utf8(range.end)?;
262        self.source_offsets(utf8)
263    }
264
265    /// Project exact filtered UTF-16 units, preserving split pairs through every edit map.
266    ///
267    /// The returned original UTF-16 range stays exact; UTF-8 covers the original scalars. Replacements cover their source and insertions retain their source boundary, as in strict projection.
268    ///
269    /// ```
270    /// use uqa_analysis::CharFilter;
271    /// let filtered = CharFilter::HTMLStrip.filter_with_offsets("<b>🙂a</b>")?;
272    /// let source = filtered.source_covering_offsets_utf16(2..3)?;
273    /// assert_eq!(source.utf16, 4..5);
274    /// assert_eq!(source.utf8, 3..7);
275    /// assert_eq!(&filtered.original()[source.utf8], "🙂");
276    /// # Ok::<(), uqa_analysis::AnalysisError>(())
277    /// ```
278    pub fn source_covering_offsets_utf16(
279        &self,
280        mut range: Range<usize>,
281    ) -> AnalysisResult<SourceOffsets> {
282        self.filtered_coordinates().validate_utf16_range(&range)?;
283        if let Some(maps) = &self.maps {
284            for map in maps.iter().rev() {
285                range = map.project_utf16(range);
286            }
287        }
288        self.original_coordinates().covering_offsets_utf16(range)
289    }
290
291    /// The final input boundary, even when filters remove all source characters.
292    pub fn final_offsets(&self) -> SourceOffsets {
293        let utf8 = self.original.len();
294        let utf16 = self.original_coordinates().utf16_len();
295        SourceOffsets {
296            utf8: utf8..utf8,
297            utf16: utf16..utf16,
298        }
299    }
300
301    pub(crate) fn apply_edited(
302        &mut self,
303        edited: Option<EditedText>,
304        budget: &MemoryBudget,
305        poll: &mut dyn FnMut() -> AnalysisResult<()>,
306    ) -> AnalysisResult<()> {
307        let Some(edited) = edited else {
308            return Ok(());
309        };
310        poll()?;
311        let text = edited.text.into_shared()?;
312        let map = edited.map.into_shared()?;
313        EditMaps::push(&mut self.maps, map, budget, poll)?;
314        self.text = SourceText::Owned(text);
315        self.filtered_coordinates.take();
316        Ok(())
317    }
318
319    pub(crate) fn prepare_coordinates(
320        &self,
321        budget: &MemoryBudget,
322        poll: &mut dyn FnMut() -> AnalysisResult<()>,
323    ) -> AnalysisResult<()> {
324        poll()?;
325        if self.original_coordinates.get().is_none() {
326            let coordinates =
327                TextCoordinates::new_budgeted(self.original, budget, poll)?.into_shared()?;
328            self.original_coordinates
329                .set(coordinates)
330                .expect("uninitialized original coordinates");
331        }
332        if self.maps.is_some() && self.filtered_coordinates.get().is_none() {
333            let coordinates =
334                TextCoordinates::new_budgeted(self.as_str(), budget, poll)?.into_shared()?;
335            self.filtered_coordinates
336                .set(coordinates)
337                .expect("uninitialized filtered coordinates");
338        }
339        Ok(())
340    }
341
342    pub(crate) fn filtered_utf16(&self, range: Range<usize>) -> AnalysisResult<Range<usize>> {
343        Ok(self.filtered_coordinates().offsets(range)?.utf16)
344    }
345
346    #[cfg(any(feature = "nori", feature = "kuromoji"))]
347    pub(crate) fn projection(&self) -> Arc<Budgeted<SourceProjection>> {
348        let budget = MemoryBudget::new(usize::MAX);
349        self.projection_budgeted(&budget, &mut || Ok(()))
350            .expect("unbounded source retention")
351    }
352
353    #[cfg(any(feature = "nori", feature = "kuromoji"))]
354    pub(crate) fn projection_budgeted(
355        &self,
356        budget: &MemoryBudget,
357        poll: &mut dyn FnMut() -> AnalysisResult<()>,
358    ) -> AnalysisResult<Arc<Budgeted<SourceProjection>>> {
359        self.prepare_coordinates(budget, poll)?;
360        let source = crate::allocation::copy_text(self.original, budget, poll)?.into_shared()?;
361        let projection = SourceProjection {
362            source,
363            maps: self.maps.clone(),
364            original: self.original_coordinates().clone(),
365            filtered: self.filtered_coordinates().clone(),
366        };
367        Ok(Budgeted::new(projection, budget.empty_reservation()).into_shared()?)
368    }
369
370    fn original_coordinates(&self) -> &Arc<Budgeted<TextCoordinates>> {
371        self.original_coordinates.get_or_init(|| {
372            let budget = MemoryBudget::new(usize::MAX);
373            TextCoordinates::new_budgeted(self.original, &budget, &mut || Ok(()))
374                .and_then(|coordinates| Ok(coordinates.into_shared()?))
375                .expect("unbounded original source coordinates")
376        })
377    }
378
379    fn filtered_coordinates(&self) -> &Arc<Budgeted<TextCoordinates>> {
380        if self.maps.is_none() {
381            return self.original_coordinates();
382        }
383        self.filtered_coordinates.get_or_init(|| {
384            let budget = MemoryBudget::new(usize::MAX);
385            TextCoordinates::new_budgeted(self.as_str(), &budget, &mut || Ok(()))
386                .and_then(|coordinates| Ok(coordinates.into_shared()?))
387                .expect("unbounded filtered source coordinates")
388        })
389    }
390}
391
392/// Shared source provenance for tokens composed after the original input borrow ends.
393#[cfg(any(feature = "nori", feature = "kuromoji"))]
394#[derive(Debug, Clone, PartialEq, Eq)]
395pub(crate) struct SourceProjection {
396    source: Arc<Budgeted<String>>,
397    maps: Option<Arc<EditMaps>>,
398    original: Arc<Budgeted<TextCoordinates>>,
399    filtered: Arc<Budgeted<TextCoordinates>>,
400}
401
402#[cfg(any(feature = "nori", feature = "kuromoji"))]
403impl SourceProjection {
404    #[cfg(any(feature = "nori", feature = "kuromoji"))]
405    pub fn filtered_len(&self) -> usize {
406        self.filtered.utf16_len()
407    }
408
409    #[cfg(any(feature = "nori", feature = "kuromoji"))]
410    pub(crate) fn project_with_control(
411        &self,
412        mut range: Range<usize>,
413        poll: &mut dyn FnMut() -> AnalysisResult<()>,
414    ) -> AnalysisResult<SourceOffsets> {
415        poll()?;
416        self.filtered.validate_utf16_range(&range)?;
417        if let Some(maps) = &self.maps {
418            for map in maps.iter().rev() {
419                poll()?;
420                range = map.project_utf16(range);
421            }
422        }
423        self.original.covering_offsets_utf16(range)
424    }
425
426    #[cfg(any(feature = "nori", feature = "kuromoji"))]
427    pub(crate) fn is_verbatim_with_control(
428        &self,
429        term: &crate::TokenTerm,
430        offsets: &SourceOffsets,
431        poll: &mut dyn FnMut() -> AnalysisResult<()>,
432    ) -> AnalysisResult<bool> {
433        poll()?;
434        let Some(text) = term.as_str() else {
435            return Ok(false);
436        };
437        let Some(source) = self.source.get(offsets.utf8.clone()) else {
438            return Ok(false);
439        };
440        if text.len() != source.len() {
441            return Ok(false);
442        }
443        for (text, source) in text
444            .as_bytes()
445            .chunks(1024)
446            .zip(source.as_bytes().chunks(1024))
447        {
448            poll()?;
449            if text != source {
450                return Ok(false);
451            }
452        }
453        let mut length = 0;
454        for (index, unit) in text.chars().enumerate() {
455            if index % 1024 == 0 {
456                poll()?;
457            }
458            length += unit.len_utf16();
459        }
460        Ok(length == offsets.utf16.len())
461    }
462}
463
464fn validate_order(range: &Range<usize>) -> AnalysisResult<()> {
465    if range.start > range.end {
466        return Err(AnalysisError::InvalidTextSpan {
467            start: range.start,
468            end: range.end,
469        });
470    }
471    Ok(())
472}
473
474fn validate_utf8_range(text: &str, range: &Range<usize>) -> AnalysisResult<()> {
475    validate_order(range)?;
476    for offset in [range.start, range.end] {
477        if !text.is_char_boundary(offset) {
478            return Err(AnalysisError::InvalidTextOffset {
479                coordinate: "UTF-8",
480                offset,
481                length: text.len(),
482            });
483        }
484    }
485    Ok(())
486}