Skip to content

Commit 646b884

Browse files
juleswritescodejuleswritescode
authored
[3] refactor: improve & simplify tree-sitter context-tracking (#595)
1 parent 962fcfd commit 646b884

File tree

8 files changed

+263
-564
lines changed

8 files changed

+263
-564
lines changed

‎crates/pgls_completions/src/providers/columns.rs‎

Lines changed: 1 addition & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -58,8 +58,7 @@ mod tests {
5858
use crate::{
5959
CompletionItem, CompletionItemKind, complete,
6060
test_helper::{
61-
CompletionAssertion, assert_complete_results, assert_no_complete_results,
62-
get_test_deps, get_test_params,
61+
CompletionAssertion, assert_complete_results, get_test_deps, get_test_params,
6362
},
6463
};
6564

@@ -717,18 +716,6 @@ mod tests {
717716
&pool,
718717
)
719718
.await;
720-
721-
// no completions in the values list!
722-
assert_no_complete_results(
723-
format!(
724-
"insert into instruments (id, name) values ({})",
725-
QueryWithCursorPosition::cursor_marker()
726-
)
727-
.as_str(),
728-
None,
729-
&pool,
730-
)
731-
.await;
732719
}
733720

734721
#[sqlx::test(migrator = "pgls_test_utils::MIGRATIONS")]

‎crates/pgls_completions/src/providers/helper.rs‎

Lines changed: 14 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
use pgls_text_size::{TextRange, TextSize};
1+
use pgls_text_size::TextRange;
22
use pgls_treesitter::TreesitterContext;
33

44
use crate::{is_sanitized_token_with_quote, remove_sanitized_token};
@@ -9,29 +9,25 @@ pub(crate) fn node_text_surrounded_by_quotes(ctx: &TreesitterContext) -> bool {
99
}
1010

1111
pub(crate) fn get_range_to_replace(ctx: &TreesitterContext) -> TextRange {
12-
match ctx.node_under_cursor.as_ref() {
13-
Some(node) => {
14-
let content = ctx.get_node_under_cursor_content().unwrap_or("".into());
15-
let content = content.as_str();
12+
let node = &ctx.node_under_cursor;
13+
let content = ctx.get_node_under_cursor_content().unwrap_or("".into());
14+
let content = content.as_str();
1615

17-
let sanitized = remove_sanitized_token(content);
18-
let length = sanitized.len();
16+
let sanitized = remove_sanitized_token(content);
17+
let length = sanitized.len();
1918

20-
let mut start = node.start_byte();
21-
let mut end = start + length;
19+
let mut start = node.start_byte();
20+
let mut end = start + length;
2221

23-
if sanitized.starts_with('"') && sanitized.ends_with('"') {
24-
start += 1;
22+
if sanitized.starts_with('"') && sanitized.ends_with('"') {
23+
start += 1;
2524

26-
if sanitized.len() > 1 {
27-
end -= 1;
28-
}
29-
}
30-
31-
TextRange::new(start.try_into().unwrap(), end.try_into().unwrap())
25+
if sanitized.len() > 1 {
26+
end -= 1;
3227
}
33-
None => TextRange::empty(TextSize::new(0)),
3428
}
29+
30+
TextRange::new(start.try_into().unwrap(), end.try_into().unwrap())
3531
}
3632

3733
pub(crate) fn only_leading_quote(ctx: &TreesitterContext) -> bool {

‎crates/pgls_completions/src/relevance/filtering.rs‎

Lines changed: 83 additions & 101 deletions
Original file line numberDiff line numberDiff line change
@@ -33,11 +33,7 @@ impl CompletionFilter<'_> {
3333
return None;
3434
}
3535

36-
let current_node_kind = ctx
37-
.node_under_cursor
38-
.as_ref()
39-
.map(|n| n.kind())
40-
.unwrap_or("");
36+
let current_node_kind = ctx.node_under_cursor.kind();
4137

4238
if current_node_kind.starts_with("keyword_")
4339
|| current_node_kind == "="
@@ -70,30 +66,24 @@ impl CompletionFilter<'_> {
7066
}
7167
}
7268

73-
if ctx
74-
.node_under_cursor
75-
.as_ref()
76-
.is_some_and(|n| n.kind() == "any_identifier")
77-
&& ctx.matches_ancestor_history(&["alias"])
69+
if ctx.node_under_cursor.kind() == "any_identifier"
70+
&& ctx.history_ends_with(&["alias", "any_identifier"])
7871
{
7972
return None;
8073
}
8174

8275
// No autocompletions if there are two identifiers without a separator.
83-
if ctx.node_under_cursor.as_ref().is_some_and(|node| {
84-
node.prev_sibling().is_some_and(|p| {
85-
(p.kind() == "any_identifier" || p.kind() == "object_reference")
86-
&& node.kind() == "any_identifier"
87-
})
76+
if ctx.node_under_cursor.prev_sibling().is_some_and(|p| {
77+
(p.kind() == "any_identifier" || p.kind() == "object_reference")
78+
&& ctx.node_under_cursor.kind() == "any_identifier"
8879
}) {
8980
return None;
9081
}
9182

9283
// no completions if we're right after an asterisk:
9384
// `select * {}`
94-
if ctx.node_under_cursor.as_ref().is_some_and(|node| {
95-
node.prev_sibling()
96-
.is_some_and(|p| (p.kind() == "all_fields") && node.kind() == "any_identifier")
85+
if ctx.node_under_cursor.prev_sibling().is_some_and(|p| {
86+
(p.kind() == "all_fields") && ctx.node_under_cursor.kind() == "any_identifier"
9787
}) {
9888
return None;
9989
}
@@ -102,7 +92,7 @@ impl CompletionFilter<'_> {
10292
}
10393

10494
fn check_specific_node_type(&self, ctx: &TreesitterContext) -> Option<()> {
105-
let kind = ctx.node_under_cursor.as_ref().map(|n| n.kind())?;
95+
let kind = ctx.node_under_cursor.kind();
10696

10797
let is_allowed = match kind {
10898
"column_identifier" => matches!(self.data, CompletionRelevanceData::Column(_)),
@@ -112,67 +102,55 @@ impl CompletionFilter<'_> {
112102
"table_identifier" => matches!(self.data, CompletionRelevanceData::Table(_)),
113103
"policy_identifier" => matches!(self.data, CompletionRelevanceData::Policy(_)),
114104

115-
"any_identifier" => {
116-
if false || ctx.matches_ancestor_history(&["insert_values", "object_reference"]) {
117-
false
118-
} else {
119-
match self.data {
120-
CompletionRelevanceData::Column(_) => {
121-
ctx.node_under_cursor_is_within_field_name(&[
122-
"object_reference_1of1",
123-
"object_reference_2of2",
124-
"object_reference_3of3",
125-
"column_reference_1of1",
126-
"column_reference_2of2",
127-
"column_reference_3of3",
128-
]) && !ctx
129-
.node_under_cursor_is_within_field_name(&["binary_expr_right"])
130-
&& !ctx.matches_ancestor_history(&[
131-
"insert_values",
132-
"object_reference",
133-
])
134-
}
105+
"any_identifier" => match self.data {
106+
CompletionRelevanceData::Column(_) => {
107+
ctx.node_under_cursor_is_within_field(&[
108+
"object_reference_1of1",
109+
"object_reference_2of2",
110+
"object_reference_3of3",
111+
"column_reference_1of1",
112+
"column_reference_2of2",
113+
"column_reference_3of3",
114+
]) && !ctx.node_under_cursor_is_within_field(&["binary_expr_right"])
115+
}
135116

136-
CompletionRelevanceData::Schema(_) => ctx
137-
.node_under_cursor_is_within_field_name(&[
138-
"object_reference_1of1",
139-
"object_reference_1of2",
140-
"object_reference_1of3",
141-
"type_reference_1of1",
142-
"table_reference_1of1",
143-
"column_reference_1of1",
144-
"column_reference_1of2",
145-
"function_reference_1of1",
146-
]),
117+
CompletionRelevanceData::Schema(_) => ctx.node_under_cursor_is_within_field(&[
118+
"object_reference_1of1",
119+
"object_reference_1of2",
120+
"object_reference_1of3",
121+
"type_reference_1of1",
122+
"table_reference_1of1",
123+
"column_reference_1of1",
124+
"column_reference_1of2",
125+
"function_reference_1of1",
126+
]),
127+
128+
CompletionRelevanceData::Function(f) => {
129+
ctx.node_under_cursor_is_within_field(&[
130+
"object_reference_1of1",
131+
"object_reference_2of2",
132+
"function_reference_1of1",
133+
]) && !(ctx.history_ends_with(&[
134+
"check_or_using_clause",
135+
"binary_expression",
136+
"object_reference",
137+
"any_identifier",
138+
]) && matches!(f.kind, ProcKind::Aggregate))
139+
}
147140

148-
CompletionRelevanceData::Function(f) => {
149-
ctx.node_under_cursor_is_within_field_name(&[
150-
"object_reference_1of1",
151-
"object_reference_2of2",
152-
"function_reference_1of1",
153-
]) && !(ctx.matches_ancestor_history(&[
154-
"check_or_using_clause",
155-
"binary_expression",
156-
"object_reference",
157-
]) && matches!(f.kind, ProcKind::Aggregate))
158-
}
141+
CompletionRelevanceData::Table(_) => ctx.node_under_cursor_is_within_field(&[
142+
"object_reference_1of1",
143+
"object_reference_1of2",
144+
"object_reference_2of2",
145+
"object_reference_2of3",
146+
"table_reference_1of1",
147+
"column_reference_1of1",
148+
"column_reference_1of2",
149+
"column_reference_2of2",
150+
]),
159151

160-
CompletionRelevanceData::Table(_) => ctx
161-
.node_under_cursor_is_within_field_name(&[
162-
"object_reference_1of1",
163-
"object_reference_1of2",
164-
"object_reference_2of2",
165-
"object_reference_2of3",
166-
"table_reference_1of1",
167-
"column_reference_1of1",
168-
"column_reference_1of2",
169-
"column_reference_2of2",
170-
]),
171-
172-
_ => false,
173-
}
174-
}
175-
}
152+
_ => false,
153+
},
176154

177155
_ => false,
178156
};
@@ -189,21 +167,27 @@ impl CompletionFilter<'_> {
189167
WrappingClause::From | WrappingClause::Update => true,
190168

191169
WrappingClause::RevokeStatement | WrappingClause::GrantStatement => ctx
192-
.matches_ancestor_history(&["grantable_on_table", "object_reference"]),
170+
.history_ends_with(&[
171+
"grantable_on_table",
172+
"object_reference",
173+
"any_identifier",
174+
]),
193175

194176
WrappingClause::Join { on_node: None } => true,
195-
WrappingClause::Join { on_node: Some(on) } => ctx
196-
.node_under_cursor
197-
.as_ref()
198-
.is_some_and(|cn| cn.start_byte() < on.end_byte()),
177+
WrappingClause::Join { on_node: Some(on) } => {
178+
ctx.node_under_cursor.start_byte() < on.end_byte()
179+
}
199180

200181
WrappingClause::Insert => {
201182
ctx.wrapping_node_kind
202183
.as_ref()
203184
.is_none_or(|n| n != &WrappingNode::List)
204185
&& (ctx.before_cursor_matches_kind(&["keyword_into"])
205186
|| (ctx.before_cursor_matches_kind(&["."])
206-
&& ctx.matches_ancestor_history(&["object_reference"])))
187+
&& ctx.history_ends_with(&[
188+
"object_reference",
189+
"any_identifier",
190+
])))
207191
}
208192

209193
WrappingClause::DropTable | WrappingClause::AlterTable => ctx
@@ -216,7 +200,7 @@ impl CompletionFilter<'_> {
216200
WrappingClause::CreatePolicy
217201
| WrappingClause::AlterPolicy
218202
| WrappingClause::DropPolicy => {
219-
ctx.matches_ancestor_history(&["object_reference"])
203+
ctx.history_ends_with(&["object_reference", "any_identifier"])
220204
&& ctx.before_cursor_matches_kind(&["keyword_on", "."])
221205
}
222206

@@ -239,10 +223,9 @@ impl CompletionFilter<'_> {
239223

240224
// We can complete columns in JOIN cluases, but only if we are after the
241225
// ON node in the "ON u.id = posts.user_id" part.
242-
WrappingClause::Join { on_node: Some(on) } => ctx
243-
.node_under_cursor
244-
.as_ref()
245-
.is_some_and(|cn| cn.start_byte() >= on.end_byte()),
226+
WrappingClause::Join { on_node: Some(on) } => {
227+
ctx.node_under_cursor.start_byte() >= on.end_byte()
228+
}
246229

247230
// we are in a JOIN, but definitely not after an ON
248231
WrappingClause::Join { on_node: None } => false,
@@ -256,7 +239,7 @@ impl CompletionFilter<'_> {
256239
WrappingClause::Where => {
257240
ctx.before_cursor_matches_kind(&["keyword_and", "keyword_where"])
258241
|| (ctx.before_cursor_matches_kind(&["field_qualifier"])
259-
&& ctx.matches_ancestor_history(&["field"]))
242+
&& ctx.history_ends_with(&["field", "any_identifier"]))
260243
}
261244

262245
WrappingClause::CheckOrUsingClause => {
@@ -294,11 +277,12 @@ impl CompletionFilter<'_> {
294277
| WrappingClause::Delete => true,
295278

296279
WrappingClause::RevokeStatement | WrappingClause::GrantStatement => {
297-
(ctx.matches_ancestor_history(&[
280+
(ctx.history_ends_with(&[
298281
"grantable_on_table",
299282
"object_reference",
283+
"any_identifier",
300284
]) && !ctx.has_any_qualifier())
301-
|| ctx.matches_ancestor_history(&["grantable_on_all"])
285+
|| ctx.history_ends_with(&["grantable_on_all", "any_identifier"])
302286
}
303287

304288
WrappingClause::Where => {
@@ -343,20 +327,18 @@ impl CompletionFilter<'_> {
343327
.before_cursor_matches_kind(&["keyword_role", "keyword_authorization"]),
344328

345329
WrappingClause::RevokeStatement | WrappingClause::GrantStatement => {
346-
ctx.matches_ancestor_history(&["role_specification"])
347-
|| ctx.node_under_cursor.as_ref().is_some_and(|k| {
348-
k.kind() == "any_identifier"
349-
&& ctx.before_cursor_matches_kind(&[
350-
"keyword_grant",
351-
"keyword_revoke",
352-
"keyword_for",
353-
])
354-
})
330+
ctx.history_ends_with(&["role_specification", "any_identifier"])
331+
|| (ctx.node_under_cursor.kind() == "any_identifier"
332+
&& ctx.before_cursor_matches_kind(&[
333+
"keyword_grant",
334+
"keyword_revoke",
335+
"keyword_for",
336+
]))
355337
}
356338

357339
WrappingClause::AlterPolicy | WrappingClause::CreatePolicy => {
358340
ctx.before_cursor_matches_kind(&["keyword_to"])
359-
&& ctx.matches_ancestor_history(&["policy_to_role"])
341+
&& ctx.history_ends_with(&["policy_to_role", "any_identifier"])
360342
}
361343

362344
_ => false,

‎crates/pgls_completions/src/relevance/scoring.rs‎

Lines changed: 4 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -87,11 +87,8 @@ impl CompletionScore<'_> {
8787
WrappingClause::Delete => 10,
8888
WrappingClause::From => 5,
8989
WrappingClause::Join { on_node }
90-
if on_node.is_none_or(|on| {
91-
ctx.node_under_cursor
92-
.as_ref()
93-
.is_none_or(|n| n.end_byte() < on.start_byte())
94-
}) =>
90+
if on_node
91+
.is_none_or(|on| ctx.node_under_cursor.end_byte() < on.start_byte()) =>
9592
{
9693
5
9794
}
@@ -110,11 +107,8 @@ impl CompletionScore<'_> {
110107
WrappingClause::Where => 10,
111108
WrappingClause::CheckOrUsingClause => 0,
112109
WrappingClause::Join { on_node }
113-
if on_node.is_some_and(|on| {
114-
ctx.node_under_cursor
115-
.as_ref()
116-
.is_some_and(|n| n.start_byte() > on.end_byte())
117-
}) =>
110+
if on_node
111+
.is_some_and(|on| ctx.node_under_cursor.start_byte() > on.end_byte()) =>
118112
{
119113
// Users will probably join on primary keys
120114
if col.is_primary_key { 20 } else { 10 }

0 commit comments

Comments
 (0)