Rewrite MutableSpan and OutputSpan subscript for BitwiseCopyable elements

Author: meg-gupta

stdlib/public/core/Span/MutableSpan.swift

@@ -170,6 +170,38 @@ extension MutableSpan where Element: BitwiseCopyable {
   }
 }
 
+@available(SwiftCompatibilitySpan 5.0, *)
+@_originallyDefinedIn(module: "Swift;CompatibilitySpan", SwiftCompatibilitySpan 6.2)
+extension MutableSpan where Element: BitwiseCopyable {
+  @_alwaysEmitIntoClient
+  public subscript(_ position: Index) -> Element {
+    get {
+      _checkIndex(position)
+      return unsafe self[unchecked: position]
+    }
+    @lifetime(self: copy self)
+    set {
+      _checkIndex(position)
+      unsafe self[unchecked: position] = newValue
+    }
+  }
+
+  @unsafe
+  @_alwaysEmitIntoClient
+  public subscript(unchecked position: Index) -> Element {
+    get {
+      let elementOffset = position &* MemoryLayout<Element>.stride
+      let address = unsafe _start().advanced(by: elementOffset)
+      return unsafe address.loadUnaligned(as: Element.self)
+    }
+    @lifetime(self: copy self)
+    set {
+      let elementOffset = position &* MemoryLayout<Element>.stride
+      unsafe _start().storeBytes(of: newValue, toByteOffset: elementOffset, as: Element.self)
+    }
+  }
+}
+
 @available(SwiftCompatibilitySpan 5.0, *)
 @_originallyDefinedIn(module: "Swift;CompatibilitySpan", SwiftCompatibilitySpan 6.2)
 extension Span where Element: ~Copyable {

stdlib/public/core/Span/OutputSpan.swift

@@ -182,6 +182,39 @@ extension OutputSpan {
   }
 }
 
+
+@available(SwiftCompatibilitySpan 5.0, *)
+@_originallyDefinedIn(module: "Swift;CompatibilitySpan", SwiftCompatibilitySpan 6.2)
+extension OutputSpan where Element: BitwiseCopyable {
+  @_alwaysEmitIntoClient
+  public subscript(_ position: Index) -> Element {
+    get {
+      _checkIndex(position)
+      return unsafe self[unchecked: position]
+    }
+    @lifetime(self: copy self)
+    set {
+      _checkIndex(position)
+      unsafe self[unchecked: position] = newValue
+    }
+  }
+
+  @unsafe
+  @_alwaysEmitIntoClient
+  public subscript(unchecked position: Index) -> Element {
+    get {
+      let elementOffset = position &* MemoryLayout<Element>.stride
+      let address = unsafe _start().advanced(by: elementOffset)
+      return unsafe address.loadUnaligned(as: Element.self)
+    }
+    @lifetime(self: copy self)
+    set {
+      let elementOffset = position &* MemoryLayout<Element>.stride
+      unsafe _start().storeBytes(of: newValue, toByteOffset: elementOffset, as: Element.self)
+    }
+  }
+}
+
 @available(SwiftCompatibilitySpan 5.0, *)
 @_originallyDefinedIn(module: "Swift;CompatibilitySpan", SwiftCompatibilitySpan 6.2)
 extension OutputSpan where Element: ~Copyable {

test/SILOptimizer/mutable_span_bounds_check_tests.swift

@@ -171,3 +171,69 @@ public func sorted(_ span: borrowing MutableSpan<Int>) -> Bool {
   return true
 }
 
+public protocol P {
+   mutating func mutate(_ other: Self)
+}
+
+// TODO: Ensure end_cow_mutation_addr is not inserted within the loop
+// This happens because MutableSpan's subscript is implemented with unsafe addressors which introduce mark_dependence [unresolved].
+// This triggers lifetime-dependence-scope-fixup which unnecessarily inserts end_cow_mutation_addr due to the conservative analysis of `mayHaveMutableSpan` on generic types.
+// TODO: Improve `mayHaveMutableSpan` to avoid this. 
+// CHECK-LABEL: sil @$s31mutable_span_bounds_check_tests0a1_B7_doubleyys11MutableSpanVyxGzAA1PRzlF :
+// CHECK-SIL: bb3({{.*}}):
+// TODO-CHECK-NOT: end_cow_mutation
+// CHECK-LABEL: } // end sil function '$s31mutable_span_bounds_check_tests0a1_B7_doubleyys11MutableSpanVyxGzAA1PRzlF'
+public func mutable_span_double<T: P>(_ ms: inout MutableSpan<T>) {
+  for i in ms.indices {
+    ms[i].mutate(ms[i])
+  }
+}
+
+// CHECK-LABEL: sil @$s31mutable_span_bounds_check_tests0a1_B23_double_bitwisecopyableyys11MutableSpanVyxGzs15BitwiseCopyableRzAA1PRzlF :
+// CHECK-SIL: bb3({{.*}}):
+// CHECK-NOT: end_cow_mutation
+// CHECK-LABEL: } // end sil function '$s31mutable_span_bounds_check_tests0a1_B23_double_bitwisecopyableyys11MutableSpanVyxGzs15BitwiseCopyableRzAA1PRzlF'
+public func mutable_span_double_bitwisecopyable<T: P & BitwiseCopyable>(_ ms: inout MutableSpan<T>) {
+  for i in ms.indices {
+    ms[i].mutate(ms[i])
+  }
+}
+
+// CHECK-LABEL: sil @$s31mutable_span_bounds_check_tests0a1_B22_equal_bitwisecopyableySbs11MutableSpanVyxGz_Sits15BitwiseCopyableRzSQRzlF :
+// CHECK-NOT: copy_addr
+// CHECK-LABEL: } // end sil function '$s31mutable_span_bounds_check_tests0a1_B22_equal_bitwisecopyableySbs11MutableSpanVyxGz_Sits15BitwiseCopyableRzSQRzlF'
+public func mutable_span_equal_bitwisecopyable<T : Equatable & BitwiseCopyable>(_ ms: inout MutableSpan<T>, _ index: Int) -> Bool {
+  return ms[index] == ms[index]
+}
+
+// CHECK-LABEL: sil @$s31mutable_span_bounds_check_tests0a1_B28_get_element_bitwisecopyableyxs11MutableSpanVyxGz_Sits15BitwiseCopyableRzlF :
+// CHECK-NOT: copy_addr
+// CHECK-LABEL: } // end sil function '$s31mutable_span_bounds_check_tests0a1_B28_get_element_bitwisecopyableyxs11MutableSpanVyxGz_Sits15BitwiseCopyableRzlF'
+public func mutable_span_get_element_bitwisecopyable<T : BitwiseCopyable>(_ ms: inout MutableSpan<T>, _ index: Int) -> T {
+  return ms[index]
+}
+
+extension Int : P {
+  public mutating func mutate(_ other: Int) {
+    self += other
+  }
+}
+
+// CHECK-LABEL: sil @$s31mutable_span_bounds_check_tests17specializedCalleryySaySiGzF : $@convention(thin) (@inout Array<Int>) -> () {
+// CHECK-NOT: cond_fail "index out of bounds"
+// CHECK-LABEL: } // end sil function '$s31mutable_span_bounds_check_tests17specializedCalleryySaySiGzF'
+public func specializedCaller(_ array: inout Array<Int>) {
+  var mut = array.mutableSpan
+  mutable_span_double_bitwisecopyable(&mut)
+}
+
+// CHECK-LABEL: sil @$s31mutable_span_bounds_check_tests07output_B25_init_all_bitwisecopyableyys10OutputSpanVyxGz_xts15BitwiseCopyableRzAA1PRzlF :
+// CHECK-SIL: bb3({{.*}}):
+// CHECK-NOT: end_cow_mutation
+// CHECK-LABEL: } // end sil function '$s31mutable_span_bounds_check_tests07output_B25_init_all_bitwisecopyableyys10OutputSpanVyxGz_xts15BitwiseCopyableRzAA1PRzlF'
+public func output_span_init_all_bitwisecopyable<T: P & BitwiseCopyable>(_ os: inout OutputSpan<T>, _ value: T) {
+  for i in os.indices {
+    os[i] = value
+  }
+}
+