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