// SPDX-License-Identifier: GPL-2.0
/*
 * The back-end-agnostic part of Just-In-Time compiler for eBPF bytecode.
 *
 * Copyright (c) 2024 Synopsys Inc.
 * Author: Shahab Vahedi <shahab@synopsys.com>
 */
#include <linux/bug.h>
#include "bpf_jit.h"

/*
 * Check for the return value. A pattern used often in this file.
 * There must be a "ret" variable of type "int" in the scope.
 */
#define CHECK_RET(cmd)                  \
        do {                            \
                ret = (cmd);            \
                if (ret < 0)            \
                        return ret;     \
        } while (0)

#ifdef ARC_BPF_JIT_DEBUG
/* Dumps bytes in /var/log/messages at KERN_INFO level (4). */
static void dump_bytes(const u8 *buf, u32 len, const char *header)
{
        u8 line[64];
        size_t i, j;

        pr_info("-----------------[ %s ]-----------------\n", header);

        for (i = 0, j = 0; i < len; i++) {
                /* Last input byte? */
                if (i == len - 1) {
                        j += scnprintf(line + j, 64 - j, "0x%02x", buf[i]);
                        pr_info("%s\n", line);
                        break;
                }
                /* End of line? */
                else if (i % 8 == 7) {
                        j += scnprintf(line + j, 64 - j, "0x%02x", buf[i]);
                        pr_info("%s\n", line);
                        j = 0;
                } else {
                        j += scnprintf(line + j, 64 - j, "0x%02x, ", buf[i]);
                }
        }
}
#endif /* ARC_BPF_JIT_DEBUG */

/********************* JIT context ***********************/

/*
 * buf:         Translated instructions end up here.
 * len:         The length of whole block in bytes.
 * index:       The offset at which the _next_ instruction may be put.
 */
struct jit_buffer {
        u8      *buf;
        u32     len;
        u32     index;
};

/*
 * This is a subset of "struct jit_context" that its information is deemed
 * necessary for the next extra pass to come.
 *
 * bpf_header:  Needed to finally lock the region.
 * bpf2insn:    Used to find the translation for instructions of interest.
 *
 * Things like "jit.buf" and "jit.len" can be retrieved respectively from
 * "prog->bpf_func" and "prog->jited_len".
 */
struct arc_jit_data {
        struct bpf_binary_header *bpf_header;
        u32                      *bpf2insn;
};

/*
 * The JIT pertinent context that is used by different functions.
 *
 * prog:                The current eBPF program being handled.
 * orig_prog:           The original eBPF program before any possible change.
 * jit:                 The JIT buffer and its length.
 * bpf_header:          The JITed program header. "jit.buf" points inside it.
 * emit:                If set, opcodes are written to memory; else, a dry-run.
 * do_zext:             If true, 32-bit sub-regs must be zero extended.
 * bpf2insn:            Maps BPF insn indices to their counterparts in jit.buf.
 * bpf2insn_valid:      Indicates if "bpf2ins" is populated with the mappings.
 * jit_data:            A piece of memory to transfer data to the next pass.
 * arc_regs_clobbered:  Each bit status determines if that arc reg is clobbered.
 * save_blink:          Whether ARC's "blink" register needs to be saved.
 * frame_size:          Derived from "prog->aux->stack_depth".
 * epilogue_offset:     Used by early "return"s in the code to jump here.
 * need_extra_pass:     A forecast if an "extra_pass" will occur.
 * is_extra_pass:       Indicates if the current pass is an extra pass.
 * user_bpf_prog:       True, if VM opcodes come from a real program.
 * blinded:             True if "constant blinding" step returned a new "prog".
 * success:             Indicates if the whole JIT went OK.
 */
struct jit_context {
        struct bpf_prog                 *prog;
        struct bpf_prog                 *orig_prog;
        struct jit_buffer               jit;
        struct bpf_binary_header        *bpf_header;
        bool                            emit;
        bool                            do_zext;
        u32                             *bpf2insn;
        bool                            bpf2insn_valid;
        struct arc_jit_data             *jit_data;
        u32                             arc_regs_clobbered;
        bool                            save_blink;
        u16                             frame_size;
        u32                             epilogue_offset;
        bool                            need_extra_pass;
        bool                            is_extra_pass;
        bool                            user_bpf_prog;
        bool                            blinded;
        bool                            success;
};

/*
 * If we're in ARC_BPF_JIT_DEBUG mode and the debug level is right, dump the
 * input BPF stream. "bpf_jit_dump()" is not fully suited for this purpose.
 */
static void vm_dump(const struct bpf_prog *prog)
{
#ifdef ARC_BPF_JIT_DEBUG
        if (bpf_jit_enable > 1)
                dump_bytes((u8 *)prog->insns, 8 * prog->len, " VM  ");
#endif
}

/*
 * If the right level of debug is set, dump the bytes. There are 2 variants
 * of this function:
 *
 * 1. Use the standard bpf_jit_dump() which is meant only for JITed code.
 * 2. Use the dump_bytes() to match its "vm_dump()" instance.
 */
static void jit_dump(const struct jit_context *ctx)
{
#ifdef ARC_BPF_JIT_DEBUG
        u8 header[8];
#endif
        const int pass = ctx->is_extra_pass ? 2 : 1;

        if (bpf_jit_enable <= 1 || !ctx->prog->jited)
                return;

#ifdef ARC_BPF_JIT_DEBUG
        scnprintf(header, sizeof(header), "JIT:%d", pass);
        dump_bytes(ctx->jit.buf, ctx->jit.len, header);
        pr_info("\n");
#else
        bpf_jit_dump(ctx->prog->len, ctx->jit.len, pass, ctx->jit.buf);
#endif
}

/* Initialise the context so there's no garbage. */
static int jit_ctx_init(struct jit_context *ctx, struct bpf_prog *prog)
{
        memset(ctx, 0, sizeof(*ctx));

        ctx->orig_prog = prog;

        /* If constant blinding was requested but failed, scram. */
        ctx->prog = bpf_jit_blind_constants(prog);
        if (IS_ERR(ctx->prog))
                return PTR_ERR(ctx->prog);
        ctx->blinded = (ctx->prog != ctx->orig_prog);

        /* If the verifier doesn't zero-extend, then we have to do it. */
        ctx->do_zext = !ctx->prog->aux->verifier_zext;

        ctx->is_extra_pass = ctx->prog->jited;
        ctx->user_bpf_prog = ctx->prog->is_func;

        return 0;
}

/*
 * Only after the first iteration of normal pass (the dry-run),
 * there are valid offsets in ctx->bpf2insn array.
 */
static inline bool offsets_available(const struct jit_context *ctx)
{
        return ctx->bpf2insn_valid;
}

/*
 * "*mem" should be freed when there is no "extra pass" to come,
 * or the compilation terminated abruptly. A few of such memory
 * allocations are: ctx->jit_data and ctx->bpf2insn.
 */
static inline void maybe_free(struct jit_context *ctx, void **mem)
{
        if (*mem) {
                if (!ctx->success || !ctx->need_extra_pass) {
                        kfree(*mem);
                        *mem = NULL;
                }
        }
}

/*
 * Free memories based on the status of the context.
 *
 * A note about "bpf_header": On successful runs, "bpf_header" is
 * not freed, because "jit.buf", a sub-array of it, is returned as
 * the "bpf_func". However, "bpf_header" is lost and nothing points
 * to it. This should not cause a leakage, because apparently
 * "bpf_header" can be revived by "bpf_jit_binary_hdr()". This is
 * how "bpf_jit_free()" in "kernel/bpf/core.c" releases the memory.
 */
static void jit_ctx_cleanup(struct jit_context *ctx)
{
        if (ctx->blinded) {
                /* if all went well, release the orig_prog. */
                if (ctx->success)
                        bpf_jit_prog_release_other(ctx->prog, ctx->orig_prog);
                else
                        bpf_jit_prog_release_other(ctx->orig_prog, ctx->prog);
        }

        maybe_free(ctx, (void **)&ctx->bpf2insn);
        maybe_free(ctx, (void **)&ctx->jit_data);

        if (!ctx->bpf2insn)
                ctx->bpf2insn_valid = false;

        /* Freeing "bpf_header" is enough. "jit.buf" is a sub-array of it. */
        if (!ctx->success && ctx->bpf_header) {
                bpf_jit_binary_free(ctx->bpf_header);
                ctx->bpf_header = NULL;
                ctx->jit.buf    = NULL;
                ctx->jit.index  = 0;
                ctx->jit.len    = 0;
        }

        ctx->emit = false;
        ctx->do_zext = false;
}

/*
 * Analyse the register usage and record the frame size.
 * The register usage is determined by consulting the back-end.
 */
static void analyze_reg_usage(struct jit_context *ctx)
{
        size_t i;
        u32 usage = 0;
        const struct bpf_insn *insn = ctx->prog->insnsi;

        for (i = 0; i < ctx->prog->len; i++) {
                u8 bpf_reg;
                bool call;

                bpf_reg = insn[i].dst_reg;
                call = (insn[i].code == (BPF_JMP | BPF_CALL)) ? true : false;
                usage |= mask_for_used_regs(bpf_reg, call);
        }

        ctx->arc_regs_clobbered = usage;
        ctx->frame_size = ctx->prog->aux->stack_depth;
}

/* Verify that no instruction will be emitted when there is no buffer. */
static inline int jit_buffer_check(const struct jit_context *ctx)
{
        if (ctx->emit) {
                if (!ctx->jit.buf) {
                        pr_err("bpf-jit: inconsistence state; no "
                               "buffer to emit instructions.\n");
                        return -EINVAL;
                } else if (ctx->jit.index > ctx->jit.len) {
                        pr_err("bpf-jit: estimated JIT length is less "
                               "than the emitted instructions.\n");
                        return -EFAULT;
                }
        }
        return 0;
}

/* On a dry-run (emit=false), "jit.len" is growing gradually. */
static inline void jit_buffer_update(struct jit_context *ctx, u32 n)
{
        if (!ctx->emit)
                ctx->jit.len += n;
        else
                ctx->jit.index += n;
}

/* Based on "emit", determine the address where instructions are emitted. */
static inline u8 *effective_jit_buf(const struct jit_context *ctx)
{
        return ctx->emit ? (ctx->jit.buf + ctx->jit.index) : NULL;
}

/* Prologue based on context variables set by "analyze_reg_usage()". */
static int handle_prologue(struct jit_context *ctx)
{
        int ret;
        u8 *buf = effective_jit_buf(ctx);
        u32 len = 0;

        CHECK_RET(jit_buffer_check(ctx));

        len = arc_prologue(buf, ctx->arc_regs_clobbered, ctx->frame_size);
        jit_buffer_update(ctx, len);

        return 0;
}

/* The counter part for "handle_prologue()". */
static int handle_epilogue(struct jit_context *ctx)
{
        int ret;
        u8 *buf = effective_jit_buf(ctx);
        u32 len = 0;

        CHECK_RET(jit_buffer_check(ctx));

        len = arc_epilogue(buf, ctx->arc_regs_clobbered, ctx->frame_size);
        jit_buffer_update(ctx, len);

        return 0;
}

/* Tell which number of the BPF instruction we are dealing with. */
static inline s32 get_index_for_insn(const struct jit_context *ctx,
                                     const struct bpf_insn *insn)
{
        return (insn - ctx->prog->insnsi);
}

/*
 * In most of the cases, the "offset" is read from "insn->off". However,
 * if it is an unconditional BPF_JMP32, then it comes from "insn->imm".
 *
 * (Courtesy of "cpu=v4" support)
 */
static inline s32 get_offset(const struct bpf_insn *insn)
{
        if ((BPF_CLASS(insn->code) == BPF_JMP32) &&
            (BPF_OP(insn->code) == BPF_JA))
                return insn->imm;
        else
                return insn->off;
}

/*
 * Determine to which number of the BPF instruction we're jumping to.
 *
 * The "offset" is interpreted as the "number" of BPF instructions
 * from the _next_ BPF instruction. e.g.:
 *
 *  4 means 4 instructions after  the next insn
 *  0 means 0 instructions after  the next insn -> fallthrough.
 * -1 means 1 instruction  before the next insn -> jmp to current insn.
 *
 *  Another way to look at this, "offset" is the number of instructions
 *  that exist between the current instruction and the target instruction.
 *
 *  It is worth noting that a "mov r,i64", which is 16-byte long, is
 *  treated as two instructions long, therefore "offset" needn't be
 *  treated specially for those. Everything is uniform.
 */
static inline s32 get_target_index_for_insn(const struct jit_context *ctx,
                                            const struct bpf_insn *insn)
{
        return (get_index_for_insn(ctx, insn) + 1) + get_offset(insn);
}

/* Is there an immediate operand encoded in the "insn"? */
static inline bool has_imm(const struct bpf_insn *insn)
{
        return BPF_SRC(insn->code) == BPF_K;
}

/* Is the last BPF instruction? */
static inline bool is_last_insn(const struct bpf_prog *prog, u32 idx)
{
        return idx == (prog->len - 1);
}

/*
 * Invocation of this function, conditionally signals the need for
 * an extra pass. The conditions that must be met are:
 *
 * 1. The current pass itself shouldn't be an extra pass.
 * 2. The stream of bytes being JITed must come from a user program.
 */
static inline void set_need_for_extra_pass(struct jit_context *ctx)
{
        if (!ctx->is_extra_pass)
                ctx->need_extra_pass = ctx->user_bpf_prog;
}

/*
 * Check if the "size" is valid and then transfer the control to
 * the back-end for the swap.
 */
static int handle_swap(u8 *buf, u8 rd, u8 size, u8 endian,
                       bool force, bool do_zext, u8 *len)
{
        /* Sanity check on the size. */
        switch (size) {
        case 16:
        case 32:
        case 64:
                break;
        default:
                pr_err("bpf-jit: invalid size for swap.\n");
                return -EINVAL;
        }

        *len = gen_swap(buf, rd, size, endian, force, do_zext);

        return 0;
}

/* Checks if the (instruction) index is in valid range. */
static inline bool check_insn_idx_valid(const struct jit_context *ctx,
                                        const s32 idx)
{
        return (idx >= 0 && idx < ctx->prog->len);
}

/*
 * Decouple the back-end from BPF by converting BPF conditions
 * to internal enum. ARC_CC_* start from 0 and are used as index
 * to an array. BPF_J* usage must end after this conversion.
 */
static int bpf_cond_to_arc(const u8 op, u8 *arc_cc)
{
        switch (op) {
        case BPF_JA:
                *arc_cc = ARC_CC_AL;
                break;
        case BPF_JEQ:
                *arc_cc = ARC_CC_EQ;
                break;
        case BPF_JGT:
                *arc_cc = ARC_CC_UGT;
                break;
        case BPF_JGE:
                *arc_cc = ARC_CC_UGE;
                break;
        case BPF_JSET:
                *arc_cc = ARC_CC_SET;
                break;
        case BPF_JNE:
                *arc_cc = ARC_CC_NE;
                break;
        case BPF_JSGT:
                *arc_cc = ARC_CC_SGT;
                break;
        case BPF_JSGE:
                *arc_cc = ARC_CC_SGE;
                break;
        case BPF_JLT:
                *arc_cc = ARC_CC_ULT;
                break;
        case BPF_JLE:
                *arc_cc = ARC_CC_ULE;
                break;
        case BPF_JSLT:
                *arc_cc = ARC_CC_SLT;
                break;
        case BPF_JSLE:
                *arc_cc = ARC_CC_SLE;
                break;
        default:
                pr_err("bpf-jit: can't handle condition 0x%02X\n", op);
                return -EINVAL;
        }
        return 0;
}

/*
 * Check a few things for a supposedly "jump" instruction:
 *
 * 0. "insn" is a "jump" instruction, but not the "call/exit" variant.
 * 1. The current "insn" index is in valid range.
 * 2. The index of target instruction is in valid range.
 */
static int check_bpf_jump(const struct jit_context *ctx,
                          const struct bpf_insn *insn)
{
        const u8 class = BPF_CLASS(insn->code);
        const u8 op = BPF_OP(insn->code);

        /* Must be a jmp(32) instruction that is not a "call/exit". */
        if ((class != BPF_JMP && class != BPF_JMP32) ||
            (op == BPF_CALL || op == BPF_EXIT)) {
                pr_err("bpf-jit: not a jump instruction.\n");
                return -EINVAL;
        }

        if (!check_insn_idx_valid(ctx, get_index_for_insn(ctx, insn))) {
                pr_err("bpf-jit: the bpf jump insn is not in prog.\n");
                return -EINVAL;
        }

        if (!check_insn_idx_valid(ctx, get_target_index_for_insn(ctx, insn))) {
                pr_err("bpf-jit: bpf jump label is out of range.\n");
                return -EINVAL;
        }

        return 0;
}

/*
 * Based on input "insn", consult "ctx->bpf2insn" to get the
 * related index (offset) of the translation in JIT stream.
 */
static u32 get_curr_jit_off(const struct jit_context *ctx,
                            const struct bpf_insn *insn)
{
        const s32 idx = get_index_for_insn(ctx, insn);
#ifdef ARC_BPF_JIT_DEBUG
        BUG_ON(!offsets_available(ctx) || !check_insn_idx_valid(ctx, idx));
#endif
        return ctx->bpf2insn[idx];
}

/*
 * The input "insn" must be a jump instruction.
 *
 * Based on input "insn", consult "ctx->bpf2insn" to get the
 * related JIT index (offset) of "target instruction" that
 * "insn" would jump to.
 */
static u32 get_targ_jit_off(const struct jit_context *ctx,
                            const struct bpf_insn *insn)
{
        const s32 tidx = get_target_index_for_insn(ctx, insn);
#ifdef ARC_BPF_JIT_DEBUG
        BUG_ON(!offsets_available(ctx) || !check_insn_idx_valid(ctx, tidx));
#endif
        return ctx->bpf2insn[tidx];
}

/*
 * This function will return 0 for a feasible jump.
 *
 * Consult the back-end to check if it finds it feasible to emit
 * the necessary instructions based on "cond" and the displacement
 * between the "from_off" and the "to_off".
 */
static int feasible_jit_jump(u32 from_off, u32 to_off, u8 cond, bool j32)
{
        int ret = 0;

        if (j32) {
                if (!check_jmp_32(from_off, to_off, cond))
                        ret = -EFAULT;
        } else {
                if (!check_jmp_64(from_off, to_off, cond))
                        ret = -EFAULT;
        }

        if (ret != 0)
                pr_err("bpf-jit: the JIT displacement is not OK.\n");

        return ret;
}

/*
 * This jump handler performs the following steps:
 *
 * 1. Compute ARC's internal condition code from BPF's
 * 2. Determine the bitness of the operation (32 vs. 64)
 * 3. Sanity check on BPF stream
 * 4. Sanity check on what is supposed to be JIT's displacement
 * 5. And finally, emit the necessary instructions
 *
 * The last two steps are performed through the back-end.
 * The value of steps 1 and 2 are necessary inputs for the back-end.
 */
static int handle_jumps(const struct jit_context *ctx,
                        const struct bpf_insn *insn,
                        u8 *len)
{
        u8 cond;
        int ret = 0;
        u8 *buf = effective_jit_buf(ctx);
        const bool j32 = (BPF_CLASS(insn->code) == BPF_JMP32) ? true : false;
        const u8 rd = insn->dst_reg;
        u8 rs = insn->src_reg;
        u32 curr_off = 0, targ_off = 0;

        *len = 0;

        /* Map the BPF condition to internal enum. */
        CHECK_RET(bpf_cond_to_arc(BPF_OP(insn->code), &cond));

        /* Sanity check on the BPF byte stream. */
        CHECK_RET(check_bpf_jump(ctx, insn));

        /*
         * Move the immediate into a temporary register _now_ for 2 reasons:
         *
         * 1. "gen_jmp_{32,64}()" deal with operands in registers.
         *
         * 2. The "len" parameter will grow so that the current jit offset
         *    (curr_off) will have increased to a point where the necessary
         *    instructions can be inserted by "gen_jmp_{32,64}()".
         */
        if (has_imm(insn) && cond != ARC_CC_AL) {
                if (j32) {
                        *len += mov_r32_i32(BUF(buf, *len), JIT_REG_TMP,
                                            insn->imm);
                } else {
                        *len += mov_r64_i32(BUF(buf, *len), JIT_REG_TMP,
                                            insn->imm);
                }
                rs = JIT_REG_TMP;
        }

        /* If the offsets are known, check if the branch can occur. */
        if (offsets_available(ctx)) {
                curr_off = get_curr_jit_off(ctx, insn) + *len;
                targ_off = get_targ_jit_off(ctx, insn);

                /* Sanity check on the back-end side. */
                CHECK_RET(feasible_jit_jump(curr_off, targ_off, cond, j32));
        }

        if (j32) {
                *len += gen_jmp_32(BUF(buf, *len), rd, rs, cond,
                                   curr_off, targ_off);
        } else {
                *len += gen_jmp_64(BUF(buf, *len), rd, rs, cond,
                                   curr_off, targ_off);
        }

        return ret;
}

/* Jump to translated epilogue address. */
static int handle_jmp_epilogue(struct jit_context *ctx,
                               const struct bpf_insn *insn, u8 *len)
{
        u8 *buf = effective_jit_buf(ctx);
        u32 curr_off = 0, epi_off = 0;

        /* Check the offset only if the data is available. */
        if (offsets_available(ctx)) {
                curr_off = get_curr_jit_off(ctx, insn);
                epi_off = ctx->epilogue_offset;

                if (!check_jmp_64(curr_off, epi_off, ARC_CC_AL)) {
                        pr_err("bpf-jit: epilogue offset is not valid.\n");
                        return -EINVAL;
                }
        }

        /* Jump to "epilogue offset" (rd and rs don't matter). */
        *len = gen_jmp_64(buf, 0, 0, ARC_CC_AL, curr_off, epi_off);

        return 0;
}

/* Try to get the resolved address and generate the instructions. */
static int handle_call(struct jit_context *ctx,
                       const struct bpf_insn *insn,
                       u8 *len)
{
        int  ret;
        bool in_kernel_func, fixed = false;
        u64  addr = 0;
        u8  *buf = effective_jit_buf(ctx);

        ret = bpf_jit_get_func_addr(ctx->prog, insn, ctx->is_extra_pass,
                                    &addr, &fixed);
        if (ret < 0) {
                pr_err("bpf-jit: can't get the address for call.\n");
                return ret;
        }
        in_kernel_func = (fixed ? true : false);

        /* No valuable address retrieved (yet). */
        if (!fixed && !addr)
                set_need_for_extra_pass(ctx);

        *len = gen_func_call(buf, (ARC_ADDR)addr, in_kernel_func);

        if (insn->src_reg != BPF_PSEUDO_CALL) {
                /* Assigning ABI's return reg to JIT's return reg. */
                *len += arc_to_bpf_return(BUF(buf, *len));
        }

        return 0;
}

/*
 * Try to generate instructions for loading a 64-bit immediate.
 * These sort of instructions are usually associated with the 64-bit
 * relocations: R_BPF_64_64. Therefore, signal the need for an extra
 * pass if the circumstances are right.
 */
static int handle_ld_imm64(struct jit_context *ctx,
                           const struct bpf_insn *insn,
                           u8 *len)
{
        const s32 idx = get_index_for_insn(ctx, insn);
        u8 *buf = effective_jit_buf(ctx);

        /* We're about to consume 2 VM instructions. */
        if (is_last_insn(ctx->prog, idx)) {
                pr_err("bpf-jit: need more data for 64-bit immediate.\n");
                return -EINVAL;
        }

        *len = mov_r64_i64(buf, insn->dst_reg, insn->imm, (insn + 1)->imm);

        if (bpf_pseudo_func(insn))
                set_need_for_extra_pass(ctx);

        return 0;
}

/*
 * Handles one eBPF instruction at a time. To make this function faster,
 * it does not call "jit_buffer_check()". Else, it would call it for every
 * instruction. As a result, it should not be invoked directly. Only
 * "handle_body()", that has already executed the "check", may call this
 * function.
 *
 * If the "ret" value is negative, something has went wrong. Else,
 * it mostly holds the value 0 and rarely 1. Number 1 signals
 * the loop in "handle_body()" to skip the next instruction, because
 * it has been consumed as part of a 64-bit immediate value.
 */
static int handle_insn(struct jit_context *ctx, u32 idx)
{
        const struct bpf_insn *insn = &ctx->prog->insnsi[idx];
        const u8  code = insn->code;
        const u8  dst  = insn->dst_reg;
        const u8  src  = insn->src_reg;
        const s16 off  = insn->off;
        const s32 imm  = insn->imm;
        u8 *buf = effective_jit_buf(ctx);
        u8  len = 0;
        int ret = 0;

        switch (code) {
        /* dst += src (32-bit) */
        case BPF_ALU | BPF_ADD | BPF_X:
                len = add_r32(buf, dst, src);
                break;
        /* dst += imm (32-bit) */
        case BPF_ALU | BPF_ADD | BPF_K:
                len = add_r32_i32(buf, dst, imm);
                break;
        /* dst -= src (32-bit) */
        case BPF_ALU | BPF_SUB | BPF_X:
                len = sub_r32(buf, dst, src);
                break;
        /* dst -= imm (32-bit) */
        case BPF_ALU | BPF_SUB | BPF_K:
                len = sub_r32_i32(buf, dst, imm);
                break;
        /* dst = -dst (32-bit) */
        case BPF_ALU | BPF_NEG:
                len = neg_r32(buf, dst);
                break;
        /* dst *= src (32-bit) */
        case BPF_ALU | BPF_MUL | BPF_X:
                len = mul_r32(buf, dst, src);
                break;
        /* dst *= imm (32-bit) */
        case BPF_ALU | BPF_MUL | BPF_K:
                len = mul_r32_i32(buf, dst, imm);
                break;
        /* dst /= src (32-bit) */
        case BPF_ALU | BPF_DIV | BPF_X:
                len = div_r32(buf, dst, src, off == 1);
                break;
        /* dst /= imm (32-bit) */
        case BPF_ALU | BPF_DIV | BPF_K:
                len = div_r32_i32(buf, dst, imm, off == 1);
                break;
        /* dst %= src (32-bit) */
        case BPF_ALU | BPF_MOD | BPF_X:
                len = mod_r32(buf, dst, src, off == 1);
                break;
        /* dst %= imm (32-bit) */
        case BPF_ALU | BPF_MOD | BPF_K:
                len = mod_r32_i32(buf, dst, imm, off == 1);
                break;
        /* dst &= src (32-bit) */
        case BPF_ALU | BPF_AND | BPF_X:
                len = and_r32(buf, dst, src);
                break;
        /* dst &= imm (32-bit) */
        case BPF_ALU | BPF_AND | BPF_K:
                len = and_r32_i32(buf, dst, imm);
                break;
        /* dst |= src (32-bit) */
        case BPF_ALU | BPF_OR | BPF_X:
                len = or_r32(buf, dst, src);
                break;
        /* dst |= imm (32-bit) */
        case BPF_ALU | BPF_OR | BPF_K:
                len = or_r32_i32(buf, dst, imm);
                break;
        /* dst ^= src (32-bit) */
        case BPF_ALU | BPF_XOR | BPF_X:
                len = xor_r32(buf, dst, src);
                break;
        /* dst ^= imm (32-bit) */
        case BPF_ALU | BPF_XOR | BPF_K:
                len = xor_r32_i32(buf, dst, imm);
                break;
        /* dst <<= src (32-bit) */
        case BPF_ALU | BPF_LSH | BPF_X:
                len = lsh_r32(buf, dst, src);
                break;
        /* dst <<= imm (32-bit) */
        case BPF_ALU | BPF_LSH | BPF_K:
                len = lsh_r32_i32(buf, dst, imm);
                break;
        /* dst >>= src (32-bit) [unsigned] */
        case BPF_ALU | BPF_RSH | BPF_X:
                len = rsh_r32(buf, dst, src);
                break;
        /* dst >>= imm (32-bit) [unsigned] */
        case BPF_ALU | BPF_RSH | BPF_K:
                len = rsh_r32_i32(buf, dst, imm);
                break;
        /* dst >>= src (32-bit) [signed] */
        case BPF_ALU | BPF_ARSH | BPF_X:
                len = arsh_r32(buf, dst, src);
                break;
        /* dst >>= imm (32-bit) [signed] */
        case BPF_ALU | BPF_ARSH | BPF_K:
                len = arsh_r32_i32(buf, dst, imm);
                break;
        /* dst = src (32-bit) */
        case BPF_ALU | BPF_MOV | BPF_X:
                len = mov_r32(buf, dst, src, (u8)off);
                break;
        /* dst = imm32 (32-bit) */
        case BPF_ALU | BPF_MOV | BPF_K:
                len = mov_r32_i32(buf, dst, imm);
                break;
        /* dst = swap(dst) */
        case BPF_ALU   | BPF_END | BPF_FROM_LE:
        case BPF_ALU   | BPF_END | BPF_FROM_BE:
        case BPF_ALU64 | BPF_END | BPF_FROM_LE: {
                CHECK_RET(handle_swap(buf, dst, imm, BPF_SRC(code),
                                      BPF_CLASS(code) == BPF_ALU64,
                                      ctx->do_zext, &len));
                break;
        }
        /* dst += src (64-bit) */
        case BPF_ALU64 | BPF_ADD | BPF_X:
                len = add_r64(buf, dst, src);
                break;
        /* dst += imm32 (64-bit) */
        case BPF_ALU64 | BPF_ADD | BPF_K:
                len = add_r64_i32(buf, dst, imm);
                break;
        /* dst -= src (64-bit) */
        case BPF_ALU64 | BPF_SUB | BPF_X:
                len = sub_r64(buf, dst, src);
                break;
        /* dst -= imm32 (64-bit) */
        case BPF_ALU64 | BPF_SUB | BPF_K:
                len = sub_r64_i32(buf, dst, imm);
                break;
        /* dst = -dst (64-bit) */
        case BPF_ALU64 | BPF_NEG:
                len = neg_r64(buf, dst);
                break;
        /* dst *= src (64-bit) */
        case BPF_ALU64 | BPF_MUL | BPF_X:
                len = mul_r64(buf, dst, src);
                break;
        /* dst *= imm32 (64-bit) */
        case BPF_ALU64 | BPF_MUL | BPF_K:
                len = mul_r64_i32(buf, dst, imm);
                break;
        /* dst &= src (64-bit) */
        case BPF_ALU64 | BPF_AND | BPF_X:
                len = and_r64(buf, dst, src);
                break;
        /* dst &= imm32 (64-bit) */
        case BPF_ALU64 | BPF_AND | BPF_K:
                len = and_r64_i32(buf, dst, imm);
                break;
        /* dst |= src (64-bit) */
        case BPF_ALU64 | BPF_OR | BPF_X:
                len = or_r64(buf, dst, src);
                break;
        /* dst |= imm32 (64-bit) */
        case BPF_ALU64 | BPF_OR | BPF_K:
                len = or_r64_i32(buf, dst, imm);
                break;
        /* dst ^= src (64-bit) */
        case BPF_ALU64 | BPF_XOR | BPF_X:
                len = xor_r64(buf, dst, src);
                break;
        /* dst ^= imm32 (64-bit) */
        case BPF_ALU64 | BPF_XOR | BPF_K:
                len = xor_r64_i32(buf, dst, imm);
                break;
        /* dst <<= src (64-bit) */
        case BPF_ALU64 | BPF_LSH | BPF_X:
                len = lsh_r64(buf, dst, src);
                break;
        /* dst <<= imm32 (64-bit) */
        case BPF_ALU64 | BPF_LSH | BPF_K:
                len = lsh_r64_i32(buf, dst, imm);
                break;
        /* dst >>= src (64-bit) [unsigned] */
        case BPF_ALU64 | BPF_RSH | BPF_X:
                len = rsh_r64(buf, dst, src);
                break;
        /* dst >>= imm32 (64-bit) [unsigned] */
        case BPF_ALU64 | BPF_RSH | BPF_K:
                len = rsh_r64_i32(buf, dst, imm);
                break;
        /* dst >>= src (64-bit) [signed] */
        case BPF_ALU64 | BPF_ARSH | BPF_X:
                len = arsh_r64(buf, dst, src);
                break;
        /* dst >>= imm32 (64-bit) [signed] */
        case BPF_ALU64 | BPF_ARSH | BPF_K:
                len = arsh_r64_i32(buf, dst, imm);
                break;
        /* dst = src (64-bit) */
        case BPF_ALU64 | BPF_MOV | BPF_X:
                len = mov_r64(buf, dst, src, (u8)off);
                break;
        /* dst = imm32 (sign extend to 64-bit) */
        case BPF_ALU64 | BPF_MOV | BPF_K:
                len = mov_r64_i32(buf, dst, imm);
                break;
        /* dst = imm64 */
        case BPF_LD | BPF_DW | BPF_IMM:
                CHECK_RET(handle_ld_imm64(ctx, insn, &len));
                /* Tell the loop to skip the next instruction. */
                ret = 1;
                break;
        /* dst = *(size *)(src + off) */
        case BPF_LDX | BPF_MEM | BPF_W:
        case BPF_LDX | BPF_MEM | BPF_H:
        case BPF_LDX | BPF_MEM | BPF_B:
        case BPF_LDX | BPF_MEM | BPF_DW:
                len = load_r(buf, dst, src, off, BPF_SIZE(code), false);
                break;
        case BPF_LDX | BPF_MEMSX | BPF_W:
        case BPF_LDX | BPF_MEMSX | BPF_H:
        case BPF_LDX | BPF_MEMSX | BPF_B:
                len = load_r(buf, dst, src, off, BPF_SIZE(code), true);
                break;
        /* *(size *)(dst + off) = src */
        case BPF_STX | BPF_MEM | BPF_W:
        case BPF_STX | BPF_MEM | BPF_H:
        case BPF_STX | BPF_MEM | BPF_B:
        case BPF_STX | BPF_MEM | BPF_DW:
                len = store_r(buf, src, dst, off, BPF_SIZE(code));
                break;
        case BPF_ST | BPF_MEM | BPF_W:
        case BPF_ST | BPF_MEM | BPF_H:
        case BPF_ST | BPF_MEM | BPF_B:
        case BPF_ST | BPF_MEM | BPF_DW:
                len = store_i(buf, imm, dst, off, BPF_SIZE(code));
                break;
        case BPF_JMP   | BPF_JA:
        case BPF_JMP   | BPF_JEQ  | BPF_X:
        case BPF_JMP   | BPF_JEQ  | BPF_K:
        case BPF_JMP   | BPF_JNE  | BPF_X:
        case BPF_JMP   | BPF_JNE  | BPF_K:
        case BPF_JMP   | BPF_JSET | BPF_X:
        case BPF_JMP   | BPF_JSET | BPF_K:
        case BPF_JMP   | BPF_JGT  | BPF_X:
        case BPF_JMP   | BPF_JGT  | BPF_K:
        case BPF_JMP   | BPF_JGE  | BPF_X:
        case BPF_JMP   | BPF_JGE  | BPF_K:
        case BPF_JMP   | BPF_JSGT | BPF_X:
        case BPF_JMP   | BPF_JSGT | BPF_K:
        case BPF_JMP   | BPF_JSGE | BPF_X:
        case BPF_JMP   | BPF_JSGE | BPF_K:
        case BPF_JMP   | BPF_JLT  | BPF_X:
        case BPF_JMP   | BPF_JLT  | BPF_K:
        case BPF_JMP   | BPF_JLE  | BPF_X:
        case BPF_JMP   | BPF_JLE  | BPF_K:
        case BPF_JMP   | BPF_JSLT | BPF_X:
        case BPF_JMP   | BPF_JSLT | BPF_K:
        case BPF_JMP   | BPF_JSLE | BPF_X:
        case BPF_JMP   | BPF_JSLE | BPF_K:
        case BPF_JMP32 | BPF_JA:
        case BPF_JMP32 | BPF_JEQ  | BPF_X:
        case BPF_JMP32 | BPF_JEQ  | BPF_K:
        case BPF_JMP32 | BPF_JNE  | BPF_X:
        case BPF_JMP32 | BPF_JNE  | BPF_K:
        case BPF_JMP32 | BPF_JSET | BPF_X:
        case BPF_JMP32 | BPF_JSET | BPF_K:
        case BPF_JMP32 | BPF_JGT  | BPF_X:
        case BPF_JMP32 | BPF_JGT  | BPF_K:
        case BPF_JMP32 | BPF_JGE  | BPF_X:
        case BPF_JMP32 | BPF_JGE  | BPF_K:
        case BPF_JMP32 | BPF_JSGT | BPF_X:
        case BPF_JMP32 | BPF_JSGT | BPF_K:
        case BPF_JMP32 | BPF_JSGE | BPF_X:
        case BPF_JMP32 | BPF_JSGE | BPF_K:
        case BPF_JMP32 | BPF_JLT  | BPF_X:
        case BPF_JMP32 | BPF_JLT  | BPF_K:
        case BPF_JMP32 | BPF_JLE  | BPF_X:
        case BPF_JMP32 | BPF_JLE  | BPF_K:
        case BPF_JMP32 | BPF_JSLT | BPF_X:
        case BPF_JMP32 | BPF_JSLT | BPF_K:
        case BPF_JMP32 | BPF_JSLE | BPF_X:
        case BPF_JMP32 | BPF_JSLE | BPF_K:
                CHECK_RET(handle_jumps(ctx, insn, &len));
                break;
        case BPF_JMP | BPF_CALL:
                CHECK_RET(handle_call(ctx, insn, &len));
                break;

        case BPF_JMP | BPF_EXIT:
                /* If this is the last instruction, epilogue will follow. */
                if (is_last_insn(ctx->prog, idx))
                        break;
                CHECK_RET(handle_jmp_epilogue(ctx, insn, &len));
                break;
        default:
                pr_err("bpf-jit: can't handle instruction code 0x%02X\n", code);
                return -EOPNOTSUPP;
        }

        if (BPF_CLASS(code) == BPF_ALU) {
                /*
                 * Skip the "swap" instructions. Even 64-bit swaps are of type
                 * BPF_ALU (and not BPF_ALU64). Therefore, for the swaps, one
                 * has to look at the "size" of the operations rather than the
                 * ALU type. "gen_swap()" specifically takes care of that.
                 */
                if (BPF_OP(code) != BPF_END && ctx->do_zext)
                        len += zext(BUF(buf, len), dst);
        }

        jit_buffer_update(ctx, len);

        return ret;
}

static int handle_body(struct jit_context *ctx)
{
        int ret;
        bool populate_bpf2insn = false;
        const struct bpf_prog *prog = ctx->prog;

        CHECK_RET(jit_buffer_check(ctx));

        /*
         * Record the mapping for the instructions during the dry-run.
         * Doing it this way allows us to have the mapping ready for
         * the jump instructions during the real compilation phase.
         */
        if (!ctx->emit)
                populate_bpf2insn = true;

        for (u32 i = 0; i < prog->len; i++) {
                /* During the dry-run, jit.len grows gradually per BPF insn. */
                if (populate_bpf2insn)
                        ctx->bpf2insn[i] = ctx->jit.len;

                CHECK_RET(handle_insn(ctx, i));
                if (ret > 0) {
                        /* "ret" is 1 if two (64-bit) chunks were consumed. */
                        ctx->bpf2insn[i + 1] = ctx->bpf2insn[i];
                        i++;
                }
        }

        /* If bpf2insn had to be populated, then it is done at this point. */
        if (populate_bpf2insn)
                ctx->bpf2insn_valid = true;

        return 0;
}

/*
 * Initialize the memory with "unimp_s" which is the mnemonic for
 * "unimplemented" instruction and always raises an exception.
 *
 * The instruction is 2 bytes. If "size" is odd, there is not much
 * that can be done about the last byte in "area". Because, the
 * CPU always fetches instructions in two bytes. Therefore, the
 * byte beyond the last one is going to accompany it during a
 * possible fetch. In the most likely case of a little endian
 * system, that beyond-byte will become the major opcode and
 * we have no control over its initialisation.
 */
static void fill_ill_insn(void *area, unsigned int size)
{
        const u16 unimp_s = 0x79e0;

        if (size & 1) {
                *((u8 *)area + (size - 1)) = 0xff;
                size -= 1;
        }

        memset16(area, unimp_s, size >> 1);
}

/* Piece of memory that can be allocated at the beginning of jit_prepare(). */
static int jit_prepare_early_mem_alloc(struct jit_context *ctx)
{
        ctx->bpf2insn = kcalloc(ctx->prog->len, sizeof(ctx->jit.len),
                                GFP_KERNEL);

        if (!ctx->bpf2insn) {
                pr_err("bpf-jit: could not allocate memory for "
                       "mapping of the instructions.\n");
                return -ENOMEM;
        }

        return 0;
}

/*
 * Memory allocations that rely on parameters known at the end of
 * jit_prepare().
 */
static int jit_prepare_final_mem_alloc(struct jit_context *ctx)
{
        const size_t alignment = sizeof(u32);

        ctx->bpf_header = bpf_jit_binary_alloc(ctx->jit.len, &ctx->jit.buf,
                                               alignment, fill_ill_insn);
        if (!ctx->bpf_header) {
                pr_err("bpf-jit: could not allocate memory for translation.\n");
                return -ENOMEM;
        }

        if (ctx->need_extra_pass) {
                ctx->jit_data = kzalloc_obj(*ctx->jit_data);
                if (!ctx->jit_data)
                        return -ENOMEM;
        }

        return 0;
}

/*
 * The first phase of the translation without actually emitting any
 * instruction. It helps in getting a forecast on some aspects, such
 * as the length of the whole program or where the epilogue starts.
 *
 * Whenever the necessary parameters are known, memories are allocated.
 */
static int jit_prepare(struct jit_context *ctx)
{
        int ret;

        /* Dry run. */
        ctx->emit = false;

        CHECK_RET(jit_prepare_early_mem_alloc(ctx));

        /* Get the length of prologue section after some register analysis. */
        analyze_reg_usage(ctx);
        CHECK_RET(handle_prologue(ctx));

        CHECK_RET(handle_body(ctx));

        /* Record at which offset epilogue begins. */
        ctx->epilogue_offset = ctx->jit.len;

        /* Process the epilogue section now. */
        CHECK_RET(handle_epilogue(ctx));

        CHECK_RET(jit_prepare_final_mem_alloc(ctx));

        return 0;
}

/*
 * jit_compile() is the real compilation phase. jit_prepare() is
 * invoked before jit_compile() as a dry-run to make sure everything
 * will go OK and allocate the necessary memory.
 *
 * In the end, jit_compile() checks if it has produced the same number
 * of instructions as jit_prepare() would.
 */
static int jit_compile(struct jit_context *ctx)
{
        int ret;

        /* Let there be code. */
        ctx->emit = true;

        CHECK_RET(handle_prologue(ctx));

        CHECK_RET(handle_body(ctx));

        CHECK_RET(handle_epilogue(ctx));

        if (ctx->jit.index != ctx->jit.len) {
                pr_err("bpf-jit: divergence between the phases; "
                       "%u vs. %u (bytes).\n",
                       ctx->jit.len, ctx->jit.index);
                return -EFAULT;
        }

        return 0;
}

/*
 * Calling this function implies a successful JIT. A successful
 * translation is signaled by setting the right parameters:
 *
 * prog->jited=1, prog->jited_len=..., prog->bpf_func=...
 */
static int jit_finalize(struct jit_context *ctx)
{
        struct bpf_prog *prog = ctx->prog;

        /* We're going to need this information for the "do_extra_pass()". */
        if (ctx->need_extra_pass) {
                ctx->jit_data->bpf_header = ctx->bpf_header;
                ctx->jit_data->bpf2insn = ctx->bpf2insn;
                prog->aux->jit_data = (void *)ctx->jit_data;
        } else {
                /*
                 * If things seem finalised, then mark the JITed memory
                 * as R-X and flush it.
                 */
                if (bpf_jit_binary_lock_ro(ctx->bpf_header)) {
                        pr_err("bpf-jit: Could not lock the JIT memory.\n");
                        return -EFAULT;
                }
                flush_icache_range((unsigned long)ctx->bpf_header,
                                   (unsigned long)
                                   BUF(ctx->jit.buf, ctx->jit.len));
                prog->aux->jit_data = NULL;
                bpf_prog_fill_jited_linfo(prog, ctx->bpf2insn);
        }

        ctx->success = true;
        prog->bpf_func = (void *)ctx->jit.buf;
        prog->jited_len = ctx->jit.len;
        prog->jited = 1;

        jit_ctx_cleanup(ctx);
        jit_dump(ctx);

        return 0;
}

/*
 * A lenient verification for the existence of JIT context in "prog".
 * Apparently the JIT internals, namely jit_subprogs() in bpf/verifier.c,
 * may request for a second compilation although nothing needs to be done.
 */
static inline int check_jit_context(const struct bpf_prog *prog)
{
        if (!prog->aux->jit_data) {
                pr_notice("bpf-jit: no jit data for the extra pass.\n");
                return 1;
        } else {
                return 0;
        }
}

/* Reuse the previous pass's data. */
static int jit_resume_context(struct jit_context *ctx)
{
        struct arc_jit_data *jdata =
                (struct arc_jit_data *)ctx->prog->aux->jit_data;

        if (!jdata) {
                pr_err("bpf-jit: no jit data for the extra pass.\n");
                return -EINVAL;
        }

        ctx->jit.buf = (u8 *)ctx->prog->bpf_func;
        ctx->jit.len = ctx->prog->jited_len;
        ctx->bpf_header = jdata->bpf_header;
        ctx->bpf2insn = (u32 *)jdata->bpf2insn;
        ctx->bpf2insn_valid = ctx->bpf2insn ? true : false;
        ctx->jit_data = jdata;

        return 0;
}

/*
 * Patch in the new addresses. The instructions of interest are:
 *
 * - call
 * - ld r64, imm64
 *
 * For "call"s, it resolves the addresses one more time through the
 * handle_call().
 *
 * For 64-bit immediate loads, it just retranslates them, because the BPF
 * core in kernel might have changed the value since the normal pass.
 */
static int jit_patch_relocations(struct jit_context *ctx)
{
        const u8 bpf_opc_call = BPF_JMP | BPF_CALL;
        const u8 bpf_opc_ldi64 = BPF_LD | BPF_DW | BPF_IMM;
        const struct bpf_prog *prog = ctx->prog;
        int ret;

        ctx->emit = true;
        for (u32 i = 0; i < prog->len; i++) {
                const struct bpf_insn *insn = &prog->insnsi[i];
                u8 dummy;
                /*
                 * Adjust "ctx.jit.index", so "gen_*()" functions below
                 * can use it for their output addresses.
                 */
                ctx->jit.index = ctx->bpf2insn[i];

                if (insn->code == bpf_opc_call) {
                        CHECK_RET(handle_call(ctx, insn, &dummy));
                } else if (insn->code == bpf_opc_ldi64) {
                        CHECK_RET(handle_ld_imm64(ctx, insn, &dummy));
                        /* Skip the next instruction. */
                        ++i;
                }
        }
        return 0;
}

/*
 * A normal pass that involves a "dry-run" phase, jit_prepare(),
 * to get the necessary data for the real compilation phase,
 * jit_compile().
 */
static struct bpf_prog *do_normal_pass(struct bpf_prog *prog)
{
        struct jit_context ctx;

        /* Bail out if JIT is disabled. */
        if (!prog->jit_requested)
                return prog;

        if (jit_ctx_init(&ctx, prog)) {
                jit_ctx_cleanup(&ctx);
                return prog;
        }

        /* Get the lengths and allocate buffer. */
        if (jit_prepare(&ctx)) {
                jit_ctx_cleanup(&ctx);
                return prog;
        }

        if (jit_compile(&ctx)) {
                jit_ctx_cleanup(&ctx);
                return prog;
        }

        if (jit_finalize(&ctx)) {
                jit_ctx_cleanup(&ctx);
                return prog;
        }

        return ctx.prog;
}

/*
 * If there are multi-function BPF programs that call each other,
 * their translated addresses are not known all at once. Therefore,
 * an extra pass is needed to consult the bpf_jit_get_func_addr()
 * again to get the newly translated addresses in order to resolve
 * the "call"s.
 */
static struct bpf_prog *do_extra_pass(struct bpf_prog *prog)
{
        struct jit_context ctx;

        /* Skip if there's no context to resume from. */
        if (check_jit_context(prog))
                return prog;

        if (jit_ctx_init(&ctx, prog)) {
                jit_ctx_cleanup(&ctx);
                return prog;
        }

        if (jit_resume_context(&ctx)) {
                jit_ctx_cleanup(&ctx);
                return prog;
        }

        if (jit_patch_relocations(&ctx)) {
                jit_ctx_cleanup(&ctx);
                return prog;
        }

        if (jit_finalize(&ctx)) {
                jit_ctx_cleanup(&ctx);
                return prog;
        }

        return ctx.prog;
}

/*
 * This function may be invoked twice for the same stream of BPF
 * instructions. The "extra pass" happens, when there are
 * (re)locations involved that their addresses are not known
 * during the first run.
 */
struct bpf_prog *bpf_int_jit_compile(struct bpf_prog *prog)
{
        vm_dump(prog);

        /* Was this program already translated? */
        if (!prog->jited)
                return do_normal_pass(prog);
        else
                return do_extra_pass(prog);

        return prog;
}