text: /*++
  
  Module Name:
  
      GzipDataWriter.cpp
  
  Abstract:
  
      File writer that compresses data into zip format.
  
  Environment:
  
      User mode service.
  
      Not thread safe.
  
  --*/
  
  #include "stdafx.h"
  #include "GzipDataWriter.h"
  #include "BigAlloc.h"
  #include "VariableSizeVector.h"
  #include "ParallelTask.h"
  #include "RangeSplitter.h"
  #include "Bam.h"
  #include "zlib.h"
  #include "exit.h"
  #include "Error.h"
  
  using std::min;
  using std::max;
  using std::pair;
  
  class GzipWriterFilterSupplier;
  
  class GzipCompressWorkerManager : public ParallelWorkerManager
  {
  public:
      GzipCompressWorkerManager(GzipWriterFilterSupplier* i_filterSupplier)
          : filterSupplier(i_filterSupplier), buffer(NULL),
          chunkSize(i_filterSupplier->chunkSize), bam(i_filterSupplier->bamFormat)
      {}
  
      virtual ~GzipCompressWorkerManager();
  
      virtual void initialize(void* i_writer);
  
      virtual ParallelWorker* createWorker();
  
      virtual void beginStep();
  
      virtual void finishStep();
  
  private:
      VariableSizeVector<size_t> sizes;
      volatile int nChunks;
      const size_t chunkSize;
      const bool bam;
      FileEncoder* encoder;
      GzipWriterFilterSupplier* filterSupplier;
      char* input;
      size_t inputSize;
      size_t bufferSizeInit;
      size_t bufferSize;
      size_t inputUsed;
      char* buffer;
      VariableSizeVector< pair<_uint64,_uint64> > translation;
  
      friend class GzipCompressWorker;
  };
  
  class GzipCompressWorker : public ParallelWorker
  {
  public:
      GzipCompressWorker() : heap(NULL) {}
  
      virtual ~GzipCompressWorker() { delete heap; }
  
      virtual void step();
  
      static size_t compressChunk(z_stream& zstream, bool bamFormat, char* toBuffer, size_t toSize, char* fromBuffer, size_t fromUsed);
  
  private:
      z_stream zstream;
      ThreadHeap* heap;
  };
  
  // used for case where each thread compresses by itself
  
  class GzipWriterFilter : public DataWriter::Filter
  {
  public:
      GzipWriterFilter(GzipWriterFilterSupplier* i_supplier);
  
      ~GzipWriterFilter();
  
      virtual void onAdvance(DataWriter* writer, size_t batchOffset, char* data, GenomeDistance bytes, GenomeLocation location);
  
      virtual size_t onNextBatch(DataWriter* writer, size_t offset, size_t bytes, bool lastBatch = false, bool* needMoreBuffer = NULL, size_t* fromBufferUsed = NULL);
  
  private:
  
      GzipWriterFilterSupplier* supplier;
      // if doing inline compression, filled in with minimally initialized objects
      GzipCompressWorkerManager* manager;
      ParallelWorker* worker;
      FileEncoder* encoder;
  };
  
  GzipCompressWorkerManager::~GzipCompressWorkerManager()
  {
      if (buffer != NULL) {
          BigDealloc(buffer);
          buffer = NULL;
      }
  }
  
      void
  GzipCompressWorkerManager::initialize(
      void* i_encoder)
  {
      encoder = (FileEncoder*) i_encoder;
  }
  
      ParallelWorker*
  GzipCompressWorkerManager::createWorker()
  {
      return new GzipCompressWorker();
  }
  
      void
  GzipCompressWorkerManager::beginStep()
  {
      if (filterSupplier->closing) {
          nChunks = 0;
          return;
      }
      encoder->getEncodeBatch(&input, &inputSize, &inputUsed);
      nChunks = (int) ((inputUsed + chunkSize - 1) / chunkSize);
      sizes.clear();
      sizes.extend(nChunks);
  
      if (buffer == NULL) {
          buffer = (char*) BigAlloc(inputSize);
          if (buffer == NULL) {
              WriteErrorMessage("Unable to allocate %lld bytes for gzip compression buffer\n", inputSize);
              soft_exit(1);
          }
          bufferSizeInit = inputSize;
          bufferSize = inputSize;
      }
      else if (inputSize > bufferSize) {
  #ifdef VALIDATE_WRITE
          fprintf(stderr, "CompressManager::beginStep() prevInputSize: %lld inputSize: %lld\n", bufferSize, inputSize);
  #endif
          BigDealloc(buffer);
          buffer = NULL;
          buffer = (char*)BigAlloc(inputSize);
          if (buffer == NULL) {
              WriteErrorMessage("Unable to allocate %lld bytes for gzip compression buffer\n", inputSize);
              soft_exit(1);
          }
          bufferSize = inputSize;
      }
  }
  
      void
  GzipCompressWorkerManager::finishStep()
  {
      if (filterSupplier->closing) {
          return;
      }
      size_t toUsed = 0, logicalOffset, physicalOffset;
      encoder->getOffsets(&logicalOffset, &physicalOffset);
      for (int i = 0; i < nChunks; i++) {
          translation.push_back(pair<_uint64,_uint64>(logicalOffset, physicalOffset + toUsed));
          _ASSERT(i * chunkSize < inputUsed);
          _ASSERT(sizes[i] <= chunkSize);
          size_t logicalChunk = min(chunkSize, inputUsed - i * chunkSize);
          logicalOffset += logicalChunk;
          // _ASSERT(((BgzfHeader*)(buffer + i * chunkSize))->validate(sizes[i], logicalChunk));
          if (!((BgzfHeader*)(buffer + i * chunkSize))->validate(sizes[i], logicalChunk)) {
              WriteErrorMessage("BGZF Header validation failed. size:%lld, logicalChunk:%lld\n", sizes[i], logicalChunk);
              soft_exit(1);
          }
          memcpy(input + toUsed, buffer + i * chunkSize, sizes[i]);
          toUsed += sizes[i];
      }
      // _ASSERT(BgzfHeader::validate(input, toUsed));
      if (!BgzfHeader::validate(input, toUsed)) {
          WriteErrorMessage("BGZF Header validation failed. toUsed:%lld\n", toUsed);
          soft_exit(1);
      }
      encoder->setEncodedBatchSize(toUsed);
      filterSupplier->addTranslations(&translation);
      translation.clear();
  
      //
      // Shrink buffers to reduce memory consumption
      //
      if (bufferSize > bufferSizeInit) {
          size_t newBufferSize = bufferSizeInit;
          char* newBuffer = (char*)BigAlloc(newBufferSize);
          if (newBuffer == NULL) {
              WriteErrorMessage("Unable to allocate %lld bytes for write buffer\n", newBufferSize);
              soft_exit(1);
          }
  #ifdef VALIDATE_WRITE
          fprintf(stderr, "Shrinking compression buffer from %lld to %lld bytes\n", bufferSize, newBufferSize);
  #endif
          BigDealloc(buffer);
          buffer = newBuffer;
          bufferSize = newBufferSize;
      }
  }
  
      void
  GzipCompressWorker::step()
  {
      GzipCompressWorkerManager* supplier = (GzipCompressWorkerManager*) getManager();
      if (heap == NULL) {
          heap = new ThreadHeap(supplier->chunkSize * 8); // appears to use 4*chunkSize per run
          zstream.zalloc = zalloc;
          zstream.zfree = zfree;
          zstream.opaque = heap;
      }
      //fprintf(stderr, "zip task thread %d begin. nChunks %d\n", GetCurrentThreadId(), supplier->nChunks);
      _int64 start = timeInMillis();
      int begin = (getThreadNum() * supplier->nChunks) / getNumThreads();
      int end = ((1 + getThreadNum()) * supplier->nChunks) / getNumThreads();
      for (int i = begin; i < end; i++) {
          size_t bytes = min(supplier->chunkSize, supplier->inputUsed - i * supplier->chunkSize);
          supplier->sizes[i] = compressChunk(zstream, supplier->bam,
              supplier->buffer + i * supplier->chunkSize, supplier->chunkSize,
              supplier->input + i * supplier->chunkSize, bytes);
          _ASSERT(supplier->sizes[i] <= supplier->chunkSize); // can't grow!
      }
  }
  
  
      size_t
  GzipCompressWorker::compressChunk(
      z_stream& zstream,
      bool bamFormat,
      char* toBuffer,
      size_t toSize,
      char* fromBuffer,
      size_t fromUsed)
  {
      if (bamFormat && fromUsed > BAM_BLOCK) {
          WriteErrorMessage("exceeded BAM chunk size\n");
          soft_exit(1);
      }
      if (zstream.opaque != NULL) {
          ((ThreadHeap*)zstream.opaque)->reset();
      }
      // set up BAM header structure
      gz_header header;
      _uint8 bamExtraData[6];
      if (bamFormat) {
          header.text = false;
          header.time = 0;
          header.xflags = 0;
          header.os = 0;
          header.extra = bamExtraData;
          header.extra_len = 6;
          header.extra_max = 6;
          header.name = NULL;
          header.name_max = 0;
          header.comment = NULL;
          header.comm_max = 0;
          header.hcrc = false;
          header.done = true;
          bamExtraData[0] = 'B';
          bamExtraData[1] = 'C';
          bamExtraData[2] = 2;
          bamExtraData[3] = 0;
          bamExtraData[4] = 3; // will be filled in later
          bamExtraData[5] = 7; // will be filled in later
      }
  
      if (fromUsed > 0xffffffff || toSize > 0xffffffff) {
          WriteErrorMessage("GZipDataWriter: fromUsed or toSize too big\n");
          soft_exit(1);
      }
  
      // based on sample code at http://www.lemoda.net/c/zlib-open-write/index.html
      const int windowBits = 15;
      const int GZIP_ENCODING = 16;
      zstream.next_in = (Bytef*) fromBuffer;
      zstream.avail_in = (uInt)fromUsed;
      zstream.next_out = (Bytef*) toBuffer;
      zstream.avail_out = (uInt)toSize;
      uInt oldAvail;
      int status;
  
      status = deflateInit2(&zstream, Z_DEFAULT_COMPRESSION, Z_DEFLATED, windowBits | GZIP_ENCODING, 8, Z_DEFAULT_STRATEGY);
      if (status < 0) {
          WriteErrorMessage("GzipWriterFilter: deflateInit2 failed with %d\n", status);
          soft_exit(1);
      }
      if (bamFormat) {
          status = deflateSetHeader(&zstream, &header);
          if (status != Z_OK) {
              WriteErrorMessage("GzipWriterFilter: defaultSetHeader failed with %d\n", status);
              soft_exit(1);
          }
      }
      oldAvail = zstream.avail_out;
      status = deflate(&zstream, Z_FINISH);
      if (status < 0 && status != Z_BUF_ERROR) {
          WriteErrorMessage("GzipWriterFilter: deflate failed with %d\n", status);
          soft_exit(1);
      }
  
      // make sure it all got written out in a single compressed block
      if (zstream.avail_in != 0) {
          WriteErrorMessage("GzipWriterFilter: default failed to read all input\n");
          soft_exit(1);
      }
      if (zstream.avail_out == oldAvail) {
          WriteErrorMessage("GzipWriterFilter: default failed to write output\n");
          soft_exit(1);
      }
      status = deflateEnd(&zstream);
      if (status < 0) {
          WriteErrorMessage("GzipWriterFilter: deflateEnd failed with %d\n", status);
          soft_exit(1);
      }
  
      size_t toUsed = toSize - zstream.avail_out;
      if (bamFormat) {
          // backpatch compressed block size into gzip header
          if (toUsed >= BAM_BLOCK) {
              WriteErrorMessage("exceeded BAM chunk size\n");
              soft_exit(1);
          }
          * (_uint16*) (toBuffer + 16) = (_uint16) (toUsed - 1);
      }
      return toUsed;
  }
  
  GzipWriterFilter::GzipWriterFilter(GzipWriterFilterSupplier* i_supplier)
      : DataWriter::Filter(DataWriter::ResizeFilter), supplier(i_supplier), manager(NULL), worker(NULL)
  {}
  
  GzipWriterFilter::~GzipWriterFilter()
  {
      if (manager != NULL) {
          delete manager;
          manager = NULL;
      }
  
      if (worker != NULL) {
          delete worker;
          worker = NULL;
      }
  } // ~GzipWriterFilter
  
      void
  GzipWriterFilter::onAdvance(
      DataWriter* writer,
      size_t batchOffset,
      char* data,
      GenomeDistance bytes,
      GenomeLocation location)
  {
      // nothing
  }
  
      size_t
  GzipWriterFilter::onNextBatch(
      DataWriter* writer,
      size_t offset,
      size_t bytes,
      bool lastBatch,
      bool* needMoreBuffer,
      size_t* fromBufferUsed)
  {
      
      // 
      // Nothing to write
      //
      if (bytes == 0 || *needMoreBuffer) {
          return 0;
      }
  
      char* fromBuffer;
      size_t fromSize, fromUsed, physicalOffset, logicalOffset;
      writer->getBatch(-1, &fromBuffer, &fromSize, &fromUsed, &physicalOffset, NULL, &logicalOffset);
      if (fromUsed == 0 || supplier->multiThreaded || supplier->closing) {
          if (fromBufferUsed != NULL) {
              *fromBufferUsed = min<long long>(fromUsed, bytes);
          }
          return min<long long>(fromUsed, bytes);
      }
      // do compress buffer synchronously in-place
      if (manager == NULL) {
          manager = new GzipCompressWorkerManager(supplier);
          worker = manager->createWorker();
          encoder = new FileEncoder(0, false, manager);
          encoder->initialize((AsyncDataWriter*) writer);
          manager->initialize(encoder);
          manager->configure(worker, 0, 1);
      }
      encoder->setupEncode(-1);
      manager->beginStep();
      worker->step();
      manager->finishStep();
      writer->getBatch(-1, &fromBuffer, &fromSize, &fromUsed, &physicalOffset, NULL, &logicalOffset);
      if (fromBufferUsed != NULL) {
          *fromBufferUsed = min<long long>(fromUsed, bytes);
      }
      return min<long long>(fromUsed, bytes);
  }
  
      GzipWriterFilterSupplier*
  DataWriterSupplier::gzip(
      bool bamFormat,
      size_t chunkSize,
      int numThreads,
      bool bindToProcessors,
      bool multiThreaded)
  {
      return new GzipWriterFilterSupplier(bamFormat, chunkSize, numThreads, bindToProcessors, multiThreaded);
  }
  
      DataWriter::Filter*
  GzipWriterFilterSupplier::getFilter()
  {
      return new GzipWriterFilter(this);
  }
  
      void
  GzipWriterFilterSupplier::onClosing(
      DataWriterSupplier* supplier)
  {
      if (bamFormat) {
          closing = true;
          DataWriter* writer = supplier->getWriter();
          // write empty block as BAM end of file marker
          static _uint8 eof[] = {
              0x1f, 0x8b, 0x08, 0x04, 0x00, 0x00, 0x00, 0x00, 0x00, 0xff, 0x06, 0x00, 0x42, 0x43,
              0x02, 0x00, 0x1b, 0x00, 0x03, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00
          };
          char* buffer;
          size_t bytes;
          if (! (writer->getBuffer(&buffer, &bytes) && bytes >= sizeof(eof))) {
              WriteErrorMessage("no space to write eof marker\n");
              soft_exit(1);
          }
          memcpy(buffer, eof, sizeof(eof));
          writer->advance(sizeof(eof));
  
          // add final translation for last empty block
          writer->nextBatch();
          char* ignore;
          pair<_uint64,_uint64> last;
          size_t used;
          writer->getBatch(-1, &ignore, NULL, &used, (size_t*) &last.second, NULL, (size_t*) &last.first);
          last.second += used;
          translation.push_back(last);
  
          writer->close();
          delete writer;
      }
  
      // sort translations
      std::sort(translation.begin(), translation.end(), translationComparator);
  }
  
      void
  GzipWriterFilterSupplier::addTranslations(
      VariableSizeVector< pair<_uint64,_uint64> >* moreTranslations)
  {
      AcquireExclusiveLock(&lock);
      translation.append(moreTranslations);
      ReleaseExclusiveLock(&lock);
  }
  
  
      bool
  GzipWriterFilterSupplier::translate(
      _uint64 logical,
      _uint64* o_physical,
      _uint64* o_logicalDelta)
  {
      pair<_uint64,_uint64> value;
      value.first = logical;
      value.second = 0; //ignored
      pair<_uint64,_uint64>* upper = std::upper_bound(translation.begin(), translation.end(), value, translationComparator);
      if (upper == translation.begin()) {
          return false;
      }
      upper--;
      *o_physical = upper->second;
      *o_logicalDelta = logical - upper->first;
      return true;
  }
  
      bool
  GzipWriterFilterSupplier::translationComparator(
      const pair<_uint64,_uint64>& a,
      const pair<_uint64,_uint64>& b)
  {
      return a.first < b.first;
  }
  
      FileEncoder*
  FileEncoder::gzip(
      GzipWriterFilterSupplier* filterSupplier,
      int numThreads,
      bool bindToProcessor,
      size_t chunkSize,
      bool bam)
  {
      return new FileEncoder(numThreads, bindToProcessor, new GzipCompressWorkerManager(filterSupplier));
  }
  