scudo: Split setRandomTag in two. NFCI.

Separate the IRG part from the STZG part since we will need to use
the latter on its own for some upcoming changes.

Differential Revision: https://reviews.llvm.org/D92880

compiler-rt/lib/scudo/standalone/memtag.h

@@ -21,6 +21,9 @@
 
 namespace scudo {
 
+void setRandomTag(void *Ptr, uptr Size, uptr ExcludeMask, uptr *TaggedBegin,
+                  uptr *TaggedEnd);
+
 #if defined(__aarch64__) || defined(SCUDO_FUZZ)
 
 inline constexpr bool archSupportsMemoryTagging() { return true; }
@@ -91,37 +94,32 @@ class ScopedDisableMemoryTagChecks {
   }
 };
 
-inline void setRandomTag(void *Ptr, uptr Size, uptr ExcludeMask,
-                         uptr *TaggedBegin, uptr *TaggedEnd) {
-  void *End;
+inline uptr selectRandomTag(uptr Ptr, uptr ExcludeMask) {
+  uptr TaggedPtr;
   __asm__ __volatile__(
-      R"(
-    .arch_extension mte
-
-    // Set a random tag for Ptr in TaggedPtr. This needs to happen even if
-    // Size = 0 so that TaggedPtr ends up pointing at a valid address.
-    irg %[TaggedPtr], %[Ptr], %[ExcludeMask]
-    mov %[Cur], %[TaggedPtr]
-
-    // Skip the loop if Size = 0. We don't want to do any tagging in this case.
-    cbz %[Size], 2f
-
-    // Set the memory tag of the region
-    // [TaggedPtr, TaggedPtr + roundUpTo(Size, 16))
-    // to the pointer tag stored in TaggedPtr.
-    add %[End], %[TaggedPtr], %[Size]
-
-  1:
-    stzg %[Cur], [%[Cur]], #16
-    cmp %[Cur], %[End]
-    b.lt 1b
+      ".arch_extension mte; irg %[TaggedPtr], %[Ptr], %[ExcludeMask]"
+      : [TaggedPtr] "=r"(TaggedPtr)
+      : [Ptr] "r"(Ptr), [ExcludeMask] "r"(ExcludeMask));
+  return TaggedPtr;
+}
 
-  2:
-  )"
-      :
-      [TaggedPtr] "=&r"(*TaggedBegin), [Cur] "=&r"(*TaggedEnd), [End] "=&r"(End)
-      : [Ptr] "r"(Ptr), [Size] "r"(Size), [ExcludeMask] "r"(ExcludeMask)
-      : "memory");
+inline uptr storeTags(uptr Begin, uptr End) {
+  DCHECK(Begin % 16 == 0);
+  if (Begin != End) {
+    __asm__ __volatile__(
+        R"(
+      .arch_extension mte
+
+    1:
+      stzg %[Cur], [%[Cur]], #16
+      cmp %[Cur], %[End]
+      b.lt 1b
+    )"
+        : [Cur] "+&r"(Begin)
+        : [End] "r"(End)
+        : "memory");
+  }
+  return Begin;
 }
 
 inline void *prepareTaggedChunk(void *Ptr, uptr Size, uptr ExcludeMask,
@@ -224,13 +222,15 @@ struct ScopedDisableMemoryTagChecks {
   ScopedDisableMemoryTagChecks() {}
 };
 
-inline void setRandomTag(void *Ptr, uptr Size, uptr ExcludeMask,
-                         uptr *TaggedBegin, uptr *TaggedEnd) {
+inline uptr selectRandomTag(uptr Ptr, uptr ExcludeMask) {
   (void)Ptr;
-  (void)Size;
   (void)ExcludeMask;
-  (void)TaggedBegin;
-  (void)TaggedEnd;
+  UNREACHABLE("memory tagging not supported");
+}
+
+inline uptr storeTags(uptr Begin, uptr End) {
+  (void)Begin;
+  (void)End;
   UNREACHABLE("memory tagging not supported");
 }
 
@@ -257,6 +257,12 @@ inline uptr loadTag(uptr Ptr) {
 
 #endif
 
+inline void setRandomTag(void *Ptr, uptr Size, uptr ExcludeMask,
+                         uptr *TaggedBegin, uptr *TaggedEnd) {
+  *TaggedBegin = selectRandomTag(reinterpret_cast<uptr>(Ptr), ExcludeMask);
+  *TaggedEnd = storeTags(*TaggedBegin, *TaggedBegin + Size);
+}
+
 } // namespace scudo
 
 #endif