[MemProf] Add helper to access the back (last) call stack id

[tejohnson]

Summary

This is split out of D140908 as suggested.

[snehasish]

lgtm

llvm/include/llvm/Analysis/MemoryProfileInfo.h

141

nit: I prefer the prior version with an early check since it's easier to read than a nested ternary.

if (!N) {
 Iter = nullptr;
 return;
}

Iter = End ? N->StackIdIndices.end() : N->StackIdIndices.begin();

Also writing it out made it clear that the change in this diff is ensuring that Iter is initialized to nullptr.

[tejohnson]

Address comments

llvm/include/llvm/Analysis/MemoryProfileInfo.h

struct CallStackIterator {
     void operator++() { ++Iter; }
   };
 
   bool empty() const { return N == nullptr; }
 
   CallStackIterator begin() const;
   CallStackIterator end() const { return CallStackIterator(N, /*End*/ true); }
   CallStackIterator beginAfterSharedPrefix(CallStack &Other);
+  uint64_t back() const;
 
 private:
   const NodeT *N = nullptr;
 };
 
 template <class NodeT, class IteratorT>
 CallStack<NodeT, IteratorT>::CallStackIterator::CallStackIterator(
     const NodeT *N, bool End)
     : N(N) {
-  if (!N)
-    return;
-  Iter = End ? N->StackIdIndices.end() : N->StackIdIndices.begin();
+  Iter =
+      N ? (End ? N->StackIdIndices.end() : N->StackIdIndices.begin()) : nullptr;
 }
 
 template <class NodeT, class IteratorT>
 uint64_t CallStack<NodeT, IteratorT>::CallStackIterator::operator*() {
   assert(Iter != N->StackIdIndices.end());
   return *Iter;
 }
 
 template <class NodeT, class IteratorT>
+uint64_t CallStack<NodeT, IteratorT>::back() const {
+  assert(N);
+  return N->StackIdIndices.back();
+}
+
+template <class NodeT, class IteratorT>
 typename CallStack<NodeT, IteratorT>::CallStackIterator
 CallStack<NodeT, IteratorT>::begin() const {
   return CallStackIterator(N, /*End*/ false);
 }
 
 template <class NodeT, class IteratorT>
 typename CallStack<NodeT, IteratorT>::CallStackIterator
 CallStack<NodeT, IteratorT>::beginAfterSharedPrefix(CallStack &Other) {
   CallStackIterator Cur = begin();
   for (CallStackIterator OtherCur = Other.begin();
        Cur != end() && OtherCur != Other.end(); ++Cur, ++OtherCur)
     assert(*Cur == *OtherCur);
   return Cur;
 }
 
 /// Specializations for iterating through IR metadata stack contexts.
 template <>
 CallStack<MDNode, MDNode::op_iterator>::CallStackIterator::CallStackIterator(
     const MDNode *N, bool End);
 template <>
 uint64_t CallStack<MDNode, MDNode::op_iterator>::CallStackIterator::operator*();
+template <> uint64_t CallStack<MDNode, MDNode::op_iterator>::back() const;
 
 } // end namespace memprof
 } // end namespace llvm
 
 #endif

llvm/lib/Analysis/MemoryProfileInfo.cpp

 template <>
 uint64_t
 CallStack<MDNode, MDNode::op_iterator>::CallStackIterator::operator*() {
   assert(Iter != N->op_end());
   ConstantInt *StackIdCInt = mdconst::dyn_extract<ConstantInt>(*Iter);
   assert(StackIdCInt);
   return StackIdCInt->getZExtValue();
 }
+
+template <> uint64_t CallStack<MDNode, MDNode::op_iterator>::back() const {
+  assert(N);
+  return mdconst::dyn_extract<ConstantInt>(N->operands().back())
+      ->getZExtValue();
+}

llvm/unittests/Analysis/MemoryProfileInfoTest.cpp

CallStack<MDNode, MDNode::op_iterator> InstCallsite(
       Call->getMetadata(LLVMContext::MD_callsite));
 
   MDNode *MemProfMD = Call->getMetadata(LLVMContext::MD_memprof);
   bool First = true;
   for (auto &MIBOp : MemProfMD->operands()) {
     auto *MIBMD = cast<const MDNode>(MIBOp);
     MDNode *StackNode = getMIBStackNode(MIBMD);
     CallStack<MDNode, MDNode::op_iterator> StackContext(StackNode);
+    uint64_t ExpectedBack = First ? 4 : 5;
+    EXPECT_EQ(StackContext.back(), ExpectedBack);
     std::vector<uint64_t> StackIds;
     for (auto ContextIter = StackContext.beginAfterSharedPrefix(InstCallsite);
          ContextIter != StackContext.end(); ++ContextIter)
       StackIds.push_back(*ContextIter);
     if (First) {
       std::vector<uint64_t> Expected = {2, 3, 4};
       EXPECT_EQ(ArrayRef(StackIds), ArrayRef(Expected));
     } else {
)Summary");
 
   auto *AllocSummary =
       cast<FunctionSummary>(Index->getGlobalValueSummary(/*guid=*/23));
   for (auto &AI : AllocSummary->allocs()) {
     bool First = true;
     for (auto &MIB : AI.MIBs) {
       CallStack<MIBInfo, SmallVector<unsigned>::const_iterator> StackContext(
           &MIB);
+      uint64_t ExpectedBack = First ? 4 : 5;
+      EXPECT_EQ(Index->getStackIdAtIndex(StackContext.back()), ExpectedBack);
       std::vector<uint64_t> StackIds;
       for (auto StackIdIndex : StackContext)
         StackIds.push_back(Index->getStackIdAtIndex(StackIdIndex));
       if (First) {
         std::vector<uint64_t> Expected = {1, 2, 3, 4};
         EXPECT_EQ(ArrayRef(StackIds), ArrayRef(Expected));
       } else {
         std::vector<uint64_t> Expected = {1, 2, 3, 5};
         EXPECT_EQ(ArrayRef(StackIds), ArrayRef(Expected));
       }
       First = false;
     }
   }
 }
 } // end anonymous namespace