kbtree: allow iterators to start at arbitrary position

3 files changed, 83 insertions, 23 deletions

diff --git a/src/nvim/buffer.c b/src/nvim/buffer.c
index 8491e77e73..580eb07d6b 100644
--- a/src/nvim/buffer.c
+++ b/src/nvim/buffer.c
@@ -74,6 +74,12 @@
 #include "nvim/os/time.h"
 #include "nvim/os/input.h"
 
+typedef enum {
+  kBLSUnchanged = 0,
+  kBLSChanged = 1,
+  kBLSDeleted = 2,
+} BufhlLineStatus;
+
 #ifdef INCLUDE_GENERATED_DECLARATIONS
 # include "buffer.c.generated.h"
 #endif
@@ -5228,21 +5234,21 @@ void bufhl_clear_line_range(buf_T *buf,
     return;
   }
   linenr_T first_changed = MAXLNUM, last_changed = -1;
-  // TODO: implement kb_itr_interval and jump directly to the first line
+
   kbitr_t itr;
-  BufhlLine *l;
-  for (kb_itr_first(bufhl, buf->b_bufhl_info, &itr);
-       kb_itr_valid(&itr);
-       kb_itr_next(bufhl, buf->b_bufhl_info, &itr)) {
+  BufhlLine *l, t = {line_start};
+  if (!kb_itr_get(bufhl, buf->b_bufhl_info, &t, &itr)) {
+    kb_itr_next(bufhl, buf->b_bufhl_info, &itr);
+  }
+  for (; kb_itr_valid(&itr); kb_itr_next(bufhl, buf->b_bufhl_info, &itr)) {
     l = kb_itr_key(BufhlLine *, &itr);
     linenr_T line = l->line;
-    if (line < line_start) {
-      continue;
-    } else if (line > line_end) {
+    if (line > line_end) {
       break;
     }
     if (line_start <= line && line <= line_end) {
-      if (bufhl_clear_line(buf->b_bufhl_info, src_id, line)) {
+      BufhlLineStatus status = bufhl_clear_line(buf->b_bufhl_info, src_id, line);
+      if (status != kBLSUnchanged) {
         if (line > last_changed) {
           last_changed = line;
         }
@@ -5250,6 +5256,9 @@ void bufhl_clear_line_range(buf_T *buf,
           first_changed = line;
         }
       }
+      if (status == kBLSDeleted) {
+        kb_del_itr(bufhl, buf->b_bufhl_info, &itr);
+      }
     }
   }
 
@@ -5264,8 +5273,8 @@ void bufhl_clear_line_range(buf_T *buf,
 /// @param bufhl_info The highlight info for the buffer
 /// @param src_id Highlight source group to clear, or -1 to clear all groups.
 /// @param lnum Linenr where the highlight should be cleared
-static bool bufhl_clear_line(bufhl_info_T *bufhl_info, int src_id,
-                             linenr_T lnum)
+static BufhlLineStatus bufhl_clear_line(bufhl_info_T *bufhl_info, int src_id,
+                                        linenr_T lnum)
 {
   BufhlLine *lineinfo = bufhl_tree_ref(bufhl_info, lnum, false);
   size_t oldsize = kv_size(lineinfo->items);
@@ -5286,9 +5295,9 @@ static bool bufhl_clear_line(bufhl_info_T *bufhl_info, int src_id,
 
   if (kv_size(lineinfo->items) == 0) {
     kv_destroy(lineinfo->items);
-    kb_del(bufhl, bufhl_info, lineinfo);
+    return kBLSDeleted;
   }
-  return kv_size(lineinfo->items) != oldsize;
+  return kv_size(lineinfo->items) != oldsize ? kBLSChanged : kBLSUnchanged;
 }
 
 /// Remove all highlights and free the highlight data
@@ -5315,15 +5324,19 @@ void bufhl_mark_adjust(buf_T* buf,
   // we need to detect this case and
 
   kbitr_t itr;
-  BufhlLine *l;
-  for (kb_itr_first(bufhl, buf->b_bufhl_info, &itr);
-       kb_itr_valid(&itr);
-       kb_itr_next(bufhl, buf->b_bufhl_info, &itr)) {
+  BufhlLine *l, t = {line1};
+  if (!kb_itr_get(bufhl, buf->b_bufhl_info, &t, &itr)) {
+    kb_itr_next(bufhl, buf->b_bufhl_info, &itr);
+  }
+  for (; kb_itr_valid(&itr); kb_itr_next(bufhl, buf->b_bufhl_info, &itr)) {
     l = kb_itr_key(BufhlLine *, &itr);
     if (l->line >= line1 && l->line <= line2) {
       if (amount == MAXLNUM) {
-        bufhl_clear_line(buf->b_bufhl_info, -1, l->line);
-        continue;
+        if (bufhl_clear_line(buf->b_bufhl_info, -1, l->line) == kBLSDeleted) {
+          kb_del_itr(bufhl, buf->b_bufhl_info, &itr);
+        } else {
+          assert(false);
+        }
       } else {
         l->line += amount;
       }
diff --git a/src/nvim/bufhl_defs.h b/src/nvim/bufhl_defs.h
index 0667a82a62..c8b81e7927 100644
--- a/src/nvim/bufhl_defs.h
+++ b/src/nvim/bufhl_defs.h
@@ -28,7 +28,7 @@ typedef struct {
   colnr_T valid_to;
 } bufhl_lineinfo_T;
 
-#define BUFHL_CMP(a, b) (((b)->line - (a)->line))
+#define BUFHL_CMP(a, b) (((a)->line - (b)->line))
 KBTREE_INIT(bufhl, BufhlLine *, BUFHL_CMP)
 typedef kbtree_t(bufhl) bufhl_info_T;
 #endif  // NVIM_BUFHL_DEFS_H
diff --git a/src/nvim/lib/kbtree.h b/src/nvim/lib/kbtree.h
index b58c3f97a4..ec83d4d124 100644
--- a/src/nvim/lib/kbtree.h
+++ b/src/nvim/lib/kbtree.h
@@ -57,7 +57,9 @@ typedef struct {
 		int	off_key, off_ptr, ilen, elen;		\
 		int	n, t;								\
 		int	n_keys, n_nodes;					\
-	} kbtree_##name##_t;
+	} kbtree_##name##_t; \
+
+
 
 #define __KB_INIT(name, key_t)											\
 	static inline kbtree_##name##_t *kb_init_##name(unsigned int size)							\
@@ -75,7 +77,7 @@ typedef struct {
 		b->root = (kbnode_t*)calloc(1, (unsigned int)b->ilen);						\
 		++b->n_nodes;													\
 		return b;														\
-	}
+	} \
 
 #define __kb_destroy(b) do {											\
 		int i;                                                          \
@@ -346,7 +348,48 @@ typedef struct {
 			if (itr->p < itr->stack) return 0; \
 			if (itr->p->x && itr->p->i < itr->p->x->n) return 1; \
 		} \
-	}
+	} \
+	static inline int kb_itr_prev_##name(kbtree_##name##_t *b, kbitr_t *itr) \
+	{ \
+		if (itr->p < itr->stack) return 0; \
+		for (;;) { \
+			while (itr->p->x && itr->p->i >= 0) { \
+				itr->p[1].x = itr->p->x->is_internal? __KB_PTR(b, itr->p->x)[itr->p->i] : 0; \
+				itr->p[1].i = itr->p[1].x ? itr->p[1].x->n : -1; \
+				++itr->p; \
+			} \
+			--itr->p; \
+			if (itr->p < itr->stack) return 0; \
+			--itr->p->i; \
+			if (itr->p->x && itr->p->i >= 0) return 1; \
+		} \
+	} \
+	static int kb_itr_getp_##name(kbtree_##name##_t *b, const key_t * __restrict k, kbitr_t *itr) \
+	{ \
+		int i, r = 0; \
+		itr->p = itr->stack; \
+		itr->p->x = b->root; \
+		while (itr->p->x) { \
+			i = __kb_getp_aux_##name(itr->p->x, k, &r); \
+			itr->p->i = i; \
+			if (i >= 0 && r == 0) return 1; \
+			++itr->p->i; \
+			itr->p[1].x = itr->p->x->is_internal? __KB_PTR(b, itr->p->x)[i + 1] : 0; \
+			++itr->p; \
+		} \
+		return 0; \
+	} \
+	static int kb_itr_get_##name(kbtree_##name##_t *b, const key_t k, kbitr_t *itr) \
+	{																	\
+		return kb_itr_getp_##name(b,&k,itr); \
+	} \
+	static inline void kb_del_itr_##name(kbtree_##name##_t *b, kbitr_t *itr) \
+	{ \
+		key_t k = kb_itr_key(key_t, itr); \
+		kb_delp_##name(b, &k); \
+		kb_itr_getp_##name(b, &k, itr); \
+	} 
+
 
 #define KBTREE_INIT(name, key_t, __cmp)			\
 	__KB_TREE_T(name)							\
@@ -373,7 +416,11 @@ typedef struct {
 #define kb_intervalp(name, b, k, l, u) kb_intervalp_##name(b, k, l, u)
 
 #define kb_itr_first(name, b, i) kb_itr_first_##name(b, i)
+#define kb_itr_get(name, b, k, i) kb_itr_get_##name(b, k, i)
+#define kb_itr_getp(name, b, k, i) kb_itr_getp_##name(b, k, i)
 #define kb_itr_next(name, b, i) kb_itr_next_##name(b, i)
+#define kb_itr_prev(name, b, i) kb_itr_prev_##name(b, i)
+#define kb_del_itr(name, b, i) kb_del_itr_##name(b, i)
 #define kb_itr_key(type, itr) __KB_KEY(type, (itr)->p->x)[(itr)->p->i]
 #define kb_itr_valid(itr) ((itr)->p >= (itr)->stack)