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

squawk_ide/
expand_selection.rs
