Merge pull request swiftlang#32843 from LucianoPAlmeida/keypath-root-diagnostic-improvement

[Diagnostics][Sema] Improving diagnostics for optional key path application root/base mismatch

Author: LucianoPAlmeida

include/swift/AST/DiagnosticsSema.def

@@ -1077,7 +1077,9 @@ NOTE(optional_base_chain,none,
      "base values", (DeclNameRef))
 NOTE(optional_base_remove_optional_for_keypath_root, none,
      "use unwrapped type %0 as key path root", (Type))
-
+NOTE(optional_keypath_application_base, none,
+     "use '?' to access key path subscript only for non-'nil' base values", ())
+    
 ERROR(missing_unwrap_optional_try,none,
       "value of optional type %0 not unwrapped; did you mean to use 'try!' "
       "or chain with '?'?",

lib/Sema/CSDiagnostics.cpp

@@ -6493,3 +6493,19 @@ void MissingRawValueFailure::fixIt(InFlightDiagnostic &diagnostic) const {
 
   diagnostic.fixItInsertAfter(range.End, fix);
 }
+
+bool MissingOptionalUnwrapKeyPathFailure::diagnoseAsError() {
+  emitDiagnostic(diag::optional_not_unwrapped, getFromType(),
+                 getFromType()->lookThroughSingleOptionalType());
+  
+  emitDiagnostic(diag::optional_keypath_application_base)
+      .fixItInsertAfter(getLoc(), "?");
+  emitDiagnostic(diag::unwrap_with_force_value)
+      .fixItInsertAfter(getLoc(), "!");
+  return true;
+}
+
+SourceLoc MissingOptionalUnwrapKeyPathFailure::getLoc() const {
+  auto *SE = castToExpr<SubscriptExpr>(getAnchor());
+  return SE->getBase()->getEndLoc();
+}

lib/Sema/CSDiagnostics.h

@@ -2191,6 +2191,24 @@ class MissingRawValueFailure final : public AbstractRawRepresentableFailure {
   void fixIt(InFlightDiagnostic &diagnostic) const override;
 };
 
+/// Diagnose a key path optional base that should be unwraped in order to 
+/// apply key path subscript.
+///
+/// \code
+/// func f(_ bar: Bar? , keyPath: KeyPath<Bar, Int>) {
+///   bar[keyPath: keyPath]
+/// }
+/// \endcode
+class MissingOptionalUnwrapKeyPathFailure final : public ContextualFailure {
+public:
+  MissingOptionalUnwrapKeyPathFailure(const Solution &solution, Type lhs,
+                                      Type rhs, ConstraintLocator *locator)
+      : ContextualFailure(solution, lhs, rhs, locator) {}
+
+  bool diagnoseAsError() override;
+  SourceLoc getLoc() const override;
+};
+
 } // end namespace constraints
 } // end namespace swift
 

lib/Sema/CSFix.cpp

@@ -1508,3 +1508,36 @@ bool SpecifyKeyPathRootType::diagnose(const Solution &solution,
 
   return failure.diagnose(asNote);
 }
+
+bool UnwrapOptionalBaseKeyPathApplication::diagnose(const Solution &solution,
+                                                    bool asNote) const {
+  MissingOptionalUnwrapKeyPathFailure failure(solution, getFromType(),
+                                              getToType(), getLocator());
+  return failure.diagnose(asNote);
+}
+
+UnwrapOptionalBaseKeyPathApplication *
+UnwrapOptionalBaseKeyPathApplication::attempt(ConstraintSystem &cs, Type baseTy,
+                                              Type rootTy,
+                                              ConstraintLocator *locator) {
+  if(baseTy->hasTypeVariable() || rootTy->hasTypeVariable())
+    return nullptr;
+
+  if (!isExpr<SubscriptExpr>(locator->getAnchor()))
+    return nullptr;
+  
+  // Only diagnose this if base is an optional type and we only have a
+  // single level of optionality so we can safely suggest unwrapping.
+  auto nonOptionalTy = baseTy->getOptionalObjectType();
+  if (!nonOptionalTy || nonOptionalTy->getOptionalObjectType())
+    return nullptr;
+
+  auto result =
+      cs.matchTypes(nonOptionalTy, rootTy, ConstraintKind::Subtype,
+                    ConstraintSystem::TypeMatchFlags::TMF_ApplyingFix, locator);
+  if (result.isFailure())
+    return nullptr;
+
+  return new (cs.getAllocator())
+      UnwrapOptionalBaseKeyPathApplication(cs, baseTy, rootTy, locator);
+}

lib/Sema/CSFix.h

@@ -268,7 +268,9 @@ enum class FixKind : uint8_t {
 
   /// Specify key path root type when it cannot be infered from context.
   SpecifyKeyPathRootType,
-
+  
+  /// Unwrap optional base on key path application.
+  UnwrapOptionalBaseKeyPathApplication,
 };
 
 class ConstraintFix {
@@ -1918,6 +1920,32 @@ class SpecifyKeyPathRootType final : public ConstraintFix {
                                           ConstraintLocator *locator);
 };
 
+/// Diagnose missing unwrap of optional base type on key path application.
+///
+/// \code
+/// func f(_ bar: Bar? , keyPath: KeyPath<Bar, Int>) {
+///   bar[keyPath: keyPath]
+/// }
+/// \endcode
+class UnwrapOptionalBaseKeyPathApplication final : public ContextualMismatch {
+protected:
+  UnwrapOptionalBaseKeyPathApplication(ConstraintSystem &cs, Type lhs, Type rhs,
+                                       ConstraintLocator *locator)
+      : ContextualMismatch(cs, FixKind::UnwrapOptionalBaseKeyPathApplication,
+                           lhs, rhs, locator) {}
+
+public:
+  std::string getName() const override {
+    return "force unwrap base on key path application";
+  }
+
+  bool diagnose(const Solution &solution, bool asNote = false) const override;
+
+  static UnwrapOptionalBaseKeyPathApplication *
+  attempt(ConstraintSystem &cs, Type baseTy, Type rootTy,
+          ConstraintLocator *locator);
+};
+
 } // end namespace constraints
 } // end namespace swift
 

lib/Sema/CSSimplify.cpp

@@ -3774,6 +3774,14 @@ bool ConstraintSystem::repairFailures(
   }
 
   case ConstraintLocator::KeyPathRoot: {
+    // The root mismatch is from base U? to U or a subtype of U in keypath 
+    // application so let's suggest an unwrap the optional fix.
+    if (auto unwrapFix = UnwrapOptionalBaseKeyPathApplication::attempt(
+            *this, lhs, rhs, getConstraintLocator(locator))) {
+      conversionsOrFixes.push_back(unwrapFix);
+      break;
+    }
+
     conversionsOrFixes.push_back(AllowKeyPathRootTypeMismatch::create(
         *this, lhs, rhs, getConstraintLocator(locator)));
 
@@ -9742,6 +9750,7 @@ ConstraintSystem::SolutionKind ConstraintSystem::simplifyFixConstraint(
   case FixKind::CoerceToCheckedCast:
   case FixKind::SpecifyObjectLiteralTypeImport:
   case FixKind::AllowKeyPathRootTypeMismatch:
+  case FixKind::UnwrapOptionalBaseKeyPathApplication:
   case FixKind::AllowCoercionToForceCast:
   case FixKind::SpecifyKeyPathRootType: {
     return recordFix(fix) ? SolutionKind::Error : SolutionKind::Solved;

localization/diagnostics/en.yaml

@@ -4292,6 +4292,10 @@
   msg: >-
     use unwrapped type %0 as key path root
 
+- id: optional_keypath_application_base 
+  msg: >-
+    use '?' to access key path subscript only for non-'nil' base values
+
 - id: missing_unwrap_optional_try
   msg: >-
     value of optional type %0 not unwrapped; did you mean to use 'try!' or

test/Constraints/keypath.swift

@@ -152,3 +152,30 @@ func fSR11184_O(_ c: SR11184!, _ kp: ReferenceWritableKeyPath<SR11184, String?>,
   c![keyPath: kp] = str // OK
   c?[keyPath: kp] = str // OK
 }
+
+class KeyPathBase {}
+class KeyPathBaseSubtype: KeyPathBase {}
+class AnotherBase {}
+class AnotherComposeBase {
+  var member: KeyPathBase?
+}
+
+func key_path_root_mismatch<T>(_ base: KeyPathBase?, subBase: KeyPathBaseSubtype?, _ abase: AnotherComposeBase,
+                               _ kp: KeyPath<KeyPathBase, T>, _ kpa: KeyPath<AnotherBase, T>) {
+  let _ : T = base[keyPath: kp] // expected-error {{value of optional type 'KeyPathBase?' must be unwrapped to a value of type 'KeyPathBase'}}
+  // expected-note@-1 {{force-unwrap using '!' to abort execution if the optional value contains 'nil'}} {{19-19=!}}
+  // expected-note@-2 {{use '?' to access key path subscript only for non-'nil' base values}} {{19-19=?}}
+  let _ : T = base[keyPath: kpa] // expected-error {{key path with root type 'AnotherBase' cannot be applied to a base of type 'KeyPathBase?'}}
+
+  // Chained root mismatch
+  let _ : T = abase.member[keyPath: kp] // expected-error {{value of optional type 'KeyPathBase?' must be unwrapped to a value of type 'KeyPathBase'}}
+  // expected-note@-1 {{force-unwrap using '!' to abort execution if the optional value contains 'nil'}} {{27-27=!}}
+  // expected-note@-2 {{use '?' to access key path subscript only for non-'nil' base values}} {{27-27=?}}
+  let _ : T = abase.member[keyPath: kpa] // expected-error {{key path with root type 'AnotherBase' cannot be applied to a base of type 'KeyPathBase?'}}
+
+  let _ : T = subBase[keyPath: kp] // expected-error {{value of optional type 'KeyPathBaseSubtype?' must be unwrapped to a value of type 'KeyPathBaseSubtype'}}
+  // expected-note@-1 {{force-unwrap using '!' to abort execution if the optional value contains 'nil'}} {{22-22=!}}
+  // expected-note@-2 {{use '?' to access key path subscript only for non-'nil' base values}} {{22-22=?}}
+  let _ : T = subBase[keyPath: kpa] // expected-error {{key path with root type 'AnotherBase' cannot be applied to a base of type 'KeyPathBaseSubtype?'}}
+
+}

test/expr/unary/keypath/keypath.swift

@@ -452,20 +452,24 @@ func testKeyPathSubscriptExistentialBase(concreteBase: inout B,
   _ = concreteBase[keyPath: rkp]
   _ = concreteBase[keyPath: pkp]
 
-  concreteBase[keyPath: kp] = s // expected-error{{}}
-  concreteBase[keyPath: wkp] = s // expected-error{{}}
+  concreteBase[keyPath: kp] = s // expected-error {{cannot assign through subscript: 'kp' is a read-only key path}}
+  concreteBase[keyPath: wkp] = s // expected-error {{key path with root type 'P' cannot be applied to a base of type 'B'}}
   concreteBase[keyPath: rkp] = s
-  concreteBase[keyPath: pkp] = s // expected-error{{}}
+  // TODO(diagnostics): Improve this diagnostic message because concreteBase is mutable, the problem is related to assign 
+  // through PartialKeyPath.
+  concreteBase[keyPath: pkp] = s // expected-error {{cannot assign through subscript: 'concreteBase' is immutable}}
 
   _ = existentialBase[keyPath: kp]
   _ = existentialBase[keyPath: wkp]
   _ = existentialBase[keyPath: rkp]
   _ = existentialBase[keyPath: pkp]
 
-  existentialBase[keyPath: kp] = s // expected-error{{}}
+  existentialBase[keyPath: kp] = s // expected-error {{cannot assign through subscript: 'kp' is a read-only key path}}
   existentialBase[keyPath: wkp] = s
   existentialBase[keyPath: rkp] = s
-  existentialBase[keyPath: pkp] = s // expected-error{{}}
+  // TODO(diagnostics): Improve this diagnostic message because existentialBase is mutable, the problem is related to assign 
+  // through PartialKeyPath.
+  existentialBase[keyPath: pkp] = s // expected-error {{cannot assign through subscript: 'existentialBase' is immutable}}
 }
 
 struct AA {