#!/usr/bin/env bash

NUM_ITERATION=10

if [ "$EUID" -ne 0 ]; then
    echo "Error: Please run as root (sudo) to access /proc/vmallocinfo."
    exit 1
fi

TMP_BEFORE=$(mktemp)
TMP_AFTER=$(mktemp)

# Function to parse and aggregate vmallocinfo by caller/module
parse_vmalloc() {
    # Typical format: 0xffffc90000000000-0xffffc90000002000    8192 pcpu_alloc_alloc+0x3b/0x240 vmalloc
    # We aggregate the size (column 2) grouped by the caller function (column 3)
    awk '
    {
        # Remove offset details like +0x3b/0x240 to group by base function name
        split($3, caller, "+")
        # Keep track of caller type if present (e.g. vmalloc, ioremap)
        type = $4 ? " (" $4 ")" : ""
        key = caller[1] type
        sum[key] += $2
    }
    END {
        for (k in sum) {
            print k "," sum[k]
        }
    }' /proc/vmallocinfo
}

parse_vmalloc > "$TMP_BEFORE"
./amx_64 --exec-test $NUM_ITERATION
parse_vmalloc > "$TMP_AFTER"


# Join both files on the caller name and compute the differences
LEAKED=$(awk -F, '
NR == FNR {
    before[$1] = $2
    next
}
{
    after[$1] = $2
    diff = $2 - before[$1]
    if ($1 ~ /__xfd_enable_feature/ && diff > 1024*4*3*$NUM_ITERATION) {
        found=1
        #print "%-40s %+15d\n", $1, diff
    }
    delete before[$1] # Mark as processed
} END {
    print found+0
}' "$TMP_BEFORE" "$TMP_AFTER")

if [ "$LEAKED" == "1" ]; then
    echo "[FAIL] memory leak"
else
    echo "[OK] no memory leak"
fi

rm -f "$TMP_BEFORE" "$TMP_AFTER"