[sema] Follow-up on the raw-representable casting fixits, produce a fixit for optional source types as well.

For optional source types we use the form "expr.map{ T(rawValue: $0) }".

Author: akyrtzi

lib/Sema/CSDiag.cpp

@@ -3159,8 +3159,37 @@ static void tryRawRepresentableFixIts(InFlightDiagnostic &diag,
                                       KnownProtocolKind kind,
                                       Expr *expr) {
   // The following fixes apply for optional destination types as well.
+  bool toTypeIsOptional = !toType->getAnyOptionalObjectType().isNull();
   toType = toType->lookThroughAllAnyOptionalTypes();
 
+  Type fromTypeUnwrapped = fromType->getAnyOptionalObjectType();
+  bool fromTypeIsOptional = !fromTypeUnwrapped.isNull();
+  if (fromTypeIsOptional)
+    fromType = fromTypeUnwrapped;
+
+  auto fixIt = [&](StringRef convWrapBefore, StringRef convWrapAfter) {
+    SourceRange exprRange = expr->getSourceRange();
+    if (fromTypeIsOptional && toTypeIsOptional) {
+      // Use optional's map function to convert conditionally, like so:
+      //   expr.map{ T(rawValue: $0) }
+      bool needsParens = !expr->canAppendCallParentheses();
+      std::string mapCodeFix;
+      if (needsParens) {
+        diag.fixItInsert(exprRange.Start, "(");
+        mapCodeFix += ")";
+      }
+      mapCodeFix += ".map{ ";
+      mapCodeFix += convWrapBefore;
+      mapCodeFix += "$0";
+      mapCodeFix += convWrapAfter;
+      mapCodeFix += " }";
+      diag.fixItInsertAfter(exprRange.End, mapCodeFix);
+    } else if (!fromTypeIsOptional) {
+      diag.fixItInsert(exprRange.Start, convWrapBefore);
+      diag.fixItInsertAfter(exprRange.End, convWrapAfter);
+    }
+  };
+
   if (isLiteralConvertibleType(fromType, kind, CS)) {
     if (auto rawTy = isRawRepresentable(toType, kind, CS)) {
       // Produce before/after strings like 'Result(rawValue: RawType(<expr>))'
@@ -3173,9 +3202,7 @@ static void tryRawRepresentableFixIts(InFlightDiagnostic &diag,
         convWrapBefore += "(";
         convWrapAfter += ")";
       }
-      SourceRange exprRange = expr->getSourceRange();
-      diag.fixItInsert(exprRange.Start, convWrapBefore);
-      diag.fixItInsertAfter(exprRange.End, convWrapAfter);
+      fixIt(convWrapBefore, convWrapAfter);
       return;
     }
   }
@@ -3189,9 +3216,7 @@ static void tryRawRepresentableFixIts(InFlightDiagnostic &diag,
         convWrapBefore += "(";
         convWrapAfter += ")";
       }
-      SourceRange exprRange = expr->getSourceRange();
-      diag.fixItInsert(exprRange.Start, convWrapBefore);
-      diag.fixItInsertAfter(exprRange.End, convWrapAfter);
+      fixIt(convWrapBefore, convWrapAfter);
       return;
     }
   }

test/FixCode/fixits-apply.swift

@@ -56,6 +56,27 @@ func testMask5(a: Int) {
 func testMask6(a: Int) {
   sendItOpt(a)
 }
+func testMask7(a: Int?) {
+  sendItOpt(a)
+}
+func testMask8(a: UInt64?) {
+  sendItOpt(a)
+}
+func testMask9(a: Any) {
+  sendItOpt(a as? Int)
+}
+func testMask10(a: Int?) {
+  sendIt(a) // no fix, nullability mismatch.
+}
+func testMask11(a: MyEventMask2?) {
+  testMask7(a: a)
+}
+func testMask12(a: MyEventMask2?) {
+  testMask8(a: a)
+}
+func testMask13(a: MyEventMask2?) {
+  testMask1(a: a) // no fix, nullability mismatch.
+}
 
 enum MyEnumType : UInt32 {
   case invalid

test/FixCode/fixits-apply.swift.result

@@ -56,6 +56,27 @@ func testMask5(a: Int) {
 func testMask6(a: Int) {
   sendItOpt(MyEventMask2(rawValue: UInt64(a)))
 }
+func testMask7(a: Int?) {
+  sendItOpt(a.map{ MyEventMask2(rawValue: UInt64($0)) })
+}
+func testMask8(a: UInt64?) {
+  sendItOpt(a.map{ MyEventMask2(rawValue: $0) })
+}
+func testMask9(a: Any) {
+  sendItOpt((a as? Int).map{ MyEventMask2(rawValue: UInt64($0)) })
+}
+func testMask10(a: Int?) {
+  sendIt(a) // no fix, nullability mismatch.
+}
+func testMask11(a: MyEventMask2?) {
+  testMask7(a: a.map{ UInt64($0.rawValue) })
+}
+func testMask12(a: MyEventMask2?) {
+  testMask8(a: a.map{ $0.rawValue })
+}
+func testMask13(a: MyEventMask2?) {
+  testMask1(a: a) // no fix, nullability mismatch.
+}
 
 enum MyEnumType : UInt32 {
   case invalid