ReleaseDevirtualizer: improve finding the last release of an allocation

With OSSA it can happen more easily that the final release is not immediately located before the related dealloc_stack_ref.
Therefore do a more precise check (using escape analysis) if any instruction between a release and a dealloc_stack_ref might (implicitly) release the allocated object.

This is part of fixing regressions when enabling OSSA modules:
rdar://140229560

SwiftCompilerSources/Sources/Optimizer/FunctionPasses/ReleaseDevirtualizer.swift

@@ -34,47 +34,18 @@ import SIL
 let releaseDevirtualizerPass = FunctionPass(name: "release-devirtualizer") {
   (function: Function, context: FunctionPassContext) in
 
-  for block in function.blocks {
-    // The last `release_value`` or `strong_release`` instruction before the
-    // deallocation.
-    var lastRelease: RefCountingInst?
-
-    for instruction in block.instructions {
-      switch instruction {
-      case let dealloc as DeallocStackRefInst:
-        if let lastRel = lastRelease {
-          // We only do the optimization for stack promoted object, because for
-          // these we know that they don't have associated objects, which are
-          // _not_ released by the deinit method.
-          if !context.continueWithNextSubpassRun(for: lastRel) {
-            return
-          }
-          tryDevirtualizeRelease(of: dealloc.allocRef, lastRelease: lastRel, context)
-          lastRelease = nil
-        }
-      case let strongRelease as StrongReleaseInst:
-        lastRelease = strongRelease
-      case let releaseValue as ReleaseValueInst where releaseValue.value.type.containsSingleReference(in: function):
-        lastRelease = releaseValue
-      case is DeallocRefInst, is BeginDeallocRefInst:
-        lastRelease = nil
-      default:
-        if instruction.mayRelease {
-          lastRelease = nil
-        }
+  for inst in function.instructions {
+    if let dealloc = inst as? DeallocStackRefInst {
+      if !context.continueWithNextSubpassRun(for: dealloc) {
+        return
       }
+      tryDevirtualizeRelease(of: dealloc, context)
     }
   }
 }
 
-/// Tries to de-virtualize the final release of a stack-promoted object.
-private func tryDevirtualizeRelease(
-  of allocRef: AllocRefInstBase,
-  lastRelease: RefCountingInst,
-  _ context: FunctionPassContext
-) {
-  var downWalker = FindReleaseWalker(release: lastRelease)
-  guard let pathToRelease = downWalker.getPathToRelease(from: allocRef) else {
+private func tryDevirtualizeRelease(of dealloc: DeallocStackRefInst, _ context: FunctionPassContext) {
+  guard let (lastRelease, pathToRelease) = findLastRelease(of: dealloc, context) else {
     return
   }
 
@@ -82,6 +53,7 @@ private func tryDevirtualizeRelease(
     return
   }
 
+  let allocRef = dealloc.allocRef
   var upWalker = FindAllocationWalker(allocation: allocRef)
   if upWalker.walkUp(value: lastRelease.operand.value, path: pathToRelease) == .abortWalk {
     return
@@ -120,9 +92,55 @@ private func tryDevirtualizeRelease(
   context.erase(instruction: lastRelease)
 }
 
+private func findLastRelease(
+  of dealloc: DeallocStackRefInst,
+  _ context: FunctionPassContext
+) -> (lastRelease: RefCountingInst, pathToRelease: SmallProjectionPath)? {
+  let allocRef = dealloc.allocRef
+
+  // Search for the final release in the same basic block of the dealloc.
+  for instruction in ReverseInstructionList(first: dealloc.previous) {
+    switch instruction {
+    case let strongRelease as StrongReleaseInst:
+      if let pathToRelease = getPathToRelease(from: allocRef, to: strongRelease) {
+        return (strongRelease, pathToRelease)
+      }
+    case let releaseValue as ReleaseValueInst:
+      if releaseValue.value.type.containsSingleReference(in: dealloc.parentFunction) {
+        if let pathToRelease = getPathToRelease(from: allocRef, to: releaseValue) {
+          return (releaseValue, pathToRelease)
+        }
+      }
+    case is BeginDeallocRefInst, is DeallocRefInst:
+      // Check if the last release was already de-virtualized.
+      if allocRef.escapes(to: instruction, context) {
+        return nil
+      }
+    default:
+      break
+    }
+    if instruction.mayRelease && allocRef.escapes(to: instruction, context) {
+      // This instruction may release the allocRef, which means that any release we find
+      // earlier in the block is not guaranteed to be the final release.
+      return nil
+    }
+  }
+  return nil
+}
+
+// If the release is a release_value it might release a struct which _contains_ the allocated object.
+// Return a projection path to the contained object in this case.
+private func getPathToRelease(from allocRef: AllocRefInstBase, to release: RefCountingInst) -> SmallProjectionPath? {
+  var downWalker = FindReleaseWalker(release: release)
+  if downWalker.walkDownUses(ofValue: allocRef, path: SmallProjectionPath()) == .continueWalk {
+    return downWalker.result
+  }
+  return nil
+}
+
 private struct FindReleaseWalker : ValueDefUseWalker {
   private let release: RefCountingInst
-  private var result: SmallProjectionPath? = nil
+  private(set) var result: SmallProjectionPath? = nil
 
   var walkDownCache = WalkerCache<SmallProjectionPath>()
 
@@ -130,13 +148,6 @@ private struct FindReleaseWalker : ValueDefUseWalker {
     self.release = release
   }
 
-  mutating func getPathToRelease(from allocRef: AllocRefInstBase) -> SmallProjectionPath? {
-    if walkDownUses(ofValue: allocRef, path: SmallProjectionPath()) == .continueWalk {
-      return result
-    }
-    return nil
-  }
-
   mutating func leafUse(value: Operand, path: SmallProjectionPath) -> WalkResult {
     if value.instruction == release {
       if let existingResult = result {
@@ -149,6 +160,23 @@ private struct FindReleaseWalker : ValueDefUseWalker {
   }
 }
 
+private extension AllocRefInstBase {
+  func escapes(to instruction: Instruction, _ context: FunctionPassContext) -> Bool {
+    return self.isEscaping(using: EscapesToInstructionVisitor(target: instruction), context)
+  }
+}
+
+private struct EscapesToInstructionVisitor : EscapeVisitor {
+  let target: Instruction
+
+  mutating func visitUse(operand: Operand, path: EscapePath) -> UseResult {
+    if operand.instruction == target {
+      return .abort
+    }
+    return .continueWalk
+  }
+}
+
 // Up-walker to find the root of a release instruction.
 private struct FindAllocationWalker : ValueUseDefWalker {
   private let allocInst: AllocRefInstBase

test/SILOptimizer/devirt_release.sil

@@ -18,6 +18,18 @@ struct Str {
   @_hasStorage var b: B
 }
 
+sil @unknown_func : $@convention(thin) () -> ()
+sil @unknown_releaseing_func : $@convention(thin) (@owned B) -> ()
+
+sil @$s4test1BCfD : $@convention(method) (@owned B) -> () {
+// Note that the destructor is not destroying, but it must block the optimization.
+[global: ]
+bb0(%0 : $B):
+  dealloc_ref %0
+  %r = tuple ()
+  return %r
+}
+
 // CHECK-LABEL: sil @devirtualize_object
 // CHECK: [[A:%[0-9]+]] = alloc_ref
 // CHECK-NEXT: [[BD:%.*]] = begin_dealloc_ref [[A]] : $B of [[A]]
@@ -71,37 +83,69 @@ bb0:
   return %r : $()
 }
 
-// CHECK-LABEL: sil @dont_devirtualize_wrong_release
-// CHECK: alloc_ref
-// CHECK-NEXT: strong_release
-// CHECK-NEXT: strong_release
-// CHECK-NEXT: dealloc_ref
-// CHECK: return
-// CHECK: } // end sil function 'dont_devirtualize_wrong_release'
-sil @dont_devirtualize_wrong_release : $@convention(thin) (@owned B) -> () {
-bb0(%0 : $B):
-  %1 = alloc_ref $B
+// CHECK-LABEL: sil @dont_devirtualize_devirtualized_not_inlined :
+// CHECK:         alloc_ref
+// CHECK-NEXT:    strong_retain
+// CHECK-NEXT:    strong_release
+// CHECK-NEXT:    begin_dealloc_ref
+// CHECK: } // end sil function 'dont_devirtualize_devirtualized_not_inlined'
+sil @dont_devirtualize_devirtualized_not_inlined : $@convention(thin) () -> () {
+bb0:
+  %1 = alloc_ref [stack] $B
+  strong_retain %1 : $B
   strong_release %1 : $B
-  strong_release %0 : $B
-  dealloc_ref %1 : $B
+  %6 = begin_dealloc_ref %1 : $B of %1 : $B
+  %7 = function_ref @$s4test1BCfD : $@convention(method) (@owned B) -> ()
+  %8 = apply %7(%6) : $@convention(method) (@owned B) -> ()
+  dealloc_stack_ref %1 : $B
   %r = tuple ()
   return %r : $()
 }
 
-// CHECK-LABEL: sil @dont_devirtualize_unknown_release
-// CHECK: alloc_ref
-// CHECK-NEXT: strong_release
-// CHECK: [[F:%[0-9]+]] = function_ref @unknown_func
-// CHECK-NEXT: apply [[F]]()
-// CHECK-NEXT: dealloc_ref
-// CHECK: } // end sil function 'dont_devirtualize_unknown_release'
-sil @dont_devirtualize_unknown_release : $@convention(thin) (@owned B) -> () {
+// CHECK-LABEL: sil @devirtualize_with_other_release_in_between :
+// CHECK:         %1 = alloc_ref [stack] $B
+// CHECK-NOT:     release
+// CHECK:         apply
+// CHECK-NEXT:    strong_release %0
+// CHECK-NEXT:    dealloc_stack_ref %1
+// CHECK:       } // end sil function 'devirtualize_with_other_release_in_between'
+sil @devirtualize_with_other_release_in_between : $@convention(thin) (@owned B) -> () {
 bb0(%0 : $B):
-  %1 = alloc_ref $B
+  %1 = alloc_ref [stack] $B
+  strong_release %1 : $B
+  strong_release %0 : $B
+  dealloc_stack_ref %1 : $B
+  %r = tuple ()
+  return %r : $()
+}
+
+// CHECK-LABEL: sil @devirtualize_with_apply_in_between :
+// CHECK:         alloc_ref
+// CHECK-NEXT:    begin_dealloc_ref
+// CHECK-NOT:     release
+// CHECK:       } // end sil function 'devirtualize_with_apply_in_between'
+sil @devirtualize_with_apply_in_between : $@convention(thin) (@owned B) -> () {
+bb0(%0 : $B):
+  %1 = alloc_ref [stack] $B
   strong_release %1 : $B
   %f = function_ref @unknown_func : $@convention(thin) () -> ()
   %a = apply %f() : $@convention(thin) () -> ()
-  dealloc_ref %1 : $B
+  dealloc_stack_ref %1 : $B
+  %r = tuple ()
+  return %r : $()
+}
+
+// CHECK-LABEL: sil @dont_devirtualize_unknown_release :
+// CHECK:         alloc_ref
+// CHECK-NEXT:    strong_release
+// CHECK:       } // end sil function 'dont_devirtualize_unknown_release'
+sil @dont_devirtualize_unknown_release : $@convention(thin) (@owned B) -> () {
+bb0(%0 : $B):
+  %1 = alloc_ref [stack] $B
+  strong_release %1 : $B
+  %f = function_ref @unknown_releaseing_func : $@convention(thin) (@owned B) -> ()
+  %a = apply %f(%1) : $@convention(thin) (@owned B) -> ()
+  dealloc_stack_ref %1 : $B
   %r = tuple ()
   return %r : $()
 }
@@ -205,10 +249,24 @@ bb3(%a : $B):
   return %r : $()
 }
 
-
-sil @unknown_func : $@convention(thin) () -> ()
-
-sil @$s4test1BCfD : $@convention(method) (@owned B) -> ()
+// CHECK-LABEL: sil @devirtualize_double_release :
+// CHECK-NOT:     release
+// CHECK: } // end sil function 'devirtualize_double_release'
+sil @devirtualize_double_release : $@convention(thin) () -> () {
+bb0:
+  %0 = alloc_ref [stack] $B
+  %1 = end_init_let_ref %0
+  %2 = alloc_ref [stack] $B
+  %3 = end_init_let_ref %2
+  strong_release %1
+  strong_release %3
+  %4 = function_ref @unknown_func : $@convention(thin) () -> ()
+  %5 = apply %4() : $@convention(thin) () -> ()
+  dealloc_stack_ref %2
+  dealloc_stack_ref %0
+  %r = tuple ()
+  return %r
+}
 
 sil_vtable B {
   #B.deinit!deallocator: @$s4test1BCfD