squawk_ide/
expand_selection.rs

1// via https://github.com/rust-lang/rust-analyzer/blob/8d75311400a108d7ffe17dc9c38182c566952e6e/crates/ide/src/extend_selection.rs#L1C1-L1C1
2//
3// Permission is hereby granted, free of charge, to any
4// person obtaining a copy of this software and associated
5// documentation files (the "Software"), to deal in the
6// Software without restriction, including without
7// limitation the rights to use, copy, modify, merge,
8// publish, distribute, sublicense, and/or sell copies of
9// the Software, and to permit persons to whom the Software
10// is furnished to do so, subject to the following
11// conditions:
12//
13// The above copyright notice and this permission notice
14// shall be included in all copies or substantial portions
15// of the Software.
16//
17// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF
18// ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED
19// TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A
20// PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT
21// SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY
22// CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
23// OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR
24// IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
25// DEALINGS IN THE SOFTWARE.
26
27// NOTE: this is pretty much copied as is from rust analyzer with some
28// simplifications. I imagine there's more we can do to adapt it for SQL.
29
30use rowan::{Direction, NodeOrToken, TextRange, TextSize};
31use squawk_syntax::{
32    SyntaxKind, SyntaxNode, SyntaxToken,
33    ast::{self, AstToken},
34};
35
36use crate::tokens::is_string_or_comment;
37
38const DELIMITED_LIST_KINDS: &[SyntaxKind] = &[
39    SyntaxKind::ALTER_OPTION_LIST,
40    SyntaxKind::ARG_LIST,
41    SyntaxKind::ATTRIBUTE_LIST,
42    SyntaxKind::BEGIN_FUNC_OPTION_LIST,
43    SyntaxKind::COLUMN_LIST,
44    SyntaxKind::CONFLICT_INDEX_ITEM_LIST,
45    SyntaxKind::CONSTRAINT_EXCLUSION_LIST,
46    SyntaxKind::COPY_OPTION_LIST,
47    SyntaxKind::DROP_OP_CLASS_OPTION_LIST,
48    SyntaxKind::FDW_OPTION_LIST,
49    SyntaxKind::FUNCTION_SIG_LIST,
50    SyntaxKind::GROUP_BY_LIST,
51    SyntaxKind::JSON_TABLE_COLUMN_LIST,
52    SyntaxKind::OPERATOR_CLASS_OPTION_LIST,
53    SyntaxKind::OPTION_ITEM_LIST,
54    SyntaxKind::OP_SIG_LIST,
55    SyntaxKind::PARAM_LIST,
56    SyntaxKind::PARTITION_ITEM_LIST,
57    SyntaxKind::PARTITION_LIST,
58    SyntaxKind::RETURNING_OPTION_LIST,
59    SyntaxKind::REVOKE_COMMAND_LIST,
60    SyntaxKind::ROLE_REF_LIST,
61    SyntaxKind::ROW_LIST,
62    SyntaxKind::XML_ATTRIBUTE_LIST,
63    SyntaxKind::XML_NAMESPACE_LIST,
64    SyntaxKind::SET_COLUMN_LIST,
65    SyntaxKind::SET_EXPR_LIST,
66    SyntaxKind::SET_OPTIONS_LIST,
67    SyntaxKind::SORT_BY_LIST,
68    SyntaxKind::TABLE_AND_COLUMNS_LIST,
69    SyntaxKind::TABLE_ARG_LIST,
70    SyntaxKind::TABLE_LIST,
71    SyntaxKind::TARGET_LIST,
72    SyntaxKind::TRANSACTION_MODE_LIST,
73    SyntaxKind::VACUUM_OPTION_LIST,
74    SyntaxKind::VARIANT_LIST,
75    SyntaxKind::XML_TABLE_COLUMN_LIST,
76];
77
78pub fn extend_selection(root: &SyntaxNode, range: TextRange) -> TextRange {
79    try_extend_selection(root, range).unwrap_or(range)
80}
81
82fn try_extend_selection(root: &SyntaxNode, range: TextRange) -> Option<TextRange> {
83    if range.is_empty() {
84        let offset = range.start();
85        let mut leaves = root.token_at_offset(offset);
86        // Make sure that if we're on the whitespace at the start of a line, we
87        // expand to the node on that line instead of the previous one
88        if leaves.clone().all(|it| it.kind() == SyntaxKind::WHITESPACE) {
89            return Some(extend_ws(root, leaves.next()?, offset));
90        }
91        let leaf_range = match root.token_at_offset(offset) {
92            rowan::TokenAtOffset::None => return None,
93            rowan::TokenAtOffset::Single(l) => {
94                if is_string_or_comment(l.kind()) {
95                    extend_single_word_in_comment_or_string(&l, offset)
96                        .unwrap_or_else(|| l.text_range())
97                } else {
98                    l.text_range()
99                }
100            }
101            rowan::TokenAtOffset::Between(l, r) => pick_best(l, r).text_range(),
102        };
103        return Some(leaf_range);
104    }
105
106    let node = match root.covering_element(range) {
107        NodeOrToken::Token(token) => {
108            if token.text_range() != range {
109                return Some(token.text_range());
110            }
111            if let Some(comment) = ast::Comment::cast(token.clone())
112                && let Some(range) = extend_comments(comment)
113            {
114                return Some(range);
115            }
116            token.parent()?
117        }
118        NodeOrToken::Node(node) => node,
119    };
120
121    if node.text_range() != range {
122        return Some(node.text_range());
123    }
124
125    let node = shallowest_node(&node);
126
127    if node
128        .parent()
129        .is_some_and(|n| DELIMITED_LIST_KINDS.contains(&n.kind()))
130    {
131        if let Some(range) = extend_list_item(&node) {
132            return Some(range);
133        }
134    }
135
136    node.parent().map(|it| it.text_range())
137}
138
139/// Find the shallowest node with same range, which allows us to traverse siblings.
140fn shallowest_node(node: &SyntaxNode) -> SyntaxNode {
141    node.ancestors()
142        .take_while(|n| n.text_range() == node.text_range())
143        .last()
144        .unwrap()
145}
146
147/// Expand to the current word instead the full text range of the node.
148fn extend_single_word_in_comment_or_string(
149    leaf: &SyntaxToken,
150    offset: TextSize,
151) -> Option<TextRange> {
152    let text: &str = leaf.text();
153    let cursor_position: u32 = (offset - leaf.text_range().start()).into();
154
155    let (before, after) = text.split_at(cursor_position as usize);
156
157    fn non_word_char(c: char) -> bool {
158        !(c.is_alphanumeric() || c == '_')
159    }
160
161    let start_idx = before.rfind(non_word_char)? as u32;
162    let end_idx = after.find(non_word_char).unwrap_or(after.len()) as u32;
163
164    // FIXME: use `ceil_char_boundary` from `std::str` when it gets stable
165    // https://github.com/rust-lang/rust/issues/93743
166    fn ceil_char_boundary(text: &str, index: u32) -> u32 {
167        (index..)
168            .find(|&index| text.is_char_boundary(index as usize))
169            .unwrap_or(text.len() as u32)
170    }
171
172    let from: TextSize = ceil_char_boundary(text, start_idx + 1).into();
173    let to: TextSize = (cursor_position + end_idx).into();
174
175    let range = TextRange::new(from, to);
176    if range.is_empty() {
177        None
178    } else {
179        Some(range + leaf.text_range().start())
180    }
181}
182
183fn extend_comments(comment: ast::Comment) -> Option<TextRange> {
184    let prev = adj_comments(&comment, Direction::Prev);
185    let next = adj_comments(&comment, Direction::Next);
186    if prev != next {
187        Some(TextRange::new(
188            prev.syntax().text_range().start(),
189            next.syntax().text_range().end(),
190        ))
191    } else {
192        None
193    }
194}
195
196fn adj_comments(comment: &ast::Comment, dir: Direction) -> ast::Comment {
197    let mut res = comment.clone();
198    for element in comment.syntax().siblings_with_tokens(dir) {
199        let Some(token) = element.as_token() else {
200            break;
201        };
202        if let Some(c) = ast::Comment::cast(token.clone()) {
203            res = c
204        } else if token.kind() != SyntaxKind::WHITESPACE || token.text().contains("\n\n") {
205            break;
206        }
207    }
208    res
209}
210
211fn extend_ws(root: &SyntaxNode, ws: SyntaxToken, offset: TextSize) -> TextRange {
212    let ws_text = ws.text();
213    let suffix = TextRange::new(offset, ws.text_range().end()) - ws.text_range().start();
214    let prefix = TextRange::new(ws.text_range().start(), offset) - ws.text_range().start();
215    let ws_suffix = &ws_text[suffix];
216    let ws_prefix = &ws_text[prefix];
217    if ws_text.contains('\n')
218        && !ws_suffix.contains('\n')
219        && let Some(node) = ws.next_sibling_or_token()
220    {
221        let start = match ws_prefix.rfind('\n') {
222            Some(idx) => ws.text_range().start() + TextSize::from((idx + 1) as u32),
223            None => node.text_range().start(),
224        };
225        let end = if root.text().char_at(node.text_range().end()) == Some('\n') {
226            node.text_range().end() + TextSize::of('\n')
227        } else {
228            node.text_range().end()
229        };
230        return TextRange::new(start, end);
231    }
232    ws.text_range()
233}
234
235fn pick_best(l: SyntaxToken, r: SyntaxToken) -> SyntaxToken {
236    return if priority(&r) > priority(&l) { r } else { l };
237    fn priority(n: &SyntaxToken) -> usize {
238        match n.kind() {
239            SyntaxKind::WHITESPACE => 0,
240            // TODO: we can probably include more here, rust analyzer includes a
241            // handful of keywords
242            SyntaxKind::IDENT => 2,
243            _ => 1,
244        }
245    }
246}
247
248/// Extend list item selection to include nearby delimiter and whitespace.
249fn extend_list_item(node: &SyntaxNode) -> Option<TextRange> {
250    fn is_single_line_ws(node: &SyntaxToken) -> bool {
251        node.kind() == SyntaxKind::WHITESPACE && !node.text().contains('\n')
252    }
253
254    fn nearby_comma(node: &SyntaxNode, dir: Direction) -> Option<SyntaxToken> {
255        node.siblings_with_tokens(dir)
256            .skip(1)
257            .find(|node| match node {
258                NodeOrToken::Node(_) => true,
259                NodeOrToken::Token(it) => !is_single_line_ws(it),
260            })
261            .and_then(|it| it.into_token())
262            .filter(|node| node.kind() == SyntaxKind::COMMA)
263    }
264
265    if let Some(comma) = nearby_comma(node, Direction::Next) {
266        // Include any following whitespace when delimiter is after list item.
267        let final_node = comma
268            .next_sibling_or_token()
269            .and_then(|n| n.into_token())
270            .filter(is_single_line_ws)
271            .unwrap_or(comma);
272
273        return Some(TextRange::new(
274            node.text_range().start(),
275            final_node.text_range().end(),
276        ));
277    }
278
279    if let Some(comma) = nearby_comma(node, Direction::Prev) {
280        return Some(TextRange::new(
281            comma.text_range().start(),
282            node.text_range().end(),
283        ));
284    }
285
286    None
287}
288
289#[cfg(test)]
290mod tests {
291    use super::*;
292    use crate::test_utils::fixture;
293    use insta::assert_debug_snapshot;
294    use squawk_syntax::{SourceFile, ast::AstNode};
295
296    fn expand(sql: &str) -> Vec<String> {
297        let (offset, sql) = fixture(sql);
298        let parse = SourceFile::parse(&sql);
299        let file = parse.tree();
300        let root = file.syntax();
301
302        let mut range = TextRange::empty(offset);
303        let mut results = vec![];
304
305        for _ in 0..20 {
306            let new_range = extend_selection(root, range);
307            if new_range == range {
308                break;
309            }
310            range = new_range;
311            results.push(sql[range].to_string());
312        }
313
314        results
315    }
316
317    #[test]
318    fn simple() {
319        assert_debug_snapshot!(expand(r#"select $01 + 1"#), @r#"
320        [
321            "1",
322            "1 + 1",
323            "select 1 + 1",
324        ]
325        "#);
326    }
327
328    #[test]
329    fn word_in_string_string() {
330        assert_debug_snapshot!(expand(r"
331select 'some stret$0ched out words in a string'
332"), @r#"
333        [
334            "stretched",
335            "'some stretched out words in a string'",
336            "select 'some stretched out words in a string'",
337            "\nselect 'some stretched out words in a string'\n",
338        ]
339        "#);
340    }
341
342    #[test]
343    fn string() {
344        assert_debug_snapshot!(expand(r"
345select b'foo$0 bar'
346'buzz';
347"), @r#"
348        [
349            "foo",
350            "b'foo bar'",
351            "b'foo bar'\n'buzz'",
352            "select b'foo bar'\n'buzz'",
353            "\nselect b'foo bar'\n'buzz';\n",
354        ]
355        "#);
356    }
357
358    #[test]
359    fn dollar_string() {
360        assert_debug_snapshot!(expand(r"
361select $$foo$0 bar$$;
362"), @r#"
363        [
364            "foo",
365            "$$foo bar$$",
366            "select $$foo bar$$",
367            "\nselect $$foo bar$$;\n",
368        ]
369        "#);
370    }
371
372    #[test]
373    fn comment_muli_line() {
374        assert_debug_snapshot!(expand(r"
375-- foo bar
376-- buzz$0
377-- boo
378select 1
379"), @r#"
380        [
381            "-- buzz",
382            "-- foo bar\n-- buzz\n-- boo",
383            "\n-- foo bar\n-- buzz\n-- boo\nselect 1\n",
384        ]
385        "#);
386    }
387
388    #[test]
389    fn comment() {
390        assert_debug_snapshot!(expand(r"
391-- foo bar$0
392select 1
393"), @r#"
394        [
395            "-- foo bar",
396            "\n-- foo bar\nselect 1\n",
397        ]
398        "#);
399
400        assert_debug_snapshot!(expand(r"
401/* foo bar$0 */
402select 1
403"), @r#"
404        [
405            "bar",
406            "/* foo bar */",
407            "\n/* foo bar */\nselect 1\n",
408        ]
409        "#);
410    }
411
412    #[test]
413    fn create_table_with_comment() {
414        assert_debug_snapshot!(expand(r"
415-- foo bar buzz
416create table t(
417  x int$0,
418  y text
419);
420"), @r#"
421        [
422            "int",
423            "x int",
424            "x int,",
425            "(\n  x int,\n  y text\n)",
426            "-- foo bar buzz\ncreate table t(\n  x int,\n  y text\n)",
427            "\n-- foo bar buzz\ncreate table t(\n  x int,\n  y text\n);\n",
428        ]
429        "#);
430    }
431
432    #[test]
433    fn column_list() {
434        assert_debug_snapshot!(expand(r#"create table t($0x int)"#), @r#"
435        [
436            "x",
437            "x int",
438            "(x int)",
439            "create table t(x int)",
440        ]
441        "#);
442
443        assert_debug_snapshot!(expand(r#"create table t($0x int, y int)"#), @r#"
444        [
445            "x",
446            "x int",
447            "x int, ",
448            "(x int, y int)",
449            "create table t(x int, y int)",
450        ]
451        "#);
452
453        assert_debug_snapshot!(expand(r#"create table t(x int, $0y int)"#), @r#"
454        [
455            "y",
456            "y int",
457            ", y int",
458            "(x int, y int)",
459            "create table t(x int, y int)",
460        ]
461        "#);
462    }
463
464    #[test]
465    fn start_of_line_whitespace_select() {
466        assert_debug_snapshot!(expand(r#"    
467select 1;
468
469$0    select 2;"#), @r#"
470        [
471            "    select 2",
472            "    \nselect 1;\n\n    select 2;",
473        ]
474        "#);
475    }
476
477    #[test]
478    fn select_list() {
479        assert_debug_snapshot!(expand(r#"select x$0, y from t"#), @r#"
480        [
481            "x",
482            "x, ",
483            "x, y",
484            "select x, y",
485            "select x, y from t",
486        ]
487        "#);
488
489        assert_debug_snapshot!(expand(r#"select x, y$0 from t"#), @r#"
490        [
491            "y",
492            ", y",
493            "x, y",
494            "select x, y",
495            "select x, y from t",
496        ]
497        "#);
498    }
499
500    #[test]
501    fn expand_whitespace() {
502        assert_debug_snapshot!(expand(r#"select 1 + 
503$0
5041;"#), @r#"
505        [
506            " \n\n",
507            "1 + \n\n1",
508            "select 1 + \n\n1",
509            "select 1 + \n\n1;",
510        ]
511        "#);
512    }
513
514    #[test]
515    fn function_args() {
516        assert_debug_snapshot!(expand(r#"select f(1$0, 2)"#), @r#"
517        [
518            "1",
519            "1, ",
520            "(1, 2)",
521            "f(1, 2)",
522            "select f(1, 2)",
523        ]
524        "#);
525    }
526
527    #[test]
528    fn prefer_idents() {
529        assert_debug_snapshot!(expand(r#"select foo$0+bar"#), @r#"
530        [
531            "foo",
532            "foo+bar",
533            "select foo+bar",
534        ]
535        "#);
536
537        assert_debug_snapshot!(expand(r#"select foo+$0bar"#), @r#"
538        [
539            "bar",
540            "foo+bar",
541            "select foo+bar",
542        ]
543        "#);
544    }
545
546    #[test]
547    fn list_variants() {
548        let delimited_ws_list_kinds = &[
549            SyntaxKind::CREATE_DATABASE_OPTION_LIST,
550            SyntaxKind::FUNC_OPTION_LIST,
551            SyntaxKind::ROLE_OPTION_LIST,
552            SyntaxKind::SEQUENCE_OPTION_LIST,
553            SyntaxKind::TRIGGER_EVENT_LIST,
554            SyntaxKind::XML_COLUMN_OPTION_LIST,
555            SyntaxKind::WHEN_CLAUSE_LIST,
556        ];
557
558        let unhandled_list_kinds = (0..SyntaxKind::__LAST as u16)
559            .map(SyntaxKind::from)
560            .filter(|kind| {
561                format!("{:?}", kind).ends_with("_LIST") && !delimited_ws_list_kinds.contains(kind)
562            })
563            .filter(|kind| !DELIMITED_LIST_KINDS.contains(kind))
564            .collect::<Vec<_>>();
565
566        assert_eq!(
567            unhandled_list_kinds,
568            vec![],
569            "We shouldn't have any unhandled list kinds"
570        )
571    }
572}
squawk_ide/expand_selection.rs

squawk_ide/
expand_selection.rs
