FastSort/src/sort/sam_io.cpp

#include "sam_io.h"

#include <htslib/sam.h>
#include <libdeflate.h>
#include <stdint.h>
#include <stdio.h>
#include <zlib.h>
#include <spdlog/spdlog.h>

#include "const_val.h"

namespace nsgv {
extern bool gIsBigEndian;
};

static const uint8_t g_magic[19] = "\037\213\010\4\0\0\0\0\0\377\6\0\102\103\2\0\0\0";

size_t fileSize(FILE *fp) {
    spdlog::info("test file start");
    size_t bufSize = 64L * 1024 * 1024;
    uint8_t *testBuf = (uint8_t *)malloc(bufSize);
    size_t readState = 0;
    size_t fileSize = 0;
    while ((readState = fread(testBuf, 1, bufSize, fp)) > 0) {
        fileSize += readState;
    }
    spdlog::info("test file size: {}", fileSize);
    return fileSize;
}

static inline void packInt16(uint8_t* buffer, uint16_t value) {
    buffer[0] = value;
    buffer[1] = value >> 8;
}

static inline void packInt32(uint8_t* buffer, uint32_t value) {
    buffer[0] = value;
    buffer[1] = value >> 8;
    buffer[2] = value >> 16;
    buffer[3] = value >> 24;
}

// dlen本身长度一定得够存放压缩后的数据，否则会返回错误
int bgzfCompress(void* _dst, size_t* dlen, const void* src, size_t slen, int level) {
    // fprintf(stderr, "bgzfCompress, slen: %zu\n", slen);
    if (slen == 0) {
        // EOF block
        if (*dlen < 28)
            return -1;
        memcpy(_dst, "\037\213\010\4\0\0\0\0\0\377\6\0\102\103\2\0\033\0\3\0\0\0\0\0\0\0\0\0", 28);
        *dlen = 28;
        return 0;
    }

    uint8_t* dst = (uint8_t*)_dst;

    if (level == 0) {
        // Uncompressed data
        if (*dlen < slen + 5 + BLOCK_HEADER_LENGTH + BLOCK_FOOTER_LENGTH)
            return -1;
        dst[BLOCK_HEADER_LENGTH] = 1;                     // BFINAL=1, BTYPE=00; see RFC1951
        u16_to_le(slen, &dst[BLOCK_HEADER_LENGTH + 1]);   // length
        u16_to_le(~slen, &dst[BLOCK_HEADER_LENGTH + 3]);  // ones-complement length
        memcpy(dst + BLOCK_HEADER_LENGTH + 5, src, slen);
        *dlen = slen + 5 + BLOCK_HEADER_LENGTH + BLOCK_FOOTER_LENGTH;

    } else {
        level = level > 0 ? level : 6;  // libdeflate doesn't honour -1 as default
        // NB levels go up to 12 here.
        int lvl_map[] = {0, 1, 2, 3, 5, 6, 7, 8, 10, 12};
        level = lvl_map[level > 9 ? 9 : level];
        struct libdeflate_compressor* z = libdeflate_alloc_compressor(level);
        if (!z)
            return -1;

        // Raw deflate
        size_t clen = libdeflate_deflate_compress(z, src, slen, dst + BLOCK_HEADER_LENGTH, *dlen - BLOCK_HEADER_LENGTH - BLOCK_FOOTER_LENGTH);

        if (clen <= 0) {
            hts_log_error("Call to libdeflate_deflate_compress failed");
            libdeflate_free_compressor(z);
            return -1;
        }

        *dlen = clen + BLOCK_HEADER_LENGTH + BLOCK_FOOTER_LENGTH;

        libdeflate_free_compressor(z);
    }

    // write the header
    memcpy(dst, g_magic, BLOCK_HEADER_LENGTH);  // the last two bytes are a place holder for the length of the block
    packInt16(&dst[16], *dlen - 1);             // write the compressed length; -1 to fit 2 bytes

    // write the footer
    uint32_t crc = libdeflate_crc32(0, src, slen);
    packInt32((uint8_t*)&dst[*dlen - 8], crc);
    packInt32((uint8_t*)&dst[*dlen - 4], slen);
    return 0;
}

int bgzfUncompress(uint8_t *dst, size_t *dlen, const uint8_t *src, size_t slen, uint32_t expected_crc) {
    struct libdeflate_decompressor *z = libdeflate_alloc_decompressor();
    if (!z) {
        hts_log_error("Call to libdeflate_alloc_decompressor failed");
        return -1;
    }

    int ret = libdeflate_deflate_decompress(z, src, slen, dst, *dlen, dlen);
    libdeflate_free_decompressor(z);

    if (ret != LIBDEFLATE_SUCCESS) {
        hts_log_error("Inflate operation failed: %d", ret);
        return -1;
    }

    uint32_t crc = libdeflate_crc32(0, (unsigned char *)dst, *dlen);
#ifdef FUZZING_BUILD_MODE_UNSAFE_FOR_PRODUCTION
    // Pretend the CRC was OK so the fuzzer doesn't have to get it right
    crc = expected_crc;
#endif
    if (crc != expected_crc) {
        hts_log_error("CRC32 checksum mismatch");
        return -2;
    }

    return 0;
}

// 读取并解压一个gz block
size_t readUncompressOneBlock(FILE *fpr, uint8_t *fBuf, DataBuffer *uDataPtr) {
    DataBuffer &uData = *uDataPtr;
    size_t readState = 0;
    int blockLen = 0;                 // 每个gz block的大小
    size_t dlen = SINGLE_BLOCK_SIZE;  // gz block解压之后的真正大小
    readState = fread(fBuf, 1, BLOCK_HEADER_LENGTH, fpr);  // 先读取一个gz block的头部
    blockLen = unpackInt16(&fBuf[16]) + 1;                // block的字节数
    readState = fread(&fBuf[BLOCK_HEADER_LENGTH], 1, blockLen - BLOCK_HEADER_LENGTH, fpr);
    // 解压gz block
    uint32_t crc = le_to_u32(fBuf + blockLen - 8);
    size_t newDataSize = uData.maxLen;
    while (uData.curLen + SINGLE_BLOCK_SIZE > newDataSize) newDataSize *= 2;
    uData.ReAllocMem(newDataSize);  // 需要重新开辟空间
    int ret = bgzfUncompress(&uData.data[uData.curLen], &dlen, (Bytef *)fBuf + BLOCK_HEADER_LENGTH,
                         blockLen - BLOCK_HEADER_LENGTH, crc);
    if (ret < 0) {
        spdlog::error("Error decompressing gz block");
        exit(0);
    }
    uData.curLen += dlen;
    spdlog::info("gz block size: {}, uncompressed size: {}", blockLen, dlen);
    return readState;
}

// 解析sam头部信息
void parseSamHeader(FILE *fpr, HeaderBuf &hdrBuf) {
    const int kMaxBlockSize = SINGLE_BLOCK_SIZE;
    sam_hdr_t *header = NULL;
    uint8_t *fBuf = (uint8_t*)malloc(kMaxBlockSize);
    DataBuffer uData; // 用来放解压缩后的数据，待解析
    int magicLen;
    int32_t i, nameLen, numNames = 0;

    header = sam_hdr_init();       // 初始化header
    uData.AllocMem(kMaxBlockSize); // 初始化解压数据的buffer
    readUncompressOneBlock(fpr, fBuf, &uData); // 读取第一个gz block

    // 解析header
    magicLen = 4;
    if (memcmp(uData.data, "BAM\1", magicLen)) {
        spdlog::error("Invalid BAM binary header");
        return;
    }
    uData.readPos += magicLen; // 在解压数据中读取了magicLen个字节
    header->l_text = le_to_u32(uData.data + uData.readPos); uData.readPos += 4; // 解析l_text的长度
    header->text = (char *)malloc(header->l_text + 1);
    header->text[header->l_text] = 0; // 在text末尾添加'\0'

    while (uData.readPos + header->l_text > uData.curLen) { // header内容超过了一个block，继续读
        readUncompressOneBlock(fpr, fBuf, &uData);          // 读取包含header内容的gz block
    }

    memcpy(header->text, &uData.data[uData.readPos], header->l_text);
    uData.readPos += header->l_text;

    memcpy(&header->n_targets, &uData.data[uData.readPos], 4); // n_targets
    uData.readPos += 4;
    spdlog::info("num target: {}", header->n_targets);
    if (nsgv::gIsBigEndian) ed_swap_4p(&header->n_targets);
    if (header->n_targets > 0) {
        header->target_name = (char **)calloc(header->n_targets, sizeof(char *));
        header->target_len = (uint32_t *)calloc(header->n_targets, sizeof(uint32_t));
    } else {
        header->target_name = NULL;
        header->target_len = NULL;
    }

    for (int i = 0; i != header->n_targets; ++i) {
        while (uData.readPos + 4 > uData.curLen) readUncompressOneBlock(fpr, fBuf, &uData);
        memcpy(&nameLen, &uData.data[uData.readPos], 4);
        uData.readPos += 4;
        if (nsgv::gIsBigEndian) ed_swap_4p(&nameLen);
        header->target_name[i] = (char *)malloc(nameLen);
        ++numNames;
        while (uData.readPos + nameLen > uData.curLen) readUncompressOneBlock(fpr, fBuf, &uData);
        memcpy(header->target_name[i], &uData.data[uData.readPos], nameLen);
        uData.readPos += nameLen;
        if (header->target_name[i][nameLen - 1] != '\0') {
            /* Fix missing NUL-termination.  Is this being too nice?
               We could alternatively bail out with an error. */
            char *newName = (char*)realloc(header->target_name[i], nameLen + 1);
            header->target_name[i] = newName;
            header->target_name[i][nameLen] = '\0';
        }
        while (uData.readPos + 4 > uData.curLen) readUncompressOneBlock(fpr, fBuf, &uData);
        memcpy(&header->target_len[i], &uData.data[uData.readPos], 4);
        uData.readPos += 4;
        if (nsgv::gIsBigEndian) ed_swap_4p(&header->target_len[i]);

        // spdlog::info("nameLen {}, {}, {}", nameLen, header->target_len[i], header->target_name[i]);
    }
    // spdlog::info("test res: {} {} {} {}", header->l_text, dlen, header->n_targets, header->text);

    spdlog::info("udata: readPos: {}, curLen: {}, maxLen: {}", uData.readPos, uData.curLen, uData.maxLen);

    hdrBuf.header = header;
    if (uData.readPos < uData.curLen) { // 如果还有数据，则将其保存在hdrBuf中，用以后续的bam解析
        hdrBuf.data = (uint8_t *)malloc(uData.curLen - uData.readPos);
        memcpy(hdrBuf.data, &uData.data[uData.readPos], uData.curLen - uData.readPos);
        hdrBuf.dataLen = uData.curLen - uData.readPos;
    }

    free(fBuf);
}