[RFC PATCH 0/7] mm: store zswap entries in swap table (Pointer entry)

[RFC PATCH 0/7] mm: store zswap entries in swap table (Pointer entry)

[Baoquan He]

Currently zswap uses a per-swap-device xarray to map swap offsets to
zswap_entry pointers. This introduces an additional tree walk on every
zswap entry access.

This series replaces the xarray with a "Pointer" entry type stored
directly in the swap cluster table. The new swp_tb_t variant embeds
a zswap_entry pointer plus a entry count, eliminating the xarray and
its lookup overhead entirely.

Design overview
---------------
A swap table slot can now be one of five types:

  NULL:    |---------------- 0 ---------------|
  Shadow:  |SWAP_COUNT|Z|---- SHADOW_VAL ---|1|
  PFN:     |SWAP_COUNT|Z|------ PFN -------|10|
  Pointer: |----------- Pointer ----------|100|
  Bad:     |------------- 1 -------------|1000|

When zswap compresses a page, it saves the original swap table entry
(PFN or Shadow) in zswap_entry->swp_tb_val, then writes the Pointer
entry into the swap table.  The swap count and flags remain accessible
via helper functions that read them from swp_tb_val.

The zswap_entry is threaded through the swapin call chain via a new
void **zentry parameter, so zswap_load() receives it directly instead
of doing an xarray lookup.

Performance
-----------
[~]# free -h
               total        used        free      shared  buff/cache   available
Mem:           3.8Gi       155Mi       3.4Gi       656Ki       315Mi       3.7Gi
Swap:          4.0Gi        54Mi       3.9Gi

Scalability benchmark: 'usemem -O -y -x -n 1 5G', 10 runs under cgroup
MemoryMax=2G MemorySwapMax=3G.  zswap enabled with zstd compressor.

  Kernel A: cfb8731f5396 (xarray-based lookup, baseline)
  Kernel B: (Pointer entry)

5G test — the gap widens significantly as swap pressure grows:

     Metric           xarray(A)       Pointer(B)
    --------------------------------------------------
    Avg throughput      709,556 KB/s    792,681 KB/s  (+11.7%)
    Avg free time       167,994 µs     132,896 µs   (-20.9%)

  Free time trend across 10 rounds at 5G:

    Pointer: 126→124→126→126→124→129→138→145→142→148  (gradual, +17%)
    xarray:  136→139→136→160→162→161→159→230→219→178  (collapse, +31%)

  xarray degrades sharply in the later rounds (last 3 avg free time
  209 ms vs Pointer's 145 ms, 1.44× worse) as the radix tree
  insertion/deletion cost grows non-linearly under sustained swap
  churn.  Pointer entry avoids this entirely by looking up entries
  directly in the swap cluster table, which is already hot in cache.

Baoquan He (7):
  mm/swap_table: add Pointer type support for zswap entry storage
  mm/zswap: add swp_tb_val to zswap_entry and Pointer accessor helpers
  mm/swap_state: do zswap conversion in __swap_cache_do_del_folio
  mm/zswap, mm/swap_state: handle Pointer entries in add_folio path
  mm/swapfile: handle Pointer entries in count/dup/put paths
  mm/zswap, mm/swap_state: migrate zswap to use swap_table for entry
    indexing
  mm/zswap, mm/swap_state: replace xarray indexing with swap table
    lookups

 include/linux/zswap.h |   4 +-
 mm/page_io.c          |   4 +-
 mm/swap.h             |  52 +++++-
 mm/swap_state.c       | 117 +++++++++++---
 mm/swap_table.h       |  35 ++++-
 mm/swapfile.c         | 107 ++++++++++---
 mm/zswap.c            | 358 +++++++++++++++++++++++++++++-------------
 7 files changed, 508 insertions(+), 169 deletions(-)

-- 
2.54.0

[RFC PATCH 1/7] mm/swap_table: add Pointer type support for zswap entry storage

[Baoquan He]

Add the Pointer type infrastructure to swap table entries, enabling
direct storage of zswap_entry pointers in swap_table entries. The
0b100 marker in the low 3 bits distinguishes pointer entries from
NULL (0), Shadow (0b1), and PFN (0b10) entries.

Add helpers:
- SWP_TB_POINTER_MARK: the 0b100 marker constant
- swp_tb_is_pointer(): type check for Pointer entries
- pointer_to_swp_tb(): encode a pointer into a swap table entry
- swp_tb_to_pointer(): decode a swap table entry back to a pointer

Signed-off-by: Baoquan He <baoquan.he@linux.dev>
---
 mm/swap_table.h | 35 +++++++++++++++++++++++++++++++----
 1 file changed, 31 insertions(+), 4 deletions(-)

diff --git a/mm/swap_table.h b/mm/swap_table.h
index e6613e62f8d0..6318bf8aa76b 100644
--- a/mm/swap_table.h
+++ b/mm/swap_table.h
@@ -28,7 +28,7 @@ struct swap_memcg_table {
  * NULL:     |---------------- 0 ---------------| - Free slot
  * Shadow:   |SWAP_COUNT|Z|---- SHADOW_VAL ---|1| - Swapped out slot
  * PFN:      |SWAP_COUNT|Z|------ PFN -------|10| - Cached slot
- * Pointer:  |----------- Pointer ----------|100| - (Unused)
+ * Pointer:  |----------- Pointer ----------|100| - Zswap compressed entry
  * Bad:      |------------- 1 -------------|1000| - Bad slot
  *
  * COUNT is `SWP_TB_COUNT_BITS` long, Z is the `SWP_TB_ZERO_FLAG` bit,
@@ -49,9 +49,11 @@ struct swap_memcg_table {
  * - PFN: Swap slot is in use, and cached. Memcg info is recorded on the page
  *   struct.
  *
- * - Pointer: Unused yet. `0b100` is reserved for potential pointer usage
- *   because only the lower three bits can be used as a marker for 8 bytes
- *   aligned pointers.
+ * - Pointer: Stores a direct pointer to a zswap_entry when the swap slot's
+ *   compressed data resides in zswap. The swap count (if any) is stored in
+ *   the zswap_entry struct rather than inline, since the pointer occupies
+ *   all upper bits. `0b100` is used as the marker because zswap_entry
+ *   pointers are at least 8-byte aligned, leaving the lower three bits free.
  *
  * - Bad: Swap slot is reserved, protects swap header or holes on swap devices.
  */
@@ -78,6 +80,9 @@ struct swap_memcg_table {
 /* The first flag is zero bit (SWAP_TABLE_HAS_ZEROFLAG) */
 #define SWP_TB_ZERO_FLAG	BIT(BITS_PER_LONG - SWP_TB_FLAGS_BITS)
 
+/* Pointer: zswap_entry pointer stored directly, ends with 0b100 */
+#define SWP_TB_POINTER_MARK	0b100UL
+
 /* Bad slot: ends with 0b1000 and rests of bits are all 1 */
 #define SWP_TB_BAD		((~0UL) << 3)
 
@@ -166,6 +171,11 @@ static inline bool swp_tb_is_bad(unsigned long swp_tb)
 	return swp_tb == SWP_TB_BAD;
 }
 
+static inline bool swp_tb_is_pointer(unsigned long swp_tb)
+{
+	return (swp_tb & (BIT(3) - 1)) == SWP_TB_POINTER_MARK;
+}
+
 static inline bool swp_tb_is_countable(unsigned long swp_tb)
 {
 	return (swp_tb_is_shadow(swp_tb) || swp_tb_is_folio(swp_tb) ||
@@ -188,6 +198,23 @@ static inline void *swp_tb_to_shadow(unsigned long swp_tb)
 	return (void *)(swp_tb & ~SWP_TB_FLAGS_MASK);
 }
 
+static inline unsigned long pointer_to_swp_tb(void *ptr)
+{
+	unsigned long val = (unsigned long)ptr;
+
+	BUILD_BUG_ON(sizeof(unsigned long) != sizeof(void *));
+	/* Pointers are 8-byte aligned, low 3 bits must be clear */
+	VM_WARN_ON_ONCE(val & (BIT(3) - 1));
+	return val | SWP_TB_POINTER_MARK;
+}
+
+static inline void *swp_tb_to_pointer(unsigned long swp_tb)
+{
+	VM_WARN_ON(!swp_tb_is_pointer(swp_tb));
+	/* Clear the low 3-bit marker to recover the original pointer */
+	return (void *)(swp_tb & ~((unsigned long)(BIT(3) - 1)));
+}
+
 static inline unsigned char __swp_tb_get_count(unsigned long swp_tb)
 {
 	VM_WARN_ON(!swp_tb_is_countable(swp_tb));
-- 
2.54.0

[RFC PATCH 2/7] mm/zswap: add swp_tb_val to zswap_entry and Pointer accessor helpers

[Baoquan He]

Add 'swp_tb_val' to struct zswap_entry to preserve the original swap
table entry value when a Shadow or PFN entry is replaced by a Pointer
entry.

Add accessor helpers that read and modify swap count and flags through
entry->swp_tb_val:
- zswap_swp_tb_get_count(): extract swap count from the saved entry
- zswap_swp_tb_get_flags(): extract flags from the saved entry
- zswap_swp_tb_dup_count(): increment swap count, returns -ENOMEM on overflow
- zswap_swp_tb_put_count(): decrement swap count

The saved swap table entry value uses the Shadow/PFN format, so the
existing __swp_tb_get_count/flags and __swp_tb_mk_count helpers work
directly on swp_tb_val without any format conversion.

Signed-off-by: Baoquan He <baoquan.he@linux.dev>
---
 mm/swap.h  | 27 +++++++++++++++++
 mm/zswap.c | 87 ++++++++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 114 insertions(+)

diff --git a/mm/swap.h b/mm/swap.h
index 44ab8e1e595b..afe2200bd160 100644
--- a/mm/swap.h
+++ b/mm/swap.h
@@ -210,6 +210,33 @@ static inline void swap_cluster_unlock_irq(struct swap_cluster_info *ci)
 
 extern int swap_retry_table_alloc(swp_entry_t entry, gfp_t gfp);
 
+#ifdef CONFIG_ZSWAP
+unsigned char zswap_swp_tb_get_count(unsigned long swp_tb);
+unsigned char zswap_swp_tb_get_flags(unsigned long swp_tb);
+int zswap_swp_tb_dup_count(unsigned long swp_tb,
+			   struct swap_cluster_info *ci, unsigned int ci_off);
+void zswap_swp_tb_put_count(unsigned long swp_tb,
+			    struct swap_cluster_info *ci, unsigned int ci_off);
+#else
+static inline unsigned char zswap_swp_tb_get_count(unsigned long swp_tb)
+{
+	return 0;
+}
+static inline unsigned char zswap_swp_tb_get_flags(unsigned long swp_tb)
+{
+	return 0;
+}
+static inline int zswap_swp_tb_dup_count(unsigned long swp_tb,
+					 struct swap_cluster_info *ci,
+					 unsigned int ci_off)
+{
+	return 0;
+}
+static inline void zswap_swp_tb_put_count(unsigned long swp_tb,
+					  struct swap_cluster_info *ci,
+					  unsigned int ci_off) {}
+#endif
+
 /*
  * Below are the core routines for doing swap for a folio.
  * All helpers requires the folio to be locked, and a locked folio
diff --git a/mm/zswap.c b/mm/zswap.c
index 761cd699e0a3..81115ca4d10a 100644
--- a/mm/zswap.c
+++ b/mm/zswap.c
@@ -38,6 +38,7 @@
 #include <linux/zsmalloc.h>
 
 #include "swap.h"
+#include "swap_table.h"
 #include "internal.h"
 
 /*********************************
@@ -182,6 +183,10 @@ static struct shrinker *zswap_shrinker;
  *              writeback logic. The entry is only reclaimed by the writeback
  *              logic if referenced is unset. See comments in the shrinker
  *              section for context.
+ * swp_tb_val - the original swap table entry value saved when this zswap_entry
+ *              replaced a Shadow or PFN entry. It preserves the swap count,
+ *              zero flag, and working set shadow information (SHADOW_VAL)
+ *              that the Shadow entry carried.
  * pool - the zswap_pool the entry's data is in
  * handle - zsmalloc allocation handle that stores the compressed page data
  * objcg - the obj_cgroup that the compressed memory is charged to
@@ -191,6 +196,7 @@ struct zswap_entry {
 	swp_entry_t swpentry;
 	unsigned int length;
 	bool referenced;
+	unsigned long swp_tb_val;
 	struct zswap_pool *pool;
 	unsigned long handle;
 	struct obj_cgroup *objcg;
@@ -1646,6 +1652,87 @@ void zswap_invalidate(swp_entry_t swp)
 		zswap_entry_free(entry);
 }
 
+/*
+ * Read the swap count from a Pointer-type swap table entry. The count
+ * and flags are stored in zswap_entry->swp_tb_val which preserves the
+ * original swap table entry value (Shadow or PFN format).
+ */
+unsigned char zswap_swp_tb_get_count(unsigned long swp_tb)
+{
+	struct zswap_entry *entry = swp_tb_to_pointer(swp_tb);
+
+	return __swp_tb_get_count(entry->swp_tb_val);
+}
+
+/*
+ * Read the flags from a Pointer-type swap table entry.
+ */
+unsigned char zswap_swp_tb_get_flags(unsigned long swp_tb)
+{
+	struct zswap_entry *entry = swp_tb_to_pointer(swp_tb);
+
+	return __swp_tb_get_flags(entry->swp_tb_val);
+}
+
+/*
+ * Increment the swap count of a Pointer-type swap table entry.
+ * The count is stored in zswap_entry->swp_tb_val for counts below
+ * SWP_TB_COUNT_MAX, and overflows into the cluster's extend_table
+ * just like normal Shadow/PFN entries do.  Returns 0 on success,
+ * -ENOMEM if count would overflow and no extend_table is available.
+ */
+int zswap_swp_tb_dup_count(unsigned long swp_tb,
+			   struct swap_cluster_info *ci, unsigned int ci_off)
+{
+	struct zswap_entry *entry = swp_tb_to_pointer(swp_tb);
+	unsigned char count = __swp_tb_get_count(entry->swp_tb_val);
+
+	if (count < SWP_TB_COUNT_MAX) {
+		entry->swp_tb_val = __swp_tb_mk_count(entry->swp_tb_val,
+						      count + 1);
+		return 0;
+	}
+	/* count == MAX, overflow into extend_table */
+	if (!ci->extend_table)
+		return -ENOMEM;
+	count = ci->extend_table[ci_off];
+	if (count == 0)
+		count = SWP_TB_COUNT_MAX;
+	count++;
+	ci->extend_table[ci_off] = count;
+	return 0;
+}
+
+/*
+ * Decrement the swap count of a Pointer-type swap table entry.
+ * Handles both inline count (in swp_tb_val) and overflowed count
+ * (in ci->extend_table), matching the logic used by normal entries.
+ */
+void zswap_swp_tb_put_count(unsigned long swp_tb,
+			    struct swap_cluster_info *ci, unsigned int ci_off)
+{
+	struct zswap_entry *entry = swp_tb_to_pointer(swp_tb);
+	unsigned char count = __swp_tb_get_count(entry->swp_tb_val);
+
+	VM_WARN_ON_ONCE(count == 0);
+
+	if (count == SWP_TB_COUNT_MAX) {
+		count = ci->extend_table[ci_off];
+		/* Overflow starts with SWP_TB_COUNT_MAX */
+		VM_WARN_ON_ONCE(count < SWP_TB_COUNT_MAX);
+		count--;
+		if (count == (SWP_TB_COUNT_MAX - 1)) {
+			ci->extend_table[ci_off] = 0;
+			entry->swp_tb_val = __swp_tb_mk_count(entry->swp_tb_val,
+							      count);
+		} else {
+			ci->extend_table[ci_off] = count;
+		}
+		return;
+	}
+	entry->swp_tb_val = __swp_tb_mk_count(entry->swp_tb_val, count - 1);
+}
+
 int zswap_swapon(int type, unsigned long nr_pages)
 {
 	struct xarray *trees, *tree;
-- 
2.54.0

[RFC PATCH 3/7] mm/swap_state: do zswap conversion in __swap_cache_do_del_folio

[Baoquan He]

Add zswap_try_convert_to_pointer(), called from __swap_cache_do_del_folio
under cluster lock. When a folio is evicted from swap cache and a Shadow
entry is written, this helper checks if zswap has a compressed copy of
the slot and replaces the Shadow with a Pointer entry.

Signed-off-by: Baoquan He <baoquan.he@linux.dev>
---
 mm/swap.h       |  5 +++++
 mm/swap_state.c |  7 +++++++
 mm/zswap.c      | 35 +++++++++++++++++++++++++++++++++++
 3 files changed, 47 insertions(+)

diff --git a/mm/swap.h b/mm/swap.h
index afe2200bd160..1b07a08e5277 100644
--- a/mm/swap.h
+++ b/mm/swap.h
@@ -211,6 +211,8 @@ static inline void swap_cluster_unlock_irq(struct swap_cluster_info *ci)
 extern int swap_retry_table_alloc(swp_entry_t entry, gfp_t gfp);
 
 #ifdef CONFIG_ZSWAP
+void zswap_try_convert_to_pointer(struct swap_cluster_info *ci,
+				  unsigned int ci_off, int type, pgoff_t offset);
 unsigned char zswap_swp_tb_get_count(unsigned long swp_tb);
 unsigned char zswap_swp_tb_get_flags(unsigned long swp_tb);
 int zswap_swp_tb_dup_count(unsigned long swp_tb,
@@ -218,6 +220,9 @@ int zswap_swp_tb_dup_count(unsigned long swp_tb,
 void zswap_swp_tb_put_count(unsigned long swp_tb,
 			    struct swap_cluster_info *ci, unsigned int ci_off);
 #else
+static inline void zswap_try_convert_to_pointer(struct swap_cluster_info *ci,
+						unsigned int ci_off, int type,
+						pgoff_t offset) {}
 static inline unsigned char zswap_swp_tb_get_count(unsigned long swp_tb)
 {
 	return 0;
diff --git a/mm/swap_state.c b/mm/swap_state.c
index 6fd6e3415b71..20fb57adbb7f 100644
--- a/mm/swap_state.c
+++ b/mm/swap_state.c
@@ -278,6 +278,13 @@ static void __swap_cache_do_del_folio(struct swap_cluster_info *ci,
 		/* If shadow is NULL, we set an empty shadow. */
 		__swap_table_set(ci, ci_off, shadow_to_swp_tb(shadow,
 				 __swp_tb_get_flags(old_tb)));
+		/*
+		 * If zswap has a compressed copy of this slot, convert the
+		 * just-written Shadow to a Pointer entry referencing the
+		 * zswap_entry, so the data can be found via the swap table.
+		 */
+		zswap_try_convert_to_pointer(ci, ci_off, si->type,
+					     swp_offset(entry) + ci_off - ci_start);
 	} while (++ci_off < ci_end);
 
 	folio->swap.val = 0;
diff --git a/mm/zswap.c b/mm/zswap.c
index 81115ca4d10a..f6e8259ab56a 100644
--- a/mm/zswap.c
+++ b/mm/zswap.c
@@ -1652,6 +1652,41 @@ void zswap_invalidate(swp_entry_t swp)
 		zswap_entry_free(entry);
 }
 
+/*
+ * This is called after the Shadow entry has been written. We replace the
+ * just-written Shadow with a Pointer entry pointing to the zswap_entry.
+ *
+ * Context: cluster lock must be held. Cannot sleep.
+ */
+void zswap_try_convert_to_pointer(struct swap_cluster_info *ci,
+				  unsigned int ci_off, int type, pgoff_t offset)
+{
+	struct xarray *tree;
+	struct zswap_entry *entry;
+	unsigned long swp_tb;
+
+	tree = &zswap_trees[type][offset >> ZSWAP_ADDRESS_SPACE_SHIFT];
+	if (xa_empty(tree))
+		return;
+
+	entry = xa_load(tree, offset);
+	if (!entry)
+		return;
+
+	swp_tb = __swap_table_get(ci, ci_off);
+	/* Entry must be a Shadow (just written by __swap_cache_do_del_folio) */
+	if (!swp_tb_is_countable(swp_tb))
+		return;
+
+	/*
+	 * Save the complete original swap table entry value, which includes
+	 * the swap count, zero flag, and working set shadow information.
+	 */
+	entry->swp_tb_val = swp_tb;
+
+	__swap_table_set(ci, ci_off, pointer_to_swp_tb(entry));
+}
+
 /*
  * Read the swap count from a Pointer-type swap table entry. The count
  * and flags are stored in zswap_entry->swp_tb_val which preserves the
-- 
2.54.0

[RFC PATCH 4/7] mm/zswap, mm/swap_state: handle Pointer entries in add_folio path

[Baoquan He]

Update __swap_cache_add_check and __swap_cache_do_add_folio to handle
Pointer-type swap table entries.

Add zswap_swp_tb_get_shadow() helper that extracts the xa_value shadow
from the saved swap table entry value in zswap_entry->swp_tb_val.

Signed-off-by: Baoquan He <baoquan.he@linux.dev>
---
 mm/swap.h       |   5 +++
 mm/swap_state.c | 112 ++++++++++++++++++++++++++++++++++++------------
 mm/zswap.c      |  22 ++++++++++
 3 files changed, 111 insertions(+), 28 deletions(-)

diff --git a/mm/swap.h b/mm/swap.h
index 1b07a08e5277..c23f54dc32bd 100644
--- a/mm/swap.h
+++ b/mm/swap.h
@@ -213,6 +213,7 @@ extern int swap_retry_table_alloc(swp_entry_t entry, gfp_t gfp);
 #ifdef CONFIG_ZSWAP
 void zswap_try_convert_to_pointer(struct swap_cluster_info *ci,
 				  unsigned int ci_off, int type, pgoff_t offset);
+void *zswap_swp_tb_get_shadow(unsigned long swp_tb);
 unsigned char zswap_swp_tb_get_count(unsigned long swp_tb);
 unsigned char zswap_swp_tb_get_flags(unsigned long swp_tb);
 int zswap_swp_tb_dup_count(unsigned long swp_tb,
@@ -223,6 +224,10 @@ void zswap_swp_tb_put_count(unsigned long swp_tb,
 static inline void zswap_try_convert_to_pointer(struct swap_cluster_info *ci,
 						unsigned int ci_off, int type,
 						pgoff_t offset) {}
+static inline void *zswap_swp_tb_get_shadow(unsigned long swp_tb)
+{
+	return NULL;
+}
 static inline unsigned char zswap_swp_tb_get_count(unsigned long swp_tb)
 {
 	return 0;
diff --git a/mm/swap_state.c b/mm/swap_state.c
index 20fb57adbb7f..620dc763eb3b 100644
--- a/mm/swap_state.c
+++ b/mm/swap_state.c
@@ -175,16 +175,29 @@ static int __swap_cache_add_check(struct swap_cluster_info *ci,
 	old_tb = __swap_table_get(ci, ci_off);
 	if (swp_tb_is_folio(old_tb))
 		return -EEXIST;
-	if (!__swp_tb_get_count(old_tb))
+	if (swp_tb_is_pointer(old_tb)) {
+		if (!zswap_swp_tb_get_count(old_tb))
+			return -ENOENT;
+	} else if (!__swp_tb_get_count(old_tb)) {
 		return -ENOENT;
+	}
 	if (shadowp && swp_tb_is_shadow(old_tb))
 		*shadowp = swp_tb_to_shadow(old_tb);
+	else if (shadowp && swp_tb_is_pointer(old_tb))
+		*shadowp = zswap_swp_tb_get_shadow(old_tb);
 	if (memcg_id)
 		*memcg_id = __swap_cgroup_get(ci, ci_off);
 
 	if (nr == 1)
 		return 0;
 
+	/*
+	 * For multi-slot ranges, Pointer entries are not expected since
+	 * zswap does not support large folios. Return -EBUSY for safety.
+	 */
+	if (swp_tb_is_pointer(old_tb))
+		return -EBUSY;
+
 	is_zero = __swap_table_test_zero(ci, ci_off);
 	ci_off = round_down(ci_off, nr);
 	ci_end = ci_off + nr;
@@ -201,7 +214,8 @@ static int __swap_cache_add_check(struct swap_cluster_info *ci,
 }
 
 static void __swap_cache_do_add_folio(struct swap_cluster_info *ci,
-				      struct folio *folio, swp_entry_t entry)
+				      struct folio *folio, swp_entry_t entry,
+				      void **zentry)
 {
 	unsigned int ci_off = swp_cluster_offset(entry), ci_end;
 	unsigned long nr_pages = folio_nr_pages(folio);
@@ -216,7 +230,26 @@ static void __swap_cache_do_add_folio(struct swap_cluster_info *ci,
 	do {
 		old_tb = __swap_table_get(ci, ci_off);
 		VM_WARN_ON_ONCE(swp_tb_is_folio(old_tb));
-		__swap_table_set(ci, ci_off, pfn_to_swp_tb(pfn, __swp_tb_get_flags(old_tb)));
+		/*
+		 * If the old entry is a Pointer (zswap compressed data),
+		 * extract the zswap_entry for the caller and preserve
+		 * its flags. The Pointer entry has no inline flags.
+		 */
+		if (swp_tb_is_pointer(old_tb)) {
+			unsigned char cur_count, cur_flags;
+			unsigned long new_tb;
+
+			cur_count = zswap_swp_tb_get_count(old_tb);
+			cur_flags = zswap_swp_tb_get_flags(old_tb);
+			if (zentry && !*zentry)
+				*zentry = swp_tb_to_pointer(old_tb);
+			new_tb = pfn_to_swp_tb(pfn, cur_flags);
+			new_tb = __swp_tb_mk_count(new_tb, cur_count);
+			__swap_table_set(ci, ci_off, new_tb);
+		} else {
+			__swap_table_set(ci, ci_off,
+				pfn_to_swp_tb(pfn, __swp_tb_get_flags(old_tb)));
+		}
 	} while (++ci_off < ci_end);
 
 	folio_ref_add(folio, nr_pages);
@@ -239,11 +272,12 @@ static void __swap_cache_do_add_folio(struct swap_cluster_info *ci,
  * that holds the entries.
  */
 void __swap_cache_add_folio(struct swap_cluster_info *ci,
-			    struct folio *folio, swp_entry_t entry)
+			    struct folio *folio, swp_entry_t entry,
+			    void **zentry)
 {
 	unsigned long nr_pages = folio_nr_pages(folio);
 
-	__swap_cache_do_add_folio(ci, folio, entry);
+	__swap_cache_do_add_folio(ci, folio, entry, zentry);
 	node_stat_mod_folio(folio, NR_FILE_PAGES, nr_pages);
 	lruvec_stat_mod_folio(folio, NR_SWAPCACHE, nr_pages);
 }
@@ -269,22 +303,28 @@ static void __swap_cache_do_del_folio(struct swap_cluster_info *ci,
 	ci_off = ci_start;
 	do {
 		old_tb = __swap_table_get(ci, ci_off);
-		WARN_ON_ONCE(!swp_tb_is_folio(old_tb) ||
-			     swp_tb_to_folio(old_tb) != folio);
-		if (__swp_tb_get_count(old_tb))
-			folio_swapped = true;
-		else
-			need_free = true;
-		/* If shadow is NULL, we set an empty shadow. */
-		__swap_table_set(ci, ci_off, shadow_to_swp_tb(shadow,
-				 __swp_tb_get_flags(old_tb)));
 		/*
-		 * If zswap has a compressed copy of this slot, convert the
-		 * just-written Shadow to a Pointer entry referencing the
-		 * zswap_entry, so the data can be found via the swap table.
+		 * If zswap already stored data for this slot
+		 * (zswap_store wrote a Pointer entry), leave it
+		 * in place. The count and flags are stored in the
+		 * zswap_entry struct.
 		 */
-		zswap_try_convert_to_pointer(ci, ci_off, si->type,
-					     swp_offset(entry) + ci_off - ci_start);
+		if (swp_tb_is_pointer(old_tb)) {
+			if (zswap_swp_tb_get_count(old_tb))
+				folio_swapped = true;
+			else
+				need_free = true;
+		} else {
+			WARN_ON_ONCE(!swp_tb_is_folio(old_tb) ||
+				     swp_tb_to_folio(old_tb) != folio);
+			if (__swp_tb_get_count(old_tb))
+				folio_swapped = true;
+			else
+				need_free = true;
+			/* If shadow is NULL, we set an empty shadow. */
+			__swap_table_set(ci, ci_off, shadow_to_swp_tb(shadow,
+					 __swp_tb_get_flags(old_tb)));
+		}
 	} while (++ci_off < ci_end);
 
 	folio->swap.val = 0;
@@ -295,7 +335,14 @@ static void __swap_cache_do_del_folio(struct swap_cluster_info *ci,
 	} else if (need_free) {
 		ci_off = ci_start;
 		do {
-			if (!__swp_tb_get_count(__swap_table_get(ci, ci_off)))
+			unsigned long tb = __swap_table_get(ci, ci_off);
+			unsigned char cnt;
+
+			if (swp_tb_is_pointer(tb))
+				cnt = zswap_swp_tb_get_count(tb);
+			else
+				cnt = __swp_tb_get_count(tb);
+			if (!cnt)
 				__swap_cluster_free_entries(si, ci, ci_off, 1);
 		} while (++ci_off < ci_end);
 	}
@@ -410,7 +457,8 @@ void __swap_cache_replace_folio(struct swap_cluster_info *ci,
 static struct folio *__swap_cache_alloc(struct swap_cluster_info *ci,
 					swp_entry_t targ_entry, gfp_t gfp,
 					unsigned int order, struct vm_fault *vmf,
-					struct mempolicy *mpol, pgoff_t ilx)
+					struct mempolicy *mpol, pgoff_t ilx,
+					void **zentry)
 {
 	int err;
 	swp_entry_t entry;
@@ -457,7 +505,7 @@ static struct folio *__swap_cache_alloc(struct swap_cluster_info *ci,
 
 	__folio_set_locked(folio);
 	__folio_set_swapbacked(folio);
-	__swap_cache_do_add_folio(ci, folio, entry);
+	__swap_cache_do_add_folio(ci, folio, entry, zentry);
 	spin_unlock(&ci->lock);
 
 	if (mem_cgroup_swapin_charge_folio(folio, memcg_id,
@@ -514,7 +562,8 @@ static struct folio *__swap_cache_alloc(struct swap_cluster_info *ci,
  */
 struct folio *swap_cache_alloc_folio(swp_entry_t targ_entry, gfp_t gfp,
 				     unsigned long orders, struct vm_fault *vmf,
-				     struct mempolicy *mpol, pgoff_t ilx)
+				     struct mempolicy *mpol, pgoff_t ilx,
+				     void **zentry)
 {
 	int order, err;
 	struct folio *ret;
@@ -527,9 +576,12 @@ struct folio *swap_cache_alloc_folio(swp_entry_t targ_entry, gfp_t gfp,
 	if (WARN_ON_ONCE(!orders || (1UL << order) > SWAPFILE_CLUSTER))
 		return ERR_PTR(-EINVAL);
 
+	if (zentry)
+		*zentry = NULL;
+
 	do {
 		ret = __swap_cache_alloc(ci, targ_entry, gfp, order,
-					 vmf, mpol, ilx);
+					 vmf, mpol, ilx, zentry);
 		if (!IS_ERR(ret))
 			break;
 		err = PTR_ERR(ret);
@@ -645,18 +697,20 @@ static struct folio *swap_cache_read_folio(swp_entry_t entry, gfp_t gfp,
 					   struct swap_iocb **plug, bool readahead)
 {
 	struct folio *folio;
+	void *zentry;
 
 	do {
 		folio = swap_cache_get_folio(entry);
 		if (folio)
 			return folio;
-		folio = swap_cache_alloc_folio(entry, gfp, BIT(0), NULL, mpol, ilx);
+		folio = swap_cache_alloc_folio(entry, gfp, BIT(0), NULL, mpol, ilx,
+					       &zentry);
 	} while (PTR_ERR(folio) == -EEXIST);
 
 	if (IS_ERR_OR_NULL(folio))
 		return NULL;
 
-	swap_read_folio(folio, plug);
+	swap_read_folio(folio, plug, zentry);
 	if (readahead) {
 		folio_set_readahead(folio);
 		count_vm_event(SWAP_RA);
@@ -685,18 +739,20 @@ struct folio *swapin_sync(swp_entry_t entry, gfp_t gfp, unsigned long orders,
 			   struct vm_fault *vmf, struct mempolicy *mpol, pgoff_t ilx)
 {
 	struct folio *folio;
+	void *zentry;
 
 	do {
 		folio = swap_cache_get_folio(entry);
 		if (folio)
 			return folio;
-		folio = swap_cache_alloc_folio(entry, gfp, orders, vmf, mpol, ilx);
+		folio = swap_cache_alloc_folio(entry, gfp, orders, vmf, mpol, ilx,
+					       &zentry);
 	} while (PTR_ERR(folio) == -EEXIST);
 
 	if (IS_ERR(folio))
 		return folio;
 
-	swap_read_folio(folio, NULL);
+	swap_read_folio(folio, NULL, zentry);
 	return folio;
 }
 
diff --git a/mm/zswap.c b/mm/zswap.c
index f6e8259ab56a..1bb484d1ac87 100644
--- a/mm/zswap.c
+++ b/mm/zswap.c
@@ -1687,6 +1687,28 @@ void zswap_try_convert_to_pointer(struct swap_cluster_info *ci,
 	__swap_table_set(ci, ci_off, pointer_to_swp_tb(entry));
 }
 
+/*
+ * Extract the working set shadow from a Pointer-type swap table entry.
+ * The shadow is recovered from zswap_entry->swp_tb_val which preserves
+ * the original Shadow entry value. The shadow can then be used by
+ * workingset_refault() for refault distance calculation.
+ */
+void *zswap_swp_tb_get_shadow(unsigned long swp_tb)
+{
+	struct zswap_entry *entry = swp_tb_to_pointer(swp_tb);
+	unsigned long saved = entry->swp_tb_val;
+
+	/*
+	 * swp_tb_val preserves the original swap table entry.  Only
+	 * Shadow entries carry workingset information; PFN and NULL
+	 * entries have no shadow to extract.
+	 */
+	if (!swp_tb_is_shadow(saved))
+		return NULL;
+
+	return (void *)(saved & ~SWP_TB_FLAGS_MASK);
+}
+
 /*
  * Read the swap count from a Pointer-type swap table entry. The count
  * and flags are stored in zswap_entry->swp_tb_val which preserves the
-- 
2.54.0

[RFC PATCH 5/7] mm/swapfile: handle Pointer entries in count/dup/put paths

[Baoquan He]

Update swap count, duplicate, and put operations to handle Pointer-type
swap table entries where the count is stored in zswap_entry->swp_tb_val
rather than inline in the entry bits.

Signed-off-by: Baoquan He <baoquan.he@linux.dev>
---
 mm/swapfile.c | 95 ++++++++++++++++++++++++++++++++++++++++-----------
 mm/zswap.c    |  5 +++
 2 files changed, 81 insertions(+), 19 deletions(-)

diff --git a/mm/swapfile.c b/mm/swapfile.c
index a602e5820513..15572383cafa 100644
--- a/mm/swapfile.c
+++ b/mm/swapfile.c
@@ -1536,25 +1536,51 @@ static void __swap_cluster_put_entry(struct swap_cluster_info *ci,
 
 	lockdep_assert_held(&ci->lock);
 	swp_tb = __swap_table_get(ci, ci_off);
+
+	/* Pointer entries store count in the zswap_entry struct */
+	if (swp_tb_is_pointer(swp_tb)) {
+		zswap_swp_tb_put_count(swp_tb, ci, ci_off);
+		return;
+	}
+
 	count = __swp_tb_get_count(swp_tb);
 
 	VM_WARN_ON_ONCE(count <= 0);
 	VM_WARN_ON_ONCE(count > SWP_TB_COUNT_MAX);
-
-	if (count == SWP_TB_COUNT_MAX) {
-		count = ci->extend_table[ci_off];
-		/* Overflow starts with SWP_TB_COUNT_MAX */
-		VM_WARN_ON_ONCE(count < SWP_TB_COUNT_MAX);
-		count--;
-		if (count == (SWP_TB_COUNT_MAX - 1)) {
-			ci->extend_table[ci_off] = 0;
-			__swap_table_set(ci, ci_off, __swp_tb_mk_count(swp_tb, count));
+		if (count == SWP_TB_COUNT_MAX) {
+			if (!ci->extend_table) {
+				/*
+				 * MAX without extend_table: MAX is the real
+				 * count.  Decrement and write back inline.
+				 */
+				count--;
+				__swap_table_set(ci, ci_off,
+					__swp_tb_mk_count(swp_tb, count));
+			} else {
+				/*
+				 * MAX with extend_table: real count lives in
+				 * ci->extend_table[ci_off].  Pull, decrement
+				 * it, and either move it back inline (if now
+				 * fits) or write it back to the extend table.
+				 */
+				count = ci->extend_table[ci_off];
+				/* Overflow starts with SWP_TB_COUNT_MAX */
+				VM_WARN_ON_ONCE(count < SWP_TB_COUNT_MAX);
+				count--;
+				if (count == (SWP_TB_COUNT_MAX - 1)) {
+					ci->extend_table[ci_off] = 0;
+					__swap_table_set(ci, ci_off,
+						__swp_tb_mk_count(swp_tb, count));
+				} else {
+					ci->extend_table[ci_off] = count;
+				}
+			}
 		} else {
-			ci->extend_table[ci_off] = count;
+			/* count < MAX: simple inline decrement */
+			count--;
+			__swap_table_set(ci, ci_off,
+				__swp_tb_mk_count(swp_tb, count));
 		}
-	} else {
-		__swap_table_set(ci, ci_off, __swp_tb_mk_count(swp_tb, --count));
-	}
 
 	/*
 	 * `SWP_TB_COUNT_MAX - 1` triggers extend table allocation. If the
@@ -1596,7 +1622,18 @@ static void swap_put_entries_cluster(struct swap_info_struct *si,
 	ci_end = ci_off + nr;
 	do {
 		swp_tb = __swap_table_get(ci, ci_off);
-		if (swp_tb_get_count(swp_tb) == 1) {
+		/*
+		 * Pointer entries store count in the zswap_entry;
+		 * handle them before swp_tb_get_count (which returns
+		 * -EINVAL for non-countable types).
+		 */
+		if (swp_tb_is_pointer(swp_tb)) {
+			if (zswap_swp_tb_get_count(swp_tb) == 1) {
+				if (ci_batch == -1)
+					ci_batch = ci_off;
+				continue;
+			}
+		} else if (swp_tb_get_count(swp_tb) == 1) {
 			/* count == 1 and non-cached slots will be batch freed. */
 			if (!swp_tb_is_folio(swp_tb)) {
 				if (ci_batch == -1)
@@ -1647,6 +1684,11 @@ static int __swap_cluster_dup_entry(struct swap_cluster_info *ci,
 	/* Bad or special slots can't be handled */
 	if (WARN_ON_ONCE(swp_tb_is_bad(swp_tb)))
 		return -EINVAL;
+
+	/* Pointer entries store count in the zswap_entry struct */
+	if (swp_tb_is_pointer(swp_tb))
+		return zswap_swp_tb_dup_count(swp_tb, ci, ci_off);
+
 	count = __swp_tb_get_count(swp_tb);
 	/* Must be either cached or have a count already */
 	if (WARN_ON_ONCE(!count && !swp_tb_is_folio(swp_tb)))
@@ -1923,13 +1965,19 @@ void __swap_cluster_free_entries(struct swap_info_struct *si,
 		old_tb = __swap_table_get(ci, ci_off);
 		/*
 		 * Freeing is done after release of the last swap count
-		 * ref, or after swap cache is dropped
+		 * ref, or after swap cache is dropped. A Pointer entry
+		 * means zswap has a compressed copy; the xarray still
+		 * holds the entry and zswap_invalidate will free it.
 		 */
-		VM_WARN_ON(!swp_tb_is_shadow(old_tb) || __swp_tb_get_count(old_tb) > 1);
+		if (swp_tb_is_pointer(old_tb))
+			VM_WARN_ON(zswap_swp_tb_get_count(old_tb) > 1);
+		else
+			VM_WARN_ON(!swp_tb_is_shadow(old_tb) ||
+				   __swp_tb_get_count(old_tb) > 1);
 
 		/* Resetting the slot to NULL also clears the inline flags. */
 		__swap_table_set(ci, ci_off, null_to_swp_tb());
-		if (!SWAP_TABLE_HAS_ZEROFLAG)
+		if (!SWAP_TABLE_HAS_ZEROFLAG && !swp_tb_is_pointer(old_tb))
 			__swap_table_clear_zero(ci, ci_off);
 
 		/*
@@ -1961,8 +2009,11 @@ int __swap_count(swp_entry_t entry)
 {
 	struct swap_cluster_info *ci = __swap_entry_to_cluster(entry);
 	unsigned int ci_off = swp_cluster_offset(entry);
+	unsigned long swp_tb = __swap_table_get(ci, ci_off);
 
-	return swp_tb_get_count(__swap_table_get(ci, ci_off));
+	if (swp_tb_is_pointer(swp_tb))
+		return zswap_swp_tb_get_count(swp_tb);
+	return swp_tb_get_count(swp_tb);
 }
 
 /**
@@ -1980,6 +2031,8 @@ bool swap_entry_swapped(struct swap_info_struct *si, swp_entry_t entry)
 	swp_tb = swap_table_get(ci, offset % SWAPFILE_CLUSTER);
 	swap_cluster_unlock(ci);
 
+	if (swp_tb_is_pointer(swp_tb))
+		return zswap_swp_tb_get_count(swp_tb) > 0;
 	return swp_tb_get_count(swp_tb) > 0;
 }
 
@@ -2556,8 +2609,12 @@ static int unuse_pte_range(struct vm_area_struct *vma, pmd_t *pmd,
 		if (!folio) {
 			swp_tb = swap_table_get(__swap_entry_to_cluster(entry),
 						swp_cluster_offset(entry));
-			if (swp_tb_get_count(swp_tb) <= 0)
+			if (swp_tb_is_pointer(swp_tb)) {
+				if (zswap_swp_tb_get_count(swp_tb) <= 0)
+					continue;
+			} else if (swp_tb_get_count(swp_tb) <= 0) {
 				continue;
+			}
 			return -ENOMEM;
 		}
 
diff --git a/mm/zswap.c b/mm/zswap.c
index 1bb484d1ac87..58259673ce2b 100644
--- a/mm/zswap.c
+++ b/mm/zswap.c
@@ -1774,6 +1774,11 @@ void zswap_swp_tb_put_count(unsigned long swp_tb,
 	VM_WARN_ON_ONCE(count == 0);
 
 	if (count == SWP_TB_COUNT_MAX) {
+		if (!ci->extend_table) {
+			count--;
+			entry->swp_tb_val = __swp_tb_mk_count(entry->swp_tb_val, count);
+			return;
+		}
 		count = ci->extend_table[ci_off];
 		/* Overflow starts with SWP_TB_COUNT_MAX */
 		VM_WARN_ON_ONCE(count < SWP_TB_COUNT_MAX);
-- 
2.54.0

[RFC PATCH 6/7] mm/zswap, mm/swap_state: migrate zswap to use swap_table for entry indexing

[Baoquan He]

Transition zswap to use swap table entries as the primary index for
locating compressed data. Key changes:

Signed-off-by: Baoquan He <baoquan.he@linux.dev>
---
 mm/swap_state.c | 62 ++++++++++++++---------------------
 mm/zswap.c      | 87 ++++++++++++++++++++++++++++++++++++++++---------
 2 files changed, 97 insertions(+), 52 deletions(-)

diff --git a/mm/swap_state.c b/mm/swap_state.c
index 620dc763eb3b..b3d5bedfaf07 100644
--- a/mm/swap_state.c
+++ b/mm/swap_state.c
@@ -214,8 +214,7 @@ static int __swap_cache_add_check(struct swap_cluster_info *ci,
 }
 
 static void __swap_cache_do_add_folio(struct swap_cluster_info *ci,
-				      struct folio *folio, swp_entry_t entry,
-				      void **zentry)
+				      struct folio *folio, swp_entry_t entry)
 {
 	unsigned int ci_off = swp_cluster_offset(entry), ci_end;
 	unsigned long nr_pages = folio_nr_pages(folio);
@@ -232,24 +231,15 @@ static void __swap_cache_do_add_folio(struct swap_cluster_info *ci,
 		VM_WARN_ON_ONCE(swp_tb_is_folio(old_tb));
 		/*
 		 * If the old entry is a Pointer (zswap compressed data),
-		 * extract the zswap_entry for the caller and preserve
-		 * its flags. The Pointer entry has no inline flags.
+		 * extract the flags from the zswap_entry. The Pointer
+		 * entry has no inline flags.
 		 */
-		if (swp_tb_is_pointer(old_tb)) {
-			unsigned char cur_count, cur_flags;
-			unsigned long new_tb;
-
-			cur_count = zswap_swp_tb_get_count(old_tb);
-			cur_flags = zswap_swp_tb_get_flags(old_tb);
-			if (zentry && !*zentry)
-				*zentry = swp_tb_to_pointer(old_tb);
-			new_tb = pfn_to_swp_tb(pfn, cur_flags);
-			new_tb = __swp_tb_mk_count(new_tb, cur_count);
-			__swap_table_set(ci, ci_off, new_tb);
-		} else {
+		if (swp_tb_is_pointer(old_tb))
+			__swap_table_set(ci, ci_off,
+				pfn_to_swp_tb(pfn, zswap_swp_tb_get_flags(old_tb)));
+		else
 			__swap_table_set(ci, ci_off,
 				pfn_to_swp_tb(pfn, __swp_tb_get_flags(old_tb)));
-		}
 	} while (++ci_off < ci_end);
 
 	folio_ref_add(folio, nr_pages);
@@ -272,12 +262,11 @@ static void __swap_cache_do_add_folio(struct swap_cluster_info *ci,
  * that holds the entries.
  */
 void __swap_cache_add_folio(struct swap_cluster_info *ci,
-			    struct folio *folio, swp_entry_t entry,
-			    void **zentry)
+			    struct folio *folio, swp_entry_t entry)
 {
 	unsigned long nr_pages = folio_nr_pages(folio);
 
-	__swap_cache_do_add_folio(ci, folio, entry, zentry);
+	__swap_cache_do_add_folio(ci, folio, entry);
 	node_stat_mod_folio(folio, NR_FILE_PAGES, nr_pages);
 	lruvec_stat_mod_folio(folio, NR_SWAPCACHE, nr_pages);
 }
@@ -324,6 +313,14 @@ static void __swap_cache_do_del_folio(struct swap_cluster_info *ci,
 			/* If shadow is NULL, we set an empty shadow. */
 			__swap_table_set(ci, ci_off, shadow_to_swp_tb(shadow,
 					 __swp_tb_get_flags(old_tb)));
+			/*
+			 * If zswap has a compressed copy of this slot,
+			 * convert the just-written Shadow to a Pointer
+			 * entry referencing the zswap_entry.
+			 */
+			zswap_try_convert_to_pointer(ci, ci_off, si->type,
+						     swp_offset(entry) +
+						     ci_off - ci_start);
 		}
 	} while (++ci_off < ci_end);
 
@@ -457,8 +454,7 @@ void __swap_cache_replace_folio(struct swap_cluster_info *ci,
 static struct folio *__swap_cache_alloc(struct swap_cluster_info *ci,
 					swp_entry_t targ_entry, gfp_t gfp,
 					unsigned int order, struct vm_fault *vmf,
-					struct mempolicy *mpol, pgoff_t ilx,
-					void **zentry)
+					struct mempolicy *mpol, pgoff_t ilx)
 {
 	int err;
 	swp_entry_t entry;
@@ -505,7 +501,7 @@ static struct folio *__swap_cache_alloc(struct swap_cluster_info *ci,
 
 	__folio_set_locked(folio);
 	__folio_set_swapbacked(folio);
-	__swap_cache_do_add_folio(ci, folio, entry, zentry);
+	__swap_cache_do_add_folio(ci, folio, entry);
 	spin_unlock(&ci->lock);
 
 	if (mem_cgroup_swapin_charge_folio(folio, memcg_id,
@@ -562,8 +558,7 @@ static struct folio *__swap_cache_alloc(struct swap_cluster_info *ci,
  */
 struct folio *swap_cache_alloc_folio(swp_entry_t targ_entry, gfp_t gfp,
 				     unsigned long orders, struct vm_fault *vmf,
-				     struct mempolicy *mpol, pgoff_t ilx,
-				     void **zentry)
+				     struct mempolicy *mpol, pgoff_t ilx)
 {
 	int order, err;
 	struct folio *ret;
@@ -576,12 +571,9 @@ struct folio *swap_cache_alloc_folio(swp_entry_t targ_entry, gfp_t gfp,
 	if (WARN_ON_ONCE(!orders || (1UL << order) > SWAPFILE_CLUSTER))
 		return ERR_PTR(-EINVAL);
 
-	if (zentry)
-		*zentry = NULL;
-
 	do {
 		ret = __swap_cache_alloc(ci, targ_entry, gfp, order,
-					 vmf, mpol, ilx, zentry);
+					 vmf, mpol, ilx);
 		if (!IS_ERR(ret))
 			break;
 		err = PTR_ERR(ret);
@@ -697,20 +689,18 @@ static struct folio *swap_cache_read_folio(swp_entry_t entry, gfp_t gfp,
 					   struct swap_iocb **plug, bool readahead)
 {
 	struct folio *folio;
-	void *zentry;
 
 	do {
 		folio = swap_cache_get_folio(entry);
 		if (folio)
 			return folio;
-		folio = swap_cache_alloc_folio(entry, gfp, BIT(0), NULL, mpol, ilx,
-					       &zentry);
+		folio = swap_cache_alloc_folio(entry, gfp, BIT(0), NULL, mpol, ilx);
 	} while (PTR_ERR(folio) == -EEXIST);
 
 	if (IS_ERR_OR_NULL(folio))
 		return NULL;
 
-	swap_read_folio(folio, plug, zentry);
+	swap_read_folio(folio, plug);
 	if (readahead) {
 		folio_set_readahead(folio);
 		count_vm_event(SWAP_RA);
@@ -739,20 +729,18 @@ struct folio *swapin_sync(swp_entry_t entry, gfp_t gfp, unsigned long orders,
 			   struct vm_fault *vmf, struct mempolicy *mpol, pgoff_t ilx)
 {
 	struct folio *folio;
-	void *zentry;
 
 	do {
 		folio = swap_cache_get_folio(entry);
 		if (folio)
 			return folio;
-		folio = swap_cache_alloc_folio(entry, gfp, orders, vmf, mpol, ilx,
-					       &zentry);
+		folio = swap_cache_alloc_folio(entry, gfp, orders, vmf, mpol, ilx);
 	} while (PTR_ERR(folio) == -EEXIST);
 
 	if (IS_ERR(folio))
 		return folio;
 
-	swap_read_folio(folio, NULL, zentry);
+	swap_read_folio(folio, NULL);
 	return folio;
 }
 
diff --git a/mm/zswap.c b/mm/zswap.c
index 58259673ce2b..f5f4756c9ce6 100644
--- a/mm/zswap.c
+++ b/mm/zswap.c
@@ -1404,7 +1404,11 @@ static bool zswap_store_page(struct page *page,
 			     struct zswap_pool *pool)
 {
 	swp_entry_t page_swpentry = page_swap_entry(page);
-	struct zswap_entry *entry, *old;
+	struct zswap_entry *entry, *old = NULL;
+	struct swap_cluster_info *ci;
+	struct swap_info_struct *si;
+	pgoff_t offset = swp_offset(page_swpentry);
+	unsigned long swp_tb;
 
 	/* allocate entry */
 	entry = zswap_entry_cache_alloc(GFP_KERNEL, page_to_nid(page));
@@ -1427,20 +1431,42 @@ static bool zswap_store_page(struct page *page,
 		goto store_failed;
 	}
 
+	si = __swap_entry_to_info(page_swpentry);
+	ci = __swap_offset_to_cluster(si, offset);
+	spin_lock(&ci->lock);
+	swp_tb = __swap_table_get(ci, offset % SWAPFILE_CLUSTER);
+	if (swp_tb_is_countable(swp_tb)) {
+		/*
+		 * Save the complete original swap table entry value
+		 * (Shadow or PFN), preserving swap count, zero flag,
+		 * and working set shadow information.
+		 */
+		entry->swp_tb_val = swp_tb;
+	} else if (swp_tb_is_pointer(swp_tb)) {
+		/*
+		 * Replacing an existing Pointer entry (e.g. from a
+		 * redirtied folio). Inherit the saved swap table value.
+		 */
+		old = swp_tb_to_pointer(swp_tb);
+		entry->swp_tb_val = old->swp_tb_val;
+	}
+	__swap_table_set(ci, offset % SWAPFILE_CLUSTER,
+			 pointer_to_swp_tb(entry));
+	spin_unlock(&ci->lock);
 	/*
-	 * We may have had an existing entry that became stale when
-	 * the folio was redirtied and now the new version is being
-	 * swapped out. Get rid of the old.
+	 * If we had an existing zswap entry from the swap table
+	 * (e.g. from a previous redirtied folio), get rid of it
+	 * outside the lock.
 	 */
 	if (old)
 		zswap_entry_free(old);
 
 	/*
-	 * The entry is successfully compressed and stored in the tree, there is
-	 * no further possibility of failure. Grab refs to the pool and objcg,
-	 * charge zswap memory, and increment zswap_stored_pages.
-	 * The opposite actions will be performed by zswap_entry_free()
-	 * when the entry is removed from the tree.
+	 * The entry is successfully compressed and stored in the swap table,
+	 * there is no further possibility of failure. Grab refs to the pool
+	 * and objcg, charge zswap memory, and increment zswap_stored_pages.
+	 * The opposite actions will be performed by zswap_entry_put()
+	 * when the entry is removed from the swap table.
 	 */
 	zswap_pool_get(pool);
 	if (objcg) {
@@ -1452,8 +1478,8 @@ static bool zswap_store_page(struct page *page,
 		atomic_long_inc(&zswap_stored_incompressible_pages);
 
 	/*
-	 * We finish initializing the entry while it's already in xarray.
-	 * This is safe because:
+	 * We finish initializing the entry while it's already in the
+	 * swap table. This is safe because:
 	 *
 	 * 1. Concurrent stores and invalidations are excluded by folio lock.
 	 *
@@ -1553,12 +1579,25 @@ bool zswap_store(struct folio *folio)
 		pgoff_t offset = swp_offset(swp);
 		struct zswap_entry *entry;
 		struct xarray *tree;
+		struct swap_cluster_info *ci;
+		struct swap_info_struct *sis;
+		unsigned long swp_tb;
 
+		sis = __swap_type_to_info(type);
 		for (index = 0; index < nr_pages; ++index) {
 			tree = swap_zswap_tree(swp_entry(type, offset + index));
 			entry = xa_erase(tree, offset + index);
-			if (entry)
-				zswap_entry_free(entry);
+			if (!entry)
+				continue;
+			ci = __swap_offset_to_cluster(sis, offset + index);
+			spin_lock(&ci->lock);
+			swp_tb = __swap_table_get(ci, (offset + index) % SWAPFILE_CLUSTER);
+			if (swp_tb_is_pointer(swp_tb) &&
+			    swp_tb_to_pointer(swp_tb) == entry)
+				__swap_table_set(ci, (offset + index) % SWAPFILE_CLUSTER,
+						 entry->swp_tb_val);
+			spin_unlock(&ci->lock);
+			zswap_entry_free(entry);
 		}
 	}
 
@@ -1642,14 +1681,32 @@ void zswap_invalidate(swp_entry_t swp)
 {
 	pgoff_t offset = swp_offset(swp);
 	struct xarray *tree = swap_zswap_tree(swp);
+	struct swap_cluster_info *ci;
+	struct swap_info_struct *si;
+	unsigned long swp_tb;
 	struct zswap_entry *entry;
 
 	if (xa_empty(tree))
 		return;
 
 	entry = xa_erase(tree, offset);
-	if (entry)
-		zswap_entry_free(entry);
+	if (!entry)
+		return;
+
+	/*
+	 * Also clear the swap table Pointer entry if present.
+	 * This is needed because zswap_store now writes Pointer
+	 * entries to both the xarray and the swap table.
+	 */
+	si = __swap_type_to_info(swp_type(swp));
+	ci = __swap_offset_to_cluster(si, offset);
+	spin_lock(&ci->lock);
+	swp_tb = __swap_table_get(ci, offset % SWAPFILE_CLUSTER);
+	if (swp_tb_is_pointer(swp_tb) && swp_tb_to_pointer(swp_tb) == entry)
+		__swap_table_set(ci, offset % SWAPFILE_CLUSTER, entry->swp_tb_val);
+	spin_unlock(&ci->lock);
+
+	zswap_entry_free(entry);
 }
 
 /*
-- 
2.54.0

[RFC PATCH 7/7] mm/zswap, mm/swap_state: replace xarray indexing with swap table lookups

[Baoquan He]

Remove the xarray-based zswap entry lookup infrastructure entirely and
use the swap table Pointer entries as the sole index for locating
compressed zswap data.

Thread the zswap_entry through the swapin call chain via a new
void **zentry parameter.

Update all declarations in swap.h and include/linux/zswap.h to match
the new signatures.

Signed-off-by: Baoquan He <baoquan.he@linux.dev>
---
 include/linux/zswap.h |   4 +-
 mm/page_io.c          |   4 +-
 mm/swap.h             |  25 +++--
 mm/swap_state.c       |  68 ++++++++----
 mm/swapfile.c         |  16 ++-
 mm/zswap.c            | 252 +++++++++++++++---------------------------
 6 files changed, 165 insertions(+), 204 deletions(-)

diff --git a/include/linux/zswap.h b/include/linux/zswap.h
index 30c193a1207e..4f55c8d0c4c5 100644
--- a/include/linux/zswap.h
+++ b/include/linux/zswap.h
@@ -26,7 +26,7 @@ struct zswap_lruvec_state {
 
 unsigned long zswap_total_pages(void);
 bool zswap_store(struct folio *folio);
-int zswap_load(struct folio *folio);
+int zswap_load(struct folio *folio, void *zentry);
 void zswap_invalidate(swp_entry_t swp);
 int zswap_swapon(int type, unsigned long nr_pages);
 void zswap_swapoff(int type);
@@ -44,7 +44,7 @@ static inline bool zswap_store(struct folio *folio)
 	return false;
 }
 
-static inline int zswap_load(struct folio *folio)
+static inline int zswap_load(struct folio *folio, void *zentry)
 {
 	return -ENOENT;
 }
diff --git a/mm/page_io.c b/mm/page_io.c
index b23f494fcc83..1ce8201ed86e 100644
--- a/mm/page_io.c
+++ b/mm/page_io.c
@@ -658,7 +658,7 @@ static void swap_read_folio_bdev_async(struct folio *folio,
 	submit_bio(bio);
 }
 
-void swap_read_folio(struct folio *folio, struct swap_iocb **plug)
+void swap_read_folio(struct folio *folio, struct swap_iocb **plug, void *zentry)
 {
 	struct swap_info_struct *sis = __swap_entry_to_info(folio->swap);
 	bool synchronous = sis->flags & SWP_SYNCHRONOUS_IO;
@@ -686,7 +686,7 @@ void swap_read_folio(struct folio *folio, struct swap_iocb **plug)
 		goto finish;
 	}
 
-	if (zswap_load(folio) != -ENOENT)
+	if (zswap_load(folio, zentry) != -ENOENT)
 		goto finish;
 
 	/* We have to read from slower devices. Increase zswap protection. */
diff --git a/mm/swap.h b/mm/swap.h
index c23f54dc32bd..84e13ba6e605 100644
--- a/mm/swap.h
+++ b/mm/swap.h
@@ -210,9 +210,8 @@ static inline void swap_cluster_unlock_irq(struct swap_cluster_info *ci)
 
 extern int swap_retry_table_alloc(swp_entry_t entry, gfp_t gfp);
 
+struct zswap_entry;
 #ifdef CONFIG_ZSWAP
-void zswap_try_convert_to_pointer(struct swap_cluster_info *ci,
-				  unsigned int ci_off, int type, pgoff_t offset);
 void *zswap_swp_tb_get_shadow(unsigned long swp_tb);
 unsigned char zswap_swp_tb_get_count(unsigned long swp_tb);
 unsigned char zswap_swp_tb_get_flags(unsigned long swp_tb);
@@ -220,10 +219,8 @@ int zswap_swp_tb_dup_count(unsigned long swp_tb,
 			   struct swap_cluster_info *ci, unsigned int ci_off);
 void zswap_swp_tb_put_count(unsigned long swp_tb,
 			    struct swap_cluster_info *ci, unsigned int ci_off);
+void zswap_entry_free(struct zswap_entry *entry);
 #else
-static inline void zswap_try_convert_to_pointer(struct swap_cluster_info *ci,
-						unsigned int ci_off, int type,
-						pgoff_t offset) {}
 static inline void *zswap_swp_tb_get_shadow(unsigned long swp_tb)
 {
 	return NULL;
@@ -245,6 +242,7 @@ static inline int zswap_swp_tb_dup_count(unsigned long swp_tb,
 static inline void zswap_swp_tb_put_count(unsigned long swp_tb,
 					  struct swap_cluster_info *ci,
 					  unsigned int ci_off) {}
+static inline void zswap_entry_free(struct zswap_entry *entry) {}
 #endif
 
 /*
@@ -278,7 +276,7 @@ extern void __swap_cluster_free_entries(struct swap_info_struct *si,
 /* linux/mm/page_io.c */
 int sio_pool_init(void);
 struct swap_iocb;
-void swap_read_folio(struct folio *folio, struct swap_iocb **plug);
+void swap_read_folio(struct folio *folio, struct swap_iocb **plug, void *zentry);
 void __swap_read_unplug(struct swap_iocb *plug);
 static inline void swap_read_unplug(struct swap_iocb *plug)
 {
@@ -340,13 +338,16 @@ static inline bool folio_matches_swap_entry(const struct folio *folio,
 bool swap_cache_has_folio(swp_entry_t entry);
 struct folio *swap_cache_get_folio(swp_entry_t entry);
 void *swap_cache_get_shadow(swp_entry_t entry);
+bool folio_maybe_swapped(struct folio *folio);
 void swap_cache_del_folio(struct folio *folio);
 struct folio *swap_cache_alloc_folio(swp_entry_t target_entry, gfp_t gfp_mask,
 				     unsigned long orders, struct vm_fault *vmf,
-				     struct mempolicy *mpol, pgoff_t ilx);
+				     struct mempolicy *mpol, pgoff_t ilx,
+				     void **zentry);
 /* Below helpers require the caller to lock and pass in the swap cluster. */
 void __swap_cache_add_folio(struct swap_cluster_info *ci,
-			    struct folio *folio, swp_entry_t entry);
+			    struct folio *folio, swp_entry_t entry,
+			    void **zentry);
 void __swap_cache_del_folio(struct swap_cluster_info *ci,
 			    struct folio *folio, swp_entry_t entry, void *shadow);
 void __swap_cache_replace_folio(struct swap_cluster_info *ci,
@@ -418,7 +419,8 @@ static inline void folio_put_swap(struct folio *folio, struct page *page)
 {
 }
 
-static inline void swap_read_folio(struct folio *folio, struct swap_iocb **plug)
+static inline void swap_read_folio(struct folio *folio, struct swap_iocb **plug,
+				   void *zentry)
 {
 }
 
@@ -490,6 +492,11 @@ static inline void *swap_cache_get_shadow(swp_entry_t entry)
 	return NULL;
 }
 
+static inline bool folio_maybe_swapped(struct folio *folio)
+{
+	return false;
+}
+
 static inline void swap_cache_del_folio(struct folio *folio)
 {
 }
diff --git a/mm/swap_state.c b/mm/swap_state.c
index b3d5bedfaf07..60d937c6b392 100644
--- a/mm/swap_state.c
+++ b/mm/swap_state.c
@@ -214,7 +214,8 @@ static int __swap_cache_add_check(struct swap_cluster_info *ci,
 }
 
 static void __swap_cache_do_add_folio(struct swap_cluster_info *ci,
-				      struct folio *folio, swp_entry_t entry)
+				      struct folio *folio, swp_entry_t entry,
+				      void **zentry)
 {
 	unsigned int ci_off = swp_cluster_offset(entry), ci_end;
 	unsigned long nr_pages = folio_nr_pages(folio);
@@ -231,15 +232,32 @@ static void __swap_cache_do_add_folio(struct swap_cluster_info *ci,
 		VM_WARN_ON_ONCE(swp_tb_is_folio(old_tb));
 		/*
 		 * If the old entry is a Pointer (zswap compressed data),
-		 * extract the flags from the zswap_entry. The Pointer
-		 * entry has no inline flags.
+		 * extract the zswap_entry for the caller and preserve
+		 * its flags. The Pointer entry has no inline flags.
 		 */
-		if (swp_tb_is_pointer(old_tb))
+		if (swp_tb_is_pointer(old_tb)) {
+			unsigned char pflags = zswap_swp_tb_get_flags(old_tb);
+			/*
+			 * If swp_tb_val carries SWP_TB_COUNT_MAX, the
+			 * real overflow count lives in ci->extend_table.
+			 * Pull it so the new PFN carries the correct count
+			 * and won't later crash on a missing extend_table.
+			 */
+			if (zswap_swp_tb_get_count(old_tb) == SWP_TB_COUNT_MAX) {
+				unsigned char real = SWP_TB_COUNT_MAX - 1;
+				if (ci->extend_table && ci->extend_table[ci_off])
+					real = ci->extend_table[ci_off];
+				pflags = (real << (SWP_TB_FLAGS_BITS - SWP_TB_COUNT_BITS)) |
+					 (pflags & ((1 << (SWP_TB_FLAGS_BITS - SWP_TB_COUNT_BITS)) - 1));
+			}
+			if (zentry && !*zentry)
+				*zentry = swp_tb_to_pointer(old_tb);
 			__swap_table_set(ci, ci_off,
-				pfn_to_swp_tb(pfn, zswap_swp_tb_get_flags(old_tb)));
-		else
+				pfn_to_swp_tb(pfn, pflags));
+		} else {
 			__swap_table_set(ci, ci_off,
 				pfn_to_swp_tb(pfn, __swp_tb_get_flags(old_tb)));
+		}
 	} while (++ci_off < ci_end);
 
 	folio_ref_add(folio, nr_pages);
@@ -262,11 +280,12 @@ static void __swap_cache_do_add_folio(struct swap_cluster_info *ci,
  * that holds the entries.
  */
 void __swap_cache_add_folio(struct swap_cluster_info *ci,
-			    struct folio *folio, swp_entry_t entry)
+			    struct folio *folio, swp_entry_t entry,
+			    void **zentry)
 {
 	unsigned long nr_pages = folio_nr_pages(folio);
 
-	__swap_cache_do_add_folio(ci, folio, entry);
+	__swap_cache_do_add_folio(ci, folio, entry, zentry);
 	node_stat_mod_folio(folio, NR_FILE_PAGES, nr_pages);
 	lruvec_stat_mod_folio(folio, NR_SWAPCACHE, nr_pages);
 }
@@ -313,14 +332,6 @@ static void __swap_cache_do_del_folio(struct swap_cluster_info *ci,
 			/* If shadow is NULL, we set an empty shadow. */
 			__swap_table_set(ci, ci_off, shadow_to_swp_tb(shadow,
 					 __swp_tb_get_flags(old_tb)));
-			/*
-			 * If zswap has a compressed copy of this slot,
-			 * convert the just-written Shadow to a Pointer
-			 * entry referencing the zswap_entry.
-			 */
-			zswap_try_convert_to_pointer(ci, ci_off, si->type,
-						     swp_offset(entry) +
-						     ci_off - ci_start);
 		}
 	} while (++ci_off < ci_end);
 
@@ -454,7 +465,8 @@ void __swap_cache_replace_folio(struct swap_cluster_info *ci,
 static struct folio *__swap_cache_alloc(struct swap_cluster_info *ci,
 					swp_entry_t targ_entry, gfp_t gfp,
 					unsigned int order, struct vm_fault *vmf,
-					struct mempolicy *mpol, pgoff_t ilx)
+					struct mempolicy *mpol, pgoff_t ilx,
+					void **zentry)
 {
 	int err;
 	swp_entry_t entry;
@@ -501,7 +513,7 @@ static struct folio *__swap_cache_alloc(struct swap_cluster_info *ci,
 
 	__folio_set_locked(folio);
 	__folio_set_swapbacked(folio);
-	__swap_cache_do_add_folio(ci, folio, entry);
+	__swap_cache_do_add_folio(ci, folio, entry, zentry);
 	spin_unlock(&ci->lock);
 
 	if (mem_cgroup_swapin_charge_folio(folio, memcg_id,
@@ -558,7 +570,8 @@ static struct folio *__swap_cache_alloc(struct swap_cluster_info *ci,
  */
 struct folio *swap_cache_alloc_folio(swp_entry_t targ_entry, gfp_t gfp,
 				     unsigned long orders, struct vm_fault *vmf,
-				     struct mempolicy *mpol, pgoff_t ilx)
+				     struct mempolicy *mpol, pgoff_t ilx,
+				     void **zentry)
 {
 	int order, err;
 	struct folio *ret;
@@ -571,9 +584,12 @@ struct folio *swap_cache_alloc_folio(swp_entry_t targ_entry, gfp_t gfp,
 	if (WARN_ON_ONCE(!orders || (1UL << order) > SWAPFILE_CLUSTER))
 		return ERR_PTR(-EINVAL);
 
+	if (zentry)
+		*zentry = NULL;
+
 	do {
 		ret = __swap_cache_alloc(ci, targ_entry, gfp, order,
-					 vmf, mpol, ilx);
+					 vmf, mpol, ilx, zentry);
 		if (!IS_ERR(ret))
 			break;
 		err = PTR_ERR(ret);
@@ -689,18 +705,20 @@ static struct folio *swap_cache_read_folio(swp_entry_t entry, gfp_t gfp,
 					   struct swap_iocb **plug, bool readahead)
 {
 	struct folio *folio;
+	void *zentry;
 
 	do {
 		folio = swap_cache_get_folio(entry);
 		if (folio)
 			return folio;
-		folio = swap_cache_alloc_folio(entry, gfp, BIT(0), NULL, mpol, ilx);
+		folio = swap_cache_alloc_folio(entry, gfp, BIT(0), NULL, mpol, ilx,
+					       &zentry);
 	} while (PTR_ERR(folio) == -EEXIST);
 
 	if (IS_ERR_OR_NULL(folio))
 		return NULL;
 
-	swap_read_folio(folio, plug);
+	swap_read_folio(folio, plug, zentry);
 	if (readahead) {
 		folio_set_readahead(folio);
 		count_vm_event(SWAP_RA);
@@ -729,18 +747,20 @@ struct folio *swapin_sync(swp_entry_t entry, gfp_t gfp, unsigned long orders,
 			   struct vm_fault *vmf, struct mempolicy *mpol, pgoff_t ilx)
 {
 	struct folio *folio;
+	void *zentry;
 
 	do {
 		folio = swap_cache_get_folio(entry);
 		if (folio)
 			return folio;
-		folio = swap_cache_alloc_folio(entry, gfp, orders, vmf, mpol, ilx);
+		folio = swap_cache_alloc_folio(entry, gfp, orders, vmf, mpol, ilx,
+					       &zentry);
 	} while (PTR_ERR(folio) == -EEXIST);
 
 	if (IS_ERR(folio))
 		return folio;
 
-	swap_read_folio(folio, NULL);
+	swap_read_folio(folio, NULL, zentry);
 	return folio;
 }
 
diff --git a/mm/swapfile.c b/mm/swapfile.c
index 15572383cafa..18925e462571 100644
--- a/mm/swapfile.c
+++ b/mm/swapfile.c
@@ -935,7 +935,8 @@ static bool __swap_cluster_alloc_entries(struct swap_info_struct *si,
 		nr_pages = 1 << order;
 		swap_cluster_assert_empty(ci, ci_off, nr_pages, false);
 		__swap_cache_add_folio(ci, folio, swp_entry(si->type,
-							    ci_off + cluster_offset(si, ci)));
+							    ci_off + cluster_offset(si, ci)),
+				       NULL);
 	} else if (IS_ENABLED(CONFIG_HIBERNATION)) {
 		order = 0;
 		nr_pages = 1;
@@ -1969,11 +1970,18 @@ void __swap_cluster_free_entries(struct swap_info_struct *si,
 		 * means zswap has a compressed copy; the xarray still
 		 * holds the entry and zswap_invalidate will free it.
 		 */
-		if (swp_tb_is_pointer(old_tb))
+		if (swp_tb_is_pointer(old_tb)) {
 			VM_WARN_ON(zswap_swp_tb_get_count(old_tb) > 1);
-		else
+			/*
+			 * Free the zswap entry now under ci->lock while
+			 * the Pointer is still in the swap table.  Once
+			 * we null the slot, the entry would be leaked.
+			 */
+			zswap_entry_free(swp_tb_to_pointer(old_tb));
+		} else {
 			VM_WARN_ON(!swp_tb_is_shadow(old_tb) ||
 				   __swp_tb_get_count(old_tb) > 1);
+		}
 
 		/* Resetting the slot to NULL also clears the inline flags. */
 		__swap_table_set(ci, ci_off, null_to_swp_tb());
@@ -2073,7 +2081,7 @@ int swp_swapcount(swp_entry_t entry)
  *
  * Context: Caller must ensure the folio is locked and in the swap cache.
  */
-static bool folio_maybe_swapped(struct folio *folio)
+bool folio_maybe_swapped(struct folio *folio)
 {
 	swp_entry_t entry = folio->swap;
 	struct swap_cluster_info *ci;
diff --git a/mm/zswap.c b/mm/zswap.c
index f5f4756c9ce6..f9762d0a596a 100644
--- a/mm/zswap.c
+++ b/mm/zswap.c
@@ -203,9 +203,6 @@ struct zswap_entry {
 	struct list_head lru;
 };
 
-static struct xarray *zswap_trees[MAX_SWAPFILES];
-static unsigned int nr_zswap_trees[MAX_SWAPFILES];
-
 /* RCU-protected iteration */
 static LIST_HEAD(zswap_pools);
 /* protects zswap_pools list modification */
@@ -231,15 +228,6 @@ static bool zswap_has_pool;
 * helpers and fwd declarations
 **********************************/
 
-/* One swap address space for each 64M swap space */
-#define ZSWAP_ADDRESS_SPACE_SHIFT 14
-#define ZSWAP_ADDRESS_SPACE_PAGES (1 << ZSWAP_ADDRESS_SPACE_SHIFT)
-static inline struct xarray *swap_zswap_tree(swp_entry_t swp)
-{
-	return &zswap_trees[swp_type(swp)][swp_offset(swp)
-		>> ZSWAP_ADDRESS_SPACE_SHIFT];
-}
-
 #define zswap_pool_debug(msg, p)			\
 	pr_debug("%s pool %s\n", msg, (p)->tfm_name)
 
@@ -768,7 +756,7 @@ static void zswap_entry_cache_free(struct zswap_entry *entry)
  * Carries out the common pattern of freeing an entry's zsmalloc allocation,
  * freeing the entry itself, and decrementing the number of stored pages.
  */
-static void zswap_entry_free(struct zswap_entry *entry)
+void zswap_entry_free(struct zswap_entry *entry)
 {
 	zswap_lru_del(&zswap_list_lru, entry);
 	zs_free(entry->pool->zs_pool, entry->handle);
@@ -993,11 +981,12 @@ static bool zswap_decompress(struct zswap_entry *entry, struct folio *folio)
 static int zswap_writeback_entry(struct zswap_entry *entry,
 				 swp_entry_t swpentry)
 {
-	struct xarray *tree;
 	pgoff_t offset = swp_offset(swpentry);
 	struct folio *folio;
 	struct mempolicy *mpol;
 	struct swap_info_struct *si;
+	struct swap_cluster_info *ci;
+	unsigned long swp_tb;
 	int ret = 0;
 
 	/* try to allocate swap cache folio */
@@ -1005,9 +994,25 @@ static int zswap_writeback_entry(struct zswap_entry *entry,
 	if (!si)
 		return -EEXIST;
 
+	/*
+	 * Verify the swap table Pointer entry still references this
+	 * zswap_entry. If the entry was invalidated or replaced by a
+	 * concurrent store, the Pointer will be gone or point elsewhere.
+	 */
+	ci = __swap_offset_to_cluster(si, offset);
+	spin_lock(&ci->lock);
+	swp_tb = __swap_table_get(ci, offset % SWAPFILE_CLUSTER);
+	if (!swp_tb_is_pointer(swp_tb) ||
+	    swp_tb_to_pointer(swp_tb) != entry) {
+		spin_unlock(&ci->lock);
+		put_swap_device(si);
+		return -ENOMEM;
+	}
+	spin_unlock(&ci->lock);
+
 	mpol = get_task_policy(current);
 	folio = swap_cache_alloc_folio(swpentry, GFP_KERNEL, BIT(0), NULL, mpol,
-				       NO_INTERLEAVE_INDEX);
+				       NO_INTERLEAVE_INDEX, NULL);
 	put_swap_device(si);
 
 	/*
@@ -1021,27 +1026,15 @@ static int zswap_writeback_entry(struct zswap_entry *entry,
 		return PTR_ERR(folio);
 
 	/*
-	 * folio is locked, and the swapcache is now secured against
-	 * concurrent swapping to and from the slot, and concurrent
-	 * swapoff so we can safely dereference the zswap tree here.
-	 * Verify that the swap entry hasn't been invalidated and recycled
-	 * behind our backs, to avoid overwriting a new swap folio with
-	 * old compressed data. Only when this is successful can the entry
-	 * be dereferenced.
+	 * The folio is now locked and in swap cache, so concurrent
+	 * stores are blocked. __swap_cache_do_add_folio already
+	 * overwrote the Pointer entry with a PFN.
 	 */
-	tree = swap_zswap_tree(swpentry);
-	if (entry != xa_load(tree, offset)) {
-		ret = -ENOMEM;
-		goto out;
-	}
-
 	if (!zswap_decompress(entry, folio)) {
 		ret = -EIO;
 		goto out;
 	}
 
-	xa_erase(tree, offset);
-
 	count_vm_event(ZSWPWB);
 	if (entry->objcg)
 		count_objcg_events(entry->objcg, ZSWPWB, 1);
@@ -1111,36 +1104,22 @@ static enum lru_status shrink_memcg_cb(struct list_head *item, struct list_lru_o
 	}
 
 	/*
-	 * As soon as we drop the LRU lock, the entry can be freed by
-	 * a concurrent invalidation. This means the following:
-	 *
-	 * 1. We extract the swp_entry_t to the stack, allowing
-	 *    zswap_writeback_entry() to pin the swap entry and
-	 *    then validate the zswap entry against that swap entry's
-	 *    tree using pointer value comparison. Only when that
-	 *    is successful can the entry be dereferenced.
+	 * After we drop the LRU lock, the entry may be freed by a
+	 * concurrent swap slot free.  This is safe because
+	 * zswap_writeback_entry() rechecks the swap table under
+	 * ci->lock before dereferencing entry; if the Pointer is
+	 * gone it bails out.
 	 *
-	 * 2. Usually, objects are taken off the LRU for reclaim. In
-	 *    this case this isn't possible, because if reclaim fails
-	 *    for whatever reason, we have no means of knowing if the
-	 *    entry is alive to put it back on the LRU.
-	 *
-	 *    So rotate it before dropping the lock. If the entry is
-	 *    written back or invalidated, the free path will unlink
-	 *    it. For failures, rotation is the right thing as well.
-	 *
-	 *    Temporary failures, where the same entry should be tried
-	 *    again immediately, almost never happen for this shrinker.
-	 *    We don't do any trylocking; -ENOMEM comes closest,
-	 *    but that's extremely rare and doesn't happen spuriously
-	 *    either. Don't bother distinguishing this case.
+	 * Rotate the entry before dropping the lock: if it gets
+	 * written back or invalidated the free path will unlink
+	 * it from the LRU.  For failures rotation is fine.
 	 */
 	list_move_tail(item, &l->list);
 
 	/*
-	 * Once the lru lock is dropped, the entry might get freed. The
-	 * swpentry is copied to the stack, and entry isn't deref'd again
-	 * until the entry is verified to still be alive in the tree.
+	 * The swpentry is copied to the stack so we can pin the swap
+	 * device and validate the entry's swap table Pointer under
+	 * ci->lock in zswap_writeback_entry().
 	 */
 	swpentry = entry->swpentry;
 
@@ -1420,17 +1399,13 @@ static bool zswap_store_page(struct page *page,
 	if (!zswap_compress(page, entry, pool))
 		goto compress_failed;
 
-	old = xa_store(swap_zswap_tree(page_swpentry),
-		       swp_offset(page_swpentry),
-		       entry, GFP_KERNEL);
-	if (xa_is_err(old)) {
-		int err = xa_err(old);
-
-		WARN_ONCE(err != -ENOMEM, "unexpected xarray error: %d\n", err);
-		zswap_reject_alloc_fail++;
-		goto store_failed;
-	}
-
+	/*
+	 * Store the zswap_entry pointer in the swap table as a
+	 * Pointer-type entry, replacing the PFN or Shadow entry.
+	 * The full swap table entry value (including working set
+	 * shadow) is preserved in entry->swp_tb_val.
+	 * Folio lock serializes concurrent stores.
+	 */
 	si = __swap_entry_to_info(page_swpentry);
 	ci = __swap_offset_to_cluster(si, offset);
 	spin_lock(&ci->lock);
@@ -1465,7 +1440,7 @@ static bool zswap_store_page(struct page *page,
 	 * The entry is successfully compressed and stored in the swap table,
 	 * there is no further possibility of failure. Grab refs to the pool
 	 * and objcg, charge zswap memory, and increment zswap_stored_pages.
-	 * The opposite actions will be performed by zswap_entry_put()
+	 * The opposite actions will be performed by zswap_entry_free()
 	 * when the entry is removed from the swap table.
 	 */
 	zswap_pool_get(pool);
@@ -1498,8 +1473,6 @@ static bool zswap_store_page(struct page *page,
 
 	return true;
 
-store_failed:
-	zs_free(pool->zs_pool, entry->handle);
 compress_failed:
 	zswap_entry_cache_free(entry);
 	return false;
@@ -1578,26 +1551,33 @@ bool zswap_store(struct folio *folio)
 		unsigned type = swp_type(swp);
 		pgoff_t offset = swp_offset(swp);
 		struct zswap_entry *entry;
-		struct xarray *tree;
 		struct swap_cluster_info *ci;
 		struct swap_info_struct *sis;
 		unsigned long swp_tb;
 
 		sis = __swap_type_to_info(type);
 		for (index = 0; index < nr_pages; ++index) {
-			tree = swap_zswap_tree(swp_entry(type, offset + index));
-			entry = xa_erase(tree, offset + index);
-			if (!entry)
-				continue;
 			ci = __swap_offset_to_cluster(sis, offset + index);
 			spin_lock(&ci->lock);
 			swp_tb = __swap_table_get(ci, (offset + index) % SWAPFILE_CLUSTER);
-			if (swp_tb_is_pointer(swp_tb) &&
-			    swp_tb_to_pointer(swp_tb) == entry)
+			if (swp_tb_is_pointer(swp_tb)) {
+				entry = swp_tb_to_pointer(swp_tb);
+				/*
+				 * Restore the original swap table entry
+				 * (Shadow or PFN) that was saved when the
+				 * Pointer entry was created, rather than
+				 * zeroing the slot.  This preserves the swap
+				 * count and flags so the swap cache folio
+				 * can be correctly cleaned up later.
+				 */
 				__swap_table_set(ci, (offset + index) % SWAPFILE_CLUSTER,
 						 entry->swp_tb_val);
+			} else {
+				entry = NULL;
+			}
 			spin_unlock(&ci->lock);
-			zswap_entry_free(entry);
+			if (entry)
+				zswap_entry_free(entry);
 		}
 	}
 
@@ -1607,6 +1587,11 @@ bool zswap_store(struct folio *folio)
 /**
  * zswap_load() - load a folio from zswap
  * @folio: folio to load
+ * @zentry: zswap_entry pointer extracted from swap table, or NULL
+ *
+ * The zswap_entry was extracted from the swap table Pointer entry by
+ * __swap_cache_do_add_folio before it overwrote the entry with a PFN.
+ * If @zentry is NULL, the slot was not backed by zswap.
  *
  * Return: 0 on success, with the folio unlocked and marked up-to-date, or one
  * of the following error codes:
@@ -1624,12 +1609,9 @@ bool zswap_store(struct folio *folio)
  *  -ENOENT: if the swapped out content was not in zswap. The folio remains
  *  locked on return.
  */
-int zswap_load(struct folio *folio)
+int zswap_load(struct folio *folio, void *zentry)
 {
-	swp_entry_t swp = folio->swap;
-	pgoff_t offset = swp_offset(swp);
-	struct xarray *tree = swap_zswap_tree(swp);
-	struct zswap_entry *entry;
+	struct zswap_entry *entry = zentry;
 
 	VM_WARN_ON_ONCE(!folio_test_locked(folio));
 	VM_WARN_ON_ONCE(!folio_test_swapcache(folio));
@@ -1647,7 +1629,6 @@ int zswap_load(struct folio *folio)
 		return -EINVAL;
 	}
 
-	entry = xa_load(tree, offset);
 	if (!entry)
 		return -ENOENT;
 
@@ -1670,7 +1651,6 @@ int zswap_load(struct folio *folio)
 	 * compression work.
 	 */
 	folio_mark_dirty(folio);
-	xa_erase(tree, offset);
 	zswap_entry_free(entry);
 
 	folio_unlock(folio);
@@ -1680,68 +1660,42 @@ int zswap_load(struct folio *folio)
 void zswap_invalidate(swp_entry_t swp)
 {
 	pgoff_t offset = swp_offset(swp);
-	struct xarray *tree = swap_zswap_tree(swp);
 	struct swap_cluster_info *ci;
 	struct swap_info_struct *si;
 	unsigned long swp_tb;
 	struct zswap_entry *entry;
 
-	if (xa_empty(tree))
-		return;
-
-	entry = xa_erase(tree, offset);
-	if (!entry)
-		return;
-
-	/*
-	 * Also clear the swap table Pointer entry if present.
-	 * This is needed because zswap_store now writes Pointer
-	 * entries to both the xarray and the swap table.
-	 */
 	si = __swap_type_to_info(swp_type(swp));
 	ci = __swap_offset_to_cluster(si, offset);
-	spin_lock(&ci->lock);
+	/*
+	 * Do a lockless peek first.  If ci->lock is already held by
+	 * the caller (e.g. __swap_cluster_free_entries), and the
+	 * Pointer has already been cleared, we can bail out without
+	 * self-deadlocking.
+	 */
 	swp_tb = __swap_table_get(ci, offset % SWAPFILE_CLUSTER);
-	if (swp_tb_is_pointer(swp_tb) && swp_tb_to_pointer(swp_tb) == entry)
-		__swap_table_set(ci, offset % SWAPFILE_CLUSTER, entry->swp_tb_val);
-	spin_unlock(&ci->lock);
-
-	zswap_entry_free(entry);
-}
-
-/*
- * This is called after the Shadow entry has been written. We replace the
- * just-written Shadow with a Pointer entry pointing to the zswap_entry.
- *
- * Context: cluster lock must be held. Cannot sleep.
- */
-void zswap_try_convert_to_pointer(struct swap_cluster_info *ci,
-				  unsigned int ci_off, int type, pgoff_t offset)
-{
-	struct xarray *tree;
-	struct zswap_entry *entry;
-	unsigned long swp_tb;
-
-	tree = &zswap_trees[type][offset >> ZSWAP_ADDRESS_SPACE_SHIFT];
-	if (xa_empty(tree))
-		return;
-
-	entry = xa_load(tree, offset);
-	if (!entry)
-		return;
-
-	swp_tb = __swap_table_get(ci, ci_off);
-	/* Entry must be a Shadow (just written by __swap_cache_do_del_folio) */
-	if (!swp_tb_is_countable(swp_tb))
+	if (!swp_tb_is_pointer(swp_tb))
 		return;
 
-	/*
-	 * Save the complete original swap table entry value, which includes
-	 * the swap count, zero flag, and working set shadow information.
-	 */
-	entry->swp_tb_val = swp_tb;
+	spin_lock(&ci->lock);
+	swp_tb = __swap_table_get(ci, offset % SWAPFILE_CLUSTER);
+	if (swp_tb_is_pointer(swp_tb)) {
+		entry = swp_tb_to_pointer(swp_tb);
+		/*
+		 * Restore the original swap table entry (Shadow or PFN)
+		 * that was saved in swp_tb_val when the Pointer entry
+		 * was created.  This preserves swap count and flags
+		 * for the subsequent slot-free path.
+		 */
+		__swap_table_set(ci, offset % SWAPFILE_CLUSTER,
+				 entry->swp_tb_val);
+	} else {
+		entry = NULL;
+	}
+	spin_unlock(&ci->lock);
 
-	__swap_table_set(ci, ci_off, pointer_to_swp_tb(entry));
+	if (entry)
+		zswap_entry_free(entry);
 }
 
 /*
@@ -1854,39 +1808,11 @@ void zswap_swp_tb_put_count(unsigned long swp_tb,
 
 int zswap_swapon(int type, unsigned long nr_pages)
 {
-	struct xarray *trees, *tree;
-	unsigned int nr, i;
-
-	nr = DIV_ROUND_UP(nr_pages, ZSWAP_ADDRESS_SPACE_PAGES);
-	trees = kvzalloc_objs(*tree, nr);
-	if (!trees) {
-		pr_err("alloc failed, zswap disabled for swap type %d\n", type);
-		return -ENOMEM;
-	}
-
-	for (i = 0; i < nr; i++)
-		xa_init(trees + i);
-
-	nr_zswap_trees[type] = nr;
-	zswap_trees[type] = trees;
 	return 0;
 }
 
 void zswap_swapoff(int type)
 {
-	struct xarray *trees = zswap_trees[type];
-	unsigned int i;
-
-	if (!trees)
-		return;
-
-	/* try_to_unuse() invalidated all the entries already */
-	for (i = 0; i < nr_zswap_trees[type]; i++)
-		WARN_ON_ONCE(!xa_empty(trees + i));
-
-	kvfree(trees);
-	nr_zswap_trees[type] = 0;
-	zswap_trees[type] = NULL;
 }
 
 /*********************************
-- 
2.54.0