Implement reabstraction of inouts in witness and reabstraction thunks

include/swift/Basic/DiverseStack.h

@@ -366,10 +366,14 @@ template <class T> class DiverseStackImpl : private DiverseStackBase {
   void pop(stable_iterator stable_iter) {
     iterator iter = find(stable_iter);
     checkIterator(iter);
+#ifndef NDEBUG
     while (Begin != iter.Ptr) {
       pop();
       checkIterator(iter);
     }
+#else
+    Begin = iter.Ptr;
+#endif
   }
 };
 

lib/SILGen/Cleanup.cpp

@@ -75,45 +75,52 @@ void CleanupManager::popAndEmitCleanup(CleanupHandle handle,
 
 void CleanupManager::emitCleanups(CleanupsDepth depth, CleanupLocation loc,
                                   ForUnwind_t forUnwind, bool popCleanups) {
-  auto begin = stack.stable_begin();
-  while (begin != depth) {
-    auto iter = stack.find(begin);
-
+  auto cur = stack.stable_begin();
+#ifndef NDEBUG
+  auto topOfStack = cur;
+#endif
+  while (cur != depth) {
+    // Copy the cleanup off the stack if it needs to be emitted.
+    // This is necessary both because we might need to pop the cleanup and
+    // because the cleanup might push other cleanups that will invalidate
+    // references onto the stack.
+    auto iter = stack.find(cur);
     Cleanup &stackCleanup = *iter;
+    Optional<CleanupBuffer> copiedCleanup;
+    if (stackCleanup.isActive() && SGF.B.hasValidInsertionPoint()) {
+      copiedCleanup.emplace(stackCleanup);
+    }
 
-    // Copy it off the cleanup stack in case the cleanup pushes a new cleanup
-    // and the backing storage is re-allocated.
-    CleanupBuffer buffer(stackCleanup);
-    Cleanup &cleanup = buffer.getCopy();
-
-    // Advance stable iterator.
-    begin = stack.stabilize(++iter);
+    // Advance the iterator.
+    cur = stack.stabilize(++iter);
 
-    // Pop now.
-    if (popCleanups)
+    // Pop now if that was requested.
+    if (popCleanups) {
       stack.pop();
 
-    if (cleanup.isActive() && SGF.B.hasValidInsertionPoint())
-      cleanup.emit(SGF, loc, forUnwind);
+#ifndef NDEBUG
+      topOfStack = stack.stable_begin();
+#endif
+    }
 
-    stack.checkIterator(begin);
+    // Emit the cleanup.
+    if (copiedCleanup) {
+      copiedCleanup->getCopy().emit(SGF, loc, forUnwind);
+#ifndef NDEBUG
+      if (hasAnyActiveCleanups(stack.stable_begin(), topOfStack)) {
+        copiedCleanup->getCopy().dump(SGF);
+        llvm_unreachable("cleanup left active cleanups on stack");
+      }
+#endif
+    }
+
+    stack.checkIterator(cur);
   }
 }
 
-/// Leave a scope, with all its cleanups.
+/// Leave a scope, emitting all the cleanups that are currently active.
 void CleanupManager::endScope(CleanupsDepth depth, CleanupLocation loc) {
-  stack.checkIterator(depth);
-
-  // FIXME: Thread a branch through the cleanups if there are any active
-  // cleanups and we have a valid insertion point.
-
-  if (!::hasAnyActiveCleanups(stack.begin(), stack.find(depth))) {
-    return;
-  }
-
-  // Iteratively mark cleanups dead and pop them.
-  // Maybe we'd get better results if we marked them all dead in one shot?
-  emitCleanups(depth, loc, NotForUnwind);
+  emitCleanups(depth, loc, NotForUnwind, /*popCleanups*/ true);
 }
 
 bool CleanupManager::hasAnyActiveCleanups(CleanupsDepth from,

lib/SILGen/Cleanup.h

@@ -136,8 +136,7 @@ class LLVM_LIBRARY_VISIBILITY CleanupManager {
 
   void popTopDeadCleanups(CleanupsDepth end);
   void emitCleanups(CleanupsDepth depth, CleanupLocation l,
-                    ForUnwind_t forUnwind,
-                    bool popCleanups=true);
+                    ForUnwind_t forUnwind, bool popCleanups);
   void endScope(CleanupsDepth depth, CleanupLocation l);
 
   Cleanup &initCleanup(Cleanup &cleanup, size_t allocSize, CleanupState state);

lib/SILGen/SILGenLValue.cpp

@@ -47,6 +47,8 @@ struct LValueWritebackCleanup : Cleanup {
 
   void emit(SILGenFunction &SGF, CleanupLocation loc,
             ForUnwind_t forUnwind) override {
+    FullExpr scope(SGF.Cleanups, loc);
+
     // TODO: honor forUnwind!
     auto &evaluation = *SGF.FormalEvalContext.find(Depth);
     assert(evaluation.getKind() == FormalAccess::Exclusive);

lib/SILGen/SILGenPoly.cpp

@@ -827,6 +827,53 @@ static CanType getSingleTupleElement(CanType type) {
 }
 
 namespace {
+  class TranslateIndirect : public Cleanup {
+    AbstractionPattern InputOrigType, OutputOrigType;
+    CanType InputSubstType, OutputSubstType;
+    SILValue Input, Output;
+
+  public:
+    TranslateIndirect(AbstractionPattern inputOrigType, CanType inputSubstType,
+                      AbstractionPattern outputOrigType, CanType outputSubstType,
+                      SILValue input, SILValue output)
+      : InputOrigType(inputOrigType), OutputOrigType(outputOrigType),
+        InputSubstType(inputSubstType), OutputSubstType(outputSubstType),
+        Input(input), Output(output) {
+      assert(input->getType().isAddress());
+      assert(output->getType().isAddress());
+    }
+
+    void emit(SILGenFunction &SGF, CleanupLocation loc,
+              ForUnwind_t forUnwind) override {
+      FullExpr scope(SGF.Cleanups, loc);
+
+      // Re-assert ownership of the input value.
+      auto inputMV = SGF.emitManagedBufferWithCleanup(Input);
+
+      // Set up an initialization of the output buffer.
+      auto &outputTL = SGF.getTypeLowering(Output->getType());
+      auto outputInit = SGF.useBufferAsTemporary(Output, outputTL);
+
+      // Transform into the output buffer.
+      auto mv = SGF.emitTransformedValue(loc, inputMV,
+                                         InputOrigType, InputSubstType,
+                                         OutputOrigType, OutputSubstType,
+                                         SGFContext(outputInit.get()));
+      emitForceInto(SGF, loc, mv, *outputInit);
+
+      // Disable the cleanup; we've kept our promise to leave the inout
+      // initialized.
+      outputInit->getManagedAddress().forward(SGF);
+    }
+
+    void dump(SILGenFunction &SGF) const override {
+      llvm::errs() << "TranslateIndirect("
+        << InputOrigType << ", " << InputSubstType << ", "
+        << OutputOrigType << ", " << OutputSubstType << ", "
+        << Output << ", " << Input << ")\n";
+    }
+  };
+
   class TranslateArguments {
     SILGenFunction &SGF;
     SILLocation Loc;
@@ -1356,13 +1403,49 @@ namespace {
                                 outputSubstType, input);
         return;
       case ParameterConvention::Indirect_Inout: {
-        // If it's inout, we need writeback.
-        llvm::errs() << "inout writeback in abstraction difference thunk "
-                        "not yet implemented\n";
-        llvm::errs() << "input value ";
-        input.getValue()->dump();
-        llvm::errs() << "output type " << SGF.getSILType(result) << "\n";
-        abort();
+        inputOrigType = inputOrigType.getWithoutSpecifierType();
+        inputSubstType = inputSubstType.getWithoutSpecifierType();
+        outputOrigType = outputOrigType.getWithoutSpecifierType();
+        outputSubstType = outputSubstType.getWithoutSpecifierType();
+
+        // Create a temporary of the right type.
+        auto &temporaryTL = SGF.getTypeLowering(result.getType());
+        auto temporary = SGF.emitTemporary(Loc, temporaryTL);
+
+        // Take ownership of the input value.  This leaves the input l-value
+        // effectively uninitialized, but we'll push a cleanup that will put
+        // a value back into it.
+        FullExpr scope(SGF.Cleanups, CleanupLocation::get(Loc));
+        auto ownedInput =
+          SGF.emitManagedBufferWithCleanup(input.getLValueAddress());
+
+        // Translate the input value into the temporary.
+        translateSingleInto(inputOrigType, inputSubstType,
+                            outputOrigType, outputSubstType,
+                            ownedInput, *temporary);
+
+        // Forward the cleanup on the temporary.  We're about to push a new
+        // cleanup that will re-assert ownership of this value.
+        auto temporaryAddr = temporary->getManagedAddress().forward(SGF);
+
+        // Leave the scope in which we did the forward translation.  This
+        // ensures that the value in the input buffer is destroyed
+        // immediately rather than (potentially) arbitrarily later
+        // at a point where we want to put new values in the input buffer.
+        scope.pop();
+
+        // Push the cleanup to perform the reverse translation.  This cleanup
+        // asserts ownership of the value of the temporary.
+        SGF.Cleanups.pushCleanup<TranslateIndirect>(outputOrigType,
+                                                    outputSubstType,
+                                                    inputOrigType,
+                                                    inputSubstType,
+                                                    temporaryAddr,
+                                                    input.getLValueAddress());
+
+        // Add the temporary as an l-value argument.
+        Outputs.push_back(ManagedValue::forLValue(temporaryAddr));
+        return;
       }
       case ParameterConvention::Indirect_In: {
         if (SGF.silConv.useLoweredAddresses()) {

lib/SILGen/Scope.cpp

@@ -120,8 +120,6 @@ RValue Scope::popPreservingValue(RValue &&rv) {
 }
 
 void Scope::popImpl() {
-  SmallVector<SILValue, 16> cleanupsToPropagateToOuterScope;
-
   cleanups.stack.checkIterator(depth);
   cleanups.stack.checkIterator(cleanups.innermostScope);
   assert(cleanups.innermostScope == depth && "popping scopes out of order");

test/SILGen/witnesses.swift

@@ -522,3 +522,37 @@ protocol EscapingReq {
 struct EscapingCovariance: EscapingReq {
   func f(_: (Int) -> Int) { }
 }
+
+protocol InoutFunctionReq {
+  associatedtype T
+  func updateFunction(x: inout () -> T)
+}
+
+// CHECK-LABEL: sil private [transparent] [thunk] @$S9witnesses13InoutFunctionVAA0bC3ReqA2aDP06updateC01xy1TQzycz_tFTW
+// CHECK:      bb0(%0 : @trivial $*@callee_guaranteed () -> @out (), %1 : @trivial $*InoutFunction):
+// CHECK-NEXT:   [[TEMP:%.*]] = alloc_stack $@callee_guaranteed () -> ()
+//   Reabstract the contents of the inout argument into the temporary.
+// CHECK-NEXT:   [[OLD_FN:%.*]] = load [take] %0
+// CHECK-NEXT:   // function_ref
+// CHECK-NEXT:   [[THUNK:%.*]] = function_ref @$SytIegr_Ieg_TR
+// CHECK-NEXT:   [[THUNKED_OLD_FN:%.*]] = partial_apply [callee_guaranteed] [[THUNK]]([[OLD_FN]])
+// CHECK-NEXT:   store [[THUNKED_OLD_FN]] to [init] [[TEMP]] :
+//   Call the function.
+// CHECK-NEXT:   [[SELF:%.*]] = load [trivial] %1 : $*InoutFunction
+// CHECK-NEXT:   // function_ref
+// CHECK-NEXT:   [[T0:%.*]] = function_ref @$S9witnesses13InoutFunctionV06updateC01xyyycz_tF :
+// CHECK-NEXT:   apply [[T0]]([[TEMP]], [[SELF]])
+// CHECK-NEXT:   [[TUPLE:%.*]] = tuple ()
+//   Reabstract the contents of the temporary back into the inout argument.
+// CHECK-NEXT:   [[NEW_FN:%.*]] = load [take] [[TEMP]]
+// CHECK-NEXT:   // function_ref
+// CHECK-NEXT:   [[THUNK:%.*]] = function_ref @$SIeg_ytIegr_TR
+// CHECK-NEXT:   [[THUNKED_NEW_FN:%.*]] = partial_apply [callee_guaranteed] [[THUNK]]([[NEW_FN]])
+// CHECK-NEXT:   store [[THUNKED_NEW_FN]] to [init] %0 :
+// CHECK-NEXT:   dealloc_stack [[TEMP]]
+// CHECK-NEXT:   return [[TUPLE]]
+// CHECK-LABEL:  } // end sil function '$S9witnesses13InoutFunctionVAA0bC3ReqA2aDP06updateC01xy1TQzycz_tFTW'
+
+struct InoutFunction : InoutFunctionReq {
+  func updateFunction(x: inout () -> ()) {}
+}