Merge pull request #11004 from jckarter/iuo-key-path

Handle IUO unwraps in key paths.

Author: jckarter

lib/SIL/SILVerifier.cpp

@@ -3766,20 +3766,20 @@ class SILVerifier : public SILVerifierBase<SILVerifier> {
           break;
         }
         case KeyPathPatternComponent::Kind::OptionalChain: {
-          require(OptionalType::get(componentTy)->isEqual(baseTy),
+          require(baseTy->getAnyOptionalObjectType()->isEqual(componentTy),
                   "chaining component should unwrap optional");
           require(leafTy->getAnyOptionalObjectType(),
                   "key path with chaining component should have optional "
                   "result");
           break;
         }
         case KeyPathPatternComponent::Kind::OptionalForce: {
-          require(OptionalType::get(componentTy)->isEqual(baseTy),
+          require(baseTy->getAnyOptionalObjectType()->isEqual(componentTy),
                   "forcing component should unwrap optional");
           break;
         }
         case KeyPathPatternComponent::Kind::OptionalWrap: {
-          require(OptionalType::get(baseTy)->isEqual(componentTy),
+          require(componentTy->getAnyOptionalObjectType()->isEqual(baseTy),
                   "wrapping component should wrap optional");
           break;
         }

lib/Sema/CSApply.cpp

@@ -3989,6 +3989,19 @@ namespace {
             }
           }
 
+          // Unwrap if we needed to look through an IUO to find the
+          // property.
+          if (auto objTy = cs.lookThroughImplicitlyUnwrappedOptionalType(
+                              baseTy->getRValueType())) {
+            if (baseTy->is<LValueType>())
+              baseTy = LValueType::get(objTy);
+            else
+              baseTy = objTy;
+
+            resolvedComponents.push_back(
+                 KeyPathExpr::Component::forOptionalForce(baseTy, SourceLoc()));
+          }
+          
           auto dc = property->getInnermostDeclContext();
           SmallVector<Substitution, 4> subs;
           if (auto sig = dc->getGenericSignatureOfContext()) {
@@ -4020,6 +4033,19 @@ namespace {
                            diag::expr_keypath_mutating_getter,
                            subscript->getFullName());
           }
+
+          // Unwrap if we needed to look through an IUO to find the
+          // subscript.
+          if (auto objTy = cs.lookThroughImplicitlyUnwrappedOptionalType(
+                              baseTy->getRValueType())) {
+            if (baseTy->is<LValueType>())
+              baseTy = LValueType::get(objTy);
+            else
+              baseTy = objTy;
+
+            resolvedComponents.push_back(
+                 KeyPathExpr::Component::forOptionalForce(baseTy, SourceLoc()));
+          }
 
           auto dc = subscript->getInnermostDeclContext();
           SmallVector<Substitution, 4> subs;

stdlib/public/core/KeyPath.swift

@@ -1020,7 +1020,9 @@ internal struct RawKeyPathComponent {
       return .continue(get(base, argument?.data.baseAddress ?? rawGet))
 
     case .optionalChain:
-      _sanityCheck(CurValue.self == Optional<NewValue>.self,
+      // TODO: IUO shouldn't be a first class type
+      _sanityCheck(CurValue.self == Optional<NewValue>.self
+                   || CurValue.self == ImplicitlyUnwrappedOptional<NewValue>.self,
                    "should be unwrapping optional value")
       _sanityCheck(_isOptional(LeafValue.self),
                    "leaf result should be optional")
@@ -1033,12 +1035,16 @@ internal struct RawKeyPathComponent {
       }
 
     case .optionalForce:
-      _sanityCheck(CurValue.self == Optional<NewValue>.self,
+      // TODO: IUO shouldn't be a first class type
+      _sanityCheck(CurValue.self == Optional<NewValue>.self
+                   || CurValue.self == ImplicitlyUnwrappedOptional<NewValue>.self,
                    "should be unwrapping optional value")
       return .continue(unsafeBitCast(base, to: Optional<NewValue>.self)!)
 
     case .optionalWrap:
-      _sanityCheck(NewValue.self == Optional<CurValue>.self,
+      // TODO: IUO shouldn't be a first class type
+      _sanityCheck(NewValue.self == Optional<CurValue>.self
+                   || CurValue.self == ImplicitlyUnwrappedOptional<CurValue>.self,
                    "should be wrapping optional value")
       return .continue(
         unsafeBitCast(base as Optional<CurValue>, to: NewValue.self))
@@ -1121,7 +1127,9 @@ internal struct RawKeyPathComponent {
       return UnsafeRawPointer(Builtin.addressof(&writeback.value))
 
     case .optionalForce:
-      _sanityCheck(CurValue.self == Optional<NewValue>.self,
+      // TODO: ImplicitlyUnwrappedOptional should not be a first-class type
+      _sanityCheck(CurValue.self == Optional<NewValue>.self
+                   || CurValue.self == ImplicitlyUnwrappedOptional<NewValue>.self,
                    "should be unwrapping an optional value")
       // Optional's layout happens to always put the payload at the start
       // address of the Optional value itself, if a value is present at all.

test/SILGen/keypaths.swift

@@ -240,3 +240,29 @@ func keyPathForStorageQualified() {
   // CHECK-SAME: settable_property $StorageQualified, id ##FinalStorageQualified.disowned
   _ = \FinalStorageQualified.disowned
 }
+
+struct IUOProperty {
+  var iuo: IUOBlob!
+}
+
+struct IUOBlob {
+  var x: Int
+  subscript(y: String) -> String {
+    get { return y }
+    set {}
+  }
+}
+
+// CHECK-LABEL: sil hidden @{{.*}}iuoKeyPaths
+func iuoKeyPaths() {
+  // CHECK: = keypath $WritableKeyPath<IUOProperty, Int>,
+  // CHECK-SAME: stored_property #IUOProperty.iuo
+  // CHECK-SAME: optional_force
+  // CHECK-SAME: stored_property #IUOBlob.x
+  _ = \IUOProperty.iuo.x
+  // CHECK: = keypath $WritableKeyPath<IUOProperty, Int>,
+  // CHECK-SAME: stored_property #IUOProperty.iuo
+  // CHECK-SAME: optional_force
+  // CHECK-SAME: stored_property #IUOBlob.x
+  _ = \IUOProperty.iuo!.x
+}

test/stdlib/KeyPath.swift

@@ -523,4 +523,27 @@ keyPath.test("writebacks nest properly") {
   expectEqual(nestedWritebackLog, 0x383736)
 }
 
+struct IUOWrapper {
+  var wrapped: IUOWrapped!
+}
+
+struct IUOWrapped {
+  var value: Int
+}
+
+keyPath.test("IUO and key paths") {
+  var subject = IUOWrapper(wrapped: IUOWrapped(value: 1989))
+  let kp1 = \IUOWrapper.wrapped.value
+
+  expectEqual(subject[keyPath: kp1], 1989)
+  subject[keyPath: kp1] = 1738
+  expectEqual(subject[keyPath: kp1], 1738)
+  expectEqual(subject.wrapped.value, 1738)
+
+  let kp2 = \IUOWrapper.wrapped!.value
+
+  expectEqual(kp1, kp2)
+  expectEqual(kp1.hashValue, kp2.hashValue)
+}
+
 runAllTests()