[asan] Simplify calculation of stack frame layout extraction calculation of stack description into separate function.

[eugenis]

looks great!

include/llvm/Transforms/Utils/ASanStackFrameLayout.h

struct ASanStackVariableDescription {
   AllocaInst *AI;      // The actual AllocaInst.
   size_t Offset;       // Offset from the beginning of the frame;
                        // set by ComputeASanStackFrameLayout.
   unsigned Line;       // Line number.
 };
 
 // Output data struct for ComputeASanStackFrameLayout.
 struct ASanStackFrameLayout {
-  size_t Granularity;
-  // Frame description, see DescribeAddressIfStack in ASan runtime.
-  SmallString<64> DescriptionString;
+  size_t Granularity;     // Shadow granularity.
   size_t FrameAlignment;  // Alignment for the entire frame.
   size_t FrameSize;       // Size of the frame in bytes.
 };
 
 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);
 
+// Compute frame description, see DescribeAddressIfStack in ASan runtime.
+SmallString<64> ComputeASanStackFrameDescription(
+    const SmallVectorImpl<ASanStackVariableDescription> &Vars);
+
 // Returns shadow bytes with marked red zones. This shadow represents the state
 // if the stack frame when all local variables are inside of the own scope.
 SmallVector<uint8_t, 64>
 GetShadowBytes(const SmallVectorImpl<ASanStackVariableDescription> &Vars,
                const ASanStackFrameLayout &Layout);
 
 // Returns shadow bytes with marked red zones and after scope. This shadow
 // represents the state if the stack frame when all local variables are outside

lib/Transforms/Instrumentation/AddressSanitizer.cpp

void FunctionStackPoisoner::processStaticAllocas() {
   for (BasicBlock::iterator I(InsBefore); I != InsBeforeB->end(); ++I)
     if (auto *AI = dyn_cast<AllocaInst>(I))
       if (NonInstrumentedStaticAllocaVec.count(AI) > 0)
         AI->moveBefore(InsBefore);
 
   // If we have a call to llvm.localescape, keep it in the entry block.
   if (LocalEscapeCall) LocalEscapeCall->moveBefore(InsBefore);
 
-  // Find static allocas with lifetime analysis.
-  DenseMap<const AllocaInst *,
-           std::pair<const ASanStackVariableDescription *, unsigned>>
-      AllocaToSVDMap;
-  for (const auto &APC : StaticAllocaPoisonCallVec) {
-    assert(APC.InsBefore);
-    assert(APC.AI);
-    assert(ASan.isInterestingAlloca(*APC.AI));
-    assert(APC.AI->isStaticAlloca());
-
-    auto &Pair = AllocaToSVDMap[APC.AI];
-    if (const DILocation *FnLoc = EntryDebugLocation.get()) {
-      if (const DILocation *LifetimeLoc = APC.InsBefore->getDebugLoc().get()) {
-        if (LifetimeLoc->getFile() == FnLoc->getFile())
-          if (unsigned Line = LifetimeLoc->getLine())
-            Pair.second = std::min(Pair.second ? Pair.second : Line, Line);
-      }
-    }
-  }
-
   SmallVector<ASanStackVariableDescription, 16> SVD;
   SVD.reserve(AllocaVec.size());
   for (AllocaInst *AI : AllocaVec) {
     ASanStackVariableDescription D = {AI->getName().data(),
                                       ASan.getAllocaSizeInBytes(*AI),
                                       0,
                                       AI->getAlignment(),
                                       AI,
                                       0,
                                       0};
-    auto It = AllocaToSVDMap.find(AI);
-    if (It != AllocaToSVDMap.end()) {
-      D.LifetimeSize = D.Size;
-      D.Line = It->second.second;
-    }
     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;
   const ASanStackFrameLayout &L =
       ComputeASanStackFrameLayout(SVD, 1ULL << Mapping.Scale, MinHeaderSize);
 
-  DEBUG(dbgs() << L.DescriptionString << " --- " << L.FrameSize << "\n");
+  // Build AllocaToSVDMap for ASanStackVariableDescription lookup.
+  DenseMap<const AllocaInst *, ASanStackVariableDescription *> AllocaToSVDMap;
+  for (auto &Desc : SVD)
+    AllocaToSVDMap[Desc.AI] = &Desc;
+
+  // Update SVD with information from lifetime intrinsics.
+  for (const auto &APC : StaticAllocaPoisonCallVec) {
+    assert(APC.InsBefore);
+    assert(APC.AI);
+    assert(ASan.isInterestingAlloca(*APC.AI));
+    assert(APC.AI->isStaticAlloca());
+
+    ASanStackVariableDescription &Desc = *AllocaToSVDMap[APC.AI];
+    Desc.LifetimeSize = Desc.Size;
+    if (const DILocation *FnLoc = EntryDebugLocation.get()) {
+      if (const DILocation *LifetimeLoc = APC.InsBefore->getDebugLoc().get()) {
+        if (LifetimeLoc->getFile() == FnLoc->getFile())
+          if (unsigned Line = LifetimeLoc->getLine())
+            Desc.Line = std::min(Desc.Line ? Desc.Line : Line, Line);
+      }
+    }
+  }
+
+  auto DescriptionString = ComputeASanStackFrameDescription(SVD);
+  DEBUG(dbgs() << 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.
   // 2) There is a returns_twice call (typically setjmp), which is
void FunctionStackPoisoner::processStaticAllocas() {
   IRB.CreateStore(ConstantInt::get(IntptrTy, kCurrentStackFrameMagic),
                   BasePlus0);
   // Write the frame description constant to redzone[1].
   Value *BasePlus1 = IRB.CreateIntToPtr(
       IRB.CreateAdd(LocalStackBase,
                     ConstantInt::get(IntptrTy, ASan.LongSize / 8)),
       IntptrPtrTy);
   GlobalVariable *StackDescriptionGlobal =
-      createPrivateGlobalForString(*F.getParent(), L.DescriptionString,
+      createPrivateGlobalForString(*F.getParent(), DescriptionString,
                                    /*AllowMerging*/ true);
   Value *Description = IRB.CreatePointerCast(StackDescriptionGlobal, IntptrTy);
   IRB.CreateStore(Description, BasePlus1);
   // Write the PC to redzone[2].
   Value *BasePlus2 = IRB.CreateIntToPtr(
       IRB.CreateAdd(LocalStackBase,
                     ConstantInt::get(IntptrTy, 2 * ASan.LongSize / 8)),
       IntptrPtrTy);
   IRB.CreateStore(IRB.CreatePointerCast(&F, IntptrTy), BasePlus2);
 
   const auto &ShadowAfterScope = GetShadowBytesAfterScope(SVD, L);
 
   // Poison the stack red zones at the entry.
   Value *ShadowBase = ASan.memToShadow(LocalStackBase, IRB);
   // As mask we must use most poisoned case: red zones and after scope.
   // As bytes we can use either the same or just red zones only.
   copyToShadow(ShadowAfterScope, ShadowAfterScope, IRB, ShadowBase);
 
   if (!StaticAllocaPoisonCallVec.empty()) {
-    // Complete AllocaToSVDMap
-    for (const auto &Desc : SVD) {
-      auto It = AllocaToSVDMap.find(Desc.AI);
-      if (It != AllocaToSVDMap.end()) {
-        It->second.first = &Desc;
-      }
-    }
-
     const auto &ShadowInScope = GetShadowBytes(SVD, L);
 
     // Poison static allocas near lifetime intrinsics.
     for (const auto &APC : StaticAllocaPoisonCallVec) {
-      // Must be already set.
-      assert(AllocaToSVDMap[APC.AI].first);
-      const auto &Desc = *AllocaToSVDMap[APC.AI].first;
+      const ASanStackVariableDescription &Desc = *AllocaToSVDMap[APC.AI];
       assert(Desc.Offset % L.Granularity == 0);
       size_t Begin = Desc.Offset / L.Granularity;
       size_t End = Begin + (APC.Size + L.Granularity - 1) / L.Granularity;
 
       IRBuilder<> IRB(APC.InsBefore);
       copyToShadow(ShadowAfterScope,
                    APC.DoPoison ? ShadowAfterScope : ShadowInScope, Begin, End,
                    IRB, ShadowBase);

lib/Transforms/Utils/ASanStackFrameLayout.cpp

ComputeASanStackFrameLayout(SmallVectorImpl<ASanStackVariableDescription> &Vars,
   assert(MinHeaderSize >= 16 && (MinHeaderSize & (MinHeaderSize - 1)) == 0 &&
          MinHeaderSize >= Granularity);
   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;
 
   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);
   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;
-    std::string Name = Vars[i].Name;
     assert((Alignment & (Alignment - 1)) == 0);
     assert(Layout.FrameAlignment >= Alignment);
     assert((Offset % Alignment) == 0);
     assert(Size > 0);
-    assert(Vars[i].LifetimeSize <= Size);
-    if (Vars[i].Line) {
-      Name += ":";
-      Name += std::to_string(Vars[i].Line);
-    }
-    StackDescription << " " << Offset << " " << Size << " " << Name.size()
-                     << " " << Name;
     size_t NextAlignment = IsLast ? Granularity
                    : std::max(Granularity, Vars[i + 1].Alignment);
-    size_t SizeWithRedzone = VarAndRedzoneSize(Vars[i].Size, NextAlignment);
+    size_t SizeWithRedzone = VarAndRedzoneSize(Size, NextAlignment);
     Vars[i].Offset = Offset;
     Offset += SizeWithRedzone;
   }
   if (Offset % MinHeaderSize) {
     Offset += MinHeaderSize - (Offset % MinHeaderSize);
   }
-  Layout.DescriptionString = StackDescription.str();
   Layout.FrameSize = Offset;
   assert((Layout.FrameSize % MinHeaderSize) == 0);
-
   return Layout;
 }
 
+SmallString<64> ComputeASanStackFrameDescription(
+    const SmallVectorImpl<ASanStackVariableDescription> &Vars) {
+  SmallString<2048> StackDescriptionStorage;
+  raw_svector_ostream StackDescription(StackDescriptionStorage);
+  StackDescription << Vars.size();
+
+  for (const auto &Var : Vars) {
+    std::string Name = Var.Name;
+    if (Var.Line) {
+      Name += ":";
+      Name += std::to_string(Var.Line);
+    }
+    StackDescription << " " << Var.Offset << " " << Var.Size << " "
+                     << Name.size() << " " << Name;
+  }
+  return StackDescription.str();
+}
+
 SmallVector<uint8_t, 64>
 GetShadowBytes(const SmallVectorImpl<ASanStackVariableDescription> &Vars,
                const ASanStackFrameLayout &Layout) {
   assert(Vars.size() > 0);
   SmallVector<uint8_t, 64> SB;
   SB.clear();
   const size_t Granularity = Layout.Granularity;
   SB.resize(Vars[0].Offset / Granularity, kAsanStackLeftRedzoneMagic);

 
 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) {
+    assert(Var.LifetimeSize <= Var.Size);
     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

 
 // 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_STREQ(ExpectedDescr, L.DescriptionString.c_str());                  \
+    EXPECT_STREQ(ExpectedDescr,                                                \
+                 ComputeASanStackFrameDescription(Vars).c_str());              \
     EXPECT_EQ(ExpectedShadow, ShadowBytesToString(GetShadowBytes(Vars, L)));   \
     EXPECT_EQ(ExpectedShadowAfterScope,                                        \
               ShadowBytesToString(GetShadowBytesAfterScope(Vars, L)));         \
   }
 
 TEST(ASanStackFrameLayout, Test) {
 #define VAR(name, size, lifetime, alignment, line)                             \
   ASanStackVariableDescription name##size##_##alignment = {                    \