[asan] Separate calculation of ShadowBytes from calculating ASanStackFrameLayout

[vitalybuka]

Summary

No functional changes, just refactoring to make D23947 simpler.

[eugenis]

LGTM

include/llvm/Transforms/Utils/ASanStackFrameLayout.h

74

delete

[vitalybuka]

Removed unused declaration

[vitalybuka]

[vitalybuka]

updated doc string

include/llvm/Transforms/Utils/ASanStackFrameLayout.h

struct ASanStackVariableDescription {
   size_t Alignment;    // Alignment of the variable (power of 2).
   AllocaInst *AI;      // The actual AllocaInst.
   size_t Offset;       // Offset from the beginning of the frame;
                        // set by ComputeASanStackFrameLayout.
 };
 
 // Output data struct for ComputeASanStackFrameLayout.
 struct ASanStackFrameLayout {
+  size_t Granularity;
   // Frame description, see DescribeAddressIfStack in ASan runtime.
   SmallString<64> DescriptionString;
-  // The contents of the shadow memory for the stack frame that we need
-  // to set at function entry.
-  SmallVector<uint8_t, 64> ShadowBytes;
   size_t FrameAlignment;  // Alignment for the entire frame.
   size_t FrameSize;       // Size of the frame in bytes.
 };
 
-void ComputeASanStackFrameLayout(
+ASanStackFrameLayout ComputeASanStackFrameLayout(
     // The array of stack variables. The elements may get reordered and changed.
     SmallVectorImpl<ASanStackVariableDescription> &Vars,
     // AddressSanitizer's shadow granularity. Usually 8, may also be 16, 32, 64.
     size_t Granularity,
     // The minimal size of the left-most redzone (header).
     // At least 4 pointer sizes, power of 2, and >= Granularity.
     // The resulting FrameSize should be multiple of MinHeaderSize.
-    size_t MinHeaderSize,
-    // The result is put here.
-    ASanStackFrameLayout *Layout);
+    size_t MinHeaderSize);
+
+// Returns shadow bytes with marked red zones.
+SmallVector<uint8_t, 64>
+GetShadowBytes(const SmallVectorImpl<ASanStackVariableDescription> &Vars,
+               const ASanStackFrameLayout &Layout);
+
+// Returns shadow bytes with marked red zones and after scope.
+SmallVector<uint8_t, 64> GetShadowBytesAfterScope(
+    // The array of stack variables. The elements may get reordered and changed.
+    const SmallVectorImpl<ASanStackVariableDescription> &Vars,
+    const ASanStackFrameLayout &Layout);
+
+// Returns shadow bytes with marked red zones and after scope.
+SmallVector<uint8_t, 64> GetShadowBytesAfterReturn(
+    // The array of stack variables. The elements may get reordered and changed.
+    const SmallVectorImpl<ASanStackVariableDescription> &Vars,
+    const ASanStackFrameLayout &Layout);
 
 } // llvm namespace
 
 #endif  // LLVM_TRANSFORMS_UTILS_ASANSTACKFRAMELAYOUT_H

lib/Transforms/Instrumentation/AddressSanitizer.cpp

ASanStackVariableDescription D = {AI->getName().data(),
                                       AI->getAlignment(),
                                       AI,
                                       0};
     SVD.push_back(D);
   }
   // Minimal header size (left redzone) is 4 pointers,
   // i.e. 32 bytes on 64-bit platforms and 16 bytes in 32-bit platforms.
   size_t MinHeaderSize = ASan.LongSize / 2;
-  ASanStackFrameLayout L;
-  ComputeASanStackFrameLayout(SVD, 1ULL << Mapping.Scale, MinHeaderSize, &L);
+  const ASanStackFrameLayout &L =
+      ComputeASanStackFrameLayout(SVD, 1ULL << Mapping.Scale, MinHeaderSize);
+  const SmallVector<uint8_t, 64> &ShadowBytes =
+      GetShadowBytesAfterScope(SVD, L);
   DEBUG(dbgs() << L.DescriptionString << " --- " << L.FrameSize << "\n");
   uint64_t LocalStackSize = L.FrameSize;
   bool DoStackMalloc = ClUseAfterReturn && !ASan.CompileKernel &&
                        LocalStackSize <= kMaxStackMallocSize;
   bool DoDynamicAlloca = ClDynamicAllocaStack;
   // Don't do dynamic alloca or stack malloc if:
   // 1) There is inline asm: too often it makes assumptions on which registers
   //    are available.
void FunctionStackPoisoner::processStaticAllocas() {
   Value *BasePlus2 = IRB.CreateIntToPtr(
       IRB.CreateAdd(LocalStackBase,
                     ConstantInt::get(IntptrTy, 2 * ASan.LongSize / 8)),
       IntptrPtrTy);
   IRB.CreateStore(IRB.CreatePointerCast(&F, IntptrTy), BasePlus2);
 
   // Poison the stack redzones at the entry.
   Value *ShadowBase = ASan.memToShadow(LocalStackBase, IRB);
-  poisonStackFrame(L.ShadowBytes, IRB, ShadowBase, true);
+  poisonStackFrame(ShadowBytes, IRB, ShadowBase, true);
 
   auto UnpoisonStack = [&](IRBuilder<> &IRB) {
     // Do this always as poisonAlloca can be disabled with
     // detect_stack_use_after_scope=0.
-    poisonStackFrame(L.ShadowBytes, IRB, ShadowBase, false);
+    poisonStackFrame(ShadowBytes, IRB, ShadowBase, false);
     if (!StaticAllocaPoisonCallVec.empty()) {
       // If we poisoned some allocas in llvm.lifetime analysis,
       // unpoison whole stack frame now.
       poisonAlloca(LocalStackBase, LocalStackSize, IRB, false);
     }
   };
 
   SmallVector<uint8_t, 64> ShadowBytesAfterReturn;

lib/Transforms/Utils/ASanStackFrameLayout.cpp

static size_t VarAndRedzoneSize(size_t Size, size_t Alignment) {
   else if (Size <= 16) Res = 32;
   else if (Size <= 128) Res = Size + 32;
   else if (Size <= 512) Res = Size + 64;
   else if (Size <= 4096) Res = Size + 128;
   else                   Res = Size + 256;
   return alignTo(Res, Alignment);
 }
 
-void
+ASanStackFrameLayout
 ComputeASanStackFrameLayout(SmallVectorImpl<ASanStackVariableDescription> &Vars,
-                            size_t Granularity, size_t MinHeaderSize,
-                            ASanStackFrameLayout *Layout) {
+                            size_t Granularity, size_t MinHeaderSize) {
   assert(Granularity >= 8 && Granularity <= 64 &&
          (Granularity & (Granularity - 1)) == 0);
   assert(MinHeaderSize >= 16 && (MinHeaderSize & (MinHeaderSize - 1)) == 0 &&
          MinHeaderSize >= Granularity);
-  size_t NumVars = Vars.size();
+  const size_t NumVars = Vars.size();
   assert(NumVars > 0);
   for (size_t i = 0; i < NumVars; i++)
     Vars[i].Alignment = std::max(Vars[i].Alignment, kMinAlignment);
 
   std::stable_sort(Vars.begin(), Vars.end(), CompareVars);
   SmallString<2048> StackDescriptionStorage;
   raw_svector_ostream StackDescription(StackDescriptionStorage);
   StackDescription << NumVars;
-  Layout->FrameAlignment = std::max(Granularity, Vars[0].Alignment);
-  SmallVector<uint8_t, 64> &SB(Layout->ShadowBytes);
-  SB.clear();
+
+  ASanStackFrameLayout Layout;
+  Layout.Granularity = Granularity;
+  Layout.FrameAlignment = std::max(Granularity, Vars[0].Alignment);
   size_t Offset = std::max(std::max(MinHeaderSize, Granularity),
      Vars[0].Alignment);
   assert((Offset % Granularity) == 0);
-  SB.insert(SB.end(), Offset / Granularity, kAsanStackLeftRedzoneMagic);
   for (size_t i = 0; i < NumVars; i++) {
     bool IsLast = i == NumVars - 1;
     size_t Alignment = std::max(Granularity, Vars[i].Alignment);
     (void)Alignment;  // Used only in asserts.
     size_t Size = Vars[i].Size;
     const char *Name = Vars[i].Name;
     assert((Alignment & (Alignment - 1)) == 0);
-    assert(Layout->FrameAlignment >= Alignment);
+    assert(Layout.FrameAlignment >= Alignment);
     assert((Offset % Alignment) == 0);
     assert(Size > 0);
     assert(Vars[i].LifetimeSize <= Size);
     StackDescription << " " << Offset << " " << Size << " " << strlen(Name)
                      << " " << Name;
     size_t NextAlignment = IsLast ? Granularity
                    : std::max(Granularity, Vars[i + 1].Alignment);
     size_t SizeWithRedzone = VarAndRedzoneSize(Vars[i].Size, NextAlignment);
-    size_t LifetimeShadowSize =
-        (Vars[i].LifetimeSize + Granularity - 1) / Granularity;
-    SB.insert(SB.end(), LifetimeShadowSize, kAsanStackUseAfterScopeMagic);
-    if (Size / Granularity >= LifetimeShadowSize) {
-      SB.insert(SB.end(), Size / Granularity - LifetimeShadowSize, 0);
-      if (Size % Granularity)
-        SB.insert(SB.end(), Size % Granularity);
-    }
-    SB.insert(SB.end(), (SizeWithRedzone - Size) / Granularity,
-        IsLast ? kAsanStackRightRedzoneMagic
-        : kAsanStackMidRedzoneMagic);
     Vars[i].Offset = Offset;
     Offset += SizeWithRedzone;
-    assert(Offset == SB.size() * Granularity);
   }
   if (Offset % MinHeaderSize) {
-    size_t ExtraRedzone = MinHeaderSize - (Offset % MinHeaderSize);
-    SB.insert(SB.end(), ExtraRedzone / Granularity,
-              kAsanStackRightRedzoneMagic);
-    Offset += ExtraRedzone;
-  }
-  Layout->DescriptionString = StackDescription.str();
-  Layout->FrameSize = Offset;
-  assert((Layout->FrameSize % MinHeaderSize) == 0);
-  assert(Layout->FrameSize / Granularity == Layout->ShadowBytes.size());
+    Offset += MinHeaderSize - (Offset % MinHeaderSize);
+  }
+  Layout.DescriptionString = StackDescription.str();
+  Layout.FrameSize = Offset;
+  assert((Layout.FrameSize % MinHeaderSize) == 0);
+
+  return Layout;
+}
+
+SmallVector<uint8_t, 64>
+GetShadowBytes(const SmallVectorImpl<ASanStackVariableDescription> &Vars,
+               const ASanStackFrameLayout &Layout) {
+  SmallVector<uint8_t, 64> SB;
+  SB.clear();
+  const size_t NumVars = Vars.size();
+  assert(NumVars > 0);
+  const size_t Granularity = Layout.Granularity;
+  SB.resize(Vars[0].Offset / Granularity, kAsanStackLeftRedzoneMagic);
+  for (const auto &Var : Vars) {
+    SB.resize(Var.Offset / Granularity, kAsanStackMidRedzoneMagic);
+
+    SB.resize(SB.size() + Var.Size / Granularity, 0);
+    if (Var.Size % Granularity)
+      SB.push_back(Var.Size % Granularity);
+  }
+  SB.resize(Layout.FrameSize / Granularity, kAsanStackRightRedzoneMagic);
+  return SB;
+}
+
+SmallVector<uint8_t, 64> GetShadowBytesAfterScope(
+    const SmallVectorImpl<ASanStackVariableDescription> &Vars,
+    const ASanStackFrameLayout &Layout) {
+  SmallVector<uint8_t, 64> SB = GetShadowBytes(Vars, Layout);
+  const size_t Granularity = Layout.Granularity;
+
+  for (const auto &Var : Vars) {
+    const size_t LifetimeShadowSize =
+        (Var.LifetimeSize + Granularity - 1) / Granularity;
+    const size_t Offset = Var.Offset / Granularity;
+    std::fill(SB.begin() + Offset, SB.begin() + Offset + LifetimeShadowSize,
+              kAsanStackUseAfterScopeMagic);
+  }
+
+  return SB;
 }
 
 } // llvm namespace

unittests/Transforms/Utils/ASanStackFrameLayoutTest.cpp

switch (ShadowBytes[i]) {
         os << "S";
         break;
       default:                            os << (unsigned)ShadowBytes[i];
     }
   }
   return os.str();
 }
 
-static void TestLayout(SmallVector<ASanStackVariableDescription, 10> Vars,
-                       size_t Granularity, size_t MinHeaderSize,
-                       const std::string &ExpectedDescr,
-                       const std::string &ExpectedShadow) {
-  ASanStackFrameLayout L;
-  ComputeASanStackFrameLayout(Vars, Granularity, MinHeaderSize, &L);
-  EXPECT_EQ(ExpectedDescr, L.DescriptionString);
-  EXPECT_EQ(ExpectedShadow, ShadowBytesToString(L.ShadowBytes));
+// Use macro to preserve line information in EXPECT_EQ output.
+#define TEST_LAYOUT(V, Granularity, MinHeaderSize, ExpectedDescr,              \
+                    ExpectedShadow, ExpectedShadowAfterScope)                  \
+  {                                                                            \
+    SmallVector<ASanStackVariableDescription, 10> Vars = V;                    \
+    ASanStackFrameLayout L =                                                   \
+        ComputeASanStackFrameLayout(Vars, Granularity, MinHeaderSize);         \
+    EXPECT_EQ(ExpectedDescr, L.DescriptionString);                             \
+    EXPECT_EQ(ExpectedShadow, ShadowBytesToString(GetShadowBytes(Vars, L)));   \
+    EXPECT_EQ(ExpectedShadowAfterScope,                                        \
+              ShadowBytesToString(GetShadowBytesAfterScope(Vars, L)));         \
   }
 
 TEST(ASanStackFrameLayout, Test) {
-#define VEC1(a) SmallVector<ASanStackVariableDescription, 10>(1, a)
-#define VEC(a)                                                                 \
-  SmallVector<ASanStackVariableDescription, 10>(a, a + sizeof(a) / sizeof(a[0]))
-
 #define VAR(name, size, lifetime, alignment)                                   \
   ASanStackVariableDescription name##size##_##alignment = {                    \
     #name #size "_" #alignment,                                                \
     size,                                                                      \
     lifetime,                                                                  \
     alignment,                                                                 \
     0,                                                                         \
     0                                                                          \
   }
 
   VAR(a, 1, 0, 1);
   VAR(p, 1, 0, 32);
   VAR(p, 1, 0, 256);
   VAR(a, 2, 0, 1);
   VAR(a, 3, 0, 1);
   VAR(a, 4, 0, 1);
   VAR(a, 7, 0, 1);
   VAR(a, 8, 8, 1);
   VAR(a, 9, 0, 1);
-  VAR(a, 16, 0, 1);
+  VAR(a, 16, 16, 1);
   VAR(a, 41, 9, 1);
   VAR(a, 105, 103, 1);
 
-  TestLayout(VEC1(a1_1), 8, 16, "1 16 1 4 a1_1", "LL1R");
-  TestLayout(VEC1(a1_1), 64, 64, "1 64 1 4 a1_1", "L1");
-  TestLayout(VEC1(p1_32), 8, 32, "1 32 1 5 p1_32", "LLLL1RRR");
-  TestLayout(VEC1(p1_32), 8, 64, "1 64 1 5 p1_32", "LLLLLLLL1RRRRRRR");
-
-  TestLayout(VEC1(a1_1), 8, 32, "1 32 1 4 a1_1", "LLLL1RRR");
-  TestLayout(VEC1(a2_1), 8, 32, "1 32 2 4 a2_1", "LLLL2RRR");
-  TestLayout(VEC1(a3_1), 8, 32, "1 32 3 4 a3_1", "LLLL3RRR");
-  TestLayout(VEC1(a4_1), 8, 32, "1 32 4 4 a4_1", "LLLL4RRR");
-  TestLayout(VEC1(a7_1), 8, 32, "1 32 7 4 a7_1", "LLLL7RRR");
-  TestLayout(VEC1(a8_1), 8, 32, "1 32 8 4 a8_1", "LLLLSRRR");
-  TestLayout(VEC1(a9_1), 8, 32, "1 32 9 4 a9_1", "LLLL01RR");
-  TestLayout(VEC1(a16_1), 8, 32, "1 32 16 5 a16_1", "LLLL00RR");
-  TestLayout(VEC1(p1_256), 8, 32, "1 256 1 6 p1_256",
+  TEST_LAYOUT({a1_1}, 8, 16, "1 16 1 4 a1_1", "LL1R", "LL1R");
+  TEST_LAYOUT({a1_1}, 64, 64, "1 64 1 4 a1_1", "L1", "L1");
+  TEST_LAYOUT({p1_32}, 8, 32, "1 32 1 5 p1_32", "LLLL1RRR", "LLLL1RRR");
+  TEST_LAYOUT({p1_32}, 8, 64, "1 64 1 5 p1_32", "LLLLLLLL1RRRRRRR",
+              "LLLLLLLL1RRRRRRR");
+
+  TEST_LAYOUT({a1_1}, 8, 32, "1 32 1 4 a1_1", "LLLL1RRR", "LLLL1RRR");
+  TEST_LAYOUT({a2_1}, 8, 32, "1 32 2 4 a2_1", "LLLL2RRR", "LLLL2RRR");
+  TEST_LAYOUT({a3_1}, 8, 32, "1 32 3 4 a3_1", "LLLL3RRR", "LLLL3RRR");
+  TEST_LAYOUT({a4_1}, 8, 32, "1 32 4 4 a4_1", "LLLL4RRR", "LLLL4RRR");
+  TEST_LAYOUT({a7_1}, 8, 32, "1 32 7 4 a7_1", "LLLL7RRR", "LLLL7RRR");
+  TEST_LAYOUT({a8_1}, 8, 32, "1 32 8 4 a8_1", "LLLL0RRR", "LLLLSRRR");
+  TEST_LAYOUT({a9_1}, 8, 32, "1 32 9 4 a9_1", "LLLL01RR", "LLLL01RR");
+  TEST_LAYOUT({a16_1}, 8, 32, "1 32 16 5 a16_1", "LLLL00RR", "LLLLSSRR");
+  TEST_LAYOUT({p1_256}, 8, 32, "1 256 1 6 p1_256",
+              "LLLLLLLLLLLLLLLLLLLLLLLLLLLLLLLL1RRR",
               "LLLLLLLLLLLLLLLLLLLLLLLLLLLLLLLL1RRR");
-  TestLayout(VEC1(a41_1), 8, 32, "1 32 41 5 a41_1", "LLLLSS0001RRRRRR");
-  TestLayout(VEC1(a105_1), 8, 32, "1 32 105 6 a105_1",
+  TEST_LAYOUT({a41_1}, 8, 32, "1 32 41 5 a41_1", "LLLL000001RRRRRR",
+              "LLLLSS0001RRRRRR");
+  TEST_LAYOUT({a105_1}, 8, 32, "1 32 105 6 a105_1", "LLLL00000000000001RRRRRR",
               "LLLLSSSSSSSSSSSSS1RRRRRR");
 
   {
-    ASanStackVariableDescription t[] = {a1_1, p1_256};
-    TestLayout(VEC(t), 8, 32, "2 256 1 6 p1_256 272 1 4 a1_1",
+    SmallVector<ASanStackVariableDescription, 10> t = {a1_1, p1_256};
+    TEST_LAYOUT(t, 8, 32, "2 256 1 6 p1_256 272 1 4 a1_1",
+                "LLLLLLLLLLLLLLLLLLLLLLLLLLLLLLLL1M1R",
                 "LLLLLLLLLLLLLLLLLLLLLLLLLLLLLLLL1M1R");
   }
 
   {
-    ASanStackVariableDescription t[] = {a1_1, a16_1, a41_1};
-    TestLayout(VEC(t), 8, 32, "3 32 1 4 a1_1 48 16 5 a16_1 80 41 5 a41_1",
-               "LLLL1M00MMSS0001RRRR");
+    SmallVector<ASanStackVariableDescription, 10> t = {a1_1, a16_1, a41_1};
+    TEST_LAYOUT(t, 8, 32, "3 32 1 4 a1_1 48 16 5 a16_1 80 41 5 a41_1",
+                "LLLL1M00MM000001RRRR", "LLLL1MSSMMSS0001RRRR");
   }
-#undef VEC1
-#undef VEC
 #undef VAR
+#undef TEST_LAYOUT
 }