Skip to main content

polyglot_stream_heal/
lib.rs

1pub mod discriminator;
2pub mod html_pda;
3pub mod json_pda;
4pub mod mermaid_pda;
5pub mod partitioner;
6
7pub use discriminator::Modality;
8pub use html_pda::{heal_html, IncrementalHtmlPda};
9pub use json_pda::{heal_json, IncrementalJsonPda};
10pub use mermaid_pda::{heal_mermaid, IncrementalMermaidPda};
11pub use partitioner::partition_visual_spec;
12
13use serde::{Deserialize, Serialize};
14
15#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
16pub struct TransducerConfig {
17    pub max_table_cols: usize,
18    pub max_chart_bars: usize,
19    pub partition_visuals: bool,
20    pub repair_nested_escaped_json: bool,
21}
22
23impl Default for TransducerConfig {
24    fn default() -> Self {
25        Self {
26            max_table_cols: 5,
27            max_chart_bars: 10,
28            partition_visuals: true,
29            repair_nested_escaped_json: true,
30        }
31    }
32}
33
34#[derive(Debug, Clone, Serialize, Deserialize)]
35pub struct HealedSegment {
36    pub index: usize,
37    pub modality: Modality,
38    pub fence_tag: String,
39    pub content: String,
40    pub is_closed: bool,
41    pub was_partitioned: bool,
42}
43
44#[derive(Debug, Clone, Serialize, Deserialize)]
45pub struct StreamSnapshot {
46    pub segments: Vec<HealedSegment>,
47    pub reconstructed_markdown: String,
48    pub total_bytes_processed: usize,
49}
50
51#[derive(Debug)]
52enum ActiveInnerPda {
53    Json(IncrementalJsonPda),
54    Mermaid(IncrementalMermaidPda),
55    Html(IncrementalHtmlPda),
56    Passthrough(String),
57}
58
59#[derive(Debug, Clone, Copy, PartialEq, Eq)]
60enum OuterFenceState {
61    ProseStartOfLine,
62    ProseMidLine,
63    ProseBackticks(u8),
64    FenceHeader { backtick_len: u8 },
65    InsideBlockStartOfLine { backtick_len: u8 },
66    InsideBlockMidLine { backtick_len: u8 },
67    InsideBlockClosingBackticks { backtick_len: u8, seen: u8 },
68}
69
70/// Hierarchical Multi-Grammar Pushdown Transducer (`IncrementalTransducer`).
71///
72/// Processes arbitrary byte or UTF-8 token chunks in strict O(|\Delta s_t|) amortized time,
73/// dispatching outer Markdown prose and 19+ fenced DSL modalities into specialized inner
74/// pushdown automata without ever re-scanning previously committed segments.
75#[derive(Debug)]
76pub struct IncrementalTransducer {
77    config: TransducerConfig,
78    outer_state: OuterFenceState,
79    utf8_tail: Vec<u8>,
80    total_bytes: usize,
81    completed_segments: Vec<HealedSegment>,
82    active_prose: String,
83    active_fence_tag: String,
84    active_modality: Modality,
85    active_inner: Option<ActiveInnerPda>,
86    pending_line_prefix: String,
87}
88
89impl Default for IncrementalTransducer {
90    fn default() -> Self {
91        Self::new(TransducerConfig::default())
92    }
93}
94
95impl IncrementalTransducer {
96    pub fn new(config: TransducerConfig) -> Self {
97        Self {
98            config,
99            outer_state: OuterFenceState::ProseStartOfLine,
100            utf8_tail: Vec::with_capacity(4),
101            total_bytes: 0,
102            completed_segments: Vec::with_capacity(8),
103            active_prose: String::with_capacity(256),
104            active_fence_tag: String::new(),
105            active_modality: Modality::Markdown,
106            active_inner: None,
107            pending_line_prefix: String::new(),
108        }
109    }
110
111    pub fn reset(&mut self) {
112        self.outer_state = OuterFenceState::ProseStartOfLine;
113        self.utf8_tail.clear();
114        self.total_bytes = 0;
115        self.completed_segments.clear();
116        self.active_prose.clear();
117        self.active_fence_tag.clear();
118        self.active_modality = Modality::Markdown;
119        self.active_inner = None;
120        self.pending_line_prefix.clear();
121    }
122
123    /// Feeds raw network/SSE bytes, safely buffering split UTF-8 codepoints across chunk boundaries.
124    pub fn push_bytes(&mut self, bytes: &[u8]) {
125        if bytes.is_empty() {
126            return;
127        }
128        if self.utf8_tail.is_empty() {
129            match std::str::from_utf8(bytes) {
130                Ok(valid) => {
131                    self.push_chunk(valid);
132                    return;
133                }
134                Err(err) => {
135                    let valid_up_to = err.valid_up_to();
136                    if valid_up_to > 0 {
137                        let valid = unsafe { std::str::from_utf8_unchecked(&bytes[..valid_up_to]) };
138                        self.push_chunk(valid);
139                    }
140                    self.utf8_tail.extend_from_slice(&bytes[valid_up_to..]);
141                    return;
142                }
143            }
144        }
145
146        self.utf8_tail.extend_from_slice(bytes);
147        let combined = std::mem::take(&mut self.utf8_tail);
148        match std::str::from_utf8(&combined) {
149            Ok(valid) => self.push_chunk(valid),
150            Err(err) => {
151                let valid_up_to = err.valid_up_to();
152                if valid_up_to > 0 {
153                    let valid = unsafe { std::str::from_utf8_unchecked(&combined[..valid_up_to]) };
154                    self.push_chunk(valid);
155                }
156                if err.error_len().is_none() {
157                    self.utf8_tail.extend_from_slice(&combined[valid_up_to..]);
158                }
159            }
160        }
161    }
162
163    /// Feeds a UTF-8 string chunk in O(|chunk|) time.
164    pub fn push_chunk(&mut self, chunk: &str) {
165        self.total_bytes += chunk.len();
166        for ch in chunk.chars() {
167            self.step_char(ch);
168        }
169    }
170
171    fn step_char(&mut self, ch: char) {
172        match self.outer_state {
173            OuterFenceState::ProseStartOfLine => {
174                if ch == '`' {
175                    self.outer_state = OuterFenceState::ProseBackticks(1);
176                } else if ch == '\n' {
177                    self.active_prose.push('\n');
178                } else if ch == ' ' || ch == '\t' {
179                    self.pending_line_prefix.push(ch);
180                } else {
181                    if !self.pending_line_prefix.is_empty() {
182                        self.active_prose.push_str(&self.pending_line_prefix);
183                        self.pending_line_prefix.clear();
184                    }
185                    self.active_prose.push(ch);
186                    self.outer_state = OuterFenceState::ProseMidLine;
187                }
188            }
189            OuterFenceState::ProseMidLine => {
190                self.active_prose.push(ch);
191                if ch == '\n' {
192                    self.outer_state = OuterFenceState::ProseStartOfLine;
193                }
194            }
195            OuterFenceState::ProseBackticks(count) => {
196                if ch == '`' {
197                    self.outer_state = OuterFenceState::ProseBackticks(count.saturating_add(1));
198                } else if count >= 3 {
199                    // Commit accumulated prose segment if non-empty
200                    self.pending_line_prefix.clear();
201                    if !self.active_prose.trim().is_empty() {
202                        let prose = std::mem::take(&mut self.active_prose);
203                        self.completed_segments.push(HealedSegment {
204                            index: self.completed_segments.len(),
205                            modality: Modality::Markdown,
206                            fence_tag: String::new(),
207                            content: prose,
208                            is_closed: true,
209                            was_partitioned: false,
210                        });
211                    } else {
212                        self.active_prose.clear();
213                    }
214                    self.active_fence_tag.clear();
215                    if ch == '\n' {
216                        self.enter_fenced_block(count);
217                    } else if ch != '\r' {
218                        self.active_fence_tag.push(ch);
219                        self.outer_state = OuterFenceState::FenceHeader {
220                            backtick_len: count,
221                        };
222                    } else {
223                        self.outer_state = OuterFenceState::FenceHeader {
224                            backtick_len: count,
225                        };
226                    }
227                } else {
228                    // Fewer than 3 backticks -> inline code span inside prose
229                    if !self.pending_line_prefix.is_empty() {
230                        self.active_prose.push_str(&self.pending_line_prefix);
231                        self.pending_line_prefix.clear();
232                    }
233                    for _ in 0..count {
234                        self.active_prose.push('`');
235                    }
236                    self.active_prose.push(ch);
237                    self.outer_state = if ch == '\n' {
238                        OuterFenceState::ProseStartOfLine
239                    } else {
240                        OuterFenceState::ProseMidLine
241                    };
242                }
243            }
244            OuterFenceState::FenceHeader { backtick_len } => {
245                if ch == '\n' {
246                    self.enter_fenced_block(backtick_len);
247                } else if ch != '\r' {
248                    self.active_fence_tag.push(ch);
249                }
250            }
251            OuterFenceState::InsideBlockStartOfLine { backtick_len } => {
252                if ch == '`' {
253                    self.outer_state = OuterFenceState::InsideBlockClosingBackticks {
254                        backtick_len,
255                        seen: 1,
256                    };
257                } else if ch == ' ' || ch == '\t' {
258                    self.pending_line_prefix.push(ch);
259                } else if ch == '\n' {
260                    if !self.pending_line_prefix.is_empty() {
261                        let prefix = std::mem::take(&mut self.pending_line_prefix);
262                        self.feed_active_inner(&prefix);
263                    }
264                    self.feed_active_inner("\n");
265                } else {
266                    if !self.pending_line_prefix.is_empty() {
267                        let prefix = std::mem::take(&mut self.pending_line_prefix);
268                        self.feed_active_inner(&prefix);
269                    }
270                    let mut buf = [0u8; 4];
271                    self.feed_active_inner(ch.encode_utf8(&mut buf));
272                    self.outer_state = OuterFenceState::InsideBlockMidLine { backtick_len };
273                }
274            }
275            OuterFenceState::InsideBlockMidLine { backtick_len } => {
276                let mut buf = [0u8; 4];
277                self.feed_active_inner(ch.encode_utf8(&mut buf));
278                if ch == '\n' {
279                    self.outer_state = OuterFenceState::InsideBlockStartOfLine { backtick_len };
280                }
281            }
282            OuterFenceState::InsideBlockClosingBackticks { backtick_len, seen } => {
283                if ch == '`' {
284                    self.outer_state = OuterFenceState::InsideBlockClosingBackticks {
285                        backtick_len,
286                        seen: seen.saturating_add(1),
287                    };
288                } else if seen >= backtick_len && (ch == '\n' || ch == '\r' || ch.is_whitespace()) {
289                    if ch == '\n' {
290                        self.pending_line_prefix.clear();
291                        self.finalize_active_block();
292                        self.outer_state = OuterFenceState::ProseStartOfLine;
293                    } else {
294                        // Wait for newline after closing fence (e.g. trailing spaces or `\r`)
295                        self.outer_state = OuterFenceState::InsideBlockClosingBackticks {
296                            backtick_len,
297                            seen,
298                        };
299                    }
300                } else {
301                    // False alarm: backticks inside block body
302                    if !self.pending_line_prefix.is_empty() {
303                        let prefix = std::mem::take(&mut self.pending_line_prefix);
304                        self.feed_active_inner(&prefix);
305                    }
306                    for _ in 0..seen {
307                        self.feed_active_inner("`");
308                    }
309                    let mut buf = [0u8; 4];
310                    self.feed_active_inner(ch.encode_utf8(&mut buf));
311                    self.outer_state = if ch == '\n' {
312                        OuterFenceState::InsideBlockStartOfLine { backtick_len }
313                    } else {
314                        OuterFenceState::InsideBlockMidLine { backtick_len }
315                    };
316                }
317            }
318        }
319    }
320
321    fn enter_fenced_block(&mut self, backtick_len: u8) {
322        let modality = Modality::from_fence_tag(&self.active_fence_tag);
323        self.active_modality = modality;
324        self.active_inner = Some(if modality.is_json_backed() {
325            ActiveInnerPda::Json(IncrementalJsonPda::new(
326                self.config.repair_nested_escaped_json,
327            ))
328        } else if modality == Modality::Mermaid {
329            ActiveInnerPda::Mermaid(IncrementalMermaidPda::new())
330        } else if modality == Modality::HtmlSandbox {
331            ActiveInnerPda::Html(IncrementalHtmlPda::new())
332        } else {
333            ActiveInnerPda::Passthrough(String::with_capacity(256))
334        });
335        self.outer_state = OuterFenceState::InsideBlockStartOfLine { backtick_len };
336    }
337
338    fn feed_active_inner(&mut self, s: &str) {
339        match &mut self.active_inner {
340            Some(ActiveInnerPda::Json(pda)) => pda.push_chunk(s),
341            Some(ActiveInnerPda::Mermaid(pda)) => pda.push_chunk(s),
342            Some(ActiveInnerPda::Html(pda)) => pda.push_chunk(s),
343            Some(ActiveInnerPda::Passthrough(buf)) => buf.push_str(s),
344            None => {}
345        }
346    }
347
348    fn render_inner_snapshot(&self, inner: &ActiveInnerPda, modality: Modality) -> (String, bool) {
349        match inner {
350            ActiveInnerPda::Json(pda) => {
351                let healed = pda.snapshot_healed();
352                if self.config.partition_visuals
353                    && matches!(modality, Modality::DynamicVisual | Modality::Json)
354                    && !healed.is_empty()
355                {
356                    if let Ok(partitioned) = partition_visual_spec(
357                        &healed,
358                        self.config.max_table_cols,
359                        self.config.max_chart_bars,
360                    ) {
361                        let changed = partitioned.len() != healed.len() || partitioned != healed;
362                        return (partitioned, changed);
363                    }
364                }
365                (healed, false)
366            }
367            ActiveInnerPda::Mermaid(pda) => (pda.snapshot_healed(), false),
368            ActiveInnerPda::Html(pda) => (pda.snapshot_healed(), false),
369            ActiveInnerPda::Passthrough(buf) => (buf.trim_end().to_string(), false),
370        }
371    }
372
373    fn finalize_active_block(&mut self) {
374        if let Some(inner) = self.active_inner.take() {
375            let (content, was_partitioned) = self.render_inner_snapshot(&inner, self.active_modality);
376            let tag = std::mem::take(&mut self.active_fence_tag);
377            self.completed_segments.push(HealedSegment {
378                index: self.completed_segments.len(),
379                modality: self.active_modality,
380                fence_tag: tag.trim().to_string(),
381                content,
382                is_closed: true,
383                was_partitioned,
384            });
385        }
386        self.active_modality = Modality::Markdown;
387    }
388
389    /// Returns ONLY the currently active (in-flight) segment healed in O(d + |active_tail|) time,
390    /// without cloning or re-serializing any previously completed segments 0..m-1 (Theorem 2).
391    pub fn snapshot_active_segment(&self) -> Option<HealedSegment> {
392        let eof_closed_fence = match self.outer_state {
393            OuterFenceState::InsideBlockClosingBackticks { backtick_len, seen } => {
394                seen >= backtick_len
395            }
396            _ => false,
397        };
398
399        if let Some(inner) = &self.active_inner {
400            let (content, was_partitioned) = self.render_inner_snapshot(inner, self.active_modality);
401            return Some(HealedSegment {
402                index: self.completed_segments.len(),
403                modality: self.active_modality,
404                fence_tag: self.active_fence_tag.trim().to_string(),
405                content,
406                is_closed: eof_closed_fence,
407                was_partitioned,
408            });
409        }
410
411        if !self.active_prose.trim().is_empty() {
412            return Some(HealedSegment {
413                index: self.completed_segments.len(),
414                modality: Modality::Markdown,
415                fence_tag: String::new(),
416                content: self.active_prose.clone(),
417                is_closed: true,
418                was_partitioned: false,
419            });
420        }
421
422        self.completed_segments.last().cloned()
423    }
424
425    /// Produces a zero-mutation healed snapshot of the entire polyglot stream in O(d + |active_tail|) time.
426    pub fn snapshot(&self) -> StreamSnapshot {
427        let mut segments = self.completed_segments.clone();
428
429        // Check if closing backticks reached EOF without a trailing newline
430        let eof_closed_fence = match self.outer_state {
431            OuterFenceState::InsideBlockClosingBackticks { backtick_len, seen } => {
432                seen >= backtick_len
433            }
434            _ => false,
435        };
436
437        if let Some(inner) = &self.active_inner {
438            let (content, was_partitioned) = self.render_inner_snapshot(inner, self.active_modality);
439            segments.push(HealedSegment {
440                index: segments.len(),
441                modality: self.active_modality,
442                fence_tag: self.active_fence_tag.trim().to_string(),
443                content,
444                is_closed: eof_closed_fence,
445                was_partitioned,
446            });
447        } else if !self.active_prose.trim().is_empty() {
448            // If the entire stream has zero code fences, check semantic fallback for bare JSON/Mermaid/HTML
449            if segments.is_empty() {
450                let inferred = Modality::infer_unfenced(&self.active_prose);
451                match inferred {
452                    Modality::Json => {
453                        let healed = heal_json(&self.active_prose);
454                        let (final_json, was_partitioned) = if self.config.partition_visuals {
455                            match partition_visual_spec(
456                                &healed,
457                                self.config.max_table_cols,
458                                self.config.max_chart_bars,
459                            ) {
460                                Ok(p) => {
461                                    let changed = p != healed;
462                                    (p, changed)
463                                }
464                                Err(_) => (healed, false),
465                            }
466                        } else {
467                            (healed, false)
468                        };
469                        segments.push(HealedSegment {
470                            index: 0,
471                            modality: Modality::Json,
472                            fence_tag: String::new(),
473                            content: final_json,
474                            is_closed: false,
475                            was_partitioned,
476                        });
477                    }
478                    Modality::Mermaid => {
479                        segments.push(HealedSegment {
480                            index: 0,
481                            modality: Modality::Mermaid,
482                            fence_tag: String::new(),
483                            content: heal_mermaid(&self.active_prose),
484                            is_closed: false,
485                            was_partitioned: false,
486                        });
487                    }
488                    Modality::HtmlSandbox => {
489                        segments.push(HealedSegment {
490                            index: 0,
491                            modality: Modality::HtmlSandbox,
492                            fence_tag: String::new(),
493                            content: heal_html(&self.active_prose),
494                            is_closed: false,
495                            was_partitioned: false,
496                        });
497                    }
498                    _ => {
499                        segments.push(HealedSegment {
500                            index: 0,
501                            modality: Modality::Markdown,
502                            fence_tag: String::new(),
503                            content: self.active_prose.clone(),
504                            is_closed: true,
505                            was_partitioned: false,
506                        });
507                    }
508                }
509            } else {
510                segments.push(HealedSegment {
511                    index: segments.len(),
512                    modality: Modality::Markdown,
513                    fence_tag: String::new(),
514                    content: self.active_prose.clone(),
515                    is_closed: true,
516                    was_partitioned: false,
517                });
518            }
519        }
520
521        let mut reconstructed = String::with_capacity(self.total_bytes + 64);
522        for (idx, seg) in segments.iter().enumerate() {
523            if idx > 0 && !reconstructed.ends_with('\n') {
524                reconstructed.push('\n');
525            }
526            if seg.modality == Modality::Markdown {
527                reconstructed.push_str(&seg.content);
528            } else if seg.fence_tag.is_empty() && segments.len() == 1 {
529                reconstructed.push_str(&seg.content);
530            } else {
531                let tag = if seg.fence_tag.is_empty() {
532                    "code"
533                } else {
534                    seg.fence_tag.as_str()
535                };
536                reconstructed.push_str("```");
537                reconstructed.push_str(tag);
538                reconstructed.push('\n');
539                reconstructed.push_str(&seg.content);
540                reconstructed.push_str("\n```\n");
541            }
542        }
543
544        StreamSnapshot {
545            segments,
546            reconstructed_markdown: reconstructed,
547            total_bytes_processed: self.total_bytes,
548        }
549    }
550}
551
552/// Stateless one-shot helper over a complete or truncated polyglot stream.
553pub fn heal_stream(raw: &str, config: TransducerConfig) -> StreamSnapshot {
554    let mut t = IncrementalTransducer::new(config);
555    t.push_chunk(raw);
556    t.snapshot()
557}
558
559#[cfg(test)]
560mod tests {
561    use super::*;
562
563    #[test]
564    fn heals_arbitrary_mid_token_json_truncations() {
565        let full = r#"{"title":"System Latency","metrics":[{"name":"p50","val":1.25,"ok":true},{"name":"p99","val":4.8,"ok":false}]}"#;
566        for cut in 1..full.len() {
567            let prefix = &full[..cut];
568            let healed = heal_json(prefix);
569            let parsed: Result<serde_json::Value, _> = serde_json::from_str(&healed);
570            assert!(
571                parsed.is_ok(),
572                "Failed JSON parse at cut={} prefix={:?} -> healed={:?}",
573                cut,
574                prefix,
575                healed
576            );
577        }
578    }
579
580    #[test]
581    fn heals_nested_escaped_tool_call_json_arguments() {
582        let truncated_tool =
583            r#"{"tool":"query_cluster","arguments":"{\"filters\":[{\"field\":\"latency\",\"op\":\">=\",\"val\":12."}"#;
584        // Cut off even before the outer string closes:
585        let mid_stream = &truncated_tool[..truncated_tool.len() - 2];
586        let healed = heal_json(mid_stream);
587        let outer: serde_json::Value = serde_json::from_str(&healed).expect("valid outer JSON");
588        let inner_str = outer["arguments"].as_str().expect("arguments string");
589        let inner: serde_json::Value =
590            serde_json::from_str(inner_str).expect("valid healed inner escaped JSON");
591        assert_eq!(inner["filters"][0]["field"], "latency");
592        assert_eq!(inner["filters"][0]["val"], 12.0);
593    }
594
595    #[test]
596    fn heals_mermaid_dyck_nested_parentheses_and_subgraphs() {
597        let raw = "flowchart LR\n  subgraph Compute Tier (v2.1)\n  A[Evaluate O(f(g(h(x))))] -- async rpc --> B[(Primary DB (Postgres: 16.2))]\n  B --> C{Latency < p99 (5ms)?";
598        let healed = heal_mermaid(raw);
599        assert!(healed.contains(r#"A["Evaluate O(f(g(h(x))))"]"#), "got:\n{}", healed);
600        assert!(
601            healed.contains(r#"B[("Primary DB (Postgres: 16.2)")]"#),
602            "got:\n{}",
603            healed
604        );
605        assert!(
606            healed.contains(r#"C{"Latency < p99 (5ms)?"}"#),
607            "got:\n{}",
608            healed
609        );
610        assert!(healed.contains("-->|async rpc|"), "got:\n{}", healed);
611        assert!(healed.ends_with("end"), "must auto-close subgraph:\n{}", healed);
612    }
613
614    #[test]
615    fn partitions_wide_tables_and_dense_charts_losslessly() {
616        let wide = serde_json::json!({
617            "sections": [
618                {
619                    "kind": "comparison_table",
620                    "title": "Engine Matrix",
621                    "columns": ["Engine", "C1", "C2", "C3", "C4", "C5", "C6", "C7", "C8"],
622                    "rows": [
623                        ["HM-PDT", "1", "2", "3", "4", "5", "6", "7", "8"],
624                        ["Regex", "a", "b", "c", "d", "e", "f", "g", "h"]
625                    ]
626                }
627            ]
628        });
629        let out_str = partition_visual_spec(&wide.to_string(), 5, 10).unwrap();
630        let out: serde_json::Value = serde_json::from_str(&out_str).unwrap();
631        let sections = out["sections"].as_array().unwrap();
632        assert_eq!(sections.len(), 2);
633        for sec in sections {
634            let cols = sec["columns"].as_array().unwrap();
635            assert!(cols.len() <= 5);
636            assert_eq!(cols[0], "Engine");
637        }
638    }
639
640    #[test]
641    fn handles_utf8_multibyte_splits_across_polyglot_fences() {
642        let doc = "Intro — αβγ 🚀\n```dynamic_visual\n{\"title\":\"Metric Δ\",\"bars\":[{\"label\":\"A\",\"value\":42}";
643        let bytes = doc.as_bytes();
644        let mut transducer = IncrementalTransducer::default();
645        for b in bytes.chunks(3) {
646            transducer.push_bytes(b);
647        }
648        let snap = transducer.snapshot();
649        assert_eq!(snap.segments.len(), 2);
650        assert_eq!(snap.segments[0].modality, Modality::Markdown);
651        assert!(snap.segments[0].content.contains("αβγ 🚀"));
652        assert_eq!(snap.segments[1].modality, Modality::DynamicVisual);
653        let parsed: serde_json::Value = serde_json::from_str(&snap.segments[1].content).unwrap();
654        assert_eq!(parsed["title"], "Metric Δ");
655    }
656}