From 19e3750f13634d2820f5852fb135c0c0f1e20e0d Mon Sep 17 00:00:00 2001
From: Max Brunsfeld <maxbrunsfeld@gmail.com>
Date: Wed, 9 May 2018 13:46:46 -0700
Subject: [PATCH] Make ts_node_next_sibling work more like ts_node_prev_sibling

Co-Authored-By: Rick Winfrey <rewinfrey@github.com>
---
 src/runtime/node.c | 51 +++++++++++++++++++++++++++++++++-------------
 1 file changed, 37 insertions(+), 14 deletions(-)

diff --git a/src/runtime/node.c b/src/runtime/node.c
index 2230d1af..99b6256d 100644
--- a/src/runtime/node.c
+++ b/src/runtime/node.c
@@ -179,28 +179,51 @@ static inline TSNode ts_node__prev_sibling(TSNode self, bool include_anonymous)
 }
 
 static inline TSNode ts_node__next_sibling(TSNode self, bool include_anonymous) {
-  TSNode node = ts_node_parent(self);
-  if (!node.subtree) return ts_node__null();
-  uint32_t end_byte = ts_node_end_byte(self);
+  uint32_t target_end_byte = ts_node_end_byte(self);
 
-  bool did_descend = true;
-  while (did_descend) {
-    did_descend = false;
+  TSNode node = ts_node_parent(self);
+  TSNode later_node = ts_node__null();
+  bool later_node_is_relevant = false;
+
+  while (node.subtree) {
+    TSNode later_child = ts_node__null();
+    bool later_child_is_relevant = false;
+    TSNode child_containing_target = ts_node__null();
 
     TSNode child;
     NodeChildIterator iterator = ts_node_child_iterator_begin(&node);
     while (ts_node_child_iterator_next(&iterator, &child)) {
-      if (iterator.position.bytes > end_byte && child.subtree != self.subtree) {
-        if (ts_node__is_relevant(child, include_anonymous)) {
-          return child;
-        }
-        if (ts_node__relevant_child_count(child, include_anonymous) > 0) {
-          node = child;
-          did_descend = true;
-          break;
+      if (iterator.position.bytes < target_end_byte) continue;
+      if (child.byte <= self.byte) {
+        if (child.subtree != self.subtree) {
+          child_containing_target = child;
         }
+      } else if (ts_node__is_relevant(child, include_anonymous)) {
+        later_child = child;
+        later_child_is_relevant = true;
+        break;
+      } else if (ts_node__relevant_child_count(child, include_anonymous) > 0) {
+        later_child = child;
+        later_child_is_relevant = false;
+        break;
       }
     }
+
+    if (child_containing_target.subtree) {
+      if (later_child.subtree) {
+        later_node = later_child;
+        later_node_is_relevant = later_child_is_relevant;
+      }
+      node = child_containing_target;
+    } else if (later_child_is_relevant) {
+      return later_child;
+    } else if (later_child.subtree) {
+      node = later_child;
+    } else if (later_node_is_relevant) {
+      return later_node;
+    } else {
+      node = later_node;
+    }
   }
 
   return ts_node__null();