// SPDX-License-Identifier: GPL-3.0-or-later

#include "rcc/CodeGen/MachinePseudoExpansion.h"

#include "rcc/CodeGen/InstructionDescriptor.h"
#include "rcc/CodeGen/LiveVariables.h"
#include "rcc/CodeGen/MachineBasicBlock.h"
#include "rcc/CodeGen/MachineFrameInfo.h"
#include "rcc/CodeGen/MachineFunction.h"
#include "rcc/CodeGen/MachineFunctionCloner.h"
#include "rcc/CodeGen/MachineInstr.h"
#include "rcc/CodeGen/MachineOperand.h"
#include "rcc/CodeGen/MachineRegisterInfo.h"
#include "rcc/CodeGen/TargetDescription.h"
#include "rcc/CodeGen/TargetInstrInfo.h"

#include <cstddef>
#include <cstdint>
#include <limits>
#include <optional>
#include <string>
#include <tuple>
#include <utility>
#include <vector>

namespace rcc {
namespace {

using CFGEdge = std::tuple<std::uint32_t, std::uint32_t, std::uint32_t>;

struct ResourceSnapshot final {
  std::size_t BlockCount = 0;
  std::size_t VirtualRegisterCount = 0;
  std::size_t FrameObjectCount = 0;
  bool FrameFrozen = false;
  std::vector<CFGEdge> Edges;
};

using BlockSizeSnapshot = std::vector<std::pair<MachineBlockStableID, std::uint64_t>>;

[[nodiscard]] Expected<std::uint64_t> blockSize(const MachineBasicBlock &Block,
                                                const MachineFunction &Function,
                                                const TargetDescription &Target,
                                                const TargetInstrInfo &BranchInfo) {
  std::uint64_t Size = 0;
  for (const MachineInstr *Instruction = Block.firstInstruction(); Instruction != nullptr;
       Instruction = Instruction->next()) {
    auto InstructionSize = BranchInfo.instructionSize(*Instruction, Function, Target);
    if (!InstructionSize)
      return std::unexpected(InstructionSize.error());
    if (Size > std::numeric_limits<std::uint64_t>::max() - *InstructionSize) {
      return makeUnexpected(ErrorCode::PseudoExpansionFailed, "target-late pseudo size overflowed");
    }
    Size += *InstructionSize;
  }
  return Size;
}

[[nodiscard]] Expected<BlockSizeSnapshot> snapshotBlockSizes(const MachineFunction &Function,
                                                             const TargetDescription &Target,
                                                             const TargetInstrInfo &BranchInfo) {
  BlockSizeSnapshot Result;
  for (const MachineBasicBlock *Block = Function.firstBlock(); Block != nullptr;
       Block = Block->next()) {
    auto Size = blockSize(*Block, Function, Target, BranchInfo);
    if (!Size)
      return std::unexpected(Size.error());
    Result.emplace_back(Block->stableID(), *Size);
  }
  return Result;
}

[[nodiscard]] ResourceSnapshot snapshotResources(const MachineFunction &Function) {
  ResourceSnapshot Result;
  Result.BlockCount = Function.blockCount();
  Result.VirtualRegisterCount = Function.registerInfo().virtualRegisters().size();
  Result.FrameObjectCount = Function.frameInfo().objects().size();
  Result.FrameFrozen = Function.frameInfo().isFrozen();
  for (const MachineBasicBlock *Block = Function.firstBlock(); Block != nullptr;
       Block = Block->next()) {
    for (const MachineSuccessor &Successor : Block->successors()) {
      Result.Edges.emplace_back(Block->stableID().value(), Successor.Block->stableID().value(),
                                Successor.Probability.rawValue());
    }
  }
  return Result;
}

[[nodiscard]] Expected<void> verifyResourceContract(const ResourceSnapshot &Before,
                                                    const MachineFunction &Function,
                                                    MachinePseudoExpansionPhase Phase) {
  const ResourceSnapshot After = snapshotResources(Function);
  if (Before.BlockCount != After.BlockCount || Before.Edges != After.Edges) {
    return makeUnexpected(ErrorCode::PseudoExpansionFailed,
                          "pseudo expansion changed CFG resources outside its owner contract");
  }
  if (Before.VirtualRegisterCount != After.VirtualRegisterCount) {
    return makeUnexpected(ErrorCode::PseudoExpansionFailed,
                          "post-selection pseudo expansion created a virtual register");
  }
  if (Phase != MachinePseudoExpansionPhase::PrePEI &&
      (Before.FrameObjectCount != After.FrameObjectCount ||
       Before.FrameFrozen != After.FrameFrozen)) {
    return makeUnexpected(ErrorCode::PseudoExpansionFailed,
                          "post-PEI pseudo expansion changed frozen frame resources");
  }
  return {};
}

[[nodiscard]] PseudoOwner ownerFor(MachinePseudoExpansionPhase Phase) noexcept {
  switch (Phase) {
  case MachinePseudoExpansionPhase::PrePEI:
    return PseudoOwner::PrePEIExpansion;
  case MachinePseudoExpansionPhase::PostRA:
    return PseudoOwner::PostRAExpansion;
  case MachinePseudoExpansionPhase::TargetLate:
    return PseudoOwner::TargetLateExpansion;
  }
  return PseudoOwner::None;
}

[[nodiscard]] const MachinePassContract &preparationContract() {
  static const MachinePassContract Contract = [] {
    MachinePassContract Value{
        .Name = "prepare-post-ra-pseudos",
        .Requires = {MachineFunctionProperty::Selected},
        .Sets = {},
        .Clears = {},
        .RequiredAnalyses = {},
        .VerificationStage = PipelineStage::Selected,
    };
    Value.ThreadSafety = PassThreadSafety::ThreadSafe;
    Value.DeclaredChanges = {MachinePassChange::Operands};
    return Value;
  }();
  return Contract;
}

[[nodiscard]] const MachinePassContract &phaseContract(MachinePseudoExpansionPhase Phase) {
  static const MachinePassContract PrePEI = [] {
    MachinePassContract Value{
        .Name = "pre-pei-target-pseudo-expansion",
        .Requires = {MachineFunctionProperty::Selected, MachineFunctionProperty::NoPHIs,
                     MachineFunctionProperty::NoVRegs, MachineFunctionProperty::TracksLiveness},
        .Sets = {},
        .Clears = {},
        .RequiredAnalyses = {},
        .VerificationStage = PipelineStage::RegisterAllocated,
    };
    Value.ThreadSafety = PassThreadSafety::ThreadSafe;
    Value.DeclaredChanges = {MachinePassChange::Instructions, MachinePassChange::Operands,
                             MachinePassChange::Frame};
    return Value;
  }();
  static const MachinePassContract PostRA = [] {
    MachinePassContract Value{
        .Name = "postrapseudos",
        .Requires = {MachineFunctionProperty::Selected, MachineFunctionProperty::NoPHIs,
                     MachineFunctionProperty::NoVRegs, MachineFunctionProperty::TracksLiveness,
                     MachineFunctionProperty::FrameFinalized,
                     MachineFunctionProperty::NoFrameIndices},
        .Sets = {},
        .Clears = {},
        .RequiredAnalyses = {},
        .VerificationStage = PipelineStage::PostRAFinalized,
    };
    Value.ThreadSafety = PassThreadSafety::ThreadSafe;
    Value.DeclaredChanges = {MachinePassChange::Instructions, MachinePassChange::Operands};
    return Value;
  }();
  static const MachinePassContract TargetLate = [] {
    MachinePassContract Value{
        .Name = "target-late-pseudo-expansion",
        .Requires = {MachineFunctionProperty::Selected, MachineFunctionProperty::NoPHIs,
                     MachineFunctionProperty::NoVRegs, MachineFunctionProperty::TracksLiveness,
                     MachineFunctionProperty::FrameFinalized,
                     MachineFunctionProperty::NoFrameIndices},
        .Sets = {},
        .Clears = {},
        .RequiredAnalyses = {},
        .VerificationStage = PipelineStage::PostRAFinalized,
    };
    Value.ThreadSafety = PassThreadSafety::ThreadSafe;
    Value.DeclaredChanges = {MachinePassChange::Instructions, MachinePassChange::Operands};
    return Value;
  }();
  switch (Phase) {
  case MachinePseudoExpansionPhase::PrePEI:
    return PrePEI;
  case MachinePseudoExpansionPhase::PostRA:
    return PostRA;
  case MachinePseudoExpansionPhase::TargetLate:
    return TargetLate;
  }
  return TargetLate;
}

[[nodiscard]] Error prefixedError(const Error &Failure, std::string Prefix) {
  Prefix.append(Failure.message());
  return Error(Failure.code(), std::move(Prefix));
}

[[nodiscard]] std::vector<MachineInstr *> layoutInstructions(MachineFunction &Function) {
  std::vector<MachineInstr *> Result;
  for (MachineBasicBlock *Block = Function.firstBlock(); Block != nullptr; Block = Block->next())
    for (MachineInstr *Instruction = Block->firstInstruction(); Instruction != nullptr;
         Instruction = Instruction->next())
      Result.push_back(Instruction);
  return Result;
}

[[nodiscard]] Expected<MachinePassResult>
prepareImplementation(MachineFunction &Function, const TargetDescription &Target,
                      const TargetPseudoExpansionInfo &Info, MachineActionReplay &Actions) {
  bool Changed = false;
  for (MachineInstr *Instruction : layoutInstructions(Function)) {
    const InstructionDescriptor *Descriptor = Target.instructionDescriptor(Instruction->opcode());
    if (Descriptor == nullptr)
      return makeUnexpected(ErrorCode::VerificationError,
                            "pseudo preparation encountered an unknown opcode");
    if (Descriptor->Owner != PseudoOwner::PostRAExpansion)
      continue;
    std::string Rule = "prepare-pseudo.";
    Rule.append(Descriptor->Name);
    auto Accepted = Actions.request(Rule, Function, Instruction->stableID());
    if (!Accepted)
      return std::unexpected(Accepted.error());
    if (!*Accepted) {
      return makeUnexpected(ErrorCode::PseudoExpansionFailed,
                            "required pseudo preparation action was rejected");
    }
    auto Prepared = Info.preparePostRAPseudo(*Instruction, Function, Target);
    if (!Prepared)
      return std::unexpected(Prepared.error());
    Changed |= *Prepared;
  }
  MachinePassResult Result;
  Result.Changed = Changed;
  if (Changed) {
    Result.Changes = {MachinePassChange::Operands};
    Result.Preserved = PreservedAnalyses::none();
  } else {
    Result.Preserved = PreservedAnalyses::all();
  }
  return Result;
}

[[nodiscard]] Expected<MachinePassResult>
expandImplementation(MachineFunction &Function, const TargetDescription &Target,
                     const TargetPseudoExpansionInfo &Info, MachinePseudoExpansionPhase Phase,
                     const TargetInstrInfo *BranchInfo, MachineActionReplay &Actions) {
  const ResourceSnapshot Before = snapshotResources(Function);
  std::optional<BlockSizeSnapshot> SizesBefore;
  if (Phase == MachinePseudoExpansionPhase::TargetLate) {
    if (BranchInfo == nullptr) {
      return makeUnexpected(ErrorCode::InvalidArgument,
                            "target-late pseudo expansion requires branch sizing");
    }
    auto Snapshot = snapshotBlockSizes(Function, Target, *BranchInfo);
    if (!Snapshot)
      return std::unexpected(Snapshot.error());
    SizesBefore = std::move(*Snapshot);
  }
  std::vector<MachineInstr *> Worklist;
  Worklist = layoutInstructions(Function);
  bool Changed = false;
  for (MachineInstr *Instruction : Worklist) {
    if (Instruction->parent() == nullptr)
      continue;
    const InstructionDescriptor *Descriptor = Target.instructionDescriptor(Instruction->opcode());
    if (Descriptor == nullptr)
      return makeUnexpected(ErrorCode::VerificationError,
                            "pseudo expansion encountered an unknown opcode");
    if (!Descriptor->Flags.has(InstructionFlag::Pseudo) || Descriptor->Owner != ownerFor(Phase)) {
      continue;
    }
    std::string Rule = "expand-pseudo.";
    Rule.append(Descriptor->Name);
    auto Accepted = Actions.request(Rule, Function, Instruction->stableID());
    if (!Accepted)
      return std::unexpected(Accepted.error());
    if (!*Accepted) {
      return makeUnexpected(ErrorCode::PseudoExpansionFailed,
                            "required pseudo expansion action was rejected");
    }
    MachineBasicBlock *ExpansionBlock = Instruction->parent();
    std::optional<std::uint64_t> ExpansionSizeBefore;
    if (Phase == MachinePseudoExpansionPhase::TargetLate) {
      auto Size = blockSize(*ExpansionBlock, Function, Target, *BranchInfo);
      if (!Size)
        return std::unexpected(Size.error());
      ExpansionSizeBefore = *Size;
    }
    auto Expanded = Info.expandPseudo(*Instruction, Phase, Function, Target);
    if (!Expanded)
      return std::unexpected(Expanded.error());
    if (!*Expanded) {
      return makeUnexpected(ErrorCode::PseudoExpansionFailed,
                            "target has no expansion for an owned executable pseudo");
    }
    if (ExpansionSizeBefore.has_value()) {
      auto SizeAfter = blockSize(*ExpansionBlock, Function, Target, *BranchInfo);
      if (!SizeAfter)
        return std::unexpected(SizeAfter.error());
      if (*SizeAfter > *ExpansionSizeBefore) {
        return makeUnexpected(ErrorCode::PseudoExpansionFailed,
                              "one target-late pseudo expansion increased its encoded block size");
      }
    }
    Changed = true;
  }

  auto Resources = verifyResourceContract(Before, Function, Phase);
  if (!Resources)
    return std::unexpected(Resources.error());
  if (SizesBefore.has_value()) {
    auto SizesAfter = snapshotBlockSizes(Function, Target, *BranchInfo);
    if (!SizesAfter)
      return std::unexpected(SizesAfter.error());
    if (SizesAfter->size() != SizesBefore->size()) {
      return makeUnexpected(ErrorCode::PseudoExpansionFailed,
                            "target-late pseudo expansion changed block layout");
    }
    for (std::size_t Index = 0; Index < SizesAfter->size(); ++Index) {
      if ((*SizesAfter)[Index].first != (*SizesBefore)[Index].first ||
          (*SizesAfter)[Index].second > (*SizesBefore)[Index].second) {
        return makeUnexpected(ErrorCode::PseudoExpansionFailed,
                              "target-late pseudo expansion increased an encoded block size");
      }
    }
  }
  for (MachineInstr *Instruction : layoutInstructions(Function)) {
    const InstructionDescriptor *Descriptor = Target.instructionDescriptor(Instruction->opcode());
    if (Descriptor != nullptr && Descriptor->Flags.has(InstructionFlag::Pseudo) &&
        Descriptor->Owner == ownerFor(Phase)) {
      return makeUnexpected(ErrorCode::PseudoExpansionFailed,
                            "pseudo expansion left an instruction owned by the completed phase");
    }
  }
  if (Changed && Function.properties().has(MachineFunctionProperty::TracksLiveness)) {
    auto Rebuilt = rebuildMachineLiveness(Function, Target);
    if (!Rebuilt)
      return std::unexpected(Rebuilt.error());
  }
  MachinePassResult Result;
  Result.Changed = Changed;
  if (Changed) {
    Result.Changes = {MachinePassChange::Instructions, MachinePassChange::Operands};
    Result.Preserved = PreservedAnalyses::none();
  } else {
    Result.Preserved = PreservedAnalyses::all();
  }
  return Result;
}

[[nodiscard]] Expected<MachinePassResult>
runPreparationTransaction(MachineFunction &Function, const TargetDescription &Target,
                          const TargetPseudoExpansionInfo &Info,
                          MachineFunctionAnalysisManager &Analyses) {
  auto Trial = cloneMachineFunction(
      Function, Target,
      MachineFunctionCloneOptions{.NewStableID = Function.stableID(), .NewName = std::nullopt});
  if (!Trial)
    return std::unexpected(prefixedError(Trial.error(), "pseudo preparation setup failed: "));
  MachineActionReplay Actions(Analyses, Trial->InstructionStableIDs);
  auto TrialResult = prepareImplementation(*Trial->ClonedFunction, Target, Info, Actions);
  if (!TrialResult)
    return std::unexpected(TrialResult.error());
  Actions.beginReplay();
  auto Result = prepareImplementation(Function, Target, Info, Actions);
  if (!Result) {
    return std::unexpected(
        prefixedError(Result.error(), "pseudo preparation diverged after successful preflight: "));
  }
  auto Replayed = Actions.finishReplay();
  if (!Replayed)
    return std::unexpected(Replayed.error());
  return Result;
}

[[nodiscard]] Expected<MachinePassResult>
runExpansionTransaction(MachineFunction &Function, const TargetDescription &Target,
                        const TargetPseudoExpansionInfo &Info, MachinePseudoExpansionPhase Phase,
                        const TargetInstrInfo *BranchInfo,
                        MachineFunctionAnalysisManager &Analyses) {
  auto Trial = cloneMachineFunction(
      Function, Target,
      MachineFunctionCloneOptions{.NewStableID = Function.stableID(), .NewName = std::nullopt});
  if (!Trial)
    return std::unexpected(prefixedError(Trial.error(), "pseudo expansion setup failed: "));
  MachineActionReplay Actions(Analyses, Trial->InstructionStableIDs);
  auto TrialResult =
      expandImplementation(*Trial->ClonedFunction, Target, Info, Phase, BranchInfo, Actions);
  if (!TrialResult)
    return std::unexpected(TrialResult.error());
  Actions.beginReplay();
  auto Result = expandImplementation(Function, Target, Info, Phase, BranchInfo, Actions);
  if (!Result) {
    return std::unexpected(
        prefixedError(Result.error(), "pseudo expansion diverged after successful preflight: "));
  }
  auto Replayed = Actions.finishReplay();
  if (!Replayed)
    return std::unexpected(Replayed.error());
  return Result;
}

} // namespace

const MachinePassContract &PreparePostRAPseudosPass::contract() const noexcept {
  return preparationContract();
}

Expected<MachinePassResult>
PreparePostRAPseudosPass::run(MachineFunction &Function, const TargetDescription &Target,
                              MachineFunctionAnalysisManager &Analyses) {
  if (!Function.properties().containsAll(contract().Requires))
    return makeUnexpected(ErrorCode::VerificationError,
                          "post-RA pseudo preparation requires selected MachineIR");
  return runPreparationTransaction(Function, Target, *Info, Analyses);
}

const MachinePassContract &PrePEITargetPseudoExpansionPass::contract() const noexcept {
  return phaseContract(MachinePseudoExpansionPhase::PrePEI);
}

Expected<MachinePassResult>
PrePEITargetPseudoExpansionPass::run(MachineFunction &Function, const TargetDescription &Target,
                                     MachineFunctionAnalysisManager &Analyses) {
  if (!Function.properties().containsAll(contract().Requires))
    return makeUnexpected(ErrorCode::VerificationError,
                          "pre-PEI pseudo expansion requires allocated tracked-liveness MIR");
  return runExpansionTransaction(Function, Target, *Info, MachinePseudoExpansionPhase::PrePEI,
                                 nullptr, Analyses);
}

const MachinePassContract &ExpandPostRAPseudosPass::contract() const noexcept {
  return phaseContract(MachinePseudoExpansionPhase::PostRA);
}

Expected<MachinePassResult> ExpandPostRAPseudosPass::run(MachineFunction &Function,
                                                         const TargetDescription &Target,
                                                         MachineFunctionAnalysisManager &Analyses) {
  if (!Function.properties().containsAll(contract().Requires))
    return makeUnexpected(ErrorCode::VerificationError,
                          "post-RA pseudo expansion requires finalized frame and liveness");
  return runExpansionTransaction(Function, Target, *Info, MachinePseudoExpansionPhase::PostRA,
                                 nullptr, Analyses);
}

const MachinePassContract &TargetLatePseudoExpansionPass::contract() const noexcept {
  return phaseContract(MachinePseudoExpansionPhase::TargetLate);
}

Expected<MachinePassResult>
TargetLatePseudoExpansionPass::run(MachineFunction &Function, const TargetDescription &Target,
                                   MachineFunctionAnalysisManager &Analyses) {
  if (!Function.properties().containsAll(contract().Requires))
    return makeUnexpected(ErrorCode::VerificationError,
                          "target-late pseudo expansion requires finalized allocated MIR");
  return runExpansionTransaction(Function, Target, *Info, MachinePseudoExpansionPhase::TargetLate,
                                 BranchInfo, Analyses);
}

} // namespace rcc