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

[Snehasish Kumar 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();
----------------
snehasish wrote:

nit: I don't think it's comon to use this API in unittests to reset the cl::opt, the only other tests which use this are CommandLineTest.cpp and ExecutionTest.cpp (apart from another usage in this file). In other unittests I've seen just the value being set to the default e.g. false in this case. 

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