From 477b6677537e89c7bdff14ce84dad6d23a6415bb Mon Sep 17 00:00:00 2001
From: Max Brunsfeld <maxbrunsfeld@gmail.com>
Date: Thu, 18 Aug 2022 13:48:47 -0700
Subject: [PATCH] Add ts_tree_root_node_with_offset API

---
 cli/src/tests/node_test.rs    | 28 ++++++++++++++++++++++++++++
 lib/binding_rust/bindings.rs  |  9 +++++++++
 lib/binding_rust/lib.rs       | 14 ++++++++++++++
 lib/include/tree_sitter/api.h | 10 ++++++++++
 lib/src/tree.c                | 10 ++++++++++
 5 files changed, 71 insertions(+)

diff --git a/cli/src/tests/node_test.rs b/cli/src/tests/node_test.rs
index 776ca2c7..6d5ed61d 100644
--- a/cli/src/tests/node_test.rs
+++ b/cli/src/tests/node_test.rs
@@ -529,6 +529,34 @@ fn test_node_edit() {
     }
 }
 
+#[test]
+fn test_root_node_with_offset() {
+    let mut parser = Parser::new();
+    parser.set_language(get_language("javascript")).unwrap();
+    let tree = parser.parse("  if (a) b", None).unwrap();
+
+    let node = tree.root_node_with_offset(6, Point::new(2, 2));
+    assert_eq!(node.byte_range(), 8..16);
+    assert_eq!(node.start_position(), Point::new(2, 4));
+    assert_eq!(node.end_position(), Point::new(2, 12));
+
+    let child = node.child(0).unwrap().child(2).unwrap();
+    assert_eq!(child.kind(), "expression_statement");
+    assert_eq!(child.byte_range(), 15..16);
+    assert_eq!(child.start_position(), Point::new(2, 11));
+    assert_eq!(child.end_position(), Point::new(2, 12));
+
+    let mut cursor = node.walk();
+    cursor.goto_first_child();
+    cursor.goto_first_child();
+    cursor.goto_next_sibling();
+    let child = cursor.node();
+    assert_eq!(child.kind(), "parenthesized_expression");
+    assert_eq!(child.byte_range(), 11..14);
+    assert_eq!(child.start_position(), Point::new(2, 7));
+    assert_eq!(child.end_position(), Point::new(2, 10));
+}
+
 #[test]
 fn test_node_is_extra() {
     let mut parser = Parser::new();
diff --git a/lib/binding_rust/bindings.rs b/lib/binding_rust/bindings.rs
index 1447d09d..0266521d 100644
--- a/lib/binding_rust/bindings.rs
+++ b/lib/binding_rust/bindings.rs
@@ -332,6 +332,15 @@ extern "C" {
     #[doc = " Get the root node of the syntax tree."]
     pub fn ts_tree_root_node(self_: *const TSTree) -> TSNode;
 }
+extern "C" {
+    #[doc = " Get the root node of the syntax tree, but with its position"]
+    #[doc = " shifted forward by the given offset."]
+    pub fn ts_tree_root_node_with_offset(
+        self_: *const TSTree,
+        offset_bytes: u32,
+        offset_point: TSPoint,
+    ) -> TSNode;
+}
 extern "C" {
     #[doc = " Get the language that was used to parse the syntax tree."]
     pub fn ts_tree_language(arg1: *const TSTree) -> *const TSLanguage;
diff --git a/lib/binding_rust/lib.rs b/lib/binding_rust/lib.rs
index f757b107..934915fe 100644
--- a/lib/binding_rust/lib.rs
+++ b/lib/binding_rust/lib.rs
@@ -708,6 +708,20 @@ impl Tree {
         Node::new(unsafe { ffi::ts_tree_root_node(self.0.as_ptr()) }).unwrap()
     }
 
+    /// Get the root node of the syntax tree, but with its position shifted
+    /// forward by the given offset.
+    #[doc(alias = "ts_tree_root_node_with_offset")]
+    pub fn root_node_with_offset(&self, offset_bytes: usize, offset_extent: Point) -> Node {
+        Node::new(unsafe {
+            ffi::ts_tree_root_node_with_offset(
+                self.0.as_ptr(),
+                offset_bytes as u32,
+                offset_extent.into(),
+            )
+        })
+        .unwrap()
+    }
+
     /// Get the language that was used to parse the syntax tree.
     #[doc(alias = "ts_tree_language")]
     pub fn language(&self) -> Language {
diff --git a/lib/include/tree_sitter/api.h b/lib/include/tree_sitter/api.h
index e2941532..727dded3 100644
--- a/lib/include/tree_sitter/api.h
+++ b/lib/include/tree_sitter/api.h
@@ -366,6 +366,16 @@ void ts_tree_delete(TSTree *self);
  */
 TSNode ts_tree_root_node(const TSTree *self);
 
+/**
+ * Get the root node of the syntax tree, but with its position
+ * shifted forward by the given offset.
+ */
+TSNode ts_tree_root_node_with_offset(
+  const TSTree *self,
+  uint32_t offset_bytes,
+  TSPoint offset_point
+);
+
 /**
  * Get the language that was used to parse the syntax tree.
  */
diff --git a/lib/src/tree.c b/lib/src/tree.c
index f2cc85ef..103ba84f 100644
--- a/lib/src/tree.c
+++ b/lib/src/tree.c
@@ -1,6 +1,7 @@
 #include "tree_sitter/api.h"
 #include "./array.h"
 #include "./get_changed_ranges.h"
+#include "./length.h"
 #include "./subtree.h"
 #include "./tree_cursor.h"
 #include "./tree.h"
@@ -37,6 +38,15 @@ TSNode ts_tree_root_node(const TSTree *self) {
   return ts_node_new(self, &self->root, ts_subtree_padding(self->root), 0);
 }
 
+TSNode ts_tree_root_node_with_offset(
+  const TSTree *self,
+  uint32_t offset_bytes,
+  TSPoint offset_extent
+) {
+  Length offset = {offset_bytes, offset_extent};
+  return ts_node_new(self, &self->root, length_add(offset, ts_subtree_padding(self->root)), 0);
+}
+
 const TSLanguage *ts_tree_language(const TSTree *self) {
   return self->language;
 }