Merge pull request #14045 from rudkx/keypaths-fix-iuo-on-master

Add support for key paths for the new IUO implementation.

Author: rudkx

lib/Sema/CSApply.cpp

@@ -4175,6 +4175,24 @@ namespace {
           continue;
         }
 
+        auto getObjectType = [](Type optionalTy) -> Type {
+          Type objectTy;
+          if (auto lvalue = optionalTy->getAs<LValueType>()) {
+            objectTy = lvalue->getObjectType()->getAnyOptionalObjectType();
+            if (optionalTy->hasUnresolvedType() && !objectTy) {
+              objectTy = optionalTy;
+            }
+            objectTy = LValueType::get(objectTy);
+          } else {
+            objectTy = optionalTy->getAnyOptionalObjectType();
+            if (optionalTy->hasUnresolvedType() && !objectTy) {
+              objectTy = optionalTy;
+            }
+          }
+          assert(objectTy);
+          return objectTy;
+        };
+
         KeyPathExpr::Component component;
         switch (auto kind = origComponent.getKind()) {
         case KeyPathExpr::Component::Kind::UnresolvedProperty: {
@@ -4220,8 +4238,9 @@ namespace {
             else
               baseTy = objTy;
 
+            auto loc = origComponent.getLoc();
             resolvedComponents.push_back(
-                 KeyPathExpr::Component::forOptionalForce(baseTy, SourceLoc()));
+                KeyPathExpr::Component::forOptionalForce(baseTy, loc));
           }
 
           cs.TC.requestMemberLayout(property);
@@ -4241,6 +4260,17 @@ namespace {
           component = KeyPathExpr::Component::forProperty(ref,
                                                        resolvedTy,
                                                        origComponent.getLoc());
+
+          baseTy = component.getComponentType();
+          resolvedComponents.push_back(component);
+
+          if (shouldForceUnwrapResult(property, locator)) {
+            auto objectTy = getObjectType(baseTy);
+            auto loc = origComponent.getLoc();
+            component = KeyPathExpr::Component::forOptionalForce(objectTy, loc);
+            baseTy = component.getComponentType();
+            resolvedComponents.push_back(component);
+          }
           break;
         }
         case KeyPathExpr::Component::Kind::UnresolvedSubscript: {
@@ -4268,8 +4298,9 @@ namespace {
             else
               baseTy = objTy;
 
+            auto loc = origComponent.getLoc();
             resolvedComponents.push_back(
-                 KeyPathExpr::Component::forOptionalForce(baseTy, SourceLoc()));
+                KeyPathExpr::Component::forOptionalForce(baseTy, loc));
           }
 
           cs.TC.requestMemberLayout(subscript);
@@ -4306,6 +4337,17 @@ namespace {
           // Save a reference to the component so we can do a post-pass to check
           // the Hashable conformance of the indexes.
           KeyPathSubscriptComponents.push_back({E, resolvedComponents.size()});
+
+          baseTy = component.getComponentType();
+          resolvedComponents.push_back(component);
+
+          if (shouldForceUnwrapResult(subscript, locator)) {
+            auto objectTy = getObjectType(baseTy);
+            auto loc = origComponent.getLoc();
+            component = KeyPathExpr::Component::forOptionalForce(objectTy, loc);
+            baseTy = component.getComponentType();
+            resolvedComponents.push_back(component);
+          }
           break;
         }
         case KeyPathExpr::Component::Kind::OptionalChain: {
@@ -4318,43 +4360,34 @@ namespace {
           }
           assert(objectTy);
 
-          component = KeyPathExpr::Component::forOptionalChain(objectTy,
-                                                        origComponent.getLoc());
+          auto loc = origComponent.getLoc();
+          component = KeyPathExpr::Component::forOptionalChain(objectTy, loc);
+
+          baseTy = component.getComponentType();
+          resolvedComponents.push_back(component);
           break;
         }
         case KeyPathExpr::Component::Kind::OptionalForce: {
-          Type objectTy;
-          if (auto lvalue = baseTy->getAs<LValueType>()) {
-            objectTy = lvalue->getObjectType()->getAnyOptionalObjectType();
-            if (baseTy->hasUnresolvedType() && !objectTy) {
-              objectTy = baseTy;
-            }
-            objectTy = LValueType::get(objectTy);
-          } else {
-            objectTy = baseTy->getAnyOptionalObjectType();
-            if (baseTy->hasUnresolvedType() && !objectTy) {
-              objectTy = baseTy;
-            }
-            assert(objectTy);
-          }
-          
-          component = KeyPathExpr::Component::forOptionalForce(objectTy,
-                                                        origComponent.getLoc());
+          auto objectTy = getObjectType(baseTy);
+          auto loc = origComponent.getLoc();
+          component = KeyPathExpr::Component::forOptionalForce(objectTy, loc);
+          baseTy = component.getComponentType();
+          resolvedComponents.push_back(component);
           break;
         }
         case KeyPathExpr::Component::Kind::Invalid:
           component = origComponent;
           component.setComponentType(leafTy);
+
+          baseTy = component.getComponentType();
+          resolvedComponents.push_back(component);
           break;
 
         case KeyPathExpr::Component::Kind::Property:
         case KeyPathExpr::Component::Kind::Subscript:
         case KeyPathExpr::Component::Kind::OptionalWrap:
           llvm_unreachable("already resolved");
         }
-
-        baseTy = component.getComponentType();
-        resolvedComponents.push_back(component);
       }
 
       // Wrap a non-optional result if there was chaining involved.
@@ -4389,7 +4422,8 @@ namespace {
 
       // The final component type ought to line up with the leaf type of the
       // key path.
-      assert(!baseTy || baseTy->getWithoutSpecifierType()->isEqual(leafTy));
+      assert(!baseTy || baseTy->hasUnresolvedType()
+             || baseTy->getWithoutSpecifierType()->isEqual(leafTy));
       return E;
     }