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_line_index::find_newline;
32use squawk_syntax::{
33    SyntaxKind, SyntaxNode, SyntaxToken,
34    ast::{self, AstToken},
35};
36
37use crate::tokens::is_string_or_comment;
38
39const DELIMITED_LIST_KINDS: &[SyntaxKind] = &[
40    SyntaxKind::ALTER_OPTION_LIST,
41    SyntaxKind::ALIAS_COLUMN_LIST,
42    SyntaxKind::ALTER_TYPE_ATTRIBUTE_ACTION_LIST,
43    SyntaxKind::ARG_LIST,
44    SyntaxKind::ATTRIBUTE_LIST,
45    SyntaxKind::BEGIN_FUNC_OPTION_LIST,
46    SyntaxKind::CHECKPOINT_OPTION_LIST,
47    SyntaxKind::COLUMN_LIST,
48    SyntaxKind::COLUMN_REF_LIST,
49    SyntaxKind::COLUMN_TARGET_LIST,
50    SyntaxKind::COMPOSITE_FIELD_LIST,
51    SyntaxKind::CONFLICT_INDEX_ITEM_LIST,
52    SyntaxKind::CONSTRAINT_COLUMN_REF_LIST,
53    SyntaxKind::CONSTRAINT_EXCLUSION_LIST,
54    SyntaxKind::COPY_OPTION_LIST,
55    SyntaxKind::DROP_OP_CLASS_OPTION_LIST,
56    SyntaxKind::EXPLAIN_OPTION_LIST,
57    SyntaxKind::FDW_OPTION_LIST,
58    SyntaxKind::FOREIGN_KEY_COLUMN_LIST,
59    SyntaxKind::FUNCTION_SIG_LIST,
60    SyntaxKind::PROCEDURE_SIG_LIST,
61    SyntaxKind::ROUTINE_SIG_LIST,
62    SyntaxKind::GRANT_ROLE_OPTION_LIST,
63    SyntaxKind::GROUP_BY_LIST,
64    SyntaxKind::JSON_TABLE_COLUMN_LIST,
65    SyntaxKind::OPERATOR_CLASS_OPTION_LIST,
66    SyntaxKind::OPTION_ALTER_OPTION_LIST,
67    SyntaxKind::OPTION_ITEM_LIST,
68    SyntaxKind::OP_SIG_LIST,
69    SyntaxKind::PARAM_LIST,
70    SyntaxKind::PARTITION_ITEM_LIST,
71    SyntaxKind::PARTITION_LIST,
72    SyntaxKind::TABLE_NAME_REF_LIST,
73    SyntaxKind::REINDEX_OPTION_LIST,
74    SyntaxKind::RELATION_LIST,
75    SyntaxKind::RETURNING_OPTION_LIST,
76    SyntaxKind::REVOKE_COMMAND_LIST,
77    SyntaxKind::ROLE_REF_LIST,
78    SyntaxKind::ROW_LIST,
79    SyntaxKind::RULE_STMT_LIST,
80    SyntaxKind::EXPR_AS_COLUMN_NAME_LIST,
81    SyntaxKind::EXPR_AS_ELEMENT_TAG_LIST,
82    SyntaxKind::EXPR_AS_PROPERTY_NAME_LIST,
83    SyntaxKind::EXPR_AS_XML_ATTR_LIST,
84    SyntaxKind::XML_NAMESPACE_LIST,
85    SyntaxKind::SET_COLUMN_LIST,
86    SyntaxKind::SET_EXPR_LIST,
87    SyntaxKind::SET_OPTIONS_LIST,
88    SyntaxKind::SORT_BY_LIST,
89    SyntaxKind::TABLE_AND_COLUMNS_LIST,
90    SyntaxKind::TABLE_ARG_LIST,
91    SyntaxKind::TABLE_LIST,
92    SyntaxKind::TARGET_LIST,
93    SyntaxKind::TRANSACTION_MODE_LIST,
94    SyntaxKind::VACUUM_OPTION_LIST,
95    SyntaxKind::VARIANT_LIST,
96    SyntaxKind::XML_TABLE_COLUMN_LIST,
97    SyntaxKind::PATH_PATTERN_LIST,
98];
99
100pub fn extend_selection(root: &SyntaxNode, range: TextRange) -> TextRange {
101    try_extend_selection(root, range).unwrap_or(range)
102}
103
104fn try_extend_selection(root: &SyntaxNode, range: TextRange) -> Option<TextRange> {
105    if range.is_empty() {
106        let offset = range.start();
107        let mut leaves = root.token_at_offset(offset);
108        // Make sure that if we're on the whitespace at the start of a line, we
109        // expand to the node on that line instead of the previous one
110        if leaves.clone().all(|it| it.kind() == SyntaxKind::WHITESPACE) {
111            return Some(extend_ws(root, leaves.next()?, offset));
112        }
113        let leaf_range = match root.token_at_offset(offset) {
114            rowan::TokenAtOffset::None => return None,
115            rowan::TokenAtOffset::Single(l) => {
116                if is_string_or_comment(l.kind()) {
117                    extend_single_word_in_comment_or_string(&l, offset)
118                        .unwrap_or_else(|| l.text_range())
119                } else {
120                    l.text_range()
121                }
122            }
123            rowan::TokenAtOffset::Between(l, r) => pick_best(l, r).text_range(),
124        };
125        return Some(leaf_range);
126    }
127
128    let node = match root.covering_element(range) {
129        NodeOrToken::Token(token) => {
130            if token.text_range() != range {
131                return Some(token.text_range());
132            }
133            if let Some(comment) = ast::Comment::cast(token.clone())
134                && let Some(range) = extend_comments(comment)
135            {
136                return Some(range);
137            }
138            token.parent()?
139        }
140        NodeOrToken::Node(node) => node,
141    };
142
143    if node.text_range() != range {
144        return Some(node.text_range());
145    }
146
147    let node = shallowest_node(&node);
148
149    if node
150        .parent()
151        .is_some_and(|n| DELIMITED_LIST_KINDS.contains(&n.kind()))
152    {
153        if let Some(range) = extend_list_item(&node) {
154            return Some(range);
155        }
156    }
157
158    node.parent().map(|it| it.text_range())
159}
160
161/// Find the shallowest node with same range, which allows us to traverse siblings.
162fn shallowest_node(node: &SyntaxNode) -> SyntaxNode {
163    node.ancestors()
164        .take_while(|n| n.text_range() == node.text_range())
165        .last()
166        .unwrap()
167}
168
169/// Expand to the current word instead the full text range of the node.
170fn extend_single_word_in_comment_or_string(
171    leaf: &SyntaxToken,
172    offset: TextSize,
173) -> Option<TextRange> {
174    let text: &str = leaf.text();
175    let cursor_position: u32 = (offset - leaf.text_range().start()).into();
176
177    let (before, after) = text.split_at(cursor_position as usize);
178
179    fn non_word_char(c: char) -> bool {
180        !(c.is_alphanumeric() || c == '_')
181    }
182
183    let start_idx = before.rfind(non_word_char)? as u32;
184    let end_idx = after.find(non_word_char).unwrap_or(after.len()) as u32;
185
186    // FIXME: use `ceil_char_boundary` from `std::str` when it gets stable
187    // https://github.com/rust-lang/rust/issues/93743
188    fn ceil_char_boundary(text: &str, index: u32) -> u32 {
189        (index..)
190            .find(|&index| text.is_char_boundary(index as usize))
191            .unwrap_or(text.len() as u32)
192    }
193
194    let from: TextSize = ceil_char_boundary(text, start_idx + 1).into();
195    let to: TextSize = (cursor_position + end_idx).into();
196
197    let range = TextRange::new(from, to);
198    if range.is_empty() {
199        None
200    } else {
201        Some(range + leaf.text_range().start())
202    }
203}
204
205fn extend_comments(comment: ast::Comment) -> Option<TextRange> {
206    let prev = adj_comments(&comment, Direction::Prev);
207    let next = adj_comments(&comment, Direction::Next);
208    if prev != next {
209        Some(TextRange::new(
210            prev.syntax().text_range().start(),
211            next.syntax().text_range().end(),
212        ))
213    } else {
214        None
215    }
216}
217
218fn adj_comments(comment: &ast::Comment, dir: Direction) -> ast::Comment {
219    let mut res = comment.clone();
220    for element in comment.syntax().siblings_with_tokens(dir) {
221        let Some(token) = element.as_token() else {
222            break;
223        };
224        if let Some(c) = ast::Comment::cast(token.clone()) {
225            res = c
226        } else if let Some(ws) = ast::Whitespace::cast(token.clone()) {
227            if ws.spans_multiple_lines() {
228                break;
229            }
230        } else {
231            break;
232        }
233    }
234    res
235}
236
237fn extend_ws(root: &SyntaxNode, ws: SyntaxToken, offset: TextSize) -> TextRange {
238    let ws_text = ws.text();
239    let suffix = TextRange::new(offset, ws.text_range().end()) - ws.text_range().start();
240    let prefix = TextRange::new(ws.text_range().start(), offset) - ws.text_range().start();
241    let ws_suffix = &ws_text[suffix];
242    let ws_prefix = &ws_text[prefix];
243    if find_newline(ws_text).is_some()
244        && find_newline(ws_suffix).is_none()
245        && let Some(node) = ws.next_sibling_or_token()
246    {
247        let start = match last_newline_end(ws_prefix) {
248            Some(idx) => ws.text_range().start() + TextSize::from(idx as u32),
249            None => node.text_range().start(),
250        };
251        let node_end = node.text_range().end();
252        let line_ending_len = match root.text().char_at(node_end) {
253            Some('\r') if root.text().char_at(node_end + TextSize::of('\r')) == Some('\n') => {
254                TextSize::of("\r\n")
255            }
256            Some('\n' | '\r') => TextSize::of('\n'),
257            _ => TextSize::default(),
258        };
259        return TextRange::new(start, node_end + line_ending_len);
260    }
261    ws.text_range()
262}
263
264fn last_newline_end(text: &str) -> Option<usize> {
265    let mut rest = text;
266    let mut offset = 0;
267    let mut last = None;
268
269    while let Some((idx, line_ending)) = find_newline(rest) {
270        offset += idx + line_ending.as_str().len();
271        last = Some(offset);
272        rest = &text[offset..];
273    }
274
275    last
276}
277
278fn pick_best(l: SyntaxToken, r: SyntaxToken) -> SyntaxToken {
279    return if priority(&r) > priority(&l) { r } else { l };
280    fn priority(n: &SyntaxToken) -> usize {
281        match n.kind() {
282            SyntaxKind::WHITESPACE => 0,
283            // TODO: we can probably include more here, rust analyzer includes a
284            // handful of keywords
285            SyntaxKind::IDENT => 2,
286            _ => 1,
287        }
288    }
289}
290
291/// Extend list item selection to include nearby delimiter and whitespace.
292fn extend_list_item(node: &SyntaxNode) -> Option<TextRange> {
293    fn is_single_line_ws(node: &SyntaxToken) -> bool {
294        node.kind() == SyntaxKind::WHITESPACE && find_newline(node.text()).is_none()
295    }
296
297    fn nearby_comma(node: &SyntaxNode, dir: Direction) -> Option<SyntaxToken> {
298        node.siblings_with_tokens(dir)
299            .skip(1)
300            .find(|node| match node {
301                NodeOrToken::Node(_) => true,
302                NodeOrToken::Token(it) => !is_single_line_ws(it),
303            })
304            .and_then(|it| it.into_token())
305            .filter(|node| node.kind() == SyntaxKind::COMMA)
306    }
307
308    if let Some(comma) = nearby_comma(node, Direction::Next) {
309        // Include any following whitespace when delimiter is after list item.
310        let final_node = comma
311            .next_sibling_or_token()
312            .and_then(|n| n.into_token())
313            .filter(is_single_line_ws)
314            .unwrap_or(comma);
315
316        return Some(TextRange::new(
317            node.text_range().start(),
318            final_node.text_range().end(),
319        ));
320    }
321
322    if let Some(comma) = nearby_comma(node, Direction::Prev) {
323        return Some(TextRange::new(
324            comma.text_range().start(),
325            node.text_range().end(),
326        ));
327    }
328
329    None
330}
331
332#[cfg(test)]
333mod tests {
334    use super::*;
335    use crate::test_utils::Fixture;
336    use insta::assert_debug_snapshot;
337    use squawk_syntax::ast::AstNode;
338
339    #[must_use]
340    fn expand(sql: &str) -> Vec<String> {
341        let fixture = Fixture::new(sql);
342        let offset = fixture.marker().offset();
343        let sql = offset.file_id.content(fixture.db()).clone();
344        let tree = crate::db::parse(fixture.db(), offset.file_id).tree();
345        let root = tree.syntax();
346
347        let mut range = TextRange::empty(offset.value);
348        let mut results = vec![];
349
350        for _ in 0..20 {
351            let new_range = extend_selection(root, range);
352            if new_range == range {
353                break;
354            }
355            range = new_range;
356            results.push(sql[range].to_string());
357        }
358
359        results
360    }
361
362    #[test]
363    fn simple() {
364        assert_debug_snapshot!(expand(r#"select $01 + 1"#), @r#"
365        [
366            "1",
367            "1 + 1",
368            "select 1 + 1",
369        ]
370        "#);
371    }
372
373    #[test]
374    fn word_in_string_string() {
375        assert_debug_snapshot!(expand(r"
376select 'some stret$0ched out words in a string'
377"), @r#"
378        [
379            "stretched",
380            "'some stretched out words in a string'",
381            "select 'some stretched out words in a string'",
382            "\nselect 'some stretched out words in a string'\n",
383        ]
384        "#);
385    }
386
387    #[test]
388    fn string() {
389        assert_debug_snapshot!(expand(r"
390select e'foo$0 bar'
391'buzz';
392"), @r#"
393        [
394            "foo",
395            "e'foo bar'",
396            "e'foo bar'\n'buzz'",
397            "select e'foo bar'\n'buzz'",
398            "select e'foo bar'\n'buzz';",
399            "\nselect e'foo bar'\n'buzz';\n",
400        ]
401        "#);
402    }
403
404    #[test]
405    fn dollar_string() {
406        assert_debug_snapshot!(expand(r"
407select $$foo$0 bar$$;
408"), @r#"
409        [
410            "foo",
411            "$$foo bar$$",
412            "select $$foo bar$$",
413            "select $$foo bar$$;",
414            "\nselect $$foo bar$$;\n",
415        ]
416        "#);
417    }
418
419    #[test]
420    fn comment_muli_line() {
421        assert_debug_snapshot!(expand(r"
422-- foo bar
423-- buzz$0
424-- boo
425select 1
426"), @r#"
427        [
428            "-- buzz",
429            "-- foo bar\n-- buzz\n-- boo",
430            "\n-- foo bar\n-- buzz\n-- boo\nselect 1\n",
431        ]
432        "#);
433    }
434
435    #[test]
436    fn comments_separated_by_cr_blank_line() {
437        assert_debug_snapshot!(expand("-- first\r\r-- sec$0ond\rselect 1;"), @r#"
438        [
439            "second",
440            "-- second",
441            "-- first\r\r-- second\rselect 1;",
442        ]
443        "#);
444    }
445
446    #[test]
447    fn comment() {
448        assert_debug_snapshot!(expand(r"
449-- foo bar$0
450select 1
451"), @r#"
452        [
453            "-- foo bar",
454            "\n-- foo bar\nselect 1\n",
455        ]
456        "#);
457
458        assert_debug_snapshot!(expand(r"
459/* foo bar$0 */
460select 1
461"), @r#"
462        [
463            "bar",
464            "/* foo bar */",
465            "\n/* foo bar */\nselect 1\n",
466        ]
467        "#);
468    }
469
470    #[test]
471    fn create_table_with_comment() {
472        assert_debug_snapshot!(expand(r"
473-- foo bar buzz
474create table t(
475  x int$0,
476  y text
477);
478"), @r#"
479        [
480            "int",
481            "x int",
482            "x int,",
483            "(\n  x int,\n  y text\n)",
484            "create table t(\n  x int,\n  y text\n);",
485            "\n-- foo bar buzz\ncreate table t(\n  x int,\n  y text\n);\n",
486        ]
487        "#);
488    }
489
490    #[test]
491    fn column_list() {
492        assert_debug_snapshot!(expand(r#"create table t($0x int)"#), @r#"
493        [
494            "x",
495            "x int",
496            "(x int)",
497            "create table t(x int)",
498        ]
499        "#);
500
501        assert_debug_snapshot!(expand(r#"create table t($0x int, y int)"#), @r#"
502        [
503            "x",
504            "x int",
505            "x int, ",
506            "(x int, y int)",
507            "create table t(x int, y int)",
508        ]
509        "#);
510
511        assert_debug_snapshot!(expand(r#"create table t(x int, $0y int)"#), @r#"
512        [
513            "y",
514            "y int",
515            ", y int",
516            "(x int, y int)",
517            "create table t(x int, y int)",
518        ]
519        "#);
520    }
521
522    #[test]
523    fn start_of_line_whitespace_select() {
524        assert_debug_snapshot!(expand(r#"    
525select 1;
526
527$0    select 2;"#), @r#"
528        [
529            "    select 2;",
530            "    \nselect 1;\n\n    select 2;",
531        ]
532        "#);
533    }
534
535    #[test]
536    fn start_of_line_whitespace_select_with_cr_line_endings() {
537        assert_debug_snapshot!(expand("select 1;\r$0    select 2;"), @r#"
538        [
539            "    select 2;",
540            "select 1;\r    select 2;",
541        ]
542        "#);
543    }
544
545    #[test]
546    fn select_list() {
547        assert_debug_snapshot!(expand(r#"select x$0, y from t"#), @r#"
548        [
549            "x",
550            "x, ",
551            "x, y",
552            "select x, y",
553            "select x, y from t",
554        ]
555        "#);
556
557        assert_debug_snapshot!(expand(r#"select x, y$0 from t"#), @r#"
558        [
559            "y",
560            ", y",
561            "x, y",
562            "select x, y",
563            "select x, y from t",
564        ]
565        "#);
566    }
567
568    #[test]
569    fn expand_whitespace() {
570        assert_debug_snapshot!(expand(r#"select 1 + 
571$0
5721;"#), @r#"
573        [
574            " \n\n",
575            "1 + \n\n1",
576            "select 1 + \n\n1",
577            "select 1 + \n\n1;",
578        ]
579        "#);
580    }
581
582    #[test]
583    fn function_args() {
584        assert_debug_snapshot!(expand(r#"select f(1$0, 2)"#), @r#"
585        [
586            "1",
587            "1, ",
588            "(1, 2)",
589            "f(1, 2)",
590            "select f(1, 2)",
591        ]
592        "#);
593    }
594
595    #[test]
596    fn prefer_idents() {
597        assert_debug_snapshot!(expand(r#"select foo$0+bar"#), @r#"
598        [
599            "foo",
600            "foo+bar",
601            "select foo+bar",
602        ]
603        "#);
604
605        assert_debug_snapshot!(expand(r#"select foo+$0bar"#), @r#"
606        [
607            "bar",
608            "foo+bar",
609            "select foo+bar",
610        ]
611        "#);
612    }
613
614    #[test]
615    fn list_variants() {
616        let delimited_ws_list_kinds = &[
617            SyntaxKind::DATABASE_OPTION_LIST,
618            SyntaxKind::FUNC_OPTION_LIST,
619            SyntaxKind::ROLE_OPTION_LIST,
620            SyntaxKind::SEQUENCE_OPTION_LIST,
621            SyntaxKind::TRIGGER_EVENT_LIST,
622            SyntaxKind::XML_COLUMN_OPTION_LIST,
623            SyntaxKind::WHEN_CLAUSE_LIST,
624            SyntaxKind::LABEL_AND_PROPERTIES_LIST,
625        ];
626
627        let unhandled_list_kinds = (0..SyntaxKind::__LAST as u16)
628            .map(SyntaxKind::from)
629            .filter(|kind| {
630                format!("{kind:?}").ends_with("_LIST") && !delimited_ws_list_kinds.contains(kind)
631            })
632            .filter(|kind| !DELIMITED_LIST_KINDS.contains(kind))
633            .collect::<Vec<_>>();
634
635        assert_eq!(
636            unhandled_list_kinds,
637            vec![],
638            "We shouldn't have any unhandled list kinds"
639        )
640    }
641}
squawk_ide/expand_selection.rs

squawk_ide/
expand_selection.rs
