Recurse Into Existential Types With Generic Structure

Before the introduction of parameterized existential types, there
was no generic structure under these types. Now that there is, we'll
need to recurse into them to pick up any latent generic types
and re-abstract them appropriately.

rdar://94320481

Author: CodaFi

include/swift/AST/Types.h

@@ -5366,8 +5366,8 @@ class ParameterizedProtocolType final : public TypeBase,
 public:
   /// Retrieve an instance of a protocol composition type with the
   /// given set of members.
-  static Type get(const ASTContext &C, ProtocolType *base,
-                  ArrayRef<Type> args);
+  static ParameterizedProtocolType *get(const ASTContext &C, ProtocolType *base,
+                                        ArrayRef<Type> args);
 
   ProtocolType *getBaseType() const {
     return Base;
@@ -5407,6 +5407,13 @@ class ParameterizedProtocolType final : public TypeBase,
                             RecursiveTypeProperties properties);
 };
 BEGIN_CAN_TYPE_WRAPPER(ParameterizedProtocolType, Type)
+  static CanParameterizedProtocolType get(const ASTContext &C,
+                                          ProtocolType *base,
+                                          ArrayRef<Type> args) {
+    return CanParameterizedProtocolType(
+        ParameterizedProtocolType::get(C, base, args));
+  }
+
   CanProtocolType getBaseType() const {
     return CanProtocolType(getPointer()->getBaseType());
   }
@@ -5473,6 +5480,11 @@ class ExistentialType final : public TypeBase {
   }
 };
 BEGIN_CAN_TYPE_WRAPPER(ExistentialType, Type)
+static CanExistentialType get(CanType constraint) {
+  return CanExistentialType(
+      ExistentialType::get(constraint, /*forceExistential*/ true)
+          ->castTo<ExistentialType>());
+}
   PROXY_CAN_TYPE_SIMPLE_GETTER(getConstraintType)
 END_CAN_TYPE_WRAPPER(ExistentialType, Type)
 

include/swift/SIL/AbstractionPattern.h

@@ -1287,6 +1287,10 @@ class AbstractionPattern {
   /// it.
   AbstractionPattern getReferenceStorageReferentType() const;
 
+  /// Give that the value being abstracted is an existential, return the
+  /// underlying constraint type.
+  AbstractionPattern getExistentialConstraintType() const;
+
   /// Given that the value being abstracted is a function type, return the
   /// abstraction pattern for the derivative function.
   ///

lib/AST/ASTContext.cpp

@@ -3327,9 +3327,9 @@ ProtocolCompositionType::build(const ASTContext &C, ArrayRef<Type> Members,
   return compTy;
 }
 
-Type ParameterizedProtocolType::get(const ASTContext &C,
-                                    ProtocolType *baseTy,
-                                    ArrayRef<Type> args) {
+ParameterizedProtocolType *ParameterizedProtocolType::get(const ASTContext &C,
+                                                          ProtocolType *baseTy,
+                                                          ArrayRef<Type> args) {
   assert(args.size() > 0);
 
   bool isCanonical = baseTy->isCanonical();

lib/AST/Type.cpp

@@ -1641,7 +1641,7 @@ CanType TypeBase::computeCanonicalType() {
     for (Type t : PPT->getArgs())
       CanArgs.push_back(t->getCanonicalType());
     auto &C = Base->getASTContext();
-    Result = ParameterizedProtocolType::get(C, Base, CanArgs).getPointer();
+    Result = ParameterizedProtocolType::get(C, Base, CanArgs);
     break;
   }
   case TypeKind::Existential: {

lib/SIL/IR/AbstractionPattern.cpp

@@ -943,6 +943,55 @@ AbstractionPattern AbstractionPattern::getReferenceStorageReferentType() const {
   llvm_unreachable("bad kind");
 }
 
+static CanType getExistentialConstraintType(CanType type) {
+  assert(type.isExistentialType());
+  if (auto *ET = type->getAs<ExistentialType>()) {
+    return CanType(ET->getConstraintType());
+  }
+  return type;
+}
+
+AbstractionPattern AbstractionPattern::getExistentialConstraintType() const {
+  switch (getKind()) {
+  case Kind::Invalid:
+    llvm_unreachable("querying invalid abstraction pattern!");
+  case Kind::ObjCMethodType:
+  case Kind::CurriedObjCMethodType:
+  case Kind::PartialCurriedObjCMethodType:
+  case Kind::CFunctionAsMethodType:
+  case Kind::CurriedCFunctionAsMethodType:
+  case Kind::PartialCurriedCFunctionAsMethodType:
+  case Kind::CXXMethodType:
+  case Kind::CurriedCXXMethodType:
+  case Kind::PartialCurriedCXXMethodType:
+  case Kind::Tuple:
+  case Kind::OpaqueFunction:
+  case Kind::OpaqueDerivativeFunction:
+  case Kind::ObjCCompletionHandlerArgumentsType:
+    llvm_unreachable("pattern for function or tuple cannot be for optional");
+
+  case Kind::Opaque:
+    return *this;
+
+  case Kind::Type:
+    if (isTypeParameterOrOpaqueArchetype())
+      return AbstractionPattern::getOpaque();
+    return AbstractionPattern(getGenericSignature(),
+                              ::getExistentialConstraintType(getType()));
+
+  case Kind::Discard:
+    return AbstractionPattern::getDiscard(
+        getGenericSignature(), ::getExistentialConstraintType(getType()));
+
+  case Kind::ClangType:
+    // This is not reflected in clang types.
+    return AbstractionPattern(getGenericSignature(),
+                              ::getExistentialConstraintType(getType()),
+                              getClangType());
+  }
+  llvm_unreachable("bad kind");
+}
+
 void AbstractionPattern::dump() const {
   print(llvm::errs());
   llvm::errs() << "\n";
@@ -1683,6 +1732,41 @@ class SubstFunctionTypePatternVisitor
     llvm_unreachable("Unimplemented!");
   }
 
+  CanType visitExistentialType(ExistentialType *exist,
+                               AbstractionPattern pattern) {
+    if (auto gp = handleTypeParameterInAbstractionPattern(pattern, exist))
+      return gp;
+
+    // Avoid walking into the constraint type if we can help it.
+    if (!exist->hasTypeParameter() && !exist->hasArchetype() &&
+        !exist->hasOpaqueArchetype()) {
+      return CanType(exist);
+    }
+
+    return CanExistentialType::get(visit(
+        exist->getConstraintType(), pattern.getExistentialConstraintType()));
+  }
+
+  CanType visitParameterizedProtocolType(ParameterizedProtocolType *ppt,
+                                         AbstractionPattern pattern) {
+    if (auto gp = handleTypeParameterInAbstractionPattern(pattern, ppt))
+      return gp;
+
+    // Recurse into the arguments of the parameterized protocol.
+    SmallVector<Type, 4> substArgs;
+    auto origPPT = pattern.getAs<ParameterizedProtocolType>();
+    for (unsigned i = 0; i < ppt->getArgs().size(); ++i) {
+      auto argTy = ppt->getArgs()[i];
+      auto origArgTy = AbstractionPattern(pattern.getGenericSignatureOrNull(),
+                                          origPPT.getArgs()[i]);
+      auto substEltTy = visit(argTy, origArgTy);
+      substArgs.push_back(substEltTy);
+    }
+
+    return CanType(ParameterizedProtocolType::get(
+        TC.Context, ppt->getBaseType(), substArgs));
+  }
+
   CanType visitTupleType(TupleType *tuple, AbstractionPattern pattern) {
     if (auto gp = handleTypeParameterInAbstractionPattern(pattern, tuple))
       return gp;

lib/SIL/IR/TypeLowering.cpp

@@ -2500,6 +2500,44 @@ TypeConverter::computeLoweredRValueType(TypeExpansionContext forExpansion,
                                              MetatypeRepresentation::Thick);
     }
 
+    CanType visitExistentialType(CanExistentialType substExistType) {
+      // Try to avoid walking into the constraint type if we can help it
+      if (!substExistType->hasTypeParameter() &&
+          !substExistType->hasArchetype() &&
+          !substExistType->hasOpaqueArchetype()) {
+        return substExistType;
+      }
+
+      return CanExistentialType::get(visit(substExistType.getConstraintType()));
+    }
+
+    CanType
+    visitParameterizedProtocolType(CanParameterizedProtocolType substPPT) {
+      bool changed = false;
+      SmallVector<Type, 4> loweredSubstArgs;
+      loweredSubstArgs.reserve(substPPT.getArgs().size());
+
+      auto origConstraint = origType.getExistentialConstraintType();
+      auto origPPT = origConstraint.getAs<ParameterizedProtocolType>();
+      for (auto i : indices(substPPT.getArgs())) {
+        auto origArgTy = AbstractionPattern(
+            origConstraint.getGenericSignatureOrNull(), origPPT.getArgs()[i]);
+        auto substArgType = substPPT.getArgs()[i];
+
+        CanType loweredSubstEltType =
+            TC.getLoweredRValueType(forExpansion, origArgTy, substArgType);
+        changed = changed || substArgType != loweredSubstEltType;
+
+        loweredSubstArgs.push_back(loweredSubstEltType);
+      }
+
+      if (!changed)
+        return substPPT;
+
+      return CanParameterizedProtocolType::get(
+          TC.Context, substPPT->getBaseType(), loweredSubstArgs);
+    }
+
     CanType visitPackType(CanPackType substPackType) {
       llvm_unreachable("");
     }

test/SILGen/parameterized_existentials.swift

@@ -22,6 +22,11 @@ struct S: Q {
   typealias Z = Float
 }
 
+struct R<T, U, V> {
+  var force: () -> any P<T, U, V>
+  // CHECK-LABEL: sil hidden [ossa] @$s13parameterized1RV5forceACyxq_q0_GAA1P_pyxq_q0_XPyc_tcfC : $@convention(method) <T, U, V> (@owned @callee_guaranteed @substituted <τ_0_0, τ_0_1, τ_0_2> () -> @out P<τ_0_0, τ_0_1, τ_0_2> for <T, U, V>, @thin R<T, U, V>.Type) -> @owned R<T, U, V> {
+}
+
 // CHECK-LABEL: sil hidden [ossa] @$s13parameterized6upcastyAA1P_pAA1SVF : $@convention(thin) (S) -> @out P {
 func upcast(_ x: S) -> any P {
   // CHECK: bb0([[RESULT_PARAM:%.*]] : $*P, [[CONCRETE_VAL:%.*]] : $S):
@@ -88,3 +93,8 @@ func upcastResult() {
 
   reuse({ () -> any Q<Int, String, Float> in S() })
 }
+
+// CHECK-LABEL: sil hidden [ossa] @$s13parameterized5forceAA1P_pyxq_q0_XPyr1_lF : $@convention(thin) <T, U, V> () -> @out P<T, U, V> {
+func force<T, U, V>() -> any P<T, U, V> {
+  return R(force: { force() }).force()
+}