lib/DxilDia/DxilDiaDataSource.cpp

///////////////////////////////////////////////////////////////////////////////
//                                                                           //
// DxilDiaDataSource.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.                                     //
//                                                                           //
// DIA API implementation for DXIL modules.                                  //
//                                                                           //
///////////////////////////////////////////////////////////////////////////////

#include "DxilDiaDataSource.h"

#include "dxc/DXIL/DxilPDB.h"
#include "dxc/DXIL/DxilUtil.h"
#include "dxc/DxilContainer/DxilContainer.h"
#include "dxc/Support/FileIOHelper.h"
#include "dxc/Support/dxcapi.impl.h"

#include "llvm/Support/FileSystem.h"
#include "llvm/Support/MSFileSystem.h"

#include "llvm/IR/DebugInfo.h"
#include "llvm/IR/LLVMContext.h"
#include "llvm/IR/Module.h"
#include "llvm/Support/MemoryBuffer.h"

#include "DxilDiaSession.h"

dxil_dia::DataSource::DataSource(IMalloc *pMalloc) : m_pMalloc(pMalloc) {}

dxil_dia::DataSource::~DataSource() {
  // These are cross-referenced, so let's be explicit.
  m_finder.reset();
  m_module.reset();
  m_context.reset();
}

STDMETHODIMP dxil_dia::DataSource::get_lastError(BSTR *pRetVal) {
  *pRetVal = nullptr;
  return S_OK;
}

namespace dxil_dia {
std::unique_ptr<llvm::MemoryBuffer>
getMemBufferFromBlob(IDxcBlob *pBlob, const llvm::Twine &BufferName) {
  llvm::StringRef Data((LPSTR)pBlob->GetBufferPointer(),
                       pBlob->GetBufferSize());
  return llvm::MemoryBuffer::getMemBufferCopy(Data, BufferName);
}

std::unique_ptr<llvm::MemoryBuffer>
getMemBufferFromStream(IStream *pStream, std::vector<char> &DataContainer,
                       const llvm::Twine &BufferName) {
  CComPtr<IDxcBlob> pBlob;
  if (SUCCEEDED(pStream->QueryInterface(&pBlob))) {
    return getMemBufferFromBlob(pBlob, BufferName);
  }

  STATSTG statstg;
  IFT(pStream->Stat(&statstg, STATFLAG_NONAME));
  size_t size = statstg.cbSize.LowPart;

  DataContainer.resize(size);
  ULONG read;
  IFT(pStream->Read(DataContainer.data(), size, &read));
  llvm::StringRef Str(DataContainer.data(), size);
  std::unique_ptr<llvm::MemoryBuffer> result(
      llvm::MemoryBuffer::getMemBuffer(Str, BufferName.str(), false));
  return result;
}
} // namespace dxil_dia

STDMETHODIMP dxil_dia::DataSource::loadDataFromIStream(IStream *pInputIStream) {
  try {
    DxcThreadMalloc TM(m_pMalloc);
    if (m_module.get() != nullptr) {
      return E_FAIL;
    }

    // Setup filesystem because bitcode reader might emit warning
    ::llvm::sys::fs::MSFileSystem *msfPtr;
    IFT(CreateMSFileSystemForDisk(&msfPtr));
    std::unique_ptr<::llvm::sys::fs::MSFileSystem> msf(msfPtr);

    ::llvm::sys::fs::AutoPerThreadSystem pts(msf.get());
    IFTLLVM(pts.error_code());

    CComPtr<IStream> pIStream = pInputIStream;
    CComPtr<IDxcBlob> pContainer;
    if (SUCCEEDED(hlsl::pdb::LoadDataFromStream(m_pMalloc, pInputIStream,
                                                &pContainer))) {
      const hlsl::DxilContainerHeader *pContainerHeader =
          hlsl::IsDxilContainerLike(pContainer->GetBufferPointer(),
                                    pContainer->GetBufferSize());
      if (!hlsl::IsValidDxilContainer(pContainerHeader,
                                      pContainer->GetBufferSize()))
        return E_FAIL;
      const hlsl::DxilPartHeader *PartHeader = hlsl::GetDxilPartByType(
          pContainerHeader, hlsl::DFCC_ShaderDebugInfoDXIL);
      if (!PartHeader)
        return E_FAIL;
      CComPtr<IDxcBlobEncoding> pPinnedBlob;
      IFR(hlsl::DxcCreateBlobWithEncodingFromPinned(
          PartHeader + 1, PartHeader->PartSize, CP_ACP, &pPinnedBlob));
      pIStream.Release();
      IFR(hlsl::CreateReadOnlyBlobStream(pPinnedBlob, &pIStream));
    }

    m_context.reset();
    m_finder.reset();

    m_context = std::make_shared<llvm::LLVMContext>();
    llvm::MemoryBuffer *pBitcodeBuffer;
    std::unique_ptr<llvm::MemoryBuffer> pEmbeddedBuffer;

    std::vector<char> DataContainer;
    std::unique_ptr<llvm::MemoryBuffer> pBuffer =
        getMemBufferFromStream(pIStream, DataContainer, "data");
    size_t bufferSize = pBuffer->getBufferSize();

    // The buffer can hold LLVM bitcode for a module, or the ILDB
    // part from a container.
    if (bufferSize < sizeof(UINT32)) {
      return DXC_E_MALFORMED_CONTAINER;
    }
    const UINT32 BC_C0DE = ((INT32)(INT8)'B' | (INT32)(INT8)'C' << 8 |
                            (INT32)0xDEC0 << 16); // BC0xc0de in big endian
    if (BC_C0DE == *(const UINT32 *)pBuffer->getBufferStart()) {
      pBitcodeBuffer = pBuffer.get();
    } else {
      if (bufferSize <= sizeof(hlsl::DxilProgramHeader)) {
        return DXC_E_MALFORMED_CONTAINER;
      }

      const hlsl::DxilProgramHeader *pDxilProgramHeader =
          (const hlsl::DxilProgramHeader *)pBuffer->getBufferStart();
      if (pDxilProgramHeader->BitcodeHeader.DxilMagic != hlsl::DxilMagicValue) {
        return DXC_E_MALFORMED_CONTAINER;
      }

      UINT32 BlobSize;
      const char *pBitcode = nullptr;
      hlsl::GetDxilProgramBitcode(pDxilProgramHeader, &pBitcode, &BlobSize);
      std::unique_ptr<llvm::MemoryBuffer> p = llvm::MemoryBuffer::getMemBuffer(
          llvm::StringRef(pBitcode, BlobSize), "data",
          false /* RequiresNullTerminator */);
      pEmbeddedBuffer.swap(p);
      pBitcodeBuffer = pEmbeddedBuffer.get();
    }

    std::string DiagStr;
    std::unique_ptr<llvm::Module> pModule =
        hlsl::dxilutil::LoadModuleFromBitcode(pBitcodeBuffer, *m_context.get(),
                                              DiagStr);
    if (!pModule.get())
      return E_FAIL;
    m_finder = std::make_shared<llvm::DebugInfoFinder>();
    m_finder->processModule(*pModule.get());
    m_module.reset(pModule.release());
  }
  CATCH_CPP_RETURN_HRESULT();
  return S_OK;
}

STDMETHODIMP dxil_dia::DataSource::openSession(IDiaSession **ppSession) {
  DxcThreadMalloc TM(m_pMalloc);
  *ppSession = nullptr;
  if (m_module.get() == nullptr)
    return E_FAIL;

  ::llvm::sys::fs::MSFileSystem *msfPtr;
  IFT(CreateMSFileSystemForDisk(&msfPtr));
  std::unique_ptr<::llvm::sys::fs::MSFileSystem> msf(msfPtr);

  ::llvm::sys::fs::AutoPerThreadSystem pts(msf.get());
  IFTLLVM(pts.error_code());

  CComPtr<Session> pSession = Session::Alloc(DxcGetThreadMallocNoRef());
  IFROOM(pSession.p);
  pSession->Init(m_context, m_module, m_finder);
  *ppSession = pSession.Detach();
  return S_OK;
}

HRESULT CreateDxcDiaDataSource(REFIID riid, LPVOID *ppv) {
  CComPtr<dxil_dia::DataSource> result =
      CreateOnMalloc<dxil_dia::DataSource>(DxcGetThreadMallocNoRef());
  if (result == nullptr) {
    *ppv = nullptr;
    return E_OUTOFMEMORY;
  }

  return result.p->QueryInterface(riid, ppv);
}