[Diag] Make fixItReplace slightly smart

When replace something with a punctuator, we often prefer adding spaces around it.
For instance,

   func foo(): bar {}
   // fix it
   func foo() -> bar {}

In this case we want to add a space before '->', but not after that.

With this change, we can simply `fixItReplace(ColonLoc, " -> ")`.
`fixItReplace()` automatically adjust the spaces around it.

Author: rintaro

lib/AST/DiagnosticEngine.cpp

@@ -176,9 +176,35 @@ InFlightDiagnostic &InFlightDiagnostic::fixItReplace(SourceRange R,
     return fixItRemove(R);
 
   assert(IsActive && "Cannot modify an inactive diagnostic");
-  if (Engine && R.isValid())
-    Engine->getActiveDiagnostic().addFixIt(
-        Diagnostic::FixIt(toCharSourceRange(Engine->SourceMgr, R), Str));
+  if (R.isInvalid() || !Engine) return *this;
+
+  auto &SM = Engine->SourceMgr;
+  auto charRange = toCharSourceRange(Engine->SourceMgr, R);
+
+  // If we're replacing with something that wants spaces around it, do a bit of
+  // extra work so that we don't suggest extra spaces.
+  if (Str.back() == ' ') {
+    auto afterChars = SM.extractText({charRange.getEnd(), 1});
+    if (!afterChars.empty() && isspace(afterChars[0])) {
+      Str = Str.drop_back();
+    }
+  }
+  if (!Str.empty() && Str.front() == ' ') {
+    bool ShouldRemove = false;
+    auto buffer = SM.findBufferContainingLoc(charRange.getStart());
+    if (SM.getLocForBufferStart(buffer) == charRange.getStart()) {
+      ShouldRemove = true;
+    } else {
+      auto beforeChars =
+        SM.extractText({charRange.getStart().getAdvancedLoc(-1), 1});
+      ShouldRemove = !beforeChars.empty() && isspace(beforeChars[0]);
+    }
+    if (ShouldRemove) {
+      Str = Str.drop_front();
+    }
+  }
+
+  Engine->getActiveDiagnostic().addFixIt(Diagnostic::FixIt(charRange, Str));
   return *this;
 }
 

lib/Parse/ParseDecl.cpp

@@ -333,9 +333,8 @@ bool Parser::parseNewDeclAttribute(DeclAttributes &Attributes, SourceLoc AtLoc,
       if (!Tok.is(tok::equal))
         return false;
 
-      bool wantSpace = (Tok.getRange().getEnd() == peekToken().getLoc());
       diagnose(Tok.getLoc(), diag::replace_equal_with_colon_for_value)
-        .fixItReplace(Tok.getLoc(), wantSpace ? ": " : ":");
+        .fixItReplace(Tok.getLoc(), ": ");
     }
     return true;
   };
@@ -1410,11 +1409,8 @@ bool Parser::parseTypeAttribute(TypeAttributes &Attributes, bool justChecking) {
       if (Attributes.has(TAK_noescape)) {
         diagnose(Loc, diag::attr_noescape_conflicts_escaping_autoclosure);
       } else {
-        StringRef replacement = " @escaping ";
-        if (autoclosureEscapingParenRange.End.getAdvancedLoc(1) != Tok.getLoc())
-          replacement = replacement.drop_back();
         diagnose(Loc, diag::attr_autoclosure_escaping_deprecated)
-            .fixItReplace(autoclosureEscapingParenRange, replacement);
+            .fixItReplace(autoclosureEscapingParenRange, " @escaping ");
       }
       Attributes.setAttr(TAK_escaping, Loc);
     } else if (Attributes.has(TAK_noescape)) {
@@ -2837,7 +2833,7 @@ parseDeclTypeAlias(Parser::ParseDeclOptions Flags, DeclAttributes &Attributes) {
       // It is a common mistake to write "typealias A : Int" instead of = Int.
       // Recognize this and produce a fixit.
       diagnose(Tok, diag::expected_equal_in_typealias)
-          .fixItReplace(Tok.getLoc(), "=");
+          .fixItReplace(Tok.getLoc(), " = ");
       consumeToken(tok::colon);
     } else {
       consumeToken(tok::equal);

lib/Parse/ParsePattern.cpp

@@ -696,7 +696,7 @@ Parser::parseFunctionSignature(Identifier SimpleName,
     if (!consumeIf(tok::arrow, arrowLoc)) {
       // FixIt ':' to '->'.
       diagnose(Tok, diag::func_decl_expected_arrow)
-          .fixItReplace(SourceRange(Tok.getLoc()), "->");
+          .fixItReplace(Tok.getLoc(), " -> ");
       arrowLoc = consumeToken(tok::colon);
     }
 
@@ -787,12 +787,10 @@ ParserResult<Pattern> Parser::parseTypedPattern() {
 
   // Now parse an optional type annotation.
   if (Tok.is(tok::colon)) {
-    SourceLoc pastEndOfPrevLoc = getEndOfPreviousLoc();
     SourceLoc colonLoc = consumeToken(tok::colon);
-    SourceLoc startOfNextLoc = Tok.getLoc();
 
     if (result.isNull())  // Recover by creating AnyPattern.
-      result = makeParserErrorResult(new (Context) AnyPattern(PreviousLoc));
+      result = makeParserErrorResult(new (Context) AnyPattern(colonLoc));
 
     ParserResult<TypeRepr> Ty = parseType();
     if (Ty.hasCodeCompletion())
@@ -824,19 +822,10 @@ ParserResult<Pattern> Parser::parseTypedPattern() {
         if (status.isSuccess()) {
           backtrack.cancelBacktrack();
 
-          // Suggest replacing ':' with '=' (ensuring proper whitespace).
-          
-          bool needSpaceBefore = (pastEndOfPrevLoc == colonLoc);
-          bool needSpaceAfter =
-            SourceMgr.getByteDistance(colonLoc, startOfNextLoc) <= 1;
-          
-          StringRef replacement = " = ";
-          if (!needSpaceBefore) replacement = replacement.drop_front();
-          if (!needSpaceAfter)  replacement = replacement.drop_back();
-          
+          // Suggest replacing ':' with '='
           diagnose(lParenLoc, diag::initializer_as_typed_pattern)
             .highlight({Ty.get()->getStartLoc(), rParenLoc})
-            .fixItReplace(colonLoc, replacement);
+            .fixItReplace(colonLoc, " = ");
           result.setIsParseError();
         }
       }

test/Parse/diagnose_initializer_as_typed_pattern.swift

@@ -12,6 +12,8 @@ let d : [X]()  // expected-error{{unexpected initializer in pattern; did you mea
 
 let e: X(), ee: Int  // expected-error{{unexpected initializer in pattern; did you mean to use '='?}} {{6-7= =}}
 
+let f:/*comment*/[X]()  // expected-error{{unexpected initializer in pattern; did you mean to use '='?}} {{6-7= = }}
+
 var _1 = 1, _2 = 2
 
 // paren follows the type, but it's part of a separate (valid) expression

test/Parse/typealias.swift

@@ -9,7 +9,10 @@ var fiveInts : FiveInts = ((4,2), (1,2,3))
 
 
 // <rdar://problem/13339798> QoI: poor diagnostic in malformed typealias
-typealias Foo : Int  // expected-error {{expected '=' in typealias declaration}} {{15-16==}}
+typealias Foo1 : Int  // expected-error {{expected '=' in typealias declaration}} {{16-17==}}
+typealias Foo2: Int  // expected-error {{expected '=' in typealias declaration}} {{15-16= =}}
+typealias Foo3 :Int  // expected-error {{expected '=' in typealias declaration}} {{16-17== }}
+typealias Foo4:/*comment*/Int  // expected-error {{expected '=' in typealias declaration}} {{15-16= = }}
 
 //===--- Tests for error recovery.
 

test/decl/func/functions.swift

@@ -49,6 +49,9 @@ func recover_colon_arrow_1() : Int { } // expected-error {{expected '->' after f
 func recover_colon_arrow_2() : { }     // expected-error {{expected '->' after function parameter tuple}} {{30-31=->}} expected-error {{expected type for function result}}
 func recover_colon_arrow_3 : Int { }   // expected-error {{expected '->' after function parameter tuple}} {{28-29=->}} expected-error {{expected '(' in argument list of function declaration}}
 func recover_colon_arrow_4 : { }       // expected-error {{expected '->' after function parameter tuple}} {{28-29=->}} expected-error {{expected '(' in argument list of function declaration}} expected-error {{expected type for function result}}
+func recover_colon_arrow_5():Int { }   // expected-error {{expected '->' after function parameter tuple}} {{29-30= -> }}
+func recover_colon_arrow_6(): Int { }  // expected-error {{expected '->' after function parameter tuple}} {{29-30= ->}}
+func recover_colon_arrow_7() :Int { }  // expected-error {{expected '->' after function parameter tuple}} {{30-31=-> }}
 
 //===--- Check that we recover if the function does not have a body, but the
 //===--- context requires the function to have a body.