12cab237bSDimitry Andric //===- SafeStackLayout.cpp - SafeStack frame layout -----------------------===//
23ca95b02SDimitry Andric //
33ca95b02SDimitry Andric //                     The LLVM Compiler Infrastructure
43ca95b02SDimitry Andric //
53ca95b02SDimitry Andric // This file is distributed under the University of Illinois Open Source
63ca95b02SDimitry Andric // License. See LICENSE.TXT for details.
73ca95b02SDimitry Andric //
83ca95b02SDimitry Andric //===----------------------------------------------------------------------===//
93ca95b02SDimitry Andric 
103ca95b02SDimitry Andric #include "SafeStackLayout.h"
112cab237bSDimitry Andric #include "SafeStackColoring.h"
122cab237bSDimitry Andric #include "llvm/IR/Value.h"
132cab237bSDimitry Andric #include "llvm/Support/CommandLine.h"
142cab237bSDimitry Andric #include "llvm/Support/Compiler.h"
153ca95b02SDimitry Andric #include "llvm/Support/Debug.h"
162cab237bSDimitry Andric #include "llvm/Support/MathExtras.h"
172cab237bSDimitry Andric #include "llvm/Support/raw_ostream.h"
182cab237bSDimitry Andric #include <algorithm>
192cab237bSDimitry Andric #include <cassert>
203ca95b02SDimitry Andric 
213ca95b02SDimitry Andric using namespace llvm;
223ca95b02SDimitry Andric using namespace llvm::safestack;
233ca95b02SDimitry Andric 
243ca95b02SDimitry Andric #define DEBUG_TYPE "safestacklayout"
253ca95b02SDimitry Andric 
263ca95b02SDimitry Andric static cl::opt<bool> ClLayout("safe-stack-layout",
273ca95b02SDimitry Andric                               cl::desc("enable safe stack layout"), cl::Hidden,
283ca95b02SDimitry Andric                               cl::init(true));
293ca95b02SDimitry Andric 
print(raw_ostream & OS)303ca95b02SDimitry Andric LLVM_DUMP_METHOD void StackLayout::print(raw_ostream &OS) {
313ca95b02SDimitry Andric   OS << "Stack regions:\n";
323ca95b02SDimitry Andric   for (unsigned i = 0; i < Regions.size(); ++i) {
333ca95b02SDimitry Andric     OS << "  " << i << ": [" << Regions[i].Start << ", " << Regions[i].End
343ca95b02SDimitry Andric        << "), range " << Regions[i].Range << "\n";
353ca95b02SDimitry Andric   }
363ca95b02SDimitry Andric   OS << "Stack objects:\n";
373ca95b02SDimitry Andric   for (auto &IT : ObjectOffsets) {
383ca95b02SDimitry Andric     OS << "  at " << IT.getSecond() << ": " << *IT.getFirst() << "\n";
393ca95b02SDimitry Andric   }
403ca95b02SDimitry Andric }
413ca95b02SDimitry Andric 
addObject(const Value * V,unsigned Size,unsigned Alignment,const StackColoring::LiveRange & Range)423ca95b02SDimitry Andric void StackLayout::addObject(const Value *V, unsigned Size, unsigned Alignment,
433ca95b02SDimitry Andric                             const StackColoring::LiveRange &Range) {
443ca95b02SDimitry Andric   StackObjects.push_back({V, Size, Alignment, Range});
45*4ba319b5SDimitry Andric   ObjectAlignments[V] = Alignment;
463ca95b02SDimitry Andric   MaxAlignment = std::max(MaxAlignment, Alignment);
473ca95b02SDimitry Andric }
483ca95b02SDimitry Andric 
AdjustStackOffset(unsigned Offset,unsigned Size,unsigned Alignment)493ca95b02SDimitry Andric static unsigned AdjustStackOffset(unsigned Offset, unsigned Size,
503ca95b02SDimitry Andric                                   unsigned Alignment) {
513ca95b02SDimitry Andric   return alignTo(Offset + Size, Alignment) - Size;
523ca95b02SDimitry Andric }
533ca95b02SDimitry Andric 
layoutObject(StackObject & Obj)543ca95b02SDimitry Andric void StackLayout::layoutObject(StackObject &Obj) {
553ca95b02SDimitry Andric   if (!ClLayout) {
563ca95b02SDimitry Andric     // If layout is disabled, just grab the next aligned address.
573ca95b02SDimitry Andric     // This effectively disables stack coloring as well.
583ca95b02SDimitry Andric     unsigned LastRegionEnd = Regions.empty() ? 0 : Regions.back().End;
593ca95b02SDimitry Andric     unsigned Start = AdjustStackOffset(LastRegionEnd, Obj.Size, Obj.Alignment);
603ca95b02SDimitry Andric     unsigned End = Start + Obj.Size;
613ca95b02SDimitry Andric     Regions.emplace_back(Start, End, Obj.Range);
623ca95b02SDimitry Andric     ObjectOffsets[Obj.Handle] = End;
633ca95b02SDimitry Andric     return;
643ca95b02SDimitry Andric   }
653ca95b02SDimitry Andric 
66*4ba319b5SDimitry Andric   LLVM_DEBUG(dbgs() << "Layout: size " << Obj.Size << ", align "
67*4ba319b5SDimitry Andric                     << Obj.Alignment << ", range " << Obj.Range << "\n");
683ca95b02SDimitry Andric   assert(Obj.Alignment <= MaxAlignment);
693ca95b02SDimitry Andric   unsigned Start = AdjustStackOffset(0, Obj.Size, Obj.Alignment);
703ca95b02SDimitry Andric   unsigned End = Start + Obj.Size;
71*4ba319b5SDimitry Andric   LLVM_DEBUG(dbgs() << "  First candidate: " << Start << " .. " << End << "\n");
723ca95b02SDimitry Andric   for (const StackRegion &R : Regions) {
73*4ba319b5SDimitry Andric     LLVM_DEBUG(dbgs() << "  Examining region: " << R.Start << " .. " << R.End
743ca95b02SDimitry Andric                       << ", range " << R.Range << "\n");
753ca95b02SDimitry Andric     assert(End >= R.Start);
763ca95b02SDimitry Andric     if (Start >= R.End) {
77*4ba319b5SDimitry Andric       LLVM_DEBUG(dbgs() << "  Does not intersect, skip.\n");
783ca95b02SDimitry Andric       continue;
793ca95b02SDimitry Andric     }
803ca95b02SDimitry Andric     if (Obj.Range.Overlaps(R.Range)) {
813ca95b02SDimitry Andric       // Find the next appropriate location.
823ca95b02SDimitry Andric       Start = AdjustStackOffset(R.End, Obj.Size, Obj.Alignment);
833ca95b02SDimitry Andric       End = Start + Obj.Size;
84*4ba319b5SDimitry Andric       LLVM_DEBUG(dbgs() << "  Overlaps. Next candidate: " << Start << " .. "
85*4ba319b5SDimitry Andric                         << End << "\n");
863ca95b02SDimitry Andric       continue;
873ca95b02SDimitry Andric     }
883ca95b02SDimitry Andric     if (End <= R.End) {
89*4ba319b5SDimitry Andric       LLVM_DEBUG(dbgs() << "  Reusing region(s).\n");
903ca95b02SDimitry Andric       break;
913ca95b02SDimitry Andric     }
923ca95b02SDimitry Andric   }
933ca95b02SDimitry Andric 
943ca95b02SDimitry Andric   unsigned LastRegionEnd = Regions.empty() ? 0 : Regions.back().End;
953ca95b02SDimitry Andric   if (End > LastRegionEnd) {
963ca95b02SDimitry Andric     // Insert a new region at the end. Maybe two.
973ca95b02SDimitry Andric     if (Start > LastRegionEnd) {
98*4ba319b5SDimitry Andric       LLVM_DEBUG(dbgs() << "  Creating gap region: " << LastRegionEnd << " .. "
993ca95b02SDimitry Andric                         << Start << "\n");
1003ca95b02SDimitry Andric       Regions.emplace_back(LastRegionEnd, Start, StackColoring::LiveRange());
1013ca95b02SDimitry Andric       LastRegionEnd = Start;
1023ca95b02SDimitry Andric     }
103*4ba319b5SDimitry Andric     LLVM_DEBUG(dbgs() << "  Creating new region: " << LastRegionEnd << " .. "
104*4ba319b5SDimitry Andric                       << End << ", range " << Obj.Range << "\n");
1053ca95b02SDimitry Andric     Regions.emplace_back(LastRegionEnd, End, Obj.Range);
1063ca95b02SDimitry Andric     LastRegionEnd = End;
1073ca95b02SDimitry Andric   }
1083ca95b02SDimitry Andric 
1093ca95b02SDimitry Andric   // Split starting and ending regions if necessary.
1106c4bc1bdSDimitry Andric   for (unsigned i = 0; i < Regions.size(); ++i) {
1116c4bc1bdSDimitry Andric     StackRegion &R = Regions[i];
1123ca95b02SDimitry Andric     if (Start > R.Start && Start < R.End) {
1133ca95b02SDimitry Andric       StackRegion R0 = R;
1143ca95b02SDimitry Andric       R.Start = R0.End = Start;
1153ca95b02SDimitry Andric       Regions.insert(&R, R0);
1163ca95b02SDimitry Andric       continue;
1173ca95b02SDimitry Andric     }
1183ca95b02SDimitry Andric     if (End > R.Start && End < R.End) {
1193ca95b02SDimitry Andric       StackRegion R0 = R;
1203ca95b02SDimitry Andric       R0.End = R.Start = End;
1213ca95b02SDimitry Andric       Regions.insert(&R, R0);
1223ca95b02SDimitry Andric       break;
1233ca95b02SDimitry Andric     }
1243ca95b02SDimitry Andric   }
1253ca95b02SDimitry Andric 
1263ca95b02SDimitry Andric   // Update live ranges for all affected regions.
1273ca95b02SDimitry Andric   for (StackRegion &R : Regions) {
1283ca95b02SDimitry Andric     if (Start < R.End && End > R.Start)
1293ca95b02SDimitry Andric       R.Range.Join(Obj.Range);
1303ca95b02SDimitry Andric     if (End <= R.End)
1313ca95b02SDimitry Andric       break;
1323ca95b02SDimitry Andric   }
1333ca95b02SDimitry Andric 
1343ca95b02SDimitry Andric   ObjectOffsets[Obj.Handle] = End;
1353ca95b02SDimitry Andric }
1363ca95b02SDimitry Andric 
computeLayout()1373ca95b02SDimitry Andric void StackLayout::computeLayout() {
1383ca95b02SDimitry Andric   // Simple greedy algorithm.
1393ca95b02SDimitry Andric   // If this is replaced with something smarter, it must preserve the property
1403ca95b02SDimitry Andric   // that the first object is always at the offset 0 in the stack frame (for
1413ca95b02SDimitry Andric   // StackProtectorSlot), or handle stack protector in some other way.
142d88c1a5aSDimitry Andric 
143d88c1a5aSDimitry Andric   // Sort objects by size (largest first) to reduce fragmentation.
144d88c1a5aSDimitry Andric   if (StackObjects.size() > 2)
145d88c1a5aSDimitry Andric     std::stable_sort(StackObjects.begin() + 1, StackObjects.end(),
146d88c1a5aSDimitry Andric                      [](const StackObject &a, const StackObject &b) {
147d88c1a5aSDimitry Andric                        return a.Size > b.Size;
148d88c1a5aSDimitry Andric                      });
149d88c1a5aSDimitry Andric 
1503ca95b02SDimitry Andric   for (auto &Obj : StackObjects)
1513ca95b02SDimitry Andric     layoutObject(Obj);
1523ca95b02SDimitry Andric 
153*4ba319b5SDimitry Andric   LLVM_DEBUG(print(dbgs()));
1543ca95b02SDimitry Andric }
155