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

squawk_ide/
expand_selection.rs
