Fix the SourceRange of an AssignExpr with an implicit source.

include/swift/AST/Expr.h

@@ -4270,7 +4270,7 @@ class AssignExpr : public Expr {
   }
   SourceLoc getEndLoc() const {
     if (!isFolded()) return EqualLoc;
-    return Src->getEndLoc();
+    return (Src->isImplicit() ? Dest->getEndLoc() : Src->getEndLoc());
   }
 
   /// True if the node has been processed by binary expression folding.

unittests/AST/CMakeLists.txt

@@ -1,5 +1,7 @@
 add_swift_unittest(SwiftASTTests
   OverrideTests.cpp
+  SourceLocTests.cpp
+  TestContext.cpp
   VersionRangeLattice.cpp
 )
 

unittests/AST/OverrideTests.cpp

@@ -10,6 +10,7 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "TestContext.h"
 #include "swift/AST/ASTContext.h"
 #include "swift/AST/DiagnosticEngine.h"
 #include "swift/AST/Module.h"
@@ -24,78 +25,7 @@
 #include "gtest/gtest.h"
 
 using namespace swift;
-
-namespace {
-/// Helper class used to set the LangOpts target before initializing the
-/// ASTContext.
-///
-/// \see TestContext
-class TestContextBase {
-public:
-  LangOptions LangOpts;
-  SearchPathOptions SearchPathOpts;
-  SourceManager SourceMgr;
-  DiagnosticEngine Diags;
-
-  TestContextBase() : Diags(SourceMgr) {
-    LangOpts.Target = llvm::Triple(llvm::sys::getProcessTriple());
-  }
-};
-
-enum ShouldDeclareOptionalTypes : bool {
-  DoNotDeclareOptionalTypes,
-  DeclareOptionalTypes
-};
-
-/// Owns an ASTContext and the associated types.
-class TestContext : public TestContextBase {
-  SourceFile *FileForLookups;
-
-  void declareOptionalType(Identifier name) {
-    auto wrapped = new (Ctx) GenericTypeParamDecl(FileForLookups,
-                                                  Ctx.getIdentifier("Wrapped"),
-                                                  SourceLoc(), /*depth*/0,
-                                                  /*index*/0);
-    auto params = GenericParamList::create(Ctx, SourceLoc(), wrapped,
-                                           SourceLoc());
-    auto decl = new (Ctx) EnumDecl(SourceLoc(), name, SourceLoc(),
-                                   /*inherited*/{}, params, FileForLookups);
-    wrapped->setDeclContext(decl);
-    FileForLookups->Decls.push_back(decl);
-  }
-
-public:
-  ASTContext Ctx;
-
-  TestContext(ShouldDeclareOptionalTypes optionals = DoNotDeclareOptionalTypes)
-      : Ctx(LangOpts, SearchPathOpts, SourceMgr, Diags) {
-    auto stdlibID = Ctx.getIdentifier(STDLIB_NAME);
-    auto *module = ModuleDecl::create(stdlibID, Ctx);
-    Ctx.LoadedModules[stdlibID] = module;
-
-    using ImplicitModuleImportKind = SourceFile::ImplicitModuleImportKind;
-    FileForLookups = new (Ctx) SourceFile(*module, SourceFileKind::Library,
-                                          /*buffer*/None,
-                                          ImplicitModuleImportKind::None);
-    module->addFile(*FileForLookups);
-
-    if (optionals == DeclareOptionalTypes) {
-      declareOptionalType(Ctx.getIdentifier("Optional"));
-      declareOptionalType(Ctx.getIdentifier("ImplicitlyUnwrappedOptional"));
-    }
-  }
-
-  template <typename Nominal>
-  Nominal *makeNominal(StringRef name,
-                       GenericParamList *genericParams = nullptr) {
-    auto result = new (Ctx) Nominal(SourceLoc(), Ctx.getIdentifier(name),
-                                    SourceLoc(), /*inherited*/{},
-                                    genericParams, FileForLookups);
-    result->setAccessibility(Accessibility::Internal);
-    return result;
-  }
-};
-} // end anonymous namespace
+using namespace swift::unittest;
 
 TEST(Override, IdenticalTypes) {
   TestContext C;

unittests/AST/SourceLocTests.cpp

@@ -0,0 +1,102 @@
+//===--- SourceLocTests.cpp - Tests for source locations of AST nodes -----===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2014 - 2016 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See http://swift.org/LICENSE.txt for license information
+// See http://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+
+#include "TestContext.h"
+#include "swift/AST/Expr.h"
+#include "gtest/gtest.h"
+
+using namespace swift;
+using namespace swift::unittest;
+
+namespace swift {
+  void PrintTo(SourceLoc loc, std::ostream *os) {
+    *os << loc.getOpaquePointerValue();
+    if (loc.isValid())
+      *os << " '" << *(char *)loc.getOpaquePointerValue() << "'";
+  }
+
+  void PrintTo(SourceRange range, std::ostream *os) {
+    PrintTo(range.Start, os);
+    *os << " - ";
+    PrintTo(range.End, os);
+  }
+} // end namespace swift
+
+TEST(SourceLoc, AssignExpr) {
+  TestContext C;
+
+  //                                                0123456789012
+  auto bufferID = C.Ctx.SourceMgr.addMemBufferCopy("aa.bb = cc.dd");
+  SourceLoc start = C.Ctx.SourceMgr.getLocForBufferStart(bufferID);
+
+  auto destBase = new (C.Ctx) UnresolvedDeclRefExpr(
+      C.Ctx.getIdentifier("aa"),
+      DeclRefKind::Ordinary,
+      DeclNameLoc(start));
+  auto dest = new (C.Ctx) UnresolvedDotExpr(
+      destBase,
+      start.getAdvancedLoc(2),
+      C.Ctx.getIdentifier("bb"),
+      DeclNameLoc(start.getAdvancedLoc(3)),
+      /*implicit*/false);
+
+  auto sourceBase = new (C.Ctx) UnresolvedDeclRefExpr(
+      C.Ctx.getIdentifier("cc"),
+      DeclRefKind::Ordinary,
+      DeclNameLoc(start.getAdvancedLoc(8)));
+  auto source = new (C.Ctx) UnresolvedDotExpr(
+      sourceBase,
+      start.getAdvancedLoc(10),
+      C.Ctx.getIdentifier("dd"),
+      DeclNameLoc(start.getAdvancedLoc(11)),
+      /*implicit*/false);
+
+  auto invalid = new (C.Ctx) UnresolvedDeclRefExpr(
+      C.Ctx.getIdentifier("invalid"),
+      DeclRefKind::Ordinary,
+      DeclNameLoc());
+
+  auto complete = new (C.Ctx) AssignExpr(dest, start.getAdvancedLoc(6), source,
+                                         /*implicit*/false);
+  EXPECT_EQ(start, complete->getStartLoc());
+  EXPECT_EQ(start.getAdvancedLoc(6), complete->getEqualLoc());
+  EXPECT_EQ(start.getAdvancedLoc(6), complete->getLoc());
+  EXPECT_EQ(start.getAdvancedLoc(11), complete->getEndLoc());
+  EXPECT_EQ(SourceRange(start, start.getAdvancedLoc(11)),
+            complete->getSourceRange());
+
+  auto invalidSource = new (C.Ctx) AssignExpr(dest, SourceLoc(), invalid,
+                                              /*implicit*/true);
+  EXPECT_EQ(start, invalidSource->getStartLoc());
+  EXPECT_EQ(SourceLoc(), invalidSource->getEqualLoc());
+  EXPECT_EQ(SourceLoc(), invalidSource->getLoc());
+  EXPECT_EQ(start.getAdvancedLoc(3), invalidSource->getEndLoc());
+  EXPECT_EQ(SourceRange(start, start.getAdvancedLoc(3)),
+            invalidSource->getSourceRange());
+
+  auto invalidDest = new (C.Ctx) AssignExpr(invalid, SourceLoc(), source,
+                                            /*implicit*/true);
+  EXPECT_EQ(start.getAdvancedLoc(8), invalidDest->getStartLoc());
+  EXPECT_EQ(SourceLoc(), invalidDest->getEqualLoc());
+  EXPECT_EQ(SourceLoc(), invalidDest->getLoc());
+  EXPECT_EQ(start.getAdvancedLoc(11), invalidDest->getEndLoc());
+  EXPECT_EQ(SourceRange(start.getAdvancedLoc(8), start.getAdvancedLoc(11)),
+            invalidDest->getSourceRange());
+
+  auto invalidAll = new (C.Ctx) AssignExpr(invalid, SourceLoc(), invalid,
+                                           /*implicit*/true);
+  EXPECT_EQ(SourceLoc(), invalidAll->getStartLoc());
+  EXPECT_EQ(SourceLoc(), invalidAll->getEqualLoc());
+  EXPECT_EQ(SourceLoc(), invalidAll->getLoc());
+  EXPECT_EQ(SourceLoc(), invalidAll->getEndLoc());
+  EXPECT_EQ(SourceRange(), invalidAll->getSourceRange());
+}

unittests/AST/TestContext.cpp

@@ -0,0 +1,52 @@
+//===--- TestContext.cpp - Helper for setting up ASTContexts --------------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2014 - 2016 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See http://swift.org/LICENSE.txt for license information
+// See http://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+
+#include "TestContext.h"
+#include "swift/AST/Module.h"
+#include "swift/Strings.h"
+
+using namespace swift;
+using namespace swift::unittest;
+
+
+static void declareOptionalType(ASTContext &ctx, SourceFile *fileForLookups,
+                                Identifier name) {
+  auto wrapped = new (ctx) GenericTypeParamDecl(fileForLookups,
+                                                ctx.getIdentifier("Wrapped"),
+                                                SourceLoc(), /*depth*/0,
+                                                /*index*/0);
+  auto params = GenericParamList::create(ctx, SourceLoc(), wrapped,
+                                         SourceLoc());
+  auto decl = new (ctx) EnumDecl(SourceLoc(), name, SourceLoc(),
+                                 /*inherited*/{}, params, fileForLookups);
+  wrapped->setDeclContext(decl);
+  fileForLookups->Decls.push_back(decl);
+}
+
+TestContext::TestContext(ShouldDeclareOptionalTypes optionals)
+    : Ctx(LangOpts, SearchPathOpts, SourceMgr, Diags) {
+  auto stdlibID = Ctx.getIdentifier(STDLIB_NAME);
+  auto *module = ModuleDecl::create(stdlibID, Ctx);
+  Ctx.LoadedModules[stdlibID] = module;
+
+  using ImplicitModuleImportKind = SourceFile::ImplicitModuleImportKind;
+  FileForLookups = new (Ctx) SourceFile(*module, SourceFileKind::Library,
+                                        /*buffer*/None,
+                                        ImplicitModuleImportKind::None);
+  module->addFile(*FileForLookups);
+
+  if (optionals == DeclareOptionalTypes) {
+    declareOptionalType(Ctx, FileForLookups, Ctx.getIdentifier("Optional"));
+    declareOptionalType(Ctx, FileForLookups,
+                        Ctx.getIdentifier("ImplicitlyUnwrappedOptional"));
+  }
+}

unittests/AST/TestContext.h

@@ -0,0 +1,63 @@
+//===--- TestContext.h - Helper for setting up ASTContexts ------*- C++ -*-===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2014 - 2016 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See http://swift.org/LICENSE.txt for license information
+// See http://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+
+#include "swift/AST/ASTContext.h"
+#include "swift/AST/DiagnosticEngine.h"
+#include "swift/Basic/LangOptions.h"
+#include "swift/Basic/SourceManager.h"
+
+namespace swift {
+namespace unittest {
+
+/// Helper class used to set the LangOpts target before initializing the
+/// ASTContext.
+///
+/// \see TestContext
+class TestContextBase {
+public:
+  LangOptions LangOpts;
+  SearchPathOptions SearchPathOpts;
+  SourceManager SourceMgr;
+  DiagnosticEngine Diags;
+
+  TestContextBase() : Diags(SourceMgr) {
+    LangOpts.Target = llvm::Triple(llvm::sys::getProcessTriple());
+  }
+};
+
+enum ShouldDeclareOptionalTypes : bool {
+  DoNotDeclareOptionalTypes,
+  DeclareOptionalTypes
+};
+
+/// Owns an ASTContext and the associated types.
+class TestContext : public TestContextBase {
+  SourceFile *FileForLookups;
+
+public:
+  ASTContext Ctx;
+
+  TestContext(ShouldDeclareOptionalTypes optionals = DoNotDeclareOptionalTypes);
+
+  template <typename Nominal>
+  Nominal *makeNominal(StringRef name,
+                       GenericParamList *genericParams = nullptr) {
+    auto result = new (Ctx) Nominal(SourceLoc(), Ctx.getIdentifier(name),
+                                    SourceLoc(), /*inherited*/{},
+                                    genericParams, FileForLookups);
+    result->setAccessibility(Accessibility::Internal);
+    return result;
+  }
+};
+
+} // end namespace unittest
+} // end namespace swift