Merge pull request #59290 from jckarter/type-expansion-wmo-private-type

SILGen: Carry WMO of type lowering context to closure captures.

Author: jckarter

include/swift/SIL/TypeLowering.h

@@ -821,6 +821,8 @@ class TypeConverter {
 
   llvm::DenseMap<AbstractClosureExpr *, Optional<AbstractionPattern>>
     ClosureAbstractionPatterns;
+  llvm::DenseMap<SILDeclRef, TypeExpansionContext>
+    CaptureTypeExpansionContexts;
 
   CanAnyFunctionType makeConstantInterfaceType(SILDeclRef constant);
 
@@ -1225,6 +1227,7 @@ class TypeConverter {
   /// the abstraction pattern is queried using this function. Once the
   /// abstraction pattern has been asked for, it may not be changed.
   Optional<AbstractionPattern> getConstantAbstractionPattern(SILDeclRef constant);
+  TypeExpansionContext getCaptureTypeExpansionContext(SILDeclRef constant);
 
   /// Set the preferred abstraction pattern for a closure.
   ///
@@ -1234,6 +1237,8 @@ class TypeConverter {
   void setAbstractionPattern(AbstractClosureExpr *closure,
                              AbstractionPattern pattern);
 
+  void setCaptureTypeExpansionContext(SILDeclRef constant,
+                                      SILModule &M);
 private:
   CanType computeLoweredRValueType(TypeExpansionContext context,
                                    AbstractionPattern origType,

lib/SIL/IR/SILFunctionType.cpp

@@ -2078,13 +2078,9 @@ static CanSILFunctionType getSILFunctionType(
     // Lower in the context of the closure. Since the set of captures is a
     // private contract between the closure and its enclosing context, we
     // don't need to keep its capture types opaque.
-    auto expansion = TypeExpansionContext::maximal(
-        constant->getAnyFunctionRef()->getAsDeclContext(), false);
-    // ...unless it's inlinable, in which case it might get inlined into
-    // some place we need to keep opaque types opaque.
-    if (constant->isSerialized())
-      expansion = TypeExpansionContext::minimal();
-    lowerCaptureContextParameters(TC, *constant, genericSig, expansion, inputs);
+    lowerCaptureContextParameters(TC, *constant, genericSig,
+                                  TC.getCaptureTypeExpansionContext(*constant),
+                                  inputs);
   }
 
   auto calleeConvention = ParameterConvention::Direct_Unowned;

lib/SIL/IR/TypeLowering.cpp

@@ -3862,6 +3862,19 @@ TypeConverter::getConstantAbstractionPattern(SILDeclRef constant) {
   return None;
 }
 
+TypeExpansionContext
+TypeConverter::getCaptureTypeExpansionContext(SILDeclRef constant) {
+  auto found = CaptureTypeExpansionContexts.find(constant);
+  if (found != CaptureTypeExpansionContexts.end()) {
+    return found->second;
+  }
+  // Insert a minimal type expansion context into the cache, so that further
+  // attempts to change it raise an error.
+  auto minimal = TypeExpansionContext::minimal();
+  CaptureTypeExpansionContexts.insert({constant, minimal});
+  return minimal;
+}
+
 void TypeConverter::setAbstractionPattern(AbstractClosureExpr *closure,
                                           AbstractionPattern pattern) {
   auto existing = ClosureAbstractionPatterns.find(closure);
@@ -3873,6 +3886,31 @@ void TypeConverter::setAbstractionPattern(AbstractClosureExpr *closure,
   }
 }
 
+void TypeConverter::setCaptureTypeExpansionContext(SILDeclRef constant,
+                                                   SILModule &M) {
+  if (!hasLoweredLocalCaptures(constant)) {
+    return;
+  }
+  
+  TypeExpansionContext context = constant.isSerialized()
+    ? TypeExpansionContext::minimal()
+    : TypeExpansionContext::maximal(constant.getAnyFunctionRef()->getAsDeclContext(),
+                                    M.isWholeModule());
+
+  auto existing = CaptureTypeExpansionContexts.find(constant);
+  if (existing != CaptureTypeExpansionContexts.end()) {
+    assert(existing->second == context
+     && "closure shouldn't be emitted with different capture type expansion contexts");
+  } else {
+    // Lower in the context of the closure. Since the set of captures is a
+    // private contract between the closure and its enclosing context, we
+    // don't need to keep its capture types opaque.
+    // The exception is if it's inlinable, in which case it might get inlined into
+    // some place we need to keep opaque types opaque.
+    CaptureTypeExpansionContexts.insert({constant, context});
+  }
+}
+
 static void countNumberOfInnerFields(unsigned &fieldsCount, TypeConverter &TC,
                                      SILType Ty,
                                      TypeExpansionContext expansion) {

lib/SILGen/SILGen.cpp

@@ -1421,6 +1421,8 @@ void SILGenModule::emitAbstractFuncDecl(AbstractFunctionDecl *AFD) {
 }
 
 void SILGenModule::emitFunction(FuncDecl *fd) {
+  Types.setCaptureTypeExpansionContext(SILDeclRef(fd), M);
+
   SILDeclRef::Loc decl = fd;
 
   emitAbstractFuncDecl(fd);

lib/SILGen/SILGenApply.cpp

@@ -1142,6 +1142,8 @@ class SILGenApply : public Lowering::ExprVisitor<SILGenApply> {
     }
 
     auto captureInfo = SGF.SGM.Types.getLoweredLocalCaptures(constant);
+    SGF.SGM.Types.setCaptureTypeExpansionContext(constant, SGF.SGM.M);
+    
     if (afd->getDeclContext()->isLocalContext() &&
         !captureInfo.hasGenericParamCaptures())
       subs = SubstitutionMap();
@@ -1207,6 +1209,9 @@ class SILGenApply : public Lowering::ExprVisitor<SILGenApply> {
   }
 
   void visitAbstractClosureExpr(AbstractClosureExpr *e) {
+    SILDeclRef constant(e);
+
+    SGF.SGM.Types.setCaptureTypeExpansionContext(constant, SGF.SGM.M);
     // Emit the closure body.
     SGF.SGM.emitClosure(e);
 
@@ -1219,7 +1224,6 @@ class SILGenApply : public Lowering::ExprVisitor<SILGenApply> {
 
     // A directly-called closure can be emitted as a direct call instead of
     // really producing a closure object.
-    SILDeclRef constant(e);
 
     auto captureInfo = SGF.SGM.M.Types.getLoweredLocalCaptures(constant);
 

lib/SILGen/SILGenExpr.cpp

@@ -2545,6 +2545,7 @@ RValue RValueEmitter::visitAbstractClosureExpr(AbstractClosureExpr *e,
   if (auto contextOrigType = C.getAbstractionPattern()) {
     SGF.SGM.Types.setAbstractionPattern(e, *contextOrigType);
   }
+  SGF.SGM.Types.setCaptureTypeExpansionContext(SILDeclRef(e), SGF.SGM.M);
 
   // Emit the closure body.
   SGF.SGM.emitClosure(e);

lib/SILGen/SILGenFunction.cpp

@@ -441,7 +441,8 @@ SILGenFunction::emitClosureValue(SILLocation loc, SILDeclRef constant,
                                  SubstitutionMap subs,
                                  bool alreadyConverted) {
   auto loweredCaptureInfo = SGM.Types.getLoweredLocalCaptures(constant);
-
+  SGM.Types.setCaptureTypeExpansionContext(constant, SGM.M);
+  
   auto constantInfo = getConstantInfo(getTypeExpansionContext(), constant);
   SILValue functionRef = emitGlobalFunctionRef(loc, constant, constantInfo);
   SILType functionTy = functionRef->getType();

test/SILGen/Inputs/opaque_result_type_captured_wmo_2.swift

@@ -0,0 +1,27 @@
+
+public protocol P {}
+
+struct PImpl: P {}
+
+public struct Wrapper<T: P, U>: P {
+  public init(value: T, extra: U) {}
+}
+
+private struct Burrito {}
+
+extension P {
+  @inlinable
+  public func wrapped<U>(extra: U) -> Wrapper<Self, U> {
+    return Wrapper(value: self, extra: extra)
+  }
+
+  public func burritoed() -> some P {
+    return wrapped(extra: Burrito())
+  }
+}
+
+public class Butz<T: P> {
+  init(_: T) {}
+}
+
+

test/SILGen/opaque_result_type_captured_wmo.swift

@@ -0,0 +1,10 @@
+// RUN: %target-swift-emit-silgen -disable-availability-checking -verify -wmo %s %S/Inputs/opaque_result_type_captured_wmo_2.swift
+func foo(s: String?) {
+  let x = PImpl()
+    .burritoed()
+    .wrapped(extra: 1)
+
+  let butz = Butz(x)
+
+  s.map { print("\($0) \(butz)") }
+}