1 files changed, 171 insertions, 0 deletions
diff --git a/backend/lib/CodeGen/MachineBranchRelaxation.cpp b/backend/lib/CodeGen/MachineBranchRelaxation.cpp
new file mode 100644
index 0000000..f6c7df1
--- /dev/null
+++ b/backend/lib/CodeGen/MachineBranchRelaxation.cpp
@@ -0,0 +1,171 @@
+// SPDX-License-Identifier: GPL-3.0-or-later
+
+#include "rcc/CodeGen/MachineBranchRelaxation.h"
+
+#include "rcc/CodeGen/MachineBasicBlock.h"
+#include "rcc/CodeGen/MachineFunction.h"
+#include "rcc/CodeGen/MachineInstr.h"
+#include "rcc/CodeGen/TargetDescription.h"
+#include "rcc/CodeGen/TargetInstrInfo.h"
+
+#include <cstddef>
+#include <cstdint>
+#include <limits>
+#include <optional>
+#include <string>
+#include <unordered_map>
+#include <utility>
+
+namespace rcc {
+namespace {
+
+struct MachineLayout final {
+  std::unordered_map<const MachineBasicBlock *, std::uint64_t> BlockOffsets;
+  std::unordered_map<const MachineInstr *, std::uint64_t> InstructionOffsets;
+};
+
+[[nodiscard]] const MachinePassContract &relaxationContract() {
+  static const MachinePassContract Contract = [] {
+    MachinePassContract Value{
+        .Name = "branch-relaxation",
+        .Requires = {},
+        .Sets = {},
+        .Clears = {},
+        .RequiredAnalyses = {},
+        .VerificationStage = std::nullopt,
+    };
+    Value.ThreadSafety = PassThreadSafety::ThreadSafe;
+    Value.DeclaredChanges = {MachinePassChange::Instructions, MachinePassChange::Operands};
+    return Value;
+  }();
+  return Contract;
+}
+
+[[nodiscard]] Expected<std::uint64_t> alignTo(std::uint64_t Value, std::uint64_t Alignment) {
+  if (Alignment == 0 || (Alignment & (Alignment - 1U)) != 0)
+    return makeUnexpected(ErrorCode::VerificationError, "block alignment is not a power of two");
+  const std::uint64_t Mask = Alignment - 1U;
+  if (Value > std::numeric_limits<std::uint64_t>::max() - Mask)
+    return makeUnexpected(ErrorCode::BranchRelaxationFailed, "machine layout offset overflowed");
+  return (Value + Mask) & ~Mask;
+}
+
+[[nodiscard]] Expected<MachineLayout> computeLayout(MachineFunction &Function,
+                                                    const TargetDescription &Target,
+                                                    const TargetInstrInfo &Info) {
+  MachineLayout Result;
+  std::uint64_t Offset = 0;
+  for (MachineBasicBlock *Block = Function.firstBlock(); Block != nullptr; Block = Block->next()) {
+    auto Aligned = alignTo(Offset, Block->attributes().alignment().bytes());
+    if (!Aligned)
+      return std::unexpected(Aligned.error());
+    Offset = *Aligned;
+    Result.BlockOffsets.emplace(Block, Offset);
+    for (MachineInstr *Instruction = Block->firstInstruction(); Instruction != nullptr;
+         Instruction = Instruction->next()) {
+      Result.InstructionOffsets.emplace(Instruction, Offset);
+      auto Size = Info.instructionSize(*Instruction, Function, Target);
+      if (!Size)
+        return std::unexpected(Size.error());
+      if (Offset > std::numeric_limits<std::uint64_t>::max() - *Size)
+        return makeUnexpected(ErrorCode::BranchRelaxationFailed, "machine layout size overflowed");
+      Offset += *Size;
+    }
+  }
+  return Result;
+}
+
+[[nodiscard]] Expected<std::int64_t> displacement(std::uint64_t Source, std::uint64_t Target) {
+  if (Target >= Source) {
+    const std::uint64_t Difference = Target - Source;
+    if (Difference > static_cast<std::uint64_t>(std::numeric_limits<std::int64_t>::max()))
+      return makeUnexpected(ErrorCode::BranchRelaxationFailed,
+                            "forward branch displacement overflowed");
+    return static_cast<std::int64_t>(Difference);
+  }
+  const std::uint64_t Difference = Source - Target;
+  if (Difference > static_cast<std::uint64_t>(std::numeric_limits<std::int64_t>::max()))
+    return makeUnexpected(ErrorCode::BranchRelaxationFailed,
+                          "backward branch displacement overflowed");
+  return -static_cast<std::int64_t>(Difference);
+}
+
+} // namespace
+
+const MachinePassContract &MachineBranchRelaxationPass::contract() const noexcept {
+  return relaxationContract();
+}
+
+Expected<MachinePassResult>
+MachineBranchRelaxationPass::run(MachineFunction &Function, const TargetDescription &Target,
+                                 MachineFunctionAnalysisManager &Analyses) {
+  return runImplementation(Function, Target, Analyses);
+}
+
+Expected<MachinePassResult>
+MachineBranchRelaxationPass::runImplementation(MachineFunction &Function,
+                                               const TargetDescription &Target,
+                                               MachineFunctionAnalysisManager &Analyses) const {
+  if (Options.ConvergenceWatchdog == 0)
+    return makeUnexpected(ErrorCode::InvalidArgument, "branch relaxation watchdog must be nonzero");
+  MachinePassResult Result;
+  for (std::size_t Iteration = 0; Iteration < Options.ConvergenceWatchdog; ++Iteration) {
+    auto Layout = computeLayout(Function, Target, *Info);
+    if (!Layout)
+      return std::unexpected(Layout.error());
+    bool Changed = false;
+    for (MachineBasicBlock *Block = Function.firstBlock(); Block != nullptr && !Changed;
+         Block = Block->next()) {
+      for (MachineInstr *Instruction = Block->firstInstruction(); Instruction != nullptr;
+           Instruction = Instruction->next()) {
+        auto BranchTarget = Info->branchTarget(*Instruction, Function, Target);
+        if (!BranchTarget)
+          return std::unexpected(BranchTarget.error());
+        if (*BranchTarget == nullptr)
+          continue;
+        const auto SourcePosition = Layout->InstructionOffsets.find(Instruction);
+        const auto TargetPosition = Layout->BlockOffsets.find(*BranchTarget);
+        if (SourcePosition == Layout->InstructionOffsets.end() ||
+            TargetPosition == Layout->BlockOffsets.end()) {
+          return makeUnexpected(ErrorCode::InternalError,
+                                "branch relaxation lost an instruction or target offset");
+        }
+        auto Offset = displacement(SourcePosition->second, TargetPosition->second);
+        if (!Offset)
+          return std::unexpected(Offset.error());
+        if (Info->isBranchOffsetInRange(*Instruction, *Offset))
+          continue;
+        auto Accepted = Analyses.requestAction("relax-branch", Function, Instruction->stableID());
+        if (!Accepted)
+          return std::unexpected(Accepted.error());
+        if (!*Accepted) {
+          return makeUnexpected(ErrorCode::BranchRelaxationFailed,
+                                "required branch relaxation action was rejected");
+        }
+        auto Relaxed = Info->relaxBranch(*Instruction, *Offset, Function, Target);
+        if (!Relaxed)
+          return std::unexpected(Relaxed.error());
+        if (!*Relaxed) {
+          return makeUnexpected(ErrorCode::BranchRelaxationFailed,
+                                "target did not relax an out-of-range branch");
+        }
+        Changed = true;
+        Result.Changed = true;
+        break;
+      }
+    }
+    if (Changed)
+      continue;
+    if (Result.Changed) {
+      Result.Changes = {MachinePassChange::Instructions, MachinePassChange::Operands};
+      Result.Preserved = PreservedAnalyses::none();
+    } else {
+      Result.Preserved = PreservedAnalyses::all();
+    }
+    return Result;
+  }
+  return makeUnexpected(ErrorCode::InternalError,
+                        "branch relaxation failed to reach a monotonic fixed point");
+}
+
+} // namespace rcc