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

[Baoquan He <baoquan.he@linux.dev>]

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