Merge pull request #67239 from jckarter/reabstracted-closure-literal-inout-reabstraction

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

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
+    }
+}