1 files changed, 162 insertions, 0 deletions
diff --git a/tensorflow/core/lib/jpeg/jpeg_handle.cc b/tensorflow/core/lib/jpeg/jpeg_handle.cc
new file mode 100644
index 0000000000..4521be0afb
--- /dev/null
+++ b/tensorflow/core/lib/jpeg/jpeg_handle.cc
@@ -0,0 +1,162 @@
+// This file implements a memory destination for libjpeg
+// The design is very similar to jdatadst.c in libjpeg
+// These functions are not meant to be used directly, see jpeg_mem.h instead.
+// We are filling out stubs required by jpeglib, those stubs are private to
+// the implementation, we are just making available JPGMemSrc, JPGMemDest
+
+#include "tensorflow/core/lib/jpeg/jpeg_handle.h"
+
+#include <setjmp.h>
+#include <stddef.h>
+
+#include "tensorflow/core/platform/logging.h"
+
+namespace tensorflow {
+namespace jpeg {
+
+void CatchError(j_common_ptr cinfo) {
+  (*cinfo->err->output_message)(cinfo);
+  jmp_buf *jpeg_jmpbuf = reinterpret_cast<jmp_buf *>(cinfo->client_data);
+  jpeg_destroy(cinfo);
+  longjmp(*jpeg_jmpbuf, 1);
+}
+
+// *****************************************************************************
+// *****************************************************************************
+// *****************************************************************************
+// Destination functions
+
+// -----------------------------------------------------------------------------
+void MemInitDestination(j_compress_ptr cinfo) {
+  MemDestMgr *dest = reinterpret_cast<MemDestMgr *>(cinfo->dest);
+  VLOG(1) << "Initializing buffer=" << dest->bufsize << " bytes";
+  dest->pub.next_output_byte = dest->buffer;
+  dest->pub.free_in_buffer = dest->bufsize;
+  dest->datacount = 0;
+  if (dest->dest) {
+    dest->dest->clear();
+  }
+}
+
+// -----------------------------------------------------------------------------
+boolean MemEmptyOutputBuffer(j_compress_ptr cinfo) {
+  MemDestMgr *dest = reinterpret_cast<MemDestMgr *>(cinfo->dest);
+  VLOG(1) << "Writing " << dest->bufsize << " bytes";
+  if (dest->dest) {
+    dest->dest->append(reinterpret_cast<char *>(dest->buffer), dest->bufsize);
+  }
+  dest->pub.next_output_byte = dest->buffer;
+  dest->pub.free_in_buffer = dest->bufsize;
+  return TRUE;
+}
+
+// -----------------------------------------------------------------------------
+void MemTermDestination(j_compress_ptr cinfo) {
+  MemDestMgr *dest = reinterpret_cast<MemDestMgr *>(cinfo->dest);
+  VLOG(1) << "Writing " << dest->bufsize - dest->pub.free_in_buffer << " bytes";
+  if (dest->dest) {
+    dest->dest->append(reinterpret_cast<char *>(dest->buffer),
+                       dest->bufsize - dest->pub.free_in_buffer);
+    VLOG(1) << "Total size= " << dest->dest->size();
+  }
+  dest->datacount = dest->bufsize - dest->pub.free_in_buffer;
+}
+
+// -----------------------------------------------------------------------------
+void SetDest(j_compress_ptr cinfo, void *buffer, int bufsize) {
+  SetDest(cinfo, buffer, bufsize, NULL);
+}
+
+// -----------------------------------------------------------------------------
+void SetDest(j_compress_ptr cinfo, void *buffer, int bufsize,
+             string *destination) {
+  MemDestMgr *dest;
+  if (cinfo->dest == NULL) {
+    cinfo->dest = reinterpret_cast<struct jpeg_destination_mgr *>(
+        (*cinfo->mem->alloc_small)(reinterpret_cast<j_common_ptr>(cinfo),
+                                   JPOOL_PERMANENT, sizeof(MemDestMgr)));
+  }
+
+  dest = reinterpret_cast<MemDestMgr *>(cinfo->dest);
+  dest->bufsize = bufsize;
+  dest->buffer = static_cast<JOCTET *>(buffer);
+  dest->dest = destination;
+  dest->pub.init_destination = MemInitDestination;
+  dest->pub.empty_output_buffer = MemEmptyOutputBuffer;
+  dest->pub.term_destination = MemTermDestination;
+}
+
+// *****************************************************************************
+// *****************************************************************************
+// *****************************************************************************
+// Source functions
+
+// -----------------------------------------------------------------------------
+void MemInitSource(j_decompress_ptr cinfo) {
+  MemSourceMgr *src = reinterpret_cast<MemSourceMgr *>(cinfo->src);
+  src->pub.next_input_byte = src->data;
+  src->pub.bytes_in_buffer = src->datasize;
+}
+
+// -----------------------------------------------------------------------------
+// We emulate the same error-handling as fill_input_buffer() from jdatasrc.c,
+// for coherency's sake.
+boolean MemFillInputBuffer(j_decompress_ptr cinfo) {
+  static const JOCTET kEOIBuffer[2] = {0xff, JPEG_EOI};
+  MemSourceMgr *src = reinterpret_cast<MemSourceMgr *>(cinfo->src);
+  if (src->pub.bytes_in_buffer == 0 && src->pub.next_input_byte == src->data) {
+    // empty file -> treated as an error.
+    ERREXIT(cinfo, JERR_INPUT_EMPTY);
+    return FALSE;
+  } else if (src->pub.bytes_in_buffer) {
+    // if there's still some data left, it's probably corrupted
+    return src->try_recover_truncated_jpeg ? TRUE : FALSE;
+  } else if (src->pub.next_input_byte != kEOIBuffer &&
+             src->try_recover_truncated_jpeg) {
+    // In an attempt to recover truncated files, we insert a fake EOI
+    WARNMS(cinfo, JWRN_JPEG_EOF);
+    src->pub.next_input_byte = kEOIBuffer;
+    src->pub.bytes_in_buffer = 2;
+    return TRUE;
+  } else {
+    // We already inserted a fake EOI and it wasn't enough, so this time
+    // it's really an error.
+    ERREXIT(cinfo, JERR_FILE_READ);
+    return FALSE;
+  }
+}
+
+// -----------------------------------------------------------------------------
+void MemTermSource(j_decompress_ptr cinfo) {}
+
+// -----------------------------------------------------------------------------
+void MemSkipInputData(j_decompress_ptr cinfo, long jump) {
+  MemSourceMgr *src = reinterpret_cast<MemSourceMgr *>(cinfo->src);
+  src->pub.bytes_in_buffer -= jump;
+  src->pub.next_input_byte += jump;
+}
+
+// -----------------------------------------------------------------------------
+void SetSrc(j_decompress_ptr cinfo, const void *data,
+            unsigned long int datasize, bool try_recover_truncated_jpeg) {
+  MemSourceMgr *src;
+
+  cinfo->src = reinterpret_cast<struct jpeg_source_mgr *>(
+      (*cinfo->mem->alloc_small)(reinterpret_cast<j_common_ptr>(cinfo),
+                                 JPOOL_PERMANENT, sizeof(MemSourceMgr)));
+
+  src = reinterpret_cast<MemSourceMgr *>(cinfo->src);
+  src->pub.init_source = MemInitSource;
+  src->pub.fill_input_buffer = MemFillInputBuffer;
+  src->pub.skip_input_data = MemSkipInputData;
+  src->pub.resync_to_restart = jpeg_resync_to_restart;
+  src->pub.term_source = MemTermSource;
+  src->data = reinterpret_cast<const unsigned char *>(data);
+  src->datasize = datasize;
+  src->pub.bytes_in_buffer = 0;
+  src->pub.next_input_byte = NULL;
+  src->try_recover_truncated_jpeg = try_recover_truncated_jpeg;
+}
+
+}  // namespace jpeg
+}  // namespace tensorflow