[llvm] [MemProf] Fix bug introduced by restructuring in optional handling (PR #139092)

[Teresa Johnson via llvm-commits]

================
@@ -530,6 +531,78 @@ declare dso_local noalias noundef i8* @malloc(i64 noundef)
   EXPECT_THAT(MemProfMD, MemprofMetadataEquals(ExpectedVals));
 }
 
+// Test to ensure that we keep optionally keep unneeded NotCold contexts.
+// Same as PruneUnneededNotColdContexts test but with the
+// MemProfKeepAllNotColdContexts set to true.
+TEST_F(MemoryProfileInfoTest, KeepUnneededNotColdContexts) {
+  LLVMContext C;
+  std::unique_ptr<Module> M = makeLLVMModule(C,
+                                             R"IR(
+target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-pc-linux-gnu"
+define i32* @test() {
+entry:
+  %call = call noalias dereferenceable_or_null(40) i8* @malloc(i64 noundef 40)
+  %0 = bitcast i8* %call to i32*
+  ret i32* %0
+}
+declare dso_local noalias noundef i8* @malloc(i64 noundef)
+)IR");
+
+  Function *Func = M->getFunction("test");
+
+  CallStackTrie Trie;
+
+  Trie.addCallStack(AllocationType::Cold, {1, 2, 3, 4});
+  Trie.addCallStack(AllocationType::Cold, {1, 2, 3, 5, 6, 7});
+  // This NotCold context is needed to know where the above two Cold contexts
+  // must be cloned from:
+  Trie.addCallStack(AllocationType::NotCold, {1, 2, 3, 5, 6, 13});
+
+  Trie.addCallStack(AllocationType::Cold, {1, 2, 3, 8, 9, 10});
+  // This NotCold context is needed to know where the above Cold context must be
+  // cloned from:
+  Trie.addCallStack(AllocationType::NotCold, {1, 2, 3, 8, 9, 14});
+  // This NotCold context is not needed since the above is sufficient (we pick
+  // the first in sorted order).
+  Trie.addCallStack(AllocationType::NotCold, {1, 2, 3, 8, 9, 15});
+
+  // None of these NotCold contexts are needed as the Cold contexts they
+  // overlap with are covered by longer overlapping NotCold contexts.
+  Trie.addCallStack(AllocationType::NotCold, {1, 2, 3, 12});
+  Trie.addCallStack(AllocationType::NotCold, {1, 2, 11});
+  Trie.addCallStack(AllocationType::NotCold, {1, 16});
+
+  // We should keep all of the above contexts, even those that are unneeded by
+  // default, because we will set the option to keep them.
+  std::vector<std::pair<AllocationType, std::vector<unsigned>>> ExpectedVals = {
+      {AllocationType::Cold, {1, 2, 3, 4}},
+      {AllocationType::Cold, {1, 2, 3, 5, 6, 7}},
+      {AllocationType::NotCold, {1, 2, 3, 5, 6, 13}},
+      {AllocationType::Cold, {1, 2, 3, 8, 9, 10}},
+      {AllocationType::NotCold, {1, 2, 3, 8, 9, 14}},
+      {AllocationType::NotCold, {1, 2, 3, 8, 9, 15}},
+      {AllocationType::NotCold, {1, 2, 3, 12}},
+      {AllocationType::NotCold, {1, 2, 11}},
+      {AllocationType::NotCold, {1, 16}}};
+
+  CallBase *Call = findCall(*Func, "call");
+  ASSERT_NE(Call, nullptr);
+
+  // Specify that all non-cold contexts should be kept.
+  MemProfKeepAllNotColdContexts = true;
+
+  Trie.buildAndAttachMIBMetadata(Call);
+
+  // Undo the manual set of the MemProfKeepAllNotColdContexts above.
+  cl::ResetAllOptionOccurrences();
----------------
teresajohnson wrote:

I use it in one other place in this file to do something similar, and landed on that solution because there is no copy constructor for cl::opt so I couldn't save and restore the default (and was concerned about hardwiring it). But I can just use if conditions to do the save and restore. Will send a follow on change to modify both usages in this file to that approach.

https://github.com/llvm/llvm-project/pull/139092