// SPDX-License-Identifier: GPL-2.0
/*
 * DMABUF mtk secure heap helper
 *
 * Copyright (C) 2023 MediaTek Inc.
 *
 */
#define pr_fmt(fmt) "dma_heap: mtk_svp_cma: "fmt

#include <linux/align.h>
#include <linux/atomic.h>
#include <linux/bitmap.h>
#include <linux/cma.h>
#include <linux/compiler.h>
#include <linux/err.h>
#include <linux/genalloc.h>
#include <linux/gfp.h>
#include <linux/list.h>
#include <linux/mm.h>
#include <linux/of_fdt.h>
#include <linux/of_reserved_mem.h>
#include <linux/slab.h>
#include <linux/vmalloc.h>

#include "mtk_sec_heap_cma_helper.h"


#ifdef MTK_DMA_SEC_HEAP_DEBUG
#define DBG_LOG(str, args...)		pr_info(str, ##args)
#else
#define DBG_LOG(str, args...)
#endif

/* The maximum aiviable number of mpu regions for SVP */
#define CMA_BLOCKS_MAX_NUM		(16)
/* CMA_BLOCKS_MAX_NUM should in range (0, MPU_REGION_MAX=32) */
static_assert(CMA_BLOCKS_MAX_NUM > 0 && CMA_BLOCKS_MAX_NUM < 32);

/* emi_mpu requires 64K align for each block */
#define MPU_REGION_ALIGN		(SZ_64K)

struct mtk_cma_memory_info {
	const char		*name;
	struct reserved_mem	*rmem;
	struct cma		*cma;

	uint32_t		cma_block_size;		/* memory size per block */
	uint32_t		cma_block_count;	/* number of blocks */
	uint32_t		cma_block_align;	/* align size of blocks */

	phys_addr_t		mem_start;		/* phy addr of resv mem */
	phys_addr_t		mem_size;
	struct gen_pool         *mem_pool;
	atomic64_t		allocated_size;

	struct mutex            lock;
	uint32_t		cma_holded_block_mask;
	phys_addr_t		cma_holded_size;
};

static struct mtk_cma_memory_info mtk_cma_memory_list[MTK_CMA_MEM_NUM] = {
	[MTK_CMA_MEM_SVP] = {
		.name = "svp_mem",
#ifdef CMA_MEM_ALIGN_FORCE_ENABLE
		.cma_block_count = CMA_BLOCKS_MAX_NUM,
		.cma_block_align = CMA_BLOCKS_MAX_NUM * MPU_REGION_ALIGN,
#endif
	},
};

static struct mtk_cma_memory_info *
mtk_secure_heap_cma_get_mem_info(enum MTK_CMA_MEMORY_ID mem_id)
{
	if (mem_id >= MTK_CMA_MEM_NUM || mem_id <= MTK_CMA_MEM_INV)
		return NULL;

	return &mtk_cma_memory_list[mem_id];
}

int mtk_secure_heap_cma_query_mem_info(struct device *dev,
				       struct mtk_sec_heap_mem_info *info)
{
	struct mtk_cma_memory_info *cma_mem = NULL;
	struct device_node *rmem_node = NULL;
	struct reserved_mem *rmem = NULL;
	enum MTK_CMA_MEMORY_ID mem_id;
	int ret = 0;

	if (!dev || !info)
		return -EINVAL;

	info->cma_mem_id = MTK_CMA_MEM_INV;

	rmem_node = of_parse_phandle(dev->of_node, "memory-region", 0);
	if (rmem_node)
		rmem = of_reserved_mem_lookup(rmem_node);
	if (!rmem) {
		dev_err(dev, "can not found a valid memory-region");
		return -ENODEV;
	}

	for (mem_id = 0; mem_id < MTK_CMA_MEM_NUM; mem_id++)
		if (mtk_cma_memory_list[mem_id].rmem == rmem)
			break;

	cma_mem = mtk_secure_heap_cma_get_mem_info(mem_id);
	if (!cma_mem) {
		dev_err(dev, "can not get memory info by id %d", mem_id);
		return -ENOENT;
	}

	/* create pool here instead of memory_init because
	 * gen_pool_create is not ready when memory_init.
	 */
	mutex_lock(&cma_mem->lock);
	if (!cma_mem->mem_pool) {
		cma_mem->mem_pool = gen_pool_create(PAGE_SHIFT, -1);
		if (!cma_mem->mem_pool) {
			mutex_unlock(&cma_mem->lock);
			pr_err("%s: failed to create gen_pool\n", __func__);
			return -EFAULT;
		}

		/* init mem pool for reserved memory */
		if (!cma_mem->cma) {
			cma_mem->cma_holded_size = cma_mem->mem_size;
			cma_mem->cma_holded_block_mask = BIT(0);
			ret = gen_pool_add(cma_mem->mem_pool, cma_mem->mem_start,
					   cma_mem->mem_size, -1);
			if (ret) {
				gen_pool_destroy(cma_mem->mem_pool);
				cma_mem->mem_pool = NULL;
				mutex_unlock(&cma_mem->lock);
				pr_err("%s: failed to add resv mem, ret %d\n",
					__func__, ret);
				return ret;
			}
		}
	}
	mutex_unlock(&cma_mem->lock);

	info->cma_mem_id = mem_id;
	info->cma_mem_start = cma_mem->mem_start;
	info->cma_mem_size = cma_mem->mem_size;
	info->cma_block_size = cma_mem->cma_block_size;
	info->cma_block_count = cma_mem->cma_block_count;

	return 0;
}

int mtk_secure_heap_cma_query_block_mask(enum MTK_CMA_MEMORY_ID mem_id,
					 uint32_t *holded_block_mask)
{
	struct mtk_cma_memory_info *cma_mem = NULL;

	if (!holded_block_mask)
		return -EINVAL;

	cma_mem = mtk_secure_heap_cma_get_mem_info(mem_id);
	if (!cma_mem)
		return -EINVAL;

	mutex_lock(&cma_mem->lock);
	*holded_block_mask = cma_mem->cma_holded_block_mask;
	mutex_unlock(&cma_mem->lock);

	return 0;
}

int mtk_secure_heap_cma_hold_blocks(enum MTK_CMA_MEMORY_ID mem_id,
				    unsigned int pagecnt)
{
	phys_addr_t block_start, block_size, buf_size = pagecnt << PAGE_SHIFT;
	struct mtk_cma_memory_info *cma_mem = NULL;
	phys_addr_t cm_mem_end;
	int ret = 0;
	size_t i;

	cma_mem = mtk_secure_heap_cma_get_mem_info(mem_id);
	if (!cma_mem)
		return -EINVAL;

	if (buf_size + atomic64_read(&cma_mem->allocated_size) >
	    cma_mem->mem_size) {
		pr_err("%s: failed! use 0x%llx mem_size %pa needed 0x%x\n",
		       __func__, atomic64_read(&cma_mem->allocated_size),
		       &cma_mem->mem_size, buf_size);
		return -ENOMEM;
	}

	if (!cma_mem->cma)
		return 0;

	cm_mem_end = cma_mem->mem_start + cma_mem->mem_size;
	mutex_lock(&cma_mem->lock);

	for (i = 0; i < cma_mem->cma_block_count; i++) {
		if (atomic64_read(&cma_mem->allocated_size) +
		    buf_size <= cma_mem->cma_holded_size)
			break;

		if (cma_mem->cma_holded_block_mask & BIT(i))
			continue;

		block_start = cma_mem->mem_start + cma_mem->cma_block_size * i;
		block_size = min(cma_mem->cma_block_size, cm_mem_end - block_start);
		ret = alloc_contig_range(PFN_DOWN(block_start),
					 PFN_DOWN(block_start + block_size),
					 MIGRATE_CMA, GFP_KERNEL);
		if (ret) {
			pr_err("%s: failed to alloc block %u, ret %d\n",
			       __func__, i, ret);
			continue;
		}

		ret = gen_pool_add(cma_mem->mem_pool, block_start, block_size, -1);
		if (ret) {
			pr_err("%s: failed to add block %u to pool, ret %d\n",
			       __func__, i, ret);
			free_contig_range(PFN_DOWN(block_start),
					  block_size >> PAGE_SHIFT);
			continue;
		}

		cma_mem->cma_holded_size += block_size;
		cma_mem->cma_holded_block_mask |= BIT(i);
	}

	if (buf_size + atomic64_read(&cma_mem->allocated_size) >
	    cma_mem->cma_holded_size) {
		pr_err("%s: failed! use 0x%llx mask 0x%x hold %pa need 0x%x\n",
		       __func__, atomic64_read(&cma_mem->allocated_size),
		       cma_mem->cma_holded_block_mask,
		       &cma_mem->cma_holded_size,  buf_size);
		ret = -ENOMEM;
	}

	mutex_unlock(&cma_mem->lock);

	DBG_LOG("%s: use 0x%llx/%pa hold %pa msk 0x%x need 0x%x\n",
		__func__, atomic64_read(&cma_mem->allocated_size),
		&cma_mem->mem_size, &cma_mem->cma_holded_size,
		cma_mem->cma_holded_block_mask, buf_size);

	return ret;
}

/* refer to gen_pool_destory() in genalloc.c */
static void mtk_secure_heap_cma_remove_chunks(struct gen_pool *pool)
{
	struct list_head *_chunk, *_next_chunk;
	struct gen_pool_chunk *chunk;
	int order = pool->min_alloc_order;
	unsigned long bit, end_bit;
	size_t chunk_size;

	list_for_each_safe(_chunk, _next_chunk, &pool->chunks) {
		chunk = list_entry(_chunk, struct gen_pool_chunk, next_chunk);
		list_del(&chunk->next_chunk);

		chunk_size = chunk->end_addr - chunk->start_addr + 1;
		end_bit = chunk_size >> order;
		bit = find_first_bit(chunk->bits, end_bit);
		WARN_ON_ONCE(bit < end_bit);

		vfree(chunk);
	}
}

void mtk_secure_heap_cma_free_blocks_all(enum MTK_CMA_MEMORY_ID mem_id)
{
	struct mtk_cma_memory_info *cma_mem = NULL;
	phys_addr_t block_start, block_size;
	phys_addr_t cm_mem_end;
	size_t i;

	cma_mem = mtk_secure_heap_cma_get_mem_info(mem_id);
	if (!cma_mem)
		return;

	DBG_LOG("%s: heap usage 0x%llx/%pa, pool avail 0x%lx/0x%lx by %s(%d)\n",
		__func__, atomic64_read(&cma_mem->allocated_size),
		&cma_mem->mem_size, gen_pool_avail(cma_mem->mem_pool),
		gen_pool_size(cma_mem->mem_pool), cma_mem->name, mem_id);

	WARN_ON_ONCE(atomic64_read(&cma_mem->allocated_size) ||
		     (gen_pool_size(cma_mem->mem_pool) -
		     gen_pool_avail(cma_mem->mem_pool)));

	if (!cma_mem->cma)
		return;

	cm_mem_end = cma_mem->mem_start + cma_mem->mem_size;
	mutex_lock(&cma_mem->lock);
	for (i = 0; i < cma_mem->cma_block_count; i++) {
		if (!(cma_mem->cma_holded_block_mask & BIT(i)))
			continue;

		block_start = cma_mem->mem_start + cma_mem->cma_block_size * i;
		block_size = min(cma_mem->cma_block_size, cm_mem_end - block_start);
		free_contig_range(PFN_DOWN(block_start), block_size >> PAGE_SHIFT);
	}
	cma_mem->cma_holded_size = 0;
	cma_mem->cma_holded_block_mask = 0;

	mtk_secure_heap_cma_remove_chunks(cma_mem->mem_pool);

	mutex_unlock(&cma_mem->lock);
}

struct page *mtk_secure_heap_cma_alloc_pages(enum MTK_CMA_MEMORY_ID mem_id,
					     unsigned int pagecount,
					     unsigned int align_size)
{
	struct mtk_cma_memory_info *cma_mem = NULL;
	struct genpool_data_align data = {
		.align = align_size,
	};
	struct page *pages = NULL;
	phys_addr_t paddr = 0;

	cma_mem = mtk_secure_heap_cma_get_mem_info(mem_id);
	if (!cma_mem)
		return NULL;

	if (!align_size)
		paddr = gen_pool_alloc(cma_mem->mem_pool, pagecount << PAGE_SHIFT);
	else
		paddr = gen_pool_alloc_algo(cma_mem->mem_pool,
					    pagecount << PAGE_SHIFT,
					    gen_pool_first_fit_align, &data);
	if (!paddr) {
		DBG_LOG("%s: failed, avai/size 0x%lx/0x%lx, need 0x%llx\n",
			__func__, gen_pool_avail(cma_mem->mem_pool),
			gen_pool_size(cma_mem->mem_pool),
			pagecount << PAGE_SHIFT);
		return NULL;
	}

	atomic64_add(((s64)pagecount) << PAGE_SHIFT, &cma_mem->allocated_size);

	DBG_LOG("%s alloc 0x%x from %s(%d) done, allocated 0x%llx\n",
		__func__, pagecount << PAGE_SHIFT, cma_mem->name, mem_id,
		atomic64_read(&cma_mem->allocated_size));

	return phys_to_page(paddr);
}

void mtk_secure_heap_cma_release_pages(enum MTK_CMA_MEMORY_ID mem_id,
				       struct page *pages,
				       unsigned int pagecount)
{
	struct mtk_cma_memory_info *cma_mem = NULL;

	cma_mem = mtk_secure_heap_cma_get_mem_info(mem_id);
	if (!cma_mem)
		return;

	gen_pool_free(cma_mem->mem_pool, page_to_phys(pages),
		      pagecount << PAGE_SHIFT);

	if (atomic64_sub_return(((s64)pagecount) << PAGE_SHIFT,
	    &cma_mem->allocated_size) < 0) {
		pr_info("%s: WARNING: allocated_size of %s(%d) underflow 0x%llx\n",
			__func__, cma_mem->name, mem_id,
			atomic64_read(&cma_mem->allocated_size));
		atomic64_set(&cma_mem->allocated_size, 0);
	}

	DBG_LOG("%s release 0x%x to %s(%d) done, allocated 0x%llx\n",
		__func__, pagecount << PAGE_SHIFT, cma_mem->name, mem_id,
		atomic64_read(&cma_mem->allocated_size));
}

static int __init
_mtk_secure_heap_cma_memory_init(struct reserved_mem *rmem,
				 enum MTK_CMA_MEMORY_ID mem_id)
{
	struct mtk_cma_memory_info *cma_mem = NULL;
	unsigned long node = rmem->fdt_node;
#ifdef CMA_MEM_ALIGN_FORCE_ENABLE
	unsigned int cma_block_align = 0;
#else
	uint32_t remainder;
	uint64_t dividend;
#endif
	int ret = 0;

	cma_mem = mtk_secure_heap_cma_get_mem_info(mem_id);
	if (!cma_mem) {
		pr_err("%s Invalid mem_id %d\n", __func__, mem_id);
		return -EINVAL;
	}

	/* property shall only be either "reusable" or "no-map" */
	if ((!of_get_flat_dt_prop(node, "reusable", NULL) &&
	     !of_get_flat_dt_prop(node, "no-map", NULL)) ||
	    (of_get_flat_dt_prop(node, "reusable", NULL) &&
	     of_get_flat_dt_prop(node, "no-map", NULL))) {
		pr_err("%s: %s incorrect property\n", __func__, rmem->name);
		return -EINVAL;
	}

#ifdef CMA_MEM_ALIGN_FORCE_ENABLE
	cma_block_align = cma_mem->cma_block_align;
	if (!rmem->base || !IS_ALIGNED(rmem->base, cma_block_align) ||
	    !rmem->size || !IS_ALIGNED(rmem->size, cma_block_align)) {
		pr_err("%s: %s rmem %pa(%pa) not 0x%x align\n",
		       &rmem->base, &rmem->size, cma_block_align);
		return -EINVAL;
	}
#endif

	if (of_get_flat_dt_prop(node, "reusable", NULL)) {
		ret = cma_init_reserved_mem(rmem->base, rmem->size, 0,
					    cma_mem->name, &cma_mem->cma);
		if (ret) {
			pr_err("%s: %s init CMA fail, ret %d\n",
				__func__, rmem->name, ret);
			cma_mem->cma = NULL;
			return ret;
		}
	}

	mutex_init(&cma_mem->lock);
	cma_mem->rmem = rmem;
	cma_mem->mem_pool = NULL;
	cma_mem->mem_start = rmem->base;
	cma_mem->mem_size = rmem->size;
	atomic64_set(&cma_mem->allocated_size, 0);

	if (!cma_mem->cma) {
		cma_mem->cma_block_size = cma_mem->mem_size;
		cma_mem->cma_block_count = 1;
	} else {
#ifdef CMA_MEM_ALIGN_FORCE_ENABLE
		cma_mem->cma_block_size = cma_mem->mem_size / CMA_BLOCKS_MAX_NUM;
		cma_mem->cma_block_count = CMA_BLOCKS_MAX_NUM;
#else
		dividend = cma_mem->mem_size + CMA_BLOCKS_MAX_NUM - 1;
		remainder = do_div(dividend, CMA_BLOCKS_MAX_NUM);
		cma_mem->cma_block_size = ALIGN(dividend, MPU_REGION_ALIGN);

		dividend = cma_mem->mem_size;
		remainder = do_div(dividend, cma_mem->cma_block_size);
		cma_mem->cma_block_count = dividend + (remainder ? 1 : 0);
#endif
	}

	DBG_LOG("%s get %s mem base %pa size %pa block num 0x%x size 0x%x for %s\n",
		__func__, cma_mem->cma ? "CMA" : "RSV", &cma_mem->mem_start,
		&cma_mem->mem_size, cma_mem->cma_block_count,
		cma_mem->cma_block_size, cma_mem->name);

	return 0;
}

static int __init mtk_secure_heap_cma_svp_mem_init(struct reserved_mem *rmem)
{
	return _mtk_secure_heap_cma_memory_init(rmem, MTK_CMA_MEM_SVP);
}

RESERVEDMEM_OF_DECLARE(mtk_svp_memory, "mediatek,svp_memory",
		       mtk_secure_heap_cma_svp_mem_init);
