Merge pull request #116 from seart-group/feature/depth

seart-group · Jan 30, 2024 · 9eee806 · 9eee806
2 parents b74cfbc + a2fa8f3
commit 9eee806
Show file tree

Hide file tree

Showing 6 changed files with 51 additions and 12 deletions.
diff --git a/lib/ch_usi_si_seart_treesitter_TreeCursor.cc b/lib/ch_usi_si_seart_treesitter_TreeCursor.cc
@@ -11,6 +11,12 @@ JNIEXPORT void JNICALL Java_ch_usi_si_seart_treesitter_TreeCursor_delete(
   __clearPointer(env, thisObject);
 }
 
+JNIEXPORT jint JNICALL Java_ch_usi_si_seart_treesitter_TreeCursor_getCurrentDepth(
+  JNIEnv* env, jobject thisObject) {
+  TSTreeCursor* treeCursor = (TSTreeCursor*)__getPointer(env, thisObject);
+  return (jint)ts_tree_cursor_current_depth(treeCursor);
+}
+
 JNIEXPORT jobject JNICALL Java_ch_usi_si_seart_treesitter_TreeCursor_getCurrentNode(
   JNIEnv* env, jobject thisObject) {
   TSTreeCursor* cursor = (TSTreeCursor*)__getPointer(env, thisObject);

diff --git a/lib/ch_usi_si_seart_treesitter_TreeCursor.h b/lib/ch_usi_si_seart_treesitter_TreeCursor.h
diff --git a/src/main/java/ch/usi/si/seart/treesitter/OffsetTreeCursor.java b/src/main/java/ch/usi/si/seart/treesitter/OffsetTreeCursor.java
@@ -54,6 +54,11 @@ public void close() {
         cursor.close();
     }
 
+    @Override
+    public int getCurrentDepth() {
+        return cursor.getCurrentDepth();
+    }
+
     @Override
     public String getCurrentFieldName() {
         return cursor.getCurrentFieldName();

diff --git a/src/main/java/ch/usi/si/seart/treesitter/TreeCursor.java b/src/main/java/ch/usi/si/seart/treesitter/TreeCursor.java
@@ -43,6 +43,15 @@ public class TreeCursor extends External implements Cloneable {
     @Override
     protected native void delete();
 
+    /**
+     * Get the depth of the cursor's current {@link Node} relative
+     * to the original {@link Node} that the cursor was constructed from.
+     *
+     * @return the relative depth of the tree cursor's current node
+     * @since 1.11.0
+     */
+    public native int getCurrentDepth();
+
     /**
      * Get the {@link Node} that the cursor is currently pointing to.
      *
@@ -160,6 +169,11 @@ public void preorderTraversal(@NotNull Consumer<Node> callback) {
 
     static class Stub extends TreeCursor {
 
+        @Override
+        public int getCurrentDepth() {
+            throw new UnsupportedOperationException();
+        }
+
         @Override
         public Node getCurrentNode() {
             throw new UnsupportedOperationException();

diff --git a/src/main/java/ch/usi/si/seart/treesitter/printer/SyntaxTreePrinter.java b/src/main/java/ch/usi/si/seart/treesitter/printer/SyntaxTreePrinter.java
@@ -31,27 +31,18 @@ protected String getFileExtension() {
     }
 
     protected void write(Consumer<String> appender) {
-        int depth = 0;
         for (;;) {
             TreeCursorNode cursorNode = cursor.getCurrentTreeCursorNode();
             if (cursorNode.isNamed()) {
+                int depth = cursor.getCurrentDepth();
                 String indent = "  ".repeat(depth);
                 appender.accept(indent);
                 appender.accept(cursorNode.toString());
                 appender.accept("\n");
             }
-            if (cursor.gotoFirstChild()) {
-                depth++;
-                continue;
-            } else if (cursor.gotoNextSibling()) {
-                continue;
-            }
+            if (cursor.gotoFirstChild() || cursor.gotoNextSibling()) continue;
             do {
-                if (cursor.gotoParent()) {
-                    depth--;
-                } else {
-                    return;
-                }
+                if (!cursor.gotoParent()) return;
             } while (!cursor.gotoNextSibling());
         }
     }

diff --git a/src/test/java/ch/usi/si/seart/treesitter/TreeCursorTest.java b/src/test/java/ch/usi/si/seart/treesitter/TreeCursorTest.java
@@ -69,6 +69,21 @@ void testWalk() {
         Assertions.assertTrue(cursor.gotoFirstChild());
     }
 
+    @Test
+    void testGetCurrentDepth() {
+        Assertions.assertEquals(0, cursor.getCurrentDepth());
+        cursor.gotoFirstChild();
+        Assertions.assertEquals(1, cursor.getCurrentDepth());
+        cursor.gotoFirstChild();
+        Assertions.assertEquals(2, cursor.getCurrentDepth());
+        cursor.gotoNextSibling();
+        Assertions.assertEquals(2, cursor.getCurrentDepth());
+        cursor.gotoParent();
+        Assertions.assertEquals(1, cursor.getCurrentDepth());
+        cursor.gotoParent();
+        Assertions.assertEquals(0, cursor.getCurrentDepth());
+    }
+
     @Test
     void testGotoFirstChildByteOffset() {
         cursor.gotoFirstChild(); // function_definition