FastSort/src/sort/phase_2.h

/*
     Description: 排序过程的第二阶段
        * 第二阶段（phase-2）：
        * 1. 打开所有中间文件，并关联上解压线程池
        * 2. phase-2-stage-1：解压中间文件，并进行归并排序，放入buf
        * 3. phase-2-stage-2：多线程进行压缩，放入buf
        * 4. phase-2-stage-3：将buf写入最终的bam文件
        * 5. 要注意计算index并写入index文件
     Copyright : All right reserved by ICT

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

#pragma once
#include <spdlog/spdlog.h>
#include <stdio.h>

#include <string>

#include "phase_1.h"
#include "sam_io.h"
#include "sort.h"
using std::string;

// 循环缓冲区
struct CircularBuffer {
    uint8_t* data = nullptr;
    size_t readIdx = 0;  // 可以读取的开始位置
    size_t writeIdx = 0;  // 可以写入的开始位置
    size_t valueSize = 0; // 有效字节
    size_t bufSize = 0;  // 缓冲区空间

    CircularBuffer() {}
    CircularBuffer(size_t initSize) {
        bufSize = initSize;
        data = (uint8_t*)malloc(bufSize);
    }
    ~CircularBuffer() {
        if (data) {
            free(data);
            data = nullptr;
        }
        readIdx = 0;
        writeIdx = 0;
        valueSize = 0;
        bufSize = 0;
    }
    void AllocMem(size_t memSize) { ReAllocMem(memSize); }

    void ReAllocMem(size_t memSize) {
        if (memSize > bufSize) {
            bufSize = memSize;
            data = (uint8_t*)realloc(data, bufSize);
        }
    }

    size_t Write(uint8_t* src, size_t len) {
        if (len == 0) return 0;
        size_t freeSpace = bufSize - valueSize;
        size_t toWrite = MIN(len, freeSpace);
        if (toWrite == 0) return 0;

        size_t firstPart = MIN(toWrite, bufSize - writeIdx);
        // 第一段：从 write_index_ 到末尾
        memcpy(data + writeIdx, src, firstPart);
        // 第二段：如果超过了末尾，则从头部继续
        if (toWrite > firstPart) {
            memcpy(data, src + firstPart, toWrite - firstPart);
        }
        writeIdx = (writeIdx + toWrite) % bufSize;
        valueSize += toWrite;
        return toWrite;
    }

    size_t Read(uint8_t* out, size_t len) {
        if (len == 0) return 0;
        size_t toRead = MIN(len, valueSize);
        if (toRead == 0) return 0;

        size_t firstPart = MIN(toRead, bufSize - readIdx);
        memcpy(out, data + readIdx, firstPart);
        if (toRead > firstPart) {
            std::memcpy(out + firstPart, data, toRead - firstPart);
        }
        readIdx = (readIdx + toRead) % bufSize;
        valueSize -= toRead;
        return toRead;
    }

    // 返回第一个连续空间的大小
    size_t FirstPartWriteSize() {
        size_t freeSpace = bufSize - valueSize;
        size_t firstPart = MIN(freeSpace, bufSize - writeIdx);
        return firstPart;
    }

    // 如果空间不连续，那么返回第二个连续空间的内存大小
    size_t SecondPartWriteSize() {
        size_t freeSpace = bufSize - valueSize;
        size_t firstPart = MIN(freeSpace, bufSize - writeIdx);
        if (firstPart == freeSpace)
            return 0;
        return freeSpace - firstPart;
    }

    // 跳过不能完整保存一个bam的空间
    void SkipWrite(size_t skipBytes) {
        valueSize += skipBytes;
        writeIdx = (writeIdx + skipBytes) % bufSize;
    }

    void SkipRead(size_t from, size_t skipBytes) {
        if (valueSize <= skipBytes) { // 直接清空
            readIdx = 0;      // 可以读取的开始位置
            writeIdx = 0;     // 可以写入的开始位置
            valueSize = 0;    // 有效字节
        } else {
            if (from == readIdx) {
                readIdx = (readIdx + skipBytes) % bufSize;
                valueSize -= skipBytes;
            } else {
                valueSize -= (bufSize - readIdx + skipBytes);
                readIdx = skipBytes;
            }
        }
    }

    // 退回一个bam
    void RevertRead(size_t from, size_t revertBytes) {
        readIdx = from - revertBytes;
        valueSize += revertBytes;
    }

    // 读取一个bam
    size_t ReadBam(uint8_t* out, size_t start, size_t len) {
        memcpy(out, data + start, len);
        valueSize -= len;
        if (readIdx != start) {
            valueSize -= bufSize - readIdx;
        }
        readIdx = start + len;
        return len;
    }

    // 是否为空
    bool Empty() const { return valueSize == 0; }
    // 是否已满
    bool Full() const { return valueSize == bufSize; }
    // 当前数据量（字节）
    size_t Size() const { return valueSize; }
    // 剩余可写空间
    size_t Free() const { return bufSize - valueSize; }
    // 总容量
    size_t Capacity() const { return bufSize; }

    void Clear() {
        readIdx = 0;    // 可以读取的开始位置
        writeIdx = 0;   // 可以写入的开始位置
        valueSize = 0;  // 有效字节
    }
};

// 循环数组
template<class T>
struct CircularArray {
    std::vector<T> arr;
    size_t readIdx = 0;    // 可以读取的开始位置
    size_t writeIdx = 0;   // 可以写入的开始位置
    size_t valueSize = 0;  // 有效字节
    size_t bufSize = 0;    // 缓冲区空间

    CircularArray() {}
    CircularArray(size_t initSize) {
        bufSize = initSize;
        arr.resize(initSize);
    }
    void AllocArr(size_t arrSize) { ReAllocArr(arrSize); }

    void ReAllocArr(size_t arrSize) {
        if (arrSize > bufSize) {
            bufSize = arrSize;
            arr.resize(arrSize);
        }
    }

    inline T& operator[](size_t idx) {
        return arr[(idx + readIdx) % bufSize];
    }

    // 推入元素（自动扩容）
    void Push(const T& value) {
        if (Full()) {
            ReAllocArr(bufSize * 1.5);
        }
        arr[writeIdx] = value;
        writeIdx = (writeIdx + 1) % bufSize;
        ++valueSize;
    }

    void Push(T&& value) {
        if (Full()) {
            ReAllocArr(bufSize * 1.5);
        }
        arr[writeIdx] = std::move(value);
        writeIdx = (writeIdx + 1) % bufSize;
        ++valueSize;
    }

    // 弹出头部元素（如果为空则返回 nullopt）
    void Pop() {
        if (Empty())
            return;
        --valueSize;
        readIdx = (readIdx + 1) % bufSize;
    }

    void Revert() {
        ++valueSize;
        readIdx = (readIdx - 1) % bufSize;
    }

    // 查看头部/尾部（不弹出）
    T* Front() {
        if (Empty())
            return nullptr;
        return &arr[readIdx];
    }

    T* Back() {
        if (Empty())
            return nullptr;
        return &arr[(writeIdx - 1) % bufSize];
    }

    // 是否为空
    bool Empty() const { return valueSize == 0; }
    // 是否已满
    bool Full() const { return valueSize == bufSize; }
    // 当前数据量（字节）
    size_t Size() const { return valueSize; }
    // 剩余可写空间
    size_t Free() const { return bufSize - valueSize; }
    // 总容量
    size_t Capacity() const { return bufSize; }


    void Clear() {
        readIdx = 0;    // 可以读取的开始位置
        writeIdx = 0;   // 可以写入的开始位置
        valueSize = 0;  // 有效字节
    }
};

// 合并过程需要的缓冲区
struct Phase2MergeBuffer {
    CircularBuffer data;
    CircularArray<OneBam> bams;
    bool initialized = false;
    size_t numProcessed = 0;

    OneBam* Front() { return bams.Front(); }
    void Pop() {
        OneBam* bp = Front();
        data.SkipRead(bp->offset, bp->wholeBamLen);
        bams.Pop();
    }

    void Revert() {
        bams.Revert();
        OneBam* b = bams.Front();
        data.RevertRead(b->offset, b->wholeBamLen);
    }

    size_t Size() { return bams.Size(); }

    bool hasSpace(size_t bamLen) { return data.FirstPartWriteSize() >= bamLen; }

    // 从一个解压后的block缓冲区拷贝多个bam到循环缓冲区
    // 返回实际拷贝的bam数量
    size_t CopyBams(DataBuffer &blockBuf, BamArr &arr, int start) {
        if (start >= arr.Size()) return 0;
        int origStart = start;
        size_t numCopied = 0;
        OneBam* b1 = &arr.Get(start);
        OneBam* b2 = arr.Back();

        size_t firstPartSize = data.FirstPartWriteSize();
        size_t secondPartSize = data.SecondPartWriteSize();
        size_t needSize = b2->offset - b1->offset + b2->wholeBamLen;
        int stop = arr.Size();
        if (needSize <= firstPartSize) {  // 在第一个连续空间里就能放下
            // 每个bam的offset需要加上diff，以对应新的buf
            int64_t diff = (int64_t)data.writeIdx - b1->offset;
            data.Write(blockBuf.data + b1->offset, needSize);
            for (int i = start; i < stop; ++i) {
                bams.Push(arr.Get(i));
                bams.Back()->offset += diff;
                bams.Back()->addr = data.data;
            }
        } else {
            stop = start;
            size_t firstNeedSize = 0;
            while (arr.Get(stop).wholeBamLen + firstNeedSize < firstPartSize) {
                firstNeedSize += arr.Get(stop).wholeBamLen;
                stop += 1;
            }
            int64_t diff = (int64_t)data.writeIdx - b1->offset;
            data.Write(blockBuf.data + b1->offset, firstNeedSize);
            for (int i = start; i < stop; ++i) {
                bams.Push(arr.Get(i));
                bams.Back()->offset += diff;
                bams.Back()->addr = data.data;
            }
            // 跳过first part不能放下完整bam的部分
            data.SkipWrite(firstPartSize - firstNeedSize);
            b1 = &arr.Get(stop);
            start = stop;
            // 拷贝第二段
            size_t secondNeedSize = needSize - firstNeedSize;
            if (secondNeedSize <= secondPartSize) {
                stop = arr.Size();
            } else {
                secondNeedSize = 0;
                while (arr.Get(stop).wholeBamLen + secondNeedSize < secondPartSize) {
                    secondNeedSize += arr.Get(stop).wholeBamLen;
                    stop += 1;
                }
            }
            diff = (int64_t)data.writeIdx - b1->offset;
            data.Write(blockBuf.data + b1->offset, secondNeedSize);
            for (int i = start; i < stop; ++i) {
                bams.Push(arr.Get(i));
                bams.Back()->offset += diff;
                bams.Back()->addr = data.data;
            }
        }
        numCopied = stop - origStart;

        return numCopied;
    }

    void InitSize(size_t _bufsize, size_t _bamsize) {
        data.AllocMem(_bufsize);
        bams.AllocArr(_bamsize);
        initialized = true;
    }
};

struct BamDataBuffer {
    uint64_t offset = 0;
    int bamOffset = 0;
    DataBuffer block;
    BamArr bams;
    void Clear() {
        offset = 0;
        block.Clear();
        bams.Clear();
    }
};

// 读取和解压需要的buffer
struct Phase2ReadBuffer : ReadBuffer {
    FastVector<BamDataBuffer> uncompressBlocks;
    int bamNum = 0;
};

// 解析
struct UncompressBuffer {
    DataBuffer blockBuf;  // 解压的block放在这里
    BamArr bamArr;        // 解析后的bam数据放在这里
    int startIdx = 0;

    size_t Size() { return bamArr.Size() - startIdx; }

    OneBam *Front() {
        if (Size() == 0)
            return nullptr;
        return &bamArr.Get(startIdx);
    }

    void Clear() {
        startIdx = 0;
        blockBuf.Clear();
        bamArr.Clear();
    }
};

struct Phase2File {
    // static constexpr int MID_FILE_BUF_SIZE = 8 * SINGLE_BLOCK_SIZE;
    static constexpr int READ_BUF_NUM = 2;         // 读入的buf数量
    static constexpr int UNCOMPRESSS_BUF_NUM = 2;  // 解析的buf数量
    static constexpr int COPY_BUF_NUM = 1; // 拷贝到merge缓冲区数量

    // 参数
    size_t readBufSize = 0;
    string fileName;

    FILE* fp = nullptr;
    Phase2ReadBuffer readData[READ_BUF_NUM];
    UncompressBuffer uncompressBuf[UNCOMPRESSS_BUF_NUM];
    DataBuffer halfBlock;  // 剩余不完整的压缩的block数据

    // 读
    volatile int readyReadBufNum = 0;  // 有几个buffer已经填满可用了, 类似lock_t的作用
    uint64_t readOrder = 0;         // 当前读取的buf索引

    // 解压
    volatile int readyUncompressBufNum = 0;
    uint64_t uncompressOrder = 0;         // 写入的buf索引，类似readOrder
    bool needUncompress = false;

    // 拷贝到归并buf
    Phase2MergeBuffer mergeData;  // 归并时候用到的数据
    uint64_t copyOrder = 0;
    volatile int readyMergeBufNum = 0;

    volatile int finish = 0; // 读入结束

    void Init() {
        for (int i = 0; i < READ_BUF_NUM; ++i) {
            readData[i].Resize(readBufSize);
        }
        halfBlock.AllocMem(SINGLE_BLOCK_SIZE);
        // 打开中间文件
        // spdlog::info("file: {}", fileName);
        fp = fopen(fileName.c_str(), "rb");
    }

    void Init(const string& _fileName, size_t _readBufSize) {
        readBufSize = _readBufSize;
        fileName = _fileName;
        Init();
    }

    // 缓冲区大小应该根据中间文件的数量进行调整

    Phase2File() {}
    Phase2File(const string &_fileName, size_t _readBufSize) {
        readBufSize = _readBufSize;
        fileName = _fileName;
    }

    ~Phase2File() {
        if (fp) {
            fclose(fp);
        }
    }
};

// 用于合并压缩的数据结构
struct Phase2MergeData {
    vector<DataBuffer> blockDataArr;     // 待压缩的数据
    vector<DataBuffer> compressDataArr;  // 压缩后的数据

    void Resize(int blockNum) {
        blockDataArr.resize(blockNum);
        compressDataArr.resize(blockNum);
    }
};

/* 第二阶段的多线程流水线参数 */
struct Phase2PipelineArg {
    // common parameters
    static const int MERGE_BUF_NUM = 2;  // 压缩的buf数量
    static const int COMPRESS_BUF_NUM = 2;  // 压缩的buf数量

    int compressBlocksThreshold = 0;  // 进行归并排序的，单次block数量阈值，参数
    int numThread = 0;  // 线程数
    int compressLevel = 3;  // 压缩级别，参数
    uint64_t numBam = 0;

    // 从phase1转移来的数据
    UncompressBlockBuffer& uncompressData;  // 所有线程共用一个，串行往这里添加解压后的block数据
    BamArr& allBams;                        // 所有线程共用一个，串行往这里添加解析后的bam数据

    // 读入中间文件
    vector<Phase2File> midFiles;
    yarn::lock_t* readSig;
    uint64_t readOrder = 0;
    volatile int readFinish = 0;

    // 解压读入的blocks
    yarn::lock_t* uncompressSig;
    uint64_t uncompressOrder = 0;
    volatile int uncompressFinish = 0;
    int curBlockNum = 0;

    // 拷贝到merge缓冲区
    yarn::lock_t* copyMerge;
    uint64_t copyMergeOrder = 0;
    int uncompressReadyNum = 0;
    volatile int copyMergeFinish = 0;

    // 将解压的数据拷贝进merge的缓冲区，归并后放入，以便后续压缩用到
    MergeCompressData mergeData[MERGE_BUF_NUM];
    yarn::lock_t* mergeSig;
    yarn::lock_t* copyCompressSig;
    uint64_t mergeOrder = 0; // 前半段
    uint64_t copyCompressOrder = 0; // 后半段，需要压缩
    volatile int mergeFinish = 0;
    volatile int copyCompressFinish = 0;

    // 压缩
    DataBuffer compressBuf[COMPRESS_BUF_NUM];  // 压缩用到的buf
    yarn::lock_t* compressSig;
    uint64_t compressOrder = 0;
    volatile int compressFinish = 0;

    // 写入文件
    FILE* outFilePtr = nullptr;  // 文件指针
    uint64_t writeOrder = 0;

    void SetThreadNum(int nThread) {
        numThread = nThread;
        // const int kThres = 1024;
        const int kThres = 2048;
        int threadBlocksThres = numThread * 32;
        compressBlocksThreshold = threadBlocksThres < kThres ? threadBlocksThres : kThres;  // 大概64M的解压缩后的数据
        for (int i = 0; i < COMPRESS_BUF_NUM; ++i) {
            mergeData[i].Resize(compressBlocksThreshold);
        }
    }

    Phase2PipelineArg(UncompressBlockBuffer& lastBuf, BamArr& lastBams, string &midFilePrefix, int midFileNum, int threadNum) : uncompressData(lastBuf), allBams(lastBams){
        SetThreadNum(threadNum);

        int allFileBufSize = 64L * 1024 * 1024;
        // int bufBlocks = (numThread * 4 + midFileNum) / (midFileNum + 1);
        // size_t kFileBufSize = bufBlocks * SINGLE_BLOCK_SIZE;
        size_t kFileBufSize = allFileBufSize  / (midFileNum + 1);

        spdlog::info("mid file buf size: {}", kFileBufSize);

        midFiles.resize(midFileNum);
        for (int i = 0; i < midFileNum; ++i) {
            midFiles[i].Init(midFilePrefix + std::to_string(i), kFileBufSize);
        }

        readSig = yarn::NEW_LOCK(0);
        uncompressSig = yarn::NEW_LOCK(0);
        copyMerge = yarn::NEW_LOCK(0);
        mergeSig = yarn::NEW_LOCK(0);
        copyCompressSig = yarn::NEW_LOCK(0);
        compressSig = yarn::NEW_LOCK(0);


    }
};

void phase2Pipeline(Phase2PipelineArg &p);