SILGen: Provide a reabstracted writeback for inout parameters emitted in a more abstract closure literal context.

When we emit closure literals in the context of a generic parameter, the parameters of the
closure may be more abstract than they would normally be for the closure's abstract type.
We handled this by reabstracting value parameters to their usual substituted representation
in the prolog, but neglected to do so for `inout` parameters too. In this case, we need to take
the initial value, reabstract it to provide a local variable with the substituted representation,
and then take the final value, reabstract it back, and write back to the original inout parameter
on function exit.

Fixes rdar://111563642.

Author: jckarter

lib/SILGen/SILGenProlog.cpp

@@ -130,6 +130,39 @@ struct LoweredParamGenerator {
   }
 };
 
+struct WritebackReabstractedInoutCleanup final : Cleanup {
+  SILValue OrigAddress, SubstAddress;
+  AbstractionPattern OrigTy;
+  CanType SubstTy;
+  WritebackReabstractedInoutCleanup(SILValue origAddress, SILValue substAddress,
+                                    AbstractionPattern origTy,
+                                    CanType substTy)
+      : OrigAddress(origAddress), SubstAddress(substAddress),
+        OrigTy(origTy), SubstTy(substTy)
+  {}
+  
+  void emit(SILGenFunction &SGF, CleanupLocation l, ForUnwind_t forUnwind)
+  override {
+    Scope s(SGF.Cleanups, l);
+    // Load the final local value coming in.
+    auto mv = SGF.emitLoad(l, SubstAddress,
+                           SGF.getTypeLowering(SubstAddress->getType()),
+                           SGFContext(), IsTake);
+    // Reabstract the value back to the original representation.
+    mv = SGF.emitSubstToOrigValue(l, mv.ensurePlusOne(SGF, l),
+                                  OrigTy, SubstTy);
+    // Write it back to the original inout parameter.
+    SGF.B.createStore(l, mv.forward(SGF), OrigAddress,
+                      StoreOwnershipQualifier::Init);
+  }
+  
+  void dump(SILGenFunction&) const override {
+    llvm::errs() << "WritebackReabstractedInoutCleanup\n";
+    OrigAddress->print(llvm::errs());
+    SubstAddress->print(llvm::errs());
+  }
+};
+
 class EmitBBArguments : public CanTypeVisitor<EmitBBArguments,
                                               /*RetTy*/ ManagedValue,
                                               /*ArgTys...*/ AbstractionPattern,
@@ -255,6 +288,23 @@ class EmitBBArguments : public CanTypeVisitor<EmitBBArguments,
       if (mv.getType().isMoveOnly() && !mv.getType().isMoveOnlyWrapped())
         mv = SGF.B.createMarkMustCheckInst(
             loc, mv, MarkMustCheckInst::CheckKind::ConsumableAndAssignable);
+
+      // If the value needs to be reabstracted, set up a shadow copy with
+      // writeback here.
+      if (argType.getASTType() != mv.getType().getASTType()) {
+        // Load the value coming in.
+        auto origBuf = mv.getValue();
+        mv = SGF.emitLoad(loc, origBuf, SGF.getTypeLowering(mv.getType()), SGFContext(), IsTake);
+        // Reabstract the value if necessary.
+        mv = SGF.emitOrigToSubstValue(loc, mv.ensurePlusOne(SGF, loc), orig, t);
+        // Store the value to a local buffer.
+        auto substBuf = SGF.emitTemporaryAllocation(loc, argType);
+        SGF.B.createStore(loc, mv.forward(SGF), substBuf, StoreOwnershipQualifier::Init);
+        // Introduce a writeback to put the final value back in the inout.
+        SGF.Cleanups.pushCleanup<WritebackReabstractedInoutCleanup>(origBuf, substBuf, orig, t);
+        mv = ManagedValue::forLValue(substBuf);
+      }
+
       return mv;
     }
 

test/SILGen/inout_reabstracted_closure_literal.swift

@@ -0,0 +1,66 @@
+// RUN: %target-swift-emit-silgen %s
+
+func bar<T>(f: (inout T) throws -> Void) {}
+
+func condition() -> Bool { fatalError() }
+func error() -> Error { fatalError() }
+
+func foo(x: @escaping () -> Void) {
+    // CHECK-LABEL: sil private [ossa] @${{.*}}3foo{{.*}}U_ :
+    // CHECK:       bb0(
+    // CHECK-SAME:    [[ARG:%.*]] = $*@callee_guaranteed @substituted <τ_0_0> () -> @out τ_0_0 for <()>,
+    // CHECK-SAME:    [[CAPTURE:%.*]] = @closureCapture @guaranteed $@callee_guaranteed () -> ()
+
+    // Reabstract the initial value of the inout parameter
+    // CHECK:         [[INITIAL:%.*]] = load [take] [[ARG]]
+    // CHECK:         [[INITIAL_CONV:%.*]] = convert_function [[INITIAL]]
+    // CHECK:         [[THUNK:%.*]] = function_ref @{{.*}}_TR
+    // CHECK:         [[INITIAL_REAB:%.*]] = partial_apply [callee_guaranteed] [[THUNK]]([[INITIAL_CONV]])
+    // CHECK:         [[REAB_ARG:%.*]] = alloc_stack
+    // CHECK:         store [[INITIAL_REAB]] to [[REAB_ARG]]
+    bar {
+        // Read from the reabstracted buffer
+        //  CHECK: [[ACCESS:%.*]] = begin_access [read] [unknown] [[REAB_ARG]]
+        //  CHECK: [[COPY]] = load [copy] [[ACCESS]]
+        _ = $0 as () -> Void
+
+        // Assign to the reabstracted buffer
+        $0 = x
+
+        // Writeback occurs no matter how we exit the function
+        // CHECK:   cond_br {{%.*}}, [[THEN1:bb[0-9]+]], [[ELSE1:bb[0-9]+]]
+        if condition() {
+            // Throwing:
+            // CHECK: [[THEN1]]:
+            // Take final reabstracted value from the reabstracted buffer
+            // CHECK:   [[FINAL_REAB:%.*]] = load [take] [[REAB_ARG]]
+            // Reabstract back to original representation and store back to
+            // original buffer
+            // CHECK:   [[THUNK_ORIG:%.*]] = function_ref @{{.*}}_TR
+            // CHECK:   [[FINAL_CONV:%.*]] = partial_apply [callee_guaranteed] [[THUNK_ORIG]]([[FINAL_REAB]])
+            // CHECK:   [[FINAL:%.*]] = convert_function [[FINAL_CONV]]
+            // CHECK:   store [[FINAL]] to [[ARG]]
+            // CHECK:   throw
+            throw error()
+        // CHECK: [[ELSE1]]:
+        // CHECK:   cond_br {{%.*}}, [[THEN2:bb[0-9]+]], [[ELSE2:bb[0-9]+]]
+        } else if condition() {
+            // Explicit return
+            // CHECK: [[THEN1]]:
+            // CHECK:   br [[EPILOG:bb[0-9]+]]
+            return
+        } else {
+            // Fall through
+        }
+        // CHECK: [[EPILOG]]:
+        // Take final reabstracted value from the reabstracted buffer
+        // CHECK:   [[FINAL_REAB:%.*]] = load [take] [[REAB_ARG]]
+        // Reabstract back to original representation and store back to
+        // original buffer
+        // CHECK:   [[THUNK_ORIG:%.*]] = function_ref @{{.*}}_TR
+        // CHECK:   [[FINAL_CONV:%.*]] = partial_apply [callee_guaranteed] [[THUNK_ORIG]]([[FINAL_REAB]])
+        // CHECK:   [[FINAL:%.*]] = convert_function [[FINAL_CONV]]
+        // CHECK:   store [[FINAL]] to [[ARG]]
+        // CHECK:   return
+    }
+}