162306a36Sopenharmony_ci// SPDX-License-Identifier: GPL-2.0-only
262306a36Sopenharmony_ci
362306a36Sopenharmony_ci#include <linux/pagemap.h>
462306a36Sopenharmony_ci#include <linux/xarray.h>
562306a36Sopenharmony_ci#include <linux/slab.h>
662306a36Sopenharmony_ci#include <linux/swap.h>
762306a36Sopenharmony_ci#include <linux/swapops.h>
862306a36Sopenharmony_ci#include <asm/mte.h>
962306a36Sopenharmony_ci
1062306a36Sopenharmony_cistatic DEFINE_XARRAY(mte_pages);
1162306a36Sopenharmony_ci
1262306a36Sopenharmony_civoid *mte_allocate_tag_storage(void)
1362306a36Sopenharmony_ci{
1462306a36Sopenharmony_ci	/* tags granule is 16 bytes, 2 tags stored per byte */
1562306a36Sopenharmony_ci	return kmalloc(MTE_PAGE_TAG_STORAGE, GFP_KERNEL);
1662306a36Sopenharmony_ci}
1762306a36Sopenharmony_ci
1862306a36Sopenharmony_civoid mte_free_tag_storage(char *storage)
1962306a36Sopenharmony_ci{
2062306a36Sopenharmony_ci	kfree(storage);
2162306a36Sopenharmony_ci}
2262306a36Sopenharmony_ci
2362306a36Sopenharmony_ciint mte_save_tags(struct page *page)
2462306a36Sopenharmony_ci{
2562306a36Sopenharmony_ci	void *tag_storage, *ret;
2662306a36Sopenharmony_ci
2762306a36Sopenharmony_ci	if (!page_mte_tagged(page))
2862306a36Sopenharmony_ci		return 0;
2962306a36Sopenharmony_ci
3062306a36Sopenharmony_ci	tag_storage = mte_allocate_tag_storage();
3162306a36Sopenharmony_ci	if (!tag_storage)
3262306a36Sopenharmony_ci		return -ENOMEM;
3362306a36Sopenharmony_ci
3462306a36Sopenharmony_ci	mte_save_page_tags(page_address(page), tag_storage);
3562306a36Sopenharmony_ci
3662306a36Sopenharmony_ci	/* lookup the swap entry.val from the page */
3762306a36Sopenharmony_ci	ret = xa_store(&mte_pages, page_swap_entry(page).val, tag_storage,
3862306a36Sopenharmony_ci		       GFP_KERNEL);
3962306a36Sopenharmony_ci	if (WARN(xa_is_err(ret), "Failed to store MTE tags")) {
4062306a36Sopenharmony_ci		mte_free_tag_storage(tag_storage);
4162306a36Sopenharmony_ci		return xa_err(ret);
4262306a36Sopenharmony_ci	} else if (ret) {
4362306a36Sopenharmony_ci		/* Entry is being replaced, free the old entry */
4462306a36Sopenharmony_ci		mte_free_tag_storage(ret);
4562306a36Sopenharmony_ci	}
4662306a36Sopenharmony_ci
4762306a36Sopenharmony_ci	return 0;
4862306a36Sopenharmony_ci}
4962306a36Sopenharmony_ci
5062306a36Sopenharmony_civoid mte_restore_tags(swp_entry_t entry, struct page *page)
5162306a36Sopenharmony_ci{
5262306a36Sopenharmony_ci	void *tags = xa_load(&mte_pages, entry.val);
5362306a36Sopenharmony_ci
5462306a36Sopenharmony_ci	if (!tags)
5562306a36Sopenharmony_ci		return;
5662306a36Sopenharmony_ci
5762306a36Sopenharmony_ci	if (try_page_mte_tagging(page)) {
5862306a36Sopenharmony_ci		mte_restore_page_tags(page_address(page), tags);
5962306a36Sopenharmony_ci		set_page_mte_tagged(page);
6062306a36Sopenharmony_ci	}
6162306a36Sopenharmony_ci}
6262306a36Sopenharmony_ci
6362306a36Sopenharmony_civoid mte_invalidate_tags(int type, pgoff_t offset)
6462306a36Sopenharmony_ci{
6562306a36Sopenharmony_ci	swp_entry_t entry = swp_entry(type, offset);
6662306a36Sopenharmony_ci	void *tags = xa_erase(&mte_pages, entry.val);
6762306a36Sopenharmony_ci
6862306a36Sopenharmony_ci	mte_free_tag_storage(tags);
6962306a36Sopenharmony_ci}
7062306a36Sopenharmony_ci
7162306a36Sopenharmony_civoid mte_invalidate_tags_area(int type)
7262306a36Sopenharmony_ci{
7362306a36Sopenharmony_ci	swp_entry_t entry = swp_entry(type, 0);
7462306a36Sopenharmony_ci	swp_entry_t last_entry = swp_entry(type + 1, 0);
7562306a36Sopenharmony_ci	void *tags;
7662306a36Sopenharmony_ci
7762306a36Sopenharmony_ci	XA_STATE(xa_state, &mte_pages, entry.val);
7862306a36Sopenharmony_ci
7962306a36Sopenharmony_ci	xa_lock(&mte_pages);
8062306a36Sopenharmony_ci	xas_for_each(&xa_state, tags, last_entry.val - 1) {
8162306a36Sopenharmony_ci		__xa_erase(&mte_pages, xa_state.xa_index);
8262306a36Sopenharmony_ci		mte_free_tag_storage(tags);
8362306a36Sopenharmony_ci	}
8462306a36Sopenharmony_ci	xa_unlock(&mte_pages);
8562306a36Sopenharmony_ci}
86