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

squawk_ide/
expand_selection.rs
