Add indexed include definition navigation

Author: Myriad-Dreamin

src/index/merged_index.cpp

@@ -38,6 +38,28 @@ struct DenseMapInfo<clice::index::Occurrence> {
     }
 };
 
+template <>
+struct DenseMapInfo<clice::index::IncludeDefinition> {
+    using R = clice::LocalSourceRange;
+    using V = clice::index::IncludeDefinition;
+
+    inline static V getEmptyKey() {
+        return V(R(-1, 0), 0);
+    }
+
+    inline static V getTombstoneKey() {
+        return V(R(-2, 0), 0);
+    }
+
+    static auto getHashValue(const V& v) {
+        return dense_hash(v.range.begin, v.range.end, v.target_path);
+    }
+
+    static bool isEqual(const V& lhs, const V& rhs) {
+        return lhs.range == rhs.range && lhs.target_path == rhs.target_path;
+    }
+};
+
 template <>
 struct DenseMapInfo<clice::index::Relation> {
     using R = clice::index::Relation;
@@ -134,12 +156,18 @@ struct MergedIndex::Impl {
     /// All merged symbol occurrences.
     llvm::DenseMap<Occurrence, roaring::Roaring> occurrences;
 
+    /// All merged include definitions.
+    llvm::DenseMap<IncludeDefinition, roaring::Roaring> include_definitions;
+
     /// All merged symbol relations.
     llvm::DenseMap<SymbolHash, llvm::DenseMap<Relation, roaring::Roaring>> relations;
 
     /// Sorted occurrences cache for fast lookup.
     std::vector<Occurrence> occurrences_cache;
 
+    /// Sorted include definitions cache for fast lookup.
+    std::vector<IncludeDefinition> include_definitions_cache;
+
     void merge(this Impl& self, std::uint32_t path_id, FileIndex& index, auto&& add_context) {
         auto hash = index.hash();
         auto hash_key = llvm::StringRef(reinterpret_cast<char*>(hash.data()), hash.size());
@@ -158,6 +186,10 @@ struct MergedIndex::Impl {
             self.occurrences[occurrence].add(canonical_id);
         }
 
+        for(auto& include: index.include_definitions) {
+            self.include_definitions[include].add(canonical_id);
+        }
+
         for(auto& [symbol_id, relations]: index.relations) {
             auto& target = self.relations[symbol_id];
             for(auto& relation: relations) {
@@ -245,6 +277,14 @@ void MergedIndex::load_in_memory(this Self& self) {
                                       read_bitmap(entry->context()));
     }
 
+    if(root->include_definitions()) {
+        for(auto entry: *root->include_definitions()) {
+            index.include_definitions.try_emplace(
+                *safe_cast<IncludeDefinition>(entry->include_definition()),
+                read_bitmap(entry->context()));
+        }
+    }
+
     for(auto entry: *root->relations()) {
         auto& relations = index.relations[entry->symbol()];
         for(auto relation_entry: *entry->relations()) {
@@ -329,6 +369,27 @@ void MergedIndex::serialize(this const Self& self, llvm::raw_ostream& out) {
                std::tuple(ro.range.begin, ro.range.end, ro.target);
     });
 
+    llvm::SmallVector<const IncludeDefinition*> include_definition_keys;
+    include_definition_keys.reserve(index->include_definitions.size());
+    auto include_definitions = transform(index->include_definitions, [&](auto&& value) {
+        auto&& [include, bitmap] = value;
+        buffer.clear();
+        buffer.resize_for_overwrite(bitmap.getSizeInBytes(false));
+        bitmap.write(buffer.data(), false);
+        include_definition_keys.emplace_back(&include);
+        return binary::CreateIncludeDefinitionEntry(
+            builder,
+            safe_cast<binary::IncludeDefinition>(&include),
+            CreateVector(builder, buffer));
+    });
+    std::ranges::sort(std::views::zip(include_definition_keys, include_definitions),
+                      [](auto lhs, auto rhs) {
+                          const auto& lo = *std::get<0>(lhs);
+                          const auto& ro = *std::get<0>(rhs);
+                          return std::tuple(lo.range.begin, lo.range.end, lo.target_path) <
+                                 std::tuple(ro.range.begin, ro.range.end, ro.target_path);
+                      });
+
     llvm::SmallVector<std::uint64_t> relation_keys;
     relation_keys.reserve(index->relations.size());
     auto relations = transform(index->relations, [&](auto&& value) {
@@ -369,7 +430,8 @@ void MergedIndex::serialize(this const Self& self, llvm::raw_ostream& out) {
                                                   CreateVector(builder, occurrences),
                                                   CreateVector(builder, relations),
                                                   removed,
-                                                  content_offset);
+                                                  content_offset,
+                                                  CreateVector(builder, include_definitions));
     builder.Finish(merged_index);
 
     out.write(safe_cast<char>(builder.GetBufferPointer()), builder.GetSize());
@@ -443,6 +505,81 @@ void MergedIndex::lookup(this const Self& self,
     }
 }
 
+void MergedIndex::lookup_include_definition(
+    this const Self& self,
+    std::uint32_t offset,
+    llvm::function_ref<bool(const IncludeDefinition&)> callback) {
+    if(self.impl) {
+        auto& index = *self.impl;
+        auto& includes = index.include_definitions_cache;
+        if(includes.empty()) {
+            for(auto& [include, _]: index.include_definitions) {
+                includes.emplace_back(include);
+            }
+            std::ranges::sort(includes,
+                              [](const IncludeDefinition& lhs, const IncludeDefinition& rhs) {
+                                  return std::tuple(lhs.range.begin,
+                                                    lhs.range.end,
+                                                    lhs.target_path) < std::tuple(rhs.range.begin,
+                                                                                  rhs.range.end,
+                                                                                  rhs.target_path);
+                              });
+        }
+
+        auto it = std::ranges::lower_bound(includes, offset, {}, [](const IncludeDefinition& i) {
+            return i.range.end;
+        });
+
+        while(it != includes.end()) {
+            if(it->range.contains(offset)) {
+                if(!index.removed.isEmpty()) {
+                    auto bitmap_it = index.include_definitions.find(*it);
+                    if(bitmap_it != index.include_definitions.end()) {
+                        auto remaining = bitmap_it->second - index.removed;
+                        if(remaining.isEmpty()) {
+                            it++;
+                            continue;
+                        }
+                    }
+                }
+
+                if(!callback(*it)) {
+                    break;
+                }
+
+                it++;
+                continue;
+            }
+
+            break;
+        }
+    } else if(self.buffer) {
+        auto index = fbs::GetRoot<binary::MergedIndex>(self.buffer->getBufferStart());
+        auto include_definitions = index->include_definitions();
+        if(!include_definitions) {
+            return;
+        }
+
+        auto it = std::ranges::lower_bound(*include_definitions, offset, {}, [](auto i) {
+            return i->include_definition()->range().end();
+        });
+
+        while(it != include_definitions->end()) {
+            auto include = safe_cast<IncludeDefinition>((*it)->include_definition());
+            if(include->range.contains(offset)) {
+                if(!callback(*include)) {
+                    break;
+                }
+
+                it++;
+                continue;
+            }
+
+            break;
+        }
+    }
+}
+
 void MergedIndex::lookup(this const Self& self,
                          SymbolHash symbol,
                          RelationKind kind,
@@ -577,6 +714,7 @@ void MergedIndex::remove(this Self& self, std::uint32_t path_id) {
 
     // Invalidate cached occurrences.
     index.occurrences_cache.clear();
+    index.include_definitions_cache.clear();
 }
 
 void MergedIndex::merge(this Self& self,
@@ -594,6 +732,7 @@ void MergedIndex::merge(this Self& self,
         context.include_locations = std::move(include_locations);
     });
     self.impl->occurrences_cache.clear();
+    self.impl->include_definitions_cache.clear();
 }
 
 void MergedIndex::merge(this Self& self,
@@ -610,6 +749,7 @@ void MergedIndex::merge(this Self& self,
         context.includes.emplace_back(include_id, canonical_id);
     });
     self.impl->occurrences_cache.clear();
+    self.impl->include_definitions_cache.clear();
 }
 
 llvm::StringRef MergedIndex::content(this const Self& self) {

src/index/merged_index.h

@@ -48,6 +48,12 @@ class MergedIndex {
                 std::uint32_t offset,
                 llvm::function_ref<bool(const Occurrence&)> callback);
 
+    /// Lookup the include definition at the corresponding offset.
+    void lookup_include_definition(
+        this const Self& self,
+        std::uint32_t offset,
+        llvm::function_ref<bool(const IncludeDefinition&)> callback);
+
     /// Lookup the relations of given symbol.
     void lookup(this const Self& self,
                 SymbolHash symbol,

src/index/schema.fbs

@@ -10,6 +10,11 @@ struct Occurrence {
     target : ulong;
 }
 
+struct IncludeDefinition {
+    range : Range;
+    target_path : uint;
+}
+
 struct Relation {
     kind : uint;
     padding : uint;
@@ -64,6 +69,13 @@ context:
     [ubyte];
 }
 
+table IncludeDefinitionEntry {
+include_definition:
+    IncludeDefinition;
+context:
+    [ubyte];
+}
+
 table RelationEntry {
 relation:
     Relation;
@@ -118,6 +130,9 @@ removed:
 
 content:
     string;
+
+include_definitions:
+    [IncludeDefinitionEntry];
 }
 
 table TUFileRelationsEntry {
@@ -134,6 +149,8 @@ occurrences:
     [Occurrence];
 relations:
     [TUFileRelationsEntry];
+include_definitions:
+    [IncludeDefinition];
 }
 
 table TUIndex {