type_roots: Add types map to `Interface` node

The types map is a mapping from cached matching keys of root types to
the node Ids for the types. The mapping will be used to hold type roots
of an STG in the `Interface` node.

PiperOrigin-RevId: 520888424
Change-Id: Ib4e9d0b6dbcef7e7c34448d52a7161048518236f

9 files changed, 100 insertions, 64 deletions

diff --git a/comparison.cc b/comparison.cc
index eda1d3f..29453b8 100644
--- a/comparison.cc
+++ b/comparison.cc
@@ -23,6 +23,7 @@
 
 #include <algorithm>
 #include <array>
+#include <map>
 #include <optional>
 #include <ostream>
 #include <sstream>
@@ -371,6 +372,46 @@ void CompareNodes(Result& result, Compare& compare, const std::vector<Id>& ids1,
   }
 }
 
+void CompareNodes(Result& result, Compare& compare,
+                  const std::map<std::string, Id>& x1,
+                  const std::map<std::string, Id>& x2) {
+  // Group diffs into removed, added and changed symbols for readability.
+  std::vector<Id> removed;
+  std::vector<Id> added;
+  std::vector<std::pair<Id, Id>> in_both;
+
+  auto it1 = x1.begin();
+  auto it2 = x2.begin();
+  const auto end1 = x1.end();
+  const auto end2 = x2.end();
+  while (it1 != end1 || it2 != end2) {
+    if (it2 == end2 || (it1 != end1 && it1->first < it2->first)) {
+      // removed
+      removed.push_back(it1->second);
+      ++it1;
+    } else if (it1 == end1 || (it2 != end2 && it1->first > it2->first)) {
+      // added
+      added.push_back(it2->second);
+      ++it2;
+    } else {
+      // in both
+      in_both.emplace_back(it1->second, it2->second);
+      ++it1;
+      ++it2;
+    }
+  }
+
+  for (const auto symbol1 : removed) {
+    result.AddEdgeDiff("", compare.Removed(symbol1));
+  }
+  for (const auto symbol2 : added) {
+    result.AddEdgeDiff("", compare.Added(symbol2));
+  }
+  for (const auto& [symbol1, symbol2] : in_both) {
+    result.MaybeAddEdgeDiff("", compare(symbol1, symbol2));
+  }
+}
+
 }  // namespace
 
 Result Compare::operator()(const BaseClass& x1, const BaseClass& x2) {
@@ -627,45 +668,8 @@ Result Compare::operator()(const ElfSymbol& x1, const ElfSymbol& x2) {
 Result Compare::operator()(const Interface& x1, const Interface& x2) {
   Result result;
   result.diff_.holds_changes = true;
-
-  // Group diffs into removed, added and changed symbols for readability.
-  std::vector<Id> removed;
-  std::vector<Id> added;
-  std::vector<std::pair<Id, Id>> in_both;
-
-  const auto& symbols1 = x1.symbols;
-  const auto& symbols2 = x2.symbols;
-  auto it1 = symbols1.begin();
-  auto it2 = symbols2.begin();
-  const auto end1 = symbols1.end();
-  const auto end2 = symbols2.end();
-  while (it1 != end1 || it2 != end2) {
-    if (it2 == end2 || (it1 != end1 && it1->first < it2->first)) {
-      // removed
-      removed.push_back(it1->second);
-      ++it1;
-    } else if (it1 == end1 || (it2 != end2 && it1->first > it2->first)) {
-      // added
-      added.push_back(it2->second);
-      ++it2;
-    } else {
-      // in both
-      in_both.emplace_back(it1->second, it2->second);
-      ++it1;
-      ++it2;
-    }
-  }
-
-  for (const auto symbol1 : removed) {
-    result.AddEdgeDiff("", Removed(symbol1));
-  }
-  for (const auto symbol2 : added) {
-    result.AddEdgeDiff("", Added(symbol2));
-  }
-  for (const auto& [symbol1, symbol2] : in_both) {
-    result.MaybeAddEdgeDiff("", (*this)(symbol1, symbol2));
-  }
-
+  CompareNodes(result, *this, x1.symbols, x2.symbols);
+  CompareNodes(result, *this, x1.types, x2.types);
   return result;
 }
 
diff --git a/equality.h b/equality.h
index f17e75f..827c484 100644
--- a/equality.h
+++ b/equality.h
@@ -22,7 +22,6 @@
 
 #include <map>
 #include <vector>
-#include <utility>
 
 #include "graph.h"
 #include "scc.h"
@@ -216,7 +215,8 @@ struct Equals {
   }
 
   bool operator()(const Interface& x1, const Interface& x2) {
-    return (*this)(x1.symbols, x2.symbols);
+    return (*this)(x1.symbols, x2.symbols)
+        && (*this)(x1.types, x2.types);
   }
 
   bool Mismatch() {
diff --git a/fidelity.cc b/fidelity.cc
index f5b5718..d25e1c4 100644
--- a/fidelity.cc
+++ b/fidelity.cc
@@ -20,6 +20,7 @@
 #include "fidelity.h"
 
 #include <algorithm>
+#include <map>
 #include <ostream>
 #include <set>
 #include <string>
@@ -73,6 +74,7 @@ struct Fidelity {
 
   void operator()(Id);
   void operator()(const std::vector<Id>&);
+  void operator()(const std::map<std::string, Id>&);
   void operator()(const Void&, Id);
   void operator()(const Variadic&, Id);
   void operator()(const PointerReference&, Id);
@@ -109,6 +111,12 @@ void Fidelity::operator()(const std::vector<Id>& x) {
   }
 }
 
+void Fidelity::operator()(const std::map<std::string, Id>& x) {
+  for (const auto& [_, id] : x) {
+    (*this)(id);
+  }
+}
+
 void Fidelity::operator()(const Void&, Id) {}
 
 void Fidelity::operator()(const Variadic&, Id) {}
@@ -188,9 +196,8 @@ void Fidelity::operator()(const ElfSymbol& x, Id) {
 }
 
 void Fidelity::operator()(const Interface& x, Id) {
-  for (const auto& [_, id] : x.symbols) {
-    (*this)(id);
-  }
+  (*this)(x.symbols);
+  (*this)(x.types);
 }
 
 template <typename T>
diff --git a/fingerprint.cc b/fingerprint.cc
index bd3494c..7a26aac 100644
--- a/fingerprint.cc
+++ b/fingerprint.cc
@@ -19,12 +19,15 @@
 
 #include "fingerprint.h"
 
+#include <map>
 #include <string>
 #include <unordered_map>
 #include <unordered_set>
 #include <utility>
 #include <vector>
 
+#include "graph.h"
+#include "hashing.h"
 #include "scc.h"
 
 namespace stg {
@@ -136,9 +139,8 @@ struct Hasher {
   }
 
   HashValue operator()(const Interface& x) {
-    for (const auto& [name, symbol] : x.symbols) {
-      todo.insert(symbol);
-    }
+    ToDo(x.symbols);
+    ToDo(x.types);
     return hash('Z');
   }
 
@@ -193,6 +195,12 @@ struct Hasher {
     }
   }
 
+  void ToDo(const std::map<std::string, Id>& ids) {
+    for (const auto& [_, id] : ids) {
+      todo.insert(id);
+    }
+  }
+
   const Graph& graph;
   std::unordered_map<Id, HashValue>& hashes;
   std::unordered_set<Id> &todo;
diff --git a/graph.h b/graph.h
index 529ca27..96ef344 100644
--- a/graph.h
+++ b/graph.h
@@ -312,8 +312,12 @@ std::ostream& operator<<(std::ostream& os, ElfSymbol::CRC crc);
 struct Interface {
   explicit Interface(const std::map<std::string, Id>& symbols)
       : symbols(symbols) {}
+  Interface(const std::map<std::string, Id>& symbols,
+            const std::map<std::string, Id>& types)
+      : symbols(symbols), types(types) {}
 
   std::map<std::string, Id> symbols;
+  std::map<std::string, Id> types;
 };
 
 std::ostream& operator<<(std::ostream& os, Primitive::Encoding encoding);
diff --git a/stg.cc b/stg.cc
index 1ed363c..d531ae9 100644
--- a/stg.cc
+++ b/stg.cc
@@ -27,6 +27,7 @@
 #include <memory>
 #include <ostream>
 #include <string>
+#include <utility>
 #include <vector>
 
 #include "deduplication.h"
@@ -45,22 +46,22 @@ namespace {
 Metrics metrics;
 
 struct GetInterface {
-  const std::map<std::string, Id>& operator()(const Interface& x) {
-    return x.symbols;
+  Interface& operator()(Interface& x) {
+    return x;
   }
 
   template <typename Node>
-  const std::map<std::string, Id>& operator()(const Node&) {
+  Interface& operator()(Node&) {
     Die() << "expected an Interface root node";
   }
 };
 
+// TODO: Implement merging for rooted types.
 Id Merge(Graph& graph, const std::vector<Id>& roots) {
   std::map<std::string, Id> symbols;
   GetInterface get;
   for (auto root : roots) {
-    for (const auto& x :
-         graph.Apply<const std::map<std::string, Id>&>(get, root)) {
+    for (const auto& x : graph.Apply<Interface&>(get, root).symbols) {
       if (!symbols.insert(x).second) {
         Die() << "merge failed with duplicate symbol: " << x.first;
       }
@@ -73,14 +74,13 @@ Id Merge(Graph& graph, const std::vector<Id>& roots) {
 void Filter(Graph& graph, Id root, const SymbolFilter& filter) {
   std::map<std::string, Id> symbols;
   GetInterface get;
-  for (const auto& x :
-       graph.Apply<const std::map<std::string, Id>&>(get, root)) {
+  auto& interface = graph.Apply<Interface&>(get, root);
+  for (const auto& x : interface.symbols) {
     if (filter(x.first)) {
       symbols.insert(x);
     }
   }
-  graph.Unset(root);
-  graph.Set<Interface>(root, symbols);
+  std::swap(interface.symbols, symbols);
 }
 
 void Write(const Graph& graph, Id root, const char* output, bool stable) {
diff --git a/substitution.h b/substitution.h
index f68392f..adcb801 100644
--- a/substitution.h
+++ b/substitution.h
@@ -130,6 +130,7 @@ struct Substitute {
 
   void operator()(Interface& x) {
     Update(x.symbols);
+    Update(x.types);
   }
 
   Graph& graph;
diff --git a/type_normalisation.cc b/type_normalisation.cc
index 74eb723..72cc94b 100644
--- a/type_normalisation.cc
+++ b/type_normalisation.cc
@@ -19,6 +19,8 @@
 
 #include "type_normalisation.h"
 
+#include <map>
+#include <string>
 #include <unordered_map>
 #include <unordered_set>
 #include <vector>
@@ -80,6 +82,12 @@ struct FindQualifiedTypesAndFunctions {
     }
   }
 
+  void operator()(const std::map<std::string, Id>& x) {
+    for (const auto& [_, id] : x) {
+      (*this)(id);
+    }
+  }
+
   void operator()(const Void&, Id) {}
 
   void operator()(const Variadic&, Id) {}
@@ -150,9 +158,8 @@ struct FindQualifiedTypesAndFunctions {
   }
 
   void operator()(const Interface& x, Id) {
-    for (auto& [_, id] : x.symbols) {
-      (*this)(id);
-    }
+    (*this)(x.symbols);
+    (*this)(x.types);
   }
 
   const Graph& graph;
diff --git a/type_resolution.cc b/type_resolution.cc
index aa187d8..bde42f6 100644
--- a/type_resolution.cc
+++ b/type_resolution.cc
@@ -21,7 +21,6 @@
 
 #include <functional>
 #include <map>
-#include <sstream>
 #include <string>
 #include <unordered_map>
 #include <unordered_set>
@@ -59,6 +58,12 @@ struct NamedTypes {
     }
   }
 
+  void operator()(const std::map<std::string, Id>& x) {
+    for (const auto& [_, id] : x) {
+      (*this)(id);
+    }
+  }
+
   // main entry point
   void operator()(Id id) {
     if (seen.Insert(id)) {
@@ -169,9 +174,8 @@ struct NamedTypes {
   }
 
   void operator()(const Interface& x, Id) {
-    for (const auto& [_, symbol] : x.symbols) {
-      (*this)(symbol);
-    }
+    (*this)(x.symbols);
+    (*this)(x.types);
   }
 
   const Graph& graph;
@@ -437,7 +441,8 @@ struct Unify {
   }
 
   bool operator()(const Interface& x1, const Interface& x2) {
-    return (*this)(x1.symbols, x2.symbols);
+    return (*this)(x1.symbols, x2.symbols)
+        && (*this)(x1.types, x2.types);
   }
 
   bool Mismatch() {