Add support for serializing VirtualMethodIndex to/from proto. This data needs to persist between inference stages. PiperOrigin-RevId: 817783981 Change-Id: Ic9dc537e6f5e7e6639338ecf1c6d98d5270b901a

diff --git a/nullability/inference/collect_evidence.cc b/nullability/inference/collect_evidence.cc
index 01aed29..6015dd9 100644
--- a/nullability/inference/collect_evidence.cc
+++ b/nullability/inference/collect_evidence.cc

@@ -170,6 +170,26 @@
   return std::move(W.Index);
 }
 
+RelatedSymbols saveVirtualMethodsMap(const RelatedVirtualMethodsMap &M) {
+  RelatedSymbols Result;
+  auto &RelatedMethods = *Result.mutable_related_symbols();
+  for (auto &[KeyMethod, MethodSet] : M) {
+    RelatedSymbols::SymbolSet &Methods = RelatedMethods[KeyMethod];
+    for (auto &Method : MethodSet)
+      Methods.add_symbols()->set_usr(Method.getKey());
+  }
+  return Result;
+}
+
+RelatedVirtualMethodsMap loadVirtualMethodsMap(const RelatedSymbols &R) {
+  RelatedVirtualMethodsMap Related;
+  for (const auto &[KeyMethod, Methods] : R.related_symbols()) {
+    llvm::StringSet<> &RelatedMethods = Related[KeyMethod];
+    for (auto &Symbol : Methods.symbols()) RelatedMethods.insert(Symbol.usr());
+  }
+  return Related;
+}
+
 VirtualMethodEvidenceFlowDirection getFlowDirection(Evidence::Kind Kind,
                                                     bool ForReturnSlot) {
   switch (Kind) {

diff --git a/nullability/inference/collect_evidence.h b/nullability/inference/collect_evidence.h
index d4f2128..89a61ca 100644
--- a/nullability/inference/collect_evidence.h
+++ b/nullability/inference/collect_evidence.h

@@ -75,6 +75,9 @@
 /// Index the relationships between virtual methods in the TU.
 VirtualMethodIndex getVirtualMethodIndex(ASTContext &Ctx, USRCache &UC);
 
+RelatedSymbols saveVirtualMethodsMap(const RelatedVirtualMethodsMap &M);
+RelatedVirtualMethodsMap loadVirtualMethodsMap(const RelatedSymbols &R);
+
 class SortedFingerprintVector {
  public:
   SortedFingerprintVector() = default;

diff --git a/nullability/inference/collect_evidence_test.cc b/nullability/inference/collect_evidence_test.cc
index e8d4e6d..6fbb615 100644
--- a/nullability/inference/collect_evidence_test.cc
+++ b/nullability/inference/collect_evidence_test.cc

@@ -205,14 +205,19 @@
   if (!Summary) return {Summary.takeError(), Results};
 
   // In the context of a pipeline, the index would be created from the AST and
-  // then serialized to proto, along with the summaries. We use it directly
-  // here, to simulate.
+  // then serialized to proto, along with the summaries. We round-trip the index
+  // here to ensure proper testing of the full save/restore flow.
   VirtualMethodIndex VMI = getVirtualMethodIndex(AST.context(), UsrCache);
+  RelatedSymbols VMIProto = saveVirtualMethodsMap(VMI.Bases);
+
+  VirtualMethodIndex PostVMI;
+  PostVMI.Overrides = std::move(VMI.Overrides);
+  PostVMI.Bases = loadVirtualMethodsMap(VMIProto);
   return {collectEvidenceFromSummary(
               *Summary,
               evidenceEmitterWithPropagation(
                   [&Results](const Evidence& E) { Results.push_back(E); },
-                  std::move(VMI)),
+                  std::move(PostVMI)),
               InputInferences, MakeSolver),
           Results};
 }

diff --git a/nullability/inference/inference.proto b/nullability/inference/inference.proto
index f0bb0b2..ec93518 100644
--- a/nullability/inference/inference.proto
+++ b/nullability/inference/inference.proto

@@ -385,3 +385,18 @@
   optional LogicalContext logical_context = 2;
   repeated NullabilityBehaviorSummary behavior_summaries = 3;
 }
+
+// A generic serialization format for a relation between symbols. Currently,
+// used as for the `RelatedVirtualMethodsMap` type.
+//
+// TODO: b/440317964 - design a compact representation. For example, we could
+// represent Symbols by integers, and separately maintain a map from integers to
+// USRs.
+message RelatedSymbols {
+  message SymbolSet {
+    repeated Symbol symbols = 1;
+  }
+  // Since map keys cannot be message fields, we use strings directly for the
+  // key. These must be USRs.
+  map<string, SymbolSet> related_symbols = 1;
+}