[Diag] Make fixItReplace slightly smart

lib/AST/DiagnosticEngine.cpp

@@ -105,6 +105,26 @@ static CharSourceRange toCharSourceRange(SourceManager &SM, SourceLoc Start,
   return CharSourceRange(SM, Start, End);
 }
 
+/// \brief Extract a character at \p Loc. If \p Loc is the end of the buffer,
+/// return '\f'.
+static char extractCharAfter(SourceManager &SM, SourceLoc Loc) {
+  auto chars = SM.extractText({Loc, 1});
+  return chars.empty() ? '\f' : chars[0];
+}
+
+/// \brief Extract a character immediately before \p Loc. If \p Loc is the
+/// start of the buffer, return '\f'.
+static char extractCharBefore(SourceManager &SM, SourceLoc Loc) {
+  // We have to be careful not to go off the front of the buffer.
+  auto bufferID = SM.findBufferContainingLoc(Loc);
+  auto bufferRange = SM.getRangeForBuffer(bufferID);
+  if (bufferRange.getStart() == Loc)
+    return '\f';
+  auto chars = SM.extractText({Loc.getAdvancedLoc(-1), 1}, bufferID);
+  assert(!chars.empty() && "Couldn't extractText with valid range");
+  return chars[0];
+}
+
 InFlightDiagnostic &InFlightDiagnostic::highlight(SourceRange R) {
   assert(IsActive && "Cannot modify an inactive diagnostic");
   if (Engine && R.isValid())
@@ -147,23 +167,10 @@ InFlightDiagnostic &InFlightDiagnostic::fixItRemove(SourceRange R) {
   // space around its hole.  Specifically, check to see there is whitespace
   // before and after the end of range.  If so, nuke the space afterward to keep
   // things consistent.
-  if (SM.extractText({charRange.getEnd(), 1}) == " ") {
-    // Check before the string, we have to be careful not to go off the front of
-    // the buffer.
-    auto bufferRange =
-      SM.getRangeForBuffer(SM.findBufferContainingLoc(charRange.getStart()));
-    bool ShouldRemove = false;
-    if (bufferRange.getStart() == charRange.getStart())
-      ShouldRemove = true;
-    else {
-      auto beforeChars =
-        SM.extractText({charRange.getStart().getAdvancedLoc(-1), 1});
-      ShouldRemove = !beforeChars.empty() && isspace(beforeChars[0]);
-    }
-    if (ShouldRemove) {
-      charRange = CharSourceRange(charRange.getStart(),
-                                  charRange.getByteLength()+1);
-    }
+  if (extractCharAfter(SM, charRange.getEnd()) == ' ' &&
+      isspace(extractCharBefore(SM, charRange.getStart()))) {
+    charRange = CharSourceRange(charRange.getStart(),
+                                charRange.getByteLength()+1);
   }
   Engine->getActiveDiagnostic().addFixIt(Diagnostic::FixIt(charRange, {}));
   return *this;
@@ -176,9 +183,23 @@ 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(SM, 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() == ' ') {
+    if (isspace(extractCharAfter(SM, charRange.getEnd())))
+      Str = Str.drop_back();
+  }
+  if (!Str.empty() && Str.front() == ' ') {
+    if (isspace(extractCharBefore(SM, charRange.getStart())))
+      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.