Remove FIXME in TypeExpansionAnalysis: It's cache needed to be keyed by the type expansion context

Author: aschwaighofer

include/swift/AST/TypeExpansionContext.h

@@ -18,9 +18,11 @@
 #define SWIFT_TYPEEXPANSIONCONTEXT_H
 
 #include "swift/AST/ResilienceExpansion.h"
+#include "llvm/ADT/DenseMap.h"
 
 namespace swift {
   class DeclContext;
+  class SILFunction;
 
 /// Describes the context in which SIL types should eventually be expanded.
 /// Required for lowering resilient types and deciding whether to look through
@@ -111,4 +113,33 @@ class TypeExpansionContext {
 
 } // namespace swift
 
+namespace llvm {
+template <> struct DenseMapInfo<swift::TypeExpansionContext> {
+  using TypeExpansionContext = swift::TypeExpansionContext;
+
+  static TypeExpansionContext getEmptyKey() {
+    return TypeExpansionContext(
+        swift::ResilienceExpansion::Minimal,
+        reinterpret_cast<swift::DeclContext *>(
+            DenseMapInfo<swift::DeclContext *>::getEmptyKey()),
+        false);
+  }
+  static TypeExpansionContext getTombstoneKey() {
+    return TypeExpansionContext(
+        swift::ResilienceExpansion::Minimal,
+        reinterpret_cast<swift::DeclContext *>(
+            DenseMapInfo<swift::DeclContext *>::getTombstoneKey()),
+        false);
+  }
+
+  static unsigned getHashValue(TypeExpansionContext val) {
+    return DenseMapInfo<uintptr_t>::getHashValue(val.getHashKey());
+  }
+
+  static bool isEqual(TypeExpansionContext LHS, TypeExpansionContext RHS) {
+    return LHS == RHS;
+  }
+};
+} // namespace llvm
+
 #endif

include/swift/SIL/TypeLowering.h

@@ -30,35 +30,6 @@ namespace clang {
   class Type;
 }
 
-namespace llvm {
-template <> struct DenseMapInfo<swift::TypeExpansionContext> {
-  using TypeExpansionContext = swift::TypeExpansionContext;
-
-  static TypeExpansionContext getEmptyKey() {
-    return TypeExpansionContext(
-        swift::ResilienceExpansion::Minimal,
-        reinterpret_cast<swift::ModuleDecl *>(
-            DenseMapInfo<swift::ModuleDecl *>::getEmptyKey()),
-        false);
-  }
-  static TypeExpansionContext getTombstoneKey() {
-    return TypeExpansionContext(
-        swift::ResilienceExpansion::Minimal,
-        reinterpret_cast<swift::ModuleDecl *>(
-            DenseMapInfo<swift::ModuleDecl *>::getTombstoneKey()),
-        false);
-  }
-
-  static unsigned getHashValue(TypeExpansionContext val) {
-    return DenseMapInfo<uintptr_t>::getHashValue(val.getHashKey());
-  }
-
-  static bool isEqual(TypeExpansionContext LHS, TypeExpansionContext RHS) {
-    return LHS == RHS;
-  }
-};
-}
-
 namespace swift {
   class AnyFunctionRef;
   enum class Bridgeability : unsigned;

include/swift/SILOptimizer/Analysis/TypeExpansionAnalysis.h

@@ -12,6 +12,7 @@
 #ifndef SWIFT_SILOPTIMIZER_ANALYSIS_TYPEEXPANSIONANALYSIS_H
 #define SWIFT_SILOPTIMIZER_ANALYSIS_TYPEEXPANSIONANALYSIS_H
 
+#include "swift/AST/TypeExpansionContext.h"
 #include "swift/SIL/Projection.h"
 #include "swift/SIL/SILType.h"
 #include "swift/SIL/SILValue.h"
@@ -22,7 +23,9 @@ namespace swift {
 
 /// This analysis determines memory effects during destruction.
 class TypeExpansionAnalysis : public SILAnalysis {
-  llvm::DenseMap<SILType, ProjectionPathList> ExpansionCache;
+  llvm::DenseMap<std::pair<SILType, TypeExpansionContext>, ProjectionPathList>
+      ExpansionCache;
+
 public:
   TypeExpansionAnalysis(SILModule *M)
       : SILAnalysis(SILAnalysisKind::TypeExpansion) {}

lib/SILOptimizer/Analysis/TypeExpansionAnalysis.cpp

@@ -28,9 +28,8 @@ const ProjectionPathList &
 TypeExpansionAnalysis::getTypeExpansion(SILType B, SILModule *Mod,
                                         TypeExpansionContext context) {
   // Check whether we have the type expansion.
-  auto Iter = ExpansionCache.find(B);
-  // TODO: FIXME: ExpansionCache.find((B,context))
-  //
+  auto key = std::make_pair(B, context);
+  auto Iter = ExpansionCache.find(key);
   //
   if (Iter != ExpansionCache.end()) {
     return Iter->second;
@@ -40,8 +39,8 @@ TypeExpansionAnalysis::getTypeExpansion(SILType B, SILModule *Mod,
   if (!shouldExpand(*Mod, B)) {
     // Push the empty projection path.
     ProjectionPath P(B);
-    ExpansionCache[B].push_back(P);
-    return ExpansionCache[B];
+    ExpansionCache[key].push_back(P);
+    return ExpansionCache[key];
   }
 
   // Flush the cache if the size of the cache is too large.
@@ -51,8 +50,8 @@ TypeExpansionAnalysis::getTypeExpansion(SILType B, SILModule *Mod,
 
   // Build the type expansion for the leaf nodes.
   ProjectionPath::expandTypeIntoLeafProjectionPaths(B, Mod, context,
-                                                    ExpansionCache[B]);
-  return ExpansionCache[B];
+                                                    ExpansionCache[key]);
+  return ExpansionCache[key];
 }
 
 SILAnalysis *swift::createTypeExpansionAnalysis(SILModule *M) {