Insert end_cow_mutation_addr for lifetime dependent values dependent on mutable addresses

meg-gupta · meg-gupta · commit f40f5cbbe5e9 · 2025-04-30T13:22:19.000-07:00
Array/ArraySlice/ContiguousArray have support for mutableSpan property.
These types are "optimized COW" types, we use compiler builtins begin_cow_mutation/end_cow_mutation to
optimize uniqueness checks. Since mutableSpan is a property and not a coroutine there is
no way to schedule the end_cow_mutaton operation at the end of the access.
This can lead to miscompiles in rare cases where we can end up using a stale storage buffer after a cow.

This PR inserts end_cow_mutation_addr to avoid this issue.

Note: We can end up with unnecessary end_cow_mutation_addr. But it is just a barrier to prevent invalid optimizations
and has no impact.
diff --git a/SwiftCompilerSources/Sources/Optimizer/FunctionPasses/LifetimeDependenceScopeFixup.swift b/SwiftCompilerSources/Sources/Optimizer/FunctionPasses/LifetimeDependenceScopeFixup.swift
@@ -124,11 +124,102 @@ let lifetimeDependenceScopeFixupPass = FunctionPass(
     }
     let args = scopeExtension.findArgumentDependencies()
 
+    // If the scope cannot be extended to the caller, this must be the outermost dependency level.
+    // Insert end_cow_mutation_addr if needed.
+    if args.isEmpty {
+      createEndCOWMutationIfNeeded(lifetimeDep: newLifetimeDep, context)
+    }
+
     // Redirect the dependence base to the function arguments. This may create additional mark_dependence instructions.
     markDep.redirectFunctionReturn(to: args, context)
   }
 }
 
+private extension Type {
+  func mayHaveMutableSpan(in function: Function, _ context: FunctionPassContext) -> Bool {
+    if hasArchetype {
+      return true
+    }
+    if isBuiltinType {
+      return false
+    }
+    // Only result types that are nominal can have a MutableSpan derived from an inout array access.
+    if nominal == nil {
+      return false
+    }
+    if nominal == context.swiftMutableSpan {
+      return true
+    }
+    if isStruct {
+      guard let fields = getNominalFields(in: function) else {
+        return false
+      }
+      return fields.contains { $0.mayHaveMutableSpan(in: function, context) }
+    }
+    if isTuple {
+      return tupleElements.contains { $0.mayHaveMutableSpan(in: function, context) }
+    }
+    if isEnum {
+      guard let cases = getEnumCases(in: function) else {
+        return true
+      }
+      return cases.contains { $0.payload?.mayHaveMutableSpan(in: function, context) ?? false }
+    }
+    // Classes cannot be ~Escapable, therefore cannot hold a MutableSpan.
+    if isClass {
+      return false
+    }
+    return false
+  }
+}
+
+/// Insert end_cow_mutation_addr for lifetime dependent values that maybe of type MutableSpan and depend on a mutable address.
+private func createEndCOWMutationIfNeeded(lifetimeDep: LifetimeDependence, _ context: FunctionPassContext) {
+  var scoped : ScopedInstruction
+
+  // Handle cases which generate mutable addresses: begin_access [modify] and yield &
+  switch lifetimeDep.scope {
+    case let .access(beginAccess):
+      if beginAccess.accessKind != .modify {
+        return
+      }
+      scoped = beginAccess
+    case let .yield(value):
+      let beginApply = value.definingInstruction as! BeginApplyInst
+      if value == beginApply.token {
+        return
+      }
+      if beginApply.convention(of: value as! MultipleValueInstructionResult) != .indirectInout {
+        return
+      }
+      scoped = beginApply
+    // None of the below cases can generate a mutable address.
+    case let .owned:
+      fallthrough
+    case let .borrowed:
+      fallthrough
+    case let .local:
+      fallthrough
+    case let .initialized:
+      fallthrough
+    case let .caller:
+      fallthrough
+    case let .global:
+      fallthrough
+    case let .unknown:
+      return
+  }
+
+  guard lifetimeDep.dependentValue.type.mayHaveMutableSpan(in: lifetimeDep.dependentValue.parentFunction, context) else {
+    return
+  }
+
+  for endInstruction in scoped.endInstructions {
+    let builder = Builder(before: endInstruction, context)
+    builder.createEndCOWMutationAddr(address: lifetimeDep.parentValue)
+  }
+}
+
 private extension MarkDependenceInstruction {
   /// Rewrite the mark_dependence base operand to ignore inner borrow scopes (begin_borrow, load_borrow).
   ///
@@ -194,7 +285,7 @@ private extension MarkDependenceAddrInst {
   }
 }
 
-/// A scope extension is a set of nested scopes and their owners. The owner is a value that represents ownerhip of
+/// A scope extension is a set of nested scopes and their owners. The owner is a value that represents ownership of
 /// the outermost scopes, which cannot be extended; it limits how far the nested scopes can be extended.
 private struct ScopeExtension {
   let context: FunctionPassContext
diff --git a/docs/SIL/Instructions.md b/docs/SIL/Instructions.md
@@ -2146,6 +2146,19 @@ not replace this reference with a not uniquely reference object.
 
 For details see [Copy-on-Write Representation](SIL.md#Copy-on-Write-Representation).
 
+### end_cow_mutation_addr
+
+```
+sil-instruction ::= 'end_cow_mutation_addr' sil-operand
+
+end_cow_mutation_addr %0 : $*T
+// %0 must be of an address $*T type
+```
+
+This instruction marks the end of mutation of a reference counted object. It
+is currently only generated in cases where is it is not possible to schedule an
+`end_cow_mutation` in the standard library automatically. Ex: Array.mutableSpan
+
 ### destroy_not_escaped_closure
 
 ```
diff --git a/stdlib/public/core/Array.swift b/stdlib/public/core/Array.swift
@@ -1738,9 +1738,7 @@ extension Array {
     @_alwaysEmitIntoClient
     mutating get {
       _makeMutableAndUnique()
-      // NOTE: We don't have the ability to schedule a call to
-      //       ContiguousArrayBuffer.endCOWMutation().
-      //       rdar://146785284 (lifetime analysis for end of mutation)
+      // LifetimeDependence analysis inserts call to Builtin.endCOWMutation.
       let pointer = unsafe _buffer.firstElementAddress
       let count = _buffer.mutableCount
       let span = unsafe MutableSpan(_unsafeStart: pointer, count: count)
diff --git a/test/SILOptimizer/lifetime_dependence/lifetime_dependence_cow.sil b/test/SILOptimizer/lifetime_dependence/lifetime_dependence_cow.sil
@@ -0,0 +1,132 @@
+// RUN: %target-sil-opt %s \
+// RUN:   -lifetime-dependence-scope-fixup -sil-verify-all \
+// RUN:   -enable-experimental-feature LifetimeDependence | %FileCheck %s
+
+// REQUIRES: swift_in_compiler
+// REQUIRES: swift_feature_LifetimeDependence
+
+import Swift
+import SwiftShims
+import Builtin
+
+struct ArrayWrapper {
+  private var a: [Int]
+  var array: [Int]
+  init(a: [Int])
+}
+
+sil hidden [ossa] @$s23lifetime_dependence_cow12ArrayWrapperV5arraySaySiGvM : $@yield_once @convention(method) (@inout ArrayWrapper) -> @yields @inout Array<Int> {
+bb0(%0 : $*ArrayWrapper):
+  debug_value %0, var, name "self", argno 1, expr op_deref 
+  %2 = begin_access [modify] [static] %0          
+  %3 = struct_element_addr %2, #ArrayWrapper.a    
+  yield %3, resume bb1, unwind bb2                
+
+bb1:                                              
+  end_access %2                                   
+  %6 = tuple ()                                   
+  return %6                                       
+
+bb2:                                              
+  end_access %2                                   
+  unwind                                          
+} 
+
+sil hidden_external [serialized] [available 9999] @$sSa11mutableSpans07MutableB0VyxGvg : $@convention(method) <τ_0_0> (@inout Array<τ_0_0>) -> @lifetime(borrow 0) @owned MutableSpan<τ_0_0>
+
+sil hidden_external [serialized] [available 10.14.4] @$ss11MutableSpanVsRi_zrlE7indicesSnySiGvg : $@convention(method) <τ_0_0 where τ_0_0 : ~Copyable> (@guaranteed MutableSpan<τ_0_0>) -> Range<Int>
+
+sil [serialized] [always_inline] @$sSlss16IndexingIteratorVyxG0B0RtzrlE04makeB0ACyF : $@convention(method) <τ_0_0 where τ_0_0 : Collection, τ_0_0.Iterator == IndexingIterator<τ_0_0>> (@in τ_0_0) -> @out IndexingIterator<τ_0_0>
+
+sil [serialized] [always_inline] @$ss16IndexingIteratorV4next7ElementQzSgyF : $@convention(method) <τ_0_0 where τ_0_0 : Collection> (@inout IndexingIterator<τ_0_0>) -> @out Optional<τ_0_0.Element>
+
+sil hidden_external [serialized] [available 10.14.4] @$ss11MutableSpanVss15BitwiseCopyableRzlEyxSicis : $@convention(method) <τ_0_0 where τ_0_0 : BitwiseCopyable> (@in τ_0_0, Int, @lifetime(copy 2) @inout MutableSpan<τ_0_0>) -> ()
+
+sil [transparent] [serialized] @$sSi22_builtinIntegerLiteralSiBI_tcfC : $@convention(method) (Builtin.IntLiteral, @thin Int.Type) -> Int
+
+sil hidden [noinline] [ossa] @$s23lifetime_dependence_cow14getMutableSpanys0eF0VySiGSaySiGzF : $@convention(thin) (@inout Array<Int>) -> @lifetime(borrow 0) @owned MutableSpan<Int> {
+bb0(%0 : $*Array<Int>):
+  debug_value %0, var, name "array", argno 1, expr op_deref 
+  %2 = begin_access [modify] [static] %0          
+  %3 = function_ref @$sSa11mutableSpans07MutableB0VyxGvg : $@convention(method) <τ_0_0> (@inout Array<τ_0_0>) -> @lifetime(borrow 0) @owned MutableSpan<τ_0_0> 
+  %4 = apply %3<Int>(%2) : $@convention(method) <τ_0_0> (@inout Array<τ_0_0>) -> @lifetime(borrow 0) @owned MutableSpan<τ_0_0> 
+  %5 = mark_dependence [nonescaping] %4 on %0     
+  end_access %2                                   
+  return %5                                       
+} 
+
+// CHECK-LABEL: sil hidden [ossa] @$s23lifetime_dependence_cow9testWriteyyAA12ArrayWrapperVzF :
+// CHECK: [[FUN:%.*]] = function_ref @$s23lifetime_dependence_cow12ArrayWrapperV5arraySaySiGvM : $@yield_once @convention(method) (@inout ArrayWrapper) -> @yields @inout Array<Int>
+// CHECK: ([[Y:%.*]], [[T:%.*]]) = begin_apply %6(%5) : $@yield_once @convention(method) (@inout ArrayWrapper) -> @yields @inout Array<Int>
+// CHECK:  end_cow_mutation_addr [[Y]]
+// CHECK:  end_apply [[T]] as $()
+// CHECK-LABEL: } // end sil function '$s23lifetime_dependence_cow9testWriteyyAA12ArrayWrapperVzF'
+sil hidden [ossa] @$s23lifetime_dependence_cow9testWriteyyAA12ArrayWrapperVzF : $@convention(thin) (@inout ArrayWrapper) -> () {
+bb0(%0 : $*ArrayWrapper):
+  debug_value %0, var, name "w", argno 1, expr op_deref 
+  %2 = alloc_box ${ var MutableSpan<Int> }, var, name "span" 
+  %3 = begin_borrow [lexical] [var_decl] %2       
+  %4 = project_box %3, 0                          
+  %5 = begin_access [modify] [unknown] %0         
+  %6 = function_ref @$s23lifetime_dependence_cow12ArrayWrapperV5arraySaySiGvM : $@yield_once @convention(method) (@inout ArrayWrapper) -> @yields @inout Array<Int> 
+  (%7, %8) = begin_apply %6(%5) : $@yield_once @convention(method) (@inout ArrayWrapper) -> @yields @inout Array<Int> 
+  %9 = function_ref @$s23lifetime_dependence_cow14getMutableSpanys0eF0VySiGSaySiGzF : $@convention(thin) (@inout Array<Int>) -> @lifetime(borrow 0) @owned MutableSpan<Int> 
+  %10 = apply %9(%7) : $@convention(thin) (@inout Array<Int>) -> @lifetime(borrow 0) @owned MutableSpan<Int> 
+  %11 = mark_dependence [unresolved] %10 on %7    
+  %12 = end_apply %8 as $()
+  end_access %5                                   
+  store %11 to [init] %4                          
+  %15 = alloc_box ${ var IndexingIterator<Range<Int>> }, var, name "$i$generator" 
+  %16 = begin_borrow [var_decl] %15               
+  %17 = project_box %16, 0                        
+  %18 = begin_access [read] [unknown] %4          
+  %19 = mark_unresolved_non_copyable_value [no_consume_or_assign] %18 
+  %20 = load_borrow [unchecked] %19               
+  %21 = function_ref @$ss11MutableSpanVsRi_zrlE7indicesSnySiGvg : $@convention(method) <τ_0_0 where τ_0_0 : ~Copyable> (@guaranteed MutableSpan<τ_0_0>) -> Range<Int> 
+  %22 = apply %21<Int>(%20) : $@convention(method) <τ_0_0 where τ_0_0 : ~Copyable> (@guaranteed MutableSpan<τ_0_0>) -> Range<Int> 
+  end_borrow %20                                  
+  end_access %18                                  
+  %25 = alloc_stack $Range<Int>                   
+  store %22 to [trivial] %25                      
+  %27 = function_ref @$sSlss16IndexingIteratorVyxG0B0RtzrlE04makeB0ACyF : $@convention(method) <τ_0_0 where τ_0_0 : Collection, τ_0_0.Iterator == IndexingIterator<τ_0_0>> (@in τ_0_0) -> @out IndexingIterator<τ_0_0> 
+  %28 = apply %27<Range<Int>>(%17, %25) : $@convention(method) <τ_0_0 where τ_0_0 : Collection, τ_0_0.Iterator == IndexingIterator<τ_0_0>> (@in τ_0_0) -> @out IndexingIterator<τ_0_0>
+  dealloc_stack %25                               
+  br bb1                                          
+
+bb1:                                              
+  %31 = alloc_stack $Optional<Int>                
+  %32 = begin_access [modify] [unknown] %17       
+  %33 = function_ref @$ss16IndexingIteratorV4next7ElementQzSgyF : $@convention(method) <τ_0_0 where τ_0_0 : Collection> (@inout IndexingIterator<τ_0_0>) -> @out Optional<τ_0_0.Element> 
+  %34 = apply %33<Range<Int>>(%31, %32) : $@convention(method) <τ_0_0 where τ_0_0 : Collection> (@inout IndexingIterator<τ_0_0>) -> @out Optional<τ_0_0.Element>
+  end_access %32                                  
+  %36 = load [trivial] %31                        
+  dealloc_stack %31                               
+  switch_enum %36, case #Optional.some!enumelt: bb2, case #Optional.none!enumelt: bb3 
+
+bb2(%39 : $Int):                                  
+  %40 = move_value [var_decl] %39                 
+  debug_value %40, let, name "i"                  
+  %42 = integer_literal $Builtin.IntLiteral, 0    
+  %43 = metatype $@thin Int.Type                  
+  
+  %44 = function_ref @$sSi22_builtinIntegerLiteralSiBI_tcfC : $@convention(method) (Builtin.IntLiteral, @thin Int.Type) -> Int 
+  %45 = apply %44(%42, %43) : $@convention(method) (Builtin.IntLiteral, @thin Int.Type) -> Int 
+  %46 = alloc_stack $Int                          
+  store %45 to [trivial] %46                      
+  %48 = begin_access [modify] [unknown] %4        
+  %49 = mark_unresolved_non_copyable_value [assignable_but_not_consumable] %48 
+  %50 = function_ref @$ss11MutableSpanVss15BitwiseCopyableRzlEyxSicis : $@convention(method) <τ_0_0 where τ_0_0 : BitwiseCopyable> (@in τ_0_0, Int, @lifetime(copy 2) @inout MutableSpan<τ_0_0>) -> () 
+  %51 = apply %50<Int>(%46, %40, %49) : $@convention(method) <τ_0_0 where τ_0_0 : BitwiseCopyable> (@in τ_0_0, Int, @lifetime(copy 2) @inout MutableSpan<τ_0_0>) -> ()
+  end_access %48                                  
+  dealloc_stack %46                               
+  extend_lifetime %40                             
+  br bb1                                          
+
+bb3:                                              
+  end_borrow %16                                  
+  destroy_value %15                               
+  end_borrow %3                                   
+  destroy_value %2                                
+  %60 = tuple ()                                  
+  return %60                                      
+} 
diff --git a/test/SILOptimizer/lifetime_dependence/lifetime_dependence_cow.swift b/test/SILOptimizer/lifetime_dependence/lifetime_dependence_cow.swift
