[Edit] Fix issue with tracking what macro argument inputs have been edited.

This was not working correctly, leading to erroneously rejecting valid edits.

llvm-svn: 247462

4 files changed, 84 insertions, 16 deletions

diff --git a/clang/include/clang/Edit/EditedSource.h b/clang/include/clang/Edit/EditedSource.h
index 8bc27e73bee..b6ec8b8f06e 100644
--- a/clang/include/clang/Edit/EditedSource.h
+++ b/clang/include/clang/Edit/EditedSource.h
@@ -10,9 +10,11 @@
 #ifndef LLVM_CLANG_EDIT_EDITEDSOURCE_H
 #define LLVM_CLANG_EDIT_EDITEDSOURCE_H
 
+#include "clang/Basic/IdentifierTable.h"
 #include "clang/Edit/FileOffset.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/StringRef.h"
+#include "llvm/ADT/TinyPtrVector.h"
 #include "llvm/Support/Allocator.h"
 #include <map>
 
@@ -39,14 +41,18 @@ class EditedSource {
   typedef std::map<FileOffset, FileEdit> FileEditsTy;
   FileEditsTy FileEdits;
 
-  llvm::DenseMap<unsigned, SourceLocation> ExpansionToArgMap;
+  llvm::DenseMap<unsigned, llvm::TinyPtrVector<IdentifierInfo*>>
+    ExpansionToArgMap;
+  SmallVector<std::pair<SourceLocation, IdentifierInfo*>, 2>
+    CurrCommitMacroArgExps;
 
+  IdentifierTable IdentTable;
   llvm::BumpPtrAllocator StrAlloc;
 
 public:
   EditedSource(const SourceManager &SM, const LangOptions &LangOpts,
                const PPConditionalDirectiveRecord *PPRec = nullptr)
-    : SourceMgr(SM), LangOpts(LangOpts), PPRec(PPRec),
+    : SourceMgr(SM), LangOpts(LangOpts), PPRec(PPRec), IdentTable(LangOpts),
       StrAlloc() { }
 
   const SourceManager &getSourceManager() const { return SourceMgr; }
@@ -76,6 +82,12 @@ private:
   StringRef getSourceText(FileOffset BeginOffs, FileOffset EndOffs,
                           bool &Invalid);
   FileEditsTy::iterator getActionForOffset(FileOffset Offs);
+  void deconstructMacroArgLoc(SourceLocation Loc,
+                              SourceLocation &ExpansionLoc,
+                              IdentifierInfo *&II);
+
+  void startingCommit();
+  void finishedCommit();
 };
 
 }
diff --git a/clang/lib/Edit/EditedSource.cpp b/clang/lib/Edit/EditedSource.cpp
index e557de92410..54659f6e4ed 100644
--- a/clang/lib/Edit/EditedSource.cpp
+++ b/clang/lib/Edit/EditedSource.cpp
@@ -23,6 +23,36 @@ void EditsReceiver::remove(CharSourceRange range) {
   replace(range, StringRef());
 }
 
+void EditedSource::deconstructMacroArgLoc(SourceLocation Loc,
+                                          SourceLocation &ExpansionLoc,
+                                          IdentifierInfo *&II) {
+  assert(SourceMgr.isMacroArgExpansion(Loc));
+  SourceLocation DefArgLoc = SourceMgr.getImmediateExpansionRange(Loc).first;
+  ExpansionLoc = SourceMgr.getImmediateExpansionRange(DefArgLoc).first;
+  SmallString<20> Buf;
+  StringRef ArgName = Lexer::getSpelling(SourceMgr.getSpellingLoc(DefArgLoc),
+                                         Buf, SourceMgr, LangOpts);
+  II = nullptr;
+  if (!ArgName.empty()) {
+    II = &IdentTable.get(ArgName);
+  }
+}
+
+void EditedSource::startingCommit() {}
+
+void EditedSource::finishedCommit() {
+  for (auto &ExpArg : CurrCommitMacroArgExps) {
+    SourceLocation ExpLoc;
+    IdentifierInfo *II;
+    std::tie(ExpLoc, II) = ExpArg;
+    auto &ArgNames = ExpansionToArgMap[ExpLoc.getRawEncoding()];
+    if (std::find(ArgNames.begin(), ArgNames.end(), II) == ArgNames.end()) {
+      ArgNames.push_back(II);
+    }
+  }
+  CurrCommitMacroArgExps.clear();
+}
+
 StringRef EditedSource::copyString(const Twine &twine) {
   SmallString<128> Data;
   return copyString(twine.toStringRef(Data));
@@ -36,15 +66,27 @@ bool EditedSource::canInsertInOffset(SourceLocation OrigLoc, FileOffset Offs) {
   }
 
   if (SourceMgr.isMacroArgExpansion(OrigLoc)) {
-    SourceLocation
-      DefArgLoc = SourceMgr.getImmediateExpansionRange(OrigLoc).first;
-    SourceLocation
-      ExpLoc = SourceMgr.getImmediateExpansionRange(DefArgLoc).first;
-    llvm::DenseMap<unsigned, SourceLocation>::iterator
-      I = ExpansionToArgMap.find(ExpLoc.getRawEncoding());
-    if (I != ExpansionToArgMap.end() && I->second != DefArgLoc)
-      return false; // Trying to write in a macro argument input that has
-                 // already been written for another argument of the same macro. 
+    IdentifierInfo *II;
+    SourceLocation ExpLoc;
+    deconstructMacroArgLoc(OrigLoc, ExpLoc, II);
+    auto I = ExpansionToArgMap.find(ExpLoc.getRawEncoding());
+    if (I != ExpansionToArgMap.end() &&
+        std::find(I->second.begin(), I->second.end(), II) != I->second.end()) {
+      // Trying to write in a macro argument input that has already been
+      // written by a previous commit for another expansion of the same macro
+      // argument name. For example:
+      //
+      // \code
+      //   #define MAC(x) ((x)+(x))
+      //   MAC(a)
+      // \endcode
+      //
+      // A commit modified the macro argument 'a' due to the first '(x)'
+      // expansion inside the macro definition, and a subsequent commit tried
+      // to modify 'a' again for the second '(x)' expansion. The edits of the
+      // second commit will be rejected.
+      return false;
+    }
   }
 
   return true;
@@ -59,11 +101,11 @@ bool EditedSource::commitInsert(SourceLocation OrigLoc,
     return true;
 
   if (SourceMgr.isMacroArgExpansion(OrigLoc)) {
-    SourceLocation
-      DefArgLoc = SourceMgr.getImmediateExpansionRange(OrigLoc).first;
-    SourceLocation
-      ExpLoc = SourceMgr.getImmediateExpansionRange(DefArgLoc).first;
-    ExpansionToArgMap[ExpLoc.getRawEncoding()] = DefArgLoc;
+    IdentifierInfo *II;
+    SourceLocation ExpLoc;
+    deconstructMacroArgLoc(OrigLoc, ExpLoc, II);
+    if (II)
+      CurrCommitMacroArgExps.emplace_back(ExpLoc, II);
   }
   
   FileEdit &FA = FileEdits[Offs];
@@ -219,6 +261,16 @@ bool EditedSource::commit(const Commit &commit) {
   if (!commit.isCommitable())
     return false;
 
+  struct CommitRAII {
+    EditedSource &Editor;
+    CommitRAII(EditedSource &Editor) : Editor(Editor) {
+      Editor.startingCommit();
+    }
+    ~CommitRAII() {
+      Editor.finishedCommit();
+    }
+  } CommitRAII(*this);
+
   for (edit::Commit::edit_iterator
          I = commit.edit_begin(), E = commit.edit_end(); I != E; ++I) {
     const edit::Commit::Edit &edit = *I;
diff --git a/clang/test/ARCMT/objcmt-subscripting-literals.m b/clang/test/ARCMT/objcmt-subscripting-literals.m
index 014c1092999..0974c3b8bb0 100644
--- a/clang/test/ARCMT/objcmt-subscripting-literals.m
+++ b/clang/test/ARCMT/objcmt-subscripting-literals.m
@@ -88,6 +88,7 @@ typedef const struct __CFString * CFStringRef;
 #define M(x) (x)
 #define PAIR(x) @#x, [NSNumber numberWithInt:(x)]
 #define TWO(x) ((x), (x))
+#define TWO_SEP(x,y) ((x), (y))
 
 @interface I {
   NSArray *ivarArr;
@@ -118,6 +119,7 @@ typedef const struct __CFString * CFStringRef;
   id o = [arr objectAtIndex:2];
   o = [dict objectForKey:@"key"];
   o = TWO([dict objectForKey:@"key"]);
+  o = TWO_SEP([dict objectForKey:@"key"], [arr objectAtIndex:2]);
   o = [NSDictionary dictionaryWithObject:[NSDictionary dictionary] forKey:@"key"];
   NSMutableArray *marr = 0;
   NSMutableDictionary *mdict = 0;
diff --git a/clang/test/ARCMT/objcmt-subscripting-literals.m.result b/clang/test/ARCMT/objcmt-subscripting-literals.m.result
index e9ff8df34dc..ed7879bb139 100644
--- a/clang/test/ARCMT/objcmt-subscripting-literals.m.result
+++ b/clang/test/ARCMT/objcmt-subscripting-literals.m.result
@@ -88,6 +88,7 @@ typedef const struct __CFString * CFStringRef;
 #define M(x) (x)
 #define PAIR(x) @#x, [NSNumber numberWithInt:(x)]
 #define TWO(x) ((x), (x))
+#define TWO_SEP(x,y) ((x), (y))
 
 @interface I {
   NSArray *ivarArr;
@@ -118,6 +119,7 @@ typedef const struct __CFString * CFStringRef;
   id o = arr[2];
   o = dict[@"key"];
   o = TWO(dict[@"key"]);
+  o = TWO_SEP(dict[@"key"], arr[2]);
   o = @{@"key": @{}};
   NSMutableArray *marr = 0;
   NSMutableDictionary *mdict = 0;