Skip to main content

cloudiful_redactor/session/
streaming_restore.rs

1use anyhow::Result;
2
3use crate::replace::{TOKEN_PREFIX, TOKEN_SUFFIX};
4use crate::{RedactionSession, RestorePermit, RestoreResult};
5
6use super::RestoreContext;
7
8const MAX_PENDING_BYTES: usize = 4096;
9
10/// Incrementally restores one decoded logical text stream.
11///
12/// Create a separate context for each logical stream. JSON string escapes must be decoded before
13/// calling `push_str`; serialize the restored text again after restoration. This type does not
14/// parse bytes, JSON framing, SSE, HTTP, or any other transport framing.
15#[derive(Debug)]
16pub struct StreamingRestoreContext<'a> {
17    restore: RestoreContext<'a>,
18    pending: String,
19}
20
21impl<'a> StreamingRestoreContext<'a> {
22    pub fn new(session: &'a RedactionSession) -> Self {
23        Self {
24            restore: RestoreContext::new(session),
25            pending: String::new(),
26        }
27    }
28
29    pub fn with_permits(session: &'a RedactionSession, permits: &[RestorePermit]) -> Result<Self> {
30        Ok(Self {
31            restore: RestoreContext::with_permits(session, permits)?,
32            pending: String::new(),
33        })
34    }
35
36    pub fn push_str(&mut self, fragment: &str) -> RestoreResult {
37        let mut result = empty_result(fragment.len());
38        for ch in fragment.chars() {
39            self.push_char(ch, &mut result);
40        }
41        result
42    }
43
44    pub fn finish(self) -> RestoreResult {
45        let mut result = empty_result(self.pending.len());
46        if self.pending.starts_with(TOKEN_PREFIX) {
47            result.unresolved_tokens.push(self.pending.clone());
48            result.validation_errors.push(format!(
49                "truncated token `{}`: missing closing `{TOKEN_SUFFIX}`",
50                self.pending
51            ));
52        }
53        result.restored_text = self.pending;
54        result
55    }
56
57    fn push_char(&mut self, ch: char, result: &mut RestoreResult) {
58        if self.pending.starts_with(TOKEN_PREFIX)
59            && self.pending.len() + ch.len_utf8() > MAX_PENDING_BYTES
60        {
61            let overflow = std::mem::take(&mut self.pending);
62            result.restored_text.push_str(&overflow);
63            result.unresolved_tokens.push(overflow.clone());
64            result.validation_errors.push(format!(
65                "token candidate exceeds {MAX_PENDING_BYTES} byte pending limit"
66            ));
67        }
68
69        self.pending.push(ch);
70        if self.pending.starts_with(TOKEN_PREFIX) {
71            if self.pending.ends_with(TOKEN_SUFFIX) {
72                let candidate = std::mem::take(&mut self.pending);
73                merge_result(result, self.restore.restore_text(&candidate));
74            }
75            return;
76        }
77
78        self.flush_non_marker_prefix(result);
79    }
80
81    fn flush_non_marker_prefix(&mut self, result: &mut RestoreResult) {
82        if TOKEN_PREFIX.starts_with(&self.pending) {
83            return;
84        }
85
86        let keep_from = self
87            .pending
88            .char_indices()
89            .map(|(index, _)| index)
90            .find(|&index| TOKEN_PREFIX.starts_with(&self.pending[index..]))
91            .unwrap_or(self.pending.len());
92        result.restored_text.push_str(&self.pending[..keep_from]);
93        self.pending.drain(..keep_from);
94    }
95}
96
97fn empty_result(capacity: usize) -> RestoreResult {
98    RestoreResult {
99        restored_text: String::with_capacity(capacity),
100        restored_count: 0,
101        skipped_tokens: Vec::new(),
102        unresolved_tokens: Vec::new(),
103        validation_errors: Vec::new(),
104    }
105}
106
107fn merge_result(target: &mut RestoreResult, source: RestoreResult) {
108    target.restored_text.push_str(&source.restored_text);
109    target.restored_count += source.restored_count;
110    target.skipped_tokens.extend(source.skipped_tokens);
111    target.unresolved_tokens.extend(source.unresolved_tokens);
112    target.validation_errors.extend(source.validation_errors);
113}