]> git.donarmstrong.com Git - bamtools.git/blob - BGZF.cpp
Added the <cctype> include to BamAux to explicitly support the toupper() fxn. It...
[bamtools.git] / BGZF.cpp
1 // ***************************************************************************\r
2 // BGZF.cpp (c) 2009 Derek Barnett, Michael Str�mberg\r
3 // Marth Lab, Department of Biology, Boston College\r
4 // All rights reserved.\r
5 // ---------------------------------------------------------------------------\r
6 // Last modified: 19 July 2010 (DB)\r
7 // ---------------------------------------------------------------------------\r
8 // BGZF routines were adapted from the bgzf.c code developed at the Broad\r
9 // Institute.\r
10 // ---------------------------------------------------------------------------\r
11 // Provides the basic functionality for reading & writing BGZF files\r
12 // ***************************************************************************\r
13 \r
14 #include <algorithm>\r
15 #include "BGZF.h"\r
16 using namespace BamTools;\r
17 using std::string;\r
18 using std::min;\r
19 \r
20 BgzfData::BgzfData(void)\r
21     : UncompressedBlockSize(DEFAULT_BLOCK_SIZE)\r
22     , CompressedBlockSize(MAX_BLOCK_SIZE)\r
23     , BlockLength(0)\r
24     , BlockOffset(0)\r
25     , BlockAddress(0)\r
26     , IsOpen(false)\r
27     , IsWriteOnly(false)\r
28     , Stream(NULL)\r
29     , UncompressedBlock(NULL)\r
30     , CompressedBlock(NULL)\r
31 {\r
32     try {\r
33         CompressedBlock   = new char[CompressedBlockSize];\r
34         UncompressedBlock = new char[UncompressedBlockSize];\r
35     } catch( std::bad_alloc& ba ) {\r
36         printf("BGZF ERROR: unable to allocate memory for our BGZF object.\n");\r
37         exit(1);\r
38     }\r
39 }\r
40 \r
41 // destructor\r
42 BgzfData::~BgzfData(void) {\r
43     if( CompressedBlock )   { delete[] CompressedBlock;   }\r
44     if( UncompressedBlock ) { delete[] UncompressedBlock; }\r
45 }\r
46 \r
47 // closes BGZF file\r
48 void BgzfData::Close(void) {\r
49 \r
50         // skip if file not open, otherwise set flag\r
51     if ( !IsOpen ) return;\r
52     IsOpen = false;\r
53 \r
54     // flush the current BGZF block\r
55     if ( IsWriteOnly ) FlushBlock();\r
56 \r
57     // write an empty block (as EOF marker)\r
58     int blockLength = DeflateBlock();\r
59     fwrite(CompressedBlock, 1, blockLength, Stream);\r
60     \r
61     // flush and close\r
62     fflush(Stream);\r
63     fclose(Stream);\r
64 }\r
65 \r
66 // compresses the current block\r
67 int BgzfData::DeflateBlock(void) {\r
68 \r
69     // initialize the gzip header\r
70     char* buffer = CompressedBlock;\r
71     memset(buffer, 0, 18);\r
72     buffer[0]  = GZIP_ID1;\r
73     buffer[1]  = (char)GZIP_ID2;\r
74     buffer[2]  = CM_DEFLATE;\r
75     buffer[3]  = FLG_FEXTRA;\r
76     buffer[9]  = (char)OS_UNKNOWN;\r
77     buffer[10] = BGZF_XLEN;\r
78     buffer[12] = BGZF_ID1;\r
79     buffer[13] = BGZF_ID2;\r
80     buffer[14] = BGZF_LEN;\r
81 \r
82     // loop to retry for blocks that do not compress enough\r
83     int inputLength = BlockOffset;\r
84     int compressedLength = 0;\r
85     unsigned int bufferSize = CompressedBlockSize;\r
86 \r
87     while(true) {\r
88         \r
89         // initialize zstream values\r
90         z_stream zs;\r
91         zs.zalloc    = NULL;\r
92         zs.zfree     = NULL;\r
93         zs.next_in   = (Bytef*)UncompressedBlock;\r
94         zs.avail_in  = inputLength;\r
95         zs.next_out  = (Bytef*)&buffer[BLOCK_HEADER_LENGTH];\r
96         zs.avail_out = bufferSize - BLOCK_HEADER_LENGTH - BLOCK_FOOTER_LENGTH;\r
97 \r
98         // initialize the zlib compression algorithm\r
99         if(deflateInit2(&zs, Z_DEFAULT_COMPRESSION, Z_DEFLATED, GZIP_WINDOW_BITS, Z_DEFAULT_MEM_LEVEL, Z_DEFAULT_STRATEGY) != Z_OK) {\r
100             printf("BGZF ERROR: zlib deflate initialization failed.\n");\r
101             exit(1);\r
102         }\r
103 \r
104         // compress the data\r
105         int status = deflate(&zs, Z_FINISH);\r
106         if(status != Z_STREAM_END) {\r
107 \r
108             deflateEnd(&zs);\r
109 \r
110             // reduce the input length and try again\r
111             if(status == Z_OK) {\r
112                 inputLength -= 1024;\r
113                 if(inputLength < 0) {\r
114                     printf("BGZF ERROR: input reduction failed.\n");\r
115                     exit(1);\r
116                 }\r
117                 continue;\r
118             }\r
119 \r
120             printf("BGZF ERROR: zlib::deflateEnd() failed.\n");\r
121             exit(1);\r
122         }\r
123 \r
124         // finalize the compression routine\r
125         if(deflateEnd(&zs) != Z_OK) {\r
126             printf("BGZF ERROR: zlib::deflateEnd() failed.\n");\r
127             exit(1);\r
128         }\r
129 \r
130         compressedLength = zs.total_out;\r
131         compressedLength += BLOCK_HEADER_LENGTH + BLOCK_FOOTER_LENGTH;\r
132         if(compressedLength > MAX_BLOCK_SIZE) {\r
133             printf("BGZF ERROR: deflate overflow.\n");\r
134             exit(1);\r
135         }\r
136 \r
137         break;\r
138     }\r
139 \r
140     // store the compressed length\r
141     BgzfData::PackUnsignedShort(&buffer[16], (unsigned short)(compressedLength - 1));\r
142 \r
143     // store the CRC32 checksum\r
144     unsigned int crc = crc32(0, NULL, 0);\r
145     crc = crc32(crc, (Bytef*)UncompressedBlock, inputLength);\r
146     BgzfData::PackUnsignedInt(&buffer[compressedLength - 8], crc);\r
147     BgzfData::PackUnsignedInt(&buffer[compressedLength - 4], inputLength);\r
148 \r
149     // ensure that we have less than a block of data left\r
150     int remaining = BlockOffset - inputLength;\r
151     if(remaining > 0) {\r
152         if(remaining > inputLength) {\r
153             printf("BGZF ERROR: after deflate, remainder too large.\n");\r
154             exit(1);\r
155         }\r
156         memcpy(UncompressedBlock, UncompressedBlock + inputLength, remaining);\r
157     }\r
158 \r
159     BlockOffset = remaining;\r
160     return compressedLength;\r
161 }\r
162 \r
163 // flushes the data in the BGZF block\r
164 void BgzfData::FlushBlock(void) {\r
165 \r
166     // flush all of the remaining blocks\r
167     while(BlockOffset > 0) {\r
168 \r
169         // compress the data block\r
170         int blockLength = DeflateBlock();\r
171 \r
172         // flush the data to our output stream\r
173         int numBytesWritten = fwrite(CompressedBlock, 1, blockLength, Stream);\r
174 \r
175         if(numBytesWritten != blockLength) {\r
176           printf("BGZF ERROR: expected to write %u bytes during flushing, but wrote %u bytes.\n", blockLength, numBytesWritten);\r
177           exit(1);\r
178         }\r
179               \r
180         BlockAddress += blockLength;\r
181     }\r
182 }\r
183 \r
184 // de-compresses the current block\r
185 int BgzfData::InflateBlock(const int& blockLength) {\r
186 \r
187     // Inflate the block in m_BGZF.CompressedBlock into m_BGZF.UncompressedBlock\r
188     z_stream zs;\r
189     zs.zalloc    = NULL;\r
190     zs.zfree     = NULL;\r
191     zs.next_in   = (Bytef*)CompressedBlock + 18;\r
192     zs.avail_in  = blockLength - 16;\r
193     zs.next_out  = (Bytef*)UncompressedBlock;\r
194     zs.avail_out = UncompressedBlockSize;\r
195 \r
196     int status = inflateInit2(&zs, GZIP_WINDOW_BITS);\r
197     if (status != Z_OK) {\r
198         printf("BGZF ERROR: could not decompress block - zlib::inflateInit() failed\n");\r
199         return -1;\r
200     }\r
201 \r
202     status = inflate(&zs, Z_FINISH);\r
203     if (status != Z_STREAM_END) {\r
204         inflateEnd(&zs);\r
205         printf("BGZF ERROR: could not decompress block - zlib::inflate() failed\n");\r
206         return -1;\r
207     }\r
208 \r
209     status = inflateEnd(&zs);\r
210     if (status != Z_OK) {\r
211         printf("BGZF ERROR: could not decompress block - zlib::inflateEnd() failed\n");\r
212         return -1;\r
213     }\r
214 \r
215     return zs.total_out;\r
216 }\r
217 \r
218 // opens the BGZF file for reading (mode is either "rb" for reading, or "wb" for writing)\r
219 bool BgzfData::Open(const string& filename, const char* mode) {\r
220 \r
221         // determine open mode\r
222     if ( strcmp(mode, "rb") == 0 ) {\r
223         IsWriteOnly = false;\r
224     } else if ( strcmp(mode, "wb") == 0) {\r
225         IsWriteOnly = true;\r
226     } else {\r
227         printf("BGZF ERROR: unknown file mode: %s\n", mode);\r
228         return false; \r
229     }\r
230 \r
231     // open Stream to read to/write from file, stdin, or stdout\r
232     // stdin/stdout option contributed by Aaron Quinlan (2010-Jan-03)\r
233     if ( (filename != "stdin") && (filename != "stdout") ) {\r
234         // read/write BGZF data to/from a file\r
235 //         Stream = fopen64(filename.c_str(), mode);\r
236         Stream = fopen(filename.c_str(), mode);\r
237     }\r
238     else if ( (filename == "stdin") && (strcmp(mode, "rb") == 0 ) ) { \r
239         // read BGZF data from stdin\r
240 //         Stream = freopen64(NULL, mode, stdin);\r
241         Stream = freopen(NULL, mode, stdin);\r
242     }\r
243     else if ( (filename == "stdout") && (strcmp(mode, "wb") == 0) ) { \r
244         // write BGZF data to stdout\r
245 //         Stream = freopen64(NULL, mode, stdout);\r
246         Stream = freopen(NULL, mode, stdout);\r
247     }\r
248 \r
249     if(!Stream) {\r
250         printf("BGZF ERROR: unable to open file %s\n", filename.c_str() );\r
251         return false;\r
252     }\r
253     \r
254     // set flag, return success\r
255     IsOpen = true;\r
256     return true;\r
257 }\r
258 \r
259 // reads BGZF data into a byte buffer\r
260 int BgzfData::Read(char* data, const unsigned int dataLength) {\r
261 \r
262    if (dataLength == 0) return 0;\r
263 \r
264    char* output = data;\r
265    unsigned int numBytesRead = 0;\r
266    while (numBytesRead < dataLength) {\r
267 \r
268        int bytesAvailable = BlockLength - BlockOffset;\r
269        if ( bytesAvailable <= 0 ) {\r
270            if (!ReadBlock()) return -1; \r
271            bytesAvailable = BlockLength - BlockOffset;\r
272            if (bytesAvailable <= 0) break;\r
273        }\r
274 \r
275        char* buffer   = UncompressedBlock;\r
276        int copyLength = min( (int)(dataLength-numBytesRead), bytesAvailable );\r
277        memcpy(output, buffer + BlockOffset, copyLength);\r
278 \r
279        BlockOffset  += copyLength;\r
280        output       += copyLength;\r
281        numBytesRead += copyLength;\r
282    }\r
283 \r
284    if ( BlockOffset == BlockLength ) {\r
285        BlockAddress = ftell64(Stream);\r
286        BlockOffset  = 0;\r
287        BlockLength  = 0;\r
288    }\r
289 \r
290    return numBytesRead;\r
291 }\r
292 \r
293 // reads a BGZF block\r
294 bool BgzfData::ReadBlock(void) {\r
295 \r
296     char    header[BLOCK_HEADER_LENGTH];\r
297     int64_t blockAddress = ftell64(Stream);\r
298     \r
299     int count = fread(header, 1, sizeof(header), Stream);\r
300     if (count == 0) {\r
301         BlockLength = 0;\r
302         return true;\r
303     }\r
304 \r
305     if (count != sizeof(header)) {\r
306         printf("BGZF ERROR: read block failed - could not read block header\n");\r
307         return false;\r
308     }\r
309 \r
310     if (!BgzfData::CheckBlockHeader(header)) {\r
311         printf("BGZF ERROR: read block failed - invalid block header\n");\r
312         return false;\r
313     }\r
314 \r
315     int blockLength = BgzfData::UnpackUnsignedShort(&header[16]) + 1;\r
316     char* compressedBlock = CompressedBlock;\r
317     memcpy(compressedBlock, header, BLOCK_HEADER_LENGTH);\r
318     int remaining = blockLength - BLOCK_HEADER_LENGTH;\r
319 \r
320     count = fread(&compressedBlock[BLOCK_HEADER_LENGTH], 1, remaining, Stream);\r
321     if (count != remaining) {\r
322         printf("BGZF ERROR: read block failed - could not read data from block\n");\r
323         return false;\r
324     }\r
325 \r
326     count = InflateBlock(blockLength);\r
327     if (count < 0) { \r
328       printf("BGZF ERROR: read block failed - could not decompress block data\n");\r
329       return false;\r
330     }\r
331 \r
332     if ( BlockLength != 0 )\r
333         BlockOffset = 0;\r
334 \r
335     BlockAddress = blockAddress;\r
336     BlockLength  = count;\r
337     return true;\r
338 }\r
339 \r
340 // seek to position in BGZF file\r
341 bool BgzfData::Seek(int64_t position) {\r
342 \r
343     int     blockOffset  = (position & 0xFFFF);\r
344     int64_t blockAddress = (position >> 16) & 0xFFFFFFFFFFFFLL;\r
345 \r
346     if (fseek64(Stream, blockAddress, SEEK_SET) != 0) {\r
347         printf("BGZF ERROR: unable to seek in file\n");\r
348         return false;\r
349     }\r
350 \r
351     BlockLength  = 0;\r
352     BlockAddress = blockAddress;\r
353     BlockOffset  = blockOffset;\r
354     return true;\r
355 }\r
356 \r
357 // get file position in BGZF file\r
358 int64_t BgzfData::Tell(void) {\r
359     return ( (BlockAddress << 16) | (BlockOffset & 0xFFFF) );\r
360 }\r
361 \r
362 // writes the supplied data into the BGZF buffer\r
363 unsigned int BgzfData::Write(const char* data, const unsigned int dataLen) {\r
364 \r
365     // initialize\r
366     unsigned int numBytesWritten = 0;\r
367     const char* input = data;\r
368     unsigned int blockLength = UncompressedBlockSize;\r
369 \r
370     // copy the data to the buffer\r
371     while(numBytesWritten < dataLen) {\r
372       \r
373         unsigned int copyLength = min(blockLength - BlockOffset, dataLen - numBytesWritten);\r
374         char* buffer = UncompressedBlock;\r
375         memcpy(buffer + BlockOffset, input, copyLength);\r
376 \r
377         BlockOffset     += copyLength;\r
378         input           += copyLength;\r
379         numBytesWritten += copyLength;\r
380 \r
381         if(BlockOffset == blockLength)\r
382             FlushBlock();\r
383     }\r
384 \r
385     return numBytesWritten;\r
386 }\r