tenant.openagents/omega

1use crate::parser::{MarkdownEvent, MarkdownTag, MarkdownTagEnd};
2use std::ops::Range;
3
4struct InlineSpan {
5    range: Range<usize>,
6    content: Range<usize>,
7    is_code: bool,
8}
9
10impl InlineSpan {
11    fn opening<'a>(&self, source: &'a str) -> &'a str {
12        source
13            .get(self.range.start..self.content.start)
14            .unwrap_or("")
15    }
16
17    fn closing<'a>(&self, source: &'a str) -> &'a str {
18        source.get(self.content.end..self.range.end).unwrap_or("")
19    }
20}
21
22fn is_inline_span_tag(tag: &MarkdownTag) -> bool {
23    matches!(
24        tag,
25        MarkdownTag::Emphasis
26            | MarkdownTag::Strong
27            | MarkdownTag::Strikethrough
28            | MarkdownTag::Superscript
29            | MarkdownTag::Subscript
30            | MarkdownTag::Link { .. }
31    )
32}
33
34fn is_inline_span_tag_end(tag: &MarkdownTagEnd) -> bool {
35    matches!(
36        tag,
37        MarkdownTagEnd::Emphasis
38            | MarkdownTagEnd::Strong
39            | MarkdownTagEnd::Strikethrough
40            | MarkdownTagEnd::Superscript
41            | MarkdownTagEnd::Subscript
42            | MarkdownTagEnd::Link
43    )
44}
45
46fn inline_code_full_range(source: &str, content: &Range<usize>) -> Range<usize> {
47    let opening_ticks = source[..content.start]
48        .bytes()
49        .rev()
50        .take_while(|&byte| byte == b'`')
51        .count();
52    let closing_ticks = source[content.end..]
53        .bytes()
54        .take_while(|&byte| byte == b'`')
55        .count();
56    let ticks = opening_ticks.min(closing_ticks);
57    content.start - ticks..content.end + ticks
58}
59
60fn collect_inline_spans(source: &str, events: &[(Range<usize>, MarkdownEvent)]) -> Vec<InlineSpan> {
61    fn note_child(stack: &mut [(Range<usize>, Option<Range<usize>>)], child: &Range<usize>) {
62        for (_, content) in stack.iter_mut() {
63            match content {
64                Some(content) => content.end = content.end.max(child.end),
65                None => *content = Some(child.clone()),
66            }
67        }
68    }
69
70    let mut spans = Vec::new();
71    let mut stack: Vec<(Range<usize>, Option<Range<usize>>)> = Vec::new();
72    for (event_range, event) in events {
73        match event {
74            MarkdownEvent::Start(tag) if is_inline_span_tag(tag) => {
75                note_child(&mut stack, event_range);
76                stack.push((event_range.clone(), None));
77            }
78            MarkdownEvent::End(tag) if is_inline_span_tag_end(tag) => {
79                if let Some((range, content)) = stack.pop() {
80                    let content = content.unwrap_or(range.clone());
81                    spans.push(InlineSpan {
82                        range,
83                        content,
84                        is_code: false,
85                    });
86                }
87                note_child(&mut stack, event_range);
88            }
89            MarkdownEvent::Code | MarkdownEvent::SubstitutedCode(_) => {
90                let range = inline_code_full_range(source, event_range);
91                note_child(&mut stack, &range);
92                spans.push(InlineSpan {
93                    range,
94                    content: event_range.clone(),
95                    is_code: true,
96                });
97            }
98            _ => note_child(&mut stack, event_range),
99        }
100    }
101    spans
102}
103
104pub(crate) fn rebalanced_markdown_for_selection(
105    source: &str,
106    events: &[(Range<usize>, MarkdownEvent)],
107    root_block_starts: &[usize],
108    selection: Range<usize>,
109) -> String {
110    let Some(selection) = snap_to_char_boundaries(source, selection) else {
111        return String::new();
112    };
113
114    let (start_events, end_events) = boundary_block_events(events, root_block_starts, &selection);
115    let mut spans = collect_inline_spans(source, start_events);
116    spans.extend(collect_inline_spans(source, end_events));
117
118    let Some(selection) = snap_out_of_delimiters(&spans, selection) else {
119        return String::new();
120    };
121
122    if selection_is_only_inside_code_spans(&spans, &selection) {
123        return source[selection].to_string();
124    }
125
126    rebalance_delimiters(source, &spans, &selection)
127}
128
129/// Returns the events of the root blocks containing each selection boundary.
130/// The second slice is empty when both boundaries share a block.
131fn boundary_block_events<'a>(
132    events: &'a [(Range<usize>, MarkdownEvent)],
133    root_block_starts: &[usize],
134    selection: &Range<usize>,
135) -> (
136    &'a [(Range<usize>, MarkdownEvent)],
137    &'a [(Range<usize>, MarkdownEvent)],
138) {
139    if root_block_starts.is_empty() {
140        return (events, &[]);
141    }
142    let start_block = root_block_index(root_block_starts, selection.start);
143    let end_block = root_block_index(root_block_starts, selection.end);
144    let start_events = root_block_events(events, root_block_starts, start_block);
145    if end_block == start_block {
146        (start_events, &[])
147    } else {
148        (
149            start_events,
150            root_block_events(events, root_block_starts, end_block),
151        )
152    }
153}
154
155fn root_block_index(root_block_starts: &[usize], offset: usize) -> usize {
156    root_block_starts
157        .partition_point(|block_start| *block_start <= offset)
158        .saturating_sub(1)
159}
160
161fn root_block_events<'a>(
162    events: &'a [(Range<usize>, MarkdownEvent)],
163    root_block_starts: &[usize],
164    block: usize,
165) -> &'a [(Range<usize>, MarkdownEvent)] {
166    let Some(&block_start) = root_block_starts.get(block) else {
167        return events;
168    };
169    let start = events.partition_point(|(range, _)| range.start < block_start);
170    let end = match root_block_starts.get(block + 1) {
171        Some(&next_block_start) => {
172            events.partition_point(|(range, _)| range.start < next_block_start)
173        }
174        None => events.len(),
175    };
176    events.get(start..end).unwrap_or(events)
177}
178
179fn snap_to_char_boundaries(source: &str, selection: Range<usize>) -> Option<Range<usize>> {
180    let mut start = selection.start.min(source.len());
181    let mut end = selection.end.min(source.len());
182    if start >= end {
183        return None;
184    }
185    while start > 0 && !source.is_char_boundary(start) {
186        start -= 1;
187    }
188    while end < source.len() && !source.is_char_boundary(end) {
189        end += 1;
190    }
191    Some(start..end)
192}
193
194/// Shrinks selection boundaries that fall inside delimiter syntax (`**`,
195/// etc.) so no delimiter is left half-selected:
196///
197/// - an end in `**bold*|*` snaps back to `**bold|**`
198/// - a start in `*|*bold**` snaps forward to `**|bold**`
199///
200/// This repeats until stable, since snapping can land inside a nested span's
201/// delimiter. Returns `None` if the selection becomes empty.
202fn snap_out_of_delimiters(spans: &[InlineSpan], selection: Range<usize>) -> Option<Range<usize>> {
203    let mut start = selection.start;
204    let mut end = selection.end;
205    loop {
206        let mut changed = false;
207        for span in spans {
208            if end > span.range.start && end <= span.content.start {
209                end = span.range.start;
210                changed = true;
211            } else if end > span.content.end && end <= span.range.end {
212                end = span.content.end;
213                changed = true;
214            }
215            if start >= span.range.start && start < span.content.start {
216                start = span.content.start;
217                changed = true;
218            } else if start >= span.content.end && start < span.range.end {
219                start = span.range.end;
220                changed = true;
221            }
222        }
223        if !changed {
224            break;
225        }
226    }
227    (start < end).then(|| start..end)
228}
229
230fn selection_is_only_inside_code_spans(spans: &[InlineSpan], selection: &Range<usize>) -> bool {
231    let contains = |span: &InlineSpan| {
232        span.content.start <= selection.start && selection.end <= span.content.end
233    };
234    spans.iter().any(|span| span.is_code && contains(span))
235        && !spans.iter().any(|span| !span.is_code && contains(span))
236}
237
238/// Re-adds delimiters cut off by the selection so the result is well-formed
239/// markdown:
240///
241/// - selecting `old te` in `**bold text**` yields `**old te**`
242/// - nested spans are reopened outermost first: selecting `alic` in
243///   `**bold _italic_**` yields `**_alic_**`
244fn rebalance_delimiters(source: &str, spans: &[InlineSpan], selection: &Range<usize>) -> String {
245    let nesting_order = |a: &&InlineSpan, b: &&InlineSpan| {
246        a.range
247            .start
248            .cmp(&b.range.start)
249            .then(b.range.end.cmp(&a.range.end))
250    };
251
252    let mut open_at_start = spans
253        .iter()
254        .filter(|span| span.content.start <= selection.start && selection.start < span.content.end)
255        .collect::<Vec<_>>();
256    open_at_start.sort_by(nesting_order);
257
258    let mut open_at_end = spans
259        .iter()
260        .filter(|span| span.content.start < selection.end && selection.end <= span.content.end)
261        .collect::<Vec<_>>();
262    open_at_end.sort_by(nesting_order);
263
264    let mut result = String::new();
265    for span in &open_at_start {
266        result.push_str(span.opening(source));
267    }
268    result.push_str(&source[selection.clone()]);
269    for span in open_at_end.iter().rev() {
270        result.push_str(span.closing(source));
271    }
272
273    result
274}
275
276#[cfg(test)]
277mod tests {
278    use super::*;
279    use crate::parser::parse_markdown_with_options;
280    use util::test::marked_text_ranges;
281
282    fn markdown_for(source: &str, selection: Range<usize>) -> String {
283        let parsed = parse_markdown_with_options(source, false, false, false);
284        rebalanced_markdown_for_selection(
285            source,
286            &parsed.events,
287            &parsed.root_block_starts,
288            selection,
289        )
290    }
291
292    fn markdown_for_marked(marked_source: &str) -> String {
293        let (source, selection) = marked_range(marked_source);
294        markdown_for(&source, selection)
295    }
296
297    #[track_caller]
298    fn assert_marked_selection(marked_source: &str, expected: &str) {
299        assert_eq!(
300            markdown_for_marked(marked_source),
301            expected,
302            "source: {marked_source:?}"
303        );
304    }
305
306    #[track_caller]
307    fn marked_range(marked_source: &str) -> (String, Range<usize>) {
308        let (source, ranges) = marked_text_ranges(marked_source, false);
309        match ranges.as_slice() {
310            [selection] => (source, selection.clone()),
311            _ => panic!("expected exactly one «» range in marked source"),
312        }
313    }
314
315    fn inline_spans(source: &str) -> Vec<InlineSpan> {
316        let parsed = parse_markdown_with_options(source, false, false, false);
317        collect_inline_spans(source, &parsed.events)
318    }
319
320    #[track_caller]
321    fn assert_snapped(marked_selection: &str, marked_expected: Option<&str>, message: &str) {
322        let (source, selection) = marked_range(marked_selection);
323        let expected = marked_expected.map(|marked_expected| {
324            let (expected_source, range) = marked_range(marked_expected);
325            assert_eq!(
326                expected_source, source,
327                "expected must be marked on the same source"
328            );
329            range
330        });
331        assert_eq!(
332            snap_out_of_delimiters(&inline_spans(&source), selection),
333            expected,
334            "{message}"
335        );
336    }
337
338    #[test]
339    fn test_snap_out_of_delimiters() {
340        assert_snapped(
341            "**«bold»** rest",
342            Some("**«bold»** rest"),
343            "boundaries already in content must be untouched",
344        );
345        assert_snapped(
346            "*«*bold»** rest",
347            Some("**«bold»** rest"),
348            "a start in the opening delimiter must advance into the content",
349        );
350        assert_snapped(
351            "**«bold*»* rest",
352            Some("**«bold»** rest"),
353            "an end in the closing delimiter must retreat to the content end",
354        );
355        assert_snapped(
356            "**bold*«* rest»",
357            Some("**bold**« rest»"),
358            "a start in the closing delimiter must advance past the whole span",
359        );
360        assert_snapped(
361            "«*»*bold** rest",
362            None,
363            "an end in the opening delimiter must retreat to before the span",
364        );
365        assert_snapped(
366            "**bold«**» rest",
367            None,
368            "a selection covering only delimiter text must collapse",
369        );
370    }
371
372    #[test]
373    fn test_snap_out_of_delimiters_cascades_through_nested_spans() {
374        assert_snapped(
375            "**`«code`*»*",
376            Some("**`«code»`**"),
377            "an end inside bold's closing `**` first snaps to code's closing \
378             backtick, so it must cascade into the code content",
379        );
380    }
381
382    #[test]
383    fn test_snap_to_char_boundaries() {
384        // `√` occupies bytes 1..4.
385        let source = "a√b";
386        assert_eq!(
387            snap_to_char_boundaries(source, 0..5),
388            Some(0..5),
389            "boundaries already on char boundaries must be untouched"
390        );
391        assert_eq!(
392            snap_to_char_boundaries(source, 0..2),
393            Some(0..4),
394            "an end mid-character must expand forward to keep the character"
395        );
396        assert_eq!(
397            snap_to_char_boundaries(source, 2..5),
398            Some(1..5),
399            "a start mid-character must expand backward to keep the character"
400        );
401        assert_eq!(
402            snap_to_char_boundaries(source, 2..3),
403            Some(1..4),
404            "a selection entirely inside a character must cover the whole character"
405        );
406        assert_eq!(
407            snap_to_char_boundaries(source, 2..10),
408            Some(1..5),
409            "an end past the source must clamp to its length"
410        );
411        assert_eq!(
412            snap_to_char_boundaries(source, 2..2),
413            None,
414            "an empty selection must collapse, even mid-character"
415        );
416        assert_eq!(
417            snap_to_char_boundaries(source, 5..10),
418            None,
419            "a selection entirely past the source must collapse"
420        );
421        assert_eq!(
422            snap_to_char_boundaries(source, Range { start: 4, end: 2 }),
423            None,
424            "a reversed selection must collapse"
425        );
426    }
427
428    fn selection_is_plain(marked_source: &str) -> bool {
429        let (source, selection) = marked_range(marked_source);
430        selection_is_only_inside_code_spans(&inline_spans(&source), &selection)
431    }
432
433    #[test]
434    fn test_selection_is_only_inside_code_spans() {
435        assert!(
436            selection_is_plain("run `«cargo» test` now"),
437            "a selection fully inside the code span's content must be plain"
438        );
439        assert!(
440            !selection_is_plain("«run `cargo» test` now"),
441            "a selection reaching outside the code span must not be plain"
442        );
443        assert!(
444            !selection_is_plain("**`«code»`**"),
445            "code nested in bold must not be plain: the bold span also contains it"
446        );
447    }
448
449    #[test]
450    fn test_markdown_for_selection_balances_inline_spans() {
451        assert_marked_selection("This is **«bold»** text in a sentence.", "**bold**");
452        assert_marked_selection("This is **«bold**» text in a sentence.", "**bold**");
453        assert_marked_selection("Th«is is **bo»ld** text in a sentence.", "is is **bo**");
454
455        assert_marked_selection("This is *«italic»* text in a sentence.", "*italic*");
456        assert_marked_selection("This is *it«al»ic* text in a sentence.", "*al*");
457
458        assert_marked_selection("T«his is `cod»e` all `in one` sentence.", "his is `cod`");
459        assert_marked_selection(
460            "This is `c«ode` all `in o»ne` sentence.",
461            "`ode` all `in o`",
462        );
463        assert_marked_selection(
464            "This is `«code` all `in one»` sentence.",
465            "`code` all `in one`",
466        );
467        assert_marked_selection(
468            "This is `«code` all `in one`» sentence.",
469            "`code` all `in one`",
470        );
471
472        // Special case for single inline code blocks
473        assert_marked_selection("This is `«code»` all `in one` sentence.", "code");
474        assert_marked_selection("This is `«code`» all `in one` sentence.", "code");
475        assert_marked_selection("This is `c«od»e` all `in one` sentence.", "od");
476    }
477
478    #[test]
479    fn test_markdown_for_selection_nested_spans() {
480        assert_marked_selection("**bo«ld wi»th `code` inside**", "**ld wi**");
481        assert_marked_selection("**bold with `c«od»e` inside**", "**`od`**");
482        assert_marked_selection("**bold with `«code»` inside**", "**`code`**");
483        assert_marked_selection(
484            "**«bold with `code` inside»**",
485            "**bold with `code` inside**",
486        );
487        assert_marked_selection(
488            "«**bold with `code` inside**»",
489            "**bold with `code` inside**",
490        );
491    }
492
493    #[test]
494    fn test_markdown_for_selection_links() {
495        assert_marked_selection(
496            "[Visit Rust's we«bsite»](https://rust.org)",
497            "[bsite](https://rust.org)",
498        );
499        assert_marked_selection(
500            "[«Visit Rust's website»](https://rust.org)",
501            "[Visit Rust's website](https://rust.org)",
502        );
503        assert_marked_selection("visit https://«example».com now", "example");
504    }
505
506    #[test]
507    fn test_markdown_for_selection_scopes_to_boundary_blocks() {
508        assert_marked_selection(
509            "**bo«ld one**\n\nmiddle `x`\n\n*ita»lic two*",
510            "**ld one**\n\nmiddle `x`\n\n*ita*",
511        );
512        assert_marked_selection("**bold** first\n\nsecond *ita«li»c* here", "*li*");
513        assert_marked_selection("one\n«\ntwo»", "\ntwo");
514        assert_marked_selection("**one*«*\n\ntw»o", "\n\ntw");
515    }
516
517    #[test]
518    fn test_root_block_index() {
519        let starts = [2, 10, 20];
520        assert_eq!(
521            root_block_index(&starts, 0),
522            0,
523            "an offset before the first block start must clamp to the first block"
524        );
525        assert_eq!(root_block_index(&starts, 2), 0);
526        assert_eq!(root_block_index(&starts, 9), 0);
527        assert_eq!(root_block_index(&starts, 10), 1);
528        assert_eq!(root_block_index(&starts, 19), 1);
529        assert_eq!(root_block_index(&starts, 20), 2);
530        assert_eq!(
531            root_block_index(&starts, 100),
532            2,
533            "an offset past the last block start must clamp to the last block"
534        );
535    }
536
537    #[test]
538    fn test_root_block_events_slices_each_block() {
539        let source = "first **a**\n\nsecond `b`\n\nthird *c*";
540        let parsed = parse_markdown_with_options(source, false, false, false);
541        let events = parsed.events.as_slice();
542        let starts = parsed.root_block_starts.as_slice();
543        assert_eq!(starts, &[0, 13, 25]);
544
545        let mut sliced_events = Vec::new();
546        for block in 0..starts.len() {
547            let slice = root_block_events(events, starts, block);
548            assert!(!slice.is_empty(), "block {block} must have events");
549            let block_end = starts.get(block + 1).copied().unwrap_or(source.len());
550            for (range, event) in slice {
551                assert!(
552                    (starts[block]..block_end).contains(&range.start),
553                    "event {event:?} at {range:?} must start within block {block}"
554                );
555            }
556            sliced_events.extend_from_slice(slice);
557        }
558        assert_eq!(
559            sliced_events, events,
560            "the per-block slices must partition all events in order"
561        );
562
563        assert_eq!(
564            root_block_events(events, starts, starts.len()),
565            events,
566            "an out-of-range block index must fall back to all events"
567        );
568    }
569
570    #[test]
571    fn test_markdown_for_selection_plain_text_and_blocks() {
572        assert_eq!(
573            markdown_for_marked("some «text»"),
574            "text",
575            "plain text must be unchanged"
576        );
577        assert_eq!(
578            markdown_for_marked("«para one\n\n- item **bold**\n- item two»"),
579            "para one\n\n- item **bold**\n- item two",
580            "selections spanning multiple blocks must keep interior syntax as-is"
581        );
582        assert_eq!(
583            markdown_for_marked("```rust\n«let x = 1;»\n```"),
584            "let x = 1;",
585            "a selection inside a fenced code block must stay plain"
586        );
587        assert_eq!(
588            markdown_for("abc", 2..100),
589            "c",
590            "out-of-bounds ends must be clamped"
591        );
592        assert_eq!(
593            markdown_for("abc", Range { start: 3, end: 2 }),
594            "",
595            "inverted ranges must yield an empty string"
596        );
597    }
598}
599