[compiler-rt] [ctxprof][nfc] Refactor `__llvm_ctx_profile_start_context` (PR #133744)

[Mircea Trofin via llvm-commits]

https://github.com/mtrofin created https://github.com/llvm/llvm-project/pull/133744

Most of the functionality will be reused with the auto-root detection mechanism (which is introduced subsequently in PR# 133147).

>From da439460a6d2ec77ba8c74d8150dd02a07339535 Mon Sep 17 00:00:00 2001
From: Mircea Trofin <mtrofin at google.com>
Date: Mon, 31 Mar 2025 09:13:41 -0700
Subject: [PATCH] [ctxprof][nfc] Refactor `__llvm_ctx_profile_start_context`

Most of the functionality will be reused with the auto-root detection
mechanism (which is introduced subsequently in PR# 133147).
---
 .../lib/ctx_profile/CtxInstrProfiling.cpp     | 39 ++++++++++---------
 1 file changed, 21 insertions(+), 18 deletions(-)

diff --git a/compiler-rt/lib/ctx_profile/CtxInstrProfiling.cpp b/compiler-rt/lib/ctx_profile/CtxInstrProfiling.cpp
index da291e0bbabdd..10a6a8c1f71e5 100644
--- a/compiler-rt/lib/ctx_profile/CtxInstrProfiling.cpp
+++ b/compiler-rt/lib/ctx_profile/CtxInstrProfiling.cpp
@@ -277,6 +277,25 @@ ContextRoot *FunctionData::getOrAllocateContextRoot() {
   return Root;
 }
 
+ContextNode *tryStartContextGivenRoot(ContextRoot *Root, GUID Guid,
+                                      uint32_t Counters, uint32_t Callsites)
+    SANITIZER_NO_THREAD_SAFETY_ANALYSIS {
+  IsUnderContext = true;
+  __sanitizer::atomic_fetch_add(&Root->TotalEntries, 1,
+                                __sanitizer::memory_order_relaxed);
+  if (!Root->FirstMemBlock) {
+    setupContext(Root, Guid, Counters, Callsites);
+  }
+  if (Root->Taken.TryLock()) {
+    __llvm_ctx_profile_current_context_root = Root;
+    onContextEnter(*Root->FirstNode);
+    return Root->FirstNode;
+  }
+  // If this thread couldn't take the lock, return scratch context.
+  __llvm_ctx_profile_current_context_root = nullptr;
+  return TheScratchContext;
+}
+
 ContextNode *getUnhandledContext(FunctionData &Data, GUID Guid,
                                  uint32_t NumCounters) {
 
@@ -369,24 +388,8 @@ ContextNode *__llvm_ctx_profile_get_context(FunctionData *Data, void *Callee,
 ContextNode *__llvm_ctx_profile_start_context(
     FunctionData *FData, GUID Guid, uint32_t Counters,
     uint32_t Callsites) SANITIZER_NO_THREAD_SAFETY_ANALYSIS {
-  IsUnderContext = true;
-
-  auto *Root = FData->getOrAllocateContextRoot();
-
-  __sanitizer::atomic_fetch_add(&Root->TotalEntries, 1,
-                                __sanitizer::memory_order_relaxed);
-
-  if (!Root->FirstMemBlock) {
-    setupContext(Root, Guid, Counters, Callsites);
-  }
-  if (Root->Taken.TryLock()) {
-    __llvm_ctx_profile_current_context_root = Root;
-    onContextEnter(*Root->FirstNode);
-    return Root->FirstNode;
-  }
-  // If this thread couldn't take the lock, return scratch context.
-  __llvm_ctx_profile_current_context_root = nullptr;
-  return TheScratchContext;
+  return tryStartContextGivenRoot(FData->getOrAllocateContextRoot(), Guid,
+                                  Counters, Callsites);
 }
 
 void __llvm_ctx_profile_release_context(FunctionData *FData)