lib/DxilCompression/DxilCompression.cpp

///////////////////////////////////////////////////////////////////////////////
//                                                                           //
// DxilCompression.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.                                     //
//                                                                           //
///////////////////////////////////////////////////////////////////////////////
//
// Helper wrapper functions for zlib deflate and inflate. Entirely
// self-contained, only depends on IMalloc interface.
//

#include "dxc/DxilCompression/DxilCompression.h"
#include "dxc/Support/Global.h"
#include "dxc/Support/WinIncludes.h"

#include "miniz.h"
typedef size_t ZlibSize_t;
typedef const Bytef ZlibInputBytesf;

namespace {
//
// A resource managment class for a zlib stream that calls the appropriate init
// and end routines.
//
class Zlib {
public:
  enum Operation { INFLATE, DEFLATE };
  Zlib(Operation Op, IMalloc *pAllocator)
      : m_Stream{}, m_Op(Op), m_Initalized(false) {
    m_Stream = {};

    if (pAllocator) {
      m_Stream.zalloc = ZAlloc;
      m_Stream.zfree = ZFree;
      m_Stream.opaque = pAllocator;
    }

    int ret = Z_ERRNO;
    if (Op == INFLATE) {
      ret = inflateInit(&m_Stream);
    } else {
      ret = deflateInit(&m_Stream, Z_DEFAULT_COMPRESSION);
    }

    if (ret != Z_OK) {
      m_InitializationResult = TranslateZlibResult(ret);
      return;
    }

    m_Initalized = true;
  }

  ~Zlib() {
    if (m_Initalized) {
      if (m_Op == INFLATE) {
        inflateEnd(&m_Stream);
      } else {
        deflateEnd(&m_Stream);
      }
    }
  }

  z_stream *GetStream() {
    if (m_Initalized)
      return &m_Stream;
    return nullptr;
  }

  hlsl::ZlibResult GetInitializationResult() const {
    return m_InitializationResult;
  }

  static hlsl::ZlibResult TranslateZlibResult(int zlibResult) {
    switch (zlibResult) {
    default:
      return hlsl::ZlibResult::InvalidData;
    case Z_MEM_ERROR:
    case Z_BUF_ERROR:
      return hlsl::ZlibResult::OutOfMemory;
    }
  }

private:
  z_stream m_Stream;
  Operation m_Op;
  bool m_Initalized;
  hlsl::ZlibResult m_InitializationResult = hlsl::ZlibResult::Success;

  static void *ZAlloc(void *context, ZlibSize_t items, ZlibSize_t size) {
    IMalloc *mallocif = (IMalloc *)context;
    return mallocif->Alloc(items * size);
  }

  static void ZFree(void *context, void *pointer) {
    IMalloc *mallocif = (IMalloc *)context;
    mallocif->Free(pointer);
  }
};

} // namespace

hlsl::ZlibResult hlsl::ZlibDecompress(IMalloc *pMalloc,
                                      const void *pCompressedBuffer,
                                      size_t BufferSizeInBytes,
                                      void *pUncompressedBuffer,
                                      size_t UncompressedBufferSize) {
  Zlib zlib(Zlib::INFLATE, pMalloc);
  z_stream *pStream = zlib.GetStream();
  if (!pStream)
    return zlib.GetInitializationResult();

  pStream->avail_in = BufferSizeInBytes;
  pStream->next_in = (ZlibInputBytesf *)pCompressedBuffer;
  pStream->next_out = (Byte *)pUncompressedBuffer;
  pStream->avail_out = UncompressedBufferSize;

  // Compression should finish in one call because of the call to deflateBound.
  int status = inflate(pStream, Z_FINISH);
  if (status != Z_STREAM_END) {
    return Zlib::TranslateZlibResult(status);
  }

  return ZlibResult::Success;
}

hlsl::ZlibResult hlsl::ZlibCompress(IMalloc *pMalloc, const void *pData,
                                    size_t pDataSize, void *pUserData,
                                    ZlibCallbackFn *Callback,
                                    size_t *pOutCompressedSize) {
  Zlib zlib(Zlib::DEFLATE, pMalloc);
  z_stream *pStream = zlib.GetStream();
  if (!pStream)
    return zlib.GetInitializationResult();

  const size_t UpperBound = deflateBound(pStream, pDataSize);
  void *pDestBuffer = Callback(pUserData, UpperBound);
  if (!pDestBuffer)
    return ZlibResult::OutOfMemory;

  pStream->next_in = (ZlibInputBytesf *)pData;
  pStream->avail_in = pDataSize;
  pStream->next_out = (Byte *)pDestBuffer;
  pStream->avail_out = UpperBound;

  int status = deflate(pStream, Z_FINISH);
  if (status != Z_STREAM_END) {
    return Zlib::TranslateZlibResult(status);
  }

  *pOutCompressedSize = pStream->total_out;
  return ZlibResult::Success;
}