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

squawk_ide/
expand_selection.rs
