[PATCH 01/15] mm/rmap: introduce anon_rmap APIs for anonymous folios

[tao]

Add a set of anon_rmap APIs to operate on the reverse mappings of
anonymous folios.

Introduce anon_rmap_for_each_vma() as a wrapper around
vma_interval_tree_foreach(), so callers no longer access the
interval tree directly.

This prepares the rmap code for upcoming ANON_VMA_LAZY support and
RCU-based lockless rmap traversal.

No functional change intended.

Signed-off-by: tao <tao.wangtao@xxxxxxxxx>
---
 include/linux/rmap.h | 68 +++++++++++++++++++++++++++++++++++++++++
 mm/rmap.c            | 73 ++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 141 insertions(+)

diff --git a/include/linux/rmap.h b/include/linux/rmap.h
index 8dc0871e5f00..c42314ea4362 100644
--- a/include/linux/rmap.h
+++ b/include/linux/rmap.h
@@ -937,6 +937,44 @@ int pfn_mkclean_range(unsigned long pfn, unsigned long nr_pages, pgoff_t pgoff,
 void remove_migration_ptes(struct folio *src, struct folio *dst,
 		enum ttu_flags flags);
 
+/* Reverse mapping handle for anonymous folio rmap helpers. */
+typedef struct anon_rmap {
+	unsigned long rmap;
+} anon_rmap_t;
+
+#define ANON_RMAP_NULL make_anon_rmap(0)
+
+static inline anon_rmap_t make_anon_rmap(const void *anon_mapping)
+{
+	return (anon_rmap_t){ .rmap = (unsigned long)anon_mapping, };
+}
+
+static inline unsigned long anon_rmap_value(anon_rmap_t anon_rmap)
+{
+	return anon_rmap.rmap;
+}
+
+static inline anon_rmap_t anon_vma_to_anon_rmap(const struct anon_vma *anon_vma)
+{
+	return make_anon_rmap(anon_vma);
+}
+
+static inline struct anon_vma *anon_rmap_to_anon_vma(anon_rmap_t anon_rmap)
+{
+	unsigned long rmap = anon_rmap_value(anon_rmap);
+
+	return (struct anon_vma *)rmap;
+}
+
+anon_rmap_t vma_get_anon_rmap(struct vm_area_struct *vma);
+void put_anon_rmap(anon_rmap_t anon_rmap);
+void anon_rmap_lock_write(anon_rmap_t anon_rmap);
+int anon_rmap_trylock_write(anon_rmap_t anon_rmap);
+void anon_rmap_unlock_write(anon_rmap_t anon_rmap);
+void anon_rmap_lock_read(anon_rmap_t anon_rmap);
+int anon_rmap_trylock_read(anon_rmap_t anon_rmap);
+void anon_rmap_unlock_read(anon_rmap_t anon_rmap);
+
 /*
  * rmap_walk_control: To control rmap traversing for specific needs
  *
@@ -969,6 +1007,36 @@ void rmap_walk_locked(struct folio *folio, struct rmap_walk_control *rwc);
 struct anon_vma *folio_lock_anon_vma_read(const struct folio *folio,
 					  struct rmap_walk_control *rwc);
 
+bool folio_maybe_same_anon_vma(const struct folio *folio,
+	const struct vm_area_struct *vma);
+anon_rmap_t folio_get_anon_rmap(const struct folio *folio);
+anon_rmap_t folio_lock_anon_rmap_read(const struct folio *folio,
+	struct rmap_walk_control *rwc);
+
+static inline struct vm_area_struct *anon_rmap_iter_first_vma(
+	anon_rmap_t anon_rmap, unsigned long start, unsigned long last,
+	struct anon_vma_chain **avc)
+{
+	struct anon_vma *anon_vma = anon_rmap_to_anon_vma(anon_rmap);
+
+	*avc = anon_vma_interval_tree_iter_first(&anon_vma->rb_root, start, last);
+	return *avc ? (*avc)->vma : NULL;
+}
+
+static inline struct vm_area_struct *anon_rmap_iter_next_vma(
+	anon_rmap_t anon_rmap, unsigned long start, unsigned long last,
+	struct anon_vma_chain **avc)
+{
+	if (!*avc)
+		return NULL;
+	*avc = anon_vma_interval_tree_iter_next(*avc, start, last);
+	return *avc ? (*avc)->vma : NULL;
+}
+
+#define anon_rmap_foreach_vma(vma, avc, anon_rmap, start, last) 	   \
+	for (vma = anon_rmap_iter_first_vma(anon_rmap, start, last, &avc); \
+	     vma; vma = anon_rmap_iter_next_vma(anon_rmap, start, last, &avc))
+
 #else	/* !CONFIG_MMU */
 
 #define anon_vma_init()		do {} while (0)
diff --git a/mm/rmap.c b/mm/rmap.c
index 78b7fb5f367c..1b2dada71778 100644
--- a/mm/rmap.c
+++ b/mm/rmap.c
@@ -701,6 +701,79 @@ struct anon_vma *folio_lock_anon_vma_read(const struct folio *folio,
 	return anon_vma;
 }
 
+anon_rmap_t vma_get_anon_rmap(struct vm_area_struct *vma)
+{
+	mmap_assert_locked(vma->vm_mm);
+	VM_BUG_ON(!vma->anon_vma);
+	get_anon_vma(vma->anon_vma);
+	return anon_vma_to_anon_rmap(vma->anon_vma);
+}
+
+void put_anon_rmap(anon_rmap_t anon_rmap)
+{
+	put_anon_vma(anon_rmap_to_anon_vma(anon_rmap));
+}
+
+void anon_rmap_lock_write(anon_rmap_t anon_rmap)
+{
+	anon_vma_lock_write(anon_rmap_to_anon_vma(anon_rmap));
+}
+
+int anon_rmap_trylock_write(anon_rmap_t anon_rmap)
+{
+	return anon_vma_trylock_write(anon_rmap_to_anon_vma(anon_rmap));
+}
+
+void anon_rmap_unlock_write(anon_rmap_t anon_rmap)
+{
+	anon_vma_unlock_write(anon_rmap_to_anon_vma(anon_rmap));
+}
+
+void anon_rmap_lock_read(anon_rmap_t anon_rmap)
+{
+	anon_vma_lock_read(anon_rmap_to_anon_vma(anon_rmap));
+}
+
+int anon_rmap_trylock_read(anon_rmap_t anon_rmap)
+{
+	return anon_vma_trylock_read(anon_rmap_to_anon_vma(anon_rmap));
+}
+
+void anon_rmap_unlock_read(anon_rmap_t anon_rmap)
+{
+	anon_vma_unlock_read(anon_rmap_to_anon_vma(anon_rmap));
+}
+
+bool folio_maybe_same_anon_vma(const struct folio *folio,
+	const struct vm_area_struct *vma)
+{
+	struct anon_vma *anon_vma;
+	struct anon_vma *tgt_anon_vma = vma->anon_vma;
+	bool same = false;
+
+	rcu_read_lock();
+	anon_vma = folio_anon_vma(folio);
+	if (anon_vma && tgt_anon_vma)
+		same = anon_vma->root == tgt_anon_vma->root;
+	rcu_read_unlock();
+	return same;
+}
+
+anon_rmap_t folio_get_anon_rmap(const struct folio *folio)
+{
+	struct anon_vma *anon_vma = folio_get_anon_vma(folio);
+
+	return anon_vma ? anon_vma_to_anon_rmap(anon_vma) : ANON_RMAP_NULL;
+}
+
+anon_rmap_t folio_lock_anon_rmap_read(const struct folio *folio,
+				      struct rmap_walk_control *rwc)
+{
+	struct anon_vma *anon_vma = folio_lock_anon_vma_read(folio, rwc);
+
+	return anon_vma ? anon_vma_to_anon_rmap(anon_vma) : ANON_RMAP_NULL;
+}
+
 #ifdef CONFIG_ARCH_WANT_BATCHED_UNMAP_TLB_FLUSH
 /*
  * Flush TLB entries for recently unmapped pages from remote CPUs. It is
-- 
2.17.1