// SPDX-License-Identifier: GPL-2.0-only
/*:
 * Hibernate support specific for RISCV
 *
 * Copyright (C) 2022 Shanghai StarFive Technology Co., Ltd.
 *
 * Author: Jee Heng Sia <jeeheng.sia@starfivetech.com>
 *
 */

#include <linux/cpu.h>
#include <linux/memblock.h>
#include <linux/pm.h>
#include <linux/sched.h>
#include <linux/suspend.h>
#include <linux/utsname.h>

#include <asm/barrier.h>
#include <asm/cacheflush.h>
#include <asm/irqflags.h>
#include <asm/kexec.h>
#include <asm/mmu_context.h>
#include <asm/page.h>
#include <asm/sections.h>
#include <asm/set_memory.h>
#include <asm/smp.h>
#include <asm/suspend.h>

#include <soc/sifive/sifive_l2_cache.h>

/*
 * The logical cpu number we should resume on, initialised to a non-cpu
 * number.
 */
static int sleep_cpu = -EINVAL;

/* CPU context to be saved */
struct suspend_context *hibernate_cpu_context;

unsigned long relocated_restore_code;

/* Pointer to the temporary resume page tables */
pgd_t *resume_pg_dir;

/*
 * Values that may not change over hibernate/resume. We put the build number
 * and date in here so that we guarantee not to resume with a different
 * kernel.
 */
struct arch_hibernate_hdr_invariants {
	char		uts_version[__NEW_UTS_LEN + 1];
};

/* These values need to be known across a hibernate/restore. */
static struct arch_hibernate_hdr {
	struct arch_hibernate_hdr_invariants invariants;
	unsigned long	hartid;
	unsigned long	saved_satp;
	unsigned long	restore_cpu_addr;
} resume_hdr;

static inline void arch_hdr_invariants(struct arch_hibernate_hdr_invariants *i)
{
	memset(i, 0, sizeof(*i));
	memcpy(i->uts_version, init_utsname()->version, sizeof(i->uts_version));
}

int pfn_is_nosave(unsigned long pfn)
{
	unsigned long nosave_begin_pfn = sym_to_pfn(&__nosave_begin);
	unsigned long nosave_end_pfn = sym_to_pfn(&__nosave_end - 1);

	return ((pfn >= nosave_begin_pfn) && (pfn <= nosave_end_pfn));
}

void notrace save_processor_state(void)
{
	WARN_ON(num_online_cpus() != 1);
}

void notrace restore_processor_state(void)
{
}


int arch_hibernation_header_save(void *addr, unsigned int max_size)
{
	struct arch_hibernate_hdr *hdr = addr;

	if (max_size < sizeof(*hdr))
		return -EOVERFLOW;

	arch_hdr_invariants(&hdr->invariants);

	hdr->hartid = cpuid_to_hartid_map(sleep_cpu);
	hdr->saved_satp = csr_read(CSR_SATP);
	hdr->restore_cpu_addr = (unsigned long) __hibernate_cpu_resume;

	return 0;
}
EXPORT_SYMBOL(arch_hibernation_header_save);

int arch_hibernation_header_restore(void *addr)
{
	struct arch_hibernate_hdr_invariants invariants;
	struct arch_hibernate_hdr *hdr = addr;
	int ret = 0;

	arch_hdr_invariants(&invariants);

	if (memcmp(&hdr->invariants, &invariants, sizeof(invariants))) {
		pr_crit("Hibernate image not generated by this kernel!\n");
		return -EINVAL;
	}

	sleep_cpu = riscv_hartid_to_cpuid(hdr->hartid);
	if (sleep_cpu < 0) {
		pr_crit("Hibernated on a CPU not known to this kernel!\n");
		sleep_cpu = -EINVAL;
		return -EINVAL;
	}
#ifdef CONFIG_SMP
	ret = bringup_hibernate_cpu(sleep_cpu);
	if (ret) {
		sleep_cpu = -EINVAL;
		return ret;
	}
#endif
	resume_hdr = *hdr;

	return ret;
}
EXPORT_SYMBOL(arch_hibernation_header_restore);

int swsusp_arch_suspend(void)
{
	int ret = 0;

	if (__cpu_suspend_enter(hibernate_cpu_context)) {
		sleep_cpu = smp_processor_id();
		suspend_save_csrs(hibernate_cpu_context);
		ret = swsusp_save();
	} else {
		suspend_restore_csrs(hibernate_cpu_context);
		flush_tlb_all();

		/* Invalidated Icache */
		flush_icache_all();

		/*
		 * Tell the hibernation core that we've just restored
		 * the memory
		 */
		in_suspend = 0;
		sleep_cpu = -EINVAL;
	}

	return ret;
}

void temp_page_mapping(pgd_t *pgdp, unsigned long va, pgprot_t prot)
{
	uintptr_t pgd_idx = pgd_index(va);
	phys_addr_t pmd_phys;
	phys_addr_t pte_phys;
	uintptr_t pmd_idx;
	uintptr_t pte_idx;
	pmd_t *pmdp;
	pte_t *ptep;

	if (pgd_val(pgdp[pgd_idx]) == 0) {
		pmdp = (pmd_t *)get_safe_page(GFP_ATOMIC);
		if (!pmdp)
			return;

		memset(pmdp, 0, PAGE_SIZE);
		pmd_phys = __pa(pmdp);
		pgdp[pgd_idx] = pfn_pgd(PFN_DOWN(pmd_phys), PAGE_TABLE);
	} else {
		pmd_phys = PFN_PHYS(_pgd_pfn(pgdp[pgd_idx]));
		pmdp = (pmd_t *) __va(pmd_phys);
	}

	pmd_idx = pmd_index(va);

	if (pmd_none(pmdp[pmd_idx])) {
		ptep = (pte_t *)get_safe_page(GFP_ATOMIC);
		if (!ptep)
			return;

		memset(ptep, 0, PAGE_SIZE);
		pte_phys = __pa(ptep);
		pmdp[pmd_idx] = pfn_pmd(PFN_DOWN(pte_phys), PAGE_TABLE);
	} else {
		pte_phys = PFN_PHYS(_pmd_pfn(pmdp[pmd_idx]));
		ptep = (pte_t *) __va(pte_phys);
	}

	pte_idx = pte_index(va);

	ptep[pte_idx] = pfn_pte(PFN_DOWN(__pa(va)), prot);
}

unsigned long relocate_restore_code(void)
{
	void *page = (void *)get_safe_page(GFP_ATOMIC);

	if (!page)
		return -ENOMEM;

	memcpy(page, core_restore_code, PAGE_SIZE);

	/* Make the page containing the relocated code executable */
	set_memory_x((unsigned long)page, 1);

	temp_page_mapping(resume_pg_dir, (unsigned long)page, PAGE_KERNEL_READ_EXEC);

	return (unsigned long)page;
}

int swsusp_arch_resume(void)
{
	unsigned long addr;

	resume_pg_dir = (pgd_t *)get_safe_page(GFP_ATOMIC);
	if (!resume_pg_dir)
		return -ENOMEM;

	for (addr = PAGE_OFFSET; addr <= (unsigned long)__va(end_linear_map); addr += PAGE_SIZE)
		temp_page_mapping(resume_pg_dir, addr, PAGE_KERNEL);

	relocated_restore_code = relocate_restore_code();
	temp_page_mapping(resume_pg_dir, (unsigned long)resume_hdr.restore_cpu_addr,
				PAGE_KERNEL_READ_EXEC);

	restore_image(resume_hdr.saved_satp, (PFN_DOWN(__pa(resume_pg_dir)) | SATP_MODE),
			resume_hdr.restore_cpu_addr, (unsigned long)hibernate_cpu_context);

	return 0;
}

#ifdef CONFIG_SMP
int hibernate_resume_nonboot_cpu_disable(void)
{
	if (sleep_cpu < 0) {
		pr_err("Failing to resume from hibernate on an unknown CPU.\n");
		return -ENODEV;
	}

	return freeze_secondary_cpus(sleep_cpu);
}
#endif

static int __init riscv_hibernate__init(void)
{
	hibernate_cpu_context = kcalloc(1, sizeof(struct suspend_context), GFP_KERNEL);

	if (WARN_ON(!hibernate_cpu_context))
		return -ENOMEM;

	return 0;
}

early_initcall(riscv_hibernate__init);