[Polly][SI] Fix a potential use-after-free

[philip.pfaffe]

Summary

There is a potential use-after-free bug in Scop::buildSchedule(Region *,
LoopStackTy &, LoopInfo &). Before, we took a reference to LoopStack.back()
which is a use after free, since back is popped off further below. This didn't
crash before by pure chance, since LoopStack is actually a vector, and the
memory isn't freed upon pop. I turned this into an iterator-based algorithm.

[Meinersbur]

LGMT, thanks.

lib/Analysis/ScopInfo.cpp

4843 ↗

(On Diff #122721)

Could you consider adding an assertion LoopData != LoopStack.rend() as well?

4861 ↗

(On Diff #122721)

[Nit] LoopStack.size() returns size_t, no "almost-always-auto" as of LLVM coding standards.

4880 ↗

(On Diff #122721)

Could you consider adding a comment on why the stack has to be popped only at the end?

[philip.pfaffe]

Address Review Comments

polly/trunk/lib/Analysis/ScopInfo.cpp

void Scop::buildSchedule(RegionNode *RN, LoopStackTy &LoopStack, LoopInfo &LI) {
   if (RN->isSubRegion()) {
     auto *LocalRegion = RN->getNodeAs<Region>();
     if (!isNonAffineSubRegion(LocalRegion)) {
       buildSchedule(LocalRegion, LoopStack, LI);
       return;
     }
   }
 
-  auto &LoopData = LoopStack.back();
-  LoopData.NumBlocksProcessed += getNumBlocksInRegionNode(RN);
+  assert(LoopStack.rbegin() != LoopStack.rend());
+  auto LoopData = LoopStack.rbegin();
+  LoopData->NumBlocksProcessed += getNumBlocksInRegionNode(RN);
 
   for (auto *Stmt : getStmtListFor(RN)) {
     auto *UDomain = isl_union_set_from_set(Stmt->getDomain().release());
     auto *StmtSchedule = isl_schedule_from_domain(UDomain);
-    LoopData.Schedule = combineInSequence(LoopData.Schedule, StmtSchedule);
+    LoopData->Schedule = combineInSequence(LoopData->Schedule, StmtSchedule);
   }
 
   // Check if we just processed the last node in this loop. If we did, finalize
   // the loop by:
   //
   //   - adding new schedule dimensions
   //   - folding the resulting schedule into the parent loop schedule
   //   - dropping the loop schedule from the LoopStack.
   //
   // Then continue to check surrounding loops, which might also have been
   // completed by this node.
-  while (LoopData.L &&
-         LoopData.NumBlocksProcessed == getNumBlocksInLoop(LoopData.L)) {
-    auto *Schedule = LoopData.Schedule;
-    auto NumBlocksProcessed = LoopData.NumBlocksProcessed;
-
-    LoopStack.pop_back();
-    auto &NextLoopData = LoopStack.back();
+  size_t Dimension = LoopStack.size();
+  while (LoopData->L &&
+         LoopData->NumBlocksProcessed == getNumBlocksInLoop(LoopData->L)) {
+    auto *Schedule = LoopData->Schedule;
+    auto NumBlocksProcessed = LoopData->NumBlocksProcessed;
+
+    assert(std::next(LoopData) != LoopStack.rend());
+    ++LoopData;
+    --Dimension;
 
     if (Schedule) {
       isl::union_set Domain = give(isl_schedule_get_domain(Schedule));
-      isl::multi_union_pw_aff MUPA = mapToDimension(Domain, LoopStack.size());
+      isl::multi_union_pw_aff MUPA = mapToDimension(Domain, Dimension);
       Schedule = isl_schedule_insert_partial_schedule(Schedule, MUPA.release());
-      NextLoopData.Schedule =
-          combineInSequence(NextLoopData.Schedule, Schedule);
+      LoopData->Schedule = combineInSequence(LoopData->Schedule, Schedule);
     }
 
-    NextLoopData.NumBlocksProcessed += NumBlocksProcessed;
-    LoopData = NextLoopData;
+    LoopData->NumBlocksProcessed += NumBlocksProcessed;
   }
+  // Now pop all loops processed up there from the LoopStack
+  LoopStack.erase(LoopStack.begin() + Dimension, LoopStack.end());
 }
 
 ArrayRef<ScopStmt *> Scop::getStmtListFor(BasicBlock *BB) const {
   auto StmtMapIt = StmtMap.find(BB);
   if (StmtMapIt == StmtMap.end())
     return {};
   return StmtMapIt->second;
 }