Skip to main content

clankerdiff_markdown/
stream.rs

1use crate::{
2    document::MarkdownDocument,
3    incremental::{IncrementalDocument, MarkdownParseStats},
4};
5use std::{
6    ops::Range,
7    sync::atomic::{AtomicU64, Ordering},
8};
9
10#[derive(Debug, Clone, PartialEq, Eq)]
11pub struct MarkdownStreamUpdate {
12    pub revision: u64,
13    pub changed_source: Range<usize>,
14    pub changed_blocks: Range<usize>,
15    pub reset: bool,
16    pub finished: bool,
17}
18
19#[derive(Debug, Clone, Copy, PartialEq, Eq)]
20pub struct MarkdownStreamChanges {
21    pub first_block: usize,
22    pub replaced: bool,
23}
24
25#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
26pub struct MarkdownStreamIdentity(u64);
27
28impl Default for MarkdownStreamIdentity {
29    fn default() -> Self {
30        static NEXT: AtomicU64 = AtomicU64::new(0);
31        Self(NEXT.fetch_add(1, Ordering::Relaxed))
32    }
33}
34
35#[derive(Debug, Default)]
36pub struct MarkdownStream {
37    inner: IncrementalDocument,
38    revision: u64,
39    source_revision: u64,
40    identity: MarkdownStreamIdentity,
41    finished: bool,
42    changes: Vec<(u64, usize)>,
43    replaced_at: u64,
44}
45
46impl Clone for MarkdownStream {
47    fn clone(&self) -> Self {
48        Self {
49            inner: self.inner.clone(),
50            revision: self.revision,
51            source_revision: self.source_revision,
52            identity: MarkdownStreamIdentity::default(),
53            finished: self.finished,
54            changes: self.changes.clone(),
55            replaced_at: self.replaced_at,
56        }
57    }
58}
59
60impl MarkdownStream {
61    #[must_use]
62    pub fn new() -> Self {
63        Self::default()
64    }
65
66    pub fn push(&mut self, chunk: &str) -> MarkdownStreamUpdate {
67        if chunk.is_empty() && !self.finished {
68            return self.unchanged();
69        }
70        let start = self.source().len();
71        let reset = std::mem::take(&mut self.finished);
72        let mut first_block = self.inner.document().blocks().len();
73        if !chunk.is_empty() {
74            first_block = self.inner.append(chunk);
75            self.source_revision = self.source_revision.wrapping_add(1);
76        }
77        self.revision = self.revision.wrapping_add(1);
78        self.record(first_block);
79        self.update(start..self.source().len(), first_block, reset)
80    }
81
82    pub fn replace(&mut self, source: impl Into<String>) -> MarkdownStreamUpdate {
83        let source = source.into();
84        self.inner.replace(&source);
85        self.source_revision = self.source_revision.wrapping_add(1);
86        self.finished = false;
87        self.revision = self.revision.wrapping_add(1);
88        self.changes.clear();
89        self.replaced_at = self.revision;
90        self.update(0..self.source().len(), 0, true)
91    }
92
93    pub fn finish(&mut self) -> MarkdownStreamUpdate {
94        if self.finished {
95            return self.unchanged();
96        }
97        self.finished = true;
98        self.revision = self.revision.wrapping_add(1);
99        let len = self.source().len();
100        self.update(len..len, self.block_count(), false)
101    }
102
103    #[must_use]
104    pub fn source(&self) -> &str {
105        self.inner.document().source()
106    }
107
108    #[must_use]
109    pub const fn document(&self) -> &MarkdownDocument {
110        self.inner.document()
111    }
112
113    #[must_use]
114    pub const fn revision(&self) -> u64 {
115        self.revision
116    }
117
118    #[must_use]
119    pub const fn source_revision(&self) -> u64 {
120        self.source_revision
121    }
122
123    #[must_use]
124    pub const fn identity(&self) -> MarkdownStreamIdentity {
125        self.identity
126    }
127
128    #[must_use]
129    pub const fn is_finished(&self) -> bool {
130        self.finished
131    }
132
133    #[must_use]
134    pub const fn settled_blocks(&self) -> usize {
135        self.inner.settled_blocks()
136    }
137
138    #[must_use]
139    pub fn open_code_block(&self) -> Option<usize> {
140        self.inner.open_code_block()
141    }
142
143    #[must_use]
144    pub const fn parse_stats(&self) -> MarkdownParseStats {
145        self.inner.stats()
146    }
147
148    #[must_use]
149    pub fn changes_since(&self, revision: u64) -> MarkdownStreamChanges {
150        if revision < self.replaced_at {
151            return MarkdownStreamChanges {
152                first_block: 0,
153                replaced: true,
154            };
155        }
156        let position = self
157            .changes
158            .partition_point(|(changed_at, _)| *changed_at <= revision);
159        MarkdownStreamChanges {
160            first_block: self
161                .changes
162                .get(position)
163                .map_or_else(|| self.block_count(), |(_, block)| *block),
164            replaced: false,
165        }
166    }
167
168    fn block_count(&self) -> usize {
169        self.inner.document().blocks().len()
170    }
171
172    fn record(&mut self, first_block: usize) {
173        if first_block >= self.block_count() {
174            return;
175        }
176        while self
177            .changes
178            .last()
179            .is_some_and(|(_, block)| *block >= first_block)
180        {
181            self.changes.pop();
182        }
183        self.changes.push((self.revision, first_block));
184    }
185
186    fn update(
187        &self,
188        changed_source: Range<usize>,
189        first_block: usize,
190        reset: bool,
191    ) -> MarkdownStreamUpdate {
192        MarkdownStreamUpdate {
193            revision: self.revision,
194            changed_source,
195            changed_blocks: first_block.min(self.block_count())..self.block_count(),
196            reset,
197            finished: self.finished,
198        }
199    }
200
201    fn unchanged(&self) -> MarkdownStreamUpdate {
202        let len = self.source().len();
203        self.update(len..len, self.block_count(), false)
204    }
205}