// SPDX-License-Identifier: GPL-2.0

#include <linux/version.h>
#include <linux/ptrace.h>
#include <uapi/linux/bpf.h>
#include <bpf/bpf_helpers.h>

/*
 * The CPU number, cstate number and pstate number are based
 * on 96boards Hikey with octa CA53 CPUs.
 *
 * Every CPU have three idle states for cstate:
 *   WFI, CPU_OFF, CLUSTER_OFF
 *
 * Every CPU have 5 operating points:
 *   208MHz, 432MHz, 729MHz, 960MHz, 1200MHz
 *
 * This code is based on these assumption and other platforms
 * need to adjust these definitions.
 */
#define MAX_CPU                 8
#define MAX_PSTATE_ENTRIES      5
#define MAX_CSTATE_ENTRIES      3

static int cpu_opps[] = { 208000, 432000, 729000, 960000, 1200000 };

/*
 * my_map structure is used to record cstate and pstate index and
 * timestamp (Idx, Ts), when new event incoming we need to update
 * combination for new state index and timestamp (Idx`, Ts`).
 *
 * Based on (Idx, Ts) and (Idx`, Ts`) we can calculate the time
 * interval for the previous state: Duration(Idx) = Ts` - Ts.
 *
 * Every CPU has one below array for recording state index and
 * timestamp, and record for cstate and pstate saperately:
 *
 * +--------------------------+
 * | cstate timestamp         |
 * +--------------------------+
 * | cstate index             |
 * +--------------------------+
 * | pstate timestamp         |
 * +--------------------------+
 * | pstate index             |
 * +--------------------------+
 */
#define MAP_OFF_CSTATE_TIME     0
#define MAP_OFF_CSTATE_IDX      1
#define MAP_OFF_PSTATE_TIME     2
#define MAP_OFF_PSTATE_IDX      3
#define MAP_OFF_NUM             4

struct {
        __uint(type, BPF_MAP_TYPE_ARRAY);
        __type(key, u32);
        __type(value, u64);
        __uint(max_entries, MAX_CPU * MAP_OFF_NUM);
} my_map SEC(".maps");

/* cstate_duration records duration time for every idle state per CPU */
struct {
        __uint(type, BPF_MAP_TYPE_ARRAY);
        __type(key, u32);
        __type(value, u64);
        __uint(max_entries, MAX_CPU * MAX_CSTATE_ENTRIES);
} cstate_duration SEC(".maps");

/* pstate_duration records duration time for every operating point per CPU */
struct {
        __uint(type, BPF_MAP_TYPE_ARRAY);
        __type(key, u32);
        __type(value, u64);
        __uint(max_entries, MAX_CPU * MAX_PSTATE_ENTRIES);
} pstate_duration SEC(".maps");

/*
 * The trace events for cpu_idle and cpu_frequency are taken from:
 * /sys/kernel/tracing/events/power/cpu_idle/format
 * /sys/kernel/tracing/events/power/cpu_frequency/format
 *
 * These two events have same format, so define one common structure.
 */
struct cpu_args {
        u64 pad;
        u32 state;
        u32 cpu_id;
};

/* calculate pstate index, returns MAX_PSTATE_ENTRIES for failure */
static u32 find_cpu_pstate_idx(u32 frequency)
{
        u32 i;

        for (i = 0; i < sizeof(cpu_opps) / sizeof(u32); i++) {
                if (frequency == cpu_opps[i])
                        return i;
        }

        return i;
}

SEC("tracepoint/power/cpu_idle")
int bpf_prog1(struct cpu_args *ctx)
{
        u64 *cts, *pts, *cstate, *pstate, prev_state, cur_ts, delta;
        u32 key, cpu, pstate_idx;
        u64 *val;

        if (ctx->cpu_id > MAX_CPU)
                return 0;

        cpu = ctx->cpu_id;

        key = cpu * MAP_OFF_NUM + MAP_OFF_CSTATE_TIME;
        cts = bpf_map_lookup_elem(&my_map, &key);
        if (!cts)
                return 0;

        key = cpu * MAP_OFF_NUM + MAP_OFF_CSTATE_IDX;
        cstate = bpf_map_lookup_elem(&my_map, &key);
        if (!cstate)
                return 0;

        key = cpu * MAP_OFF_NUM + MAP_OFF_PSTATE_TIME;
        pts = bpf_map_lookup_elem(&my_map, &key);
        if (!pts)
                return 0;

        key = cpu * MAP_OFF_NUM + MAP_OFF_PSTATE_IDX;
        pstate = bpf_map_lookup_elem(&my_map, &key);
        if (!pstate)
                return 0;

        prev_state = *cstate;
        *cstate = ctx->state;

        if (!*cts) {
                *cts = bpf_ktime_get_ns();
                return 0;
        }

        cur_ts = bpf_ktime_get_ns();
        delta = cur_ts - *cts;
        *cts = cur_ts;

        /*
         * When state doesn't equal to (u32)-1, the cpu will enter
         * one idle state; for this case we need to record interval
         * for the pstate.
         *
         *                 OPP2
         *            +---------------------+
         *     OPP1   |                     |
         *   ---------+                     |
         *                                  |  Idle state
         *                                  +---------------
         *
         *            |<- pstate duration ->|
         *            ^                     ^
         *           pts                  cur_ts
         */
        if (ctx->state != (u32)-1) {

                /* record pstate after have first cpu_frequency event */
                if (!*pts)
                        return 0;

                delta = cur_ts - *pts;

                pstate_idx = find_cpu_pstate_idx(*pstate);
                if (pstate_idx >= MAX_PSTATE_ENTRIES)
                        return 0;

                key = cpu * MAX_PSTATE_ENTRIES + pstate_idx;
                val = bpf_map_lookup_elem(&pstate_duration, &key);
                if (val)
                        __sync_fetch_and_add((long *)val, delta);

        /*
         * When state equal to (u32)-1, the cpu just exits from one
         * specific idle state; for this case we need to record
         * interval for the pstate.
         *
         *       OPP2
         *   -----------+
         *              |                          OPP1
         *              |                     +-----------
         *              |     Idle state      |
         *              +---------------------+
         *
         *              |<- cstate duration ->|
         *              ^                     ^
         *             cts                  cur_ts
         */
        } else {

                key = cpu * MAX_CSTATE_ENTRIES + prev_state;
                val = bpf_map_lookup_elem(&cstate_duration, &key);
                if (val)
                        __sync_fetch_and_add((long *)val, delta);
        }

        /* Update timestamp for pstate as new start time */
        if (*pts)
                *pts = cur_ts;

        return 0;
}

SEC("tracepoint/power/cpu_frequency")
int bpf_prog2(struct cpu_args *ctx)
{
        u64 *pts, *cstate, *pstate, cur_ts, delta;
        u32 key, cpu, pstate_idx;
        u64 *val;

        cpu = ctx->cpu_id;

        key = cpu * MAP_OFF_NUM + MAP_OFF_PSTATE_TIME;
        pts = bpf_map_lookup_elem(&my_map, &key);
        if (!pts)
                return 0;

        key = cpu * MAP_OFF_NUM + MAP_OFF_PSTATE_IDX;
        pstate = bpf_map_lookup_elem(&my_map, &key);
        if (!pstate)
                return 0;

        key = cpu * MAP_OFF_NUM + MAP_OFF_CSTATE_IDX;
        cstate = bpf_map_lookup_elem(&my_map, &key);
        if (!cstate)
                return 0;

        *pstate = ctx->state;

        if (!*pts) {
                *pts = bpf_ktime_get_ns();
                return 0;
        }

        cur_ts = bpf_ktime_get_ns();
        delta = cur_ts - *pts;
        *pts = cur_ts;

        /* When CPU is in idle, bail out to skip pstate statistics */
        if (*cstate != (u32)(-1))
                return 0;

        /*
         * The cpu changes to another different OPP (in below diagram
         * change frequency from OPP3 to OPP1), need recording interval
         * for previous frequency OPP3 and update timestamp as start
         * time for new frequency OPP1.
         *
         *                 OPP3
         *            +---------------------+
         *     OPP2   |                     |
         *   ---------+                     |
         *                                  |    OPP1
         *                                  +---------------
         *
         *            |<- pstate duration ->|
         *            ^                     ^
         *           pts                  cur_ts
         */
        pstate_idx = find_cpu_pstate_idx(*pstate);
        if (pstate_idx >= MAX_PSTATE_ENTRIES)
                return 0;

        key = cpu * MAX_PSTATE_ENTRIES + pstate_idx;
        val = bpf_map_lookup_elem(&pstate_duration, &key);
        if (val)
                __sync_fetch_and_add((long *)val, delta);

        return 0;
}

char _license[] SEC("license") = "GPL";
u32 _version SEC("version") = LINUX_VERSION_CODE;