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