[RFC PATCH v2 08/13] mm: implement anonymous folio rmap for ANON_VMA_FRACTAL

[tao]

Add the ANON_RMAP_FOREACH_VMA() macro to encapsulate rmap traversal and
hide implementation differences from external users.

During reverse mapping of an anonymous folio, traversal first handles the
folio's own anon_node, then iterates through its descendants. Since large
folios may span multiple VMAs, VMA updates are used to advance pgoff, and
pgoff is used to determine completion of the current anon_node traversal.

For KSM, the VMA is located directly using the supplied address.

Signed-off-by: tao <tao.wangtao@xxxxxxxxx>
---
 mm/internal.h | 49 +++++++++++++++++++++++++++++++++++++++++++++++++
 mm/rmap.c     | 41 +++++++++++++++++++++++++++++++++++++++++
 2 files changed, 90 insertions(+)

diff --git a/mm/internal.h b/mm/internal.h
index 970f11b1af50..b38c7b4c2a8f 100644
--- a/mm/internal.h
+++ b/mm/internal.h
@@ -296,6 +296,55 @@ static inline unsigned long anon_node_rmap_count(struct anon_node *anon_nod)
 	return atomic_long_read(&anon_nod->rbc) & ANON_RMAP_BASE_COUNT_MAX;
 }
 
+static inline unsigned long anon_node_rmap_address(struct anon_node *anon_nod,
+		unsigned long pgoff)
+{
+	if (!anon_node_rmap_count(anon_nod))
+		return 0;
+	return anon_node_rmap_base(anon_nod) + (pgoff << PAGE_SHIFT);
+}
+
+static inline struct vm_area_struct *anon_node_lookup_vma(
+	struct anon_node *anon_nod, unsigned long addr, unsigned long pgoff)
+{
+	struct vm_area_struct *vma;
+
+	if (!addr)
+		addr = anon_node_rmap_address(anon_nod, pgoff);
+	if (!anon_nod->mm || !addr)
+		return NULL;
+
+	vma = vma_lookup(anon_nod->mm, addr);
+	return (vma && vma->anon_vma == (void *)anon_nod) ? vma : NULL;
+}
+
+void vm_lock_anon_node(struct mm_struct *mm, struct anon_node *anon_nod);
+void vm_unlock_anon_node(struct anon_node *anon_nod);
+int anon_node_trylock_rmap(struct anon_node *anon_nod);
+void anon_node_lock_rmap(struct anon_node *anon_nod);
+void anon_node_unlock_rmap(struct anon_node *anon_nod);
+
+#define ANON_RMAP_FOREACH_VMA(anon_vma, addr, start, last, rmap_proc_vma) \
+do { \
+	struct vm_area_struct *vma = NULL; /* must be named vma */ \
+	struct anon_node *_anon_nod = (void *)(anon_vma); \
+	struct anon_node *_nod = _anon_nod; \
+	unsigned long _pgoff = (start); \
+	unsigned long _nr = 0, _total = 1; \
+	\
+	while (_nod) { \
+		vma = anon_node_lookup_vma(_nod, (addr), _pgoff); \
+		if (vma) \
+			(rmap_proc_vma); \
+		_pgoff = vma ? vma->vm_pgoff + vma_pages(vma) : _pgoff + 1; \
+		if (!(addr) && _pgoff <= (last)) \
+			continue; \
+		_pgoff = (start); \
+		BUG_ON(++_nr > (_total += _nod->nr_children)); \
+		_nod = anon_node_next_descendant(_anon_nod, _nod); \
+	} \
+} while (0)
+
 #endif
 
 struct anon_vma *folio_get_anon_vma(const struct folio *folio);
diff --git a/mm/rmap.c b/mm/rmap.c
index fd22576b0f58..1551e9d2121d 100644
--- a/mm/rmap.c
+++ b/mm/rmap.c
@@ -982,6 +982,47 @@ void vma_post_update_rmap_base(struct vm_area_struct *vma, unsigned long diff)
 		release_anon_nodes(release_node);
 }
 
+#define ANON_VM_LOCK_REF	VM_REFCNT_EXCLUDE_READERS_FLAG
+
+void vm_lock_anon_node(struct mm_struct *mm, struct anon_node *anon_nod)
+{
+	atomic_t *vm_lock_ref = anon_node_vm_lock_ref(anon_nod);
+	struct rw_semaphore *rmap_sem = anon_node_rmap_sem(anon_nod);
+
+	if (!(atomic_read(vm_lock_ref) & ANON_VM_LOCK_REF)) {
+		down_write_nest_lock(rmap_sem, &mm->mmap_lock);
+		BUG_ON(atomic_read(vm_lock_ref));
+		atomic_set(vm_lock_ref, ANON_VM_LOCK_REF);
+	}
+}
+
+void vm_unlock_anon_node(struct anon_node *anon_nod)
+{
+	atomic_t *vm_lock_ref = anon_node_vm_lock_ref(anon_nod);
+	struct rw_semaphore *rmap_sem = anon_node_rmap_sem(anon_nod);
+
+	if (atomic_read(vm_lock_ref) & ANON_VM_LOCK_REF) {
+		BUG_ON(atomic_read(vm_lock_ref) != ANON_VM_LOCK_REF);
+		atomic_set(vm_lock_ref, 0);
+		up_write(rmap_sem);
+	}
+}
+
+int anon_node_trylock_rmap(struct anon_node *anon_nod)
+{
+	return down_read_trylock(anon_node_rmap_sem(anon_nod));
+}
+
+void anon_node_lock_rmap(struct anon_node *anon_nod)
+{
+	down_read(anon_node_rmap_sem(anon_nod));
+}
+
+void anon_node_unlock_rmap(struct anon_node *anon_nod)
+{
+	up_read(anon_node_rmap_sem(anon_nod));
+}
+
 #endif
 
 /*
-- 
2.17.1