llvm-project

1
//===- bolt/Passes/LongJmp.cpp --------------------------------------------===//
2
//
3
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4
// See https://llvm.org/LICENSE.txt for license information.
5
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6
//
7
//===----------------------------------------------------------------------===//
8
//
9
// This file implements the LongJmpPass class.
10
//
11
//===----------------------------------------------------------------------===//
12

13
#include "bolt/Passes/LongJmp.h"
14

15
#define DEBUG_TYPE "longjmp"
16

17
using namespace llvm;
18

19
namespace opts {
20
extern cl::OptionCategory BoltOptCategory;
21
extern llvm::cl::opt<unsigned> AlignText;
22
extern cl::opt<unsigned> AlignFunctions;
23
extern cl::opt<bool> UseOldText;
24
extern cl::opt<bool> HotFunctionsAtEnd;
25

26
static cl::opt<bool> GroupStubs("group-stubs",
27
                                cl::desc("share stubs across functions"),
28
                                cl::init(true), cl::cat(BoltOptCategory));
29
}
30

31
namespace llvm {
32
namespace bolt {
33

34
constexpr unsigned ColdFragAlign = 16;
35

36
static void relaxStubToShortJmp(BinaryBasicBlock &StubBB, const MCSymbol *Tgt) {
37
  const BinaryContext &BC = StubBB.getFunction()->getBinaryContext();
38
  InstructionListType Seq;
39
  BC.MIB->createShortJmp(Seq, Tgt, BC.Ctx.get());
40
  StubBB.clear();
41
  StubBB.addInstructions(Seq.begin(), Seq.end());
42
}
43

44
static void relaxStubToLongJmp(BinaryBasicBlock &StubBB, const MCSymbol *Tgt) {
45
  const BinaryContext &BC = StubBB.getFunction()->getBinaryContext();
46
  InstructionListType Seq;
47
  BC.MIB->createLongJmp(Seq, Tgt, BC.Ctx.get());
48
  StubBB.clear();
49
  StubBB.addInstructions(Seq.begin(), Seq.end());
50
}
51

52
static BinaryBasicBlock *getBBAtHotColdSplitPoint(BinaryFunction &Func) {
53
  if (!Func.isSplit() || Func.empty())
54
    return nullptr;
55

56
  assert(!(*Func.begin()).isCold() && "Entry cannot be cold");
57
  for (auto I = Func.getLayout().block_begin(),
58
            E = Func.getLayout().block_end();
59
       I != E; ++I) {
60
    auto Next = std::next(I);
61
    if (Next != E && (*Next)->isCold())
62
      return *I;
63
  }
64
  llvm_unreachable("No hot-colt split point found");
65
}
66

67
static bool shouldInsertStub(const BinaryContext &BC, const MCInst &Inst) {
68
  return (BC.MIB->isBranch(Inst) || BC.MIB->isCall(Inst)) &&
69
         !BC.MIB->isIndirectBranch(Inst) && !BC.MIB->isIndirectCall(Inst);
70
}
71

72
std::pair<std::unique_ptr<BinaryBasicBlock>, MCSymbol *>
73
LongJmpPass::createNewStub(BinaryBasicBlock &SourceBB, const MCSymbol *TgtSym,
74
                           bool TgtIsFunc, uint64_t AtAddress) {
75
  BinaryFunction &Func = *SourceBB.getFunction();
76
  const BinaryContext &BC = Func.getBinaryContext();
77
  const bool IsCold = SourceBB.isCold();
78
  MCSymbol *StubSym = BC.Ctx->createNamedTempSymbol("Stub");
79
  std::unique_ptr<BinaryBasicBlock> StubBB = Func.createBasicBlock(StubSym);
80
  MCInst Inst;
81
  BC.MIB->createUncondBranch(Inst, TgtSym, BC.Ctx.get());
82
  if (TgtIsFunc)
83
    BC.MIB->convertJmpToTailCall(Inst);
84
  StubBB->addInstruction(Inst);
85
  StubBB->setExecutionCount(0);
86

87
  // Register this in stubs maps
88
  auto registerInMap = [&](StubGroupsTy &Map) {
89
    StubGroupTy &StubGroup = Map[TgtSym];
90
    StubGroup.insert(
91
        llvm::lower_bound(
92
            StubGroup, std::make_pair(AtAddress, nullptr),
93
            [&](const std::pair<uint64_t, BinaryBasicBlock *> &LHS,
94
                const std::pair<uint64_t, BinaryBasicBlock *> &RHS) {
95
              return LHS.first < RHS.first;
96
            }),
97
        std::make_pair(AtAddress, StubBB.get()));
98
  };
99

100
  Stubs[&Func].insert(StubBB.get());
101
  StubBits[StubBB.get()] = BC.MIB->getUncondBranchEncodingSize();
102
  if (IsCold) {
103
    registerInMap(ColdLocalStubs[&Func]);
104
    if (opts::GroupStubs && TgtIsFunc)
105
      registerInMap(ColdStubGroups);
106
    ++NumColdStubs;
107
  } else {
108
    registerInMap(HotLocalStubs[&Func]);
109
    if (opts::GroupStubs && TgtIsFunc)
110
      registerInMap(HotStubGroups);
111
    ++NumHotStubs;
112
  }
113

114
  return std::make_pair(std::move(StubBB), StubSym);
115
}
116

117
BinaryBasicBlock *LongJmpPass::lookupStubFromGroup(
118
    const StubGroupsTy &StubGroups, const BinaryFunction &Func,
119
    const MCInst &Inst, const MCSymbol *TgtSym, uint64_t DotAddress) const {
120
  const BinaryContext &BC = Func.getBinaryContext();
121
  auto CandidatesIter = StubGroups.find(TgtSym);
122
  if (CandidatesIter == StubGroups.end())
123
    return nullptr;
124
  const StubGroupTy &Candidates = CandidatesIter->second;
125
  if (Candidates.empty())
126
    return nullptr;
127
  auto Cand = llvm::lower_bound(
128
      Candidates, std::make_pair(DotAddress, nullptr),
129
      [&](const std::pair<uint64_t, BinaryBasicBlock *> &LHS,
130
          const std::pair<uint64_t, BinaryBasicBlock *> &RHS) {
131
        return LHS.first < RHS.first;
132
      });
133
  if (Cand == Candidates.end())
134
    return nullptr;
135
  if (Cand != Candidates.begin()) {
136
    const StubTy *LeftCand = std::prev(Cand);
137
    if (Cand->first - DotAddress > DotAddress - LeftCand->first)
138
      Cand = LeftCand;
139
  }
140
  int BitsAvail = BC.MIB->getPCRelEncodingSize(Inst) - 1;
141
  assert(BitsAvail < 63 && "PCRelEncodingSize is too large to use int64_t to"
142
                           "check for out-of-bounds.");
143
  int64_t MaxVal = (1ULL << BitsAvail) - 1;
144
  int64_t MinVal = -(1ULL << BitsAvail);
145
  uint64_t PCRelTgtAddress = Cand->first;
146
  int64_t PCOffset = (int64_t)(PCRelTgtAddress - DotAddress);
147

148
  LLVM_DEBUG({
149
    if (Candidates.size() > 1)
150
      dbgs() << "Considering stub group with " << Candidates.size()
151
             << " candidates. DotAddress is " << Twine::utohexstr(DotAddress)
152
             << ", chosen candidate address is "
153
             << Twine::utohexstr(Cand->first) << "\n";
154
  });
155
  return (PCOffset < MinVal || PCOffset > MaxVal) ? nullptr : Cand->second;
156
}
157

158
BinaryBasicBlock *
159
LongJmpPass::lookupGlobalStub(const BinaryBasicBlock &SourceBB,
160
                              const MCInst &Inst, const MCSymbol *TgtSym,
161
                              uint64_t DotAddress) const {
162
  const BinaryFunction &Func = *SourceBB.getFunction();
163
  const StubGroupsTy &StubGroups =
164
      SourceBB.isCold() ? ColdStubGroups : HotStubGroups;
165
  return lookupStubFromGroup(StubGroups, Func, Inst, TgtSym, DotAddress);
166
}
167

168
BinaryBasicBlock *LongJmpPass::lookupLocalStub(const BinaryBasicBlock &SourceBB,
169
                                               const MCInst &Inst,
170
                                               const MCSymbol *TgtSym,
171
                                               uint64_t DotAddress) const {
172
  const BinaryFunction &Func = *SourceBB.getFunction();
173
  const DenseMap<const BinaryFunction *, StubGroupsTy> &StubGroups =
174
      SourceBB.isCold() ? ColdLocalStubs : HotLocalStubs;
175
  const auto Iter = StubGroups.find(&Func);
176
  if (Iter == StubGroups.end())
177
    return nullptr;
178
  return lookupStubFromGroup(Iter->second, Func, Inst, TgtSym, DotAddress);
179
}
180

181
std::unique_ptr<BinaryBasicBlock>
182
LongJmpPass::replaceTargetWithStub(BinaryBasicBlock &BB, MCInst &Inst,
183
                                   uint64_t DotAddress,
184
                                   uint64_t StubCreationAddress) {
185
  const BinaryFunction &Func = *BB.getFunction();
186
  const BinaryContext &BC = Func.getBinaryContext();
187
  std::unique_ptr<BinaryBasicBlock> NewBB;
188
  const MCSymbol *TgtSym = BC.MIB->getTargetSymbol(Inst);
189
  assert(TgtSym && "getTargetSymbol failed");
190

191
  BinaryBasicBlock::BinaryBranchInfo BI{0, 0};
192
  BinaryBasicBlock *TgtBB = BB.getSuccessor(TgtSym, BI);
193
  auto LocalStubsIter = Stubs.find(&Func);
194

195
  // If already using stub and the stub is from another function, create a local
196
  // stub, since the foreign stub is now out of range
197
  if (!TgtBB) {
198
    auto SSIter = SharedStubs.find(TgtSym);
199
    if (SSIter != SharedStubs.end()) {
200
      TgtSym = BC.MIB->getTargetSymbol(*SSIter->second->begin());
201
      --NumSharedStubs;
202
    }
203
  } else if (LocalStubsIter != Stubs.end() &&
204
             LocalStubsIter->second.count(TgtBB)) {
205
    // The TgtBB and TgtSym now are the local out-of-range stub and its label.
206
    // So, we are attempting to restore BB to its previous state without using
207
    // this stub.
208
    TgtSym = BC.MIB->getTargetSymbol(*TgtBB->begin());
209
    assert(TgtSym &&
210
           "First instruction is expected to contain a target symbol.");
211
    BinaryBasicBlock *TgtBBSucc = TgtBB->getSuccessor(TgtSym, BI);
212

213
    // TgtBB might have no successor. e.g. a stub for a function call.
214
    if (TgtBBSucc) {
215
      BB.replaceSuccessor(TgtBB, TgtBBSucc, BI.Count, BI.MispredictedCount);
216
      assert(TgtBB->getExecutionCount() >= BI.Count &&
217
             "At least equal or greater than the branch count.");
218
      TgtBB->setExecutionCount(TgtBB->getExecutionCount() - BI.Count);
219
    }
220

221
    TgtBB = TgtBBSucc;
222
  }
223

224
  BinaryBasicBlock *StubBB = lookupLocalStub(BB, Inst, TgtSym, DotAddress);
225
  // If not found, look it up in globally shared stub maps if it is a function
226
  // call (TgtBB is not set)
227
  if (!StubBB && !TgtBB) {
228
    StubBB = lookupGlobalStub(BB, Inst, TgtSym, DotAddress);
229
    if (StubBB) {
230
      SharedStubs[StubBB->getLabel()] = StubBB;
231
      ++NumSharedStubs;
232
    }
233
  }
234
  MCSymbol *StubSymbol = StubBB ? StubBB->getLabel() : nullptr;
235

236
  if (!StubBB) {
237
    std::tie(NewBB, StubSymbol) =
238
        createNewStub(BB, TgtSym, /*is func?*/ !TgtBB, StubCreationAddress);
239
    StubBB = NewBB.get();
240
  }
241

242
  // Local branch
243
  if (TgtBB) {
244
    uint64_t OrigCount = BI.Count;
245
    uint64_t OrigMispreds = BI.MispredictedCount;
246
    BB.replaceSuccessor(TgtBB, StubBB, OrigCount, OrigMispreds);
247
    StubBB->setExecutionCount(StubBB->getExecutionCount() + OrigCount);
248
    if (NewBB) {
249
      StubBB->addSuccessor(TgtBB, OrigCount, OrigMispreds);
250
      StubBB->setIsCold(BB.isCold());
251
    }
252
    // Call / tail call
253
  } else {
254
    StubBB->setExecutionCount(StubBB->getExecutionCount() +
255
                              BB.getExecutionCount());
256
    if (NewBB) {
257
      assert(TgtBB == nullptr);
258
      StubBB->setIsCold(BB.isCold());
259
      // Set as entry point because this block is valid but we have no preds
260
      StubBB->getFunction()->addEntryPoint(*StubBB);
261
    }
262
  }
263
  BC.MIB->replaceBranchTarget(Inst, StubSymbol, BC.Ctx.get());
264

265
  return NewBB;
266
}
267

268
void LongJmpPass::updateStubGroups() {
269
  auto update = [&](StubGroupsTy &StubGroups) {
270
    for (auto &KeyVal : StubGroups) {
271
      for (StubTy &Elem : KeyVal.second)
272
        Elem.first = BBAddresses[Elem.second];
273
      llvm::sort(KeyVal.second, llvm::less_first());
274
    }
275
  };
276

277
  for (auto &KeyVal : HotLocalStubs)
278
    update(KeyVal.second);
279
  for (auto &KeyVal : ColdLocalStubs)
280
    update(KeyVal.second);
281
  update(HotStubGroups);
282
  update(ColdStubGroups);
283
}
284

285
void LongJmpPass::tentativeBBLayout(const BinaryFunction &Func) {
286
  const BinaryContext &BC = Func.getBinaryContext();
287
  uint64_t HotDot = HotAddresses[&Func];
288
  uint64_t ColdDot = ColdAddresses[&Func];
289
  bool Cold = false;
290
  for (const BinaryBasicBlock *BB : Func.getLayout().blocks()) {
291
    if (Cold || BB->isCold()) {
292
      Cold = true;
293
      BBAddresses[BB] = ColdDot;
294
      ColdDot += BC.computeCodeSize(BB->begin(), BB->end());
295
    } else {
296
      BBAddresses[BB] = HotDot;
297
      HotDot += BC.computeCodeSize(BB->begin(), BB->end());
298
    }
299
  }
300
}
301

302
uint64_t LongJmpPass::tentativeLayoutRelocColdPart(
303
    const BinaryContext &BC, std::vector<BinaryFunction *> &SortedFunctions,
304
    uint64_t DotAddress) {
305
  DotAddress = alignTo(DotAddress, llvm::Align(opts::AlignFunctions));
306
  for (BinaryFunction *Func : SortedFunctions) {
307
    if (!Func->isSplit())
308
      continue;
309
    DotAddress = alignTo(DotAddress, Func->getMinAlignment());
310
    uint64_t Pad =
311
        offsetToAlignment(DotAddress, llvm::Align(Func->getAlignment()));
312
    if (Pad <= Func->getMaxColdAlignmentBytes())
313
      DotAddress += Pad;
314
    ColdAddresses[Func] = DotAddress;
315
    LLVM_DEBUG(dbgs() << Func->getPrintName() << " cold tentative: "
316
                      << Twine::utohexstr(DotAddress) << "\n");
317
    DotAddress += Func->estimateColdSize();
318
    DotAddress = alignTo(DotAddress, Func->getConstantIslandAlignment());
319
    DotAddress += Func->estimateConstantIslandSize();
320
  }
321
  return DotAddress;
322
}
323

324
uint64_t LongJmpPass::tentativeLayoutRelocMode(
325
    const BinaryContext &BC, std::vector<BinaryFunction *> &SortedFunctions,
326
    uint64_t DotAddress) {
327

328
  // Compute hot cold frontier
329
  uint32_t LastHotIndex = -1u;
330
  uint32_t CurrentIndex = 0;
331
  if (opts::HotFunctionsAtEnd) {
332
    for (BinaryFunction *BF : SortedFunctions) {
333
      if (BF->hasValidIndex()) {
334
        LastHotIndex = CurrentIndex;
335
        break;
336
      }
337

338
      ++CurrentIndex;
339
    }
340
  } else {
341
    for (BinaryFunction *BF : SortedFunctions) {
342
      if (!BF->hasValidIndex()) {
343
        LastHotIndex = CurrentIndex;
344
        break;
345
      }
346

347
      ++CurrentIndex;
348
    }
349
  }
350

351
  // Hot
352
  CurrentIndex = 0;
353
  bool ColdLayoutDone = false;
354
  for (BinaryFunction *Func : SortedFunctions) {
355
    if (!BC.shouldEmit(*Func)) {
356
      HotAddresses[Func] = Func->getAddress();
357
      continue;
358
    }
359

360
    if (!ColdLayoutDone && CurrentIndex >= LastHotIndex) {
361
      DotAddress =
362
          tentativeLayoutRelocColdPart(BC, SortedFunctions, DotAddress);
363
      ColdLayoutDone = true;
364
      if (opts::HotFunctionsAtEnd)
365
        DotAddress = alignTo(DotAddress, opts::AlignText);
366
    }
367

368
    DotAddress = alignTo(DotAddress, Func->getMinAlignment());
369
    uint64_t Pad =
370
        offsetToAlignment(DotAddress, llvm::Align(Func->getAlignment()));
371
    if (Pad <= Func->getMaxAlignmentBytes())
372
      DotAddress += Pad;
373
    HotAddresses[Func] = DotAddress;
374
    LLVM_DEBUG(dbgs() << Func->getPrintName() << " tentative: "
375
                      << Twine::utohexstr(DotAddress) << "\n");
376
    if (!Func->isSplit())
377
      DotAddress += Func->estimateSize();
378
    else
379
      DotAddress += Func->estimateHotSize();
380

381
    DotAddress = alignTo(DotAddress, Func->getConstantIslandAlignment());
382
    DotAddress += Func->estimateConstantIslandSize();
383
    ++CurrentIndex;
384
  }
385
  // BBs
386
  for (BinaryFunction *Func : SortedFunctions)
387
    tentativeBBLayout(*Func);
388

389
  return DotAddress;
390
}
391

392
void LongJmpPass::tentativeLayout(
393
    const BinaryContext &BC, std::vector<BinaryFunction *> &SortedFunctions) {
394
  uint64_t DotAddress = BC.LayoutStartAddress;
395

396
  if (!BC.HasRelocations) {
397
    for (BinaryFunction *Func : SortedFunctions) {
398
      HotAddresses[Func] = Func->getAddress();
399
      DotAddress = alignTo(DotAddress, ColdFragAlign);
400
      ColdAddresses[Func] = DotAddress;
401
      if (Func->isSplit())
402
        DotAddress += Func->estimateColdSize();
403
      tentativeBBLayout(*Func);
404
    }
405

406
    return;
407
  }
408

409
  // Relocation mode
410
  uint64_t EstimatedTextSize = 0;
411
  if (opts::UseOldText) {
412
    EstimatedTextSize = tentativeLayoutRelocMode(BC, SortedFunctions, 0);
413

414
    // Initial padding
415
    if (EstimatedTextSize <= BC.OldTextSectionSize) {
416
      DotAddress = BC.OldTextSectionAddress;
417
      uint64_t Pad =
418
          offsetToAlignment(DotAddress, llvm::Align(opts::AlignText));
419
      if (Pad + EstimatedTextSize <= BC.OldTextSectionSize) {
420
        DotAddress += Pad;
421
      }
422
    }
423
  }
424

425
  if (!EstimatedTextSize || EstimatedTextSize > BC.OldTextSectionSize)
426
    DotAddress = alignTo(BC.LayoutStartAddress, opts::AlignText);
427

428
  tentativeLayoutRelocMode(BC, SortedFunctions, DotAddress);
429
}
430

431
bool LongJmpPass::usesStub(const BinaryFunction &Func,
432
                           const MCInst &Inst) const {
433
  const MCSymbol *TgtSym = Func.getBinaryContext().MIB->getTargetSymbol(Inst);
434
  const BinaryBasicBlock *TgtBB = Func.getBasicBlockForLabel(TgtSym);
435
  auto Iter = Stubs.find(&Func);
436
  if (Iter != Stubs.end())
437
    return Iter->second.count(TgtBB);
438
  return false;
439
}
440

441
uint64_t LongJmpPass::getSymbolAddress(const BinaryContext &BC,
442
                                       const MCSymbol *Target,
443
                                       const BinaryBasicBlock *TgtBB) const {
444
  if (TgtBB) {
445
    auto Iter = BBAddresses.find(TgtBB);
446
    assert(Iter != BBAddresses.end() && "Unrecognized BB");
447
    return Iter->second;
448
  }
449
  uint64_t EntryID = 0;
450
  const BinaryFunction *TargetFunc = BC.getFunctionForSymbol(Target, &EntryID);
451
  auto Iter = HotAddresses.find(TargetFunc);
452
  if (Iter == HotAddresses.end() || (TargetFunc && EntryID)) {
453
    // Look at BinaryContext's resolution for this symbol - this is a symbol not
454
    // mapped to a BinaryFunction
455
    ErrorOr<uint64_t> ValueOrError = BC.getSymbolValue(*Target);
456
    assert(ValueOrError && "Unrecognized symbol");
457
    return *ValueOrError;
458
  }
459
  return Iter->second;
460
}
461

462
Error LongJmpPass::relaxStub(BinaryBasicBlock &StubBB, bool &Modified) {
463
  const BinaryFunction &Func = *StubBB.getFunction();
464
  const BinaryContext &BC = Func.getBinaryContext();
465
  const int Bits = StubBits[&StubBB];
466
  // Already working with the largest range?
467
  if (Bits == static_cast<int>(BC.AsmInfo->getCodePointerSize() * 8))
468
    return Error::success();
469

470
  const static int RangeShortJmp = BC.MIB->getShortJmpEncodingSize();
471
  const static int RangeSingleInstr = BC.MIB->getUncondBranchEncodingSize();
472
  const static uint64_t ShortJmpMask = ~((1ULL << RangeShortJmp) - 1);
473
  const static uint64_t SingleInstrMask =
474
      ~((1ULL << (RangeSingleInstr - 1)) - 1);
475

476
  const MCSymbol *RealTargetSym = BC.MIB->getTargetSymbol(*StubBB.begin());
477
  const BinaryBasicBlock *TgtBB = Func.getBasicBlockForLabel(RealTargetSym);
478
  uint64_t TgtAddress = getSymbolAddress(BC, RealTargetSym, TgtBB);
479
  uint64_t DotAddress = BBAddresses[&StubBB];
480
  uint64_t PCRelTgtAddress = DotAddress > TgtAddress ? DotAddress - TgtAddress
481
                                                     : TgtAddress - DotAddress;
482
  // If it fits in one instruction, do not relax
483
  if (!(PCRelTgtAddress & SingleInstrMask))
484
    return Error::success();
485

486
  // Fits short jmp
487
  if (!(PCRelTgtAddress & ShortJmpMask)) {
488
    if (Bits >= RangeShortJmp)
489
      return Error::success();
490

491
    LLVM_DEBUG(dbgs() << "Relaxing stub to short jump. PCRelTgtAddress = "
492
                      << Twine::utohexstr(PCRelTgtAddress)
493
                      << " RealTargetSym = " << RealTargetSym->getName()
494
                      << "\n");
495
    relaxStubToShortJmp(StubBB, RealTargetSym);
496
    StubBits[&StubBB] = RangeShortJmp;
497
    Modified = true;
498
    return Error::success();
499
  }
500

501
  // The long jmp uses absolute address on AArch64
502
  // So we could not use it for PIC binaries
503
  if (BC.isAArch64() && !BC.HasFixedLoadAddress)
504
    return createFatalBOLTError(
505
        "BOLT-ERROR: Unable to relax stub for PIC binary\n");
506

507
  LLVM_DEBUG(dbgs() << "Relaxing stub to long jump. PCRelTgtAddress = "
508
                    << Twine::utohexstr(PCRelTgtAddress)
509
                    << " RealTargetSym = " << RealTargetSym->getName() << "\n");
510
  relaxStubToLongJmp(StubBB, RealTargetSym);
511
  StubBits[&StubBB] = static_cast<int>(BC.AsmInfo->getCodePointerSize() * 8);
512
  Modified = true;
513
  return Error::success();
514
}
515

516
bool LongJmpPass::needsStub(const BinaryBasicBlock &BB, const MCInst &Inst,
517
                            uint64_t DotAddress) const {
518
  const BinaryFunction &Func = *BB.getFunction();
519
  const BinaryContext &BC = Func.getBinaryContext();
520
  const MCSymbol *TgtSym = BC.MIB->getTargetSymbol(Inst);
521
  assert(TgtSym && "getTargetSymbol failed");
522

523
  const BinaryBasicBlock *TgtBB = Func.getBasicBlockForLabel(TgtSym);
524
  // Check for shared stubs from foreign functions
525
  if (!TgtBB) {
526
    auto SSIter = SharedStubs.find(TgtSym);
527
    if (SSIter != SharedStubs.end())
528
      TgtBB = SSIter->second;
529
  }
530

531
  int BitsAvail = BC.MIB->getPCRelEncodingSize(Inst) - 1;
532
  assert(BitsAvail < 63 && "PCRelEncodingSize is too large to use int64_t to"
533
                           "check for out-of-bounds.");
534
  int64_t MaxVal = (1ULL << BitsAvail) - 1;
535
  int64_t MinVal = -(1ULL << BitsAvail);
536

537
  uint64_t PCRelTgtAddress = getSymbolAddress(BC, TgtSym, TgtBB);
538
  int64_t PCOffset = (int64_t)(PCRelTgtAddress - DotAddress);
539

540
  return PCOffset < MinVal || PCOffset > MaxVal;
541
}
542

543
Error LongJmpPass::relax(BinaryFunction &Func, bool &Modified) {
544
  const BinaryContext &BC = Func.getBinaryContext();
545

546
  assert(BC.isAArch64() && "Unsupported arch");
547
  constexpr int InsnSize = 4; // AArch64
548
  std::vector<std::pair<BinaryBasicBlock *, std::unique_ptr<BinaryBasicBlock>>>
549
      Insertions;
550

551
  BinaryBasicBlock *Frontier = getBBAtHotColdSplitPoint(Func);
552
  uint64_t FrontierAddress = Frontier ? BBAddresses[Frontier] : 0;
553
  if (FrontierAddress)
554
    FrontierAddress += Frontier->getNumNonPseudos() * InsnSize;
555

556
  // Add necessary stubs for branch targets we know we can't fit in the
557
  // instruction
558
  for (BinaryBasicBlock &BB : Func) {
559
    uint64_t DotAddress = BBAddresses[&BB];
560
    // Stubs themselves are relaxed on the next loop
561
    if (Stubs[&Func].count(&BB))
562
      continue;
563

564
    for (MCInst &Inst : BB) {
565
      if (BC.MIB->isPseudo(Inst))
566
        continue;
567

568
      if (!shouldInsertStub(BC, Inst)) {
569
        DotAddress += InsnSize;
570
        continue;
571
      }
572

573
      // Check and relax direct branch or call
574
      if (!needsStub(BB, Inst, DotAddress)) {
575
        DotAddress += InsnSize;
576
        continue;
577
      }
578
      Modified = true;
579

580
      // Insert stubs close to the patched BB if call, but far away from the
581
      // hot path if a branch, since this branch target is the cold region
582
      // (but first check that the far away stub will be in range).
583
      BinaryBasicBlock *InsertionPoint = &BB;
584
      if (Func.isSimple() && !BC.MIB->isCall(Inst) && FrontierAddress &&
585
          !BB.isCold()) {
586
        int BitsAvail = BC.MIB->getPCRelEncodingSize(Inst) - 1;
587
        uint64_t Mask = ~((1ULL << BitsAvail) - 1);
588
        assert(FrontierAddress > DotAddress &&
589
               "Hot code should be before the frontier");
590
        uint64_t PCRelTgt = FrontierAddress - DotAddress;
591
        if (!(PCRelTgt & Mask))
592
          InsertionPoint = Frontier;
593
      }
594
      // Always put stubs at the end of the function if non-simple. We can't
595
      // change the layout of non-simple functions because it has jump tables
596
      // that we do not control.
597
      if (!Func.isSimple())
598
        InsertionPoint = &*std::prev(Func.end());
599

600
      // Create a stub to handle a far-away target
601
      Insertions.emplace_back(InsertionPoint,
602
                              replaceTargetWithStub(BB, Inst, DotAddress,
603
                                                    InsertionPoint == Frontier
604
                                                        ? FrontierAddress
605
                                                        : DotAddress));
606

607
      DotAddress += InsnSize;
608
    }
609
  }
610

611
  // Relax stubs if necessary
612
  for (BinaryBasicBlock &BB : Func) {
613
    if (!Stubs[&Func].count(&BB) || !BB.isValid())
614
      continue;
615

616
    if (auto E = relaxStub(BB, Modified))
617
      return Error(std::move(E));
618
  }
619

620
  for (std::pair<BinaryBasicBlock *, std::unique_ptr<BinaryBasicBlock>> &Elmt :
621
       Insertions) {
622
    if (!Elmt.second)
623
      continue;
624
    std::vector<std::unique_ptr<BinaryBasicBlock>> NewBBs;
625
    NewBBs.emplace_back(std::move(Elmt.second));
626
    Func.insertBasicBlocks(Elmt.first, std::move(NewBBs), true);
627
  }
628

629
  return Error::success();
630
}
631

632
Error LongJmpPass::runOnFunctions(BinaryContext &BC) {
633
  BC.outs() << "BOLT-INFO: Starting stub-insertion pass\n";
634
  std::vector<BinaryFunction *> Sorted = BC.getSortedFunctions();
635
  bool Modified;
636
  uint32_t Iterations = 0;
637
  do {
638
    ++Iterations;
639
    Modified = false;
640
    tentativeLayout(BC, Sorted);
641
    updateStubGroups();
642
    for (BinaryFunction *Func : Sorted) {
643
      if (auto E = relax(*Func, Modified))
644
        return Error(std::move(E));
645
      // Don't ruin non-simple functions, they can't afford to have the layout
646
      // changed.
647
      if (Modified && Func->isSimple())
648
        Func->fixBranches();
649
    }
650
  } while (Modified);
651
  BC.outs() << "BOLT-INFO: Inserted " << NumHotStubs
652
            << " stubs in the hot area and " << NumColdStubs
653
            << " stubs in the cold area. Shared " << NumSharedStubs
654
            << " times, iterated " << Iterations << " times.\n";
655
  return Error::success();
656
}
657
} // namespace bolt
658
} // namespace llvm
659