AST: Split off QueryOverrideSubs and LookUpConformanceInOverrideSubs

Author: slavapestov

include/swift/AST/SubstitutionMap.h

@@ -205,10 +205,10 @@ class SubstitutionMap {
   /// Variant of the above for when we have the generic signatures but not
   /// the decls for 'derived' and 'base'.
   static SubstitutionMap
-  getOverrideSubstitutions(const ClassDecl *baseClass,
-                           const ClassDecl *derivedClass,
+  getOverrideSubstitutions(const NominalTypeDecl *baseNominal,
+                           const NominalTypeDecl *derivedNominal,
                            GenericSignature baseSig,
-                           GenericParamList *derivedParams);
+                           const GenericParamList *derivedParams);
 
   /// Combine two substitution maps as follows.
   ///
@@ -312,6 +312,39 @@ class LookUpConformanceInSubstitutionMap {
                                     ProtocolDecl *conformedProtocol) const;
 };
 
+struct OverrideSubsInfo {
+  ASTContext &Ctx;
+  unsigned BaseDepth;
+  unsigned OrigDepth;
+  SubstitutionMap BaseSubMap;
+  const GenericParamList *DerivedParams;
+
+  OverrideSubsInfo(const NominalTypeDecl *baseNominal,
+                   const NominalTypeDecl *derivedNominal,
+                   GenericSignature baseSig,
+                   const GenericParamList *derivedParams);
+};
+
+struct QueryOverrideSubs {
+  OverrideSubsInfo info;
+
+  explicit QueryOverrideSubs(const OverrideSubsInfo &info)
+    : info(info) {}
+
+  Type operator()(SubstitutableType *type) const;
+};
+
+struct LookUpConformanceInOverrideSubs {
+  OverrideSubsInfo info;
+
+  explicit LookUpConformanceInOverrideSubs(const OverrideSubsInfo &info)
+    : info(info) {}
+
+  ProtocolConformanceRef operator()(CanType type,
+                                    Type substType,
+                                    ProtocolDecl *proto) const;
+};
+
 } // end namespace swift
 
 namespace llvm {

lib/AST/SubstitutionMap.cpp

@@ -523,62 +523,76 @@ SubstitutionMap::getOverrideSubstitutions(
   return getOverrideSubstitutions(baseClass, derivedClass, baseSig, derivedParams);
 }
 
-SubstitutionMap
-SubstitutionMap::getOverrideSubstitutions(const ClassDecl *baseClass,
-                                          const ClassDecl *derivedClass,
-                                          GenericSignature baseSig,
-                                          GenericParamList *derivedParams) {
-  if (baseSig.isNull())
-    return SubstitutionMap();
+OverrideSubsInfo::OverrideSubsInfo(const NominalTypeDecl *baseNominal,
+                                   const NominalTypeDecl *derivedNominal,
+                                   GenericSignature baseSig,
+                                   const GenericParamList *derivedParams)
+  : Ctx(baseSig->getASTContext()),
+    BaseDepth(0),
+    OrigDepth(0),
+    DerivedParams(derivedParams) {
+
+  if (auto baseNominalSig = baseNominal->getGenericSignature()) {
+    BaseDepth = baseNominalSig.getGenericParams().back()->getDepth() + 1;
+
+    auto derivedNominalTy = derivedNominal->getDeclaredInterfaceType();
+    BaseSubMap = derivedNominalTy->getContextSubstitutionMap(
+        baseNominal->getParentModule(), baseNominal);
+    assert(!BaseSubMap.hasArchetypes());
+  }
+
+  if (auto derivedNominalSig = derivedNominal->getGenericSignature())
+    OrigDepth = derivedNominalSig.getGenericParams().back()->getDepth() + 1;
+}
 
-  unsigned baseDepth = 0;
-  SubstitutionMap baseSubMap;
-  if (auto baseClassSig = baseClass->getGenericSignature()) {
-    baseDepth = baseClassSig.getGenericParams().back()->getDepth() + 1;
+Type QueryOverrideSubs::operator()(SubstitutableType *type) const {
+  if (auto gp = type->getAs<GenericTypeParamType>()) {
+    if (gp->getDepth() >= info.BaseDepth) {
+      assert(gp->getDepth() == info.BaseDepth);
+      if (info.DerivedParams != nullptr) {
+        return info.DerivedParams->getParams()[gp->getIndex()]
+            ->getDeclaredInterfaceType();
+      }
 
-    auto derivedClassTy = derivedClass->getDeclaredInterfaceType();
-    baseSubMap = derivedClassTy->getContextSubstitutionMap(
-        baseClass->getParentModule(), baseClass);
-    assert(!baseSubMap.hasArchetypes());
+      return GenericTypeParamType::get(
+          gp->isTypeSequence(),
+          gp->getDepth() + info.OrigDepth - info.BaseDepth,
+          gp->getIndex(), info.Ctx);
+    }
   }
 
-  unsigned origDepth = 0;
-  if (auto derivedClassSig = derivedClass->getGenericSignature())
-    origDepth = derivedClassSig.getGenericParams().back()->getDepth() + 1;
+  return Type(type).subst(info.BaseSubMap);
+}
 
-  auto &ctx = baseSig->getASTContext();
+ProtocolConformanceRef
+LookUpConformanceInOverrideSubs::operator()(CanType type,
+                                            Type substType,
+                                            ProtocolDecl *proto) const {
+  if (type->getRootGenericParam()->getDepth() >= info.BaseDepth)
+    return ProtocolConformanceRef(proto);
 
-  return get(
-    baseSig,
-    [&](SubstitutableType *type) -> Type {
-      if (auto gp = type->getAs<GenericTypeParamType>()) {
-        if (gp->getDepth() >= baseDepth) {
-          assert(gp->getDepth() == baseDepth);
-          if (derivedParams != nullptr) {
-            return derivedParams->getParams()[gp->getIndex()]
-                ->getDeclaredInterfaceType();
-          }
-
-          return GenericTypeParamType::get(gp->isTypeSequence(),
-                                           gp->getDepth() + origDepth - baseDepth,
-                                           gp->getIndex(), ctx);
-        }
-      }
+  if (auto conformance = info.BaseSubMap.lookupConformance(type, proto))
+    return conformance;
 
-      return Type(type).subst(baseSubMap);
-    },
-    [&](CanType type, Type substType, ProtocolDecl *proto) {
-      if (type->getRootGenericParam()->getDepth() >= baseDepth)
-        return ProtocolConformanceRef(proto);
+  if (substType->isTypeParameter())
+    return ProtocolConformanceRef(proto);
 
-      if (auto conformance = baseSubMap.lookupConformance(type, proto))
-        return conformance;
+  return proto->getParentModule()->lookupConformance(substType, proto);
+}
 
-      if (substType->isTypeParameter())
-        return ProtocolConformanceRef(proto);
+SubstitutionMap
+SubstitutionMap::getOverrideSubstitutions(const NominalTypeDecl *baseNominal,
+                                          const NominalTypeDecl *derivedNominal,
+                                          GenericSignature baseSig,
+                                          const GenericParamList *derivedParams) {
+  if (baseSig.isNull())
+    return SubstitutionMap();
 
-      return proto->getParentModule()->lookupConformance(substType, proto);
-    });
+  OverrideSubsInfo info(baseNominal, derivedNominal, baseSig, derivedParams);
+
+  return get(baseSig,
+             QueryOverrideSubs(info),
+             LookUpConformanceInOverrideSubs(info));
 }
 
 SubstitutionMap