/*
 * mem_bench.c — 匿名页申请耗时基准 (iOS / HarmonyOS / macOS / Linux 通用)
 *
 * 功能:
 *   按入参大小反复 mmap 一块匿名私有内存, 分别统计:
 *     - mmap   耗时 (建立映射本身, 通常只有 µs 级)
 *     - fill   耗时 (memset 整块填充 —— 真正触发缺页 + 写入的主成本;
 *                     或 touch 模式: 每页写 1 字节, 专门测首次触页缺页成本)
 *     - unmap  耗时 (munmap 释放)
 *   每轮结束后做轻量读回校验 (防止填充被优化掉, 并确认内存真实可读),
 *   最后输出 avg/min/max 汇总与吞吐率。
 *
 * 用法:
 *   mem_bench [-r n] [-m memset|touch|none] [-q] <size>[K|M|G|T]
 *
 * 构建:
 *   host      : cc -O2 -o mem_bench mem_bench.c
 *   iOS       : xcrun -sdk iphoneos clang -arch arm64 -O2 \
 *               -miphoneos-version-min=12.0 -o mem_bench mem_bench.c
 *   HarmonyOS : <OHOS_NDK>/llvm/bin/clang --target=aarch64-linux-ohos \
 *               --sysroot=<OHOS_NDK>/sysroot -O2 -o mem_bench mem_bench.c
 *
 * 说明:
 *   - 纯 POSIX C, 无第三方依赖; 仅 RSS 读取按平台分支。
 *   - 单位为二进制: K=1024, M=1024^2, G=1024^3, T=1024^4。
 *   - 每轮 memset 使用不同填充值 (0x11/0x22/...), 规避任何零页/去重优化。
 *   - iOS 为 16K 页, Linux/鸿蒙通常为 4K 页; touch 模式的 ns/page 指标
 *     可直接对比两平台的单缺页成本。
 */

#if !defined(__APPLE__)
#define _POSIX_C_SOURCE 200809L
#endif

#include <errno.h>
#include <inttypes.h>
#include <limits.h>
#include <signal.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/mman.h>
#include <sys/resource.h>
#include <time.h>
#include <unistd.h>

#if defined(__APPLE__)
#include <mach/mach.h>
#endif

#ifndef MAP_ANONYMOUS
#define MAP_ANONYMOUS MAP_ANON
#endif

enum fill_mode { FILL_MEMSET, FILL_TOUCH, FILL_NONE };

static volatile sig_atomic_t g_stop = 0;
static volatile uint64_t g_sink = 0;   /* 读回校验累加器, volatile 防优化 */
static size_t g_page = 4096;

/* ---------------- 公共小工具 ---------------- */

static void on_signal(int sig)
{
    (void)sig;
    g_stop = 1;
}

static double now_sec(void)
{
    struct timespec ts;
    clock_gettime(CLOCK_MONOTONIC, &ts);
    return (double)ts.tv_sec + (double)ts.tv_nsec * 1e-9;
}

static double to_mib(unsigned long long bytes)
{
    return (double)bytes / 1048576.0;
}

/* 解析大小参数 (与 mem_stress 相同: 二进制单位, 兼容 B/iB 后缀) */
static int parse_size(const char *s, unsigned long long *out)
{
    if (s == NULL || *s == '\0' || *s == '-' || *s == '+')
        return -1;

    errno = 0;
    char *end = NULL;
    unsigned long long v = strtoull(s, &end, 10);
    if (errno != 0 || end == s)
        return -1;

    const char *u = end;
    unsigned long long mult = 1ULL;

    if (*u == 'K' || *u == 'k' || *u == 'M' || *u == 'm' ||
        *u == 'G' || *u == 'g' || *u == 'T' || *u == 't') {
        switch (*u | 0x20) {
        case 'k': mult = 1ULL << 10; break;
        case 'm': mult = 1ULL << 20; break;
        case 'g': mult = 1ULL << 30; break;
        case 't': mult = 1ULL << 40; break;
        }
        u++;
        if (*u == 'i')
            u++;
        if (*u == 'B' || *u == 'b')
            u++;
    } else if (*u == 'B' || *u == 'b') {
        u++;
    }
    if (*u != '\0')
        return -1;
    if (v == 0)
        return -1;
    if (v > ULLONG_MAX / mult)
        return -1;

    *out = v * mult;
    return 0;
}

/* 当前驻留集 (字节); 取不到返回 0 (调用方按 0 跳过显示) */
static unsigned long long get_rss_bytes(void)
{
#if defined(__APPLE__)
    struct task_vm_info vm;
    mach_msg_type_number_t count = TASK_VM_INFO_COUNT;
    if (task_info(mach_task_self(), TASK_VM_INFO,
                  (task_info_t)&vm, &count) != KERN_SUCCESS)
        return 0;
    return (unsigned long long)vm.resident_size;
#else
    FILE *f = fopen("/proc/self/statm", "r");
    if (f == NULL)
        return 0;
    unsigned long long total = 0, resident = 0;
    int n = fscanf(f, "%llu %llu", &total, &resident);
    fclose(f);
    if (n != 2)
        return 0;
    return resident * (unsigned long long)g_page;
#endif
}

/* 峰值驻留集 (字节); Darwin 的 ru_maxrss 单位是字节, Linux 系是 KiB */
static unsigned long long get_peak_rss_bytes(void)
{
    struct rusage ru;
    if (getrusage(RUSAGE_SELF, &ru) != 0)
        return 0;
#if defined(__APPLE__)
    return (unsigned long long)ru.ru_maxrss;
#else
    return (unsigned long long)ru.ru_maxrss * 1024ULL;
#endif
}

/*
 * 严格的整数解析 (拒绝垃圾输入)。注: 不用 getopt —— BSD/musl 的 getopt
 * 不重排 argv, "size 在前、选项在后"的写法会解析失败; 手写解析保证
 * 各平台行为一致。
 */
static int parse_int(const char *s, long *out)
{
    if (s == NULL || *s == '\0')
        return -1;
    errno = 0;
    char *end = NULL;
    long v = strtol(s, &end, 10);
    if (errno != 0 || end == s || *end != '\0')
        return -1;
    *out = v;
    return 0;
}

static const char *fill_mode_name(enum fill_mode m)
{
    switch (m) {
    case FILL_MEMSET: return "memset";
    case FILL_TOUCH:  return "touch ";
    case FILL_NONE:   return "none  ";
    }
    return "?";
}

static void usage(FILE *out, const char *prog)
{
    fprintf(out,
        "mem_bench - anonymous mmap/memset timing benchmark (iOS/HarmonyOS/macOS/Linux)\n"
        "\n"
        "usage: %s [-r n] [-m mode] [-q] <size>\n"
        "\n"
        "  <size>    bytes of anonymous memory per run\n"
        "            e.g. 268435456 | 100K | 512M | 2G (binary units, B/iB suffix ok)\n"
        "  -r n      repeat runs (default 3)\n"
        "  -m mode   fill mode after mmap (default memset):\n"
        "              memset  memset the whole region (faults + bandwidth)\n"
        "              touch   write 1 byte per page (first-touch fault cost, ns/page)\n"
        "              none    mmap only (mapping setup/teardown cost)\n"
        "  -q        quiet: per-run lines suppressed\n"
        "\n"
        "examples:\n"
        "  %s 512M -r 5          time mmap + memset of 512M, 5 runs\n"
        "  %s 512M -m touch      time first-touch page faults (ns/page)\n",
        prog, prog, prog);
}

/* 汇总打印一行: name avg/min/max (value 数组单位为秒, 乘 scale 换显示单位) */
static void print_stats(const char *name, const double *v, int n,
                        double scale, const char *unit)
{
    double sum = 0.0, mn = v[0], mx = v[0];
    for (int i = 0; i < n; i++) {
        sum += v[i];
        if (v[i] < mn) mn = v[i];
        if (v[i] > mx) mx = v[i];
    }
    printf("%-7s: avg %10.3f %s | min %10.3f %s | max %10.3f %s\n",
           name, sum / (double)n * scale, unit,
           mn * scale, unit, mx * scale, unit);
}

int main(int argc, char **argv)
{
    int repeats = 3;
    int quiet = 0;
    enum fill_mode mode = FILL_MEMSET;
    const char *size_arg = NULL;

    for (int i = 1; i < argc; i++) {
        const char *a = argv[i];
        if (strcmp(a, "-q") == 0) {
            quiet = 1;
        } else if (strcmp(a, "-h") == 0 || strcmp(a, "--help") == 0) {
            usage(stdout, argv[0]);
            return 0;
        } else if (strcmp(a, "-r") == 0) {
            long v = 0;
            if (i + 1 >= argc || parse_int(argv[i + 1], &v) != 0 ||
                v < 1 || v > 10000) {
                fprintf(stderr, "error: -r expects an integer in 1..10000\n");
                return 2;
            }
            i++;
            repeats = (int)v;
        } else if (strcmp(a, "-m") == 0) {
            if (i + 1 >= argc) {
                fprintf(stderr, "error: -m expects a mode (memset|touch|none)\n");
                return 2;
            }
            i++;
            if (strcmp(argv[i], "memset") == 0)      mode = FILL_MEMSET;
            else if (strcmp(argv[i], "touch") == 0)  mode = FILL_TOUCH;
            else if (strcmp(argv[i], "none") == 0)   mode = FILL_NONE;
            else {
                fprintf(stderr, "error: unknown mode '%s' (memset|touch|none)\n",
                        argv[i]);
                return 2;
            }
        } else if (a[0] == '-' && a[1] != '\0') {
            fprintf(stderr, "error: unknown option '%s'\n", a);
            usage(stderr, argv[0]);
            return 2;
        } else {
            if (size_arg != NULL) {
                fprintf(stderr, "error: multiple sizes given ('%s' and '%s')\n",
                        size_arg, a);
                return 2;
            }
            size_arg = a;
        }
    }
    if (size_arg == NULL) {
        usage(stderr, argv[0]);
        return 2;
    }

    unsigned long long req = 0;
    if (parse_size(size_arg, &req) != 0) {
        fprintf(stderr, "error: bad size '%s' (try 512M / 2G / 65536K)\n",
                size_arg);
        return 2;
    }

    long ps = sysconf(_SC_PAGESIZE);
    if (ps > 0)
        g_page = (size_t)ps;

    if (req > (unsigned long long)(SIZE_MAX - (g_page - 1))) {
        fprintf(stderr, "error: size %llu bytes exceeds this platform's address space\n",
                req);
        return 2;
    }
    size_t size = ((size_t)req + g_page - 1) & ~((size_t)g_page - 1);
    unsigned long long pages = (unsigned long long)(size / g_page);

    setvbuf(stdout, NULL, _IOLBF, 0);

    struct sigaction sa;
    memset(&sa, 0, sizeof sa);
    sa.sa_handler = on_signal;
    sigemptyset(&sa.sa_mask);
    sa.sa_flags = 0;
    sigaction(SIGINT, &sa, NULL);
    sigaction(SIGTERM, &sa, NULL);

    double *t_mmap = malloc((size_t)repeats * sizeof(double));
    double *t_fill = malloc((size_t)repeats * sizeof(double));
    double *t_unmap = malloc((size_t)repeats * sizeof(double));
    if (t_mmap == NULL || t_fill == NULL || t_unmap == NULL) {
        fprintf(stderr, "error: out of memory\n");
        return 1;
    }

    printf("target : %llu bytes (%.2f MiB), %llu page(s) x %zu B, "
           "mode=%s, runs=%d\n",
           (unsigned long long)size, to_mib(size), pages, g_page,
           fill_mode_name(mode), repeats);

    int completed = 0;
    for (int i = 0; i < repeats && !g_stop; i++) {
        /* --- mmap --- */
        double t0 = now_sec();
        char *p = mmap(NULL, size, PROT_READ | PROT_WRITE,
                       MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
        double t1 = now_sec();
        if (p == MAP_FAILED) {
            fprintf(stderr, "error: mmap %llu bytes failed on run %d: %s\n",
                    (unsigned long long)size, i + 1, strerror(errno));
#if defined(__APPLE__)
            fprintf(stderr, "hint: iOS caps a single process's anonymous "
                            "reservation (~4GiB observed on 14.8); "
                            "try a smaller size.\n");
#endif
            break;
        }
        t_mmap[i] = t1 - t0;

        /* --- fill --- */
        unsigned char pat = (unsigned char)(0x11 * (unsigned)(i + 1)); /* 0x11,0x22,... */
        t0 = now_sec();
        if (mode == FILL_MEMSET) {
            memset(p, pat, size);
        } else if (mode == FILL_TOUCH) {
            for (size_t off = 0; off < size && !g_stop; off += g_page)
                p[off] = (char)pat;
        }
        t1 = now_sec();
        t_fill[i] = t1 - t0;

        /* 轻量读回校验: 每页读首字节累加, 防止填充/映射被优化掉 */
        if (mode != FILL_NONE) {
            uint64_t sum = 0;
            for (size_t off = 0; off < size; off += g_page)
                sum += (uint64_t)(unsigned char)p[off];
            g_sink = sum;
        }
        unsigned long long rss = get_rss_bytes();

        /* --- unmap --- */
        t0 = now_sec();
        munmap(p, size);
        t_unmap[i] = now_sec() - t0;

        completed = i + 1;
        if (!quiet) {
            printf("run %d/%d: mmap=%9.1f us | fill(%s)=%10.3f ms",
                   i + 1, repeats, t_mmap[i] * 1e6,
                   mode == FILL_MEMSET ? "memset" :
                   mode == FILL_TOUCH  ? "touch " : "none  ",
                   t_fill[i] * 1e3);
            if (mode == FILL_MEMSET && t_fill[i] > 0.0)
                printf(" (%8.1f MiB/s)", to_mib(size) / t_fill[i]);
            if (mode == FILL_TOUCH && t_fill[i] > 0.0 && pages > 0)
                printf(" (%8.1f ns/page)",
                       t_fill[i] * 1e9 / (double)pages);
            printf(" | unmap=%8.3f ms", t_unmap[i] * 1e3);
            if (rss)
                printf(" | rss=%8.2f MiB", to_mib(rss));
            printf("\n");
        }

        if (i + 1 < repeats) {
            struct timespec ts = { 0, 100 * 1000 * 1000 };  /* 100ms 间隔 */
            nanosleep(&ts, NULL);
        }
    }

    if (completed == 0) {
        fprintf(stderr, "error: no run completed\n");
        free(t_mmap); free(t_fill); free(t_unmap);
        return 1;
    }

    /* ---- 汇总 ---- */
    printf("---- summary (%d run%s, %.2f MiB, %s mode) ----\n",
           completed, completed == 1 ? "" : "s", to_mib(size),
           mode == FILL_MEMSET ? "memset" :
           mode == FILL_TOUCH  ? "touch"  : "none");
    print_stats("mmap", t_mmap, completed, 1e6, "us");
    if (mode != FILL_NONE)
        print_stats("fill", t_fill, completed, 1e3, "ms");
    print_stats("unmap", t_unmap, completed, 1e3, "ms");

    if (mode == FILL_MEMSET) {
        double tp_sum = 0.0;
        int tp_n = 0;
        for (int i = 0; i < completed; i++) {
            if (t_fill[i] > 0.0) {
                tp_sum += to_mib(size) / t_fill[i];
                tp_n++;
            }
        }
        if (tp_n > 0)
            printf("memset throughput: avg %.1f MiB/s over %d run(s)\n",
                   tp_sum / (double)tp_n, tp_n);
    }
    if (mode == FILL_TOUCH && pages > 0) {
        double fp_sum = 0.0;
        int fp_n = 0;
        for (int i = 0; i < completed; i++) {
            if (t_fill[i] > 0.0) {
                fp_sum += t_fill[i] * 1e9 / (double)pages;
                fp_n++;
            }
        }
        if (fp_n > 0)
            printf("first-touch fault: avg %.1f ns/page over %d run(s)\n",
                   fp_sum / (double)fp_n, fp_n);
    }

    printf("peak rss: %.2f MiB\n", to_mib(get_peak_rss_bytes()));
    (void)g_sink;

    free(t_mmap); free(t_fill); free(t_unmap);
    return 0;
}