Skip to content

Commit 4e5521c

Browse files
juleswritescodejuleswritescode
committed
awesome!
1 parent 6efc22d commit 4e5521c

File tree

6 files changed

+219
-61
lines changed

6 files changed

+219
-61
lines changed

‎Cargo.lock‎

Lines changed: 13 additions & 3 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎crates/pgt_completions/src/complete.rs‎

Lines changed: 123 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -21,8 +21,45 @@ pub struct CompletionParams<'a> {
2121
text = params.text,
2222
position = params.position.to_string()
2323
))]
24-
pub fn complete(params: CompletionParams) -> Vec<CompletionItem> {
25-
let ctx = CompletionContext::new(&params);
24+
pub fn complete(mut params: CompletionParams) -> Vec<CompletionItem> {
25+
let should_adjust_params = params.tree.is_some()
26+
&& (cursor_inbetween_nodes(params.tree.unwrap(), params.position)
27+
|| cursor_prepared_to_write_token_after_last_node(
28+
params.tree.unwrap(),
29+
params.position,
30+
));
31+
32+
let usable_sql = if should_adjust_params {
33+
let pos: usize = params.position.into();
34+
35+
let mut mutated_sql = String::new();
36+
37+
for (idx, c) in params.text.chars().enumerate() {
38+
if idx == pos {
39+
mutated_sql.push_str("REPLACED_TOKEN ");
40+
}
41+
mutated_sql.push(c);
42+
}
43+
44+
mutated_sql
45+
} else {
46+
params.text
47+
};
48+
49+
let usable_tree = if should_adjust_params {
50+
let mut parser = tree_sitter::Parser::new();
51+
parser
52+
.set_language(tree_sitter_sql::language())
53+
.expect("Error loading sql language");
54+
parser.parse(usable_sql.clone(), None)
55+
} else {
56+
tracing::info!("We're reusing the previous tree.");
57+
None
58+
};
59+
60+
params.text = usable_sql;
61+
62+
let ctx = CompletionContext::new(&params, usable_tree.as_ref().or(params.tree));
2663

2764
let mut builder = CompletionBuilder::new();
2865

@@ -32,3 +69,87 @@ pub fn complete(params: CompletionParams) -> Vec<CompletionItem> {
3269

3370
builder.finish()
3471
}
72+
73+
fn cursor_inbetween_nodes(tree: &tree_sitter::Tree, position: TextSize) -> bool {
74+
let mut cursor = tree.walk();
75+
let mut node = tree.root_node();
76+
77+
loop {
78+
let child_dx = cursor.goto_first_child_for_byte(position.into());
79+
if child_dx.is_none() {
80+
break;
81+
}
82+
node = cursor.node();
83+
}
84+
85+
let byte = position.into();
86+
87+
// Return true if the cursor is NOT within the node's bounds, INCLUSIVE
88+
!(node.start_byte() <= byte && node.end_byte() >= byte)
89+
}
90+
91+
fn cursor_prepared_to_write_token_after_last_node(
92+
tree: &tree_sitter::Tree,
93+
position: TextSize,
94+
) -> bool {
95+
let cursor_pos: usize = position.into();
96+
cursor_pos == tree.root_node().end_byte() + 1
97+
}
98+
99+
#[cfg(test)]
100+
mod tests {
101+
use pgt_text_size::TextSize;
102+
103+
use crate::complete::{cursor_inbetween_nodes, cursor_prepared_to_write_token_after_last_node};
104+
105+
#[test]
106+
fn test_cursor_inbetween_nodes() {
107+
let input = "select from users;";
108+
109+
let mut parser = tree_sitter::Parser::new();
110+
parser
111+
.set_language(tree_sitter_sql::language())
112+
.expect("Error loading sql language");
113+
114+
let mut tree = parser.parse(input.to_string(), None).unwrap();
115+
116+
// select | from users;
117+
assert!(cursor_inbetween_nodes(&mut tree, TextSize::new(7)));
118+
119+
// select |from users;
120+
assert!(!cursor_inbetween_nodes(&mut tree, TextSize::new(8)));
121+
122+
// select| from users;
123+
assert!(!cursor_inbetween_nodes(&mut tree, TextSize::new(6)));
124+
}
125+
126+
#[test]
127+
fn test_cursor_after_nodes() {
128+
let input = "select * from ";
129+
130+
let mut parser = tree_sitter::Parser::new();
131+
parser
132+
.set_language(tree_sitter_sql::language())
133+
.expect("Error loading sql language");
134+
135+
let mut tree = parser.parse(input.to_string(), None).unwrap();
136+
137+
// select * from|; <-- still on previous token
138+
assert!(!cursor_prepared_to_write_token_after_last_node(
139+
&mut tree,
140+
TextSize::new(14)
141+
));
142+
143+
// select * from |; <-- too far off
144+
assert!(!cursor_prepared_to_write_token_after_last_node(
145+
&mut tree,
146+
TextSize::new(16)
147+
));
148+
149+
// select * from |; <-- just right
150+
assert!(cursor_prepared_to_write_token_after_last_node(
151+
&mut tree,
152+
TextSize::new(15)
153+
));
154+
}
155+
}

‎crates/pgt_completions/src/context.rs‎

Lines changed: 13 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -50,28 +50,12 @@ impl TryFrom<String> for ClauseType {
5050

5151
pub(crate) struct CompletionContext<'a> {
5252
pub node_under_cursor: Option<tree_sitter::Node<'a>>,
53-
pub previous_node: Option<tree_sitter::Node<'a>>,
5453

5554
pub tree: Option<&'a tree_sitter::Tree>,
5655
pub text: &'a str,
5756
pub schema_cache: &'a SchemaCache,
5857
pub position: usize,
5958

60-
/// If the cursor of the user is offset to the right of the statement,
61-
/// we'll have to move it back to the last node, otherwise, tree-sitter will break.
62-
/// However, knowing that the user is typing on the "next" node lets us prioritize different completion results.
63-
/// We consider an offset of up to two characters as valid.
64-
///
65-
/// Example:
66-
///
67-
/// ```
68-
/// select * from {}
69-
/// ```
70-
///
71-
/// We'll adjust the cursor position so it lies on the "from" token – but we're looking
72-
/// for table completions.
73-
pub cursor_offset_from_end: bool,
74-
7559
pub schema_name: Option<String>,
7660
pub wrapping_clause_type: Option<ClauseType>,
7761
pub is_invocation: bool,
@@ -81,14 +65,12 @@ pub(crate) struct CompletionContext<'a> {
8165
}
8266

8367
impl<'a> CompletionContext<'a> {
84-
pub fn new(params: &'a CompletionParams) -> Self {
68+
pub fn new(params: &'a CompletionParams, usable_tree: Option<&'a tree_sitter::Tree>) -> Self {
8569
let mut ctx = Self {
86-
tree: params.tree,
70+
tree: usable_tree,
8771
text: &params.text,
8872
schema_cache: params.schema,
8973
position: usize::from(params.position),
90-
cursor_offset_from_end: false,
91-
previous_node: None,
9274
node_under_cursor: None,
9375
schema_name: None,
9476
wrapping_clause_type: None,
@@ -97,7 +79,10 @@ impl<'a> CompletionContext<'a> {
9779
mentioned_relations: HashMap::new(),
9880
};
9981

82+
tracing::warn!("gathering tree context");
10083
ctx.gather_tree_context();
84+
85+
tracing::warn!("gathering info from ts query");
10186
ctx.gather_info_from_ts_queries();
10287

10388
ctx
@@ -164,14 +149,10 @@ impl<'a> CompletionContext<'a> {
164149
* `select * from use {}` becomes `select * from use{}`.
165150
*/
166151
let current_node = cursor.node();
167-
let position_cache = self.position.clone();
168152
while cursor.goto_first_child_for_byte(self.position).is_none() && self.position > 0 {
169153
self.position -= 1;
170154
}
171155

172-
let cursor_offset = position_cache - self.position;
173-
self.cursor_offset_from_end = cursor_offset > 0 && cursor_offset <= 2;
174-
175156
self.gather_context_from_node(cursor, current_node);
176157
}
177158

@@ -223,23 +204,11 @@ impl<'a> CompletionContext<'a> {
223204

224205
// We have arrived at the leaf node
225206
if current_node.child_count() == 0 {
226-
if self.cursor_offset_from_end {
207+
if self.get_ts_node_content(current_node).unwrap() == "REPLACED_TOKEN" {
227208
self.node_under_cursor = None;
228-
self.previous_node = Some(current_node);
229209
} else {
230-
// for the previous node, either select the previous sibling,
231-
// or collect the parent's previous sibling's last child.
232-
let previous = match current_node.prev_sibling() {
233-
Some(n) => Some(n),
234-
None => {
235-
let sib_of_parent = parent_node.prev_sibling();
236-
sib_of_parent.and_then(|p| p.children(&mut cursor).last())
237-
}
238-
};
239210
self.node_under_cursor = Some(current_node);
240-
self.previous_node = previous;
241211
}
242-
243212
return;
244213
}
245214

@@ -305,7 +274,7 @@ mod tests {
305274
schema: &pgt_schema_cache::SchemaCache::default(),
306275
};
307276

308-
let ctx = CompletionContext::new(&params);
277+
let ctx = CompletionContext::new(&params, Some(&tree));
309278

310279
assert_eq!(ctx.wrapping_clause_type, expected_clause.try_into().ok());
311280
}
@@ -337,7 +306,7 @@ mod tests {
337306
schema: &pgt_schema_cache::SchemaCache::default(),
338307
};
339308

340-
let ctx = CompletionContext::new(&params);
309+
let ctx = CompletionContext::new(&params, Some(&tree));
341310

342311
assert_eq!(ctx.schema_name, expected_schema.map(|f| f.to_string()));
343312
}
@@ -371,7 +340,7 @@ mod tests {
371340
schema: &pgt_schema_cache::SchemaCache::default(),
372341
};
373342

374-
let ctx = CompletionContext::new(&params);
343+
let ctx = CompletionContext::new(&params, Some(&tree));
375344

376345
assert_eq!(ctx.is_invocation, is_invocation);
377346
}
@@ -396,7 +365,7 @@ mod tests {
396365
schema: &pgt_schema_cache::SchemaCache::default(),
397366
};
398367

399-
let ctx = CompletionContext::new(&params);
368+
let ctx = CompletionContext::new(&params, Some(&tree));
400369

401370
let node = ctx.node_under_cursor.unwrap();
402371

@@ -424,7 +393,7 @@ mod tests {
424393
schema: &pgt_schema_cache::SchemaCache::default(),
425394
};
426395

427-
let ctx = CompletionContext::new(&params);
396+
let ctx = CompletionContext::new(&params, Some(&tree));
428397

429398
let node = ctx.node_under_cursor.unwrap();
430399

@@ -450,7 +419,7 @@ mod tests {
450419
schema: &pgt_schema_cache::SchemaCache::default(),
451420
};
452421

453-
let ctx = CompletionContext::new(&params);
422+
let ctx = CompletionContext::new(&params, Some(&tree));
454423

455424
let node = ctx.node_under_cursor.unwrap();
456425

@@ -475,7 +444,7 @@ mod tests {
475444
schema: &pgt_schema_cache::SchemaCache::default(),
476445
};
477446

478-
let ctx = CompletionContext::new(&params);
447+
let ctx = CompletionContext::new(&params, Some(&tree));
479448

480449
let node = ctx.node_under_cursor.unwrap();
481450

‎crates/pgt_workspace/Cargo.toml‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@ futures = "0.3.31"
1818
globset = "0.4.16"
1919

2020
ignore = { workspace = true }
21+
itertools = { version = "0.14.0" }
2122
pgt_analyse = { workspace = true, features = ["serde"] }
2223
pgt_analyser = { workspace = true }
2324
pgt_completions = { workspace = true }

0 commit comments

Comments
 (0)