// 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(sizeof(*ctx->jit_data), GFP_KERNEL);
		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;
}