// SPDX-License-Identifier: GPL-2.0-or-later

#include <linux/blkdev.h>
#include <linux/dax.h>
#include <linux/vmalloc.h>
#include <linux/parser.h>

#include "cache_dev.h"
#include "backing_dev.h"
#include "cache.h"
#include "dm_pcache.h"

static void cache_dev_dax_exit(struct pcache_cache_dev *cache_dev)
{
	if (cache_dev->use_vmap)
		vunmap(cache_dev->mapping);
}

static int build_vmap(struct dax_device *dax_dev, long total_pages, void **vaddr)
{
	struct page **pages;
	long i = 0, chunk;
	unsigned long pfn;
	int ret;

	pages = vmalloc_array(total_pages, sizeof(struct page *));
	if (!pages)
		return -ENOMEM;

	do {
		chunk = dax_direct_access(dax_dev, i, total_pages - i,
					  DAX_ACCESS, NULL, &pfn);
		if (chunk <= 0) {
			ret = chunk ? chunk : -EINVAL;
			goto out_free;
		}

		if (!pfn_valid(pfn)) {
			ret = -EOPNOTSUPP;
			goto out_free;
		}

		while (chunk-- && i < total_pages) {
			pages[i++] = pfn_to_page(pfn);
			pfn++;
			if (!(i & 15))
				cond_resched();
		}
	} while (i < total_pages);

	*vaddr = vmap(pages, total_pages, VM_MAP, PAGE_KERNEL);
	if (!*vaddr) {
		ret = -ENOMEM;
		goto out_free;
	}

	ret = 0;

out_free:
	vfree(pages);
	return ret;
}

static int cache_dev_dax_init(struct pcache_cache_dev *cache_dev)
{
	struct dm_pcache	*pcache = CACHE_DEV_TO_PCACHE(cache_dev);
	struct dax_device	*dax_dev;
	long			total_pages, mapped_pages;
	u64			bdev_size;
	void			*vaddr;
	int			ret;
	int			id;
	unsigned long		pfn;

	dax_dev	= cache_dev->dm_dev->dax_dev;
	/* total size check */
	bdev_size = bdev_nr_bytes(cache_dev->dm_dev->bdev);
	if (bdev_size < PCACHE_CACHE_DEV_SIZE_MIN) {
		pcache_dev_err(pcache, "dax device is too small, required at least %llu",
				PCACHE_CACHE_DEV_SIZE_MIN);
		ret = -ENOSPC;
		goto out;
	}

	total_pages = bdev_size >> PAGE_SHIFT;
	/* attempt: direct-map the whole range */
	id = dax_read_lock();
	mapped_pages = dax_direct_access(dax_dev, 0, total_pages,
					 DAX_ACCESS, &vaddr, &pfn);
	if (mapped_pages < 0) {
		pcache_dev_err(pcache, "dax_direct_access failed: %ld\n", mapped_pages);
		ret = mapped_pages;
		goto unlock;
	}

	if (!pfn_valid(pfn)) {
		ret = -EOPNOTSUPP;
		goto unlock;
	}

	if (mapped_pages == total_pages) {
		/* success: contiguous direct mapping */
		cache_dev->mapping = vaddr;
	} else {
		/* need vmap fallback */
		ret = build_vmap(dax_dev, total_pages, &vaddr);
		if (ret) {
			pcache_dev_err(pcache, "vmap fallback failed: %d\n", ret);
			goto unlock;
		}

		cache_dev->mapping	= vaddr;
		cache_dev->use_vmap	= true;
	}
	dax_read_unlock(id);

	return 0;
unlock:
	dax_read_unlock(id);
out:
	return ret;
}

void cache_dev_zero_range(struct pcache_cache_dev *cache_dev, void *pos, u32 size)
{
	memset(pos, 0, size);
	dax_flush(cache_dev->dm_dev->dax_dev, pos, size);
}

static int sb_read(struct pcache_cache_dev *cache_dev, struct pcache_sb *sb)
{
	struct pcache_sb *sb_addr = CACHE_DEV_SB(cache_dev);

	if (copy_mc_to_kernel(sb, sb_addr, sizeof(struct pcache_sb)))
		return -EIO;

	return 0;
}

static void sb_write(struct pcache_cache_dev *cache_dev, struct pcache_sb *sb)
{
	struct pcache_sb *sb_addr = CACHE_DEV_SB(cache_dev);

	memcpy_flushcache(sb_addr, sb, sizeof(struct pcache_sb));
	pmem_wmb();
}

static int sb_init(struct pcache_cache_dev *cache_dev, struct pcache_sb *sb)
{
	struct dm_pcache *pcache = CACHE_DEV_TO_PCACHE(cache_dev);
	u64 nr_segs;
	u64 cache_dev_size;
	u64 magic;
	u32 flags = 0;

	magic = le64_to_cpu(sb->magic);
	if (magic)
		return -EEXIST;

	cache_dev_size = bdev_nr_bytes(file_bdev(cache_dev->dm_dev->bdev_file));
	if (cache_dev_size < PCACHE_CACHE_DEV_SIZE_MIN) {
		pcache_dev_err(pcache, "dax device is too small, required at least %llu",
				PCACHE_CACHE_DEV_SIZE_MIN);
		return -ENOSPC;
	}

	nr_segs = (cache_dev_size - PCACHE_SEGMENTS_OFF) / ((PCACHE_SEG_SIZE));

#if defined(__BYTE_ORDER) ? (__BIG_ENDIAN == __BYTE_ORDER) : defined(__BIG_ENDIAN)
	flags |= PCACHE_SB_F_BIGENDIAN;
#endif
	sb->flags = cpu_to_le32(flags);
	sb->magic = cpu_to_le64(PCACHE_MAGIC);
	sb->seg_num = cpu_to_le32(nr_segs);
	sb->crc = cpu_to_le32(crc32c(PCACHE_CRC_SEED, (void *)(sb) + 4, sizeof(struct pcache_sb) - 4));

	cache_dev_zero_range(cache_dev, CACHE_DEV_CACHE_INFO(cache_dev),
			     PCACHE_CACHE_INFO_SIZE * PCACHE_META_INDEX_MAX +
			     PCACHE_CACHE_CTRL_SIZE);

	return 0;
}

static int sb_validate(struct pcache_cache_dev *cache_dev, struct pcache_sb *sb)
{
	struct dm_pcache *pcache = CACHE_DEV_TO_PCACHE(cache_dev);
	u32 flags;
	u32 crc;

	if (le64_to_cpu(sb->magic) != PCACHE_MAGIC) {
		pcache_dev_err(pcache, "unexpected magic: %llx\n",
				le64_to_cpu(sb->magic));
		return -EINVAL;
	}

	crc = crc32c(PCACHE_CRC_SEED, (void *)(sb) + 4, sizeof(struct pcache_sb) - 4);
	if (crc != le32_to_cpu(sb->crc)) {
		pcache_dev_err(pcache, "corrupted sb: %u, expected: %u\n", crc, le32_to_cpu(sb->crc));
		return -EINVAL;
	}

	flags = le32_to_cpu(sb->flags);
#if defined(__BYTE_ORDER) ? (__BIG_ENDIAN == __BYTE_ORDER) : defined(__BIG_ENDIAN)
	if (!(flags & PCACHE_SB_F_BIGENDIAN)) {
		pcache_dev_err(pcache, "cache_dev is not big endian\n");
		return -EINVAL;
	}
#else
	if (flags & PCACHE_SB_F_BIGENDIAN) {
		pcache_dev_err(pcache, "cache_dev is big endian\n");
		return -EINVAL;
	}
#endif
	return 0;
}

static int cache_dev_init(struct pcache_cache_dev *cache_dev, u32 seg_num)
{
	cache_dev->seg_num = seg_num;
	cache_dev->seg_bitmap = kvcalloc(BITS_TO_LONGS(cache_dev->seg_num), sizeof(unsigned long), GFP_KERNEL);
	if (!cache_dev->seg_bitmap)
		return -ENOMEM;

	return 0;
}

static void cache_dev_exit(struct pcache_cache_dev *cache_dev)
{
	kvfree(cache_dev->seg_bitmap);
}

void cache_dev_stop(struct dm_pcache *pcache)
{
	struct pcache_cache_dev *cache_dev = &pcache->cache_dev;

	cache_dev_exit(cache_dev);
	cache_dev_dax_exit(cache_dev);
}

int cache_dev_start(struct dm_pcache *pcache)
{
	struct pcache_cache_dev *cache_dev = &pcache->cache_dev;
	struct pcache_sb sb;
	bool format = false;
	int ret;

	mutex_init(&cache_dev->seg_lock);

	ret = cache_dev_dax_init(cache_dev);
	if (ret) {
		pcache_dev_err(pcache, "failed to init cache_dev %s via dax way: %d.",
			       cache_dev->dm_dev->name, ret);
		goto err;
	}

	ret = sb_read(cache_dev, &sb);
	if (ret)
		goto dax_release;

	if (le64_to_cpu(sb.magic) == 0) {
		format = true;
		ret = sb_init(cache_dev, &sb);
		if (ret < 0)
			goto dax_release;
	}

	ret = sb_validate(cache_dev, &sb);
	if (ret)
		goto dax_release;

	cache_dev->sb_flags = le32_to_cpu(sb.flags);
	ret = cache_dev_init(cache_dev, le32_to_cpu(sb.seg_num));
	if (ret)
		goto dax_release;

	if (format)
		sb_write(cache_dev, &sb);

	return 0;

dax_release:
	cache_dev_dax_exit(cache_dev);
err:
	return ret;
}

int cache_dev_get_empty_segment_id(struct pcache_cache_dev *cache_dev, u32 *seg_id)
{
	int ret;

	mutex_lock(&cache_dev->seg_lock);
	*seg_id = find_next_zero_bit(cache_dev->seg_bitmap, cache_dev->seg_num, 0);
	if (*seg_id == cache_dev->seg_num) {
		ret = -ENOSPC;
		goto unlock;
	}

	__set_bit(*seg_id, cache_dev->seg_bitmap);
	ret = 0;
unlock:
	mutex_unlock(&cache_dev->seg_lock);
	return ret;
}