AST: Re-implement GenericEnvironment::mapTypeIntoContext() using TypeTransform

slavapestov · slavapestov · commit 4923521227a5 · 2024-09-13T15:19:48.000-04:00
diff --git a/include/swift/AST/GenericEnvironment.h b/include/swift/AST/GenericEnvironment.h
@@ -157,14 +157,14 @@ class alignas(1 << DeclAlignInBits) GenericEnvironment final
   friend ArchetypeType;
   friend QueryInterfaceTypeSubstitutions;
 
-  /// Add a mapping of a generic parameter to a specific type (which may be
-  /// an archetype)
-  void addMapping(GenericParamKey key, Type contextType);
+  /// Add a mapping of a type parameter to a contextual type, usually
+  /// an archetype.
+  void addMapping(CanType depType, Type contextType);
 
-  /// Retrieve the mapping for the given generic parameter, if present.
+  /// Retrieve the mapping for the given type parameter, if present.
   ///
   /// This is only useful when lazily populating a generic environment.
-  std::optional<Type> getMappingIfPresent(GenericParamKey key) const;
+  Type getMappingIfPresent(CanType depType) const;
 
 public:
   GenericSignature getGenericSignature() const {
diff --git a/lib/AST/GenericEnvironment.cpp b/lib/AST/GenericEnvironment.cpp
@@ -19,6 +19,7 @@
 #include "swift/AST/ASTContext.h"
 #include "swift/AST/ExistentialLayout.h"
 #include "swift/AST/ProtocolConformance.h"
+#include "swift/AST/TypeTransform.h"
 #include "swift/Basic/Assertions.h"
 #include "swift/Basic/Defer.h"
 
@@ -262,33 +263,49 @@ GenericEnvironment::GenericEnvironment(GenericSignature signature,
   assert(i == openedPacksBuffer.size());
 }
 
-void GenericEnvironment::addMapping(GenericParamKey key,
-                                    Type contextType) {
-  // Find the index into the parallel arrays of generic parameters and
-  // context types.
-  auto genericParams = getGenericParams();
-  unsigned index = key.findIndexIn(genericParams);
-  assert(genericParams[index] == key && "Bad generic parameter");
+class GenericEnvironment::NestedTypeStorage
+    : public llvm::DenseMap<CanType, Type> { };
+
+void GenericEnvironment::addMapping(CanType depType, Type contextType) {
+  if (auto genericParam = dyn_cast<GenericTypeParamType>(depType)) {
+    GenericParamKey key(genericParam);
 
-  // Add the mapping from the generic parameter to the context type.
-  assert(getContextTypes()[index].isNull() ||
-         getContextTypes()[index]->is<ErrorType>() &&
-         "Already recoded this mapping");
-  getContextTypes()[index] = contextType;
+    // Find the index into the parallel arrays of generic parameters and
+    // context types.
+    auto genericParams = getGenericParams();
+    unsigned index = key.findIndexIn(genericParams);
+    assert(genericParams[index] == key && "Bad generic parameter");
+
+    // Add the mapping from the generic parameter to the context type.
+    assert(getContextTypes()[index].isNull() ||
+           getContextTypes()[index]->is<ErrorType>() &&
+           "Already recoded this mapping");
+    getContextTypes()[index] = contextType;
+  } else {
+    getOrCreateNestedTypeStorage()[depType] = contextType;
+  }
 }
 
-std::optional<Type>
-GenericEnvironment::getMappingIfPresent(GenericParamKey key) const {
-  // Find the index into the parallel arrays of generic parameters and
-  // context types.
-  auto genericParams = getGenericParams();
-  unsigned index = key.findIndexIn(genericParams);
-  assert(genericParams[index] == key && "Bad generic parameter");
+Type GenericEnvironment::getMappingIfPresent(CanType depType) const {
+  if (auto genericParam = dyn_cast<GenericTypeParamType>(depType)) {
+    GenericParamKey key(genericParam);
 
-  if (auto type = getContextTypes()[index])
-    return type;
+    // Find the index into the parallel arrays of generic parameters and
+    // context types.
+    auto genericParams = getGenericParams();
+    unsigned index = key.findIndexIn(genericParams);
+    assert(genericParams[index] == key && "Bad generic parameter");
 
-  return std::nullopt;
+    return getContextTypes()[index];
+  } else {
+    auto &storage = const_cast<GenericEnvironment *>(this)
+        ->getOrCreateNestedTypeStorage();
+    auto found = storage.find(depType);
+    if (found != storage.end())
+      return found->second;
+
+    return Type();
+  }
 }
 
 Type
@@ -337,9 +354,6 @@ Type TypeBase::mapTypeOutOfContext() {
     SubstFlags::SubstitutePrimaryArchetypes);
 }
 
-class GenericEnvironment::NestedTypeStorage
-    : public llvm::DenseMap<CanType, Type> { };
-
 auto GenericEnvironment::getOrCreateNestedTypeStorage() -> NestedTypeStorage & {
   if (nestedTypeStorage)
     return *nestedTypeStorage;
@@ -355,35 +369,41 @@ auto GenericEnvironment::getOrCreateNestedTypeStorage() -> NestedTypeStorage & {
 
 Type
 GenericEnvironment::getOrCreateArchetypeFromInterfaceType(Type depType) {
+  auto canType = depType->getCanonicalType();
+
+  // Have we seen this exact type parameter before?
+  if (auto type = getMappingIfPresent(canType))
+    return type;
+
   auto genericSig = getGenericSignature();
-  auto requirements = genericSig->getLocalRequirements(depType);
 
-  if (requirements.concreteType)
-    return mapTypeIntoContext(requirements.concreteType);
+  // Reduce it.
+  auto reducedType = genericSig.getReducedType(canType);
 
-  assert(requirements.anchor && "No anchor or concrete type?");
+  // If this type parameter is equivalent to a concrete type,
+  // map the concrete type into context and cache the result.
+  if (!reducedType->isTypeParameter()) {
+    auto result = mapTypeIntoContext(reducedType);
+    addMapping(canType, result);
+    return result;
+  }
 
   auto &ctx = genericSig->getASTContext();
 
-  // First, write an ErrorType to the location where this type is cached,
-  // to catch re-entrant lookups that might arise from an invalid generic
-  // signature (eg, <X where X == Array<X>>).
-  CanDependentMemberType nestedType;
-  GenericTypeParamType *genericParam = nullptr;
-  if (auto depMemTy = requirements.anchor->getAs<DependentMemberType>()) {
-    nestedType = cast<DependentMemberType>(depMemTy->getCanonicalType());
-    auto &entry = getOrCreateNestedTypeStorage()[nestedType];
-    if (entry)
-      return entry;
-
-    entry = ErrorType::get(ctx);
-  } else {
-    genericParam = requirements.anchor->castTo<GenericTypeParamType>();
-    if (auto type = getMappingIfPresent(genericParam))
-      return *type;
-    addMapping(genericParam, ErrorType::get(ctx));
+  // If the original type parameter was not reduced, see if we have an
+  // archetype for the reduced type parameter.
+  if (canType != reducedType) {
+    if (auto type = getMappingIfPresent(reducedType)) {
+      // Cache the result.
+      addMapping(canType, type);
+      return type;
+    }
   }
 
+  // Otherwise, we're going to create a new archetype. Look up its
+  // requirements.
+  auto requirements = genericSig->getLocalRequirements(reducedType);
+
   // Substitute into the superclass.
   Type superclass = requirements.superclass;
   if (superclass && superclass->hasTypeParameter()) {
@@ -394,16 +414,18 @@ GenericEnvironment::getOrCreateArchetypeFromInterfaceType(Type depType) {
 
   Type result;
 
-  auto rootGP = requirements.anchor->getRootGenericParam();
+  auto sugaredType = genericSig->getSugaredType(reducedType);
+
+  auto rootGP = reducedType->getRootGenericParam();
   switch (getKind()) {
   case Kind::Primary:
     if (rootGP->isParameterPack()) {
-      result = PackArchetypeType::get(ctx, this, requirements.anchor,
+      result = PackArchetypeType::get(ctx, this, sugaredType,
                                       requirements.packShape,
                                       requirements.protos, superclass,
                                       requirements.layout);
     } else {
-      result = PrimaryArchetypeType::getNew(ctx, this, requirements.anchor,
+      result = PrimaryArchetypeType::getNew(ctx, this, sugaredType,
                                             requirements.protos, superclass,
                                             requirements.layout);
     }
@@ -415,19 +437,19 @@ GenericEnvironment::getOrCreateArchetypeFromInterfaceType(Type depType) {
     // represents an opaque declaration, then apply the outer substitutions.
     // It would be incorrect to build an opaque type archetype here.
     if (rootGP->getDepth() < genericSig->getMaxDepth()) {
-      result = maybeApplyOuterContextSubstitutions(requirements.anchor);
+      result = maybeApplyOuterContextSubstitutions(reducedType);
       break;
     }
 
-    result = OpaqueTypeArchetypeType::getNew(this, requirements.anchor,
+    result = OpaqueTypeArchetypeType::getNew(this, sugaredType,
                                              requirements.protos, superclass,
                                              requirements.layout);
     break;
   }
 
   case Kind::OpenedExistential: {
     if (rootGP->getDepth() < genericSig->getMaxDepth()) {
-      result = maybeApplyOuterContextSubstitutions(requirements.anchor);
+      result = maybeApplyOuterContextSubstitutions(reducedType);
       break;
     }
 
@@ -439,16 +461,16 @@ GenericEnvironment::getOrCreateArchetypeFromInterfaceType(Type depType) {
     // For compatibility, continue using the existential layout's version when
     // the interface type is a generic parameter. We should align these at
     // some point.
-    if (depType->is<GenericTypeParamType>()) {
+    if (isa<GenericTypeParamType>(reducedType)) {
       auto layout = getOpenedExistentialType()->getExistentialLayout();
       SmallVector<ProtocolDecl *, 2> protos;
       for (auto proto : layout.getProtocols())
         protos.push_back(proto);
 
-      result = OpenedArchetypeType::getNew(this, requirements.anchor, protos,
+      result = OpenedArchetypeType::getNew(this, sugaredType, protos,
                                            superclass, requirements.layout);
     } else {
-      result = OpenedArchetypeType::getNew(this, requirements.anchor,
+      result = OpenedArchetypeType::getNew(this, sugaredType,
                                            requirements.protos, superclass,
                                            requirements.layout);
     }
@@ -458,71 +480,87 @@ GenericEnvironment::getOrCreateArchetypeFromInterfaceType(Type depType) {
 
   case Kind::OpenedElement: {
     if (rootGP->getDepth() < genericSig->getMaxDepth()) {
-      result = maybeApplyOuterContextSubstitutions(requirements.anchor);
+      result = maybeApplyOuterContextSubstitutions(reducedType);
       break;
     }
 
-    result = ElementArchetypeType::getNew(this, requirements.anchor,
+    result = ElementArchetypeType::getNew(this, sugaredType,
                                           requirements.protos, superclass,
                                           requirements.layout);
     break;
   }
   }
 
-  if (genericParam)
-    addMapping(genericParam, result);
-  else
-    getOrCreateNestedTypeStorage()[nestedType] = result;
+  // Cache the result.
+  addMapping(canType, result);
+  if (canType != reducedType)
+    addMapping(reducedType, result);
 
   return result;
 }
 
 Type QueryInterfaceTypeSubstitutions::operator()(SubstitutableType *type) const{
-  if (auto gp = type->getAs<GenericTypeParamType>()) {
-    // Find the index into the parallel arrays of generic parameters and
-    // context types.
-    auto genericParams = self->getGenericParams();
-    GenericParamKey key(gp);
+  auto gp = type->castTo<GenericTypeParamType>();
 
-    // Make sure that this generic parameter is from this environment.
-    unsigned index = key.findIndexIn(genericParams);
-    if (index == genericParams.size() || genericParams[index] != key)
-      return Type();
+  // Find the index into the parallel arrays of generic parameters and
+  // context types.
+  auto genericParams = self->getGenericParams();
+  GenericParamKey key(gp);
 
-    // If the context type isn't already known, lazily create it.
-    auto mutableSelf = const_cast<GenericEnvironment *>(self);
-    Type &contextType = mutableSelf->getContextTypes()[index];
-    if (contextType)
-      return contextType;
+  // Make sure that this generic parameter is from this environment and
+  // return substitution failure if not.
+  unsigned index = key.findIndexIn(genericParams);
+  if (index == genericParams.size())
+    return Type();
 
-    auto result = mutableSelf->getOrCreateArchetypeFromInterfaceType(type);
+  // If the context type isn't already known, lazily create it.
+  if (auto contextType = self->getContextTypes()[index])
+    return contextType;
 
-    assert (!contextType ||
-            contextType->isEqual(result) ||
-            contextType->is<ErrorType>());
-    contextType = result;
-    return result;
+  return const_cast<GenericEnvironment *>(self)
+      ->getOrCreateArchetypeFromInterfaceType(gp);
+}
+
+namespace {
+
+struct MapTypeIntoContext: TypeTransform<MapTypeIntoContext> {
+  GenericEnvironment *env;
+
+  explicit MapTypeIntoContext(GenericEnvironment *env) : env(env) {}
+
+  std::optional<Type> transform(TypeBase *type, TypePosition pos) {
+    if (!type->hasTypeParameter())
+      return Type(type);
+
+    return std::nullopt;
   }
 
-  return Type();
+  Type transformGenericTypeParamType(GenericTypeParamType *param,
+                                     TypePosition pos) {
+    return env->getOrCreateArchetypeFromInterfaceType(param);
+  }
+
+  Type transformDependentMemberType(DependentMemberType *dependent,
+                                    TypePosition pos) {
+    return env->getOrCreateArchetypeFromInterfaceType(dependent);
+  }
+
+  CanType transformSILField(CanType fieldTy, TypePosition pos) {
+    return fieldTy;
+  }
+};
+
 }
 
 Type GenericEnvironment::mapTypeIntoContext(Type type) const {
   assert(!type->hasPrimaryArchetype() && "already have a contextual type");
-
-  Type result = type.subst(QueryInterfaceTypeSubstitutions(this),
-                           LookUpConformanceInModule(),
-                           SubstFlags::PreservePackExpansionLevel);
-  ASSERT(getKind() != Kind::Primary || !result->hasTypeParameter());
-  return result;
+  return MapTypeIntoContext(const_cast<GenericEnvironment *>(this))
+      .doIt(type, TypePosition::Invariant);
 }
 
 Type GenericEnvironment::mapTypeIntoContext(GenericTypeParamType *type) const {
-  auto self = const_cast<GenericEnvironment *>(this);
-  Type result = QueryInterfaceTypeSubstitutions(self)(type);
-  if (!result)
-    return ErrorType::get(type);
-  return result;
+  return const_cast<GenericEnvironment *>(this)
+      ->getOrCreateArchetypeFromInterfaceType(type);
 }
 
 namespace {
diff --git a/lib/Sema/ConstraintSystem.cpp b/lib/Sema/ConstraintSystem.cpp
@@ -2051,7 +2051,7 @@ static void bindArchetypesFromContext(
   // Find the innermost non-type context.
   for (const auto *parentDC = outerDC;
        !parentDC->isModuleScopeContext();
-       parentDC = parentDC->getParent()) {
+       parentDC = parentDC->getParentForLookup()) {
     if (parentDC->isTypeContext()) {
       if (parentDC != outerDC && parentDC->getSelfProtocolDecl()) {
         auto selfTy = parentDC->getSelfInterfaceType();
diff --git a/lib/Sema/IDETypeCheckingRequests.cpp b/lib/Sema/IDETypeCheckingRequests.cpp
@@ -203,7 +203,8 @@ static bool isMemberDeclAppliedInternal(const DeclContext *DC, Type BaseTy,
   // The context substitution map for the base type fixes the declaration's
   // outer generic parameters.
   auto substMap = BaseTy->getContextSubstitutionMap(
-      VD->getDeclContext(), genericDecl->getGenericEnvironment());
+      VD->getDeclContext(),
+      VD->getDeclContext()->getGenericEnvironmentOfContext());
 
   // The innermost generic parameters are mapped to error types.
   unsigned innerDepth = genericSig->getMaxDepth();
diff --git a/lib/Sema/TypeCheckType.cpp b/lib/Sema/TypeCheckType.cpp
@@ -1220,7 +1220,15 @@ Type TypeResolution::applyUnboundGenericArguments(
       if (result->hasTypeParameter()) {
         if (const auto contextSig = getGenericSignature()) {
           auto *genericEnv = contextSig.getGenericEnvironment();
-          return genericEnv->mapTypeIntoContext(result);
+          // FIXME: This should just use mapTypeIntoContext(), but we can't yet
+          // because we sometimes have type parameters here that are invalid for
+          // our generic signature. This can happen if the type parameter was
+          // found via unqualified lookup, but the current context's
+          // generic signature failed to build because of circularity or
+          // completion failure.
+          return result.subst(QueryInterfaceTypeSubstitutions{genericEnv},
+                              LookUpConformanceInModule(),
+                              SubstFlags::PreservePackExpansionLevel);
         }
       }
       return result;
