lean_ctx_doc/
sections.rs

1// SPDX-License-Identifier: Apache-2.0
2
3//! Splitting a markdown body into sections and joining it back.
4//!
5//! A body splits at its `## ` headings. A `## ` block larger than the limit
6//! splits further at `### `, then at blank lines, then at line ends, and a
7//! single overlong line at a character boundary. Headings inside code fences
8//! are text. The split is lossless: [`join`] of [`split`] gives the body back
9//! whenever it opens with its `# ` title line.
10
11use crate::model::Section;
12use crate::{MAX_DOC_BODY_BYTES, ORDER_GAP, SECTION_MAX_BYTES};
13
14/// Splits `body` into its title and sections no larger than `max` bytes.
15pub fn split_with(body: &str, max: usize) -> (String, Vec<Section>) {
16    let max = max.max(4);
17    let (title, rest) = take_title(body);
18    let mut chunks = Vec::new();
19    for block in blocks(rest, 2) {
20        pack_into(&mut chunks, block, max, 3);
21    }
22    (title, lay_out(chunks))
23}
24
25/// [`split_with`] at [`SECTION_MAX_BYTES`].
26pub fn split(body: &str) -> (String, Vec<Section>) {
27    split_with(body, SECTION_MAX_BYTES)
28}
29
30/// The body of `title` and `sections`, sections in `order`.
31pub fn join(title: &str, sections: &[Section]) -> String {
32    let mut ordered: Vec<&Section> = sections.iter().collect();
33    ordered.sort_by_key(|s| s.order);
34    let mut out = String::new();
35    if !title.is_empty() {
36        out.push_str("# ");
37        out.push_str(title);
38        out.push('\n');
39    }
40    for s in ordered {
41        out.push_str(&s.text);
42    }
43    out
44}
45
46/// Fresh orders for `n` sections: `ORDER_GAP`, `2 * ORDER_GAP`, ….
47pub fn renumber(n: usize) -> Vec<u32> {
48    (1..=n as u32)
49        .map(|i| i.saturating_mul(ORDER_GAP))
50        .collect()
51}
52
53/// An order strictly between `before` and `after` (either end open), or
54/// `None` when no integer fits and the page has to be renumbered.
55pub fn order_between(before: Option<u32>, after: Option<u32>) -> Option<u32> {
56    let lo = before.unwrap_or(0);
57    let Some(hi) = after else {
58        return lo.checked_add(ORDER_GAP);
59    };
60    (hi > lo.saturating_add(1)).then(|| lo + (hi - lo) / 2)
61}
62
63fn lay_out(chunks: Vec<String>) -> Vec<Section> {
64    let chunks: Vec<String> = chunks.into_iter().filter(|c| !c.is_empty()).collect();
65    renumber(chunks.len())
66        .into_iter()
67        .zip(chunks)
68        .map(|(order, text)| Section { order, text })
69        .collect()
70}
71
72/// Strips a leading `# ` title line; a body without one has an empty title.
73fn take_title(body: &str) -> (String, &str) {
74    let first_end = body.find('\n').map_or(body.len(), |i| i + 1);
75    let first = &body[..first_end];
76    match first.strip_prefix("# ") {
77        Some(t) => (
78            t.trim_end_matches(['\n', '\r']).to_owned(),
79            &body[first_end..],
80        ),
81        None => (String::new(), body),
82    }
83}
84
85/// Is `line` an ATX heading of exactly `level`?
86fn is_heading(line: &str, level: usize) -> bool {
87    let b = line.as_bytes();
88    b.len() > level && b[..level].iter().all(|c| *c == b'#') && b[level] == b' '
89}
90
91/// Tracks fenced code blocks line by line.
92#[derive(Default)]
93pub(crate) struct Fence {
94    open: Option<(u8, usize)>,
95}
96
97impl Fence {
98    /// Feeds one line; true while the line is inside (or opens/closes) a fence.
99    pub(crate) fn feed(&mut self, line: &str) -> bool {
100        let t = line.trim_start();
101        let marker = t.bytes().next().filter(|c| *c == b'`' || *c == b'~');
102        let run = marker.map_or(0, |m| t.bytes().take_while(|c| *c == m).count());
103        match (self.open, marker) {
104            (Some((m, n)), Some(c)) if c == m && run >= n && t[run..].trim().is_empty() => {
105                self.open = None;
106                true
107            }
108            (Some(_), _) => true,
109            (None, Some(c)) if run >= 3 => {
110                self.open = Some((c, run));
111                true
112            }
113            (None, _) => false,
114        }
115    }
116}
117
118/// Splits `text` before every heading of `level` outside fences.
119fn blocks(text: &str, level: usize) -> Vec<&str> {
120    let mut out = Vec::new();
121    let mut fence = Fence::default();
122    let mut start = 0;
123    let mut pos = 0;
124    for line in text.split_inclusive('\n') {
125        let fenced = fence.feed(line);
126        if !fenced && pos > start && is_heading(line, level) {
127            out.push(&text[start..pos]);
128            start = pos;
129        }
130        pos += line.len();
131    }
132    if start < text.len() {
133        out.push(&text[start..]);
134    }
135    out
136}
137
138/// Splits `text` after every blank line outside fences.
139fn paragraphs(text: &str) -> Vec<&str> {
140    let mut out = Vec::new();
141    let mut fence = Fence::default();
142    let mut start = 0;
143    let mut pos = 0;
144    for line in text.split_inclusive('\n') {
145        let fenced = fence.feed(line);
146        pos += line.len();
147        if !fenced && line.trim().is_empty() {
148            out.push(&text[start..pos]);
149            start = pos;
150        }
151    }
152    if start < text.len() {
153        out.push(&text[start..]);
154    }
155    out
156}
157
158/// Appends `block` to `chunks`, split finer until every piece fits `max`.
159/// `level` is the next heading level to split at (7 = paragraphs, 8 = lines).
160fn pack_into(chunks: &mut Vec<String>, block: &str, max: usize, level: usize) {
161    if fits(block.len(), escaped_len(block), max) {
162        chunks.push(block.to_owned());
163        return;
164    }
165    let (pieces, next): (Vec<&str>, usize) = match level {
166        3..=6 => (blocks(block, level), level + 1),
167        7 => (paragraphs(block), 8),
168        8 => (block.split_inclusive('\n').collect(), 9),
169        _ => {
170            hard_cut(chunks, block, max);
171            return;
172        }
173    };
174    if pieces.len() <= 1 {
175        pack_into(chunks, block, max, next);
176        return;
177    }
178    // Greedily pack consecutive pieces; a piece too large alone splits finer.
179    let mut acc = String::new();
180    let mut acc_escaped = 0;
181    for piece in pieces {
182        let escaped = escaped_len(piece);
183        if !fits(piece.len(), escaped, max) {
184            if !acc.is_empty() {
185                chunks.push(std::mem::take(&mut acc));
186                acc_escaped = 0;
187            }
188            pack_into(chunks, piece, max, next);
189        } else if !fits(acc.len() + piece.len(), acc_escaped + escaped, max) {
190            chunks.push(std::mem::replace(&mut acc, piece.to_owned()));
191            acc_escaped = escaped;
192        } else {
193            acc.push_str(piece);
194            acc_escaped += escaped;
195        }
196    }
197    if !acc.is_empty() {
198        chunks.push(acc);
199    }
200}
201
202/// Room the rest of a section body (`order`, the braces) takes in its JSON.
203const BODY_ENVELOPE: usize = 64;
204
205/// Does a text of `len` bytes, `escaped` bytes as a JSON string without its
206/// quotes, fit a section of `max` bytes and a doc body of
207/// [`MAX_DOC_BODY_BYTES`]?
208fn fits(len: usize, escaped: usize, max: usize) -> bool {
209    len <= max && escaped + 2 + BODY_ENVELOPE <= MAX_DOC_BODY_BYTES
210}
211
212/// The length of `text` as a JSON string, without the quotes. Additive over
213/// concatenation, so packed pieces can be summed.
214fn escaped_len(text: &str) -> usize {
215    text.chars()
216        .map(|c| match c {
217            '"' | '\\' | '\n' | '\r' | '\t' | '\u{08}' | '\u{0c}' => 2,
218            c if (c as u32) < 0x20 => 6,
219            c => c.len_utf8(),
220        })
221        .sum()
222}
223
224/// Cuts an overlong line at character boundaries.
225fn hard_cut(chunks: &mut Vec<String>, text: &str, max: usize) {
226    let mut rest = text;
227    while !fits(rest.len(), escaped_len(rest), max) {
228        let mut cut = max.min(rest.len());
229        while !rest.is_char_boundary(cut) || !fits(cut, escaped_len(&rest[..cut]), max) {
230            cut -= 1;
231        }
232        chunks.push(rest[..cut].to_owned());
233        rest = &rest[cut..];
234    }
235    if !rest.is_empty() {
236        chunks.push(rest.to_owned());
237    }
238}
239
240#[cfg(test)]
241mod tests {
242    use super::*;
243
244    fn texts(s: &[Section]) -> Vec<&str> {
245        s.iter().map(|s| s.text.as_str()).collect()
246    }
247
248    #[test]
249    fn splits_at_level_two_headings_and_round_trips() {
250        let body = "# Title\nintro\n\n## One\na\n### Sub\nb\n## Two\nc\n";
251        let (title, sections) = split(body);
252        assert_eq!(title, "Title");
253        assert_eq!(
254            texts(&sections),
255            ["intro\n\n", "## One\na\n### Sub\nb\n", "## Two\nc\n"]
256        );
257        assert_eq!(sections[1].heading(), Some("One"));
258        assert_eq!(sections[0].heading(), None);
259        assert_eq!(join(&title, &sections), body);
260    }
261
262    #[test]
263    fn headings_inside_fences_are_text() {
264        let body = "# T\n## A\n```md\n## not a heading\n```\n## B\nx\n";
265        let (_, sections) = split(body);
266        assert_eq!(sections.len(), 2);
267        assert!(sections[0].text.contains("## not a heading"));
268    }
269
270    #[test]
271    fn oversized_blocks_split_finer_and_stay_lossless() {
272        let para = "word ".repeat(30) + "\n\n";
273        let body = format!("# T\n## Big\n{}### Sub\n{}", para.repeat(4), para.repeat(4));
274        let (title, sections) = split_with(&body, 400);
275        assert!(sections.len() > 2);
276        assert!(sections.iter().all(|s| s.text.len() <= 400));
277        assert_eq!(join(&title, &sections), body);
278        let orders: Vec<u32> = sections.iter().map(|s| s.order).collect();
279        assert_eq!(orders, renumber(sections.len()));
280    }
281
282    #[test]
283    fn a_single_overlong_line_is_cut_at_char_boundaries() {
284        let body = format!("# T\n{}\n", "ä".repeat(300));
285        let (title, sections) = split_with(&body, 101);
286        assert!(sections.iter().all(|s| s.text.len() <= 101));
287        assert_eq!(join(&title, &sections), body);
288    }
289
290    #[test]
291    fn heavy_escaping_splits_under_the_json_body_limit() {
292        let body = format!("# T\n{}\n", "\"\\".repeat(9_000));
293        let (title, sections) = split(&body);
294        assert!(sections.len() >= 2, "18 000 bytes escape to 36 000");
295        for s in &sections {
296            let json = serde_json::to_string(&s.text).unwrap_or_default();
297            assert!(json.len() + BODY_ENVELOPE <= MAX_DOC_BODY_BYTES);
298            assert_eq!(json.len(), escaped_len(&s.text) + 2);
299        }
300        assert_eq!(join(&title, &sections), body);
301    }
302
303    #[test]
304    fn a_body_without_title_keeps_everything() {
305        let (title, sections) = split("plain\n## A\nx\n");
306        assert_eq!(title, "");
307        assert_eq!(join(&title, &sections), "plain\n## A\nx\n");
308    }
309
310    #[test]
311    fn order_between_finds_gaps_or_asks_for_renumbering() {
312        assert_eq!(order_between(None, None), Some(ORDER_GAP));
313        assert_eq!(order_between(Some(1024), Some(2048)), Some(1536));
314        assert_eq!(order_between(Some(5), Some(6)), None);
315        assert_eq!(order_between(Some(2048), None), Some(3072));
316        assert_eq!(order_between(None, Some(1024)), Some(512));
317    }
318}
lean_ctx_doc/sections.rs

lean_ctx_doc/
sections.rs
