[ConstraintSystem] Increase score only if members found on Optional…

include/swift/Sema/OverloadChoice.h

@@ -73,9 +73,14 @@ class OverloadChoice {
     /// optional context type, turning a "Decl" kind into
     /// "DeclViaUnwrappedOptional".
     IsDeclViaUnwrappedOptional = 0x02,
+    /// Indicates that there are viable members found on `Optional`
+    /// type and its underlying type. And current overload choice
+    /// is a backup one, which should be picked only if members
+    /// found directly on `Optional` do not match.
+    IsFallbackDeclViaUnwrappedOptional = 0x03,
     /// Indicates that this declaration was dynamic, turning a
     /// "Decl" kind into "DeclViaDynamic" kind.
-    IsDeclViaDynamic = 0x03,
+    IsDeclViaDynamic = 0x07,
   };
 
   /// The base type to be used when referencing the declaration
@@ -175,17 +180,23 @@ class OverloadChoice {
 
   /// Retrieve an overload choice for a declaration that was found
   /// by unwrapping an optional context type.
+  ///
+  /// \param isFallback Indicates that this result should be used
+  /// as a backup, if member found directly on `Optional` doesn't
+  /// match.
   static OverloadChoice
-  getDeclViaUnwrappedOptional(Type base, ValueDecl *value,
+  getDeclViaUnwrappedOptional(Type base, ValueDecl *value, bool isFallback,
                               FunctionRefKind functionRefKind) {
     OverloadChoice result;
     result.BaseAndDeclKind.setPointer(base);
-    result.BaseAndDeclKind.setInt(IsDeclViaUnwrappedOptional);
+    result.BaseAndDeclKind.setInt(isFallback
+                                      ? IsFallbackDeclViaUnwrappedOptional
+                                      : IsDeclViaUnwrappedOptional);
     result.DeclOrKind = value;
     result.TheFunctionRefKind = functionRefKind;
     return result;
   }
-  
+
   /// Retrieve an overload choice for a declaration that was found via
   /// dynamic member lookup. The `ValueDecl` is a `subscript(dynamicMember:)`
   /// method.
@@ -219,6 +230,7 @@ class OverloadChoice {
       case IsDeclViaBridge: return OverloadChoiceKind::DeclViaBridge;
       case IsDeclViaDynamic: return OverloadChoiceKind::DeclViaDynamic;
       case IsDeclViaUnwrappedOptional:
+      case IsFallbackDeclViaUnwrappedOptional:
         return OverloadChoiceKind::DeclViaUnwrappedOptional;
       default: return OverloadChoiceKind::Decl;
       }
@@ -256,6 +268,12 @@ class OverloadChoice {
     return getKind() == OverloadChoiceKind::KeyPathDynamicMemberLookup;
   }
 
+  /// Determine whether this member is a backup in case
+  /// members found directly on `Optional` didn't match.
+  bool isFallbackMemberOnUnwrappedBase() const {
+    return BaseAndDeclKind.getInt() == IsFallbackDeclViaUnwrappedOptional;
+  }
+
   /// Get the name of the overload choice.
   DeclName getName() const;
 

lib/Sema/CSSimplify.cpp

@@ -6672,8 +6672,9 @@ performMemberLookup(ConstraintKind constraintKind, DeclNameRef memberName,
 
   // Local function that turns a ValueDecl into a properly configured
   // OverloadChoice.
-  auto getOverloadChoice = [&](ValueDecl *cand, bool isBridged,
-                               bool isUnwrappedOptional) -> OverloadChoice {
+  auto getOverloadChoice =
+      [&](ValueDecl *cand, bool isBridged, bool isUnwrappedOptional,
+          bool isFallbackUnwrap = false) -> OverloadChoice {
     // If we're looking into an existential type, check whether this
     // result was found via dynamic lookup.
     if (instanceTy->isAnyObject()) {
@@ -6694,8 +6695,9 @@ performMemberLookup(ConstraintKind constraintKind, DeclNameRef memberName,
       auto ovlBaseTy = MetatypeType::get(baseTy->castTo<MetatypeType>()
                                              ->getInstanceType()
                                              ->getOptionalObjectType());
-      return OverloadChoice::getDeclViaUnwrappedOptional(ovlBaseTy, cand,
-                                                         functionRefKind);
+      return OverloadChoice::getDeclViaUnwrappedOptional(
+          ovlBaseTy, cand,
+          /*isFallback=*/isFallbackUnwrap, functionRefKind);
     }
 
     // While looking for subscript choices it's possible to find
@@ -6719,7 +6721,7 @@ performMemberLookup(ConstraintKind constraintKind, DeclNameRef memberName,
 
     return OverloadChoice(baseTy, cand, functionRefKind);
   };
-  
+
   // Add all results from this lookup.
   for (auto result : lookup)
     addChoice(getOverloadChoice(result.getValueDecl(),
@@ -6802,11 +6804,16 @@ performMemberLookup(ConstraintKind constraintKind, DeclNameRef memberName,
       }
 
       if (objectType->mayHaveMembers()) {
+        // If there are viable members directly on `Optional`, let's
+        // prioritize them and mark any results found on wrapped type
+        // as a fallback results.
+        bool isFallback = !result.ViableCandidates.empty();
         LookupResult &optionalLookup = lookupMember(objectType, memberName);
         for (auto result : optionalLookup)
           addChoice(getOverloadChoice(result.getValueDecl(),
-                                      /*bridged*/false,
-                                      /*isUnwrappedOptional=*/true));
+                                      /*bridged*/ false,
+                                      /*isUnwrappedOptional=*/true,
+                                      /*isUnwrapFallback=*/isFallback));
       }
     }
   }

lib/Sema/ConstraintSystem.cpp

@@ -2799,8 +2799,7 @@ void ConstraintSystem::resolveOverload(ConstraintLocator *locator,
     increaseScore(SK_DisfavoredOverload);
   }
 
-  if (choice.getKind() == OverloadChoiceKind::DeclViaUnwrappedOptional &&
-      locator->isLastElement<LocatorPathElt::UnresolvedMember>()) {
+  if (choice.isFallbackMemberOnUnwrappedBase()) {
     increaseScore(SK_UnresolvedMemberViaOptional);
   }
 }

test/Constraints/sr13815.swift

@@ -0,0 +1,25 @@
+// RUN: %target-typecheck-verify-swift
+
+// SR-13815
+
+enum E {
+case foo(String)
+}
+
+struct Test {
+  var bar: E?
+}
+
+struct S {
+  func evaluate(_: Test) -> [Test] {
+    return []
+  }
+
+  func test(set: Set<String>)  {
+    let flattened = set.flatMap { element in
+      evaluate(Test(bar: .foo(element)))
+    }
+
+    let _: [Test] = flattened // Ok (find .`bar` after unwrap)
+  }
+}