refactor: Use suffixes for forward decl resolution (#349)

In a cross-repo setting, the prefix for a forward declaration cannot be
determined eagerly due to lack of package ID information, so we only
serialize the suffix, and perform lookups based on the suffix.

Author: varungandhi-src

docs/Design.md

@@ -730,7 +730,7 @@ There are two options for doing this:
    and use the suffix as the lookup key.
 
 Option 1 would increase the overhead of (de)serializing indexes,
-and requires us to fiddle with the Protobuf format
+and requires us to fiddle with the Index Protobuf schema
 (to add an extra field which tracks the prefix-free symbol name).
 
 So we go with option 2 instead. Specifically, the descriptor
@@ -739,4 +739,14 @@ add it to the end of the version, either is OK).
 We would forbid the use of `$` in package names and versions,
 allowing us to quickly split a symbol name.
 
-The extra `$` could be stripped when serializing the final index.
+<!-- See NOTE(ref: symbol-string-hack-for-forward-decls) -->
+
+The extra `$` could be stripped when serializing the final index.
+
+Additionally, since Occurrences track the symbol name,
+we separate out occurrences for forward declarations
+(see [fwd_decls.proto](/proto/fwd_decls.proto))
+into a separate file with a different format,
+to avoid having to rewrite the symbol string for a large
+number of occurrences.
+Instead, those occurrences are written out at the end.

indexer/Driver.cc

@@ -1087,7 +1087,7 @@ class Driver {
           this->options.compdbPath.asStringRef());
     }
 
-    auto symbolToInfoMap = builder.populateSymbolToInfoMap();
+    auto forwardDeclResolver = builder.populateForwardDeclResolver();
 
     for (auto &paths : this->shardPaths) {
       scip::ForwardDeclIndex indexShard;
@@ -1097,7 +1097,7 @@ class Driver {
       TRACE_EVENT(tracing::indexMerging, "addForwardDeclarations", "size",
                   indexShard.forward_decls_size());
       for (auto &forwardDeclSym : *indexShard.mutable_forward_decls()) {
-        builder.addForwardDeclaration(*symbolToInfoMap,
+        builder.addForwardDeclaration(*forwardDeclResolver,
                                       std::move(forwardDeclSym));
       }
     }

indexer/Indexer.cc

@@ -353,20 +353,23 @@ void ForwardDeclMap::emit(bool deterministic, scip::ForwardDeclIndex &index) {
               this->map.size());
   scip_clang::extractTransform(
       std::move(this->map), deterministic,
-      absl::FunctionRef<void(SymbolNameRef &&, Value &&)>(
-          [&](auto &&symbol, auto &&value) {
-            scip::ForwardDecl fwdDecl{};
-            fwdDecl.set_symbol(symbol.value.data(), symbol.value.size());
-            value.docComment.addTo(*fwdDecl.mutable_documentation());
-            for (auto [path, range] : value.ranges) {
-              scip::ForwardDecl::Reference ref{};
-              range.addTo(ref);
-              auto p = path.asStringView();
-              ref.set_relative_path(p.data(), p.size());
-              *fwdDecl.add_references() = std::move(ref);
-            }
-            *index.add_forward_decls() = std::move(fwdDecl);
-          }));
+      absl::FunctionRef<void(SymbolNameRef &&, Value &&)>([&](auto &&symbol,
+                                                              auto &&value) {
+        scip::ForwardDecl fwdDecl{};
+        auto optSuffix = symbol.getPackageAgnosticSuffix();
+        ENFORCE(optSuffix.has_value(), "missing $ in symbol name {}",
+                symbol.value);
+        fwdDecl.set_suffix(optSuffix->value.data(), optSuffix->value.size());
+        value.docComment.addTo(*fwdDecl.mutable_documentation());
+        for (auto [path, range] : value.ranges) {
+          scip::ForwardDecl::Reference ref{};
+          range.addTo(ref);
+          auto p = path.asStringView();
+          ref.set_relative_path(p.data(), p.size());
+          *fwdDecl.add_references() = std::move(ref);
+        }
+        *index.add_forward_decls() = std::move(fwdDecl);
+      }));
 }
 
 TuIndexer::TuIndexer(const clang::SourceManager &sourceManager,

indexer/ScipExtras.cc

@@ -100,10 +100,53 @@ void SymbolInformationBuilder::finish(bool deterministic,
       }));
 }
 
+void ForwardDeclResolver::insert(SymbolSuffix suffix,
+                                 SymbolInformationBuilder *builder) {
+  this->docInternalMap.emplace(suffix, SymbolInfoOrBuilderPtr{builder});
+}
+
+void ForwardDeclResolver::insert(SymbolSuffix suffix,
+                                 SymbolInformation *symbolInfo) {
+  this->docInternalMap.emplace(suffix, SymbolInfoOrBuilderPtr{symbolInfo});
+}
+
+void ForwardDeclResolver::insertExternal(SymbolNameRef symbol) {
+  if (auto optSuffix = symbol.getPackageAgnosticSuffix()) {
+    this->externalsMap[*optSuffix].insert(symbol);
+  }
+}
+
+std::optional<ForwardDeclResolver::SymbolInfoOrBuilderPtr>
+ForwardDeclResolver::lookupInDocuments(SymbolSuffix suffix) const {
+  auto it = this->docInternalMap.find(suffix);
+  if (it == this->docInternalMap.end()) {
+    return {};
+  }
+  return it->second;
+}
+
+const absl::flat_hash_set<SymbolNameRef> *
+ForwardDeclResolver::lookupExternals(SymbolSuffix suffix) const {
+  auto it = this->externalsMap.find(suffix);
+  if (it == this->externalsMap.end()) {
+    return {};
+  }
+  return &it->second;
+}
+
+void ForwardDeclResolver::deleteExternals(SymbolSuffix suffix) {
+  this->externalsMap.erase(suffix);
+}
+
 SymbolNameRef SymbolNameInterner::intern(std::string &&s) {
   return SymbolNameRef{std::string_view(this->impl.save(s))};
 }
 
+SymbolNameRef SymbolNameInterner::intern(SymbolNameRef s) {
+  return SymbolNameRef{
+      std::string_view(this->impl.save(llvm::StringRef(s.value)))};
+}
+
 DocumentBuilder::DocumentBuilder(scip::Document &&first,
                                  SymbolNameInterner interner)
     : soFar(), interner(interner),
@@ -144,11 +187,12 @@ void DocumentBuilder::merge(scip::Document &&doc) {
   }
 }
 
-void DocumentBuilder::populateSymbolToInfoMap(
-    SymbolToInfoMap &symbolToInfoMap) {
+void DocumentBuilder::populateForwardDeclResolver(
+    ForwardDeclResolver &forwardDeclResolver) {
   for (auto &[symbolName, symbolInfoBuilder] : this->symbolInfos) {
-    symbolToInfoMap.emplace(symbolName,
-                            SymbolToInfoMap::mapped_type(&symbolInfoBuilder));
+    if (auto optSuffix = symbolName.getPackageAgnosticSuffix()) {
+      forwardDeclResolver.insert(*optSuffix, &symbolInfoBuilder);
+    }
   }
 }
 
@@ -251,60 +295,92 @@ void IndexBuilder::addExternalSymbol(scip::SymbolInformation &&extSym) {
   builder->mergeRelationships(std::move(*extSym.mutable_relationships()));
 }
 
-std::unique_ptr<SymbolToInfoMap> IndexBuilder::populateSymbolToInfoMap() {
+std::unique_ptr<ForwardDeclResolver>
+IndexBuilder::populateForwardDeclResolver() {
   TRACE_EVENT(scip_clang::tracing::indexMerging,
-              "IndexBuilder::populateSymbolToInfoMap", "documents.size",
+              "IndexBuilder::populateForwardDeclResolver", "documents.size",
               this->documents.size(), "multiplyIndexed.size",
               this->multiplyIndexed.size());
-  SymbolToInfoMap symbolToInfoMap;
+  ForwardDeclResolver forwardDeclResolver;
   for (auto &document : this->documents) {
     for (auto &symbolInfo : *document.mutable_symbols()) {
-      symbolToInfoMap.emplace(SymbolNameRef{symbolInfo.symbol()},
-                              SymbolToInfoMap::mapped_type(&symbolInfo));
+      if (auto optSuffix =
+              SymbolNameRef{symbolInfo.symbol()}.getPackageAgnosticSuffix()) {
+        forwardDeclResolver.insert(*optSuffix, &symbolInfo);
+      }
     }
   }
   for (auto &[_, docBuilder] : this->multiplyIndexed) {
-    docBuilder->populateSymbolToInfoMap(symbolToInfoMap);
+    docBuilder->populateForwardDeclResolver(forwardDeclResolver);
   }
-  return std::make_unique<SymbolToInfoMap>(std::move(symbolToInfoMap));
+  return std::make_unique<ForwardDeclResolver>(std::move(forwardDeclResolver));
 }
 
-void IndexBuilder::addForwardDeclaration(const SymbolToInfoMap &symbolToInfoMap,
-                                         scip::ForwardDecl &&forwardDeclSym) {
-  auto name =
-      this->interner.intern(std::move(*forwardDeclSym.mutable_symbol()));
-  auto it = symbolToInfoMap.find(name);
-  if (it == symbolToInfoMap.end()) {
-    if (!this->externalSymbols.contains(name)) {
-      scip::SymbolInformation extSym{};
-      *extSym.add_documentation() = std::move(forwardDeclSym.documentation());
-      this->addExternalSymbolUnchecked(name, std::move(extSym));
-    } else {
+void IndexBuilder::addForwardDeclaration(
+    ForwardDeclResolver &forwardDeclResolver,
+    scip::ForwardDecl &&forwardDeclSym) {
+  auto suffix = SymbolSuffix{
+      this->interner.intern(std::move(*forwardDeclSym.mutable_suffix())).value};
+  auto optSymbolInfoOrBuilderPtr =
+      forwardDeclResolver.lookupInDocuments(suffix);
+  if (!optSymbolInfoOrBuilderPtr.has_value()) {
+    if (auto *externals = forwardDeclResolver.lookupExternals(suffix)) {
       // The main index confirms that this was an external symbol, which means
       // that we must have seen the definition in an out-of-project file.
       // In this case, just throw away the information we have,
       // and rely on what we found externally as the source of truth.
+      ENFORCE(!externals->empty(),
+              "externals list for a suffix should be non-empty");
+      // TODO: Log a debug warning if externals->size() > 1
+      for (auto symbolName : *externals) {
+        this->addForwardDeclOccurrences(symbolName,
+                                        scip::ForwardDecl{forwardDeclSym});
+      }
+    } else {
+      scip::SymbolInformation extSym{};
+      // Make up a fake prefix, as we have no package information 🤷🏽
+      auto name = this->interner.intern(
+          std::move(suffix.addFakePrefix().asStringRefMut()));
+      *extSym.add_documentation() = std::move(forwardDeclSym.documentation());
+      this->addExternalSymbolUnchecked(name, std::move(extSym));
+      this->addForwardDeclOccurrences(name, std::move(forwardDeclSym));
     }
-    this->addForwardDeclOccurrences(name, std::move(forwardDeclSym));
     return;
   }
-  auto extIt = this->externalSymbols.find(name);
-  if (extIt != this->externalSymbols.end()) {
+  auto symbolInfoOrBuilderPtr = optSymbolInfoOrBuilderPtr.value();
+  if (auto *externals = forwardDeclResolver.lookupExternals(suffix)) {
     // We found the symbol in a document, so the external symbols list is
     // too pessimistic. This can happen when a TU processes a decl only via
     // a forward decl (and hence conservatively assumes it must be external),
     // but another in-project TU contains the definition.
     //
     // So remove the entry from the external symbols list.
-    extIt->second->discard();
-    this->externalSymbols.erase(extIt);
+    for (auto name : *externals) {
+      auto it = this->externalSymbols.find(name);
+      if (it != this->externalSymbols.end()) {
+        it->second->discard();
+        this->externalSymbols.erase(it);
+      }
+    }
+    forwardDeclResolver.deleteExternals(suffix);
+  }
+  SymbolNameRef name;
+  if (auto *symbolInfo =
+          symbolInfoOrBuilderPtr.dyn_cast<scip::SymbolInformation *>()) {
+    name = SymbolNameRef{std::string_view(symbolInfo->symbol())};
+  } else {
+    auto *symbolInfoBuilder =
+        symbolInfoOrBuilderPtr.get<SymbolInformationBuilder *>();
+    name = symbolInfoBuilder->name;
   }
+  name = this->interner.intern(name);
   if (!forwardDeclSym.documentation().empty()) {
     // FIXME(def: better-doc-merging): The missing documentation placeholder
     // value is due to a bug in the backend which seemingly came up randomly
     // and also got fixed somehow. We should remove the workaround if this
     // is no longer an issue.
-    if (auto *symbolInfo = it->second.dyn_cast<scip::SymbolInformation *>()) {
+    if (auto *symbolInfo =
+            symbolInfoOrBuilderPtr.dyn_cast<scip::SymbolInformation *>()) {
       bool isMissingDocumentation =
           symbolInfo->documentation_size() == 0
           || (symbolInfo->documentation_size() == 1
@@ -316,7 +392,8 @@ void IndexBuilder::addForwardDeclaration(const SymbolToInfoMap &symbolToInfoMap,
             std::move(*forwardDeclSym.mutable_documentation());
       }
     } else {
-      auto &symbolInfoBuilder = *it->second.get<SymbolInformationBuilder *>();
+      auto &symbolInfoBuilder =
+          *symbolInfoOrBuilderPtr.get<SymbolInformationBuilder *>();
       // FIXME(def: better-doc-merging): We shouldn't drop
       // documentation attached to a forward declaration.
       if (!symbolInfoBuilder.hasDocumentation()) {

indexer/ScipExtras.h

@@ -92,11 +92,14 @@ class SymbolInformationBuilder final {
   scip_clang::Bomb _bomb;
 
 public:
+  SymbolNameRef name;
+
   template <typename C1, typename C2>
   SymbolInformationBuilder(SymbolNameRef name, C1 &&docs, C2 &&rels)
       : documentation(), relationships(),
         _bomb(BOMB_INIT(
-            fmt::format("SymbolInformationBuilder for '{}'", name.value))) {
+            fmt::format("SymbolInformationBuilder for '{}'", name.value))),
+        name(name) {
     (void)name;
     this->setDocumentation(std::move(docs));
     this->mergeRelationships(std::move(rels));
@@ -120,9 +123,28 @@ class SymbolInformationBuilder final {
   void finish(bool deterministic, scip::SymbolInformation &out);
 };
 
-using SymbolToInfoMap = absl::flat_hash_map<
-    SymbolNameRef,
-    llvm::PointerUnion<SymbolInformation *, SymbolInformationBuilder *>>;
+class ForwardDeclResolver {
+  using SymbolInfoOrBuilderPtr =
+      llvm::PointerUnion<SymbolInformation *, SymbolInformationBuilder *>;
+
+  absl::flat_hash_map<SymbolSuffix, SymbolInfoOrBuilderPtr> docInternalMap;
+
+  absl::flat_hash_map<SymbolSuffix, absl::flat_hash_set<SymbolNameRef>>
+      externalsMap;
+
+public:
+  ForwardDeclResolver() = default;
+
+  void insert(SymbolSuffix, SymbolInformationBuilder *);
+  void insert(SymbolSuffix, SymbolInformation *);
+  void insertExternal(SymbolNameRef);
+
+  std::optional<SymbolInfoOrBuilderPtr> lookupInDocuments(SymbolSuffix) const;
+
+  const absl::flat_hash_set<SymbolNameRef> *lookupExternals(SymbolSuffix) const;
+
+  void deleteExternals(SymbolSuffix);
+};
 
 class SymbolNameInterner {
   llvm::UniqueStringSaver &impl;
@@ -131,6 +153,7 @@ class SymbolNameInterner {
   SymbolNameInterner(llvm::UniqueStringSaver &impl) : impl(impl) {}
 
   SymbolNameRef intern(std::string &&);
+  SymbolNameRef intern(SymbolNameRef);
 };
 
 class DocumentBuilder final {
@@ -147,7 +170,7 @@ class DocumentBuilder final {
 public:
   DocumentBuilder(scip::Document &&document, SymbolNameInterner interner);
   void merge(scip::Document &&doc);
-  void populateSymbolToInfoMap(SymbolToInfoMap &);
+  void populateForwardDeclResolver(ForwardDeclResolver &);
   void finish(bool deterministic, scip::Document &out);
 };
 
@@ -204,8 +227,8 @@ class IndexBuilder final {
   void addExternalSymbol(scip::SymbolInformation &&extSym);
 
   // The map contains interior references into IndexBuilder's state.
-  std::unique_ptr<SymbolToInfoMap> populateSymbolToInfoMap();
-  void addForwardDeclaration(const SymbolToInfoMap &, scip::ForwardDecl &&);
+  std::unique_ptr<ForwardDeclResolver> populateForwardDeclResolver();
+  void addForwardDeclaration(ForwardDeclResolver &, scip::ForwardDecl &&);
 
   void finish(bool deterministic, std::ostream &);
 

indexer/SymbolFormatter.cc

@@ -108,6 +108,14 @@ void SymbolBuilder::formatTo(std::string &buf) const {
     ::addSpaceEscaped(out, text);
     out << ' ';
   }
+  // NOTE(def: symbol-string-hack-for-forward-decls): Add a '$' prefix
+  // after the space after the version, but before any descriptors.
+  // When splitting the symbol, we check for the '$' preceded by a ' '.
+  // Strictly speaking, it is possible that the same pattern is present
+  // due to a funky filename which contains the pattern ' $', it's highly
+  // improbable to cause any incorrect lookup by not colliding with an
+  // actual symbol name.
+  out << '$';
   for (auto &descriptor : this->descriptors) {
     descriptor.formatTo(out);
   }

indexer/SymbolFormatter.h

@@ -74,8 +74,15 @@ struct SymbolBuilder {
   /// this avoids the extra work of recomputing parent symbol strings.
   static void formatContextual(std::string &buf, SymbolNameRef contextSymbol,
                                const DescriptorBuilder &descriptor);
+
+  static std::optional<scip::SymbolSuffix>
+      getPackageAgnosticSuffix(scip::SymbolNameRef);
+
+  static scip::SymbolName addFakePrefix(scip::SymbolSuffix);
 };
 
+class FileMetadataMap;
+
 class SymbolFormatter final {
   const clang::SourceManager &sourceManager;
   FileMetadataMap &fileMetadataMap;