// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
// SPDX-License-Identifier: BSD-3-Clause

#include "textflag.h"

// ABI-checking trampoline for arm64.  Sets sentinel values in the
// registers the Go ABI fixes across calls before entering the JIT function
// and checks whether they survived on return.
//
// The return trampoline (leaveJITCheckedRaw) is a raw TEXT symbol with no
// Go function declaration, so the toolchain does NOT interpose an
// ABIInternal wrapper; the JIT function RETs directly into the check
// code, which sees the registers exactly as the function left them.
//
// Go ABI on arm64 guarantees:
//   - R29 is the frame pointer (NOSPLIT frame=0 functions must not touch it).
//   - R28 is the goroutine pointer (g) and must survive across any call.
//     The assembler spells this register "g"; R28 is not accepted.
//   - R18 is the platform register and must never be written.  The Go
//     assembler offers no spelling that addresses it, so the check below
//     cannot cover it.

// Sentinel values chosen to be unlikely in normal execution.
#define SENTINEL_FP 0xDEADBEEFCAFEF00D
#define SENTINEL_G 0x0BADF00DDEADBEEF

// GLOBL holding the raw address of the leave trampoline, read by Go.
GLOBL ·leaveCheckedPtr(SB), NOPTR, $8
DATA ·leaveCheckedPtr(SB)/8, $·leaveJITCheckedRaw(SB)

// func enterJITChecked(fn uintptr, stack uintptr)
// Sets sentinels in R29, R28 and R18, switches to the prepared stack and
// branches to fn.  The prepared stack's first word must be the address of
// leaveJITCheckedRaw (read from leaveCheckedPtr).  Only R0 and R3 are used
// as scratch: caller-saved, and not among the checked registers.
TEXT ·enterJITChecked(SB), NOSPLIT, $0-16
	MOVD fn+0(FP), R0 // target (before SP switch)
	MOVD R30, savedLR(SB) // save link register
	MOVD RSP, R3 // save Go stack pointer: R3 relays the value, the
	MOVD R3, savedSP(SB) // arm64 assembler cannot store RSP to memory
	MOVD R29, savedFP(SB) // save frame pointer (vet requires save before clobber)
	MOVD g, savedG(SB) // save g
	MOVD $SENTINEL_FP, R29 // sentinel in the frame pointer
	MOVD $SENTINEL_G, g // sentinel in g
	MOVD stack+8(FP), R3 // load prepared stack pointer
	MOVD 0(R3), R30 // load leaveJITCheckedRaw into LR
	MOVD R3, RSP // SP stays on the leave slot: the kernel
	// reads its first argument at SP+8
	JMP  (R0) // branch to JIT function

	// leaveJITCheckedRaw is the raw return trampoline.  It has NO Go function
	// declaration, so no ABIInternal wrapper is generated; the JIT function's
	// RET lands here directly, seeing R29 and g exactly as the function left
	// them.  It checks the sentinels, records violations in abiResult, then
	// restores the Go stack and returns.
TEXT ·leaveJITCheckedRaw(SB), NOSPLIT, $0-0
	MOVD $0, R4 // accumulated violation bits

	// Check the frame pointer against the sentinel.
	MOVD $SENTINEL_FP, R3
	CMP  R29, R3
	BEQ  fp_ok
	MOVD $1, R5
	ORR  R5, R4, R4

fp_ok:
	// Check g against the sentinel.
	MOVD $SENTINEL_G, R3
	CMP  g, R3
	BEQ  g_ok
	MOVD $2, R5
	ORR  R5, R4, R4

g_ok:
	CBZ  R4, restore
	MOVD R4, ·abiResult(SB)

restore:
	MOVD savedSP(SB), R3 // restore Go stack pointer
	MOVD R3, RSP
	MOVD savedLR(SB), R30 // restore link register
	MOVD savedFP(SB), R29 // restore frame pointer: Go code needs it the
	// moment it resumes, violation or not
	MOVD savedG(SB), g // restore g
	RET  // return to Go caller

// Package-level storage for the saved frame pointer.  Like savedSP and
// savedLR in trampoline_arm64.s, this is assembly-side state: the amd64
// checked trampoline saves and restores the caller's frame pointer too
// (leaveJITCheckedRaw restores BP before returning), and this file mirrors
// that.
GLOBL savedFP(SB), NOPTR, $8

GLOBL savedG(SB), NOPTR, $8