lol_html/rewritable_units/
mutations.rs

1use super::StreamingHandlerSink;
2use std::error::Error as StdError;
3use std::panic::{RefUnwindSafe, UnwindSafe};
4use std::sync::Mutex;
5
6type HandlerResult = Result<(), Box<dyn StdError + Send + Sync>>;
7
8/// The type of inserted content.
9#[derive(Copy, Clone)]
10pub enum ContentType {
11    /// HTML content type. The rewriter will insert the content as is.
12    Html,
13    /// Text content type. The rewriter will HTML-escape the content before insertion:
14    ///     - `<` will be replaced with `&lt;`
15    ///     - `>` will be replaced with `&gt;`
16    ///     - `&` will be replaced with `&amp;`
17    Text,
18}
19
20pub(crate) struct MutationsInner {
21    pub content_before: DynamicString,
22    pub replacement: DynamicString,
23    pub content_after: DynamicString,
24    pub removed: bool,
25}
26
27impl MutationsInner {
28    #[inline]
29    pub fn replace(&mut self, chunk: StringChunk) {
30        self.remove();
31        self.replacement.clear();
32        self.replacement.push_back(chunk);
33    }
34
35    #[inline]
36    pub fn remove(&mut self) {
37        self.removed = true;
38    }
39}
40
41pub(crate) struct Mutations {
42    inner: Option<Box<MutationsInner>>,
43}
44
45impl Mutations {
46    #[inline]
47    #[must_use]
48    pub const fn new() -> Self {
49        Self { inner: None }
50    }
51
52    #[inline]
53    pub fn take(&mut self) -> Option<Box<MutationsInner>> {
54        self.inner.take()
55    }
56
57    #[inline]
58    pub fn if_mutated(&mut self) -> Option<&mut MutationsInner> {
59        self.inner.as_deref_mut()
60    }
61
62    #[inline]
63    pub fn mutate(&mut self) -> &mut MutationsInner {
64        #[inline(never)]
65        fn alloc_content(inner: &mut Option<Box<MutationsInner>>) -> &mut MutationsInner {
66            inner.get_or_insert_with(move || {
67                Box::new(MutationsInner {
68                    content_before: DynamicString::new(),
69                    replacement: DynamicString::new(),
70                    content_after: DynamicString::new(),
71                    removed: false,
72                })
73            })
74        }
75
76        match &mut self.inner {
77            Some(inner) => inner,
78            uninit => alloc_content(uninit),
79        }
80    }
81
82    #[inline]
83    pub fn removed(&self) -> bool {
84        self.inner.as_ref().is_some_and(|inner| inner.removed)
85    }
86}
87
88/// Part of [`DynamicString`]
89pub(crate) enum StringChunk {
90    Buffer(Box<str>, ContentType),
91    // The mutex is never actually locked, but makes the struct `Sync` without unsafe.
92    Stream(Mutex<Box<dyn StreamingHandler + Send + 'static>>),
93}
94
95impl StringChunk {
96    pub(crate) fn from_str(content: impl Into<Box<str>>, content_type: ContentType) -> Self {
97        Self::Buffer(content.into(), content_type)
98    }
99
100    #[inline]
101    pub(crate) fn stream(handler: Box<dyn StreamingHandler + Send + 'static>) -> Self {
102        Self::Stream(Mutex::new(handler))
103    }
104}
105
106/// String built from fragments or dynamic callbacks
107#[derive(Default)]
108pub(crate) struct DynamicString {
109    chunks: Vec<StringChunk>,
110}
111
112impl DynamicString {
113    #[inline]
114    pub const fn new() -> Self {
115        Self { chunks: vec![] }
116    }
117
118    #[inline]
119    pub fn clear(&mut self) {
120        self.chunks.clear();
121    }
122
123    #[inline]
124    pub fn push_front(&mut self, chunk: StringChunk) {
125        self.chunks.insert(0, chunk);
126    }
127
128    #[inline]
129    pub fn push_back(&mut self, chunk: StringChunk) {
130        self.chunks.push(chunk);
131    }
132
133    pub fn encode(self, sink: &mut StreamingHandlerSink<'_>) -> HandlerResult {
134        for chunk in self.chunks {
135            match chunk {
136                StringChunk::Buffer(content, content_type) => {
137                    sink.write_str(&content, content_type);
138                }
139                StringChunk::Stream(handler) => {
140                    // The mutex will never be locked or poisoned. This is the cheapest way to unwrap it.
141                    let (Ok(h) | Err(h)) = handler
142                        .into_inner()
143                        .map_err(std::sync::PoisonError::into_inner);
144                    h.write_all(sink)?;
145                }
146            }
147        }
148        Ok(())
149    }
150}
151
152/// A callback used to write content asynchronously.
153///
154/// Use the [`streaming!`] macro to construct it.
155#[diagnostic::on_unimplemented(
156    note = "use `streaming!` macro to create the handler",
157    label = "Must be `FnOnce(&mut StreamingHandlerSink<'_>) -> Result<(), Box<dyn std::error::Error + Send + Sync>> + Send + 'static`"
158)]
159pub trait StreamingHandler {
160    /// This method is called only once, and is expected to write content
161    /// by calling the [`sink.write_str()`](StreamingHandlerSink::write_str) one or more times.
162    ///
163    /// Multiple calls to `sink.write_str()` append more content to the output.
164    ///
165    /// See [`StreamingHandlerSink`].
166    ///
167    /// Note: if you get "implementation of `FnOnce` is not general enough" error, add explicit argument
168    /// `sink: &mut StreamingHandlerSink<'_>` to the closure.
169    fn write_all(self: Box<Self>, sink: &mut StreamingHandlerSink<'_>) -> HandlerResult;
170}
171
172impl RefUnwindSafe for StringChunk {}
173impl UnwindSafe for StringChunk {}
174
175impl<F> From<F> for Box<dyn StreamingHandler + Send + 'static>
176where
177    F: FnOnce(&mut StreamingHandlerSink<'_>) -> HandlerResult + Send + 'static,
178{
179    #[inline]
180    fn from(f: F) -> Self {
181        Box::new(f)
182    }
183}
184
185impl<F> StreamingHandler for F
186where
187    F: FnOnce(&mut StreamingHandlerSink<'_>) -> HandlerResult + Send + 'static,
188{
189    #[inline]
190    fn write_all(self: Box<F>, sink: &mut StreamingHandlerSink<'_>) -> HandlerResult {
191        (self)(sink)
192    }
193}
lol_html/rewritable_units/mutations.rs

lol_html/rewritable_units/
mutations.rs
