fix: add missing scope boundaries (#709)

Author: juleswritescode

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);
+        }
+    }
+}

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(&current_node_kind) {
+            return;
+        }
+
         self.scope_tracker.register(current_node, self.position);
 
         // prevent infinite recursion – this can happen with ERROR nodes

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(