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"

/* 多线程解压 */
static void mtUncompressBlock(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];

    auto& blockArr = p.threadBlocksWrap.threadBlocks[tid];
    auto& blockItem = blockArr.add();
    uint8_t* block = readData.startAddrArr[idx];

    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(blockItem.data, &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);
    }
    blockItem.blockId = idx + p.startBlockId;
    blockItem.blockLen = dlen;
    blockArr.blockHeap.push({blockItem.blockId, blockArr.curIdx - 1});  // 解压完成后，将block的id和在block数组里的索引加入堆中，方便后续排序和合并

#if 0
    // 放入全局缓冲区
    // spdlog::info("top id: {}, block id: {}", blockArr.blockHeap.top().blockId, p.uncompressData.nextBlockId);
    while (blockArr.blockHeap.top().blockId == p.uncompressData.nextBlockId) {
        auto& top = blockArr.blockHeap.top();
        // auto& topBlock = blockArr.blockArr[top.blockArrIdx];
        // memcpy(p.uncompressData.dataBuf + p.uncompressData.usedBufSize, topBlock.data, topBlock.blockLen);
        // p.uncompressData.startAddrArr.push_back(p.uncompressData.dataBuf + p.uncompressData.usedBufSize);
        // p.uncompressData.usedBufSize += topBlock.blockLen;
        // p.bamNum += topBlock.bamNum;

        blockArr.blockHeap.pop();
        p.uncompressData.nextBlockId += 1;
        p.uncompressData.blockNum += 1;
    }
#endif

    PROF_T_END(tid, mem_copy);
}

// 多线程解压，静态分配任务，此时用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.threadBlocksWrap.threadBlockBuf[tid];

    if (stopIdx - startIdx > blockBuf.maxLen / SINGLE_BLOCK_SIZE) {
        blockBuf.reAllocMem((stopIdx - startIdx) * SINGLE_BLOCK_SIZE);
    }

    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;
    }
    PROF_T_END(tid, mem_copy);
}

static void mtMemCopy(void* data, long idx, int tid) {
    Phase1PipelineArg& p = *(Phase1PipelineArg*)data;
    tid = idx; // 静态分配任务，此时用idx代替tid
    uint64_t offset = 0;
    for (int i = 0; i < tid; ++i) {
        offset += p.threadBlocksWrap.threadBlockBuf[i].curLen;
    }
    memcpy(p.uncompressData.dataBuf + p.uncompressData.usedBufSize + offset, p.threadBlocksWrap.threadBlockBuf[tid].data,
           p.threadBlocksWrap.threadBlockBuf[tid].curLen);
}

/* 将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();
    // kt_for(p.numThread, mtUncompressBlock, &p, blockNum);
    kt_for(p.numThread, mtUncompressBlockBatch, &p, p.numThread);
    // 串行拷贝所有blocks
    PROF_G_BEG(mem_copy);
#if 1
    kt_for(p.numThread, mtMemCopy, &p, p.numThread);
#else
    for (int i = 0; i < p.numThread; ++i) {
        memcpy(p.uncompressData.dataBuf + p.uncompressData.usedBufSize, p.threadBlocksWrap.threadBlockBuf[i].data, p.threadBlocksWrap.threadBlockBuf[i].curLen);
        // p.uncompressData.startAddrArr.push_back(p.uncompressData.dataBuf + p.uncompressData.usedBufSize);
        p.uncompressData.usedBufSize += p.threadBlocksWrap.threadBlockBuf[i].curLen;
    }
#endif

#if 0
    for (int i = 0; i < 1; ++i) {
        auto& blockArr = p.threadBlocksWrap.threadBlocks[i];
        for (int j = 0; j < blockArr.curIdx; ++j) {
            auto& blockItem = blockArr.blockArr[j];
            memcpy(p.uncompressData.dataBuf + p.uncompressData.usedBufSize, blockItem.data, blockItem.blockLen);
            p.uncompressData.startAddrArr.push_back(p.uncompressData.dataBuf + p.uncompressData.usedBufSize);
            p.uncompressData.usedBufSize += blockItem.blockLen;
            p.uncompressData.blockNum += 1;
        }
    }
#endif

    PROF_G_END(mem_copy);

    p.startBlockId += blockNum;

    if (true) { // 缓冲区满了
        spdlog::info("blocks num: {}, left: {}, uncompressed: {}", p.threadBlocksWrap.GetTotalBlockNum(), p.threadBlocksWrap.GetHeapBlockNum(),
                     p.uncompressData.blockNum);
        p.uncompressData.Clear();
        p.threadBlocksWrap.ResetBlockArr();
        p.uncompressData.nextBlockId = p.startBlockId;
        p.uncompressData.blockNum = 0;
    }

    PROF_G_END(uncompress);
}

/* phase1Uncompress step-2 解压线程 */
void* phase1Uncompress(void* data) {
    Phase1PipelineArg& p = *(Phase1PipelineArg*)data;

    /* 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;
        }

        doPhase1Uncompress(p);

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

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