From 5ab640147872d846d302744c17699a4e54fb426d Mon Sep 17 00:00:00 2001
From: Max Brunsfeld <maxbrunsfeld@gmail.com>
Date: Thu, 28 Jun 2018 16:17:16 -0700
Subject: [PATCH] Fix TreeCursor bugs

* ts_tree_cursor_goto_first_child_for_byte failed to find the child
  in some cases
* ts_tree_cursor_goto_parent did not handle aliased invisible parent
  nodes

Co-Authored-By: Ashi Krishnan <queerviolet@github.com>
---
 src/runtime/tree_cursor.c | 16 ++++++++-
 test/runtime/node_test.cc | 68 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 83 insertions(+), 1 deletion(-)

diff --git a/src/runtime/tree_cursor.c b/src/runtime/tree_cursor.c
index 83a32927..467681ff 100644
--- a/src/runtime/tree_cursor.c
+++ b/src/runtime/tree_cursor.c
@@ -153,6 +153,11 @@ int64_t ts_tree_cursor_goto_first_child_for_byte(TSTreeCursor *_self, uint32_t g
     }
   } while (did_descend);
 
+  if (self->stack.size > initial_size &&
+      ts_tree_cursor_goto_next_sibling((TSTreeCursor *)self)) {
+    return visible_child_index;
+  }
+
   self->stack.size = initial_size;
   return -1;
 }
@@ -194,7 +199,16 @@ bool ts_tree_cursor_goto_parent(TSTreeCursor *_self) {
   TreeCursor *self = (TreeCursor *)_self;
   for (unsigned i = self->stack.size - 2; i + 1 > 0; i--) {
     TreeCursorEntry *entry = &self->stack.contents[i];
-    if (entry->subtree->visible) {
+    bool is_aliased = false;
+    if (i > 0) {
+      TreeCursorEntry *parent_entry = &self->stack.contents[i - 1];
+      const TSSymbol *alias_sequence = ts_language_alias_sequence(
+        self->tree->language,
+        parent_entry->subtree->alias_sequence_id
+      );
+      is_aliased = alias_sequence && alias_sequence[entry->structural_child_index];
+    }
+    if (entry->subtree->visible || is_aliased) {
       self->stack.size = i + 1;
       return true;
     }
diff --git a/test/runtime/node_test.cc b/test/runtime/node_test.cc
index 1c59e60b..54f21c5f 100644
--- a/test/runtime/node_test.cc
+++ b/test/runtime/node_test.cc
@@ -757,6 +757,74 @@ describe("TreeCursor", [&]() {
     AssertThat(ts_tree_cursor_goto_parent(&cursor), IsTrue());
     AssertThat(ts_tree_cursor_goto_first_child_for_byte(&cursor, 0), Equals(0));
   });
+
+  it("walks the tree correctly when there are hidden leaf nodes", [&]() {
+    ts_parser_set_language(parser, load_real_language("javascript"));
+    ts_tree_cursor_delete(&cursor);
+    ts_tree_delete(tree);
+
+    tree = ts_parser_parse_string(parser, nullptr, "`abc${1}def${2}ghi`", 19);
+    cursor = ts_tree_cursor_new(ts_tree_root_node(tree));
+    TSNode node = ts_tree_cursor_current_node(&cursor);
+    AssertThat(ts_node_type(node), Equals("program"));
+
+    ts_tree_cursor_goto_first_child(&cursor);
+    node = ts_tree_cursor_current_node(&cursor);
+    AssertThat(ts_node_type(node), Equals("expression_statement"));
+
+    ts_tree_cursor_goto_first_child(&cursor);
+    node = ts_tree_cursor_current_node(&cursor);
+    AssertThat(ts_node_type(node), Equals("template_string"));
+
+    int index = ts_tree_cursor_goto_first_child_for_byte(&cursor, 9);
+    AssertThat(index, Equals(2));
+    node = ts_tree_cursor_current_node(&cursor);
+    AssertThat(ts_node_type(node), Equals("template_substitution"));
+    AssertThat(ts_node_start_byte(node), Equals(11u));
+    AssertThat(ts_node_end_byte(node), Equals(15u));
+
+    index = ts_tree_cursor_goto_first_child_for_byte(&cursor, 20);
+    AssertThat(index, Equals(-1));
+    node = ts_tree_cursor_current_node(&cursor);
+    AssertThat(ts_node_type(node), Equals("template_substitution"));
+
+    ts_tree_cursor_goto_first_child(&cursor);
+    node = ts_tree_cursor_current_node(&cursor);
+    AssertThat(ts_node_type(node), Equals("${"));
+
+    index = ts_tree_cursor_goto_first_child_for_byte(&cursor, 20);
+    AssertThat(index, Equals(-1));
+    node = ts_tree_cursor_current_node(&cursor);
+    AssertThat(ts_node_type(node), Equals("${"));
+  });
+
+  it("handles parent nodes that are aliased", [&]() {
+    ts_parser_set_language(parser, load_real_language("html"));
+    ts_tree_cursor_delete(&cursor);
+    ts_tree_delete(tree);
+
+    tree = ts_parser_parse_string(parser, nullptr, "<script></script>", 18);
+
+    cursor = ts_tree_cursor_new(ts_tree_root_node(tree));
+    TSNode node = ts_tree_cursor_current_node(&cursor);
+    AssertThat(ts_node_type(node), Equals("fragment"));
+
+    ts_tree_cursor_goto_first_child(&cursor);
+    node = ts_tree_cursor_current_node(&cursor);
+    AssertThat(ts_node_type(node), Equals("raw_element"));
+
+    ts_tree_cursor_goto_first_child(&cursor);
+    node = ts_tree_cursor_current_node(&cursor);
+    AssertThat(ts_node_type(node), Equals("start_tag"));
+
+    ts_tree_cursor_goto_first_child(&cursor);
+    node = ts_tree_cursor_current_node(&cursor);
+    AssertThat(ts_node_type(node), Equals("<"));
+
+    ts_tree_cursor_goto_parent(&cursor);
+    node = ts_tree_cursor_current_node(&cursor);
+    AssertThat(ts_node_type(node), Equals("start_tag"));
+  });
 });
 
 END_TEST