3 files changed, 46 insertions, 4 deletions

diff --git a/llvm/include/llvm/Transforms/Utils/MemorySSA.h b/llvm/include/llvm/Transforms/Utils/MemorySSA.h
index 48050000970..091dfc128a3 100644
--- a/llvm/include/llvm/Transforms/Utils/MemorySSA.h
+++ b/llvm/include/llvm/Transforms/Utils/MemorySSA.h
@@ -552,6 +552,8 @@ public:
   /// will be placed. The caller is expected to keep ordering the same as
   /// instructions.
   /// It will return the new MemoryAccess.
+  /// Note: If a MemoryAccess already exists for I, this function will make it
+  /// inaccessible and it *must* have removeMemoryAccess called on it.
   MemoryAccess *createMemoryAccessInBB(Instruction *I, MemoryAccess *Definition,
                                        const BasicBlock *BB,
                                        InsertionPlace Point);
@@ -563,6 +565,8 @@ public:
   /// used to replace an existing memory instruction. It will *not* create PHI
   /// nodes, or verify the clobbering definition.  The clobbering definition
   /// must be non-null.
+  /// Note: If a MemoryAccess already exists for I, this function will make it
+  /// inaccessible and it *must* have removeMemoryAccess called on it.
   MemoryAccess *createMemoryAccessBefore(Instruction *I,
                                          MemoryAccess *Definition,
                                          MemoryAccess *InsertPt);
diff --git a/llvm/lib/Transforms/Utils/MemorySSA.cpp b/llvm/lib/Transforms/Utils/MemorySSA.cpp
index 13ddd835c4d..3eda6a87ecd 100644
--- a/llvm/lib/Transforms/Utils/MemorySSA.cpp
+++ b/llvm/lib/Transforms/Utils/MemorySSA.cpp
@@ -1147,7 +1147,7 @@ void MemorySSA::buildMemorySSA() {
     // Insert phi node
     AccessList *Accesses = getOrCreateAccessList(BB);
     MemoryPhi *Phi = new MemoryPhi(BB->getContext(), BB, NextID++);
-    ValueToMemoryAccess.insert(std::make_pair(BB, Phi));
+    ValueToMemoryAccess[BB] = Phi;
     // Phi's always are placed at the front of the block.
     Accesses->push_front(Phi);
   }
@@ -1204,7 +1204,7 @@ MemoryPhi *MemorySSA::createMemoryPhi(BasicBlock *BB) {
   assert(!getMemoryAccess(BB) && "MemoryPhi already exists for this BB");
   AccessList *Accesses = getOrCreateAccessList(BB);
   MemoryPhi *Phi = new MemoryPhi(BB->getContext(), BB, NextID++);
-  ValueToMemoryAccess.insert(std::make_pair(BB, Phi));
+  ValueToMemoryAccess[BB] = Phi;
   // Phi's always are placed at the front of the block.
   Accesses->push_front(Phi);
   BlockNumberingValid.erase(BB);
@@ -1293,7 +1293,7 @@ MemoryUseOrDef *MemorySSA::createNewAccess(Instruction *I) {
     MUD = new MemoryDef(I->getContext(), nullptr, I, I->getParent(), NextID++);
   else
     MUD = new MemoryUse(I->getContext(), nullptr, I, I->getParent());
-  ValueToMemoryAccess.insert(std::make_pair(I, MUD));
+  ValueToMemoryAccess[I] = MUD;
   return MUD;
 }
 
@@ -1376,7 +1376,9 @@ void MemorySSA::removeFromLookups(MemoryAccess *MA) {
   } else {
     MemoryInst = MA->getBlock();
   }
-  ValueToMemoryAccess.erase(MemoryInst);
+  auto VMA = ValueToMemoryAccess.find(MemoryInst);
+  if (VMA->second == MA)
+    ValueToMemoryAccess.erase(VMA);
 
   auto AccessIt = PerBlockAccesses.find(MA->getBlock());
   std::unique_ptr<AccessList> &Accesses = AccessIt->second;
diff --git a/llvm/unittests/Transforms/Utils/MemorySSA.cpp b/llvm/unittests/Transforms/Utils/MemorySSA.cpp
index 6916ff536bc..93bda62f75e 100644
--- a/llvm/unittests/Transforms/Utils/MemorySSA.cpp
+++ b/llvm/unittests/Transforms/Utils/MemorySSA.cpp
@@ -106,6 +106,42 @@ TEST_F(MemorySSATest, CreateALoadAndPhi) {
   MSSA.verifyMemorySSA();
 }
 
+TEST_F(MemorySSATest, MoveAStore) {
+  // We create a diamond where there is a in the entry, a store on one side, and
+  // a load at the end.  After building MemorySSA, we test updating by moving
+  // the store from the side block to the entry block.
+  F = Function::Create(
+      FunctionType::get(B.getVoidTy(), {B.getInt8PtrTy()}, false),
+      GlobalValue::ExternalLinkage, "F", &M);
+  BasicBlock *Entry(BasicBlock::Create(C, "", F));
+  BasicBlock *Left(BasicBlock::Create(C, "", F));
+  BasicBlock *Right(BasicBlock::Create(C, "", F));
+  BasicBlock *Merge(BasicBlock::Create(C, "", F));
+  B.SetInsertPoint(Entry);
+  Argument *PointerArg = &*F->arg_begin();
+  StoreInst *EntryStore = B.CreateStore(B.getInt8(16), PointerArg);
+  B.CreateCondBr(B.getTrue(), Left, Right);
+  B.SetInsertPoint(Left);
+  StoreInst *SideStore = B.CreateStore(B.getInt8(16), PointerArg);
+  BranchInst::Create(Merge, Left);
+  BranchInst::Create(Merge, Right);
+  B.SetInsertPoint(Merge);
+  B.CreateLoad(PointerArg);
+  setupAnalyses();
+  MemorySSA &MSSA = Analyses->MSSA;
+
+  // Move the store
+  SideStore->moveBefore(Entry->getTerminator());
+  MemoryAccess *EntryStoreAccess = MSSA.getMemoryAccess(EntryStore);
+  MemoryAccess *SideStoreAccess = MSSA.getMemoryAccess(SideStore);
+  MemoryAccess *NewStoreAccess = MSSA.createMemoryAccessAfter(SideStore,
+                                                         EntryStoreAccess,
+                                                         EntryStoreAccess);
+  EntryStoreAccess->replaceAllUsesWith(NewStoreAccess);
+  MSSA.removeMemoryAccess(SideStoreAccess);
+  MSSA.verifyMemorySSA();
+}
+
 TEST_F(MemorySSATest, RemoveAPhi) {
   // We create a diamond where there is a store on one side, and then a load
   // after the merge point.  This enables us to test a bunch of different