[Devirtualizer] Replace begin_apply allocations.

When devirtualizing a yield_once_2 begin_apply, replace all uses of the
allocation address, just as is done with the token.

Author: nate-chandler

lib/SIL/Parser/ParseSIL.cpp

@@ -1291,7 +1291,9 @@ static std::optional<AccessorKind> getAccessorKind(StringRef ident) {
       .Case("addressor", AccessorKind::Address)
       .Case("mutableAddressor", AccessorKind::MutableAddress)
       .Case("read", AccessorKind::Read)
+      .Case("read2", AccessorKind::Read2)
       .Case("modify", AccessorKind::Modify)
+      .Case("modify2", AccessorKind::Modify2)
       .Default(std::nullopt);
 }
 

lib/SILOptimizer/Utils/Devirtualize.cpp

@@ -619,6 +619,10 @@ replaceBeginApplyInst(SILBuilder &builder, SILPassManager *pm, SILLocation loc,
   // Forward the token.
   oldBAI->getTokenResult()->replaceAllUsesWith(newBAI->getTokenResult());
 
+  if (auto *allocation = oldBAI->getCalleeAllocationResult()) {
+    allocation->replaceAllUsesWith(newBAI->getCalleeAllocationResult());
+  }
+
   auto oldYields = oldBAI->getYieldedValues();
   auto newYields = newBAI->getYieldedValues();
   assert(oldYields.size() == newYields.size());

test/SILOptimizer/devirtualize_coroutine_accessors.sil

@@ -0,0 +1,50 @@
+// RUN: %target-sil-opt \
+// RUN:     -devirtualizer \
+// RUN:     %s \
+// RUN:     -enable-sil-verify-all \
+// RUN:     -enable-experimental-feature CoroutineAccessors \
+// RUN: | %FileCheck %s
+
+// REQUIRES: swift_feature_CoroutineAccessors
+
+import Builtin
+
+struct Int {
+  var _value: Builtin.Word
+}
+
+class Base {
+  var x : Int { read }
+}
+
+class Derived : Base {
+  override var x : Int { read }
+}
+
+// CHECK-LABEL: sil [ossa] @utilize_begin_apply : {{.*}} {
+// CHECK-NOT:     class_method
+// CHECK-LABEL: } // end sil function 'utilize_begin_apply'
+sil [ossa] @utilize_begin_apply : $@convention(thin) () -> () {
+bb0:
+  %derived = alloc_ref $Derived
+  %base = upcast %derived : $Derived to $Base
+  %reader = class_method %base : $Base, #Base.x!read2 : (Base) -> () -> (), $@yield_once_2 @convention(method) (@guaranteed Base) -> @yields Int
+  (%value, %token, %allocation) = begin_apply %reader(%base) : $@yield_once_2 @convention(method) (@guaranteed Base) -> @yields Int
+  end_apply %token as $()
+  dealloc_stack %allocation : $*Builtin.SILToken
+  destroy_value %base : $Base
+  %retval = tuple ()
+  return %retval : $()
+}
+
+sil [ossa] @Base_x_read2 : $@yield_once_2 @convention(method) (@guaranteed Base) -> @yields Int
+
+sil [ossa] @Derived_x_read2 : $@yield_once_2 @convention(method) (@guaranteed Base) -> @yields Int
+
+sil_vtable Base {
+  #Base.x!read2: @Base_x_read2
+}
+
+sil_vtable Derived {
+  #Base.x!read2: @Derived_x_read2 [override]
+}