FastSort/src/sort/phase_1_uncompress.cpp

/*
     Description: 第一阶段的解压线程

     Copyright : All right reserved by ICT

     Author : Zhang Zhonghai
     Date : 2026/05/25
*/

#include "phase_1_uncompress.h"

#include <klib/kthread.h>
#include <spdlog/spdlog.h>
#include <stdint.h>
#include <zlib.h>

#include "common_data.h"
#include "const_val.h"
#include "phase_1.h"
#include "sam_io.h"
#include "sort.h"
#include "util/profiling.h"
#include "util/yarn.h"

int GetBamLen(uint8_t* dataAddr) {
    uint32_t bamLen = 0;
    memcpy(&bamLen, dataAddr, 4);
    if (nsgv::gIsBigEndian)
        ed_swap_4p(&bamLen);
    return bamLen;
}

void ParseBam(uint8_t* dataAddr, OneBam& bam) {
    uint32_t bams = 0;
    bam.bamLen = GetBamLen(dataAddr);
    dataAddr += 4;
    bam.tid = le_to_u32(dataAddr);
    bam.pos = le_to_i32(dataAddr + 4);
    uint32_t x2 = le_to_u32(dataAddr + 8);
    bam.qnameLen = x2 & 0xff;
}

// 解析一个bam并放入bamArr
void ParseAddBam(uint8_t* dataAddr, BamArr& bamArr) {
    OneBam& bam = bamArr.Add();
    ParseBam(dataAddr, bam);
}

// 返回解析bam的个数
size_t ParseAddAllBams(uint8_t* dataAddr, size_t startOffset, size_t endOffset, BamArr& bamArr, size_t* nextBamStartPtr = nullptr,
                       size_t* lastPosPtr = nullptr) {
    size_t nextBamStart = startOffset;
    size_t lastPos = 0;
    uint32_t bamLen = 0;
    uint32_t bams = 0;

    while (nextBamStart + 4 <= endOffset) {
        uint8_t* curAddr = dataAddr + nextBamStart;
        memcpy(&bamLen, curAddr, 4);
        if (nsgv::gIsBigEndian)
            ed_swap_4p(&bamLen);
        nextBamStart += 4 + bamLen;
        if (nextBamStart == endOffset) {               // 刚好解析到最后，说明这个block的内容都是完整bam
            lastPos = endOffset;                       // 继续解析当前的block
        } else if (nextBamStart > endOffset) {         // 当前bam不完整，不能解析
            lastPos = nextBamStart - (4 + bamLen);  // 记录最后一个不完整的bam起始位置
            break;
        }
        OneBam& bam = bamArr.Add();
        bam.bamLen = bamLen;
        curAddr += 4;
        bam.tid = le_to_u32(curAddr);
        bam.pos = le_to_i32(curAddr + 4);
        uint32_t x2 = le_to_u32(curAddr + 8);
        bam.qnameLen = x2 & 0xff;
        ++bams;
    }
    if (nextBamStart < endOffset) {
        lastPos = nextBamStart;  // 最后一个不完整的bam，连长度都不够数据解析
    }
    if (nextBamStartPtr) *nextBamStartPtr = nextBamStart;
    if (lastPosPtr) *lastPosPtr = lastPos;
    return bams;
}

// 检查当前内存对应的bam是否能正确解析
static bool isValidBam(uint8_t* dataAddr, uint32_t &bamLen, int maxBamLen, int maxSeqLen) {
    memcpy(&bamLen, dataAddr, 4);
    if (nsgv::gIsBigEndian)
        ed_swap_4p(&bamLen);
    int32_t tid = 0;
    int64_t pos = 0;
    int qnameLen = 0;
    uint16_t flag = 0;
    uint8_t* x = dataAddr + 4;
    uint32_t x2 = le_to_u32(x + 8);
    qnameLen = x2 & 0xff;
    if (qnameLen >= bamLen) {
        return false;
    }
    int32_t seqLen = le_to_u32(x + 16);
    if (seqLen >= bamLen) {
        return false;
    }
    tid = le_to_u32(x);
    pos = le_to_i32(x + 4);
    uint32_t x3 = le_to_u32(x + 12);
    flag = x3 >> 16;

    if ((flag & BAM_FUNMAP) && tid == -1 && pos == -1) {
        return true;
    }
    int nref = sam_hdr_nref(nsgv::gInHdr.header);
    if (tid < 0 || tid >= nref) {
        // 非法 tid
        return false;
    }
    hts_pos_t ref_len = sam_hdr_tid2len(nsgv::gInHdr.header, tid);
    if (pos < 0 || pos >= ref_len) {
        // 非法 pos：0 ≤ pos < ref_len
        return false;
    }
    // uint32_t n_cigar = x3 & 0xffff;
    if (seqLen * 10 < maxSeqLen || maxSeqLen * 10 < seqLen) {  // 猜测条件可以仔细考虑下
        //spdlog::info("invalid bam(seqlen), bamlen: {}, seqLen: {}, maxSeqLen: {}, maxSeqLen: {}", bamLen, seqLen, maxSeqLen, maxSeqLen);
        //return false;
    }
    if (bamLen * 10 < maxBamLen || maxBamLen * 10 < bamLen) {  // 猜测条件可以仔细考虑下
        //spdlog::info("invalid bam(bamlen), bamlen: {}, maxBamLen: {}, seqLen: {}, qnameLen: {}, n_cigar: {}", bamLen, maxBamLen, seqLen, qnameLen, n_cigar);
        // return false;
    }
    return true;
}

// 多线程解压，静态分配任务，此时用idx代替tid，multi-thread uncompress bam blocks
static void mtUncompressBlockBatch(void* data, long idx, int tid) {
    PROF_T_BEG(mem_copy);

    Phase1PipelineArg& p = *(Phase1PipelineArg*)data;
    ReadBuffer& readData = p.readData[p.uncompressOrder % p.READ_BUF_NUM];

    tid = idx; // 静态分配任务，此时用idx代替tid
    int startIdx = START_IDX(idx, p.numThread, readData.startAddrArr.size());
    int stopIdx = STOP_IDX(idx, p.numThread, readData.startAddrArr.size());

    auto &blockBuf = p.threadUncompressWrap.threadUncompressDataArr[tid].blockBuf;
    auto &bamArr = p.threadUncompressWrap.threadUncompressDataArr[tid].bamArr;

    // 开辟足够的内存
    if (stopIdx - startIdx > blockBuf.maxLen / SINGLE_BLOCK_SIZE) {
        blockBuf.ReAllocMem((stopIdx - startIdx) * SINGLE_BLOCK_SIZE);
    }

    // 解压block
    for (int i = startIdx; i < stopIdx; ++i) {
        uint8_t* block = readData.startAddrArr[i];
        size_t dlen = SINGLE_BLOCK_SIZE;  // 65535
        int block_length = unpackInt16(&block[16]) + 1;
        uint32_t crc = le_to_u32(block + block_length - 8);
        int ret = bgzfUncompress(blockBuf.data + blockBuf.curLen, &dlen, (Bytef*)block + BLOCK_HEADER_LENGTH, block_length - BLOCK_HEADER_LENGTH, crc);
        if (ret != 0) {
            spdlog::error("uncompress error, block id: {}, len: {}, ret: {}", idx, block_length, ret);
            exit(0);
        }
        blockBuf.curLen += dlen;
    }

    // 用readPos来表示第一个bam起始位置，默认是ngs的bam，如果是三代bam，应该不需要计算这个了，因为三代bam很长
    blockBuf.readPos = 0;
    for (int i = 0; i < SINGLE_BLOCK_SIZE; ++i) {
        uint32_t nextBamStart = i;  // 这个要注意，每次应该要计算一下
        uint32_t bamLen = 0;
        uint8_t nextPos = 0; // 是否检查下一个可能的位置
        while (nextBamStart + 4 <= SINGLE_BLOCK_SIZE) {
            if (isValidBam(blockBuf.data + nextBamStart, bamLen, p.maxBamLen, p.maxSeqLen)) {
                nextBamStart += 4 + bamLen;
            } else {
                nextPos = 1;
                break;
            }
        }
        if (!nextPos) {
            blockBuf.readPos = i;
            break;
        }
    }
    // 解析bam
    ParseAddAllBams(blockBuf.data, blockBuf.readPos, blockBuf.curLen, bamArr, nullptr, &blockBuf.lastPos);
    PROF_T_END(tid, mem_copy);
}

// 处理相邻线程的block数据，可能有bam跨越这两个线程的block（GATK的bam）
static void handleAdjacentThreadBlock(Phase1PipelineArg& p) {
    auto& uncompressData = p.uncompressData;
    auto& threadUncompressDataArr = p.threadUncompressWrap.threadUncompressDataArr;
    size_t offset = 0;  // 当前线程对应的全局数据的起始偏移量
    size_t bamOffset = p.allBams.Size(); // 当前线程解析的bam在全局数据中的偏移量

    for (int tid = 0; tid < p.numThread; ++tid) {
        threadUncompressDataArr[tid].memOffset = offset;
        threadUncompressDataArr[tid].bamOffset = bamOffset;
        auto& blockBuf = threadUncompressDataArr[tid].blockBuf;
        auto& bamArr = threadUncompressDataArr[tid].bamArr;
        auto& firstBam = threadUncompressDataArr[tid].firstBam;
        auto& lastBamBuf = threadUncompressDataArr[tid].lastBamBuf;

        bool hasLastData = false; // 上一个block里有不完整bam数据
        int lastDataLen = 0; // 上一个block里不完整bam数据的长度
        int leftDataLen = blockBuf.readPos;  // 本轮剩余的不完整的bam数据

        if (tid == 0) { // 第一个线程
            // 检查一下bam的定位是否正确
            hasLastData = uncompressData.usedBufSize != uncompressData.lastEndPos;
            if (hasLastData || blockBuf.readPos > 0) {
                lastDataLen = uncompressData.usedBufSize - uncompressData.lastEndPos;
                leftDataLen = blockBuf.readPos;  // 本轮剩余的不完整的bam数据
                lastBamBuf.MemCopy(uncompressData.dataBuf + uncompressData.lastEndPos, lastDataLen);
                lastBamBuf.MemCopy(blockBuf.data, leftDataLen);
            }
        } else {
            hasLastData = threadUncompressDataArr[tid - 1].blockBuf.curLen != threadUncompressDataArr[tid - 1].blockBuf.lastPos;
            if (hasLastData || blockBuf.readPos > 0) {  // 上一轮有遗留数据
                lastDataLen = threadUncompressDataArr[tid - 1].blockBuf.curLen - threadUncompressDataArr[tid - 1].blockBuf.lastPos;
                leftDataLen = blockBuf.readPos;  // 本轮剩余的不完整的bam数据
                lastBamBuf.MemCopy(threadUncompressDataArr[tid - 1].blockBuf.data + threadUncompressDataArr[tid - 1].blockBuf.lastPos, lastDataLen);
                lastBamBuf.MemCopy(blockBuf.data, leftDataLen);
            }
        }
        if (hasLastData || blockBuf.readPos > 0) {
            int bamLen = GetBamLen(lastBamBuf.data);
            int claculatedBamLen = lastDataLen + leftDataLen - 4;
            int additionDataLen = 0;
            if (bamLen != claculatedBamLen) {         // 猜测错了
                if (lastDataLen + leftDataLen < 4) {  // 不够解析bam长度
                    lastBamBuf.MemCopy(blockBuf.data + leftDataLen, 4);
                    leftDataLen += 4;
                    bamLen = GetBamLen(lastBamBuf.data);  // 真正的长度
                }
                additionDataLen = bamLen + 4 - (lastDataLen + leftDataLen);  // 还缺多少数据能解析出完整的bam
                if (additionDataLen > 0) {
                    lastBamBuf.MemCopy(blockBuf.data + leftDataLen, additionDataLen);
                    leftDataLen += additionDataLen;
                }
                // 重新解析
                bamArr.Clear();
                ParseAddAllBams(blockBuf.data, leftDataLen, blockBuf.curLen, bamArr, nullptr, &blockBuf.lastPos);
                spdlog::error("bam len mismatch {}: {}, {}", tid, bamLen, claculatedBamLen);
            }
            ParseAddBam(lastBamBuf.data, firstBam);
        }
        offset += threadUncompressDataArr[tid].blockBuf.curLen;
        bamOffset += threadUncompressDataArr[tid].bamArr.Size() + threadUncompressDataArr[tid].firstBam.Size();
    }
}

static void mtMemCopy(void* data, long idx, int tid) {
    Phase1PipelineArg& p = *(Phase1PipelineArg*)data;
    tid = idx; // 静态分配任务，此时用idx代替tid

    auto& threadUncompressDataArr = p.threadUncompressWrap.threadUncompressDataArr;  // 每个thread一个，用来保存解压后的block数据
    auto &uncompressData = p.uncompressData;              // 所有线程共用一个，串行往这里添加解压后的block数据

    // 拷贝bam未解析数据到全局的uncompressData里
    memcpy(uncompressData.dataBuf + uncompressData.usedBufSize + threadUncompressDataArr[tid].memOffset, threadUncompressDataArr[tid].blockBuf.data,
           threadUncompressDataArr[tid].blockBuf.curLen);
           
    // 拷贝解析的bam到全局数据里
    size_t i = 0;
    for (; i < threadUncompressDataArr[tid].firstBam.Size(); ++i) {
        p.allBams.arr[i + threadUncompressDataArr[tid].bamOffset] = threadUncompressDataArr[tid].firstBam.arr[i];
    }
    for (size_t j = 0; j < threadUncompressDataArr[tid].bamArr.Size(); ++i, ++j) {
        p.allBams.arr[i + threadUncompressDataArr[tid].bamOffset] = threadUncompressDataArr[tid].bamArr.arr[j];
    }

    if (tid == p.numThread - 1) { // 最后一个线程，更新全局uncompressData的usedBufSize
        uncompressData.usedBufSize += threadUncompressDataArr[tid].memOffset + threadUncompressDataArr[tid].blockBuf.curLen;
        uncompressData.lastEndPos = uncompressData.usedBufSize - (threadUncompressDataArr[tid].blockBuf.curLen - threadUncompressDataArr[tid].blockBuf.lastPos);
        p.allBams.curIdx += threadUncompressDataArr[tid].bamOffset + threadUncompressDataArr[tid].bamArr.Size() + threadUncompressDataArr[tid].firstBam.Size();
    }
}

/* 将gz block进行解压，并进行线程内排序 */
static void doPhase1Uncompress(Phase1PipelineArg& p, int finish = 0) {
    PROF_G_BEG(uncompress);
    uint64_t blockNum = p.readData[p.uncompressOrder % p.READ_BUF_NUM].startAddrArr.size();
    p.blockNum += blockNum;

    kt_for(p.numThread, mtUncompressBlockBatch, &p, p.numThread);
    PROF_G_END(uncompress);

    // 并行拷贝所有blocks
    PROF_G_BEG(mem_copy);
    handleAdjacentThreadBlock(p);
    p.allBams.Add(p.threadUncompressWrap.GetTotalBamNum());
    p.bamNum += p.threadUncompressWrap.GetTotalBamNum();
    kt_for(p.numThread, mtMemCopy, &p, p.numThread);
    PROF_G_END(mem_copy);

    PROF_G_BEG(parse_block);
    PROF_G_END(parse_block);

    if (true) { // 缓冲区满了
        spdlog::info("blocks num: {}, uncompressed: {}, bam num: {}, all bam num: {}, zero start blocks: {}", blockNum, p.blockNum,
                     p.threadUncompressWrap.GetTotalBamNum(), p.bamNum, p.zeroStartBlockNum);
        // p.uncompressData.Clear();
        p.uncompressData.NextRound();
        // spdlog::info("last data - 0: {}", p.uncompressData.usedBufSize - p.uncompressData.lastEndPos);
        p.threadUncompressWrap.ResetBlockArr();
        //for (size_t i = 0; i < p.allBams.Size(); ++i) {
        //    fprintf(gfp[0], "%d-%ld\n", p.allBams.arr[i].tid, p.allBams.arr[i].pos);
        //}
        p.allBams.Clear();
    }

}

/* phase1Uncompress step-2 解压线程 */
void* phase1Uncompress(void* data) {
    Phase1PipelineArg& p = *(Phase1PipelineArg*)data;
    int parseFirstBlock = 1;
    /* 2. do the work */
    while (true) {
        // previous dependency
        yarn::DEPENDENCY_NOT_TO_BE(p.readSig, 0);

        if (p.readFinish) {
            while (p.uncompressOrder < p.readOrder) {
                doPhase1Uncompress(p, 1);
                p.uncompressOrder += 1;
            }
            break;
        }
        if (parseFirstBlock) {
            parseFirstBlock = 0;
            // 计算bam的平均长度，以及第一个block里的bam个数，用来指导后续的解压和排序
            uint8_t* block = p.readData[p.uncompressOrder % p.READ_BUF_NUM].startAddrArr[0];
            size_t dlen = SINGLE_BLOCK_SIZE;  // 65535
            int block_length = unpackInt16(&block[16]) + 1;
            uint32_t crc = le_to_u32(block + block_length - 8);
            uint8_t oneBlock[SINGLE_BLOCK_SIZE];
            int ret = bgzfUncompress(oneBlock, &dlen, (Bytef*)block + BLOCK_HEADER_LENGTH, block_length - BLOCK_HEADER_LENGTH, crc);
            if (ret != 0) {
                spdlog::error("First block uncompress error, len: {}, ret: {}", block_length, ret);
                exit(0);
            }
            uint64_t nextBamStart = 0;  // 第一个block的起始bam位置
            uint32_t bamLen = 0;
            uint64_t allBamLen = 0;
            uint64_t bamNum = 0;
            /* 解析每个bam */
            while (nextBamStart + 4 <= dlen) {
                bamLen = GetBamLen(oneBlock + nextBamStart);
                p.maxBamLen = p.maxBamLen < bamLen ? bamLen : p.maxBamLen;
                uint8_t* x = oneBlock + nextBamStart + 4;
                int32_t seqLen = le_to_u32(x + 16);
                p.maxSeqLen = p.maxSeqLen < seqLen ? seqLen : p.maxSeqLen;

                nextBamStart += 4 + bamLen;
                allBamLen += bamLen;
                ++bamNum;
            }
            p.uncompressData.avgBamSize = bamNum == 0 ? 0 : allBamLen / bamNum;
            p.uncompressData.avgBamNumPerBlock = bamNum;
            spdlog::info("avg bam size: {}, avg bam num per block: {}, max bam len: {}, max seq len: {}", p.uncompressData.avgBamSize, p.uncompressData.avgBamNumPerBlock, p.maxBamLen, p.maxSeqLen);
        }
        doPhase1Uncompress(p);

        // update status
        yarn::CONSUME_SIGNAL(p.readSig);
        p.uncompressOrder += 1;
    }

    spdlog::info("uncompress order: {}", p.uncompressOrder);
    return nullptr;
}