Merge pull request #71597 from atrick/lifetime-nested

LifetimeDependenceScopeFixup: handle nested access

Author: atrick

SwiftCompilerSources/Sources/Optimizer/FunctionPasses/LifetimeDependenceInsertion.swift

@@ -181,63 +181,3 @@ private func insertMarkDependencies(value: Value, initializer: Instruction?,
     currentValue = markDep
   }
 }
-
-/*
-/// Return base values that this return value depends on.
-///
-/// For lifetime copies, walk up the dependence chain to find the
-/// dependence roots, inserting dependencies for any
-/// LifetimeDependentApply.
-private func recursivelyFindDependenceBases(of apply: LifetimeDependentApply,
-                                            _ context: FunctionPassContext)
-  -> [Value] {
-  log("Creating dependencies for \(apply.applySite)")
-  var bases: [Value] = []
-  for lifetimeArg in apply.getLifetimeArguments() {
-    switch lifetimeArg.convention {
-    case .inherit:
-      // Inherit the argument's lifetime dependence by finding the
-      // roots.  This requires that a mark_dependence [nonescaping]
-      // already be created for any earlier LifetimeDependentApply.
-      _ = LifetimeDependence.visitDependenceRoots(enclosing: lifetimeArg.value,
-                                                  context)
-      { (scope: LifetimeDependence.Scope) in
-        if let updatedScope = recursivelyUpdate(scope: scope, context) {
-          log("Inherited lifetime from \(lifetimeArg.value)")
-          log("  depends on: \(updatedScope)")
-          bases.append(updatedScope.parentValue)
-        }
-        return .continueWalk
-      }
-    case .scope:
-      // Create a new dependence on the apply's access to the argument.
-      if let scope =
-           LifetimeDependence.Scope(base: lifetimeArg.value, context) {
-        log("Scoped lifetime from \(lifetimeArg.value)")
-        log("  scope: \(scope)")
-        bases.append(scope.parentValue)
-      }
-    }
-  }
-  return bases
-}
-
-// Recursively insert dependencies, assuming no cycle of dependent applies.
-//
-// TODO: needs unit test.
-private func recursivelyUpdate(scope: LifetimeDependence.Scope,
-  _ context: FunctionPassContext) -> LifetimeDependence.Scope? {
-  if let dependentApply =
-    LifetimeDependentApply(withResult: scope.parentValue) {
-    insertDependencies(for: dependentApply, context)
-    // If a mark_dependence [nonescaping] was created for this apply,
-    // then return it as the updated dependence. Otherwise, return the
-    // original dependence.
-    if let markDep = scope.parentValue.uses.singleUse?.instruction
-      as? MarkDependenceInst {
-      return LifetimeDependence(markDep, context)?.scope
-    }
-  }
-  return scope
-}
-*/

SwiftCompilerSources/Sources/Optimizer/FunctionPasses/LifetimeDependenceScopeFixup.swift

@@ -37,18 +37,26 @@ let lifetimeDependenceScopeFixupPass = FunctionPass(
     guard let markDep = instruction as? MarkDependenceInst else {
       continue
     }
-    if let lifetimeDep = LifetimeDependence(markDep, context) {
-      fixup(dependence: lifetimeDep, context)
+    guard let lifetimeDep = LifetimeDependence(markDep, context) else {
+      continue
+    }
+    guard let beginAccess = extendAccessScopes(dependence: lifetimeDep,
+                                               context) else {
+      continue
     }
+    extendDependenceBase(dependenceInstruction: markDep,
+                         beginAccess: beginAccess, context)
   }
 }
 
-private func fixup(dependence: LifetimeDependence,
-  _ context: FunctionPassContext) {
+// Extend all access scopes that enclose `dependence` and return the
+// outermost access.
+private func extendAccessScopes(dependence: LifetimeDependence,
+  _ context: FunctionPassContext) -> BeginAccessInst? {
   log("Scope fixup for lifetime dependent instructions: \(dependence)")
 
   guard case .access(let bai) = dependence.scope else {
-    return
+    return nil
   }
   var range = InstructionRange(begin: bai, context)
   var walker = LifetimeDependenceScopeFixupWalker(bai.parentFunction, context) {
@@ -89,6 +97,29 @@ private func fixup(dependence: LifetimeDependence,
     }
     beginAccess = enclosingBeginAccess
   }
+  return beginAccess
+}
+
+/// Rewrite the mark_dependence to depend on the outermost access
+/// scope now that the nested scopes have all been extended.
+private func extendDependenceBase(dependenceInstruction: MarkDependenceInst,
+                                  beginAccess: BeginAccessInst,
+                                  _ context: FunctionPassContext) {
+  guard case let .base(accessBase) = beginAccess.address.enclosingAccessScope
+  else {
+    fatalError("this must be the outer-most access scope")
+  }
+  // If the outermost access is in the caller, then depende on the
+  // address argument.
+  let baseAddress: Value
+  switch accessBase {
+  case let .argument(arg):
+    assert(arg.type.isAddress)
+    baseAddress = arg
+  default:
+    baseAddress = beginAccess
+  }
+  dependenceInstruction.baseOperand.set(to: baseAddress, context)
 }
 
 private struct LifetimeDependenceScopeFixupWalker : LifetimeDependenceDefUseWalker {

SwiftCompilerSources/Sources/Optimizer/Utilities/LifetimeDependenceUtils.swift

@@ -233,6 +233,9 @@ extension LifetimeDependence {
 
   /// Construct LifetimeDependence from mark_dependence [unresolved]
   ///
+  /// For any LifetimeDependence constructed from a mark_dependence,
+  /// its `dependentValue` will be the result of the mark_dependence.
+  ///
   /// TODO: Add SIL verification that all mark_depedence [unresolved]
   /// have a valid LifetimeDependence.
   init?(_ markDep: MarkDependenceInst, _ context: some Context) {
@@ -485,8 +488,7 @@ extension LifetimeDependence {
   -> WalkResult {
     var useDefVisitor = UseDefVisitor(context, visitor)
     defer { useDefVisitor.deinitialize() }
-    let owner = value.ownership == .owned ? value : nil
-    return useDefVisitor.walkUp(value: value, owner)
+    return useDefVisitor.walkUp(valueOrAddress: value)
   }
 
   private struct UseDefVisitor : LifetimeDependenceUseDefWalker {
@@ -534,6 +536,18 @@ extension LifetimeDependence {
   }
 }
 
+/// Walk up the lifetime dependence
+///
+/// This uses LifetimeDependenceUseDefWalker to find the introducers
+/// of a dependence chain, which represent the value's "inherited"
+/// dependencies. This stops at an address, unless the address refers
+/// to a singly-initialized temprorary, in which case it continues to
+/// walk up the stored value.
+///
+/// This overrides LifetimeDependenceUseDefWalker to stop at a value
+/// that introduces an immutable variable: move_value [var_decl] or
+/// begin_borrow [var_decl], and to stop at an access of a mutable
+/// variable: begin_access.
 struct VariableIntroducerUseDefWalker : LifetimeDependenceUseDefWalker {
   let context: Context
   // This visited set is only really needed for instructions with
@@ -576,24 +590,36 @@ struct VariableIntroducerUseDefWalker : LifetimeDependenceUseDefWalker {
     }
     return walkUpDefault(dependent: value, owner: owner)
   }
+
+  mutating func walkUp(address: Value) -> WalkResult {
+    if let beginAccess = address.definingInstruction as? BeginAccessInst {
+      return introducer(beginAccess, nil)
+    }
+    return walkUpDefault(address: address)
+  }
 }
 
-/// Walk up dominating dependent values.
+/// Walk up the lifetime dependence chain.
 ///
-/// Find the roots of a dependence chain stopping at phis. These root
-/// LifeDependence instances are the value's "inherited"
-/// dependencies. In this example, the dependence root is copied,
-/// borrowed, and forwarded before being used as the base operand of
-/// `mark_dependence`. The dependence "root" is the parent of the
-/// outer-most dependence scope.
+/// This finds the introducers of a dependence chain. which represent
+/// the value's "inherited" dependencies. This stops at phis; all
+/// introducers dominate. This stops at addresses in general, but if
+/// the value is loaded from a singly-initialized location, then it
+/// continues walking up the value stored by the initializer. This
+/// bypasses the copies to temporary memory locations emitted by SILGen.
+///
+/// In this example, the dependence root is
+/// copied, borrowed, and forwarded before being used as the base
+/// operand of `mark_dependence`. The dependence "root" is the parent
+/// of the outer-most dependence scope.
 ///
 ///   %root = apply                  // lifetime dependence root
 ///   %copy = copy_value %root
 ///   %parent = begin_borrow %copy   // lifetime dependence parent value
 ///   %base = struct_extract %parent // lifetime dependence base value
 ///   %dependent = mark_dependence [nonescaping] %value on %base
 ///
-/// This extends the ForwardingDefUseWalker, which finds the
+/// This extends the ForwardingUseDefWalker, which finds the
 /// forward-extended lifetime introducers. Certain forward-extended
 /// lifetime introducers can inherit a lifetime dependency from their
 /// operand: namely copies, moves, and borrows. These introducers are
@@ -631,7 +657,7 @@ struct VariableIntroducerUseDefWalker : LifetimeDependenceUseDefWalker {
 ///   needWalk(for value: Value) -> Bool
 ///
 /// Start walking:
-///   walkUp(value: Value) -> WalkResult
+///   walkUp(valueOrAddress: Value) -> WalkResult
 protocol LifetimeDependenceUseDefWalker : ForwardingUseDefWalker where PathContext == Value? {
   var context: Context { get }
 
@@ -642,14 +668,28 @@ protocol LifetimeDependenceUseDefWalker : ForwardingUseDefWalker where PathConte
   mutating func needWalk(for value: Value, _ owner: Value?) -> Bool
 
   mutating func walkUp(value: Value, _ owner: Value?) -> WalkResult
+
+  mutating func walkUp(address: Value) -> WalkResult
 }
 
 // Implement ForwardingUseDefWalker
 extension LifetimeDependenceUseDefWalker {
+  mutating func walkUp(valueOrAddress: Value) -> WalkResult {
+    if valueOrAddress.type.isAddress {
+      return walkUp(address: valueOrAddress)
+    }
+    let owner = valueOrAddress.ownership == .owned ? valueOrAddress : nil
+    return walkUp(value: valueOrAddress, owner)
+  }
+
   mutating func walkUp(value: Value, _ owner: Value?) -> WalkResult {
     walkUpDefault(dependent: value, owner: owner)
   }
 
+  mutating func walkUp(address: Value) -> WalkResult {
+    walkUpDefault(address: address)
+  }
+
   // Extend ForwardingUseDefWalker to handle copies, moves, and
   // borrows. Also transitively walk up other lifetime dependencies to
   // find the roots.
@@ -691,17 +731,12 @@ extension LifetimeDependenceUseDefWalker {
     return walkUpDefault(forwarded: value, owner)
   }
 
-  private mutating func walkUp(newLifetime: Value) -> WalkResult {
-    let newOwner = newLifetime.ownership == .owned ? newLifetime : nil
-    return walkUp(value: newLifetime, newOwner)
-  }
-
   // Walk up from a load of a singly-initialized address to find the
   // dependence root of the stored value. This ignores mutable
   // variables, which require an access scope. This ignores applies
   // because an lifetime dependence will already be expressed as a
   // mark_dependence.
-  private mutating func walkUp(address: Value) -> WalkResult {
+  mutating func walkUpDefault(address: Value) -> WalkResult {
     if let (_, initializingStore) =
          address.accessBase.findSingleInitializer(context) {
       switch initializingStore {
@@ -715,6 +750,11 @@ extension LifetimeDependenceUseDefWalker {
     }
     return introducer(address, nil)
   }
+
+  private mutating func walkUp(newLifetime: Value) -> WalkResult {
+    let newOwner = newLifetime.ownership == .owned ? newLifetime : nil
+    return walkUp(value: newLifetime, newOwner)
+  }
 }
 
 /// Walk down dependent values.

test/SILOptimizer/lifetime_dependence_borrow.swift

@@ -42,3 +42,8 @@ func bv_get_borrow(container: borrowing NC) -> _borrow(container) BV {
 func bv_get_copy(container: borrowing NC) -> _copy(container) BV {
   return container.getBV()
 }
+
+// Recognize nested accesses as part of the same dependence scope.
+func bv_get_mutate(container: inout NC) -> _mutate(container) BV {
+  container.getBV()
+}

test/SILOptimizer/lifetime_dependence_inherit.swift

@@ -39,8 +39,6 @@ struct NE {
 }
 
 // Test lifetime inheritance through chained consumes.
-//
-// This requires an inherit_lifetime marker on the argument.
 func bv_derive(bv: consuming BV) -> _consume(bv) BV {
   bv.derive()
 }

test/SILOptimizer/lifetime_dependence_scope.swift

@@ -0,0 +1,48 @@
+// RUN: %target-swift-frontend %s -emit-sil \
+// RUN:   -sil-verify-all \
+// RUN:   -module-name test \
+// RUN:   -disable-experimental-parser-round-trip \
+// RUN:   -enable-experimental-feature NonescapableTypes \
+// RUN:   -Xllvm -enable-lifetime-dependence-diagnostics \
+// RUN:   2>&1 | %FileCheck %s
+
+// REQUIRES: swift_in_compiler
+
+// Test LifetimeDependenceScopeFixup.
+
+@_nonescapable
+struct BV {
+  let p: UnsafeRawPointer
+  let c: Int
+
+  public var isEmpty: Bool { c == 0 }
+
+  @_unsafeNonescapableResult
+  init(_ p: UnsafeRawPointer, _ c: Int) {
+    self.p = p
+    self.c = c
+  }
+}
+
+struct NC : ~Copyable {
+  let p: UnsafeRawPointer
+  let c: Int
+
+  // Requires a borrow.
+  borrowing func getBV() -> _borrow(self) BV {
+    BV(p, c)
+  }
+}
+
+// Rewrite the mark_dependence to depende on the incoming argument rather than the nested access.
+//
+// CHECK-LABEL: sil hidden @$s4test13bv_get_mutate9containerAA2BVVAA2NCVz_tF : $@convention(thin) (@inout NC) -> _scope(1) @owned BV {
+// CHECK: bb0(%0 : $*NC):
+// CHECK:   [[A:%.*]] = begin_access [read] [static] %0 : $*NC
+// CHECK:   [[L:%.*]] = load [[A]] : $*NC
+// CHECK:   [[R:%.*]] = apply %{{.*}}([[L]]) : $@convention(method) (@guaranteed NC) -> _scope(0) @owned BV
+// CHECK:   [[M:%.*]] = mark_dependence [nonescaping] [[R]] : $BV on %0 : $*NC
+// CHECK-LABEL: } // end sil function '$s4test13bv_get_mutate9containerAA2BVVAA2NCVz_tF'
+func bv_get_mutate(container: inout NC) -> _mutate(container) BV {
+  container.getBV()
+}

test/SILOptimizer/lifetime_dependence_todo.swift

@@ -11,18 +11,11 @@
 // REQUIRES: disabled
 
 
-// =============================================================================
-// Diagnostics that should not fail.
-
-// Recognize nested accesses as part of the same dependence scope.
-func bv_get_mutate(container: inout NC) -> _mutate(container) BV {
-  container.getBV()
-}
-
 // =============================================================================
 // Diagnostics that should fail.
 
-// Test that an unsafe dependence requires Builtin.lifetime_dependence.
+// @_unsafeResultDependsOn: Test that an unsafe dependence requires
+// Builtin.lifetime_dependence.
 //
 func bv_derive_local(bv: consuming BV) -> _consume(bv) BV {
   let bv2 = BV(bv.p, bv.i)