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

package asm

// arm64 frame mapping, matching the Go toolchain's arm64 backend.
//
// Go's arm64 functions use R29 as the frame pointer (FP) and R30 as the link
// register (LR).  R31 is the stack pointer (SP).  FP and SP in the source
// are synthetic pseudo-registers resolved against the hardware SP and the
// frame size.
//
// The autosize is the real stack adjustment: the declared local frame plus
// 8 bytes for the saved link register, rounded up to a 16-byte multiple.
// The toolchain adds an "extrasize" to align: if autosize%16 == 8, add 8;
// if autosize%16 == 0, add 16.
//
// Prologue (autosize > 0, small frame ≤ 0xf0):
//
//	MOVD.W  LR, -autosize(SP)    // pre-index: SP -= autosize, store LR at SP
//	MOVD    FP, -8(SP)            // store FP at SP-8
//	SUB     $8, SP, FP            // FP = SP - 8
//
// Prologue (autosize > 0, large frame > 0xf0):
//
//	SUB     $autosize, SP, R20    // R20 = SP - autosize
//	STP     (FP, LR), -8(R20)    // store FP,LR at R20-8
//	MOVD    R20, SP               // SP = R20
//	SUB     $8, SP, FP            // FP = SP - 8
//
// Epilogue (non-leaf, small frame):
//
//	ADD     $autosize-8, SP, FP   // restore FP
//	ADD     $autosize, SP, SP     // deallocate frame
//	MOVD    -8(SP), FP            // (actually the reverse of prologue)
// Actually:
//	MOVD    -8(SP), FP            // load FP from SP-8
//	MOVD.P  autosize(SP), LR     // post-index: load LR, SP += autosize
//
// Epilogue (non-leaf, large frame):
//	ADD     $autosize-8, SP, FP
//	ADD     $autosize, SP, SP
// Actually:
//	LDP     -8(SP), (FP, LR)     // load FP,LR
//	ADD     $autosize, SP, SP    // deallocate frame
//
// Epilogue (leaf with frame):
//	ADD     $autosize-8, SP, FP
//	ADD     $autosize, SP, SP
//
// RET always emits as BR LR (0xd65f03c0).

import (
	"strings"

	"sourcedock.dev/petrbalvin/gasm-devkit/ast"
)

// arm64FrameInfo holds the frame layout derived from a TEXT directive.
type arm64FrameInfo struct {
	autosize int  // the real SP adjustment (locals + saved LR + alignment)
	frame    int  // the declared $framesize
	args     int  // the declared -argsize
	noSplit  bool // the NOSPLIT flag
	leaf     bool // no call instructions in the body
}

// arm64ComputeFrame derives the frame layout for a TEXT function.
func arm64ComputeFrame(t *ast.Text) arm64FrameInfo {
	fi := arm64FrameInfo{
		frame: frameSize(t),
		args:  argsSize(t),
	}
	for _, f := range t.Flags {
		if f == "NOSPLIT" {
			fi.noSplit = true
		}
	}
	fi.leaf = arm64IsLeaf(t)

	if fi.frame != 0 || !fi.leaf {
		fi.autosize = fi.frame + 8 // space for the saved LR
		if fi.autosize%16 != 0 {
			// The toolchain aligns to 16: if autosize%16 == 8, add 8;
			// otherwise add whatever is needed.
			fi.autosize += 16 - (fi.autosize % 16)
		}
	}
	return fi
}

// arm64IsLeaf reports whether a function contains no call instructions
// (BL/CALL), matching the toolchain's LEAF mark.
func arm64IsLeaf(t *ast.Text) bool {
	for _, stmt := range t.Body {
		in, ok := stmt.(*ast.Instr)
		if !ok {
			continue
		}
		switch strings.ToUpper(in.Mnemonic.Text) {
		case "BL", "CALL":
			return false
		}
	}
	return true
}

// arm64Prologue returns the prologue bytes for an arm64 function.
func arm64Prologue(fi arm64FrameInfo) []byte {
	if fi.autosize == 0 {
		return nil
	}
	if fi.autosize <= 0xf0 {
		// Small frame: MOVD.W LR, -autosize(SP); MOVD FP, -8(SP); SUB $8, SP, FP
		return a64WordsLE(
			arm64PreStoreImm(3, 0, int32(-fi.autosize), 31, 30), // STR.W LR, -autosize(SP) (pre-index store)
			arm64UnscaledStore(3, 0, -8, 31, 29),                // STUR FP, [SP, #-8]
			a64AddSub(1, 1, 0, 0, 8, 31, 29),                    // SUB $8, SP, FP (op=1 for SUB)
		)
	}
	// Large frame: SUB $autosize, SP, R20; STP (FP,LR), -8(R20); ADD $0, R20, SP; SUB $8, SP, FP
	return a64WordsLE(
		a64AddSub(1, 1, 0, 0, uint32(fi.autosize), 31, 20), // SUB $autosize, SP, R20
		a64LSP(2, 0, 0, -1, 30, 20, 29),                    // STP FP, LR, [R20, #-8] (opc=2 for 64-bit pair)
		a64AddSub(1, 0, 0, 0, 0, 20, 31),                   // ADD $0, R20, SP (= MOV R20, SP)
		a64AddSub(1, 1, 0, 0, 8, 31, 29),                   // SUB $8, SP, FP (op=1 for SUB)
	)
}

// arm64Return returns the bytes for a RET: the epilogue (restore FP/LR and
// deallocate the frame when present) followed by RET (BR LR).
func arm64Return(fi arm64FrameInfo) []byte {
	var ws []uint32
	if fi.autosize != 0 {
		if fi.leaf {
			// Leaf with frame: ADD $autosize-8, SP, FP; ADD $autosize, SP, SP
			ws = append(ws,
				a64AddSub(1, 0, 0, 0, uint32(fi.autosize-8), 31, 29), // ADD $autosize-8, SP, FP
				a64AddSub(1, 0, 0, 0, uint32(fi.autosize), 31, 31),   // ADD $autosize, SP, SP
			)
		} else if fi.autosize <= 0xf0 {
			// Non-leaf small frame: LDR FP, [SP, #-8]; LDR.P LR, [SP], #autosize
			ws = append(ws,
				arm64UnscaledLoad(3, 0, -8, 31, 29),             // LDR FP, [SP, #-8]
				arm64PostLoad(3, 0, int32(fi.autosize), 31, 30), // LDR.P LR, [SP], #autosize
			)
		} else {
			// Large frame: LDP -8(SP), (FP, LR); ADD $autosize, SP, SP
			ws = append(ws,
				a64LSP(2, 0, 1, -1, 30, 31, 29),                    // LDP FP, LR, [SP, #-8] (opc=2 for 64-bit pair)
				a64AddSub(1, 0, 0, 0, uint32(fi.autosize), 31, 31), // ADD $autosize, SP, SP
			)
		}
	}
	// RET: BR LR (0xd65f03c0)
	ws = append(ws, a64UncondBranch(2, 30, 0)) // opc=2(RET), Rn=LR(30), Rd=0
	return a64WordsLE(ws...)
}

// arm64PrologueSpadjPC returns the function-relative byte offset where the
// prologue has finished decrementing SP (the delta becomes autosize).
func arm64PrologueSpadjPC(fi arm64FrameInfo) int {
	if fi.autosize == 0 {
		return 0
	}
	if fi.autosize <= 0xf0 {
		return 4 // MOVD.W instruction decrements SP
	}
	return 8 // SUB + STP + MOVD (3 instructions, SP updated at the MOVD)
}

// arm64ReturnEpilogueLen returns the byte length of the RET's epilogue up to
// (but not including) the final RET instruction.
func arm64ReturnEpilogueLen(fi arm64FrameInfo) int {
	if fi.autosize == 0 {
		return 0
	}
	if fi.leaf {
		return 8 // ADD + ADD
	}
	if fi.autosize <= 0xf0 {
		return 8 // LDR + LDR.P
	}
	return 8 // LDP + ADD
}

// arm64ResolvePseudo translates a pseudo-register memory reference into a
// hardware base register and offset.  x+N(FP) → (N + autosize + 8)(SP);
// x+N(SP) → (N + frame + 8)(SP).  Returns base = -1 for an unresolvable
// reference (SB: static data, handled by the relocation path).
//
// The Go toolchain resolves all pseudo-register references against the
// hardware stack pointer (R31/SP): FP references add autosize+8 (the
// distance from SP after the prologue to the caller's argument area),
// SP references add frame+8 (the distance to the local area).
func arm64ResolvePseudo(sym *ast.Symbol, fi arm64FrameInfo) (base int, off int32) {
	if sym == nil {
		return -1, 0
	}
	switch sym.Pseudo {
	case "FP":
		return 31, int32(sym.Offset) + int32(fi.autosize) + 8
	case "SP":
		return 31, int32(sym.Offset) + int32(fi.frame) + 8
	case "SB":
		return -1, int32(sym.Offset)
	}
	return -1, 0
}

// arm64PreStoreImm encodes a pre-index store (STR with writeback):
// size<<30 | 7<<27 | V<<26 | opc<<22 | 1<<11 | 1<<10 | imm9<<12 | Rn<<5 | Rt.
func arm64PreStoreImm(size, V int, imm9 int32, rn, rt int) uint32 {
	return uint32(size)<<30 | 7<<27 | uint32(V)<<26 | 0<<22 |
		3<<10 | (uint32(imm9)&0x1FF)<<12 | uint32(rn&31)<<5 | uint32(rt&31)
}

// arm64UnscaledStore encodes an unscaled store (STUR):
// size<<30 | 7<<27 | V<<26 | opc<<22 | 0<<11 | 0<<10 | imm9<<12 | Rn<<5 | Rt.
func arm64UnscaledStore(size, V int, imm9 int32, rn, rt int) uint32 {
	return uint32(size)<<30 | 7<<27 | uint32(V)<<26 | 0<<22 |
		(uint32(imm9)&0x1FF)<<12 | uint32(rn&31)<<5 | uint32(rt&31)
}

// arm64UnscaledLoad encodes an unscaled load (LDUR):
// size<<30 | 7<<27 | V<<26 | opc<<22 | 0<<11 | 0<<10 | imm9<<12 | Rn<<5 | Rt.
func arm64UnscaledLoad(size, V int, imm9 int32, rn, rt int) uint32 {
	return uint32(size)<<30 | 7<<27 | uint32(V)<<26 | 1<<22 |
		(uint32(imm9)&0x1FF)<<12 | uint32(rn&31)<<5 | uint32(rt&31)
}

// arm64PostLoad encodes a post-index load (LDR with post-increment):
// size<<30 | 7<<27 | V<<26 | opc<<22 | 0<<11 | 1<<10 | imm9<<12 | Rn<<5 | Rt.
func arm64PostLoad(size, V int, imm9 int32, rn, rt int) uint32 {
	return uint32(size)<<30 | 7<<27 | uint32(V)<<26 | 1<<22 |
		1<<10 | (uint32(imm9)&0x1FF)<<12 | uint32(rn&31)<<5 | uint32(rt&31)
}