tenant.openagents/omega

1use anyhow::{Context as _, Result};
2use smallvec::SmallVec;
3use std::{collections::BTreeMap, ops::Range};
4
5#[derive(Clone, Debug, Default, PartialEq)]
6pub struct Snippet {
7    pub text: String,
8    pub tabstops: Vec<TabStop>,
9}
10
11#[derive(Clone, Debug, Default, PartialEq)]
12pub struct TabStop {
13    pub ranges: SmallVec<[Range<isize>; 2]>,
14    pub choices: Option<Vec<String>>,
15}
16
17impl Snippet {
18    pub fn parse(source: &str) -> Result<Self> {
19        let mut text = String::with_capacity(source.len());
20        let mut tabstops = BTreeMap::new();
21        parse_snippet(source, false, &mut text, &mut tabstops)
22            .context("failed to parse snippet")?;
23
24        let len = text.len() as isize;
25        let final_tabstop = tabstops.remove(&0);
26        let mut tabstops = tabstops.into_values().collect::<Vec<_>>();
27
28        if let Some(final_tabstop) = final_tabstop {
29            tabstops.push(final_tabstop);
30        } else {
31            let end_tabstop = TabStop {
32                ranges: [len..len].into_iter().collect(),
33                choices: None,
34            };
35
36            if !tabstops.last().is_some_and(|t| *t == end_tabstop) {
37                tabstops.push(end_tabstop);
38            }
39        }
40
41        Ok(Snippet { text, tabstops })
42    }
43}
44
45fn parse_snippet<'a>(
46    mut source: &'a str,
47    nested: bool,
48    text: &mut String,
49    tabstops: &mut BTreeMap<usize, TabStop>,
50) -> Result<&'a str> {
51    loop {
52        match source.chars().next() {
53            None => return Ok(""),
54            Some('$') => {
55                source = parse_tabstop(&source[1..], text, tabstops)?;
56            }
57            Some('\\') => {
58                // As specified in the LSP spec (`Grammar` section),
59                // backslashes can escape some characters:
60                // https://microsoft.github.io/language-server-protocol/specifications/lsp/3.17/specification/#snippet_syntax
61                source = &source[1..];
62                if let Some(c) = source.chars().next() {
63                    if c == '$' || c == '\\' || c == '}' {
64                        text.push(c);
65                        // All escapable characters are 1 byte long:
66                        source = &source[1..];
67                    } else {
68                        text.push('\\');
69                    }
70                } else {
71                    text.push('\\');
72                }
73            }
74            Some('}') => {
75                if nested {
76                    return Ok(source);
77                } else {
78                    text.push('}');
79                    source = &source[1..];
80                }
81            }
82            Some(_) => {
83                let chunk_end = source.find(['}', '$', '\\']).unwrap_or(source.len());
84                let (chunk, rest) = source.split_at(chunk_end);
85                text.push_str(chunk);
86                source = rest;
87            }
88        }
89    }
90}
91
92fn parse_tabstop<'a>(
93    mut source: &'a str,
94    text: &mut String,
95    tabstops: &mut BTreeMap<usize, TabStop>,
96) -> Result<&'a str> {
97    let tabstop_start = text.len();
98    let tabstop_index;
99    let mut choices = None;
100
101    if source.starts_with('{') {
102        let (index, rest) = parse_int(&source[1..])?;
103        tabstop_index = index;
104        source = rest;
105
106        if source.starts_with("|") {
107            (source, choices) = parse_choices(&source[1..], text)?;
108        }
109
110        if source.starts_with(':') {
111            source = parse_snippet(&source[1..], true, text, tabstops)?;
112        }
113
114        if source.starts_with('}') {
115            source = &source[1..];
116        } else {
117            anyhow::bail!("expected a closing brace");
118        }
119    } else {
120        let (index, rest) = parse_int(source)?;
121        tabstop_index = index;
122        source = rest;
123    }
124
125    tabstops
126        .entry(tabstop_index)
127        .or_insert_with(|| TabStop {
128            ranges: Default::default(),
129            choices,
130        })
131        .ranges
132        .push(tabstop_start as isize..text.len() as isize);
133    Ok(source)
134}
135
136fn parse_int(source: &str) -> Result<(usize, &str)> {
137    let len = source
138        .find(|c: char| !c.is_ascii_digit())
139        .unwrap_or(source.len());
140    anyhow::ensure!(len > 0, "expected an integer");
141    let (prefix, suffix) = source.split_at(len);
142    Ok((prefix.parse()?, suffix))
143}
144
145fn parse_choices<'a>(
146    mut source: &'a str,
147    text: &mut String,
148) -> Result<(&'a str, Option<Vec<String>>)> {
149    let mut found_default_choice = false;
150    let mut current_choice = String::new();
151    let mut choices = Vec::new();
152
153    loop {
154        match source.chars().next() {
155            None => return Ok(("", Some(choices))),
156            Some('\\') => {
157                source = &source[1..];
158
159                if let Some(c) = source.chars().next() {
160                    current_choice.push(c);
161                    if !found_default_choice {
162                        text.push(c);
163                    }
164                    source = &source[c.len_utf8()..];
165                }
166            }
167            Some(',') => {
168                found_default_choice = true;
169                source = &source[1..];
170                choices.push(current_choice);
171                current_choice = String::new();
172            }
173            Some('|') => {
174                source = &source[1..];
175                choices.push(current_choice);
176                return Ok((source, Some(choices)));
177            }
178            Some(_) => {
179                let chunk_end = source.find([',', '|', '\\']);
180
181                anyhow::ensure!(
182                    chunk_end.is_some(),
183                    "Placeholder choice doesn't contain closing pipe-character '|'"
184                );
185
186                let (chunk, rest) = source.split_at(chunk_end.unwrap());
187
188                if !found_default_choice {
189                    text.push_str(chunk);
190                }
191
192                current_choice.push_str(chunk);
193                source = rest;
194            }
195        }
196    }
197}
198
199#[cfg(test)]
200mod tests {
201    use super::*;
202
203    #[test]
204    fn test_snippet_without_tabstops() {
205        let snippet = Snippet::parse("one-two-three").unwrap();
206        assert_eq!(snippet.text, "one-two-three");
207        assert_eq!(tabstops(&snippet), &[vec![13..13]]);
208    }
209
210    #[test]
211    fn test_snippet_with_tabstops() {
212        let snippet = Snippet::parse("one$1two").unwrap();
213        assert_eq!(snippet.text, "onetwo");
214        assert_eq!(tabstops(&snippet), &[vec![3..3], vec![6..6]]);
215        assert_eq!(tabstop_choices(&snippet), &[&None, &None]);
216
217        // Multi-digit numbers
218        let snippet = Snippet::parse("one$123-$99-two").unwrap();
219        assert_eq!(snippet.text, "one--two");
220        assert_eq!(tabstops(&snippet), &[vec![4..4], vec![3..3], vec![8..8]]);
221        assert_eq!(tabstop_choices(&snippet), &[&None, &None, &None]);
222    }
223
224    #[test]
225    fn test_snippet_with_last_tabstop_at_end() {
226        let snippet = Snippet::parse(r#"foo.$1"#).unwrap();
227
228        // If the final tabstop is already at the end of the text, don't insert
229        // an additional tabstop at the end.
230        assert_eq!(snippet.text, r#"foo."#);
231        assert_eq!(tabstops(&snippet), &[vec![4..4]]);
232        assert_eq!(tabstop_choices(&snippet), &[&None]);
233    }
234
235    #[test]
236    fn test_snippet_with_explicit_final_tabstop() {
237        let snippet = Snippet::parse(r#"<div class="$1">$0</div>"#).unwrap();
238
239        // If the final tabstop is explicitly specified via '$0', then
240        // don't insert an additional tabstop at the end.
241        assert_eq!(snippet.text, r#"<div class=""></div>"#);
242        assert_eq!(tabstops(&snippet), &[vec![12..12], vec![14..14]]);
243        assert_eq!(tabstop_choices(&snippet), &[&None, &None]);
244    }
245
246    #[test]
247    fn test_snippet_with_placeholders() {
248        let snippet = Snippet::parse("one${1:two}three${2:four}").unwrap();
249        assert_eq!(snippet.text, "onetwothreefour");
250        assert_eq!(
251            tabstops(&snippet),
252            &[vec![3..6], vec![11..15], vec![15..15]]
253        );
254        assert_eq!(tabstop_choices(&snippet), &[&None, &None, &None]);
255    }
256
257    #[test]
258    fn test_snippet_with_choice_placeholders() {
259        let snippet = Snippet::parse("type ${1|i32, u32|} = $2")
260            .expect("Should be able to unpack choice placeholders");
261
262        assert_eq!(snippet.text, "type i32 = ");
263        assert_eq!(tabstops(&snippet), &[vec![5..8], vec![11..11],]);
264        assert_eq!(
265            tabstop_choices(&snippet),
266            &[&Some(vec!["i32".to_string(), " u32".to_string()]), &None]
267        );
268
269        let snippet = Snippet::parse(r"${1|\$\{1\|one\,two\,tree\|\}|}")
270            .expect("Should be able to parse choice with escape characters");
271
272        assert_eq!(snippet.text, "${1|one,two,tree|}");
273        assert_eq!(tabstops(&snippet), &[vec![0..18], vec![18..18]]);
274        assert_eq!(
275            tabstop_choices(&snippet),
276            &[&Some(vec!["${1|one,two,tree|}".to_string(),]), &None]
277        );
278    }
279
280    #[test]
281    fn test_snippet_with_escaped_chars_in_non_default_choices() {
282        let snippet = Snippet::parse(r"${1|a,b\,c|}").unwrap();
283        assert_eq!(snippet.text, "a");
284        assert_eq!(tabstops(&snippet), &[vec![0..1], vec![1..1]]);
285        assert_eq!(
286            tabstop_choices(&snippet),
287            &[&Some(vec!["a".to_string(), "b,c".to_string()]), &None]
288        );
289
290        let snippet = Snippet::parse(r"${1|one,two\|2,three\\3|}").unwrap();
291        assert_eq!(snippet.text, "one");
292        assert_eq!(tabstops(&snippet), &[vec![0..3], vec![3..3]]);
293        assert_eq!(
294            tabstop_choices(&snippet),
295            &[
296                &Some(vec![
297                    "one".to_string(),
298                    "two|2".to_string(),
299                    r"three\3".to_string()
300                ]),
301                &None
302            ]
303        );
304    }
305
306    #[test]
307    fn test_snippet_with_nested_placeholders() {
308        let snippet = Snippet::parse(
309            "for (${1:var ${2:i} = 0; ${2:i} < ${3:${4:array}.length}; ${2:i}++}) {$0}",
310        )
311        .unwrap();
312        assert_eq!(snippet.text, "for (var i = 0; i < array.length; i++) {}");
313        assert_eq!(
314            tabstops(&snippet),
315            &[
316                vec![5..37],
317                vec![9..10, 16..17, 34..35],
318                vec![20..32],
319                vec![20..25],
320                vec![40..40],
321            ]
322        );
323        assert_eq!(
324            tabstop_choices(&snippet),
325            &[&None, &None, &None, &None, &None]
326        );
327    }
328
329    #[test]
330    fn test_snippet_parsing_with_escaped_chars() {
331        let snippet = Snippet::parse("\"\\$schema\": $1").unwrap();
332        assert_eq!(snippet.text, "\"$schema\": ");
333        assert_eq!(tabstops(&snippet), &[vec![11..11]]);
334        assert_eq!(tabstop_choices(&snippet), &[&None]);
335
336        let snippet = Snippet::parse("{a\\}").unwrap();
337        assert_eq!(snippet.text, "{a}");
338        assert_eq!(tabstops(&snippet), &[vec![3..3]]);
339        assert_eq!(tabstop_choices(&snippet), &[&None]);
340
341        // backslash not functioning as an escape
342        let snippet = Snippet::parse("a\\b").unwrap();
343        assert_eq!(snippet.text, "a\\b");
344        assert_eq!(tabstops(&snippet), &[vec![3..3]]);
345
346        // first backslash cancelling escaping that would
347        // have happened with second backslash
348        let snippet = Snippet::parse("one\\\\$1two").unwrap();
349        assert_eq!(snippet.text, "one\\two");
350        assert_eq!(tabstops(&snippet), &[vec![4..4], vec![7..7]]);
351    }
352
353    fn tabstops(snippet: &Snippet) -> Vec<Vec<Range<isize>>> {
354        snippet.tabstops.iter().map(|t| t.ranges.to_vec()).collect()
355    }
356
357    fn tabstop_choices(snippet: &Snippet) -> Vec<&Option<Vec<String>>> {
358        snippet.tabstops.iter().map(|t| &t.choices).collect()
359    }
360}
361