1*a5da256fSEvgeniy Stepanov //===-- SafeStackLayout.cpp - SafeStack frame layout -----------*- C++ -*--===//
2*a5da256fSEvgeniy Stepanov //
3*a5da256fSEvgeniy Stepanov //                     The LLVM Compiler Infrastructure
4*a5da256fSEvgeniy Stepanov //
5*a5da256fSEvgeniy Stepanov // This file is distributed under the University of Illinois Open Source
6*a5da256fSEvgeniy Stepanov // License. See LICENSE.TXT for details.
7*a5da256fSEvgeniy Stepanov //
8*a5da256fSEvgeniy Stepanov //===----------------------------------------------------------------------===//
9*a5da256fSEvgeniy Stepanov 
10*a5da256fSEvgeniy Stepanov #include "SafeStackLayout.h"
11*a5da256fSEvgeniy Stepanov 
12*a5da256fSEvgeniy Stepanov #include "llvm/IR/Instructions.h"
13*a5da256fSEvgeniy Stepanov #include "llvm/Support/Debug.h"
14*a5da256fSEvgeniy Stepanov 
15*a5da256fSEvgeniy Stepanov using namespace llvm;
16*a5da256fSEvgeniy Stepanov using namespace llvm::safestack;
17*a5da256fSEvgeniy Stepanov 
18*a5da256fSEvgeniy Stepanov #define DEBUG_TYPE "safestacklayout"
19*a5da256fSEvgeniy Stepanov 
20*a5da256fSEvgeniy Stepanov static cl::opt<bool> ClLayout("safe-stack-layout",
21*a5da256fSEvgeniy Stepanov                               cl::desc("enable safe stack layout"), cl::Hidden,
22*a5da256fSEvgeniy Stepanov                               cl::init(true));
23*a5da256fSEvgeniy Stepanov 
24*a5da256fSEvgeniy Stepanov LLVM_DUMP_METHOD void StackLayout::print(raw_ostream &OS) {
25*a5da256fSEvgeniy Stepanov   OS << "Stack regions:\n";
26*a5da256fSEvgeniy Stepanov   for (unsigned i = 0; i < Regions.size(); ++i) {
27*a5da256fSEvgeniy Stepanov     OS << "  " << i << ": [" << Regions[i].Start << ", " << Regions[i].End
28*a5da256fSEvgeniy Stepanov        << "), range " << Regions[i].Range << "\n";
29*a5da256fSEvgeniy Stepanov   }
30*a5da256fSEvgeniy Stepanov   OS << "Stack objects:\n";
31*a5da256fSEvgeniy Stepanov   for (auto &IT : ObjectOffsets) {
32*a5da256fSEvgeniy Stepanov     OS << "  at " << IT.getSecond() << ": " << *IT.getFirst() << "\n";
33*a5da256fSEvgeniy Stepanov   }
34*a5da256fSEvgeniy Stepanov }
35*a5da256fSEvgeniy Stepanov 
36*a5da256fSEvgeniy Stepanov void StackLayout::addObject(const Value *V, unsigned Size, unsigned Alignment,
37*a5da256fSEvgeniy Stepanov                             const StackColoring::LiveRange &Range) {
38*a5da256fSEvgeniy Stepanov   StackObjects.push_back({V, Size, Alignment, Range});
39*a5da256fSEvgeniy Stepanov   MaxAlignment = std::max(MaxAlignment, Alignment);
40*a5da256fSEvgeniy Stepanov }
41*a5da256fSEvgeniy Stepanov 
42*a5da256fSEvgeniy Stepanov static unsigned AdjustStackOffset(unsigned Offset, unsigned Size,
43*a5da256fSEvgeniy Stepanov                                   unsigned Alignment) {
44*a5da256fSEvgeniy Stepanov   return alignTo(Offset + Size, Alignment) - Size;
45*a5da256fSEvgeniy Stepanov }
46*a5da256fSEvgeniy Stepanov 
47*a5da256fSEvgeniy Stepanov void StackLayout::layoutObject(StackObject &Obj) {
48*a5da256fSEvgeniy Stepanov   if (!ClLayout) {
49*a5da256fSEvgeniy Stepanov     // If layout is disabled, just grab the next aligned address.
50*a5da256fSEvgeniy Stepanov     // This effectively disables stack coloring as well.
51*a5da256fSEvgeniy Stepanov     unsigned LastRegionEnd = Regions.empty() ? 0 : Regions.back().End;
52*a5da256fSEvgeniy Stepanov     unsigned Start = AdjustStackOffset(LastRegionEnd, Obj.Size, Obj.Alignment);
53*a5da256fSEvgeniy Stepanov     unsigned End = Start + Obj.Size;
54*a5da256fSEvgeniy Stepanov     Regions.emplace_back(Start, End, Obj.Range);
55*a5da256fSEvgeniy Stepanov     ObjectOffsets[Obj.Handle] = End;
56*a5da256fSEvgeniy Stepanov     return;
57*a5da256fSEvgeniy Stepanov   }
58*a5da256fSEvgeniy Stepanov 
59*a5da256fSEvgeniy Stepanov   DEBUG(dbgs() << "Layout: size " << Obj.Size << ", align " << Obj.Alignment
60*a5da256fSEvgeniy Stepanov                << ", range " << Obj.Range << "\n");
61*a5da256fSEvgeniy Stepanov   assert(Obj.Alignment <= MaxAlignment);
62*a5da256fSEvgeniy Stepanov   unsigned Start = AdjustStackOffset(0, Obj.Size, Obj.Alignment);
63*a5da256fSEvgeniy Stepanov   unsigned End = Start + Obj.Size;
64*a5da256fSEvgeniy Stepanov   DEBUG(dbgs() << "  First candidate: " << Start << " .. " << End << "\n");
65*a5da256fSEvgeniy Stepanov   for (const StackRegion &R : Regions) {
66*a5da256fSEvgeniy Stepanov     DEBUG(dbgs() << "  Examining region: " << R.Start << " .. " << R.End
67*a5da256fSEvgeniy Stepanov                  << ", range " << R.Range << "\n");
68*a5da256fSEvgeniy Stepanov     assert(End >= R.Start);
69*a5da256fSEvgeniy Stepanov     if (Start >= R.End) {
70*a5da256fSEvgeniy Stepanov       DEBUG(dbgs() << "  Does not intersect, skip.\n");
71*a5da256fSEvgeniy Stepanov       continue;
72*a5da256fSEvgeniy Stepanov     }
73*a5da256fSEvgeniy Stepanov     if (Obj.Range.Overlaps(R.Range)) {
74*a5da256fSEvgeniy Stepanov       // Find the next appropriate location.
75*a5da256fSEvgeniy Stepanov       Start = AdjustStackOffset(R.End, Obj.Size, Obj.Alignment);
76*a5da256fSEvgeniy Stepanov       End = Start + Obj.Size;
77*a5da256fSEvgeniy Stepanov       DEBUG(dbgs() << "  Overlaps. Next candidate: " << Start << " .. " << End
78*a5da256fSEvgeniy Stepanov                    << "\n");
79*a5da256fSEvgeniy Stepanov       continue;
80*a5da256fSEvgeniy Stepanov     }
81*a5da256fSEvgeniy Stepanov     if (End <= R.End) {
82*a5da256fSEvgeniy Stepanov       DEBUG(dbgs() << "  Reusing region(s).\n");
83*a5da256fSEvgeniy Stepanov       break;
84*a5da256fSEvgeniy Stepanov     }
85*a5da256fSEvgeniy Stepanov   }
86*a5da256fSEvgeniy Stepanov 
87*a5da256fSEvgeniy Stepanov   unsigned LastRegionEnd = Regions.empty() ? 0 : Regions.back().End;
88*a5da256fSEvgeniy Stepanov   if (End > LastRegionEnd) {
89*a5da256fSEvgeniy Stepanov     // Insert a new region at the end. Maybe two.
90*a5da256fSEvgeniy Stepanov     if (Start > LastRegionEnd) {
91*a5da256fSEvgeniy Stepanov       DEBUG(dbgs() << "  Creating gap region: " << LastRegionEnd << " .. "
92*a5da256fSEvgeniy Stepanov                    << Start << "\n");
93*a5da256fSEvgeniy Stepanov       Regions.emplace_back(LastRegionEnd, Start, StackColoring::LiveRange());
94*a5da256fSEvgeniy Stepanov       LastRegionEnd = Start;
95*a5da256fSEvgeniy Stepanov     }
96*a5da256fSEvgeniy Stepanov     DEBUG(dbgs() << "  Creating new region: " << LastRegionEnd << " .. " << End
97*a5da256fSEvgeniy Stepanov                  << ", range " << Obj.Range << "\n");
98*a5da256fSEvgeniy Stepanov     Regions.emplace_back(LastRegionEnd, End, Obj.Range);
99*a5da256fSEvgeniy Stepanov     LastRegionEnd = End;
100*a5da256fSEvgeniy Stepanov   }
101*a5da256fSEvgeniy Stepanov 
102*a5da256fSEvgeniy Stepanov   // Split starting and ending regions if necessary.
103*a5da256fSEvgeniy Stepanov   for (StackRegion &R : Regions) {
104*a5da256fSEvgeniy Stepanov     if (Start > R.Start && Start < R.End) {
105*a5da256fSEvgeniy Stepanov       StackRegion R0 = R;
106*a5da256fSEvgeniy Stepanov       R.Start = R0.End = Start;
107*a5da256fSEvgeniy Stepanov       Regions.insert(&R, R0);
108*a5da256fSEvgeniy Stepanov       continue;
109*a5da256fSEvgeniy Stepanov     }
110*a5da256fSEvgeniy Stepanov     if (End > R.Start && End < R.End) {
111*a5da256fSEvgeniy Stepanov       StackRegion R0 = R;
112*a5da256fSEvgeniy Stepanov       R0.End = R.Start = End;
113*a5da256fSEvgeniy Stepanov       Regions.insert(&R, R0);
114*a5da256fSEvgeniy Stepanov       break;
115*a5da256fSEvgeniy Stepanov     }
116*a5da256fSEvgeniy Stepanov   }
117*a5da256fSEvgeniy Stepanov 
118*a5da256fSEvgeniy Stepanov   // Update live ranges for all affected regions.
119*a5da256fSEvgeniy Stepanov   for (StackRegion &R : Regions) {
120*a5da256fSEvgeniy Stepanov     if (Start < R.End && End > R.Start)
121*a5da256fSEvgeniy Stepanov       R.Range.Join(Obj.Range);
122*a5da256fSEvgeniy Stepanov     if (End <= R.End)
123*a5da256fSEvgeniy Stepanov       break;
124*a5da256fSEvgeniy Stepanov   }
125*a5da256fSEvgeniy Stepanov 
126*a5da256fSEvgeniy Stepanov   ObjectOffsets[Obj.Handle] = End;
127*a5da256fSEvgeniy Stepanov }
128*a5da256fSEvgeniy Stepanov 
129*a5da256fSEvgeniy Stepanov void StackLayout::computeLayout() {
130*a5da256fSEvgeniy Stepanov   // Simple greedy algorithm.
131*a5da256fSEvgeniy Stepanov   // If this is replaced with something smarter, it must preserve the property
132*a5da256fSEvgeniy Stepanov   // that the first object is always at the offset 0 in the stack frame (for
133*a5da256fSEvgeniy Stepanov   // StackProtectorSlot), or handle stack protector in some other way.
134*a5da256fSEvgeniy Stepanov   for (auto &Obj : StackObjects)
135*a5da256fSEvgeniy Stepanov     layoutObject(Obj);
136*a5da256fSEvgeniy Stepanov 
137*a5da256fSEvgeniy Stepanov   DEBUG(print(dbgs()));
138*a5da256fSEvgeniy Stepanov }
139