[sil-mem2reg] Add a simple scope data structure and use it to simplify some code.

We have for a long time talked about creating a scope like data structure for
use in the SILOptimizer. The discussion was whether or not to reuse the
infrastructure in SILGen that does this already. There were concerns about doing
so since the code in the SILOptimizer and SILGen can work differently.

With that in mind, I added a small AssertingScope class and built on top of that
a composition SIL level class called SILOptScope that one can use to add various
cleanups. One is able to both destructively pop at end of scope and pop along
early exits.

At an implementation level, I kept it simple and:

1. Represented a scope as a stack of Optional<Cleanup> which are just a wrapper
   around a std::function. The Optional is so that we can invalidate a cleanup.
2. Based all of these scopes around the idea that the user of the scope must
   invalidate the scope by hand. If not, the scope object will assert at the end
   of its RAII scope.
3. Rather than creating a whole class hierarchy, I just used std::function
   closures to keep things simple.

Author: gottesmm

include/swift/SILOptimizer/Utils/ScopeOptUtils.h

@@ -0,0 +1,284 @@
+//===--- ScopeOptUtils.h --------------------------------------------------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2014 - 2021 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+///
+/// A simple scope like construct for use in the SILOptimizer.
+///
+//===----------------------------------------------------------------------===//
+
+#ifndef SWIFT_SILOPTIMIZER_UTILS_SCOPEOPTUTILS_H
+#define SWIFT_SILOPTIMIZER_UTILS_SCOPEOPTUTILS_H
+
+#include "swift/Basic/LLVM.h"
+#include "swift/SIL/SILBuilder.h"
+#include "swift/SIL/SILValue.h"
+#include "llvm/ADT/Optional.h"
+#include "llvm/ADT/SmallVector.h"
+#include <functional>
+
+namespace swift {
+namespace siloptimizer {
+
+template <typename Result, typename... Args>
+struct Cleanup {
+  std::function<Result(Args...)> func;
+};
+
+/// A simple stack data structure that manages a stack of std::function cleanups
+/// that can be invalidated by index.
+///
+/// It is called asserting, since we force our user to explicitly pop the scope
+/// and abort otherwise.
+template <typename Result, typename... Args>
+class AssertingScope {
+  using ScopeCleanup = Cleanup<Result, Args...>;
+  SmallVector<Optional<ScopeCleanup>, 64> cleanups;
+  bool didPop = false;
+
+protected:
+  ~AssertingScope() { assert(didPop && "Did not pop scope?!"); }
+
+  AssertingScope() {}
+
+  /// Push a new SIL function based cleanup. Meant to be very
+  /// lightweight. Returns the index of the cleanup.
+  unsigned pushScope(std::function<Result(Args...)> func) {
+    assert(!didPop && "Can not reuse scope once used?!");
+    unsigned index = cleanups.size();
+    cleanups.emplace_back(ScopeCleanup{func});
+    return index;
+  }
+
+  /// Invalidate the cleanup at the specific index. Asserts on index too big and
+  /// on Optional has value.
+  void invalidateCleanup(unsigned cleanupIndex) {
+    assert(!didPop && "Should not invalidate once popped?!");
+    assert(cleanupIndex < cleanups.size());
+    assert(cleanups[cleanupIndex].hasValue());
+    cleanups[cleanupIndex] = None;
+  }
+
+  /// Pop the scope, running all non-Optional cleanups, and setting didPop. Once
+  /// this is done the scope has been consumed and no longer has any state.
+  ///
+  /// Can only be done once the scope has been moved.
+  void pop(Args... args) && {
+    while (cleanups.size()) {
+      auto val = cleanups.pop_back_val();
+      if (val) {
+        val->func(std::forward<Args>(args)...);
+      }
+    }
+    didPop = true;
+  }
+
+  /// Can not be called once the scope is moved.
+  void nonDestructivePop(Args... args) const {
+    for (auto val : llvm::reverse(cleanups)) {
+      if (val) {
+        val->func(std::forward<Args>(args)...);
+      }
+    }
+  }
+};
+
+class ScopedValue;
+class SILOptScope;
+
+/// Part of the SILRAIIScope interface that is known to ScopedValue.
+class SILOptScopeBase : public AssertingScope<void, SILInstruction *> {
+protected:
+  SILOptScopeBase() {}
+
+public:
+  using Index = unsigned;
+
+  /// Destructive pop at end of scope.
+  void popAtEndOfScope(SILInstruction *insertPt) && {
+    std::move(*this).pop(insertPt);
+  }
+
+  /// Non-destructive pop along a control flow path that is serving as an early
+  /// exit.
+  void popAtEarlyExit(SILInstruction *insertPt) const {
+    nonDestructivePop(insertPt);
+  }
+
+private:
+  friend class ScopedValue;
+
+  /// Invalidate the passed in scoped value destructively.
+  void invalidateValue(ScopedValue &&value);
+};
+
+/// A light weight wrapper around a SILValue that also contains information
+/// about a Cleanup closure that ends the lifetime of the given value.
+///
+/// In order to make it really easy to access the underlying value, we define
+/// operator * and operator-> as returning the underlying SILValue. It also
+/// defines explicit operator bool() so one should be able to use if statements
+/// to perform conditional initialization by returning an empty ScopedValue().
+///
+/// If one realizes that one does not actually need to cleanup a value, one can
+/// destructively forward the ScopedValue by writing:
+///
+///   std::move(scopedValue).forward();
+///
+/// This will cause the cleanup in the SILOptScopeBase to be invalidated and
+/// thus will not be emitted when the scope is popped.
+class ScopedValue {
+  friend class SILOptScopeBase;
+  friend class SILOptScope;
+
+  /// A back pointer to the scope we are managed by so that we can invalidate
+  /// the cleanup associated with us. Also used to invalidate the underlying
+  /// ScopedValue by setting this to nullptr.
+  SILOptScopeBase *scope;
+
+  /// The index in our parent scope object of our cleanup. This is used to
+  /// enable invalidation.
+  Optional<unsigned> scopeIndex;
+
+  /// The underlying SILValue that we are working with.
+  SILValue value;
+
+  /// A scoped value that has a cleanup associated with it at the index \p
+  /// scopeIndex within \p scope.
+  ScopedValue(SILOptScopeBase &scope, SILValue value, unsigned scopeIndex)
+      : scope(&scope), scopeIndex(scopeIndex), value(value) {}
+
+  /// A scoped value that does not have an associated cleanup in the scope.
+  ScopedValue(SILOptScopeBase &scope, SILValue value)
+      : scope(&scope), scopeIndex(None), value(value) {}
+
+public:
+  /// Default empty constructor. Should be used in combination with operator
+  /// bool() to perform conditional initialization.
+  ScopedValue() : scope(nullptr), scopeIndex(), value() {}
+
+  /// Explicit operator that says if this is a valid scoped Value.
+  ///
+  /// Allows for conditional initialization via return value:
+  ///
+  ///   if (auto value = getScopedValue(...)) {
+  ///     ...
+  ///   }
+  explicit operator bool() const { return isValid(); }
+
+  /// Check just for the validity of the value, this doesn't
+  bool isValid() const { return scope && value; }
+
+  /// Returns true if this scoped value has a cleanup associated with it that
+  /// can be invalidated.
+  ///
+  /// Example: Result of SILOptScope::copyValue would return true, but result of
+  /// a SILOptScope::borrowValue would return false since we do not allow for
+  /// borrows to be cancelled.
+  bool hasInvalidatableCleanup() const {
+    assert(isValid());
+    return bool(scopeIndex);
+  }
+
+  /// Operator to access internal SILValue.
+  SILValue operator*() const {
+    assert(isValid());
+    return value;
+  }
+
+  /// Operator to call methods on internal SILValue.
+  SILValue operator->() const {
+    assert(isValid());
+    return value;
+  }
+
+  /// If this scope has a scoped index, invalidate the scoped value so that its
+  /// cleanup is not emitted at end of scope.
+  void invalidate() && {
+    assert(isValid());
+    // If we have an invalidated cleanup, do the invalidation.
+    if (hasInvalidatableCleanup()) {
+      scope->invalidateValue(std::move(*this));
+    }
+    // No matter what we do, make it so that if this object is used again, we
+    // get an assert.
+    scope = nullptr;
+    scopeIndex.reset();
+  }
+};
+
+inline void SILOptScopeBase::invalidateValue(ScopedValue &&value) {
+  if (value.scopeIndex)
+    invalidateCleanup(*value.scopeIndex);
+}
+
+/// A scope like stack data structure that one can use to manage cleanups for
+/// SILValues using light weight ScopedValue wrappers. This is done by using the
+/// copyValue/borrowValue methods. These copy/borrow the value and then push a
+/// destroy_value/end_borrow cleanup and return a ScopedValue containing that
+/// the user can then manipulate the value without needing to worry about the
+/// cleanup. If the user then decides a specific ScopedValue should be able to
+/// be cancelled, we allow for the ScopedValue to be invalidated
+class SILOptScope : public SILOptScopeBase {
+public:
+  SILOptScope() {}
+
+  /// Create a destroy_value cleanup for \p value. Returns a scoped value for \p
+  /// value with a live cleanup.
+  ScopedValue pushDestroyValue(SILValue value) {
+    // We need to capture value by value, not by address.
+    auto func = std::function<void(SILInstruction *)>(
+        [innerValue = value](SILInstruction *insertPt) {
+          SILBuilderWithScope builder(insertPt);
+          builder.emitDestroyValue(RegularLocation::getAutoGeneratedLocation(),
+                                   innerValue);
+        });
+    return {*this, value, pushScope(func)};
+  }
+
+  /// Copy \p value at \p insertPt and push a destroy_value cleanup. Return the
+  /// new copied value as a ScopedValue with cleanup.
+  ScopedValue copyValue(SILInstruction *insertPt, SILValue value) {
+    SILValue copy = SILBuilderWithScope(insertPt).emitCopyValueOperation(
+        insertPt->getLoc(), value);
+    return pushDestroyValue(copy);
+  }
+
+  ScopedValue pushEndBorrow(SILValue value) {
+    // We need to capture value by value, not by address.
+    auto func = std::function<void(SILInstruction *)>(
+        [innerValue = value](SILInstruction *insertPt) {
+          SILBuilderWithScope builder(insertPt);
+          builder.emitEndBorrowOperation(
+              RegularLocation::getAutoGeneratedLocation(), innerValue);
+        });
+    // We do not return the index of our end_borrow since an end_borrow should
+    // never have its cleanup cancelled.
+    pushScope(func);
+    return {*this, value};
+  }
+
+  /// Borrow \p value at \p insertPt and push an end_borrow cleanup. Return the
+  /// new borrowed value as a ScopedValue without cleanup. The cleanup is not
+  /// returned since end_borrows should not be cancellable.
+  ScopedValue borrowValue(SILInstruction *insertPt, SILValue value) {
+    if (!insertPt->getFunction()->hasOwnership() ||
+        value.getOwnershipKind().isCompatibleWith(OwnershipKind::Guaranteed))
+      return {};
+    SILValue borrow = SILBuilderWithScope(insertPt).emitBeginBorrowOperation(
+        insertPt->getLoc(), value);
+    return pushEndBorrow(borrow);
+  }
+};
+
+} // namespace siloptimizer
+} // namespace swift
+
+#endif

lib/SILOptimizer/Transforms/SILMem2Reg.cpp

@@ -21,27 +21,31 @@
 #define DEBUG_TYPE "sil-mem2reg"
 
 #include "swift/AST/DiagnosticsSIL.h"
+#include "swift/SIL/BasicBlockDatastructures.h"
 #include "swift/SIL/Dominance.h"
 #include "swift/SIL/Projection.h"
 #include "swift/SIL/SILBuilder.h"
 #include "swift/SIL/SILFunction.h"
 #include "swift/SIL/SILInstruction.h"
 #include "swift/SIL/SILModule.h"
 #include "swift/SIL/TypeLowering.h"
-#include "swift/SIL/BasicBlockDatastructures.h"
 #include "swift/SILOptimizer/Analysis/DominanceAnalysis.h"
 #include "swift/SILOptimizer/PassManager/Passes.h"
 #include "swift/SILOptimizer/PassManager/Transforms.h"
 #include "swift/SILOptimizer/Utils/CFGOptUtils.h"
 #include "swift/SILOptimizer/Utils/InstOptUtils.h"
+#include "swift/SILOptimizer/Utils/OwnershipOptUtils.h"
+#include "swift/SILOptimizer/Utils/ScopeOptUtils.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/Statistic.h"
 #include "llvm/Support/Debug.h"
 #include <algorithm>
 #include <queue>
+
 using namespace swift;
+using namespace swift::siloptimizer;
 
 STATISTIC(NumAllocStackFound,    "Number of AllocStack found");
 STATISTIC(NumAllocStackCaptured, "Number of AllocStack captured");
@@ -141,6 +145,7 @@ static void replaceLoad(LoadInst *li, SILValue newValue, AllocStackInst *asi,
   ProjectionPath projections(newValue->getType());
   SILValue op = li->getOperand();
   SILBuilderWithScope builder(li, ctx);
+  SILOptScope scope;
 
   while (op != asi) {
     assert(isa<UncheckedAddrCastInst>(op) || isa<StructElementAddrInst>(op) ||
@@ -150,29 +155,28 @@ static void replaceLoad(LoadInst *li, SILValue newValue, AllocStackInst *asi,
     op = inst->getOperand(0);
   }
 
-  SmallVector<SILValue, 4> borrowedVals;
-  for (auto iter = projections.rbegin(); iter != projections.rend(); ++iter) {
-    const Projection &projection = *iter;
-    assert(projection.getKind() == ProjectionKind::BitwiseCast ||
-           projection.getKind() == ProjectionKind::Struct ||
-           projection.getKind() == ProjectionKind::Tuple);
+  for (const auto &proj : llvm::reverse(projections)) {
+    assert(proj.getKind() == ProjectionKind::BitwiseCast ||
+           proj.getKind() == ProjectionKind::Struct ||
+           proj.getKind() == ProjectionKind::Tuple);
 
     // struct_extract and tuple_extract expect guaranteed operand ownership
-    // non-trivial RunningVal is owned. Insert borrow operation to convert
-    if (projection.getKind() == ProjectionKind::Struct ||
-        projection.getKind() == ProjectionKind::Tuple) {
-      SILValue opVal = builder.emitBeginBorrowOperation(li->getLoc(), newValue);
-      if (opVal != newValue) {
-        borrowedVals.push_back(opVal);
-        newValue = opVal;
+    // non-trivial RunningVal is owned. Insert borrow operation to convert them
+    // to guaranteed!
+    if (proj.getKind() == ProjectionKind::Struct ||
+        proj.getKind() == ProjectionKind::Tuple) {
+      if (auto opVal = scope.borrowValue(li, newValue)) {
+        assert(*opVal != newValue &&
+               "Valid value should be different from input value");
+        newValue = *opVal;
       }
     }
     newValue =
-        projection.createObjectProjection(builder, li->getLoc(), newValue)
-            .get();
+        proj.createObjectProjection(builder, li->getLoc(), newValue).get();
   }
 
   op = li->getOperand();
+
   // Replace users of the loaded value with `val`
   // If we have a load [copy], replace the users with copy_value of `val`
   if (li->getOwnershipQualifier() == LoadOwnershipQualifier::Copy) {
@@ -183,9 +187,8 @@ static void replaceLoad(LoadInst *li, SILValue newValue, AllocStackInst *asi,
     li->replaceAllUsesWith(newValue);
   }
 
-  for (auto borrowedVal : borrowedVals) {
-    builder.emitEndBorrowOperation(li->getLoc(), borrowedVal);
-  }
+  // Pop the scope so that we emit cleanups.
+  std::move(scope).popAtEndOfScope(&*builder.getInsertionPoint());
 
   // Delete the load
   li->eraseFromParent();