[CodeGen] Unify LLVM branch folding pipeline

4 files changed, 304 insertions, 424 deletions

diff --git a/backend/lib/CodeGen/CMakeLists.txt b/backend/lib/CodeGen/CMakeLists.txt
index 862caca..881e17e 100644
--- a/backend/lib/CodeGen/CMakeLists.txt
+++ b/backend/lib/CodeGen/CMakeLists.txt
@@ -55,7 +55,6 @@ set(
   MachineSink.cpp
   MachineStackSlotColoring.cpp
   MachineTailDuplication.cpp
-  MachineTailMerge.cpp
   PrologueEpilogueInsertion.cpp
   PostRAScheduler.cpp
   RegisterPressure.cpp
@@ -152,7 +151,6 @@ set(
   "${RCC_BACKEND_SOURCE_DIR}/include/rcc/CodeGen/MachineSink.h"
   "${RCC_BACKEND_SOURCE_DIR}/include/rcc/CodeGen/MachineStackSlotColoring.h"
   "${RCC_BACKEND_SOURCE_DIR}/include/rcc/CodeGen/MachineTailDuplication.h"
-  "${RCC_BACKEND_SOURCE_DIR}/include/rcc/CodeGen/MachineTailMerge.h"
   "${RCC_BACKEND_SOURCE_DIR}/include/rcc/CodeGen/PostRAScheduler.h"
   "${RCC_BACKEND_SOURCE_DIR}/include/rcc/CodeGen/PBQP/Graph.h"
   "${RCC_BACKEND_SOURCE_DIR}/include/rcc/CodeGen/PBQP/Math.h"
diff --git a/backend/lib/CodeGen/MachineBranchOptimization.cpp b/backend/lib/CodeGen/MachineBranchOptimization.cpp
index d09de97..07a67ed 100644
--- a/backend/lib/CodeGen/MachineBranchOptimization.cpp
+++ b/backend/lib/CodeGen/MachineBranchOptimization.cpp
@@ -12,11 +12,15 @@
 #include "rcc/CodeGen/MachineFunctionRecords.h"
 #include "rcc/CodeGen/MachineInstr.h"
 #include "rcc/CodeGen/MachineJumpTableInfo.h"
+#include "rcc/CodeGen/MachineMemOperand.h"
+#include "rcc/CodeGen/MachineOperand.h"
+#include "rcc/CodeGen/RegisterMask.h"
 #include "rcc/CodeGen/TargetDescription.h"
 #include "rcc/CodeGen/TargetInstrInfo.h"
 
 #include <algorithm>
 #include <cstddef>
+#include <cstdint>
 #include <optional>
 #include <ranges>
 #include <string>
@@ -31,11 +35,11 @@ namespace {
   static const MachinePassContract Contract = [] {
     MachinePassContract Value{
         .Name = "machine-branch-optimization",
-        .Requires = {MachineFunctionProperty::Selected, MachineFunctionProperty::NoPHIs},
+        .Requires = {MachineFunctionProperty::NoPHIs},
         .Sets = {},
         .Clears = {},
         .RequiredAnalyses = {},
-        .VerificationStage = PipelineStage::Selected,
+        .VerificationStage = PipelineStage::FrameFinalized,
     };
     Value.ThreadSafety = PassThreadSafety::ThreadSafe;
     Value.DeclaredChanges = {MachinePassChange::Instructions, MachinePassChange::Operands,
@@ -80,6 +84,215 @@ namespace {
          Attributes.profileID().has_value() || isJumpTableTarget(Block, Function);
 }
 
+template <typename ValueT>
+[[nodiscard]] bool payloadEquals(const MachineOperand &Left, const MachineOperand &Right) noexcept {
+  const ValueT *LeftValue = Left.getIf<ValueT>();
+  const ValueT *RightValue = Right.getIf<ValueT>();
+  return LeftValue != nullptr && RightValue != nullptr && *LeftValue == *RightValue;
+}
+
+template <typename MaskOperandT>
+[[nodiscard]] bool maskEquals(const MachineOperand &Left, const MachineOperand &Right) noexcept {
+  const MaskOperandT *LeftValue = Left.getIf<MaskOperandT>();
+  const MaskOperandT *RightValue = Right.getIf<MaskOperandT>();
+  if (LeftValue == nullptr || RightValue == nullptr)
+    return false;
+  if (!LeftValue->Mask || !RightValue->Mask)
+    return LeftValue->Mask == RightValue->Mask;
+  return *LeftValue->Mask == *RightValue->Mask;
+}
+
+[[nodiscard]] bool operandEquals(const MachineOperand &Left, const MachineOperand &Right) noexcept {
+  if (Left.kind() != Right.kind() || Left.targetFlags() != Right.targetFlags() ||
+      Left.tiedOperandIndex() != Right.tiedOperandIndex()) {
+    return false;
+  }
+  if (Left.kind() == MachineOperand::Kind::Register) {
+    const Register *LeftReg = Left.getIf<Register>();
+    const Register *RightReg = Right.getIf<Register>();
+    if (LeftReg == nullptr || RightReg == nullptr || *LeftReg != *RightReg ||
+        Left.subReg() != Right.subReg()) {
+      return false;
+    }
+    RegisterOperandFlags LeftFlags = Left.registerFlags();
+    RegisterOperandFlags RightFlags = Right.registerFlags();
+    for (const RegisterOperandFlag Flag :
+         {RegisterOperandFlag::Kill, RegisterOperandFlag::Dead, RegisterOperandFlag::Renamable}) {
+      LeftFlags.clear(Flag);
+      RightFlags.clear(Flag);
+    }
+    return LeftFlags == RightFlags;
+  }
+  switch (Left.kind()) {
+  case MachineOperand::Kind::Register:
+    return false;
+  case MachineOperand::Kind::Immediate:
+    return payloadEquals<std::int64_t>(Left, Right);
+  case MachineOperand::Kind::IntegerImmediate:
+    return payloadEquals<IntegerImmediateOperand>(Left, Right);
+  case MachineOperand::Kind::FloatingImmediate:
+    return payloadEquals<FloatingImmediateOperand>(Left, Right);
+  case MachineOperand::Kind::BasicBlock:
+    return payloadEquals<BasicBlockOperand>(Left, Right);
+  case MachineOperand::Kind::FrameIndex:
+    return payloadEquals<FrameIndexOperand>(Left, Right);
+  case MachineOperand::Kind::ConstantPoolIndex:
+    return payloadEquals<ConstantPoolOperand>(Left, Right);
+  case MachineOperand::Kind::JumpTableIndex:
+    return payloadEquals<JumpTableOperand>(Left, Right);
+  case MachineOperand::Kind::GlobalSymbol:
+    return payloadEquals<GlobalSymbolOperand>(Left, Right);
+  case MachineOperand::Kind::ExternalSymbol:
+    return payloadEquals<ExternalSymbolOperand>(Left, Right);
+  case MachineOperand::Kind::BlockAddress:
+    return payloadEquals<BlockAddressOperand>(Left, Right);
+  case MachineOperand::Kind::RegisterMask:
+    return maskEquals<RegisterMaskOperand>(Left, Right);
+  case MachineOperand::Kind::LiveOutRegisterMask:
+    return maskEquals<LiveOutRegisterMaskOperand>(Left, Right);
+  case MachineOperand::Kind::Metadata:
+    return payloadEquals<MetadataOperand>(Left, Right);
+  case MachineOperand::Kind::TargetIndex:
+    return payloadEquals<TargetIndexOperand>(Left, Right);
+  case MachineOperand::Kind::MCSymbol:
+    return payloadEquals<MCSymbolOperand>(Left, Right);
+  case MachineOperand::Kind::CFIIndex:
+    return payloadEquals<CFIIndexOperand>(Left, Right);
+  case MachineOperand::Kind::Intrinsic:
+    return payloadEquals<IntrinsicOperand>(Left, Right);
+  case MachineOperand::Kind::Predicate:
+    return payloadEquals<PredicateOperand>(Left, Right);
+  case MachineOperand::Kind::ShuffleMask:
+    return payloadEquals<ShuffleMaskOperand>(Left, Right);
+  case MachineOperand::Kind::DebugInstrRef:
+    return payloadEquals<DebugInstrRefOperand>(Left, Right);
+  case MachineOperand::Kind::LaneMask:
+    return payloadEquals<LaneMaskOperand>(Left, Right);
+  case MachineOperand::Kind::Count:
+    return false;
+  }
+  return false;
+}
+
+[[nodiscard]] Expected<bool> instructionEquals(const MachineInstr &Left, const MachineInstr &Right,
+                                               const TargetDescription &Target) {
+  const InstructionDescriptor *Descriptor = Target.instructionDescriptor(Left.opcode());
+  if (Descriptor == nullptr || Target.instructionDescriptor(Right.opcode()) == nullptr) {
+    return makeUnexpected(ErrorCode::VerificationError,
+                          "branch folding encountered an unknown instruction");
+  }
+  if (Left.opcode() != Right.opcode() || Left.flags() != Right.flags() ||
+      Left.operandCount() != Right.operandCount() ||
+      Left.memOperands().size() != Right.memOperands().size() ||
+      Left.sourceLocation() != Right.sourceLocation() ||
+      Left.debugLocation() != Right.debugLocation() || Left.origin() != Right.origin() ||
+      hasProtectedAttachment(Left) || hasProtectedAttachment(Right) ||
+      Descriptor->Flags.has(InstructionFlag::Call) ||
+      Descriptor->Flags.has(InstructionFlag::Convergent) ||
+      Descriptor->Flags.has(InstructionFlag::NotDuplicable) ||
+      Left.flags().has(MachineInstrFlag::NoMerge) ||
+      Left.flags().has(MachineInstrFlag::BundledWithPred) ||
+      Left.flags().has(MachineInstrFlag::BundledWithSucc)) {
+    return false;
+  }
+  for (std::size_t Index = 0; Index < Left.operandCount(); ++Index) {
+    const MachineOperand *LeftOperand = Left.operandAt(Index);
+    const MachineOperand *RightOperand = Right.operandAt(Index);
+    if (LeftOperand == nullptr || RightOperand == nullptr ||
+        !operandEquals(*LeftOperand, *RightOperand)) {
+      return false;
+    }
+  }
+  for (std::size_t Index = 0; Index < Left.memOperands().size(); ++Index) {
+    const MachineMemOperand *LeftMemory = Left.memOperands()[Index];
+    const MachineMemOperand *RightMemory = Right.memOperands()[Index];
+    if (LeftMemory == nullptr || RightMemory == nullptr ||
+        LeftMemory->description() != RightMemory->description()) {
+      return false;
+    }
+  }
+  return true;
+}
+
+[[nodiscard]] Expected<bool> blockEquals(const MachineBasicBlock &Left,
+                                         const MachineBasicBlock &Right,
+                                         const MachineFunction &Function,
+                                         const TargetDescription &Target) {
+  if (&Left == &Right || &Left == Function.firstBlock() || &Right == Function.firstBlock() ||
+      hasObservableBlockIdentity(Left, Function) || hasObservableBlockIdentity(Right, Function) ||
+      Left.attributes().alignment() != Right.attributes().alignment() ||
+      Left.attributes().sectionID() != Right.attributes().sectionID() ||
+      Left.attributes().callFrameSize() != Right.attributes().callFrameSize() ||
+      !std::ranges::equal(Left.liveIns(), Right.liveIns()) ||
+      !std::ranges::equal(Left.successors(), Right.successors()) ||
+      Left.instructionCount() != Right.instructionCount()) {
+    return false;
+  }
+  const MachineInstr *LeftInstruction = Left.firstInstruction();
+  const MachineInstr *RightInstruction = Right.firstInstruction();
+  while (LeftInstruction != nullptr && RightInstruction != nullptr) {
+    auto Equal = instructionEquals(*LeftInstruction, *RightInstruction, Target);
+    if (!Equal)
+      return std::unexpected(Equal.error());
+    if (!*Equal)
+      return false;
+    LeftInstruction = LeftInstruction->next();
+    RightInstruction = RightInstruction->next();
+  }
+  return LeftInstruction == nullptr && RightInstruction == nullptr;
+}
+
+[[nodiscard]] Expected<bool> canRedirectPredecessor(MachineBasicBlock &Predecessor,
+                                                    MachineBasicBlock &Duplicate,
+                                                    MachineFunction &Function,
+                                                    const TargetDescription &Target,
+                                                    const TargetInstrInfo &Info) {
+  auto Analysis = Info.analyzeBranch(Predecessor, Function, Target);
+  if (!Analysis)
+    return std::unexpected(Analysis.error());
+  if (Analysis->HasUnanalyzableTerminator)
+    return false;
+  if ((Analysis->Conditional != nullptr && Analysis->TrueTarget == &Duplicate &&
+       hasProtectedAttachment(*Analysis->Conditional)) ||
+      (Analysis->Unconditional != nullptr && Analysis->TrueTarget == &Duplicate &&
+       hasProtectedAttachment(*Analysis->Unconditional))) {
+    return false;
+  }
+  if (Analysis->TrueTarget == &Duplicate || Analysis->FalseTarget == &Duplicate)
+    return true;
+  return Analysis->Conditional == nullptr && Analysis->Unconditional == nullptr &&
+         Predecessor.successors().size() == 1 && Predecessor.next() == &Duplicate;
+}
+
+[[nodiscard]] Expected<bool>
+redirectPredecessor(MachineBasicBlock &Predecessor, MachineBasicBlock &Duplicate,
+                    MachineBasicBlock &Canonical, MachineFunction &Function,
+                    const TargetDescription &Target, const TargetInstrInfo &Info) {
+  auto Analysis = Info.analyzeBranch(Predecessor, Function, Target);
+  if (!Analysis)
+    return std::unexpected(Analysis.error());
+  auto Rewritten = Info.rewriteBranchTarget(Predecessor, Duplicate, Canonical, Function, Target);
+  if (!Rewritten)
+    return std::unexpected(Rewritten.error());
+  if (!*Rewritten) {
+    const bool IsImplicitEdge =
+        (Analysis->Conditional == nullptr && Analysis->Unconditional == nullptr &&
+         Predecessor.successors().size() == 1 && Predecessor.next() == &Duplicate) ||
+        (Analysis->Conditional != nullptr && Analysis->Unconditional == nullptr &&
+         Analysis->FalseTarget == &Duplicate);
+    if (!IsImplicitEdge)
+      return false;
+    auto Inserted = Info.insertBranch(Predecessor, Canonical, Function, Target);
+    if (!Inserted)
+      return std::unexpected(Inserted.error());
+  }
+  MachineCFGUpdater CFG = Function.cfgUpdater();
+  auto Redirected = CFG.replaceSuccessor(Predecessor, Duplicate, Canonical);
+  if (!Redirected)
+    return std::unexpected(Redirected.error());
+  return true;
+}
+
 [[nodiscard]] bool canDeleteBlock(const MachineBasicBlock &Block,
                                   const MachineFunction &Function) noexcept {
   if (hasObservableBlockIdentity(Block, Function))
@@ -288,6 +501,77 @@ namespace {
   return AnyChange;
 }
 
+[[nodiscard]] Expected<bool> mergeIdenticalTails(MachineFunction &Function,
+                                                 const TargetDescription &Target,
+                                                 const TargetInstrInfo &Info,
+                                                 std::size_t MaximumMergedBlocks) {
+  std::size_t MergedBlocks = 0;
+  bool Changed = false;
+  bool LocalChange = true;
+  while (LocalChange && MergedBlocks < MaximumMergedBlocks) {
+    LocalChange = false;
+    std::vector<MachineBasicBlock *> Blocks;
+    for (MachineBasicBlock *Block = Function.firstBlock(); Block != nullptr; Block = Block->next())
+      Blocks.push_back(Block);
+    for (std::size_t DuplicateIndex = 1; DuplicateIndex < Blocks.size(); ++DuplicateIndex) {
+      MachineBasicBlock *Duplicate = Blocks[DuplicateIndex];
+      if (Duplicate == nullptr || !Duplicate->isLinked())
+        continue;
+      for (std::size_t CanonicalIndex = 1; CanonicalIndex < DuplicateIndex; ++CanonicalIndex) {
+        MachineBasicBlock *Canonical = Blocks[CanonicalIndex];
+        if (Canonical == nullptr || !Canonical->isLinked())
+          continue;
+        auto Equal = blockEquals(*Canonical, *Duplicate, Function, Target);
+        if (!Equal)
+          return std::unexpected(Equal.error());
+        if (!*Equal)
+          continue;
+
+        const std::vector<MachineBasicBlock *> Predecessors(Duplicate->predecessors().begin(),
+                                                            Duplicate->predecessors().end());
+        bool RedirectedAll = true;
+        for (MachineBasicBlock *Predecessor : Predecessors) {
+          if (Predecessor == nullptr || Predecessor == Canonical || Predecessor == Duplicate) {
+            RedirectedAll = false;
+            break;
+          }
+          auto CanRedirect =
+              canRedirectPredecessor(*Predecessor, *Duplicate, Function, Target, Info);
+          if (!CanRedirect)
+            return std::unexpected(CanRedirect.error());
+          if (!*CanRedirect) {
+            RedirectedAll = false;
+            break;
+          }
+        }
+        if (!RedirectedAll)
+          continue;
+        for (MachineBasicBlock *Predecessor : Predecessors) {
+          auto Redirected =
+              redirectPredecessor(*Predecessor, *Duplicate, *Canonical, Function, Target, Info);
+          if (!Redirected)
+            return std::unexpected(Redirected.error());
+          if (!*Redirected) {
+            return makeUnexpected(ErrorCode::CFGOptimizationFailed,
+                                  "tail merge predecessor changed after redirect preflight");
+          }
+        }
+        MachineCFGUpdater CFG = Function.cfgUpdater();
+        auto Deleted = CFG.deleteBlock(*Duplicate);
+        if (!Deleted)
+          return std::unexpected(Deleted.error());
+        ++MergedBlocks;
+        Changed = true;
+        LocalChange = true;
+        break;
+      }
+      if (LocalChange)
+        break;
+    }
+  }
+  return Changed;
+}
+
 } // namespace
 
 const MachinePassContract &MachineBranchOptimizationPass::contract() const noexcept {
@@ -299,7 +583,7 @@ Expected<MachinePassResult> MachineBranchOptimizationPass::run(MachineFunction &
                                                                MachineFunctionAnalysisManager &) {
   if (!Function.properties().containsAll(contract().Requires)) {
     return makeUnexpected(ErrorCode::VerificationError,
-                          "branch optimization requires selected PHI-free MachineIR");
+                          "branch folding requires PHI-free MachineIR");
   }
   auto Trial = cloneMachineFunction(
       Function, Target,
@@ -331,12 +615,16 @@ MachineBranchOptimizationPass::runImplementation(MachineFunction &Function,
   auto Threaded = threadForwardingBlocks(Function, Target, *Info, Options.MaximumThreadedEdges);
   if (!Threaded)
     return std::unexpected(Threaded.error());
+  auto Merged = mergeIdenticalTails(Function, Target, *Info, Options.MaximumMergedBlocks);
+  if (!Merged)
+    return std::unexpected(Merged.error());
   auto RemovedAfterThreading = removeUnreachableBlocks(Function);
   if (!RemovedAfterThreading)
     return std::unexpected(RemovedAfterThreading.error());
 
   MachinePassResult Result;
-  Result.Changed = *Folded || *RemovedBeforeThreading || *Threaded || *RemovedAfterThreading;
+  Result.Changed =
+      *Folded || *RemovedBeforeThreading || *Threaded || *Merged || *RemovedAfterThreading;
   if (!Result.Changed) {
     Result.Preserved = PreservedAnalyses::all();
     return Result;
diff --git a/backend/lib/CodeGen/MachineOptimizationPipeline.cpp b/backend/lib/CodeGen/MachineOptimizationPipeline.cpp
index 817b2ca..7634ecd 100644
--- a/backend/lib/CodeGen/MachineOptimizationPipeline.cpp
+++ b/backend/lib/CodeGen/MachineOptimizationPipeline.cpp
@@ -21,7 +21,6 @@
 #include "rcc/CodeGen/MachineSink.h"
 #include "rcc/CodeGen/MachineStackSlotColoring.h"
 #include "rcc/CodeGen/MachineTailDuplication.h"
-#include "rcc/CodeGen/MachineTailMerge.h"
 #include "rcc/CodeGen/PostRAScheduler.h"
 #include "rcc/CodeGen/PrologueEpilogueInsertion.h"
 
@@ -69,11 +68,8 @@ addSelectedSSAOptimizationPipeline(MachineFunctionPassManager &Pipeline,
 
 Expected<void> addPreRegisterAllocationOptimizationPipeline(
     MachineFunctionPassManager &Pipeline, const MachineOptimizationPipelineTargetInfo &TargetInfo) {
-  auto Added = addPass<MachineBranchOptimizationPass>(Pipeline, TargetInfo.InstrInfo);
-  if (!Added)
-    return std::unexpected(Added.error());
-  Added = addPass<MachineCombinerPass>(Pipeline, TargetInfo.Combiner,
-                                       MachineCombinerStage::PreRegisterAllocation);
+  auto Added = addPass<MachineCombinerPass>(Pipeline, TargetInfo.Combiner,
+                                            MachineCombinerStage::PreRegisterAllocation);
   if (!Added)
     return std::unexpected(Added.error());
   return addPass<MachineSchedulerPass>(Pipeline, TargetInfo.SchedulingModel);
@@ -102,28 +98,28 @@ addPostRegisterAllocationPipeline(MachineFunctionPassManager &Pipeline,
   Added = addPass<PrologueEpilogueInsertionPass>(Pipeline, TargetInfo.FrameLowering);
   if (!Added)
     return std::unexpected(Added.error());
-  Added = addPass<ExpandPostRAPseudosPass>(Pipeline, TargetInfo.PseudoExpansion);
+  Added = addPass<MachineBranchOptimizationPass>(Pipeline, TargetInfo.InstrInfo);
   if (!Added)
     return std::unexpected(Added.error());
-  Added = addPass<PostRASchedulerPass>(Pipeline, TargetInfo.SchedulingModel);
+  Added = addPass<MachineTailDuplicationPass>(Pipeline, TargetInfo.InstrInfo);
   if (!Added)
     return std::unexpected(Added.error());
-  Added = addPass<MachineTailDuplicationPass>(Pipeline, TargetInfo.InstrInfo);
+  Added = addPass<MachineCopyPropagationPass>(
+      Pipeline, TargetInfo.InstrInfo,
+      MachineCopyPropagationOptions{.UseTargetCopyInstructions =
+                                        Options.UseTargetCopyInstructions});
   if (!Added)
     return std::unexpected(Added.error());
-  Added = addPass<MachineTailMergePass>(Pipeline, TargetInfo.InstrInfo);
+  Added = addPass<ExpandPostRAPseudosPass>(Pipeline, TargetInfo.PseudoExpansion);
   if (!Added)
     return std::unexpected(Added.error());
-  Added = addPass<MachineBlockPlacementPass>(Pipeline);
+  Added = addPass<PostRASchedulerPass>(Pipeline, TargetInfo.SchedulingModel);
   if (!Added)
     return std::unexpected(Added.error());
-  Added = addPass<MachineFallthroughPass>(Pipeline, TargetInfo.InstrInfo);
+  Added = addPass<MachineBlockPlacementPass>(Pipeline);
   if (!Added)
     return std::unexpected(Added.error());
-  Added = addPass<MachineCopyPropagationPass>(
-      Pipeline, TargetInfo.InstrInfo,
-      MachineCopyPropagationOptions{.UseTargetCopyInstructions =
-                                        Options.UseTargetCopyInstructions});
+  Added = addPass<MachineFallthroughPass>(Pipeline, TargetInfo.InstrInfo);
   if (!Added)
     return std::unexpected(Added.error());
   Added = addPass<MachineBranchRelaxationPass>(Pipeline, TargetInfo.InstrInfo);
diff --git a/backend/lib/CodeGen/MachineTailMerge.cpp b/backend/lib/CodeGen/MachineTailMerge.cpp
deleted file mode 100644
index 7467993..0000000
--- a/backend/lib/CodeGen/MachineTailMerge.cpp
+++ /dev/null
@@ -1,402 +0,0 @@
-// SPDX-License-Identifier: GPL-3.0-or-later
-
-#include "rcc/CodeGen/MachineTailMerge.h"
-
-#include "rcc/CodeGen/InstructionDescriptor.h"
-#include "rcc/CodeGen/LiveVariables.h"
-#include "rcc/CodeGen/MachineBasicBlock.h"
-#include "rcc/CodeGen/MachineBasicBlockAttributes.h"
-#include "rcc/CodeGen/MachineCFGUpdater.h"
-#include "rcc/CodeGen/MachineFunction.h"
-#include "rcc/CodeGen/MachineFunctionCloner.h"
-#include "rcc/CodeGen/MachineInstr.h"
-#include "rcc/CodeGen/MachineJumpTableInfo.h"
-#include "rcc/CodeGen/MachineMemOperand.h"
-#include "rcc/CodeGen/MachineOperand.h"
-#include "rcc/CodeGen/RegisterMask.h"
-#include "rcc/CodeGen/TargetDescription.h"
-#include "rcc/CodeGen/TargetInstrInfo.h"
-
-#include <algorithm>
-#include <cstddef>
-#include <optional>
-#include <ranges>
-#include <string>
-#include <utility>
-#include <vector>
-
-namespace rcc {
-namespace {
-
-[[nodiscard]] const MachinePassContract &tailMergeContract() {
-  static const MachinePassContract Contract = [] {
-    MachinePassContract Value{
-        .Name = "machine-tail-merge",
-        .Requires = {MachineFunctionProperty::Selected, MachineFunctionProperty::NoPHIs,
-                     MachineFunctionProperty::TiedOpsRewritten, MachineFunctionProperty::NoVRegs,
-                     MachineFunctionProperty::TracksLiveness,
-                     MachineFunctionProperty::FrameFinalized,
-                     MachineFunctionProperty::NoFrameIndices},
-        .Sets = {},
-        .Clears = {},
-        .RequiredAnalyses = {},
-        .VerificationStage = PipelineStage::FrameFinalized,
-    };
-    Value.ThreadSafety = PassThreadSafety::ThreadSafe;
-    Value.DeclaredChanges = {MachinePassChange::Instructions, MachinePassChange::Operands,
-                             MachinePassChange::CFG, MachinePassChange::BlockLayout};
-    return Value;
-  }();
-  return Contract;
-}
-
-[[nodiscard]] bool hasProtectedAttachment(const MachineInstr &Instruction) noexcept {
-  return Instruction.debugInstrNumber() != 0 || Instruction.preInstrSymbol().has_value() ||
-         Instruction.postInstrSymbol().has_value() || Instruction.heapAllocMarker().has_value() ||
-         Instruction.pcSections().has_value() || Instruction.mmra().has_value() ||
-         Instruction.cfiType().has_value() || Instruction.deactivationSymbol().has_value();
-}
-
-[[nodiscard]] bool isJumpTableTarget(const MachineBasicBlock &Block,
-                                     const MachineFunction &Function) noexcept {
-  for (const MachineJumpTableEntry &Entry : Function.jumpTableInfo().tables()) {
-    const MachineJumpTableDesc *Description = Entry.description();
-    if (Description != nullptr &&
-        std::ranges::find(Description->Targets, &Block) != Description->Targets.end()) {
-      return true;
-    }
-  }
-  return false;
-}
-
-[[nodiscard]] bool hasProtectedIdentity(const MachineBasicBlock &Block,
-                                        const MachineFunction &Function) noexcept {
-  const MachineBasicBlockAttributes &Attributes = Block.attributes();
-  return Attributes.rawFlags() != 0 || !Attributes.sectionID().isDefault() ||
-         Attributes.profileID().has_value() || isJumpTableTarget(Block, Function);
-}
-
-template <typename ValueT>
-[[nodiscard]] bool payloadEquals(const MachineOperand &Left, const MachineOperand &Right) noexcept {
-  const ValueT *LeftValue = Left.getIf<ValueT>();
-  const ValueT *RightValue = Right.getIf<ValueT>();
-  return LeftValue != nullptr && RightValue != nullptr && *LeftValue == *RightValue;
-}
-
-template <typename MaskOperandT>
-[[nodiscard]] bool maskEquals(const MachineOperand &Left, const MachineOperand &Right) noexcept {
-  const MaskOperandT *LeftValue = Left.getIf<MaskOperandT>();
-  const MaskOperandT *RightValue = Right.getIf<MaskOperandT>();
-  if (LeftValue == nullptr || RightValue == nullptr)
-    return false;
-  if (!LeftValue->Mask || !RightValue->Mask)
-    return LeftValue->Mask == RightValue->Mask;
-  return *LeftValue->Mask == *RightValue->Mask;
-}
-
-[[nodiscard]] bool operandEquals(const MachineOperand &Left, const MachineOperand &Right) noexcept {
-  if (Left.kind() != Right.kind() || Left.targetFlags() != Right.targetFlags() ||
-      Left.tiedOperandIndex() != Right.tiedOperandIndex()) {
-    return false;
-  }
-  if (Left.kind() == MachineOperand::Kind::Register) {
-    const Register *LeftReg = Left.getIf<Register>();
-    const Register *RightReg = Right.getIf<Register>();
-    if (LeftReg == nullptr || RightReg == nullptr || *LeftReg != *RightReg ||
-        Left.subReg() != Right.subReg()) {
-      return false;
-    }
-    RegisterOperandFlags LeftFlags = Left.registerFlags();
-    RegisterOperandFlags RightFlags = Right.registerFlags();
-    for (const RegisterOperandFlag Flag :
-         {RegisterOperandFlag::Kill, RegisterOperandFlag::Dead, RegisterOperandFlag::Renamable}) {
-      LeftFlags.clear(Flag);
-      RightFlags.clear(Flag);
-    }
-    return LeftFlags == RightFlags;
-  }
-  switch (Left.kind()) {
-  case MachineOperand::Kind::Register:
-    return false;
-  case MachineOperand::Kind::Immediate:
-    return payloadEquals<std::int64_t>(Left, Right);
-  case MachineOperand::Kind::IntegerImmediate:
-    return payloadEquals<IntegerImmediateOperand>(Left, Right);
-  case MachineOperand::Kind::FloatingImmediate:
-    return payloadEquals<FloatingImmediateOperand>(Left, Right);
-  case MachineOperand::Kind::BasicBlock:
-    return payloadEquals<BasicBlockOperand>(Left, Right);
-  case MachineOperand::Kind::FrameIndex:
-    return payloadEquals<FrameIndexOperand>(Left, Right);
-  case MachineOperand::Kind::ConstantPoolIndex:
-    return payloadEquals<ConstantPoolOperand>(Left, Right);
-  case MachineOperand::Kind::JumpTableIndex:
-    return payloadEquals<JumpTableOperand>(Left, Right);
-  case MachineOperand::Kind::GlobalSymbol:
-    return payloadEquals<GlobalSymbolOperand>(Left, Right);
-  case MachineOperand::Kind::ExternalSymbol:
-    return payloadEquals<ExternalSymbolOperand>(Left, Right);
-  case MachineOperand::Kind::BlockAddress:
-    return payloadEquals<BlockAddressOperand>(Left, Right);
-  case MachineOperand::Kind::RegisterMask:
-    return maskEquals<RegisterMaskOperand>(Left, Right);
-  case MachineOperand::Kind::LiveOutRegisterMask:
-    return maskEquals<LiveOutRegisterMaskOperand>(Left, Right);
-  case MachineOperand::Kind::Metadata:
-    return payloadEquals<MetadataOperand>(Left, Right);
-  case MachineOperand::Kind::TargetIndex:
-    return payloadEquals<TargetIndexOperand>(Left, Right);
-  case MachineOperand::Kind::MCSymbol:
-    return payloadEquals<MCSymbolOperand>(Left, Right);
-  case MachineOperand::Kind::CFIIndex:
-    return payloadEquals<CFIIndexOperand>(Left, Right);
-  case MachineOperand::Kind::Intrinsic:
-    return payloadEquals<IntrinsicOperand>(Left, Right);
-  case MachineOperand::Kind::Predicate:
-    return payloadEquals<PredicateOperand>(Left, Right);
-  case MachineOperand::Kind::ShuffleMask:
-    return payloadEquals<ShuffleMaskOperand>(Left, Right);
-  case MachineOperand::Kind::DebugInstrRef:
-    return payloadEquals<DebugInstrRefOperand>(Left, Right);
-  case MachineOperand::Kind::LaneMask:
-    return payloadEquals<LaneMaskOperand>(Left, Right);
-  case MachineOperand::Kind::Count:
-    return false;
-  }
-  return false;
-}
-
-[[nodiscard]] Expected<bool> instructionEquals(const MachineInstr &Left, const MachineInstr &Right,
-                                               const TargetDescription &Target) {
-  const InstructionDescriptor *Descriptor = Target.instructionDescriptor(Left.opcode());
-  if (Descriptor == nullptr || Target.instructionDescriptor(Right.opcode()) == nullptr) {
-    return makeUnexpected(ErrorCode::VerificationError,
-                          "tail merge encountered an unknown instruction");
-  }
-  if (Left.opcode() != Right.opcode() || Left.flags() != Right.flags() ||
-      Left.operandCount() != Right.operandCount() ||
-      Left.memOperands().size() != Right.memOperands().size() ||
-      Left.sourceLocation() != Right.sourceLocation() ||
-      Left.debugLocation() != Right.debugLocation() || Left.origin() != Right.origin() ||
-      hasProtectedAttachment(Left) || hasProtectedAttachment(Right) ||
-      Descriptor->Flags.has(InstructionFlag::Call) ||
-      Descriptor->Flags.has(InstructionFlag::Convergent) ||
-      Descriptor->Flags.has(InstructionFlag::NotDuplicable) ||
-      Left.flags().has(MachineInstrFlag::NoMerge) ||
-      Left.flags().has(MachineInstrFlag::BundledWithPred) ||
-      Left.flags().has(MachineInstrFlag::BundledWithSucc)) {
-    return false;
-  }
-  for (std::size_t Index = 0; Index < Left.operandCount(); ++Index) {
-    const MachineOperand *LeftOperand = Left.operandAt(Index);
-    const MachineOperand *RightOperand = Right.operandAt(Index);
-    if (LeftOperand == nullptr || RightOperand == nullptr ||
-        !operandEquals(*LeftOperand, *RightOperand)) {
-      return false;
-    }
-  }
-  for (std::size_t Index = 0; Index < Left.memOperands().size(); ++Index) {
-    const MachineMemOperand *LeftMemory = Left.memOperands()[Index];
-    const MachineMemOperand *RightMemory = Right.memOperands()[Index];
-    if (LeftMemory == nullptr || RightMemory == nullptr ||
-        LeftMemory->description() != RightMemory->description()) {
-      return false;
-    }
-  }
-  return true;
-}
-
-[[nodiscard]] Expected<bool> blockEquals(const MachineBasicBlock &Left,
-                                         const MachineBasicBlock &Right,
-                                         const MachineFunction &Function,
-                                         const TargetDescription &Target) {
-  if (&Left == &Right || &Left == Function.firstBlock() || &Right == Function.firstBlock() ||
-      hasProtectedIdentity(Left, Function) || hasProtectedIdentity(Right, Function) ||
-      Left.attributes().alignment() != Right.attributes().alignment() ||
-      Left.attributes().sectionID() != Right.attributes().sectionID() ||
-      Left.attributes().callFrameSize() != Right.attributes().callFrameSize() ||
-      !std::ranges::equal(Left.liveIns(), Right.liveIns()) ||
-      !std::ranges::equal(Left.successors(), Right.successors()) ||
-      Left.instructionCount() != Right.instructionCount()) {
-    return false;
-  }
-  const MachineInstr *LeftInstruction = Left.firstInstruction();
-  const MachineInstr *RightInstruction = Right.firstInstruction();
-  while (LeftInstruction != nullptr && RightInstruction != nullptr) {
-    auto Equal = instructionEquals(*LeftInstruction, *RightInstruction, Target);
-    if (!Equal)
-      return std::unexpected(Equal.error());
-    if (!*Equal)
-      return false;
-    LeftInstruction = LeftInstruction->next();
-    RightInstruction = RightInstruction->next();
-  }
-  return LeftInstruction == nullptr && RightInstruction == nullptr;
-}
-
-[[nodiscard]] Expected<bool> canRedirectPredecessor(MachineBasicBlock &Predecessor,
-                                                    MachineBasicBlock &Duplicate,
-                                                    MachineFunction &Function,
-                                                    const TargetDescription &Target,
-                                                    const TargetInstrInfo &Info) {
-  auto Analysis = Info.analyzeBranch(Predecessor, Function, Target);
-  if (!Analysis)
-    return std::unexpected(Analysis.error());
-  if (Analysis->HasUnanalyzableTerminator)
-    return false;
-  if ((Analysis->Conditional != nullptr && Analysis->TrueTarget == &Duplicate &&
-       hasProtectedAttachment(*Analysis->Conditional)) ||
-      (Analysis->Unconditional != nullptr && Analysis->TrueTarget == &Duplicate &&
-       hasProtectedAttachment(*Analysis->Unconditional))) {
-    return false;
-  }
-  if (Analysis->TrueTarget == &Duplicate || Analysis->FalseTarget == &Duplicate)
-    return true;
-  return Analysis->Conditional == nullptr && Analysis->Unconditional == nullptr &&
-         Predecessor.successors().size() == 1 && Predecessor.next() == &Duplicate;
-}
-
-[[nodiscard]] Expected<bool>
-redirectPredecessor(MachineBasicBlock &Predecessor, MachineBasicBlock &Duplicate,
-                    MachineBasicBlock &Canonical, MachineFunction &Function,
-                    const TargetDescription &Target, const TargetInstrInfo &Info) {
-  auto Analysis = Info.analyzeBranch(Predecessor, Function, Target);
-  if (!Analysis)
-    return std::unexpected(Analysis.error());
-  auto Rewritten = Info.rewriteBranchTarget(Predecessor, Duplicate, Canonical, Function, Target);
-  if (!Rewritten)
-    return std::unexpected(Rewritten.error());
-  if (!*Rewritten) {
-    const bool IsImplicitEdge =
-        (Analysis->Conditional == nullptr && Analysis->Unconditional == nullptr &&
-         Predecessor.successors().size() == 1 && Predecessor.next() == &Duplicate) ||
-        (Analysis->Conditional != nullptr && Analysis->Unconditional == nullptr &&
-         Analysis->FalseTarget == &Duplicate);
-    if (!IsImplicitEdge)
-      return false;
-    auto Inserted = Info.insertBranch(Predecessor, Canonical, Function, Target);
-    if (!Inserted)
-      return std::unexpected(Inserted.error());
-  }
-  MachineCFGUpdater CFG = Function.cfgUpdater();
-  auto Redirected = CFG.replaceSuccessor(Predecessor, Duplicate, Canonical);
-  if (!Redirected)
-    return std::unexpected(Redirected.error());
-  return true;
-}
-
-} // namespace
-
-const MachinePassContract &MachineTailMergePass::contract() const noexcept {
-  return tailMergeContract();
-}
-
-Expected<MachinePassResult> MachineTailMergePass::run(MachineFunction &Function,
-                                                      const TargetDescription &Target,
-                                                      MachineFunctionAnalysisManager &) {
-  if (!Function.properties().containsAll(contract().Requires)) {
-    return makeUnexpected(ErrorCode::VerificationError,
-                          "tail merge requires finalized allocated tracked-liveness MIR");
-  }
-  auto Trial = cloneMachineFunction(
-      Function, Target,
-      MachineFunctionCloneOptions{.NewStableID = Function.stableID(), .NewName = std::nullopt});
-  if (!Trial) {
-    return std::unexpected(
-        withErrorContext(Trial.error(), "tail merge transaction setup failed: "));
-  }
-  auto TrialResult = runImplementation(*Trial->ClonedFunction, Target);
-  if (!TrialResult)
-    return std::unexpected(TrialResult.error());
-  auto Result = runImplementation(Function, Target);
-  if (!Result) {
-    return std::unexpected(withErrorContext(
-        Result.error(), "tail merge diverged after a successful transaction preflight: "));
-  }
-  return Result;
-}
-
-Expected<MachinePassResult>
-MachineTailMergePass::runImplementation(MachineFunction &Function,
-                                        const TargetDescription &Target) const {
-  MachinePassResult Result;
-  std::size_t MergedBlocks = 0;
-  bool LocalChange = true;
-  while (LocalChange && MergedBlocks < Options.MaximumMergedBlocks) {
-    LocalChange = false;
-    std::vector<MachineBasicBlock *> Blocks;
-    for (MachineBasicBlock *Block = Function.firstBlock(); Block != nullptr;
-         Block = Block->next()) {
-      Blocks.push_back(Block);
-    }
-    for (std::size_t DuplicateIndex = 1; DuplicateIndex < Blocks.size(); ++DuplicateIndex) {
-      MachineBasicBlock *Duplicate = Blocks[DuplicateIndex];
-      if (Duplicate == nullptr || !Duplicate->isLinked())
-        continue;
-      for (std::size_t CanonicalIndex = 1; CanonicalIndex < DuplicateIndex; ++CanonicalIndex) {
-        MachineBasicBlock *Canonical = Blocks[CanonicalIndex];
-        if (Canonical == nullptr || !Canonical->isLinked())
-          continue;
-        auto Equal = blockEquals(*Canonical, *Duplicate, Function, Target);
-        if (!Equal)
-          return std::unexpected(Equal.error());
-        if (!*Equal)
-          continue;
-
-        const std::vector<MachineBasicBlock *> Predecessors(Duplicate->predecessors().begin(),
-                                                            Duplicate->predecessors().end());
-        bool RedirectedAll = true;
-        for (MachineBasicBlock *Predecessor : Predecessors) {
-          if (Predecessor == nullptr || Predecessor == Canonical || Predecessor == Duplicate) {
-            RedirectedAll = false;
-            break;
-          }
-          auto CanRedirect =
-              canRedirectPredecessor(*Predecessor, *Duplicate, Function, Target, *Info);
-          if (!CanRedirect)
-            return std::unexpected(CanRedirect.error());
-          if (!*CanRedirect) {
-            RedirectedAll = false;
-            break;
-          }
-        }
-        if (!RedirectedAll)
-          continue;
-        for (MachineBasicBlock *Predecessor : Predecessors) {
-          auto Redirected =
-              redirectPredecessor(*Predecessor, *Duplicate, *Canonical, Function, Target, *Info);
-          if (!Redirected)
-            return std::unexpected(Redirected.error());
-          if (!*Redirected) {
-            return makeUnexpected(ErrorCode::CFGOptimizationFailed,
-                                  "tail merge predecessor changed after redirect preflight");
-          }
-        }
-        MachineCFGUpdater CFG = Function.cfgUpdater();
-        auto Deleted = CFG.deleteBlock(*Duplicate);
-        if (!Deleted)
-          return std::unexpected(Deleted.error());
-        ++MergedBlocks;
-        Result.Changed = true;
-        LocalChange = true;
-        break;
-      }
-      if (LocalChange)
-        break;
-    }
-  }
-
-  if (!Result.Changed) {
-    Result.Preserved = PreservedAnalyses::all();
-    return Result;
-  }
-  auto Rebuilt = rebuildMachineLiveness(Function, Target);
-  if (!Rebuilt)
-    return std::unexpected(Rebuilt.error());
-  Result.Changes = {MachinePassChange::Instructions, MachinePassChange::Operands,
-                    MachinePassChange::CFG, MachinePassChange::BlockLayout};
-  Result.Preserved = PreservedAnalyses::none();
-  return Result;
-}
-
-} // namespace rcc