///////////////////////////////////////////////////////////////////////////////
//                                                                           //
// DxilRootSignatureConvert.cpp                                              //
// Copyright (C) Microsoft Corporation. All rights reserved.                 //
// This file is distributed under the University of Illinois Open Source     //
// License. See LICENSE.TXT for details.                                     //
//                                                                           //
// Convert root signature structures.                                        //
//                                                                           //
///////////////////////////////////////////////////////////////////////////////

#include "dxc/DXIL/DxilConstants.h"
#include "dxc/DxilRootSignature/DxilRootSignature.h"
#include "dxc/Support/FileIOHelper.h"
#include "dxc/Support/Global.h"
#include "dxc/Support/WinFunctions.h"
#include "dxc/Support/WinIncludes.h"
#include "dxc/dxcapi.h"

#include "llvm/IR/DiagnosticPrinter.h"
#include "llvm/Support/raw_ostream.h"

#include <algorithm>
#include <set>
#include <string>
#include <utility>
#include <vector>

#include "DxilRootSignatureHelper.h"

using namespace llvm;
using std::string;

namespace hlsl {

using namespace root_sig_helper;

//////////////////////////////////////////////////////////////////////////////

template <typename IN_DXIL_ROOT_SIGNATURE_DESC,
          typename OUT_DXIL_ROOT_SIGNATURE_DESC,
          typename OUT_DXIL_ROOT_PARAMETER, typename OUT_DXIL_ROOT_DESCRIPTOR,
          typename OUT_DXIL_DESCRIPTOR_RANGE>
void ConvertRootSignatureTemplate(const IN_DXIL_ROOT_SIGNATURE_DESC &DescIn,
                                  DxilRootSignatureVersion DescVersionOut,
                                  OUT_DXIL_ROOT_SIGNATURE_DESC &DescOut) {
  const IN_DXIL_ROOT_SIGNATURE_DESC *pDescIn = &DescIn;
  OUT_DXIL_ROOT_SIGNATURE_DESC *pDescOut = &DescOut;

  // Root signature descriptor.
  pDescOut->Flags = pDescIn->Flags;
  pDescOut->NumParameters = 0;
  pDescOut->NumStaticSamplers = 0;
  // Intialize all pointers early so that clean up works properly.
  pDescOut->pParameters = nullptr;
  pDescOut->pStaticSamplers = nullptr;

  // Root signature parameters.
  if (pDescIn->NumParameters > 0) {
    pDescOut->pParameters = new OUT_DXIL_ROOT_PARAMETER[pDescIn->NumParameters];
    pDescOut->NumParameters = pDescIn->NumParameters;
    memset((void *)pDescOut->pParameters, 0,
           pDescOut->NumParameters * sizeof(OUT_DXIL_ROOT_PARAMETER));
  }

  for (unsigned iRP = 0; iRP < pDescIn->NumParameters; iRP++) {
    const auto &ParamIn = pDescIn->pParameters[iRP];
    OUT_DXIL_ROOT_PARAMETER &ParamOut =
        (OUT_DXIL_ROOT_PARAMETER &)pDescOut->pParameters[iRP];

    ParamOut.ParameterType = ParamIn.ParameterType;
    ParamOut.ShaderVisibility = ParamIn.ShaderVisibility;

    switch (ParamIn.ParameterType) {
    case DxilRootParameterType::DescriptorTable: {
      ParamOut.DescriptorTable.pDescriptorRanges = nullptr;
      unsigned NumRanges = ParamIn.DescriptorTable.NumDescriptorRanges;
      if (NumRanges > 0) {
        ParamOut.DescriptorTable.pDescriptorRanges =
            new OUT_DXIL_DESCRIPTOR_RANGE[NumRanges];
        ParamOut.DescriptorTable.NumDescriptorRanges = NumRanges;
      }

      for (unsigned i = 0; i < NumRanges; i++) {
        const auto &RangeIn = ParamIn.DescriptorTable.pDescriptorRanges[i];
        OUT_DXIL_DESCRIPTOR_RANGE &RangeOut =
            (OUT_DXIL_DESCRIPTOR_RANGE &)
                ParamOut.DescriptorTable.pDescriptorRanges[i];

        RangeOut.RangeType = RangeIn.RangeType;
        RangeOut.NumDescriptors = RangeIn.NumDescriptors;
        RangeOut.BaseShaderRegister = RangeIn.BaseShaderRegister;
        RangeOut.RegisterSpace = RangeIn.RegisterSpace;
        RangeOut.OffsetInDescriptorsFromTableStart =
            RangeIn.OffsetInDescriptorsFromTableStart;
        DxilDescriptorRangeFlags Flags = GetFlags(RangeIn);
        SetFlags(RangeOut, Flags);
      }
      break;
    }
    case DxilRootParameterType::Constants32Bit: {
      ParamOut.Constants.Num32BitValues = ParamIn.Constants.Num32BitValues;
      ParamOut.Constants.ShaderRegister = ParamIn.Constants.ShaderRegister;
      ParamOut.Constants.RegisterSpace = ParamIn.Constants.RegisterSpace;
      break;
    }
    case DxilRootParameterType::CBV:
    case DxilRootParameterType::SRV:
    case DxilRootParameterType::UAV: {
      ParamOut.Descriptor.ShaderRegister = ParamIn.Descriptor.ShaderRegister;
      ParamOut.Descriptor.RegisterSpace = ParamIn.Descriptor.RegisterSpace;
      DxilRootDescriptorFlags Flags = GetFlags(ParamIn.Descriptor);
      SetFlags(ParamOut.Descriptor, Flags);
      break;
    }
    default:
      IFT(E_FAIL);
    }
  }

  // Static samplers.
  if (pDescIn->NumStaticSamplers > 0) {
    pDescOut->pStaticSamplers =
        new DxilStaticSamplerDesc[pDescIn->NumStaticSamplers];
    pDescOut->NumStaticSamplers = pDescIn->NumStaticSamplers;
    memcpy((void *)pDescOut->pStaticSamplers, pDescIn->pStaticSamplers,
           pDescOut->NumStaticSamplers * sizeof(DxilStaticSamplerDesc));
  }
}

void ConvertRootSignature(
    const DxilVersionedRootSignatureDesc *pRootSignatureIn,
    DxilRootSignatureVersion RootSignatureVersionOut,
    const DxilVersionedRootSignatureDesc **ppRootSignatureOut) {
  IFTBOOL(pRootSignatureIn != nullptr && ppRootSignatureOut != nullptr,
          E_INVALIDARG);
  *ppRootSignatureOut = nullptr;

  if (pRootSignatureIn->Version == RootSignatureVersionOut) {
    // No conversion. Return the original root signature pointer; no cloning.
    *ppRootSignatureOut = pRootSignatureIn;
    return;
  }

  DxilVersionedRootSignatureDesc *pRootSignatureOut = nullptr;

  try {
    pRootSignatureOut = new DxilVersionedRootSignatureDesc();
    memset(pRootSignatureOut, 0, sizeof(*pRootSignatureOut));

    // Convert root signature.
    switch (RootSignatureVersionOut) {
    case DxilRootSignatureVersion::Version_1_0:
      switch (pRootSignatureIn->Version) {
      case DxilRootSignatureVersion::Version_1_1:
        pRootSignatureOut->Version = DxilRootSignatureVersion::Version_1_0;
        ConvertRootSignatureTemplate<DxilRootSignatureDesc1,
                                     DxilRootSignatureDesc, DxilRootParameter,
                                     DxilRootDescriptor, DxilDescriptorRange>(
            pRootSignatureIn->Desc_1_1, DxilRootSignatureVersion::Version_1_0,
            pRootSignatureOut->Desc_1_0);
        break;
      default:
        IFT(E_INVALIDARG);
      }
      break;

    case DxilRootSignatureVersion::Version_1_1:
      switch (pRootSignatureIn->Version) {
      case DxilRootSignatureVersion::Version_1_0:
        pRootSignatureOut->Version = DxilRootSignatureVersion::Version_1_1;
        ConvertRootSignatureTemplate<DxilRootSignatureDesc,
                                     DxilRootSignatureDesc1, DxilRootParameter1,
                                     DxilRootDescriptor1, DxilDescriptorRange1>(
            pRootSignatureIn->Desc_1_0, DxilRootSignatureVersion::Version_1_1,
            pRootSignatureOut->Desc_1_1);
        break;
      default:
        IFT(E_INVALIDARG);
      }
      break;

    default:
      IFT(E_INVALIDARG);
      break;
    }
  } catch (...) {
    DeleteRootSignature(pRootSignatureOut);
    throw;
  }

  *ppRootSignatureOut = pRootSignatureOut;
}

} // namespace hlsl
