commit c466672d1e616b243cd0fa1791df56a9dc8c87df
author James Molloy <jmolloy@google.com> Tue Jul 24 09:48:31 2018 -0700
committer jpienaar <jpienaar@google.com> Fri Mar 29 12:40:47 2019 -0700
tree 1492395c0e95a39b65594653b87c081808bfc83d
parent 9361fb329d7bf9d050b84d0e5bdf592e325b85f4

[mlir] Fix a use-after-free iterator error found by asan

While fixing this the parser-affine-map.mlir test started failing due to ordering of the printed affine maps. Even the existing CHECK-DAGs weren't enough to disambiguate; a partial match on one line precluded a total match on a following line.

The fix for this was easy - print the affine maps in reference order rather than in DenseMap iteration order.

PiperOrigin-RevId: 205843770

lib/IR/AsmPrinter.cpp

diff --git a/lib/IR/AsmPrinter.cpp b/lib/IR/AsmPrinter.cpp
index 55cc98c..4a21b31 100644
--- a/lib/IR/AsmPrinter.cpp
+++ b/lib/IR/AsmPrinter.cpp
@@ -63,14 +63,13 @@
     return it->second;
   }
 
-  const DenseMap<const AffineMap *, int> &getAffineMapIds() const {
-    return affineMapIds;
-  }
+  ArrayRef<const AffineMap *> getAffineMapIds() const { return affineMapsById; }
 
 private:
   void recordAffineMapReference(const AffineMap *affineMap) {
     if (affineMapIds.count(affineMap) == 0) {
-      affineMapIds[affineMap] = nextAffineMapId++;
+      affineMapIds[affineMap] = affineMapsById.size();
+      affineMapsById.push_back(affineMap);
     }
   }
 
@@ -84,7 +83,7 @@
   void visitOperation(const Operation *op);
 
   DenseMap<const AffineMap *, int> affineMapIds;
-  int nextAffineMapId = 0;
+  std::vector<const AffineMap *> affineMapsById;
 };
 } // end anonymous namespace
 
@@ -228,10 +227,10 @@
 }
 
 void ModulePrinter::print(const Module *module) {
-  for (const auto &mapAndId : state.getAffineMapIds()) {
-    printAffineMapId(mapAndId.second);
+  for (const auto &map : state.getAffineMapIds()) {
+    printAffineMapId(state.getAffineMapId(map));
     os << " = ";
-    mapAndId.first->print(os);
+    map->print(os);
     os << '\n';
   }
   for (auto *fn : module->functionList)