From a0c5990ddef04ea9ad00bc8ae6833e403cfb7b54 Mon Sep 17 00:00:00 2001
From: Kovid Goyal <kovid@kovidgoyal.net>
Date: Tue, 18 May 2021 20:25:03 +0530
Subject: [PATCH] Annotate search results with their toc nodes

---
 src/pyj/read_book/search_worker.pyj | 49 ++++++++++++++++++++++++++---
 1 file changed, 44 insertions(+), 5 deletions(-)

diff --git a/src/pyj/read_book/search_worker.pyj b/src/pyj/read_book/search_worker.pyj
index 86cc20d740..2b2bbf45ab 100644
--- a/src/pyj/read_book/search_worker.pyj
+++ b/src/pyj/read_book/search_worker.pyj
@@ -70,7 +70,7 @@ class ToCOffsetMap:
         if not matches and self.previous_toc_node:
             matches.push(self.previous_toc_node)
         ans = v'[]'
-        if matches:
+        if matches.length:
             ancestors = v'[]'
             node = matches[-1]
             parent = self.parent_map[node.id]
@@ -80,11 +80,48 @@ class ToCOffsetMap:
             if len(ancestors) > 1:
                 ancestors.pop()  # root node
                 for v'var i = ancestors.length; i-- > 0;':
-                    ans.push(ancestors[i])
-            ans.push(node)
+                    ans.push(ancestors[i].id)
+            ans.push(node.id)
         return ans
 
 
+def toc_offset_map_for_name(name):
+    if wc.toc_offset_map_cache[name]:
+        return wc.toc_offset_map_cache[name]
+    anchor_map = wc.text_cache[name][1]
+    toc_data = wc.current_book.toc_data
+    idx = toc_data.spine_idx_map[name]
+    toc_nodes = toc_data.spine_toc_map[name]
+    if not idx? or idx < 0:
+        wc.toc_offset_map_cache[name] = ToCOffsetMap()
+        return wc.toc_offset_map_cache[name]
+    offset_map = {}
+    for node in toc_nodes:
+        node_id = node.id
+        if node_id?:
+            aid = node.frag
+            offset = anchor_map[aid] or 0
+            offset_map[node_id] = offset
+    prev_toc_node = None
+    for v'var i = idx; i-- > 0;':
+        spine_name = wc.current_book.spine[i]
+        ptn = toc_data.spine_toc_map[spine_name]
+        if ptn?:
+            prev_toc_node = ptn[-1]
+            break
+    wc.toc_offset_map_cache[name] = ToCOffsetMap(toc_nodes, offset_map, prev_toc_node, toc_data.parent_map)
+    return wc.toc_offset_map_cache[name]
+
+
+def toc_nodes_for_search_result(sr):
+    sidx = sr.spine_idx
+    name = wc.current_book.spine[sidx]
+    if not name:
+        return v'[]'
+    tmap = toc_offset_map_for_name(name)
+    return tmap.toc_nodes_for_offset(sr.offset)
+
+
 class Worker:
 
     def __init__(self):
@@ -95,9 +132,9 @@ class Worker:
         self.current_query = None
         self.current_query_id = None
         self.current_book = None
-        self.text_cache = {}
         self.regex = None
         self.result_num = 0
+        self.clear_caches()
 
     @property
     def initialize_error_msg(self):
@@ -105,6 +142,7 @@ class Worker:
 
     def clear_caches(self):
         self.text_cache = {}
+        self.toc_offset_map_cache = {}
 
 
 wc = Worker()
@@ -134,10 +172,11 @@ def search_in_text_of(name):
         match_counts[q] = match_counts[q] or 0
         wc.result_num += 1
         result = {
-            'file_name': name, 'spine_idx': spine_idx, 'index': match_counts[q],
+            'file_name': name, 'spine_idx': spine_idx, 'index': match_counts[q], 'offset': start,
             'text': text, 'before': before, 'after': after, 'mode': wc.current_query.query.mode,
             'q': q, 'result_num': wc.result_num, 'on_discovery': wc.current_query_id,
         }
+        result.toc_nodes = toc_nodes_for_search_result(result)
         self.postMessage({'type': 'search_result', 'id': wc.current_query_id, 'result': result})
         match_counts[q] += 1