FastSort/src/sort/phase_1_compress.cpp

/*
     Description: 第一阶段的压缩

     Copyright : All right reserved by ICT

     Author : Zhang Zhonghai
     Date : 2026/06/02
*/

#include "phase_1_compress.h"

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

#include <algorithm>
#include <string>

#include "common_data.h"
#include "const_val.h"
#include "phase_1.h"
#include "phase_1_compress.h"
#include "phase_1_write.h"
#include "sort.h"
#include "util/profiling.h"

// for test
uint64_t block_num = 0;
uint64_t bam_num = 0;

void CheckBam(uint8_t *addr, int len) {
    int bamLen = 0;
    memcpy(&bamLen, addr, 4);
    if (nsgv::gIsBigEndian)
        ed_swap_4p(&bamLen);
    if (bamLen + 4 != len) {
        spdlog::error("bam error: {}-{}", bamLen, len);
        exit(0);
    }
}

static void mtCompressBlock(void* data, long idx, int tid) {
    Phase1PipelineArg& p = *(Phase1PipelineArg*)data;
    MergeCompressData& mergeCompressData = p.mergeCompressData[p.compressOrder % p.COMPRESS_BUF_NUM];
    auto& bams = mergeCompressData.blockDataArr[idx].bamPtrArr;
    auto& blockData = mergeCompressData.blockDataArr[idx].blockBuf;
    auto& compressData = mergeCompressData.compressDataArr[idx];

    // spdlog::info("bam size: {}", bams.Size());
    bam_num += bams.Size();
    for (int i = 0; i < bams.Size(); ++i) {
        const OneBam* bp = bams.arr[i];
        blockData.MemCopy(p.uncompressData.dataBuf + bp->offset, bp->wholeBamLen);
        // check bam
        // CheckBam(p.uncompressData.dataBuf + bp->offset, bp->wholeBamLen);
    }
    compressData.ReAllocMem(SINGLE_BLOCK_SIZE);  // 压缩后的block数据不会超过单个block的大小
    compressData.curLen = SINGLE_BLOCK_SIZE;
    // spdlog::info("block bytes: {}", blockData.curLen);
    bgzfCompress(compressData.data, &compressData.curLen, blockData.data, blockData.curLen, p.compressLevel);
    //spdlog::info("block bytes: {}, compressed bytes: {}", blockData.curLen, compressData.curLen);
}

static void doCompress(Phase1PipelineArg& p) {
    PROF_G_BEG(compress);
    DataBuffer& compressBuf = p.compressBuf[p.compressOrder % p.COMPRESS_BUF_NUM];
    MergeCompressData& mergeCompressData = p.mergeCompressData[p.compressOrder % p.MERGE_BUF_NUM];
    kt_for(p.numThread, mtCompressBlock, &p, mergeCompressData.Size());
    // kt_for(1, mtCompressBlock, &p, mergeCompressData..Size());
    PROF_G_END(compress);

    compressBuf.Clear();
    for (int i = 0; i < mergeCompressData.Size(); ++i) {  // 如果太慢，可以考虑并行拷贝，先计算偏移量，然后多线程拷贝
        compressBuf.MemCopy(mergeCompressData.compressDataArr[i].data, mergeCompressData.compressDataArr[i].curLen);
    }
    block_num += mergeCompressData.Size();
    // spdlog::info("block num: {}, {}, bam num: {}", mergeCompressData.Size(), block_num, bam_num);
}

/* phase1Compress step- 压缩线程 */
void* phase1Compress(void* data) {
    Phase1PipelineArg& p = *(Phase1PipelineArg*)data;
    /* do the work */
    while (true) {
        // previous dependency
        yarn::DEPENDENCY_NOT_TO_BE(p.mergeSig, 0);
        yarn::DEPENDENCY_NOT_TO_BE(p.compressSig, p.COMPRESS_BUF_NUM);

        if (p.mergeFinish) {
            while (p.compressOrder < p.mergeOrder) {
                yarn::DEPENDENCY_NOT_TO_BE(p.compressSig, p.COMPRESS_BUF_NUM);
                doCompress(p);
                yarn::UPDATE_SIG_ORDER(p.compressSig, p.compressOrder);
            }
            yarn::SIGNAL_FINISH(p.compressSig, p.compressFinish);
            break;
        }

        doCompress(p);

        // update status
        yarn::CONSUME_SIGNAL(p.mergeSig);
        yarn::UPDATE_SIG_ORDER(p.compressSig, p.compressOrder);
    }

    spdlog::info("End compress order: {}, blocks: {}, bams: {}", p.compressOrder, block_num, bam_num);
    return nullptr;
}