/* SPDX-License-Identifier: GPL-2.0+ */
/*
 * mmu.c - VMSAv8-64 stage 1 identity map for EL2.
 *
 * One level 0 table plus the subtables for the low 1GB of MMIO and
 * the RAM region. everything is identity mapped, the bootloader
 * never needs a different VA view, it just needs caching rules that
 * let the payload start from an architecture-defined state.
 *
 * The descriptor layouts are from the manual (DDI 0487), level 0/1/2
 * and level 3 formats at D5-2444 and D5-2447, attribute fields at
 * D5-2451, MAIR at D5-2476. feature bits come from the ID registers,
 * never hardcoded, the PA size from ID_AA64MMFR0_EL1.PARange per
 * "Address size configuration" D5-2399.
 *
 * Copyright (C) 2026 Bradley Morgan <brads@mainlining.org>
 */

#include <asm/mmu.h>

/* 4KB granule, 3 level tables below level 0 for 1GB blocks */
#define L0_ENTRIES	512
#define L1_ENTRIES	512
#define L2_ENTRIES	512



/*
 * MAIR: attr 0 normal writeback cacheable read allocate, attr 1
 * device nGnRE. encodings straight from D5-2476, B2-122 for the
 * memory types.
 */
#define TB_MAIR_EL2_VAL	0x04ffULL

static uint64_t l0_table[L0_ENTRIES] __attribute__((aligned(4096)));
static uint64_t ram_l1[L1_ENTRIES] __attribute__((aligned(4096)));
static uint64_t ram_l2[L2_ENTRIES] __attribute__((aligned(4096)));

/*
 * Device and normal descriptor templates, upper attributes from
 * D5-2451, the AF is set by hand, hardware page table walks without
 * hardware access flag update will fault otherwise.
 */
#define DEV_DESC(x)	(TB_DESC_BLOCK | TB_DESC_AF | TB_DESC_XN | \
			 TB_DESC_SH_IS | \
			 ((uint64_t)TB_ATTR_DEVICE << 2) | (x))
#define RAM_DESC(x)	(TB_DESC_BLOCK | TB_DESC_AF | TB_DESC_SH_IS | \
			 ((uint64_t)TB_ATTR_NORMAL << 2) | (x))

static void build_identity_map(void)
{
	int i;

	/*
	 * one level 1 table under l0[0], covering the low 512GB. the
	 * MMIO hole and RAM are both in it, device block at index 0
	 * (0..1GB) and the RAM table at index 1 (1GB..2GB).
	 */
	l0_table[0] = TB_DESC_TABLE |
		((uint64_t)(uintptr_t)ram_l1 & ~0xfffULL);

	/* low 1GB, device nGnRE, non executable */
	ram_l1[0] = DEV_DESC(TB_MAP_MMIO_BASE);

	/*
	 * RAM, 0x40000000 for 128MB on qemu virt, normal writeback.
	 * the level 2 table splits the 1GB into 2MB blocks so the map
	 * can be carved later.
	 */
	for (i = 0; i < TB_MAP_RAM_SIZE / (2ULL << 20); i++)
		ram_l2[i] = RAM_DESC(TB_MAP_RAM_BASE + (i * (2ULL << 20)));

	ram_l1[1] = TB_DESC_TABLE |
		((uint64_t)(uintptr_t)ram_l2 & ~0xfffULL);
}

/*
 * clean the table memory to the point of coherency. the tables were
 * written with the dcache off, the page table walker reads them as
 * memory the TCR walk attributes describe, and a dirty line sitting
 * in the cache would never reach RAM. dc cvac is by cache line, walk
 * every page of table memory.
 */
static void tb_clean_tables(void)
{
	uint64_t addr;
	uint64_t tables[] = { (uint64_t)(uintptr_t)l0_table,
			      (uint64_t)(uintptr_t)ram_l1,
			      (uint64_t)(uintptr_t)ram_l2 };
	int i;

	for (i = 0; i < 3; i++) {
		for (addr = tables[i]; addr < tables[i] + 4096; addr += 64) {
			asm volatile("dc cvac, %0" :: "r" (addr) : "memory");
		}
	}

	asm volatile("dsb sy" ::: "memory");
}

static uint64_t read_parange(void)
{
	uint64_t ips;

	asm volatile("mrs %0, id_aa64mmfr0_el1" : "=r" (ips));
	return (ips >> 0) & 0xf;
}

/*
 * EL aware enable. the EL1&0 regime registers at EL1, the EL2 regime
 * registers at EL2, one code path per the manual, one translation
 * regime per exception level (D1-2146).
 */
int tb_mmu_enable(void)
{
	uint64_t tcr, mair;
	uint64_t el;

	build_identity_map();
	tb_clean_tables();

	asm volatile("mrs %0, CurrentEL" : "=r" (el));
	el >>= 2;

	/* tcr value and PA size, D5-2399 address size configuration */
	tcr = TB_TCR_T0SZ_48 | TB_TCR_SH0_IS | TB_TCR_TG0_4K |
	      TB_TCR_IRGN0_WB | TB_TCR_ORGN0_WB | TB_TCR_IPS(read_parange());
	mair = TB_MAIR_EL2_VAL;

	if (el == 2) {
		asm volatile(
			"dsb	sy\n"
			"msr	ttbr0_el2, %1\n"
			"msr	tcr_el2, %2\n"
			"msr	mair_el2, %3\n"
			"isb\n"
			"tlbi	alle2\n"
			"dsb	sy\n"
			"ic	iallu\n"
			"dsb	sy\n"
			"isb\n"
			: "=r" (tcr)
			: "r" (l0_table), "r" (tcr), "r" (mair)
			: "memory");
		asm volatile(
			"mrs	x0, sctlr_el2\n"
			"orr	x0, x0, #1\n"
			"msr	sctlr_el2, x0\n"
			"isb\n"
			::: "x0", "memory");
	} else {
		asm volatile(
			"dsb	sy\n"
			"msr	ttbr0_el1, %1\n"
			"msr	tcr_el1, %2\n"
			"msr	mair_el1, %3\n"
			"isb\n"
			"tlbi	vmalle1\n"
			"dsb	sy\n"
			"ic	iallu\n"
			"dsb	sy\n"
			"isb\n"
			: "=r" (tcr)
			: "r" (l0_table), "r" (tcr), "r" (mair)
			: "memory");
		asm volatile(
			"mrs	x0, sctlr_el1\n"
			"orr	x0, x0, #1\n"
			"msr	sctlr_el1, x0\n"
			"isb\n"
			::: "x0", "memory");
	}

	return 0;
}

void tb_mmu_disable(void)
{
	uint64_t el;

	asm volatile("mrs %0, CurrentEL" : "=r" (el));
	el >>= 2;

	if (el == 2) {
		asm volatile(
			"mrs	x0, sctlr_el2\n"
			"bic	x0, x0, #1\n"
			"msr	sctlr_el2, x0\n"
			"dsb	sy\n"
			"tlbi	alle2\n"
			"dsb	sy\n"
			"isb\n"
			::: "x0", "memory");
	} else {
		asm volatile(
			"mrs	x0, sctlr_el1\n"
			"bic	x0, x0, #1\n"
			"msr	sctlr_el1, x0\n"
			"dsb	sy\n"
			"tlbi	vmalle1\n"
			"dsb	sy\n"
			"isb\n"
			::: "x0", "memory");
	}
}