//===-- llvm/CodeGen/MachineModuleInfo.cpp ----------------------*- C++ -*-===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//

#include "llvm/CodeGen/MachineModuleSlotTracker.h"
#include "llvm/CodeGen/MachineFunction.h"
#include "llvm/CodeGen/MachineInstr.h"
#include "llvm/CodeGen/MachineModuleInfo.h"
#include "llvm/CodeGen/MachineOperand.h"
#include "llvm/IR/DebugInfoMetadata.h"
#include "llvm/IR/Module.h"

using namespace llvm;

void MachineModuleSlotTracker::processMachineFunctionMetadata(
    AbstractSlotTrackerStorage *AST, const MachineFunction &MF) {
  // Create metadata created within the backend.
  for (const MachineBasicBlock &MBB : MF)
    for (const MachineInstr &MI : MBB.instrs()) {
      if (MDNode *N = MI.getHeapAllocMarker())
        AST->createMetadataSlot(N);
      if (MDNode *N = MI.getPCSections())
        AST->createMetadataSlot(N);
      if (MDNode *N = MI.getMMRAMetadata())
        AST->createMetadataSlot(N);

      for (const MachineOperand &MO : MI.operands())
        if (MO.isMetadata())
          AST->createMetadataSlot(MO.getMetadata());

      for (const MachineMemOperand *MMO : MI.memoperands()) {
        AAMDNodes AAInfo = MMO->getAAInfo();
        if (AAInfo.TBAA)
          AST->createMetadataSlot(AAInfo.TBAA);
        if (AAInfo.TBAAStruct)
          AST->createMetadataSlot(AAInfo.TBAAStruct);
        if (AAInfo.Scope)
          AST->createMetadataSlot(AAInfo.Scope);
        if (AAInfo.NoAlias)
          AST->createMetadataSlot(AAInfo.NoAlias);
        if (AAInfo.NoAliasAddrSpace)
          AST->createMetadataSlot(AAInfo.NoAliasAddrSpace);
        if (const MDNode *N = MMO->getRanges())
          AST->createMetadataSlot(N);
        if (const MDNode *N = MMO->getMemCacheHint())
          AST->createMetadataSlot(N);
      }
    }

  for (const MachineFunction::VariableDbgInfo &DebugVar :
       MF.getVariableDbgInfo()) {
    AST->createMetadataSlot(DebugVar.Var);
  }
}

void MachineModuleSlotTracker::processMachineModule(
    AbstractSlotTrackerStorage *AST, const Module *M,
    bool ShouldInitializeAllMetadata) {
  if (ShouldInitializeAllMetadata) {
    for (const Function &F : *M) {
      if (&F != &TheFunction)
        continue;
      MDNStartSlot = AST->getNextMetadataSlot();
      if (TheMF)
        processMachineFunctionMetadata(AST, *TheMF);
      MDNEndSlot = AST->getNextMetadataSlot();
      break;
    }
  }
}

void MachineModuleSlotTracker::processMachineFunction(
    AbstractSlotTrackerStorage *AST, const Function *F,
    bool ShouldInitializeAllMetadata) {
  if (!ShouldInitializeAllMetadata && F == &TheFunction) {
    MDNStartSlot = AST->getNextMetadataSlot();
    if (TheMF)
      processMachineFunctionMetadata(AST, *TheMF);
    MDNEndSlot = AST->getNextMetadataSlot();
  }
}

void MachineModuleSlotTracker::collectMachineMDNodes(
    MachineMDNodeListType &L) const {
  collectMDNodes(L, MDNStartSlot, MDNEndSlot);
}

MachineModuleSlotTracker::MachineModuleSlotTracker(
    MFGetterFnT Fn, const MachineFunction *MF, bool ShouldInitializeAllMetadata)
    : ModuleSlotTracker(MF->getFunction().getParent(),
                        ShouldInitializeAllMetadata),
      TheFunction(MF->getFunction()), TheMF(Fn(MF->getFunction())) {
  setProcessHook([this](AbstractSlotTrackerStorage *AST, const Module *M,
                        bool ShouldInitializeAllMetadata) {
    this->processMachineModule(AST, M, ShouldInitializeAllMetadata);
  });
  setProcessHook([this](AbstractSlotTrackerStorage *AST, const Function *F,
                        bool ShouldInitializeAllMetadata) {
    this->processMachineFunction(AST, F, ShouldInitializeAllMetadata);
  });
}

MachineModuleSlotTracker::~MachineModuleSlotTracker() = default;