外部排序：当数据装不进内存

#include <limits.h>
#include <string.h>

#define MAX_WAYS 1024

typedef int64_t sort_key_t;

typedef struct {
    int loser[MAX_WAYS];       /* 内部节点：败者来源的路编号 */
    sort_key_t keys[MAX_WAYS]; /* 叶子节点：各路当前键值 */
    int k;                     /* 归并路数 */
} LoserTree;

/* 初始化败者树 */
void loser_tree_init(LoserTree *lt, int k) {
    lt->k = k;
    /* 所有内部节点初始指向一个虚拟的"负无穷"选手 */
    for (int i = 0; i < k; i++) {
        lt->loser[i] = -1;
    }
    /* 所有叶子初始为正无穷（哨兵） */
    for (int i = 0; i < k; i++) {
        lt->keys[i] = INT64_MAX;
    }
}

/* 从叶子 idx 到根的路径上调整败者 */
void loser_tree_adjust(LoserTree *lt, int idx) {
    int k = lt->k;
    int winner = idx;
    /* 从叶子对应的内部节点开始，向根上溯 */
    int parent = (k + idx) / 2;

    while (parent > 0) {
        if (lt->loser[parent] == -1) {
            /* 首次填充：直接设为败者，winner 上升 */
            lt->loser[parent] = winner;
            /* winner 还没有对手，由初始化保证 */
            return;
        }
        int loser = lt->loser[parent];
        if (lt->keys[winner] > lt->keys[loser]) {
            /* winner 输了，留下来做败者，原败者成为新 winner */
            lt->loser[parent] = winner;
            winner = loser;
        }
        /* 否则 winner 赢了，继续上升 */
        parent /= 2;
    }
    /* 到达根：winner 就是全局最小值 */
    lt->loser[0] = winner;
}

/* 构建败者树：逐一插入各路的首元素 */
void loser_tree_build(LoserTree *lt, sort_key_t initial_keys[], int k) {
    lt->k = k;
    for (int i = 0; i < k; i++) {
        lt->loser[i] = -1;
    }
    for (int i = 0; i < k; i++) {
        lt->keys[i] = initial_keys[i];
        loser_tree_adjust(lt, i);
    }
}

/* 获取当前最小值所在的路编号 */
static inline int loser_tree_winner(const LoserTree *lt) {
    return lt->loser[0];
}

/* 弹出最小值，用该路的下一个元素替换，重新调整 */
void loser_tree_replace(LoserTree *lt, int idx, sort_key_t new_key) {
    lt->keys[idx] = new_key;
    /* 从 idx 叶子开始，沿路径上溯做调整 */
    int k = lt->k;
    int winner = idx;
    int parent = (k + idx) / 2;

    while (parent > 0) {
        int loser = lt->loser[parent];
        if (lt->keys[winner] > lt->keys[loser]) {
            lt->loser[parent] = winner;
            winner = loser;
        }
        parent /= 2;
    }
    lt->loser[0] = winner;
}

typedef struct {
    char *buf[2];       /* 两个输出缓冲区 */
    int active;         /* 当前正在填充的缓冲区 (0 或 1) */
    size_t pos;         /* 当前缓冲区的写入位置 */
    size_t buf_size;    /* 缓冲区大小 */
    int write_fd;       /* 输出文件描述符 */
    /* 异步写状态 */
    struct aiocb aio;   /* POSIX AIO 控制块 */
    int write_pending;  /* 是否有异步写操作进行中 */
} OutputBuffer;

/* 写入一条记录到输出缓冲区 */
int output_buffer_write(OutputBuffer *ob, const void *record, size_t len) {
    if (ob->pos + len > ob->buf_size) {
        /* 当前缓冲区满，等待上一次异步写完成 */
        if (ob->write_pending) {
            while (aio_error(&ob->aio) == EINPROGRESS) {
                /* 自旋等待或 yield */
            }
            ob->write_pending = 0;
        }
        /* 启动异步写 */
        memset(&ob->aio, 0, sizeof(struct aiocb));
        ob->aio.aio_fildes = ob->write_fd;
        ob->aio.aio_buf = ob->buf[ob->active];
        ob->aio.aio_nbytes = ob->pos;
        aio_write(&ob->aio);
        ob->write_pending = 1;

        /* 切换缓冲区 */
        ob->active ^= 1;
        ob->pos = 0;
    }
    memcpy(ob->buf[ob->active] + ob->pos, record, len);
    ob->pos += len;
    return 0;
}

/* 伪代码：SQLite 的多线程 PMA 生成 */
/* 主线程 */
while (has_more_records()) {
    btree_insert(pSorter->pBtree, record);
    if (btree_size(pSorter->pBtree) >= threshold) {
        /* 将当前 B-tree 交给后台线程写出 */
        hand_off_to_bg_thread(pSorter->pBtree);
        /* 主线程立即开始新的 B-tree */
        pSorter->pBtree = btree_new();
    }
}
/* 后台线程 */
void bg_thread_func(BTree *bt) {
    traverse_and_write_pma(bt);
    btree_free(bt);
}

/* 从 PMA 读取器中取下一条记录 */
static int vdbePmaReaderNext(PmaReader *pReader) {
    int rc = SQLITE_OK;
    if (pReader->iReadOff >= pReader->iEof) {
        /* 当前 PMA 已读完 */
        pReader->pKey = 0;
        return SQLITE_OK;
    }
    /* 读取记录长度（变长整数） */
    rc = vdbePmaReadVarint(pReader, &nRec);
    /* 读取记录内容 */
    rc = vdbePmaReadBlob(pReader, (int)nRec, &pReader->pKey);
    pReader->nKey = (int)nRec;
    return rc;
}

/* 多路归并的一步：找到当前最小的 PMA 并推进 */
static int vdbeSorterNext(VdbeSorter *pSorter) {
    int rc;
    int iPrev = pSorter->aTree[1]; /* 败者树的总冠军 */
    /* 推进该路的 PMA 读取器 */
    rc = vdbePmaReaderNext(&pSorter->aReader[iPrev]);
    if (rc != SQLITE_OK) return rc;
    /* 重新调整败者树 */
    rc = vdbeSorterDoCompare(pSorter, iPrev);
    return rc;
}

/* RocksDB compaction 的核心归并循环（简化） */
void DoCompaction(CompactionState *state) {
    MergingIterator *merge_iter = NewMergingIterator(
        state->input_sstables,
        state->num_inputs
    );

    for (merge_iter->SeekToFirst();
         merge_iter->Valid();
         merge_iter->Next()) {
        Slice key = merge_iter->key();
        Slice value = merge_iter->value();

        /* 处理删除标记（tombstone）、版本控制等 */
        if (should_drop(key, value)) continue;

        /* 写入输出 SSTable */
        state->builder->Add(key, value);

        /* 如果输出 SSTable 达到目标大小，切分 */
        if (state->builder->FileSize() >= target_file_size) {
            FinishOutputFile(state);
            StartNewOutputFile(state);
        }
    }
}

/* 多队列 I/O 提交（概念代码） */
void submit_io_requests(InputBuffer bufs[], int k) {
    for (int i = 0; i < k; i++) {
        if (bufs[i].needs_refill) {
            /* 每个 run 的 I/O 提交到不同的队列 */
            /* io_uring 会自动利用多队列 */
            struct io_uring_sqe *sqe = io_uring_get_sqe(&ring);
            io_uring_prep_readv(sqe, bufs[i].fd,
                                &bufs[i].iov, 1,
                                bufs[i].offset);
            sqe->user_data = i;
        }
    }
    io_uring_submit(&ring);
}

/*
 * external_sort.c — 多路外部归并排序（败者树实现）
 *
 * 编译：gcc -O2 -o external_sort external_sort.c -lpthread
 * 使用：./external_sort input.bin output.bin
 *
 * 数据格式：二进制文件，每条记录为一个 int64_t
 */

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stdint.h>
#include <errno.h>
#include <limits.h>
#include <unistd.h>
#include <fcntl.h>
#include <sys/stat.h>

/* ===== 配置参数 ===== */
#define MEMORY_LIMIT    (256 * 1024 * 1024)   /* 256MB 可用内存 */
#define BLOCK_SIZE      (4 * 1024 * 1024)     /* 4MB I/O 块大小 */
#define MAX_MERGE_WAYS  512                    /* 最大归并路数 */
#define RECORD_SIZE     sizeof(int64_t)        /* 记录大小 */

/* 每个 run 最多容纳的记录数 */
#define RECORDS_PER_RUN (MEMORY_LIMIT / RECORD_SIZE)

/* 每个 I/O 块容纳的记录数 */
#define RECORDS_PER_BLK (BLOCK_SIZE / RECORD_SIZE)

/* ===== 错误处理宏 ===== */
#define CHECK(cond, msg) do { \
    if (!(cond)) { \
        fprintf(stderr, "ERROR [%s:%d]: %s (errno=%d: %s)\n", \
                __FILE__, __LINE__, msg, errno, strerror(errno)); \
        exit(EXIT_FAILURE); \
    } \
} while(0)

/* ===== 比较函数 ===== */
static int cmp_int64(const void *a, const void *b) {
    int64_t va = *(const int64_t *)a;
    int64_t vb = *(const int64_t *)b;
    if (va < vb) return -1;
    if (va > vb) return  1;
    return 0;
}

/* ===== 败者树 ===== */
typedef struct {
    int loser[MAX_MERGE_WAYS]; /* 内部节点：败者的路编号 */
    int64_t keys[MAX_MERGE_WAYS]; /* 各路当前键值 */
    int k;                     /* 归并路数 */
} LoserTree;

static void lt_init(LoserTree *lt, int k) {
    lt->k = k;
    for (int i = 0; i < k; i++) {
        lt->loser[i] = -1;
        lt->keys[i] = INT64_MAX;
    }
}

static void lt_adjust(LoserTree *lt, int s) {
    int k = lt->k;
    int winner = s;
    int parent = (k + s) / 2;

    while (parent > 0) {
        if (lt->loser[parent] == -1) {
            lt->loser[parent] = winner;
            return;
        }
        if (lt->keys[winner] > lt->keys[lt->loser[parent]]) {
            int tmp = winner;
            winner = lt->loser[parent];
            lt->loser[parent] = tmp;
        }
        parent /= 2;
    }
    lt->loser[0] = winner;
}

static void lt_build(LoserTree *lt, int64_t initial[], int k) {
    lt_init(lt, k);
    for (int i = k - 1; i >= 0; i--) {
        lt->keys[i] = initial[i];
        lt_adjust(lt, i);
    }
}

static inline int lt_winner(const LoserTree *lt) {
    return lt->loser[0];
}

static void lt_update(LoserTree *lt, int s, int64_t new_key) {
    lt->keys[s] = new_key;
    int k = lt->k;
    int winner = s;
    int parent = (k + s) / 2;

    while (parent > 0) {
        if (lt->keys[winner] > lt->keys[lt->loser[parent]]) {
            int tmp = winner;
            winner = lt->loser[parent];
            lt->loser[parent] = tmp;
        }
        parent /= 2;
    }
    lt->loser[0] = winner;
}

/* ===== Run 信息 ===== */
typedef struct {
    char path[256]; /* 临时文件路径 */
    size_t count;   /* 该 run 中的记录数 */
} RunInfo;

/* ===== 阶段一：生成初始 runs ===== */
static int generate_runs(const char *input_path,
                         RunInfo runs[], int *num_runs) {
    FILE *fin = fopen(input_path, "rb");
    CHECK(fin != NULL, "无法打开输入文件");

    /* 获取文件大小 */
    fseek(fin, 0, SEEK_END);
    long file_size = ftell(fin);
    fseek(fin, 0, SEEK_SET);
    size_t total_records = (size_t)file_size / RECORD_SIZE;

    /* 分配内存缓冲区 */
    int64_t *buffer = (int64_t *)malloc(MEMORY_LIMIT);
    CHECK(buffer != NULL, "内存分配失败");

    int run_count = 0;
    size_t records_remaining = total_records;

    fprintf(stderr, "总记录数: %zu, 每个 run 记录数: %zu\n",
            total_records, (size_t)RECORDS_PER_RUN);

    while (records_remaining > 0) {
        /* 读取一批记录到内存 */
        size_t batch = records_remaining < RECORDS_PER_RUN
                       ? records_remaining : RECORDS_PER_RUN;
        size_t nread = fread(buffer, RECORD_SIZE, batch, fin);
        CHECK(nread == batch, "读取输入文件失败");

        /* 内部排序 */
        qsort(buffer, batch, RECORD_SIZE, cmp_int64);

        /* 写入临时文件 */
        snprintf(runs[run_count].path, sizeof(runs[run_count].path),
                 "run_%04d.tmp", run_count);
        runs[run_count].count = batch;

        FILE *frun = fopen(runs[run_count].path, "wb");
        CHECK(frun != NULL, "无法创建临时 run 文件");

        size_t nwrite = fwrite(buffer, RECORD_SIZE, batch, frun);
        CHECK(nwrite == batch, "写入 run 文件失败");
        fclose(frun);

        fprintf(stderr, "  生成 run %d: %zu 条记录\n",
                run_count, batch);

        run_count++;
        records_remaining -= batch;
    }

    free(buffer);
    fclose(fin);

    *num_runs = run_count;
    fprintf(stderr, "阶段一完成: 共 %d 个 runs\n", run_count);
    return 0;
}

/* ===== Run 读取器 ===== */
typedef struct {
    FILE *fp;
    int64_t *buf;            /* 读取缓冲区 */
    size_t buf_capacity;     /* 缓冲区容量（记录数） */
    size_t buf_count;        /* 缓冲区有效记录数 */
    size_t buf_pos;          /* 当前读取位置 */
    size_t remaining;        /* 文件中剩余的记录数 */
    int exhausted;           /* 是否已耗尽 */
} RunReader;

static int run_reader_open(RunReader *rr, const RunInfo *ri,
                           size_t buf_capacity) {
    rr->fp = fopen(ri->path, "rb");
    CHECK(rr->fp != NULL, "无法打开 run 文件");

    rr->buf = (int64_t *)malloc(buf_capacity * RECORD_SIZE);
    CHECK(rr->buf != NULL, "读取缓冲区分配失败");

    rr->buf_capacity = buf_capacity;
    rr->buf_count = 0;
    rr->buf_pos = 0;
    rr->remaining = ri->count;
    rr->exhausted = 0;
    return 0;
}

static int run_reader_refill(RunReader *rr) {
    if (rr->remaining == 0) {
        rr->exhausted = 1;
        return -1;
    }
    size_t to_read = rr->remaining < rr->buf_capacity
                     ? rr->remaining : rr->buf_capacity;
    size_t nread = fread(rr->buf, RECORD_SIZE, to_read, rr->fp);
    CHECK(nread == to_read, "读取 run 缓冲区失败");

    rr->buf_count = nread;
    rr->buf_pos = 0;
    rr->remaining -= nread;
    return 0;
}

static int run_reader_next(RunReader *rr, int64_t *key) {
    if (rr->buf_pos >= rr->buf_count) {
        if (run_reader_refill(rr) < 0) return -1;
    }
    *key = rr->buf[rr->buf_pos++];
    return 0;
}

static void run_reader_close(RunReader *rr) {
    if (rr->fp) fclose(rr->fp);
    free(rr->buf);
    rr->fp = NULL;
    rr->buf = NULL;
}

/* ===== 阶段二：多路归并 ===== */
static int merge_runs(const char *output_path,
                      RunInfo runs[], int num_runs) {
    if (num_runs == 0) return 0;

    /* 如果只有一个 run，直接重命名 */
    if (num_runs == 1) {
        rename(runs[0].path, output_path);
        return 0;
    }

    int k = num_runs;
    if (k > MAX_MERGE_WAYS) {
        fprintf(stderr, "WARN: run 数量 %d 超过最大归并路数 %d，"
                "需要多趟归并（此实现暂不支持）\n",
                k, MAX_MERGE_WAYS);
        return -1;
    }

    /* 为每路分配缓冲区 */
    size_t buf_per_way = (MEMORY_LIMIT / RECORD_SIZE) / (k + 1);
    if (buf_per_way < 1) buf_per_way = 1;

    fprintf(stderr, "阶段二: %d 路归并，每路缓冲 %zu 条记录\n",
            k, buf_per_way);

    /* 打开所有 run 读取器 */
    RunReader *readers = (RunReader *)calloc(k, sizeof(RunReader));
    CHECK(readers != NULL, "读取器数组分配失败");

    for (int i = 0; i < k; i++) {
        run_reader_open(&readers[i], &runs[i], buf_per_way);
    }

    /* 读取每路的第一条记录，构建败者树 */
    int64_t initial_keys[MAX_MERGE_WAYS];
    for (int i = 0; i < k; i++) {
        if (run_reader_next(&readers[i], &initial_keys[i]) < 0) {
            initial_keys[i] = INT64_MAX; /* 空 run 用哨兵值 */
        }
    }

    LoserTree lt;
    lt_build(&lt, initial_keys, k);

    /* 打开输出文件 */
    FILE *fout = fopen(output_path, "wb");
    CHECK(fout != NULL, "无法打开输出文件");

    /* 输出缓冲区 */
    size_t out_buf_cap = buf_per_way;
    int64_t *out_buf = (int64_t *)malloc(out_buf_cap * RECORD_SIZE);
    CHECK(out_buf != NULL, "输出缓冲区分配失败");
    size_t out_pos = 0;

    /* 归并主循环 */
    size_t total_merged = 0;
    while (1) {
        int w = lt_winner(&lt);
        if (lt.keys[w] == INT64_MAX) break; /* 所有路耗尽 */

        /* 写入输出缓冲区 */
        out_buf[out_pos++] = lt.keys[w];
        if (out_pos >= out_buf_cap) {
            size_t nw = fwrite(out_buf, RECORD_SIZE, out_pos, fout);
            CHECK(nw == out_pos, "写入输出文件失败");
            out_pos = 0;
        }

        /* 从胜出的路读取下一条记录 */
        int64_t next_key;
        if (run_reader_next(&readers[w], &next_key) < 0) {
            next_key = INT64_MAX; /* 该路耗尽 */
        }
        lt_update(&lt, w, next_key);

        total_merged++;
        if (total_merged % 10000000 == 0) {
            fprintf(stderr, "  已归并 %zu 条记录\r", total_merged);
        }
    }

    /* 刷出剩余的输出缓冲区 */
    if (out_pos > 0) {
        size_t nw = fwrite(out_buf, RECORD_SIZE, out_pos, fout);
        CHECK(nw == out_pos, "刷出输出缓冲区失败");
    }

    fprintf(stderr, "\n阶段二完成: 共归并 %zu 条记录\n", total_merged);

    /* 清理 */
    free(out_buf);
    fclose(fout);
    for (int i = 0; i < k; i++) {
        run_reader_close(&readers[i]);
        remove(runs[i].path); /* 删除临时文件 */
    }
    free(readers);

    return 0;
}

/* ===== 验证排序结果 ===== */
static int verify_sorted(const char *path) {
    FILE *f = fopen(path, "rb");
    CHECK(f != NULL, "无法打开文件进行验证");

    int64_t prev = INT64_MIN;
    int64_t cur;
    size_t count = 0;

    while (fread(&cur, RECORD_SIZE, 1, f) == 1) {
        if (cur < prev) {
            fprintf(stderr, "验证失败: 位置 %zu, %ld > %ld\n",
                    count, (long)prev, (long)cur);
            fclose(f);
            return -1;
        }
        prev = cur;
        count++;
    }

    fclose(f);
    fprintf(stderr, "验证通过: %zu 条记录有序\n", count);
    return 0;
}

/* ===== 主函数 ===== */
int main(int argc, char *argv[]) {
    if (argc != 3) {
        fprintf(stderr, "用法: %s <输入文件> <输出文件>\n", argv[0]);
        return EXIT_FAILURE;
    }

    const char *input_path = argv[1];
    const char *output_path = argv[2];

    RunInfo runs[MAX_MERGE_WAYS];
    int num_runs = 0;

    fprintf(stderr, "==== 外部排序开始 ====\n");
    fprintf(stderr, "内存限制: %d MB, 块大小: %d MB\n",
            MEMORY_LIMIT / (1024*1024), BLOCK_SIZE / (1024*1024));

    /* 阶段一 */
    int rc = generate_runs(input_path, runs, &num_runs);
    CHECK(rc == 0, "生成初始 runs 失败");

    /* 阶段二 */
    rc = merge_runs(output_path, runs, num_runs);
    CHECK(rc == 0, "归并排序失败");

    /* 验证 */
    rc = verify_sorted(output_path);
    CHECK(rc == 0, "排序结果验证失败");

    fprintf(stderr, "==== 外部排序完成 ====\n");
    return EXIT_SUCCESS;
}