• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
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