1// SPDX-License-Identifier: GPL-2.0-only
  2
  3#include <linux/pagemap.h>
  4#include <linux/xarray.h>
  5#include <linux/slab.h>
  6#include <linux/swap.h>
  7#include <linux/swapops.h>
  8#include <asm/mte.h>
  9
 10static DEFINE_XARRAY(mte_pages);
 11
 12void *mte_allocate_tag_storage(void)
 13{
 14	/* tags granule is 16 bytes, 2 tags stored per byte */
 15	return kmalloc(MTE_PAGE_TAG_STORAGE, GFP_KERNEL);
 16}
 17
 18void mte_free_tag_storage(char *storage)
 19{
 20	kfree(storage);
 21}
 22
 23int mte_save_tags(struct page *page)
 24{
 25	void *tag_storage, *ret;
 26
 27	if (!page_mte_tagged(page))
 28		return 0;
 29
 30	tag_storage = mte_allocate_tag_storage();
 31	if (!tag_storage)
 32		return -ENOMEM;
 33
 34	mte_save_page_tags(page_address(page), tag_storage);
 35
 36	/* lookup the swap entry.val from the page */
 37	ret = xa_store(&mte_pages, page_swap_entry(page).val, tag_storage,
 38		       GFP_KERNEL);
 39	if (WARN(xa_is_err(ret), "Failed to store MTE tags")) {
 40		mte_free_tag_storage(tag_storage);
 41		return xa_err(ret);
 42	} else if (ret) {
 43		/* Entry is being replaced, free the old entry */
 44		mte_free_tag_storage(ret);
 45	}
 46
 47	return 0;
 48}
 49
 50void mte_restore_tags(swp_entry_t entry, struct page *page)
 51{
 52	void *tags = xa_load(&mte_pages, entry.val);
 53
 54	if (!tags)
 55		return;
 56
 57	if (try_page_mte_tagging(page)) {
 58		mte_restore_page_tags(page_address(page), tags);
 59		set_page_mte_tagged(page);
 60	}
 61}
 62
 63void mte_invalidate_tags(int type, pgoff_t offset)
 64{
 65	swp_entry_t entry = swp_entry(type, offset);
 66	void *tags = xa_erase(&mte_pages, entry.val);
 67
 68	mte_free_tag_storage(tags);
 69}
 70
 71static inline void __mte_invalidate_tags(struct page *page)
 72{
 73	swp_entry_t entry = page_swap_entry(page);
 74
 75	mte_invalidate_tags(swp_type(entry), swp_offset(entry));
 76}
 77
 78void mte_invalidate_tags_area(int type)
 79{
 80	swp_entry_t entry = swp_entry(type, 0);
 81	swp_entry_t last_entry = swp_entry(type + 1, 0);
 82	void *tags;
 83
 84	XA_STATE(xa_state, &mte_pages, entry.val);
 85
 86	xa_lock(&mte_pages);
 87	xas_for_each(&xa_state, tags, last_entry.val - 1) {
 88		__xa_erase(&mte_pages, xa_state.xa_index);
 89		mte_free_tag_storage(tags);
 90	}
 91	xa_unlock(&mte_pages);
 92}
 93
 94int arch_prepare_to_swap(struct folio *folio)
 95{
 96	long i, nr;
 97	int err;
 98
 99	if (!system_supports_mte())
100		return 0;
101
102	nr = folio_nr_pages(folio);
103
104	for (i = 0; i < nr; i++) {
105		err = mte_save_tags(folio_page(folio, i));
106		if (err)
107			goto out;
108	}
109	return 0;
110
111out:
112	while (i--)
113		__mte_invalidate_tags(folio_page(folio, i));
114	return err;
115}
116
117void arch_swap_restore(swp_entry_t entry, struct folio *folio)
118{
119	long i, nr;
120
121	if (!system_supports_mte())
122		return;
123
124	nr = folio_nr_pages(folio);
125
126	for (i = 0; i < nr; i++) {
127		mte_restore_tags(entry, folio_page(folio, i));
128		entry.val++;
129	}
130}