diff --git a/crates/pgls_treesitter/src/context/ancestors.rs b/crates/pgls_treesitter/src/context/ancestors.rs index a99f1054d..67fc70be4 100644 --- a/crates/pgls_treesitter/src/context/ancestors.rs +++ b/crates/pgls_treesitter/src/context/ancestors.rs @@ -3,7 +3,15 @@ pub struct Scope { pub ancestors: AncestorTracker, } -static SCOPE_BOUNDARIES: &[&str] = &["statement", "ERROR", "program"]; +static SCOPE_BOUNDARIES: &[&str] = &[ + "statement", + "ERROR", + "program", + "block", + "transaction", + "psql_meta_command", + "copy_data_stream", +]; #[derive(Debug)] pub struct ScopeTracker { @@ -22,7 +30,7 @@ impl ScopeTracker { self.scopes .last_mut() - .unwrap_or_else(|| panic!("Unhandled node kind: {}", node.kind())) + .unwrap_or_else(|| panic!("No top-level grammar-rule found. Please create an issue with the entire Postgres file, noting cursor/hover position.")) .ancestors .register(node, position); } @@ -105,3 +113,104 @@ impl AncestorTracker { true } } + +#[cfg(test)] +mod tests { + use crate::context::{TreeSitterContextParams, TreesitterContext}; + + fn get_tree(input: &str) -> tree_sitter::Tree { + let mut parser = tree_sitter::Parser::new(); + parser + .set_language(&pgls_treesitter_grammar::LANGUAGE.into()) + .expect("Couldn't set language"); + + parser.parse(input, None).expect("Unable to parse tree") + } + + fn assert_no_panic_for_all_positions(sql: &str) { + let tree = get_tree(sql); + for pos in 0..sql.len() { + let params = TreeSitterContextParams { + position: (pos as u32).into(), + text: sql, + tree: &tree, + }; + let _ = TreesitterContext::new(params); + } + } + + #[test] + fn scope_boundary_block() { + assert_no_panic_for_all_positions("BEGIN; SELECT 1; END;"); + } + + #[test] + fn scope_boundary_transaction() { + assert_no_panic_for_all_positions("BEGIN TRANSACTION; SELECT 1; COMMIT;"); + assert_no_panic_for_all_positions("BEGIN; INSERT INTO t VALUES (1); ROLLBACK;"); + } + + #[test] + fn scope_boundary_psql_meta_command() { + assert_no_panic_for_all_positions("\\dt\n\\d users"); + } + + #[test] + fn scope_boundary_copy_data_stream() { + assert_no_panic_for_all_positions("COPY t FROM STDIN;\n1\tAlice\n\\.\n"); + } + + #[test] + fn scope_boundary_comment() { + assert_no_panic_for_all_positions("-- a comment\nSELECT 1;"); + } + + #[test] + fn issue_704_regression() { + let statements = vec![ + r#" + CREATE OR REPLACE FUNCTION my_schema.my_function1( + pi_1 character varying, + pi_2 character varying, + pi_3 jsonb, + OUT po_1 integer, + OUT po_2 integer, + OUT result integer + ) + RETURNS record + LANGUAGE plpgsql + AS $function$ + "# + .trim(), + + r#" + CREATE OR REPLACE FUNCTION my_schema.my_function2( + pi_1 character varying, + pi_2 character varying, + pi_3 jsonb, + OUT po_1 integer, + OUT po_2 integer, + OUT result integer + ) + RETURNS record + LANGUAGE plpgsql + AS $function$ + DECLARE + BEGIN + -- Function logic goes here + -- For example, you can perform some operations using the input parameters and set the output parameters accordingly + + -- Example logic (replace with actual implementation): + po_1 := length(pi_1); -- Set po_1 to the length of pi_1 + po_2 := length(pi_2); -- Set po_2 to the length of pi_2 + result := po_1 + po_2; -- Set result to the sum of po_1 and po_2 + END; + $function$; + "#.trim(), + ]; + + for stmt in statements { + assert_no_panic_for_all_positions(stmt); + } + } +} diff --git a/crates/pgls_treesitter/src/context/mod.rs b/crates/pgls_treesitter/src/context/mod.rs index 54e80f15d..8819ca78f 100644 --- a/crates/pgls_treesitter/src/context/mod.rs +++ b/crates/pgls_treesitter/src/context/mod.rs @@ -305,6 +305,10 @@ impl<'a> TreesitterContext<'a> { let parent_node_kind = parent_node.kind(); let current_node_kind = current_node.kind(); + if ["comment", "marginalia"].contains(¤t_node_kind) { + return; + } + self.scope_tracker.register(current_node, self.position); // prevent infinite recursion – this can happen with ERROR nodes diff --git a/crates/pgls_treesitter_grammar/grammar.js b/crates/pgls_treesitter_grammar/grammar.js index 896cf6b9a..3a93c597d 100644 --- a/crates/pgls_treesitter_grammar/grammar.js +++ b/crates/pgls_treesitter_grammar/grammar.js @@ -84,6 +84,8 @@ module.exports = grammar({ rules: { program: ($) => choice( + // NOTE: if you add a new top-level statement, make sure to define it + // as a top-level boundary in treesitter context seq( repeat( choice(