/* Tail call offset to jump into */
#ifdef CONFIG_HIVE
#if IS_ENABLED(CONFIG_ARM64_BTI_KERNEL)
#define PROLOGUE_OFFSET 8 + 6
#else
#define PROLOGUE_OFFSET 7 + 6
#endif
#else
#define PROLOGUE_OFFSET (BTI_INSNS + 2 + 8)
#endif

static int build_prologue(struct jit_ctx *ctx, bool ebpf_from_cbpf)
{
	const struct bpf_prog *prog = ctx->prog;
	const u8 r6 = bpf2a64[BPF_REG_6];
	const u8 r7 = bpf2a64[BPF_REG_7];
	const u8 r8 = bpf2a64[BPF_REG_8];
	const u8 r9 = bpf2a64[BPF_REG_9];
	const u8 fp = bpf2a64[BPF_REG_FP];
	const u8 tcc = bpf2a64[TCALL_CNT];
	const u8 fpb = bpf2a64[FP_BOTTOM];
	#ifdef CONFIG_HIVE
	const u8 base = bpf2a64[BPF_REG_BASE];
	u64 sfi_base, sfi_stack;
	int idx_before;
	#endif
	const int idx0 = ctx->idx;
	int cur_offset;

	/*
	 * BPF prog stack layout
	 *
	 *                         high
	 * original A64_SP =>   0:+-----+ BPF prologue
	 *                        |FP/LR|
	 * current A64_FP =>  -16:+-----+
	 *                        | ... | callee saved registers
	 * BPF fp register => -64:+-----+ <= (BPF_FP)
	 *                        |     |
	 *                        | ... | BPF prog stack
	 *                        |     |
	 *                        +-----+ <= (BPF_FP - prog->aux->stack_depth)
	 *                        |RSVD | padding
	 * current A64_SP =>      +-----+ <= (BPF_FP - ctx->stack_size)
	 *                        |     |
	 *                        | ... | Function call stack
	 *                        |     |
	 *                        +-----+
	 *                          low
	 *
	 */

	/* bpf function may be invoked by 3 instruction types:
	 * 1. bl, attached via freplace to bpf prog via short jump
	 * 2. br, attached via freplace to bpf prog via long jump
	 * 3. blr, working as a function pointer, used by emit_call.
	 * So BTI_JC should used here to support both br and blr.
	 */
	emit_bti(A64_BTI_JC, ctx);

	emit(A64_MOV(1, A64_R(9), A64_LR), ctx);
	emit(A64_NOP, ctx);

	/* Save FP and LR registers to stay align with ARM64 AAPCS */
	emit(A64_PUSH(A64_FP, A64_LR, A64_SP), ctx);
	emit(A64_MOV(1, A64_FP, A64_SP), ctx);

	/* Save callee-saved registers */
	emit(A64_PUSH(r6, r7, A64_SP), ctx);
	emit(A64_PUSH(r8, r9, A64_SP), ctx);
	emit(A64_PUSH(fp, tcc, A64_SP), ctx);
	emit(A64_PUSH(fpb, A64_R(28), A64_SP), ctx);
	#ifdef CONFIG_HIVE
	emit(A64_PUSH(A64_R(27), base, A64_SP), ctx);
	#endif

	/* Set up BPF prog stack base register */
	#ifdef CONFIG_HIVE
		sfi_base = (u64)prog->shadow_region_addr;
		sfi_stack = (u64)(prog->shadow_stack_addr);

		fbpf_log("fp=%016llx, base=%016llx\n", sfi_stack, sfi_base);
		
		idx_before = ctx->idx;
		emit_a64_mov_i64(fp, sfi_stack, ctx); // 3 insns
		while (ctx->idx - idx_before < 3) {
			emit(A64_NOP, ctx);
		}
		idx_before = ctx->idx;
		emit_a64_mov_i64(base, sfi_base, ctx); // 3 insns
		while (ctx->idx - idx_before < 3) {
			emit(A64_NOP, ctx);
		}	
	#else
		emit(A64_MOV(1, fp, A64_SP), ctx);
	#endif

	if (!ebpf_from_cbpf) {
		/* Initialize tail_call_cnt */
		emit(A64_MOVZ(1, tcc, 0, 0), ctx);

		cur_offset = ctx->idx - idx0;
		if (cur_offset != PROLOGUE_OFFSET) {
			#ifdef CONFIG_HIVE
			pr_err("PROLOGUE_OFFSET = %d, expected %d!\n",
				    cur_offset, PROLOGUE_OFFSET);
			#else
			pr_err_once("PROLOGUE_OFFSET = %d, expected %d!\n",
				    cur_offset, PROLOGUE_OFFSET);
			#endif
			return -1;
		}

		/* BTI landing pad for the tail call, done with a BR */
		emit_bti(A64_BTI_J, ctx);
	}

	emit(A64_SUB_I(1, fpb, fp, ctx->fpb_offset), ctx);

	ctx->stack_size = STACK_ALIGN(prog->aux->stack_depth);

	/* Set up function call stack */
	emit(A64_SUB_I(1, A64_SP, A64_SP, ctx->stack_size), ctx);
	return 0;
}