Shift spans between highlights properly

include/editor.h

@@ -3,9 +3,11 @@
 #include "./rope.h"
 #include "./ui.h"
 #include "./utils.h"
+#include <algorithm>
 #include <cstdint>
 #include <map>
 #include <shared_mutex>
+#include <unordered_map>
 
 struct Highlight {
   uint32_t fg;
@@ -22,37 +24,71 @@ struct Span {
   bool operator<(const Span &other) const { return start < other.start; }
 };
 
+struct Spans {
+  std::vector<Span> spans;
+  Queue<std::pair<uint32_t, int64_t>> edits;
+  bool mid_parse = false;
+  std::shared_mutex mtx;
+
+  void apply(uint32_t x, int64_t y) {
+    std::unique_lock lock(mtx);
+    auto it = std::lower_bound(
+        spans.begin(), spans.end(), Span{.start = x, .end = 0, .hl = nullptr},
+        [](auto &a, auto &b) { return a.start < b.start; });
+    while (it != spans.begin()) {
+      auto prev = std::prev(it);
+      if (prev->end <= x)
+        break;
+      it = prev;
+    }
+    while (it != spans.end()) {
+      if (it->start < x && it->end > x) {
+        it->end += y;
+      } else if (it->start > x) {
+        it->start += y;
+        it->end += y;
+      }
+      if (it->end <= it->start)
+        it = spans.erase(it);
+      else
+        ++it;
+    }
+  }
+};
+
 struct SpanCursor {
-  const std::vector<Span> &spans;
+  Spans &spans;
   size_t index = 0;
   std::vector<Span *> active;
+  std::shared_lock<std::shared_mutex> lock;
 
-  SpanCursor(const std::vector<Span> &s) : spans(s) {}
+  SpanCursor(Spans &s) : spans(s) {}
   Highlight *get_highlight(uint32_t byte_offset) {
     for (int i = (int)active.size() - 1; i >= 0; i--)
       if (active[i]->end <= byte_offset)
         active.erase(active.begin() + i);
-    while (index < spans.size() && spans[index].start <= byte_offset) {
+    while (index < spans.spans.size() &&
-      if (spans[index].end > byte_offset)
+           spans.spans[index].start <= byte_offset) {
-        active.push_back(const_cast<Span *>(&spans[index]));
+      if (spans.spans[index].end > byte_offset)
+        active.push_back(const_cast<Span *>(&spans.spans[index]));
       index++;
     }
     Highlight *best = nullptr;
     int max_prio = -1;
-    for (auto *s : active) {
+    for (auto *s : active)
       if (s->hl->priority > max_prio) {
         max_prio = s->hl->priority;
         best = s->hl;
       }
-    }
     return best;
   }
   void sync(uint32_t byte_offset) {
+    lock = std::shared_lock(spans.mtx);
     active.clear();
-    size_t left = 0, right = spans.size();
+    size_t left = 0, right = spans.spans.size();
     while (left < right) {
       size_t mid = (left + right) / 2;
-      if (spans[mid].start <= byte_offset)
+      if (spans.spans[mid].start <= byte_offset)
         left = mid + 1;
       else
         right = mid;
@@ -60,9 +96,9 @@ struct SpanCursor {
     index = left;
     while (left > 0) {
       left--;
-      if (spans[left].end > byte_offset)
+      if (spans.spans[left].end > byte_offset)
-        active.push_back(const_cast<Span *>(&spans[left]));
+        active.push_back(const_cast<Span *>(&spans.spans[left]));
-      else if (byte_offset - spans[left].end > 1000)
+      else if (byte_offset - spans.spans[left].end > 1000)
         break;
     }
   }
@@ -72,7 +108,6 @@ struct Editor {
   const char *filename;                 // Filename of the editor
   Knot *root;                           // A rope
   std::shared_mutex knot_mtx;           // A mutex
-  std::shared_mutex span_mtx;           // A mutex
   Coord cursor;                         // position of the cursor
   uint32_t cursor_preffered;            // preffered visual column
   Coord selection;                      // position of the selection
@@ -87,7 +122,7 @@ struct Editor {
   Queue<TSInputEdit> edit_queue;        // Tree-sitter edit queue
   std::vector<Highlight> query_map;     // Tree-sitter query map
   std::vector<int8_t> folded;           // folded lines indexed by line number
-  std::vector<Span> spans;              // Highlighted spans
+  Spans spans;                          // Highlighted spans
   std::map<uint32_t, bool> folded_node; // maps content hash to fold state
                                         // - built by tree-sitter helpers
 };

src/editor.cc

@@ -382,7 +382,6 @@ void render_editor(Editor *editor) {
   uint32_t line_index = editor->scroll.row;
   SpanCursor span_cursor(editor->spans);
   std::shared_lock knot_lock(editor->knot_mtx);
-  std::shared_lock span_lock(editor->span_mtx);
   LineIterator *it = begin_l_iter(editor->root, line_index);
   if (!it)
     return;

src/main.cc

@@ -75,6 +75,9 @@ void handle_editor_event(Editor *editor, KeyEvent event) {
       };
       editor->edit_queue.push(edit);
     }
+    editor->spans.apply(pos, 1);
+    if (editor->spans.mid_parse)
+      editor->spans.edits.push({pos, 1});
     cursor_right(editor, 1);
   }
   if (event.key_type == KEY_CHAR && event.c == '\t') {
@@ -96,6 +99,9 @@ void handle_editor_event(Editor *editor, KeyEvent event) {
       };
       editor->edit_queue.push(edit);
     }
+    editor->spans.apply(pos, 2);
+    if (editor->spans.mid_parse)
+      editor->spans.edits.push({pos, 2});
     cursor_right(editor, 2);
   }
   if (event.key_type == KEY_CHAR && (event.c == '\n' || event.c == '\r')) {
@@ -118,6 +124,9 @@ void handle_editor_event(Editor *editor, KeyEvent event) {
       };
       editor->edit_queue.push(edit);
     }
+    editor->spans.apply(pos + 1, 1);
+    if (editor->spans.mid_parse)
+      editor->spans.edits.push({pos + 1, 1});
     cursor_right(editor, 1);
   }
   if (event.key_type == KEY_CHAR && event.c == 0x7F) {
@@ -143,6 +152,9 @@ void handle_editor_event(Editor *editor, KeyEvent event) {
       };
       editor->edit_queue.push(edit);
     }
+    editor->spans.apply(start, start - pos);
+    if (editor->spans.mid_parse)
+      editor->spans.edits.push({start, start - pos});
   }
   ensure_scroll(editor);
 }

src/ts.cc

@@ -1,6 +1,7 @@
 #include "../include/ts.h"
 #include "../include/editor.h"
 #include "../include/rope.h"
+#include <algorithm>
 #include <cstdint>
 #include <fstream>
 #include <regex>
@@ -177,10 +178,10 @@ void ts_collect_spans(Editor *editor) {
       .decode = nullptr,
   };
   TSTree *tree, *copy = nullptr;
-  std::unique_lock lock_0(editor->knot_mtx);
+  std::unique_lock knot_mtx(editor->knot_mtx);
   if (editor->tree)
     copy = ts_tree_copy(editor->tree);
-  lock_0.unlock();
+  knot_mtx.unlock();
   if (!running)
     return;
   std::vector<TSInputEdit> edits;
@@ -189,21 +190,17 @@ void ts_collect_spans(Editor *editor) {
     while (editor->edit_queue.pop(edit)) {
       edits.push_back(edit);
       ts_tree_edit(copy, &edits.back());
-    }
+    };
+  editor->spans.mid_parse = true;
   tree = ts_parser_parse(editor->parser, copy, tsinput);
   if (copy)
     ts_tree_delete(copy);
-  while (editor->edit_queue.pop(edit)) {
+  knot_mtx.lock();
-    edits.push_back(edit);
-    ts_tree_edit(tree, &edits.back());
-  }
-  lock_0.lock();
   if (editor->tree)
     ts_tree_delete(editor->tree);
   editor->tree = tree;
   copy = ts_tree_copy(tree);
-  lock_0.unlock();
+  knot_mtx.unlock();
-  std::shared_lock lock_1(editor->span_mtx);
   TSQueryCursor *cursor = ts_query_cursor_new();
   ts_query_cursor_exec(cursor, editor->query, ts_tree_root_node(copy));
   std::vector<Span> new_spans;
@@ -232,12 +229,16 @@ void ts_collect_spans(Editor *editor) {
       free(load.prev);
     return;
   }
-  lock_1.unlock();
   std::sort(new_spans.begin(), new_spans.end());
-  std::unique_lock lock_2(editor->span_mtx);
+  std::unique_lock span_mtx(editor->spans.mtx);
-  editor->spans.swap(new_spans);
+  editor->spans.mid_parse = false;
-  lock_2.unlock();
+  editor->spans.spans.swap(new_spans);
+  span_mtx.unlock();
+  std::pair<uint32_t, int64_t> span_edit;
+  while (editor->spans.edits.pop(span_edit))
+    editor->spans.apply(span_edit.first, span_edit.second);
   ts_query_cursor_delete(cursor);
+  ts_tree_delete(copy);
   if (load.prev)
     free(load.prev);
 }