SILCombine: fix keypath optimization with optional chaining and classes

For optional chaining a swift_enum is created.
If the sub-projection is ending a begin_access in the some-branch, it also needs to be ended in the none-branch.

Fixes a compiler crash and/or a miscompile.

https://bugs.swift.org/browse/SR-14534
rdar://77224220

Author: eeckstein

lib/SILOptimizer/Utils/KeyPathProjector.cpp

@@ -230,7 +230,7 @@ class GettablePropertyProjector : public ComponentProjector {
       auto addr = builder.createAllocStack(loc, type);
 
       assertHasNoContext();
-      assert(getter->getArguments().size() == 2);
+      assert(getter->getConventions().getNumSILArguments());
 
       auto ref = builder.createFunctionRef(loc, getter);
       builder.createApply(loc, ref, subs, {addr, parentValue});
@@ -308,8 +308,8 @@ class SettablePropertyProjector : public GettablePropertyProjector {
           auto addr = builder.createAllocStack(loc, type);
 
           assertHasNoContext();
-          assert(getter->getArguments().size() == 2);
-          assert(setter->getArguments().size() == 2);
+          assert(getter->getConventions().getNumSILArguments());
+          assert(setter->getConventions().getNumSILArguments());
 
           // If this is a modify, we need to call the getter and
           // store the result in the writeback buffer.
@@ -468,10 +468,11 @@ class OptionalChainProjector : public ComponentProjector {
 public:
   OptionalChainProjector(const KeyPathPatternComponent &component,
                          std::unique_ptr<KeyPathProjector> parent,
+                         BeginAccessInst *&beginAccess,
                          SILValue optionalChainResult,
                          SILLocation loc, SILBuilder &builder)
         : ComponentProjector(component, std::move(parent), loc, builder),
-          optionalChainResult(optionalChainResult) {}
+          optionalChainResult(optionalChainResult), beginAccess(beginAccess) {}
 
   void project(AccessType accessType,
                std::function<void(SILValue addr)> callback) override {
@@ -506,6 +507,7 @@ class OptionalChainProjector : public ComponentProjector {
 
       // Unwrap the optional.
       auto objAddr = builder.createUncheckedTakeEnumDataAddr(loc, tempAddr, someDecl, objType);
+      BeginAccessInst *origBeginAccess = beginAccess;
 
       // at the end of the projection, callback will store a value in optionalChainResult
       callback(objAddr);
@@ -516,6 +518,12 @@ class OptionalChainProjector : public ComponentProjector {
       builder.createBranch(loc, continuation);
       // else, store nil in the result
       builder.setInsertionPoint(ifNone);
+
+      // If the sub-projection ended the access in the some-branch, we also
+      // have to end the access in the none-branch.
+      if (origBeginAccess && origBeginAccess != beginAccess)
+        builder.createEndAccess(loc, origBeginAccess, /*aborted*/ false);
+
       builder.createInjectEnumAddr(loc, optionalChainResult, noneDecl);
 
       builder.createBranch(loc, continuation);
@@ -526,6 +534,7 @@ class OptionalChainProjector : public ComponentProjector {
 
 private:
   SILValue optionalChainResult;
+  BeginAccessInst *&beginAccess;
 };
 
 /// A projector to handle a complete key path.
@@ -646,7 +655,8 @@ class CompleteKeyPathProjector : public KeyPathProjector {
         break;
       case KeyPathPatternComponent::Kind::OptionalChain:
         projector = std::make_unique<OptionalChainProjector>
-            (comp, std::move(parent), optionalChainResult, loc, builder);
+            (comp, std::move(parent), beginAccess, optionalChainResult, loc,
+             builder);
         break;
     }
 

test/SILOptimizer/sil_combine.sil

@@ -1,4 +1,4 @@
-// RUN: %target-sil-opt -enable-objc-interop -enforce-exclusivity=none -enable-sil-verify-all %s -sil-combine | %FileCheck %s
+// RUN: %target-sil-opt -enable-objc-interop -enable-sil-verify-all %s -sil-combine | %FileCheck %s
 
 // Declare this SIL to be canonical because some tests break raw SIL
 // conventions. e.g. address-type block args. -enforce-exclusivity=none is also
@@ -4226,3 +4226,42 @@ bb0:
   return %16 : $Int
 } // end sil function 'test_global_value'
 
+final class Kp1 {
+  @_hasStorage @_hasInitialValue final var b: Kp2? { get set }
+  @objc deinit
+  init()
+}
+
+class Kp2 {
+  @_hasStorage @_hasInitialValue var s: String { get set }
+  @objc deinit
+  init()
+}
+
+sil @kpgetter : $@convention(thin) (@in_guaranteed Kp2) -> @out String
+sil @kpsetter : $@convention(thin) (@in_guaranteed String, @in_guaranteed Kp2) -> ()
+sil @swift_getAtKeyPath : $@convention(thin) <τ_0_0, τ_0_1> (@in_guaranteed τ_0_0, @guaranteed KeyPath<τ_0_0, τ_0_1>) -> @out τ_0_1
+
+// Test keypath optimization for optional chaining with accesses scopes.
+// E.g.   \.a?.b   (on a class)
+//
+// CHECK-LABEL: sil @test_optional_chaining_keypath
+// CHECK:    [[BA:%[0-9]+]] = begin_access
+// CHECK:    switch_enum {{.*}} case #Optional.none!enumelt: bb2, case #Optional.some!enumelt: bb1
+// CHECK:  bb1:
+// CHECK:    apply
+// CHECK:    end_access [[BA]]
+// CHECK:  bb2:
+// CHECK:    end_access [[BA]]
+// CHECK:  bb3:
+// CHECK:     } // end sil function 'test_optional_chaining_keypath'
+sil @test_optional_chaining_keypath : $@convention(thin) (@in_guaranteed Kp1) -> @out Optional<String> {
+bb0(%0 : $*Optional<String>, %1 : $*Kp1):
+  %kp = keypath $KeyPath<Kp1, Optional<String>>, (root $Kp1; stored_property #Kp1.b : $Optional<Kp2>; optional_chain : $Kp2; settable_property $String,  id #Kp2.s!getter : (Kp2) -> () -> String, getter @kpgetter : $@convention(thin) (@in_guaranteed Kp2) -> @out String, setter @kpsetter : $@convention(thin) (@in_guaranteed String, @in_guaranteed Kp2) -> (); optional_wrap : $Optional<String>)
+  %f = function_ref @swift_getAtKeyPath : $@convention(thin) <τ_0_0, τ_0_1> (@in_guaranteed τ_0_0, @guaranteed KeyPath<τ_0_0, τ_0_1>) -> @out τ_0_1
+  %a = apply %f<Kp1, Optional<String>>(%0, %1, %kp) : $@convention(thin) <τ_0_0, τ_0_1> (@in_guaranteed τ_0_0, @guaranteed KeyPath<τ_0_0, τ_0_1>) -> @out τ_0_1
+  strong_release %kp : $KeyPath<Kp1, Optional<String>>
+  %r = tuple ()
+  return %r : $()
+}
+