[NFC] [HWASan] [MTE] factor out threadlong increment (#110340)

Author: fmayer

llvm/include/llvm/Transforms/Utils/MemoryTaggingSupport.h

@@ -99,6 +99,8 @@ Value *getPC(const Triple &TargetTriple, IRBuilder<> &IRB);
 Value *getAndroidSlotPtr(IRBuilder<> &IRB, int Slot);
 
 void annotateDebugRecords(AllocaInfo &Info, unsigned int Tag);
+Value *incrementThreadLong(IRBuilder<> &IRB, Value *ThreadLong,
+                           unsigned int Inc);
 
 } // namespace memtag
 } // namespace llvm

llvm/lib/Target/AArch64/AArch64StackTagging.cpp

@@ -506,19 +506,8 @@ Instruction *AArch64StackTagging::insertBaseTaggedPointer(
     Value *RecordPtr = IRB.CreateIntToPtr(ThreadLong, IRB.getPtrTy(0));
     IRB.CreateStore(PC, RecordPtr);
     IRB.CreateStore(TaggedFP, IRB.CreateConstGEP1_64(IntptrTy, RecordPtr, 1));
-    // Update the ring buffer. Top byte of ThreadLong defines the size of the
-    // buffer in pages, it must be a power of two, and the start of the buffer
-    // must be aligned by twice that much. Therefore wrap around of the ring
-    // buffer is simply Addr &= ~((ThreadLong >> 56) << 12).
-    // The use of AShr instead of LShr is due to
-    //   https://bugs.llvm.org/show_bug.cgi?id=39030
-    // Runtime library makes sure not to use the highest bit.
-    Value *WrapMask = IRB.CreateXor(
-        IRB.CreateShl(IRB.CreateAShr(ThreadLong, 56), 12, "", true, true),
-        ConstantInt::get(IntptrTy, (uint64_t)-1));
-    Value *ThreadLongNew = IRB.CreateAnd(
-        IRB.CreateAdd(ThreadLong, ConstantInt::get(IntptrTy, 16)), WrapMask);
-    IRB.CreateStore(ThreadLongNew, SlotPtr);
+
+    IRB.CreateStore(memtag::incrementThreadLong(IRB, ThreadLong, 16), SlotPtr);
   }
   return Base;
 }

llvm/lib/Transforms/Instrumentation/HWAddressSanitizer.cpp

@@ -1421,35 +1421,7 @@ void HWAddressSanitizer::emitPrologue(IRBuilder<> &IRB, bool WithFrameRecord) {
           IRB.CreateIntToPtr(ThreadLongMaybeUntagged, IRB.getPtrTy(0));
       IRB.CreateStore(FrameRecordInfo, RecordPtr);
 
-      // Update the ring buffer. Top byte of ThreadLong defines the size of the
-      // buffer in pages, it must be a power of two, and the start of the buffer
-      // must be aligned by twice that much. Therefore wrap around of the ring
-      // buffer is simply Addr &= ~((ThreadLong >> 56) << 12).
-      // The use of AShr instead of LShr is due to
-      //   https://bugs.llvm.org/show_bug.cgi?id=39030
-      // Runtime library makes sure not to use the highest bit.
-      //
-      // Mechanical proof of this address calculation can be found at:
-      // https://github.com/google/sanitizers/blob/master/hwaddress-sanitizer/prove_hwasanwrap.smt2
-      //
-      // Example of the wrap case for N = 1
-      // Pointer:   0x01AAAAAAAAAAAFF8
-      //                     +
-      //            0x0000000000000008
-      //                     =
-      //            0x01AAAAAAAAAAB000
-      //                     &
-      // WrapMask:  0xFFFFFFFFFFFFF000
-      //                     =
-      //            0x01AAAAAAAAAAA000
-      //
-      // Then the WrapMask will be a no-op until the next wrap case.
-      Value *WrapMask = IRB.CreateXor(
-          IRB.CreateShl(IRB.CreateAShr(ThreadLong, 56), 12, "", true, true),
-          ConstantInt::get(IntptrTy, (uint64_t)-1));
-      Value *ThreadLongNew = IRB.CreateAnd(
-          IRB.CreateAdd(ThreadLong, ConstantInt::get(IntptrTy, 8)), WrapMask);
-      IRB.CreateStore(ThreadLongNew, SlotPtr);
+      IRB.CreateStore(memtag::incrementThreadLong(IRB, ThreadLong, 8), SlotPtr);
       break;
     }
     case none: {

llvm/lib/Transforms/Utils/MemoryTaggingSupport.cpp

@@ -338,5 +338,39 @@ void annotateDebugRecords(AllocaInfo &Info, unsigned int Tag) {
   llvm::for_each(Info.DbgVariableRecords, AnnotateDbgRecord);
 }
 
+Value *incrementThreadLong(IRBuilder<> &IRB, Value *ThreadLong,
+                           unsigned int Inc) {
+  // Update the ring buffer. Top byte of ThreadLong defines the size of the
+  // buffer in pages, it must be a power of two, and the start of the buffer
+  // must be aligned by twice that much. Therefore wrap around of the ring
+  // buffer is simply Addr &= ~((ThreadLong >> 56) << 12).
+  // The use of AShr instead of LShr is due to
+  //   https://bugs.llvm.org/show_bug.cgi?id=39030
+  // Runtime library makes sure not to use the highest bit.
+  //
+  // Mechanical proof of this address calculation can be found at:
+  // https://github.com/google/sanitizers/blob/master/hwaddress-sanitizer/prove_hwasanwrap.smt2
+  //
+  // Example of the wrap case for N = 1
+  // Pointer:   0x01AAAAAAAAAAAFF8
+  //                     +
+  //            0x0000000000000008
+  //                     =
+  //            0x01AAAAAAAAAAB000
+  //                     &
+  // WrapMask:  0xFFFFFFFFFFFFF000
+  //                     =
+  //            0x01AAAAAAAAAAA000
+  //
+  // Then the WrapMask will be a no-op until the next wrap case.
+  assert((4096 % Inc) == 0);
+  Value *WrapMask = IRB.CreateXor(
+      IRB.CreateShl(IRB.CreateAShr(ThreadLong, 56), 12, "", true, true),
+      ConstantInt::get(ThreadLong->getType(), (uint64_t)-1));
+  return IRB.CreateAnd(
+      IRB.CreateAdd(ThreadLong, ConstantInt::get(ThreadLong->getType(), Inc)),
+      WrapMask);
+}
+
 } // namespace memtag
 } // namespace llvm