Merge pull request swiftlang#81043 from meg-gupta/fixcow

Insert end_cow_mutation_addr for lifetime dependent values dependent on mutable addresses

Author: meg-gupta

SwiftCompilerSources/Sources/AST/Type.swift

@@ -182,6 +182,7 @@ extension TypeProperties {
   public var hasLocalArchetype: Bool { rawType.bridged.hasLocalArchetype() }
   public var isEscapable: Bool { rawType.bridged.isEscapable() }
   public var isNoEscape: Bool { rawType.bridged.isNoEscape() }
+  public var isBuiltinType: Bool { rawType.bridged.isBuiltinType() }
   public var archetypeRequiresClass: Bool { rawType.bridged.archetypeRequiresClass() }
 
   public var representationOfMetatype: AST.`Type`.MetatypeRepresentation {

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

SwiftCompilerSources/Sources/Optimizer/PassManager/Context.swift

@@ -326,6 +326,10 @@ struct FunctionPassContext : MutatingContext {
     _bridged.getSwiftArrayDecl().getAs(NominalTypeDecl.self)
   }
 
+  var swiftMutableSpan: NominalTypeDecl {
+    _bridged.getSwiftMutableSpanDecl().getAs(NominalTypeDecl.self)
+  }
+
   func loadFunction(name: StaticString, loadCalleesRecursively: Bool) -> Function? {
     return name.withUTF8Buffer { (nameBuffer: UnsafeBufferPointer<UInt8>) in
       let nameStr = BridgedStringRef(data: nameBuffer.baseAddress, count: nameBuffer.count)

SwiftCompilerSources/Sources/Optimizer/TestPasses/MemBehaviorDumper.swift

@@ -63,6 +63,7 @@ private extension Instruction {
          is BeginAccessInst,
          is EndAccessInst,
          is EndCOWMutationInst,
+         is EndCOWMutationAddrInst,
          is CopyValueInst,
          is DestroyValueInst,
          is StrongReleaseInst,

SwiftCompilerSources/Sources/Optimizer/Utilities/AddressUtils.swift

@@ -131,7 +131,7 @@ extension AddressUseVisitor {
          is DestroyAddrInst, is DeallocStackInst, 
          is DeinitExistentialAddrInst,
          is IsUniqueInst, is MarkFunctionEscapeInst,
-         is PackElementSetInst:
+         is PackElementSetInst, is EndCOWMutationAddrInst:
       return leafAddressUse(of: operand)
 
     case is LoadInst, is LoadUnownedInst,  is LoadWeakInst, is ValueMetatypeInst, is ExistentialMetatypeInst,

SwiftCompilerSources/Sources/SIL/Builder.swift

@@ -575,6 +575,11 @@ public struct Builder {
     return notifyNew(endMutation.getAs(EndCOWMutationInst.self))
   }
 
+  public func createEndCOWMutationAddr(address: Value) -> EndCOWMutationAddrInst {
+    let endMutation = bridged.createEndCOWMutationAddr(address.bridged)
+    return notifyNew(endMutation.getAs(EndCOWMutationAddrInst.self))
+  }
+
   public func createMarkDependence(value: Value, base: Value, kind: MarkDependenceKind) -> MarkDependenceInst {
     let markDependence = bridged.createMarkDependence(value.bridged, base.bridged,
                                                       BridgedInstruction.MarkDependenceKind(rawValue: kind.rawValue)!)

SwiftCompilerSources/Sources/SIL/Instruction.swift

@@ -1242,6 +1242,10 @@ final public class EndCOWMutationInst : SingleValueInstruction, UnaryInstruction
   public var doKeepUnique: Bool { bridged.EndCOWMutationInst_doKeepUnique() }
 }
 
+final public class EndCOWMutationAddrInst : Instruction, UnaryInstruction {
+  public var address: Value { operand.value }
+}
+
 final public
 class ClassifyBridgeObjectInst : SingleValueInstruction, UnaryInstruction {}
 

SwiftCompilerSources/Sources/SIL/Registration.swift

@@ -215,6 +215,7 @@ public func registerSILClasses() {
   register(MoveValueInst.self)
   register(DropDeinitInst.self)
   register(EndCOWMutationInst.self)
+  register(EndCOWMutationAddrInst.self)
   register(ClassifyBridgeObjectInst.self)
   register(PartialApplyInst.self)
   register(ApplyInst.self)

docs/SIL/Instructions.md

@@ -2146,6 +2146,22 @@ 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 an address. The address could be
+an opaque archetype, a struct, tuple or enum type and the end_cow_mutation_addr
+will apply to all members contained within it.
+It is currently only generated in cases where we maybe deriving a MutableSpan from
+`%0` since it is not possible to schedule an `end_cow_mutation` in the standard
+library automatically for Array.mutableSpan etc.
+
 ### destroy_not_escaped_closure
 
 ```

include/swift/AST/ASTBridging.h

@@ -3106,6 +3106,7 @@ struct BridgedASTType {
   SWIFT_IMPORT_UNSAFE BRIDGED_INLINE BridgedASTType getBuiltinVectorElementType() const;
   BRIDGED_INLINE bool isBuiltinFixedWidthInteger(SwiftInt width) const;
   BRIDGED_INLINE bool isOptional() const;
+  BRIDGED_INLINE bool isBuiltinType() const;
   SWIFT_IMPORT_UNSAFE BRIDGED_INLINE OptionalBridgedDeclObj getNominalOrBoundGenericNominal() const;
   BRIDGED_INLINE TraitResult canBeClass() const;
   SWIFT_IMPORT_UNSAFE BRIDGED_INLINE OptionalBridgedDeclObj getAnyNominal() const;