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

TypeConverter doesn't know by itself what SILModule it's currently lowering on
behalf of, so the existing code forming the TypeExpansionContext for opaque types
incorrectly set the isWholeModule flag to always false. This created a miscompile
when a public API contained a closure that captured a value involving private types
from another file in the same module because of mismatched type expansion contexts
inside and outside the closure. Fixes rdar://93821679

Author: jckarter

include/swift/SIL/TypeLowering.h

@@ -752,6 +752,8 @@ class TypeConverter {
 
   llvm::DenseMap<AbstractClosureExpr *, Optional<AbstractionPattern>>
     ClosureAbstractionPatterns;
+  llvm::DenseMap<SILDeclRef, TypeExpansionContext>
+    CaptureTypeExpansionContexts;
 
   CanAnyFunctionType makeConstantInterfaceType(SILDeclRef constant);
 
@@ -1156,6 +1158,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.
   ///
@@ -1165,6 +1168,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

@@ -3623,6 +3623,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);
@@ -3634,6 +3647,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

@@ -1139,6 +1139,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();
@@ -1204,6 +1206,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);
 
@@ -1216,7 +1221,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

@@ -2542,6 +2542,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)") }
+}