BoundsCheckOpts: merge OutputSpan.append capacity checks

include/swift/AST/SemanticAttrs.def

@@ -165,6 +165,11 @@ SEMANTICS_ATTR(DELETE_IF_UNUSED, "sil.optimizer.delete_if_unused")
 SEMANTICS_ATTR(USE_FRAME_POINTER, "use_frame_pointer")
 
 SEMANTICS_ATTR(FIXED_STORAGE_CHECK_INDEX, "fixed_storage.check_index")
+// Like `fixed_storage.check_index` but validates against `capacity`, which
+// must be an invariant property of the storage. The callee must depend only
+// on `let` fields of self, so the optimizer may merge successive checks
+// across mutations of `self`.
+SEMANTICS_ATTR(FIXED_STORAGE_CHECK_CAPACITY, "fixed_storage.check_capacity")
 SEMANTICS_ATTR(FIXED_STORAGE_GET_COUNT, "fixed_storage.get_count")
 
 SEMANTICS_ATTR(NO_SIL_VERIFICATION, "sil.verify_none")

include/swift/SIL/InstWrappers.h

@@ -387,7 +387,17 @@ class ForwardingOperation {
   bool visitForwardedValues(function_ref<bool(SILValue)> visitor);
 };
 
-enum class FixedStorageSemanticsCallKind { None, CheckIndex, GetCount };
+enum class FixedStorageSemanticsCallKind {
+  None,
+  CheckIndex,
+  /// Check that an index is in `0..<capacity`. Distinct from `CheckIndex`
+  /// because `capacity` is required to be invariant for the lifetime of the
+  /// storage. The callee must depend only on `let` fields of self, so
+  /// successive checks across mutations of `self` may be merged and the self
+  /// argument can be substituted across them.
+  CheckCapacity,
+  GetCount
+};
 
 struct FixedStorageSemanticsCall {
   ApplyInst *apply = nullptr;
@@ -407,6 +417,10 @@ struct FixedStorageSemanticsCall {
         apply = applyInst;
         kind = FixedStorageSemanticsCallKind::CheckIndex;
         break;
+      } else if (attr == "fixed_storage.check_capacity") {
+        apply = applyInst;
+        kind = FixedStorageSemanticsCallKind::CheckCapacity;
+        break;
       } else if (attr == "fixed_storage.get_count") {
         apply = applyInst;
         kind = FixedStorageSemanticsCallKind::GetCount;

lib/SIL/Utils/InstWrappers.cpp

@@ -104,6 +104,7 @@ bool ForwardingOperation::visitForwardedValues(
 bool swift::isFixedStorageSemanticsCallKind(SILFunction *function) {
   for (auto &attr : function->getSemanticsAttrs()) {
     if (attr == "fixed_storage.check_index" ||
+        attr == "fixed_storage.check_capacity" ||
         attr == "fixed_storage.get_count") {
       return true;
     }

lib/SILOptimizer/LoopTransforms/BoundsCheckOpts.cpp

@@ -1231,10 +1231,47 @@ static bool canOptimize(FixedStorageSemanticsCall call) {
   return true;
 }
 
+/// Whether a fixed-storage semantic call kind participates in
+/// `hoistFixedStorageBoundsChecksInBlock`'s clustering.
+static bool isMergeableCheckKind(FixedStorageSemanticsCallKind kind) {
+  switch (kind) {
+  case FixedStorageSemanticsCallKind::CheckIndex:
+  case FixedStorageSemanticsCallKind::CheckCapacity:
+    return true;
+  case FixedStorageSemanticsCallKind::None:
+  case FixedStorageSemanticsCallKind::GetCount:
+    return false;
+  }
+}
+
+/// Compute a merge-group key for a fixed-storage bounds check.
+///
+/// `CheckIndex` calls validate the index against `count`, which may be mutated
+/// in place (e.g. `OutputSpan`). For those, the key must be the literal SSA
+/// self value, so checks separated by mutations stay in distinct groups.
+///
+/// `CheckCapacity` calls validate against `capacity`, which is invariant. When
+/// self is loaded from an addressable location (an inout, a stack slot), reach
+/// through the load to the underlying address so successive loads of the same
+/// storage merge.
+static SILValue
+getFixedStorageMergeKey(const FixedStorageSemanticsCall &call) {
+  auto self = call.getSelf();
+  if (call.getKind() == FixedStorageSemanticsCallKind::CheckCapacity) {
+    if (auto *inst = self->getDefiningInstruction()) {
+      LoadOperation load{inst};
+      if (load)
+        return load.getOperand();
+    }
+  }
+  return self;
+}
+
 bool BoundsCheckOpts::hoistFixedStorageBoundsChecksInBlock(SILBasicBlock &block) {
   bool changed = false;
 
-  // Map to track fixed storage checks grouped by self value.
+  // Map to track fixed storage checks grouped by self value (or, for
+  // CheckCapacity, the underlying address self is loaded from).
   llvm::MapVector<SILValue, SmallVector<FixedStorageSemanticsCall, 4>>
       checksToMerge;
 
@@ -1245,8 +1282,7 @@ bool BoundsCheckOpts::hoistFixedStorageBoundsChecksInBlock(SILBasicBlock &block)
   for (auto &inst : block) {
     FixedStorageSemanticsCall fixedStorageCall(&inst);
     if (!fixedStorageCall ||
-        fixedStorageCall.getKind() !=
-            FixedStorageSemanticsCallKind::CheckIndex ||
+        !isMergeableCheckKind(fixedStorageCall.getKind()) ||
         fixedStorageCall->getNumArguments() < 2) {
       continue;
     }
@@ -1255,14 +1291,14 @@ bool BoundsCheckOpts::hoistFixedStorageBoundsChecksInBlock(SILBasicBlock &block)
       continue;
     }
 
-    auto selfValue = fixedStorageCall.getSelf();
+    auto mergeKey = getFixedStorageMergeKey(fixedStorageCall);
     auto indexValue = fixedStorageCall.getIndex();
 
-    // Add this check to the list for the self value
-    checksToMerge[selfValue].push_back(fixedStorageCall);
+    // Add this check to the list for the merge key
+    checksToMerge[mergeKey].push_back(fixedStorageCall);
 
     LLVM_DEBUG(llvm::dbgs() << "  Found fixed storage check: " << inst
-                            << " with self: " << *selfValue
+                            << " with merge key: " << *mergeKey
                             << " and index: " << *indexValue);
   }
 
@@ -1278,16 +1314,57 @@ bool BoundsCheckOpts::hoistFixedStorageBoundsChecksInBlock(SILBasicBlock &block)
   // Second pass: process each self group to place checks with same self next to
   // each other, preserving their original relative order.
   for (auto &entry : checksToMerge) {
+    auto mergeKey = entry.first;
     auto &checks = entry.second;
     if (checks.size() <= 1) {
       continue;
     }
 
     SILInstruction *insertAfter = *checks[0];
+    SILValue firstSelf = checks[0].getSelf();
+
+    // For CheckCapacity, two checks with different SSA selves were grouped
+    // because they share an underlying address (mergeKey). Substituting
+    // `firstSelf` for the moved check is sound only if nothing between
+    // the previous clustered check and `check[i]` could have replaced the
+    // storage at that address — a whole-storage write changes invariant
+    // fields the substitution depends on (`capacity`).
+    auto canReplaceStorage = [&](SILInstruction *inst) {
+      // Any opaque call may replace storage reachable through its arguments.
+      // Recognized fixed-storage semantic calls are read-only by contract.
+      if (FullApplySite::isa(inst)) {
+        return !static_cast<bool>(FixedStorageSemanticsCall(inst));
+      }
+      // Whole-storage stores; field stores go through `struct_element_addr`
+      // and have a different destination SSA value.
+      if (auto *store = dyn_cast<StoreInst>(inst))
+        return store->getDest() == mergeKey;
+      if (auto *copy = dyn_cast<CopyAddrInst>(inst))
+        return copy->getDest() == mergeKey;
+      return false;
+    };
 
     for (unsigned i = 1, e = checks.size(); i < e; ++i) {
       auto check = checks[i];
 
+      if (check.getKind() == FixedStorageSemanticsCallKind::CheckCapacity) {
+        bool storageReplaced = false;
+        for (auto it = std::next(insertAfter->getIterator()),
+                  end = check->getIterator();
+             it != end; ++it) {
+          if (canReplaceStorage(&*it)) {
+            storageReplaced = true;
+            break;
+          }
+        }
+        if (storageReplaced) {
+          LLVM_DEBUG(llvm::dbgs()
+                     << "      Skipping merge across storage replacement: "
+                     << *check.apply);
+          continue;
+        }
+      }
+
       auto indexValue = check.getIndex();
       auto clonedValue =
           cloneFixedStorageIndex(indexValue, *checks[0], instIndices);
@@ -1296,6 +1373,14 @@ bool BoundsCheckOpts::hoistFixedStorageBoundsChecksInBlock(SILBasicBlock &block)
       }
       check.getIndexOperand().set(*clonedValue);
       check->getCalleeOperand()->set(checks[0]->getCallee());
+      // For CheckCapacity, self may be a fresh load following intervening
+      // stores; replace it with the first check's self so the moved apply
+      // does not reference a value defined later in the block. This is sound
+      // because capacity is invariant.
+      if (check.getKind() == FixedStorageSemanticsCallKind::CheckCapacity &&
+          check.getSelf() != firstSelf) {
+        check->getSelfArgumentOperand().set(firstSelf);
+      }
       check->moveAfter(insertAfter);
       insertAfter = *check;
 

stdlib/public/core/Span/OutputSpan.swift

@@ -292,13 +292,25 @@ extension OutputSpan where Element: ~Copyable {
 @available(SwiftCompatibilitySpan 5.0, *)
 @_originallyDefinedIn(module: "Swift;CompatibilitySpan", SwiftCompatibilitySpan 6.2)
 extension OutputSpan where Element: ~Copyable {
+  // SILOptimizer looks for fixed_storage.check_capacity semantics for bounds
+  // check optimizations. Distinct from `_checkIndex` because `capacity` is
+  // invariant for the lifetime of the storage, so successive checks across
+  // mutations of `self` may be merged. Only called with `_count`, which is
+  // non-negative by invariant, so no lower-bound check is needed.
+  @_semantics("fixed_storage.check_capacity")
+  @inline(__always)
+  @_alwaysEmitIntoClient
+  internal func _checkCanAppend(_ index: Int) {
+    _precondition(index < capacity, "OutputSpan capacity overflow")
+  }
+
   /// Append a single element to this span.
   ///
   /// - Parameter value: The element to append.
   @_alwaysEmitIntoClient
   @_lifetime(self: copy self)
   public mutating func append(_ value: consuming Element) {
-    _precondition(_count < capacity, "OutputSpan capacity overflow")
+    _checkCanAppend(_count)
     unsafe _tail().initializeMemory(as: Element.self, to: value)
     _count &+= 1
   }

test/SILOptimizer/mutable_span_bounds_check_tests.swift

@@ -180,3 +180,43 @@ public func outputspan_get_element(_ v: borrowing OutputSpan<Int>, _ i: Int) ->
   return v[i]
 }
 
+// Sequential `append`s on the same `inout OutputSpan` should fold into a
+// single merged capacity check covering all four indices, instead of one
+// check per `append`. Mirrors the `span_4_sum` test for `Span` subscripts.
+//
+// CHECK-SIL-LABEL: sil @$s31mutable_span_bounds_check_tests19outputspan_4_appendyys10OutputSpanVys5UInt8VGz_A4FtF :
+// CHECK-SIL: cond_fail {{.*}}, "OutputSpan capacity overflow"
+// CHECK-SIL-NOT: cond_fail {{.*}}, "OutputSpan capacity overflow"
+// CHECK-SIL-LABEL: } // end sil function '$s31mutable_span_bounds_check_tests19outputspan_4_appendyys10OutputSpanVys5UInt8VGz_A4FtF'
+@_lifetime(output: copy output)
+public func outputspan_4_append(
+  _ output: inout OutputSpan<UInt8>,
+  _ a: UInt8, _ b: UInt8, _ c: UInt8, _ d: UInt8
+) {
+  output.append(a); output.append(b); output.append(c); output.append(d)
+}
+
+@inline(never)
+@_lifetime(output: copy output)
+internal func _mutateOutputSpan(_ output: inout OutputSpan<UInt8>) {}
+
+// An opaque `@inout` callee between two `append`s could replace the storage,
+// so the second capacity check must NOT be merged with the first. Both
+// checks must remain.
+//
+// CHECK-SIL-LABEL: sil @$s31mutable_span_bounds_check_tests31outputspan_append_across_mutateyys10OutputSpanVys5UInt8VGz_A3FtF :
+// CHECK-SIL: cond_fail {{.*}}, "OutputSpan capacity overflow"
+// CHECK-SIL: function_ref{{.*}}_mutateOutputSpan
+// CHECK-SIL: cond_fail {{.*}}, "OutputSpan capacity overflow"
+// CHECK-SIL-LABEL: } // end sil function '$s31mutable_span_bounds_check_tests31outputspan_append_across_mutateyys10OutputSpanVys5UInt8VGz_A3FtF'
+@_lifetime(output: copy output)
+public func outputspan_append_across_mutate(
+  _ output: inout OutputSpan<UInt8>,
+  _ a: UInt8, _ b: UInt8, _ c: UInt8
+) {
+  output.append(a)
+  _mutateOutputSpan(&output)
+  output.append(b)
+  output.append(c)
+}
+

test/SILOptimizer/outputspan_append_capacity_check_traps.swift

@@ -0,0 +1,46 @@
+// RUN: %target-run-simple-swift(-O -enable-experimental-feature Lifetimes)
+// RUN: %target-run-simple-swift(-Onone -enable-experimental-feature Lifetimes)
+
+// REQUIRES: executable_test
+// REQUIRES: swift_feature_Lifetimes
+
+import StdlibUnittest
+
+// Regression test for BoundsCheckOpts CheckCapacity merging. Two
+// `OutputSpan.append` capacity checks separated by an `@inout` callee that
+// fills the span must not be merged using the first call's loaded count.
+//
+// The bug: `getFixedStorageMergeKey` looks through `load_borrow` for
+// CheckCapacity, so two checks against the same `inout` span land in one
+// merge group. If `cloneFixedStorageIndex` clones the second check's index
+// (a fresh `load` of `_count`) back past the intervening callee, the cloned
+// load reads stale memory and the check passes when it should trap.
+
+let suite = TestSuite("OutputSpan bounds check optimization")
+
+@inline(never)
+func fillRemaining(_ span: inout OutputSpan<Int>) {
+  while !span.isFull {
+    span.append(0)
+  }
+}
+
+@inline(never)
+func appendAfterFull(_ span: inout OutputSpan<Int>) {
+  span.append(1)
+  fillRemaining(&span)
+  // After `fillRemaining`, count == capacity. This append must trap.
+  span.append(99)
+}
+
+suite.test("append after inout callee fills span") {
+  expectCrashLater()
+  let buffer = UnsafeMutableBufferPointer<Int>.allocate(capacity: 2)
+  defer { buffer.deallocate() }
+  var span = unsafe OutputSpan(buffer: buffer, initializedCount: 0)
+  appendAfterFull(&span)
+  expectUnreachable("third append should have trapped on capacity overflow")
+  _ = unsafe span.finalize(for: buffer)
+}
+
+runAllTests()