//===- DXContainerEmitter.cpp - Convert YAML to a DXContainer -------------===//
//
// 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
//
//===----------------------------------------------------------------------===//
///
/// \file
/// Binary emitter for yaml to DXContainer binary
///
//===----------------------------------------------------------------------===//

#include "llvm/BinaryFormat/DXContainer.h"
#include "llvm/MC/DXContainerPSVInfo.h"
#include "llvm/MC/DXContainerRootSignature.h"
#include "llvm/ObjectYAML/ObjectYAML.h"
#include "llvm/ObjectYAML/yaml2obj.h"
#include "llvm/Support/Errc.h"
#include "llvm/Support/Error.h"
#include "llvm/Support/raw_ostream.h"

using namespace llvm;

namespace {
class DXContainerWriter {
public:
  DXContainerWriter(DXContainerYAML::Object &ObjectFile)
      : ObjectFile(ObjectFile) {}

  Error write(raw_ostream &OS);

private:
  DXContainerYAML::Object &ObjectFile;

  Error computePartOffsets();
  Error validatePartOffsets();
  Error validateSize(uint32_t Computed);

  void writeHeader(raw_ostream &OS);
  void writeParts(raw_ostream &OS);
};
} // namespace

Error DXContainerWriter::validateSize(uint32_t Computed) {
  if (!ObjectFile.Header.FileSize)
    ObjectFile.Header.FileSize = Computed;
  else if (*ObjectFile.Header.FileSize < Computed)
    return createStringError(errc::result_out_of_range,
                             "File size specified is too small.");
  return Error::success();
}

Error DXContainerWriter::validatePartOffsets() {
  if (ObjectFile.Parts.size() != ObjectFile.Header.PartOffsets->size())
    return createStringError(
        errc::invalid_argument,
        "Mismatch between number of parts and part offsets.");
  uint32_t RollingOffset =
      sizeof(dxbc::Header) + (ObjectFile.Header.PartCount * sizeof(uint32_t));
  for (auto I : llvm::zip(ObjectFile.Parts, *ObjectFile.Header.PartOffsets)) {
    if (RollingOffset > std::get<1>(I))
      return createStringError(errc::invalid_argument,
                               "Offset mismatch, not enough space for data.");
    RollingOffset =
        std::get<1>(I) + sizeof(dxbc::PartHeader) + std::get<0>(I).Size;
  }
  if (Error Err = validateSize(RollingOffset))
    return Err;

  return Error::success();
}

Error DXContainerWriter::computePartOffsets() {
  if (ObjectFile.Header.PartOffsets)
    return validatePartOffsets();
  uint32_t RollingOffset =
      sizeof(dxbc::Header) + (ObjectFile.Header.PartCount * sizeof(uint32_t));
  ObjectFile.Header.PartOffsets = std::vector<uint32_t>();
  for (const auto &Part : ObjectFile.Parts) {
    ObjectFile.Header.PartOffsets->push_back(RollingOffset);
    RollingOffset += sizeof(dxbc::PartHeader) + Part.Size;
  }
  if (Error Err = validateSize(RollingOffset))
    return Err;

  return Error::success();
}

void DXContainerWriter::writeHeader(raw_ostream &OS) {
  dxbc::Header Header;
  memcpy(Header.Magic, "DXBC", 4);
  memcpy(Header.FileHash.Digest, ObjectFile.Header.Hash.data(), 16);
  Header.Version.Major = ObjectFile.Header.Version.Major;
  Header.Version.Minor = ObjectFile.Header.Version.Minor;
  Header.FileSize = *ObjectFile.Header.FileSize;
  Header.PartCount = ObjectFile.Parts.size();
  if (sys::IsBigEndianHost)
    Header.swapBytes();
  OS.write(reinterpret_cast<char *>(&Header), sizeof(Header));
  SmallVector<uint32_t> Offsets(ObjectFile.Header.PartOffsets->begin(),
                                ObjectFile.Header.PartOffsets->end());
  if (sys::IsBigEndianHost)
    for (auto &O : Offsets)
      sys::swapByteOrder(O);
  OS.write(reinterpret_cast<char *>(Offsets.data()),
           Offsets.size() * sizeof(uint32_t));
}

void DXContainerWriter::writeParts(raw_ostream &OS) {
  uint32_t RollingOffset =
      sizeof(dxbc::Header) + (ObjectFile.Header.PartCount * sizeof(uint32_t));
  for (auto I : llvm::zip(ObjectFile.Parts, *ObjectFile.Header.PartOffsets)) {
    if (RollingOffset < std::get<1>(I)) {
      uint32_t PadBytes = std::get<1>(I) - RollingOffset;
      OS.write_zeros(PadBytes);
    }
    DXContainerYAML::Part P = std::get<0>(I);
    RollingOffset = std::get<1>(I) + sizeof(dxbc::PartHeader);
    uint32_t PartSize = P.Size;

    OS.write(P.Name.c_str(), 4);
    if (sys::IsBigEndianHost)
      sys::swapByteOrder(P.Size);
    OS.write(reinterpret_cast<const char *>(&P.Size), sizeof(uint32_t));

    dxbc::PartType PT = dxbc::parsePartType(P.Name);

    uint64_t DataStart = OS.tell();
    switch (PT) {
    case dxbc::PartType::DXIL: {
      if (!P.Program)
        continue;
      dxbc::ProgramHeader Header;
      Header.Version = dxbc::ProgramHeader::getVersion(P.Program->MajorVersion,
                                                       P.Program->MinorVersion);
      Header.Unused = 0;
      Header.ShaderKind = P.Program->ShaderKind;
      memcpy(Header.Bitcode.Magic, "DXIL", 4);
      Header.Bitcode.MajorVersion = P.Program->DXILMajorVersion;
      Header.Bitcode.MinorVersion = P.Program->DXILMinorVersion;
      Header.Bitcode.Unused = 0;

      // Compute the optional fields if needed...
      if (P.Program->DXILOffset)
        Header.Bitcode.Offset = *P.Program->DXILOffset;
      else
        Header.Bitcode.Offset = sizeof(dxbc::BitcodeHeader);

      if (P.Program->DXILSize)
        Header.Bitcode.Size = *P.Program->DXILSize;
      else
        Header.Bitcode.Size = P.Program->DXIL ? P.Program->DXIL->size() : 0;

      if (P.Program->Size)
        Header.Size = *P.Program->Size;
      else
        Header.Size = sizeof(dxbc::ProgramHeader) + Header.Bitcode.Size;

      uint32_t BitcodeOffset = Header.Bitcode.Offset;
      if (sys::IsBigEndianHost)
        Header.swapBytes();
      OS.write(reinterpret_cast<const char *>(&Header),
               sizeof(dxbc::ProgramHeader));
      if (P.Program->DXIL) {
        if (BitcodeOffset > sizeof(dxbc::BitcodeHeader)) {
          uint32_t PadBytes = BitcodeOffset - sizeof(dxbc::BitcodeHeader);
          OS.write_zeros(PadBytes);
        }
        OS.write(reinterpret_cast<char *>(P.Program->DXIL->data()),
                 P.Program->DXIL->size());
      }
      break;
    }
    case dxbc::PartType::SFI0: {
      // If we don't have any flags we can continue here and the data will be
      // zeroed out.
      if (!P.Flags.has_value())
        continue;
      uint64_t Flags = P.Flags->getEncodedFlags();
      if (sys::IsBigEndianHost)
        sys::swapByteOrder(Flags);
      OS.write(reinterpret_cast<char *>(&Flags), sizeof(uint64_t));
      break;
    }
    case dxbc::PartType::HASH: {
      if (!P.Hash.has_value())
        continue;
      dxbc::ShaderHash Hash = {0, {0}};
      if (P.Hash->IncludesSource)
        Hash.Flags |= static_cast<uint32_t>(dxbc::HashFlags::IncludesSource);
      memcpy(&Hash.Digest[0], &P.Hash->Digest[0], 16);
      if (sys::IsBigEndianHost)
        Hash.swapBytes();
      OS.write(reinterpret_cast<char *>(&Hash), sizeof(dxbc::ShaderHash));
      break;
    }
    case dxbc::PartType::PSV0: {
      if (!P.Info.has_value())
        continue;
      mcdxbc::PSVRuntimeInfo PSV;
      memcpy(&PSV.BaseData, &P.Info->Info, sizeof(dxbc::PSV::v3::RuntimeInfo));
      PSV.Resources = P.Info->Resources;
      PSV.EntryName = P.Info->EntryName;

      for (auto El : P.Info->SigInputElements)
        PSV.InputElements.push_back(mcdxbc::PSVSignatureElement{
            El.Name, El.Indices, El.StartRow, El.Cols, El.StartCol,
            El.Allocated, El.Kind, El.Type, El.Mode, El.DynamicMask,
            El.Stream});

      for (auto El : P.Info->SigOutputElements)
        PSV.OutputElements.push_back(mcdxbc::PSVSignatureElement{
            El.Name, El.Indices, El.StartRow, El.Cols, El.StartCol,
            El.Allocated, El.Kind, El.Type, El.Mode, El.DynamicMask,
            El.Stream});

      for (auto El : P.Info->SigPatchOrPrimElements)
        PSV.PatchOrPrimElements.push_back(mcdxbc::PSVSignatureElement{
            El.Name, El.Indices, El.StartRow, El.Cols, El.StartCol,
            El.Allocated, El.Kind, El.Type, El.Mode, El.DynamicMask,
            El.Stream});

      static_assert(PSV.OutputVectorMasks.size() == PSV.InputOutputMap.size());
      for (unsigned I = 0; I < PSV.OutputVectorMasks.size(); ++I) {
        PSV.OutputVectorMasks[I].insert(PSV.OutputVectorMasks[I].begin(),
                                        P.Info->OutputVectorMasks[I].begin(),
                                        P.Info->OutputVectorMasks[I].end());
        PSV.InputOutputMap[I].insert(PSV.InputOutputMap[I].begin(),
                                     P.Info->InputOutputMap[I].begin(),
                                     P.Info->InputOutputMap[I].end());
      }

      PSV.PatchOrPrimMasks.insert(PSV.PatchOrPrimMasks.begin(),
                                  P.Info->PatchOrPrimMasks.begin(),
                                  P.Info->PatchOrPrimMasks.end());
      PSV.InputPatchMap.insert(PSV.InputPatchMap.begin(),
                               P.Info->InputPatchMap.begin(),
                               P.Info->InputPatchMap.end());
      PSV.PatchOutputMap.insert(PSV.PatchOutputMap.begin(),
                                P.Info->PatchOutputMap.begin(),
                                P.Info->PatchOutputMap.end());

      PSV.finalize(static_cast<Triple::EnvironmentType>(
          Triple::Pixel + P.Info->Info.ShaderStage));
      PSV.write(OS, P.Info->Version);
      break;
    }
    case dxbc::PartType::ISG1:
    case dxbc::PartType::OSG1:
    case dxbc::PartType::PSG1: {
      mcdxbc::Signature Sig;
      if (P.Signature.has_value()) {
        for (const auto &Param : P.Signature->Parameters) {
          Sig.addParam(Param.Stream, Param.Name, Param.Index, Param.SystemValue,
                       Param.CompType, Param.Register, Param.Mask,
                       Param.ExclusiveMask, Param.MinPrecision);
        }
      }
      Sig.write(OS);
      break;
    }
    case dxbc::PartType::Unknown:
      break; // Skip any handling for unrecognized parts.
    case dxbc::PartType::RTS0:
      if (!P.RootSignature.has_value())
        continue;

      mcdxbc::RootSignatureDesc RS;
      RS.Flags = P.RootSignature->getEncodedFlags();
      RS.Version = P.RootSignature->Version;
      RS.RootParameterOffset = P.RootSignature->RootParametersOffset;
      RS.NumStaticSamplers = P.RootSignature->NumStaticSamplers;
      RS.StaticSamplersOffset = P.RootSignature->StaticSamplersOffset;

      for (DXContainerYAML::RootParameterLocationYaml &L :
           P.RootSignature->Parameters.Locations) {
        dxbc::RTS0::v1::RootParameterHeader Header{L.Header.Type, L.Header.Visibility,
                                         L.Header.Offset};

        switch (L.Header.Type) {
        case llvm::to_underlying(dxbc::RootParameterType::Constants32Bit): {
          const DXContainerYAML::RootConstantsYaml &ConstantYaml =
              P.RootSignature->Parameters.getOrInsertConstants(L);
          dxbc::RTS0::v1::RootConstants Constants;
          Constants.Num32BitValues = ConstantYaml.Num32BitValues;
          Constants.RegisterSpace = ConstantYaml.RegisterSpace;
          Constants.ShaderRegister = ConstantYaml.ShaderRegister;
          RS.ParametersContainer.addParameter(Header, Constants);
          break;
        }
        case llvm::to_underlying(dxbc::RootParameterType::CBV):
        case llvm::to_underlying(dxbc::RootParameterType::SRV):
        case llvm::to_underlying(dxbc::RootParameterType::UAV): {
          const DXContainerYAML::RootDescriptorYaml &DescriptorYaml =
              P.RootSignature->Parameters.getOrInsertDescriptor(L);

          dxbc::RTS0::v2::RootDescriptor Descriptor;
          Descriptor.RegisterSpace = DescriptorYaml.RegisterSpace;
          Descriptor.ShaderRegister = DescriptorYaml.ShaderRegister;
          if (RS.Version > 1)
            Descriptor.Flags = DescriptorYaml.getEncodedFlags();
          RS.ParametersContainer.addParameter(Header, Descriptor);
          break;
        }
        case llvm::to_underlying(dxbc::RootParameterType::DescriptorTable): {
          const DXContainerYAML::DescriptorTableYaml &TableYaml =
              P.RootSignature->Parameters.getOrInsertTable(L);
          mcdxbc::DescriptorTable Table;
          for (const auto &R : TableYaml.Ranges) {

            dxbc::RTS0::v2::DescriptorRange Range;
            Range.RangeType = R.RangeType;
            Range.NumDescriptors = R.NumDescriptors;
            Range.BaseShaderRegister = R.BaseShaderRegister;
            Range.RegisterSpace = R.RegisterSpace;
            Range.OffsetInDescriptorsFromTableStart =
                R.OffsetInDescriptorsFromTableStart;
            if (RS.Version > 1)
              Range.Flags = R.getEncodedFlags();
            Table.Ranges.push_back(Range);
          }
          RS.ParametersContainer.addParameter(Header, Table);
          break;
        }
        default:
          // Handling invalid parameter type edge case. We intentionally let
          // obj2yaml/yaml2obj parse and emit invalid dxcontainer data, in order
          // for that to be used as a testing tool more effectively.
          RS.ParametersContainer.addInvalidParameter(Header);
        }
      }

      for (const auto &Param : P.RootSignature->samplers()) {
        dxbc::RTS0::v1::StaticSampler NewSampler;
        NewSampler.Filter = Param.Filter;
        NewSampler.AddressU = Param.AddressU;
        NewSampler.AddressV = Param.AddressV;
        NewSampler.AddressW = Param.AddressW;
        NewSampler.MipLODBias = Param.MipLODBias;
        NewSampler.MaxAnisotropy = Param.MaxAnisotropy;
        NewSampler.ComparisonFunc = Param.ComparisonFunc;
        NewSampler.BorderColor = Param.BorderColor;
        NewSampler.MinLOD = Param.MinLOD;
        NewSampler.MaxLOD = Param.MaxLOD;
        NewSampler.ShaderRegister = Param.ShaderRegister;
        NewSampler.RegisterSpace = Param.RegisterSpace;
        NewSampler.ShaderVisibility = Param.ShaderVisibility;

        RS.StaticSamplers.push_back(NewSampler);
      }

      RS.write(OS);
      break;
    }
    uint64_t BytesWritten = OS.tell() - DataStart;
    RollingOffset += BytesWritten;
    if (BytesWritten < PartSize)
      OS.write_zeros(PartSize - BytesWritten);
    RollingOffset += PartSize;
  }
}

Error DXContainerWriter::write(raw_ostream &OS) {
  if (Error Err = computePartOffsets())
    return Err;
  writeHeader(OS);
  writeParts(OS);
  return Error::success();
}

namespace llvm {
namespace yaml {

bool yaml2dxcontainer(DXContainerYAML::Object &Doc, raw_ostream &Out,
                      ErrorHandler EH) {
  DXContainerWriter Writer(Doc);
  if (Error Err = Writer.write(Out)) {
    handleAllErrors(std::move(Err),
                    [&](const ErrorInfoBase &Err) { EH(Err.message()); });
    return false;
  }
  return true;
}

} // namespace yaml
} // namespace llvm