1use 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#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
30pub struct SourceOffsets {
31 pub utf8: Range<usize>,
32 pub utf16: Range<usize>,
33}
34
35#[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 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 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#[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 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 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 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 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#[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}