[SR-16058][Sema] Consider wrapping type variable layers of optionality when warning about checked casts

Author: LucianoPAlmeida

lib/Sema/CSSimplify.cpp

@@ -7467,13 +7467,20 @@ static ConstraintFix *maybeWarnAboutExtraneousCast(
   // we need to store the difference as a signed integer.
   int extraOptionals = fromOptionals.size() - toOptionals.size();
 
-  // "from" expression could be a type variable with value-to-optional
-  // restrictions that we have to account for optionality mismatch.
+  // "from" expression could be a type variable wrapped in an optional e.g.
+  // Optional<$T0>. So when that is the case we have to add this additional
+  // optionality levels to from type.
   const auto subExprType = cs.getType(castExpr->getSubExpr());
-  if (cs.hasConversionRestriction(fromType, subExprType,
-                                  ConversionRestrictionKind::ValueToOptional)) {
-    extraOptionals++;
-    origFromType = OptionalType::get(origFromType);
+  if (subExprType->getOptionalObjectType()) {
+    SmallVector<Type, 4> subExprOptionals;
+    const auto unwrappedSubExprType =
+        subExprType->lookThroughAllOptionalTypes(subExprOptionals);
+    if (unwrappedSubExprType->is<TypeVariableType>()) {
+      extraOptionals += subExprOptionals.size();
+      for (size_t i = 0; i != subExprOptionals.size(); ++i) {
+        origFromType = OptionalType::get(origFromType);
+      }
+    }
   }
 
   // Removing the optionality from to type when the force cast expr is an IUO.

test/Constraints/casts.swift

@@ -680,3 +680,16 @@ func test_SR_15562() {
     // expected-error@-1{{left side of mutating operator has immutable type 'Int'}}
   }
 }
+
+// SR-16058
+extension Dictionary {
+  func SR16058(_: Key) -> Value?? { nil }
+}
+func SR_16058_tests() { 
+  let dict: [Int: String?] = [:]
+  let foo: Int? = 1
+  let _: String? = foo.flatMap { dict[$0] } as? String  // OK
+
+  // More than one optionality wrapping
+  let _: String? = foo.flatMap { dict.SR16058(_: $0) } as? String // OK
+}