1 /* Copyright 2016 The TensorFlow Authors. All Rights Reserved. 2 3 Licensed under the Apache License, Version 2.0 (the "License"); 4 you may not use this file except in compliance with the License. 5 You may obtain a copy of the License at 6 7 http://www.apache.org/licenses/LICENSE-2.0 8 9 Unless required by applicable law or agreed to in writing, software 10 distributed under the License is distributed on an "AS IS" BASIS, 11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 See the License for the specific language governing permissions and 13 limitations under the License. 14 ==============================================================================*/ 15 16 #ifndef TENSORFLOW_LIB_IO_ZLIB_INPUTSTREAM_H_ 17 #define TENSORFLOW_LIB_IO_ZLIB_INPUTSTREAM_H_ 18 19 #include <string> 20 21 #include "tensorflow/core/lib/core/status.h" 22 #include "tensorflow/core/lib/io/inputstream_interface.h" 23 #include "tensorflow/core/lib/io/zlib_compression_options.h" 24 #include "tensorflow/core/platform/env.h" 25 #include "tensorflow/core/platform/macros.h" 26 #include "tensorflow/core/platform/types.h" 27 28 namespace tensorflow { 29 namespace io { 30 31 // Forward declare some members of zlib.h, which is only included in the 32 // .cc file. 33 struct ZStreamDef; 34 35 // An ZlibInputStream provides support for reading from a stream compressed 36 // using zlib (http://www.zlib.net/). Buffers the contents of the file. 37 // 38 // A given instance of an ZlibInputStream is NOT safe for concurrent use 39 // by multiple threads 40 class ZlibInputStream : public InputStreamInterface { 41 public: 42 // Create a ZlibInputStream for `input_stream` with a buffer of size 43 // `input_buffer_bytes` bytes for reading contents from `input_stream` and 44 // another buffer with size `output_buffer_bytes` for caching decompressed 45 // contents. 46 // 47 // Takes ownership of `input_stream` iff `owns_input_stream` is true. 48 ZlibInputStream(InputStreamInterface* input_stream, size_t input_buffer_bytes, 49 size_t output_buffer_bytes, 50 const ZlibCompressionOptions& zlib_options, 51 bool owns_input_stream); 52 53 // Equivalent to the previous constructor with owns_input_stream=false. 54 ZlibInputStream(InputStreamInterface* input_stream, size_t input_buffer_bytes, 55 size_t output_buffer_bytes, 56 const ZlibCompressionOptions& zlib_options); 57 58 ~ZlibInputStream(); 59 60 // Reads bytes_to_read bytes into *result, overwriting *result. 61 // 62 // Return Status codes: 63 // OK: If successful. 64 // OUT_OF_RANGE: If there are not enough bytes to read before 65 // the end of the stream. 66 // ABORTED: If inflate() fails, we return the error code with the 67 // error message in `z_stream_->msg`. 68 // others: If reading from stream failed. 69 Status ReadNBytes(int64 bytes_to_read, string* result) override; 70 71 int64 Tell() const override; 72 73 Status Reset() override; 74 75 private: 76 void InitZlibBuffer(); 77 78 const bool owns_input_stream_; 79 InputStreamInterface* input_stream_; 80 size_t input_buffer_capacity_; // Size of z_stream_input_ 81 size_t output_buffer_capacity_; // Size of z_stream_output_ 82 char* next_unread_byte_; // Next unread byte in z_stream_output_ 83 84 ZlibCompressionOptions const zlib_options_; 85 86 std::unique_ptr<ZStreamDef> z_stream_def_; 87 88 // Reads data from `input_stream_` and tries to fill up `z_stream_input_` if 89 // enough unread data is left in `input_stream_`. 90 // 91 // Looks up z_stream_->next_in to check how much data in z_stream_input_ 92 // has already been read. The used data is removed and new data is added to 93 // after any unread data in z_stream_input_. 94 // After this call z_stream_->next_in points to the start of z_stream_input_ 95 // and z_stream_->avail_in stores the number of readable bytes in 96 // z_stream_input_. 97 // 98 // Returns OutOfRange error if NO data could be read from stream. Note that 99 // this won't return an OutOfRange if there wasn't sufficient data in stream 100 // to completely fill up z_stream_input_. 101 Status ReadFromStream(); 102 103 // Calls `inflate()` and returns DataLoss Status if it failed. 104 Status Inflate(); 105 106 // Starts reading bytes at `next_unread_byte_` till either `bytes_to_read` 107 // bytes have been read or `z_stream_->next_out` is reached. 108 // Returns the number of bytes read and advances the `next_unread_byte_` 109 // pointer to the next location to read from. 110 size_t ReadBytesFromCache(size_t bytes_to_read, string* result); 111 112 // The number of unread bytes in z_stream_output_. 113 // 114 // z_stream_output_ --> 115 // 116 // [RRRRRRRRRRRRRRRRRRUUUUUUUUUUUUUU000000000000000000] 117 // ^ ^ 118 // next_unread_byte_ z_stream_->next_out 119 // 120 // R: Read bytes 121 // U: Unread bytes 122 // 0: garbage bytes where new output will be written 123 // 124 // Returns the size of [next_unread_byte_, z_stream_->next_out) 125 size_t NumUnreadBytes() const; 126 127 // Number of *uncompressed* bytes that have been read from this stream. 128 int64 bytes_read_; 129 130 TF_DISALLOW_COPY_AND_ASSIGN(ZlibInputStream); 131 }; 132 133 } // namespace io 134 } // namespace tensorflow 135 136 #endif // TENSORFLOW_LIB_IO_ZLIB_INPUTSTREAM_H_ 137