[Sema] Improve diagnostics for key path root type inferred as option accessing wrapped member

include/swift/AST/DiagnosticsSema.def

@@ -1082,13 +1082,22 @@ NOTE(unwrap_with_guard,none,
 ERROR(optional_base_not_unwrapped,none,
       "value of optional type %0 must be unwrapped to refer to member %1 of "
       "wrapped base type %2", (Type, DeclNameRef, Type))
+ERROR(invalid_optional_infered_keypath_root, none,
+      "key path root inferred as optional type %0 must be unwrapped to refer to member %1 "
+      "of unwrapped type %2", (Type, DeclNameRef, Type))
 NOTE(optional_base_chain,none,
      "chain the optional using '?' to access member %0 only for non-'nil' "
      "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", ())
+NOTE(optional_key_path_root_base_chain, none,
+     "chain the optional using '?.' to access unwrapped type member %0",
+     (DeclNameRef))
+NOTE(optional_key_path_root_base_unwrap, none,
+     "unwrap the optional using '!.' to access unwrapped type member %0",
+     (DeclNameRef))
 
 ERROR(missing_unwrap_optional_try,none,
       "value of optional type %0 not unwrapped; did you mean to use 'try!' "

lib/Sema/CSDiagnostics.cpp

@@ -1071,9 +1071,6 @@ bool MemberAccessOnOptionalBaseFailure::diagnoseAsError() {
     return false;
 
   auto sourceRange = getSourceRange();
-
-  emitDiagnostic(diag::optional_base_not_unwrapped,
-                 baseType, Member, unwrappedBaseType);
 
   auto componentPathElt =
       locator->getLastElementAs<LocatorPathElt::KeyPathComponent>();
@@ -1082,12 +1079,25 @@ bool MemberAccessOnOptionalBaseFailure::diagnoseAsError() {
     // let's emit a tailored note suggesting to use its unwrapped type.
     auto *keyPathExpr = castToExpr<KeyPathExpr>(getAnchor());
     if (auto rootType = keyPathExpr->getRootType()) {
+      emitDiagnostic(diag::optional_base_not_unwrapped, baseType, Member,
+                     unwrappedBaseType);
+
       emitDiagnostic(diag::optional_base_remove_optional_for_keypath_root,
                      unwrappedBaseType)
           .fixItReplace(rootType->getSourceRange(),
                         unwrappedBaseType.getString());
+    } else {
+      emitDiagnostic(diag::invalid_optional_infered_keypath_root, baseType,
+                     Member, unwrappedBaseType);
+      emitDiagnostic(diag::optional_key_path_root_base_chain, Member)
+          .fixItInsert(sourceRange.End, "?.");
+      emitDiagnostic(diag::optional_key_path_root_base_unwrap, Member)
+          .fixItInsert(sourceRange.End, "!.");
     }
   } else {
+    emitDiagnostic(diag::optional_base_not_unwrapped, baseType, Member,
+                   unwrappedBaseType);
+
     // FIXME: It would be nice to immediately offer "base?.member ?? defaultValue"
     // for non-optional results where that would be appropriate. For the moment
     // always offering "?" means that if the user chooses chaining, we'll end up

localization/diagnostics/en.yaml

@@ -3136,6 +3136,9 @@
 - id: optional_base_not_unwrapped
   msg: "value of optional type %0 must be unwrapped to refer to member %1 of wrapped base type %2"
 
+- id: invalid_optional_infered_keypath_root
+  msg: "key path root inferred as optional type %0 must be unwrapped to refer to member %1 of unwrapped type %2"
+
 - id: optional_base_chain
   msg: "chain the optional using '?' to access member %0 only for non-'nil' base values"
 
@@ -3145,6 +3148,12 @@
 - id: optional_keypath_application_base
   msg: "use '?' to access key path subscript only for non-'nil' base values"
 
+- id: optional_key_path_root_base_chain
+  msg: "chain the optional using '?.' to access unwrapped type member %0"
+
+- id: optional_key_path_root_base_unwrap
+  msg: "unwrap the optional using '!.' to access unwrapped type member %0"
+
 - id: missing_unwrap_optional_try
   msg: "value of optional type %0 not unwrapped; did you mean to use 'try!' or chain with '?'?"
 

test/expr/unary/keypath/keypath.swift

@@ -1006,9 +1006,12 @@ func testMemberAccessOnOptionalKeyPathComponent() {
   // expected-error@-1 {{value of optional type '(Int, Int)?' must be unwrapped to refer to member '0' of wrapped base type '(Int, Int)'}}
   // expected-note@-2 {{use unwrapped type '(Int, Int)' as key path root}}{{4-15=(Int, Int)}}
 
-  // TODO(diagnostics) Improve diagnostics refering to key path root not able to be infered as an optional type.
-  SR5688_KP(\.count)
-  // expected-error@-1 {{value of optional type 'String?' must be unwrapped to refer to member 'count' of wrapped base type 'String'}}
+  SR5688_KP(\.count) // expected-error {{key path root inferred as optional type 'String?' must be unwrapped to refer to member 'count' of unwrapped type 'String'}}
+  // expected-note@-1 {{chain the optional using '?.' to access unwrapped type member 'count'}} {{15-15=?.}}
+  // expected-note@-2 {{unwrap the optional using '!.' to access unwrapped type member 'count'}} {{15-15=!.}}
+  let _ : KeyPath<String?, Int> = \.count // expected-error {{key path root inferred as optional type 'String?' must be unwrapped to refer to member 'count' of unwrapped type 'String'}}
+  // expected-note@-1 {{chain the optional using '?.' to access unwrapped type member 'count'}} {{37-37=?.}}
+  // expected-note@-2 {{unwrap the optional using '!.' to access unwrapped type member 'count'}} {{37-37=!.}}
 }
 
 func testSyntaxErrors() {