Re: [PATCH v4 07/10] KVM: arm64: Convert pkvm_mappings to interval tree

[Marc Zyngier <maz@kernel.org>]

On Fri, 09 May 2025 14:17:03 +0100,
Vincent Donnefort <vdonnefort@google.com> wrote:
> 
> From: Quentin Perret <qperret@google.com>
> 
> In preparation for supporting stage-2 huge mappings for np-guest, let's
> convert pgt.pkvm_mappings to an interval tree.
> 
> No functional change intended.
> 
> Suggested-by: Vincent Donnefort <vdonnefort@google.com>
> Signed-off-by: Quentin Perret <qperret@google.com>
> Signed-off-by: Vincent Donnefort <vdonnefort@google.com>
> 
> diff --git a/arch/arm64/include/asm/kvm_pgtable.h b/arch/arm64/include/asm/kvm_pgtable.h
> index 6b9d274052c7..1b43bcd2a679 100644
> --- a/arch/arm64/include/asm/kvm_pgtable.h
> +++ b/arch/arm64/include/asm/kvm_pgtable.h
> @@ -413,7 +413,7 @@ static inline bool kvm_pgtable_walk_lock_held(void)
>   */
>  struct kvm_pgtable {
>  	union {
> -		struct rb_root					pkvm_mappings;
> +		struct rb_root_cached				pkvm_mappings;
>  		struct {
>  			u32					ia_bits;
>  			s8					start_level;
> diff --git a/arch/arm64/include/asm/kvm_pkvm.h b/arch/arm64/include/asm/kvm_pkvm.h
> index d91bfcf2db56..da75d41c948c 100644
> --- a/arch/arm64/include/asm/kvm_pkvm.h
> +++ b/arch/arm64/include/asm/kvm_pkvm.h
> @@ -173,6 +173,7 @@ struct pkvm_mapping {
>  	struct rb_node node;
>  	u64 gfn;
>  	u64 pfn;
> +	u64 __subtree_last;	/* Internal member for interval tree */
>  };
>  
>  int pkvm_pgtable_stage2_init(struct kvm_pgtable *pgt, struct kvm_s2_mmu *mmu,
> diff --git a/arch/arm64/kvm/pkvm.c b/arch/arm64/kvm/pkvm.c
> index 057874bbe3e1..6febddbec69e 100644
> --- a/arch/arm64/kvm/pkvm.c
> +++ b/arch/arm64/kvm/pkvm.c
> @@ -5,6 +5,7 @@
>   */
>  
>  #include <linux/init.h>
> +#include <linux/interval_tree_generic.h>
>  #include <linux/kmemleak.h>
>  #include <linux/kvm_host.h>
>  #include <asm/kvm_mmu.h>
> @@ -256,80 +257,63 @@ static int __init finalize_pkvm(void)
>  }
>  device_initcall_sync(finalize_pkvm);
>  
> -static int cmp_mappings(struct rb_node *node, const struct rb_node *parent)
> +static u64 __pkvm_mapping_start(struct pkvm_mapping *m)
>  {
> -	struct pkvm_mapping *a = rb_entry(node, struct pkvm_mapping, node);
> -	struct pkvm_mapping *b = rb_entry(parent, struct pkvm_mapping, node);
> -
> -	if (a->gfn < b->gfn)
> -		return -1;
> -	if (a->gfn > b->gfn)
> -		return 1;
> -	return 0;
> +	return m->gfn * PAGE_SIZE;
>  }
>  
> -static struct rb_node *find_first_mapping_node(struct rb_root *root, u64 gfn)
> +static u64 __pkvm_mapping_end(struct pkvm_mapping *m)
>  {
> -	struct rb_node *node = root->rb_node, *prev = NULL;
> -	struct pkvm_mapping *mapping;
> -
> -	while (node) {
> -		mapping = rb_entry(node, struct pkvm_mapping, node);
> -		if (mapping->gfn == gfn)
> -			return node;
> -		prev = node;
> -		node = (gfn < mapping->gfn) ? node->rb_left : node->rb_right;
> -	}
> -
> -	return prev;
> +	return (m->gfn + 1) * PAGE_SIZE - 1;
>  }
>  
> -/*
> - * __tmp is updated to rb_next(__tmp) *before* entering the body of the loop to allow freeing
> - * of __map inline.
> - */
> +INTERVAL_TREE_DEFINE(struct pkvm_mapping, node, u64, __subtree_last,
> +		     __pkvm_mapping_start, __pkvm_mapping_end, static,
> +		     pkvm_mapping);
> +
>  #define for_each_mapping_in_range_safe(__pgt, __start, __end, __map)				\
> -	for (struct rb_node *__tmp = find_first_mapping_node(&(__pgt)->pkvm_mappings,		\
> -							     ((__start) >> PAGE_SHIFT));	\
> +	for (struct pkvm_mapping *__tmp = pkvm_mapping_iter_first(&(__pgt)->pkvm_mappings,	\
> +								  __start, __end - 1);		\
>  	     __tmp && ({									\
> -				__map = rb_entry(__tmp, struct pkvm_mapping, node);		\
> -				__tmp = rb_next(__tmp);						\
> +				__map = __tmp;							\
> +				__tmp = pkvm_mapping_iter_next(__map, __start, __end - 1);	\
>  				true;								\
>  		       });									\
> -	    )											\
> -		if (__map->gfn < ((__start) >> PAGE_SHIFT))					\
> -			continue;								\
> -		else if (__map->gfn >= ((__end) >> PAGE_SHIFT))					\
> -			break;									\
> -		else
> +	    )

The removal of the comment worries me a bit. Is this iterator still
safe wrt freeing of the iterator in the loop?

Thanks,

	M.

-- 
Without deviation from the norm, progress is not possible.