/*
 * Copyright (c) 2021 BayLibre SAS
 * Written by: Nicolas Pitre
 *
 * SPDX-License-Identifier: Apache-2.0
 */

#include <zephyr/kernel.h>
#include <zephyr/kernel_structs.h>
#include <kernel_arch_interface.h>
#include <zephyr/arch/cpu.h>
#include <zephyr/sys/barrier.h>
#include <zephyr/sys/atomic.h>

/* to be found in fpu.S */
extern void z_arm64_fpu_save(struct z_arm64_fp_context *saved_fp_context);
extern void z_arm64_fpu_restore(struct z_arm64_fp_context *saved_fp_context);
extern void z_arm64_sve_save(struct z_arm64_fp_context *saved_fp_context);
extern void z_arm64_sve_restore(struct z_arm64_fp_context *saved_fp_context);

#define FPU_DEBUG 0

#if FPU_DEBUG

/*
 * Debug traces have to be produced without printk() or any other functions
 * using a va_list as va_start() always copy the FPU registers that could be
 * used to pass float arguments, and that triggers an FPU access trap.
 */

#include <string.h>

static char *dbg_prefix(char *buf, char *msg, struct k_thread *th)
{
	strcpy(buf, "CPU# exc# ");
	buf[3] = '0' + _current_cpu->id;
	buf[8] = '0' + arch_exception_depth();
	strcat(buf, _current->name);
	strcat(buf, ": ");
	strcat(buf, msg);
	if (th != NULL) {
		strcat(buf, " ");
		strcat(buf, th->name);
	}
	return buf + strlen(buf);
}

static void DBG(char *msg, struct k_thread *th)
{
	char buf[80], *p;
	unsigned int v;

	p = dbg_prefix(buf, msg, th);

	if (th == NULL) {
		th = _current;
	}
	v = *(unsigned char *)&th->arch.saved_fp_context.neon;
	*p++ = ' ';
	*p++ = ((v >> 4) < 10) ? ((v >> 4) + '0') : ((v >> 4) - 10 + 'a');
	*p++ = ((v & 15) < 10) ? ((v & 15) + '0') : ((v & 15) - 10 + 'a');
	*p++ = '\n';
	*p = 0;

	k_str_out(buf, p - buf);
}

static void DBG_PC(char *msg, uintptr_t pc)
{
	char buf[80], *p;
	uintptr_t addr = pc;
	int i;

	p = dbg_prefix(buf, msg, NULL);
	strcpy(p, " PC=0x");
	p += 6;

	/* Format PC address as hex */
	for (i = (sizeof(uintptr_t) * 2) - 1; i >= 0; i--) {
		unsigned int nibble = (addr >> (i * 4)) & 0xf;
		*p++ = (nibble < 10) ? (nibble + '0') : (nibble - 10 + 'a');
	}
	*p++ = '\n';
	*p = 0;

	k_str_out(buf, p - buf);
}

#else

static inline void DBG(char *msg, struct k_thread *t) { }
static inline void DBG_PC(char *msg, uintptr_t pc) { }

#endif /* FPU_DEBUG */

#ifdef CONFIG_ARM64_SVE

/* Get current SVE vector length */
static inline uint32_t z_arm64_sve_get_vl(void)
{
	uint32_t vl;

	__asm__("rdvl %0, #1" : "=r"(vl));
	return vl;
}

#define USE_SVE(t) ((t) && (t)->arch.saved_fp_context.sve.simd_mode == SIMD_SVE)

/* Convert NEON V registers to SVE Z registers in place */
static void convert_Vx_to_Zx(struct z_arm64_fp_context *context)
{
	uint32_t vl = z_arm64_sve_get_vl();

	if (CONFIG_ARM64_SVE_VL_MAX <= 16 || vl <= 16) {
		return;
	}

	/*
	 * Since it's a union, we need to extend each 128-bit NEON register
	 * to the full SVE vector length, working backwards to avoid overwriting
	 * data we still need to copy.
	 */
	for (int i = 31; i >= 0; i--) {
		/* Copy the 128-bit NEON value to the low 128 bits of the Z register */
		*(__int128 *)&context->sve.z_regs[i * vl] = context->neon.v_regs[i];

		/* Zero the upper part of the Z register (beyond 128 bits) */
		memset(&context->sve.z_regs[i * vl + 16], 0, vl - 16);
	}
}

#else
#define USE_SVE(t) false
#endif

/*
 * Flush FPU content and disable access.
 * This is called locally and also from flush_fpu_ipi_handler().
 */
void arch_flush_local_fpu(void)
{
	__ASSERT(read_daif() & DAIF_IRQ_BIT, "must be called with IRQs disabled");

	struct k_thread *owner = atomic_ptr_get(&_current_cpu->arch.fpu_owner);

	if (owner != NULL) {
		uint64_t cpacr = read_cpacr_el1();

		/* turn on FPU access */
		cpacr |= CPACR_EL1_FPEN;
		if (USE_SVE(owner)) {
			cpacr |= CPACR_EL1_ZEN;
		}
		write_cpacr_el1(cpacr);
		barrier_isync_fence_full();

		/* save current owner's content */
		if (USE_SVE(owner)) {
			z_arm64_sve_save(&owner->arch.saved_fp_context);
		} else {
			z_arm64_fpu_save(&owner->arch.saved_fp_context);
		}

		/* make sure content made it to memory before releasing */
		barrier_dsync_fence_full();

		/* release ownership */
		atomic_ptr_clear(&_current_cpu->arch.fpu_owner);
		DBG("disable", owner);

		/* disable FPU access */
		cpacr &= ~(CPACR_EL1_FPEN | CPACR_EL1_ZEN);
		write_cpacr_el1(cpacr);
		barrier_isync_fence_full();
	}
}

#ifdef CONFIG_SMP
static void flush_owned_fpu(struct k_thread *thread)
{
	__ASSERT(read_daif() & DAIF_IRQ_BIT, "must be called with IRQs disabled");

	int i;

	/* search all CPUs for the owner we want */
	unsigned int num_cpus = arch_num_cpus();

	for (i = 0; i < num_cpus; i++) {
		if (atomic_ptr_get(&_kernel.cpus[i].arch.fpu_owner) != thread) {
			continue;
		}
		/* we found it live on CPU i */
		if (i == _current_cpu->id) {
			arch_flush_local_fpu();
		} else {
			/* the FPU context is live on another CPU */
			arch_flush_fpu_ipi(i);

			/*
			 * Wait for it only if this is about the thread
			 * currently running on this CPU. Otherwise the
			 * other CPU running some other thread could regain
			 * ownership the moment it is removed from it and
			 * we would be stuck here.
			 *
			 * Also, if this is for the thread running on this
			 * CPU, then we preemptively flush any live context
			 * on this CPU as well since we're likely to
			 * replace it, and this avoids a deadlock where
			 * two CPUs want to pull each other's FPU context.
			 */
			if (thread == _current) {
				arch_flush_local_fpu();
				while (atomic_ptr_get(&_kernel.cpus[i].arch.fpu_owner) == thread) {
					barrier_dsync_fence_full();
				}
			}
		}
		break;
	}
}
#endif

void z_arm64_fpu_enter_exc(void)
{
	__ASSERT(read_daif() & DAIF_IRQ_BIT, "must be called with IRQs disabled");

	/* always deny FPU access whenever an exception is entered */
	write_cpacr_el1(read_cpacr_el1() & ~(CPACR_EL1_FPEN | CPACR_EL1_ZEN));
	barrier_isync_fence_full();
}

/*
 * Simulate some FPU store instructions.
 *
 * In many cases, the FPU trap is triggered by va_start() that copies
 * the content of FP registers used for floating point argument passing
 * into the va_list object in case there were actual float arguments from
 * the caller. In practice this is almost never the case, especially if
 * FPU access is disabled and we're trapped while in exception context.
 * Rather than flushing the FPU context to its owner and enabling access
 * just to let the corresponding STR instructions execute, we simply
 * simulate them and leave the FPU access disabled. This also avoids the
 * need for disabling interrupts in syscalls and IRQ handlers as well.
 */
static bool simulate_str_q_insn(struct arch_esf *esf)
{
	/*
	 * Support only the "FP in exception" cases for now.
	 * We know there is no saved FPU context to check nor any
	 * userspace stack memory to validate in that case.
	 */
	if (arch_exception_depth() <= 1) {
		return false;
	}

	uint32_t *pc = (uint32_t *)esf->elr;
	/* The original (interrupted) sp is the top of the esf structure */
	uintptr_t sp = (uintptr_t)esf + sizeof(*esf);

	for (;;) {
		uint32_t insn = *pc;

		/*
		 * We're looking for STR (immediate, SIMD&FP) of the form:
		 *
		 *  STR Q<n>, [SP, #<pimm>]
		 *
		 * where 0 <= <n> <= 7 and <pimm> is a 12-bits multiple of 16.
		 */
		if ((insn & 0xffc003f8) != 0x3d8003e0) {
			break;
		}

		uint32_t pimm = (insn >> 10) & 0xfff;

		/* Zero the location as the above STR would have done */
		*(__int128 *)(sp + pimm * 16) = 0;

		/* move to the next instruction */
		pc++;
	}

	/* did we do something? */
	if (pc != (uint32_t *)esf->elr) {
		DBG_PC("simulated", esf->elr);
		/* resume execution past the simulated instructions */
		esf->elr = (uintptr_t)pc;
		return true;
	}

	return false;
}

/*
 * Process the FPU trap.
 *
 * This usually means that FP regs belong to another thread. Save them
 * to that thread's save area and restore the current thread's content.
 *
 * We also get here when FP regs are used while in exception as FP access
 * is always disabled by default in that case. If so we save the FPU content
 * to the owning thread and simply enable FPU access. Exceptions should be
 * short and don't have persistent register contexts when they're done so
 * there is nothing to save/restore for that context... as long as we
 * don't get interrupted that is. To ensure that we mask interrupts to
 * the triggering exception context.
 */
void z_arm64_fpu_trap(struct arch_esf *esf, uint32_t exception_class)
{
	__ASSERT(read_daif() & DAIF_IRQ_BIT, "must be called with IRQs disabled");

	/* check if a quick simulation can do it */
	if (!(IS_ENABLED(CONFIG_ARM64_SVE) && exception_class == 0x19) &&
	    simulate_str_q_insn(esf)) {
		return;
	}

	DBG_PC("trap entry", esf->elr);

	struct k_thread *owner = atomic_ptr_get(&_current_cpu->arch.fpu_owner);
	uint64_t cpacr = read_cpacr_el1();

	/* turn on FPU access */
	cpacr |= CPACR_EL1_FPEN;
	if (USE_SVE(owner)) {
		cpacr |= CPACR_EL1_ZEN;
	}
	write_cpacr_el1(cpacr);
	barrier_isync_fence_full();

	/* save current owner's content  if any */
	if (owner) {
		if (USE_SVE(owner)) {
			z_arm64_sve_save(&owner->arch.saved_fp_context);
			DBG("sve_save", owner);
		} else {
			z_arm64_fpu_save(&owner->arch.saved_fp_context);
			DBG("fpu_save", owner);
		}
		barrier_dsync_fence_full();
		atomic_ptr_clear(&_current_cpu->arch.fpu_owner);
	}

	if (IS_ENABLED(CONFIG_ARM64_SVE) && exception_class == 0x19 &&
	    !(cpacr & CPACR_EL1_ZEN)) {
		/* SVE trap - also enable SVE access */
		cpacr |= CPACR_EL1_ZEN;
		write_cpacr_el1(cpacr);
		barrier_isync_fence_full();
	}

	if (arch_exception_depth() > 1) {
		/*
		 * We were already in exception when the FPU access trap.
		 * We give it access and prevent any further IRQ recursion
		 * by disabling IRQs as we wouldn't be able to preserve the
		 * interrupted exception's FPU context.
		 */
		esf->spsr |= DAIF_IRQ_BIT;
		return;
	}

#ifdef CONFIG_SMP
	/*
	 * Make sure the FPU context we need isn't live on another CPU.
	 * The current CPU's FPU context is NULL at this point.
	 */
	flush_owned_fpu(_current);
#endif

	/* become new owner */
	atomic_ptr_set(&_current_cpu->arch.fpu_owner, _current);

#ifdef CONFIG_ARM64_SVE
	if (exception_class == 0x19) {
		/* SVE trap */
		if (_current->arch.saved_fp_context.sve.simd_mode == SIMD_NEON) {
			/* upgrade from Neon to SVE before loading regs */
			convert_Vx_to_Zx(&_current->arch.saved_fp_context);
		}
		_current->arch.saved_fp_context.sve.simd_mode = SIMD_SVE;
	} else if (_current->arch.saved_fp_context.sve.simd_mode != SIMD_SVE) {
		/* not SVE trap and context is not SVE either */
		if ((cpacr & CPACR_EL1_ZEN) != 0) {
			/* disable SVE access leaving only FP */
			cpacr &= ~CPACR_EL1_ZEN;
			write_cpacr_el1(cpacr);
			barrier_isync_fence_full();
		}
		_current->arch.saved_fp_context.sve.simd_mode = SIMD_NEON;
	} else if ((cpacr & CPACR_EL1_ZEN) == 0) {
		/*
		 * Not SVE trap but context is SVE and CPACR_EL1_ZEN not set.
		 * This is an edge case that happens when previous owner
		 * didn't use SVE and we are not using SVE right now either
		 * although we did in the past. We're about to restore an
		 * SVE context so make sure SVE access is enabled.
		 */
		cpacr |= CPACR_EL1_ZEN;
		write_cpacr_el1(cpacr);
		barrier_isync_fence_full();
	}
#endif

	/* restore our content */
	if (USE_SVE(_current)) {
		z_arm64_sve_restore(&_current->arch.saved_fp_context);
		DBG("sve_restore", NULL);
	} else {
		z_arm64_fpu_restore(&_current->arch.saved_fp_context);
		DBG("fpu_restore", NULL);
	}
}

/*
 * Perform lazy FPU context switching by simply granting or denying
 * access to FP regs based on FPU ownership before leaving the last
 * exception level in case of exceptions, or during a thread context
 * switch with the exception level of the new thread being 0.
 * If current thread doesn't own the FP regs then it will trap on its
 * first access and then the actual FPU context switching will occur.
 */
static void fpu_access_update(unsigned int exc_update_level)
{
	__ASSERT(read_daif() & DAIF_IRQ_BIT, "must be called with IRQs disabled");

	uint64_t cpacr = read_cpacr_el1();

	if (arch_exception_depth() == exc_update_level) {
		/* We're about to execute non-exception code */
		if (atomic_ptr_get(&_current_cpu->arch.fpu_owner) == _current) {
			/* turn on FPU access */
			cpacr |= CPACR_EL1_FPEN;
			if (USE_SVE(_current)) {
				cpacr |= CPACR_EL1_ZEN;
			}
			write_cpacr_el1(cpacr);
		} else {
			/* deny FPU access */
			write_cpacr_el1(cpacr & ~(CPACR_EL1_FPEN | CPACR_EL1_ZEN));
		}
	} else {
		/*
		 * Any new exception level should always trap on FPU
		 * access as we want to make sure IRQs are disabled before
		 * granting it access (see z_arm64_fpu_trap() documentation).
		 */
		write_cpacr_el1(cpacr & ~(CPACR_EL1_FPEN | CPACR_EL1_ZEN));
	}
	barrier_isync_fence_full();
}

/*
 * This is called on every exception exit except for z_arm64_fpu_trap().
 * In that case the exception level of interest is 1 (soon to be 0).
 */
void z_arm64_fpu_exit_exc(void)
{
	fpu_access_update(1);
}

/*
 * This is called from z_arm64_context_switch(). FPU access may be granted
 * only if exception level is 0. If we switch to a thread that is still in
 * some exception context then FPU access would be re-evaluated at exception
 * exit time via z_arm64_fpu_exit_exc().
 */
void z_arm64_fpu_thread_context_switch(void)
{
	fpu_access_update(0);
}

int arch_float_disable(struct k_thread *thread)
{
	if (thread != NULL) {
		unsigned int key = arch_irq_lock();

#ifdef CONFIG_SMP
		flush_owned_fpu(thread);
#else
		if (thread == atomic_ptr_get(&_current_cpu->arch.fpu_owner)) {
			arch_flush_local_fpu();
		}
#endif

		arch_irq_unlock(key);
	}

	return 0;
}

int arch_float_enable(struct k_thread *thread, unsigned int options)
{
	ARG_UNUSED(thread);
	ARG_UNUSED(options);
	/* floats always gets enabled automatically at the moment */
	return 0;
}
