[PATCH 7/7] iommu/arm-smmu-v3: Use the new rb tree helpers

[Mostafa Saleh]

Hi Jason,

On Mon, Jun 03, 2024 at 07:31:33PM -0300, Jason Gunthorpe wrote:
> Since v5.12 the rbtree has gained some simplifying helpers aimed at making
> rb tree users write less convoluted boiler plate code. Instead the caller
> provides a single comparison function and the helpers generate the prior
> open-coded stuff.
> 
> Update smmu->streams to use rb_find_add() and rb_find().
> 
> Signed-off-by: Jason Gunthorpe <jgg at nvidia.com>
Reviewed-by: Mostafa Saleh <smostafa at google.com>

> ---
>  drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3.c | 68 ++++++++++-----------
>  1 file changed, 31 insertions(+), 37 deletions(-)
> 
> diff --git a/drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3.c b/drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3.c
> index 7a6c9aac4cd450..25bae0b05a488c 100644
> --- a/drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3.c
> +++ b/drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3.c
> @@ -1667,26 +1667,37 @@ static int arm_smmu_init_l2_strtab(struct arm_smmu_device *smmu, u32 sid)
>  	return 0;
>  }
>  
> +static int arm_smmu_streams_cmp_key(const void *lhs, const struct rb_node *rhs)
> +{
> +	struct arm_smmu_stream *stream_rhs =
> +		rb_entry(rhs, struct arm_smmu_stream, node);
> +	const u32 *sid_lhs = lhs;
> +
> +	if (*sid_lhs < stream_rhs->id)
> +		return -1;
> +	if (*sid_lhs > stream_rhs->id)
> +		return 1;
> +	return 0;
> +}
> +
> +static int arm_smmu_streams_cmp_node(struct rb_node *lhs,
> +				     const struct rb_node *rhs)
> +{
> +	return arm_smmu_streams_cmp_key(
> +		&rb_entry(lhs, struct arm_smmu_stream, node)->id, rhs);
> +}
> +
>  static struct arm_smmu_master *
>  arm_smmu_find_master(struct arm_smmu_device *smmu, u32 sid)
>  {
>  	struct rb_node *node;
> -	struct arm_smmu_stream *stream;
>  
>  	lockdep_assert_held(&smmu->streams_mutex);
>  
> -	node = smmu->streams.rb_node;
> -	while (node) {
> -		stream = rb_entry(node, struct arm_smmu_stream, node);
> -		if (stream->id < sid)
> -			node = node->rb_right;
> -		else if (stream->id > sid)
> -			node = node->rb_left;
> -		else
> -			return stream->master;
> -	}
> -
> -	return NULL;
> +	node = rb_find(&sid, &smmu->streams, arm_smmu_streams_cmp_key);
> +	if (!node)
> +		return NULL;
> +	return rb_entry(node, struct arm_smmu_stream, node)->master;
>  }
>  
>  /* IRQ and event handlers */
> @@ -2795,8 +2806,6 @@ static int arm_smmu_insert_master(struct arm_smmu_device *smmu,
>  {
>  	int i;
>  	int ret = 0;
> -	struct arm_smmu_stream *new_stream, *cur_stream;
> -	struct rb_node **new_node, *parent_node = NULL;
>  	struct iommu_fwspec *fwspec = dev_iommu_fwspec_get(master->dev);
>  
>  	master->streams = kcalloc(fwspec->num_ids, sizeof(*master->streams),
> @@ -2807,9 +2816,9 @@ static int arm_smmu_insert_master(struct arm_smmu_device *smmu,
>  
>  	mutex_lock(&smmu->streams_mutex);
>  	for (i = 0; i < fwspec->num_ids; i++) {
> +		struct arm_smmu_stream *new_stream = &master->streams[i];
>  		u32 sid = fwspec->ids[i];
>  
> -		new_stream = &master->streams[i];
>  		new_stream->id = sid;
>  		new_stream->master = master;
>  
> @@ -2818,28 +2827,13 @@ static int arm_smmu_insert_master(struct arm_smmu_device *smmu,
>  			break;
>  
>  		/* Insert into SID tree */
> -		new_node = &(smmu->streams.rb_node);
> -		while (*new_node) {
> -			cur_stream = rb_entry(*new_node, struct arm_smmu_stream,
> -					      node);
> -			parent_node = *new_node;
> -			if (cur_stream->id > new_stream->id) {
> -				new_node = &((*new_node)->rb_left);
> -			} else if (cur_stream->id < new_stream->id) {
> -				new_node = &((*new_node)->rb_right);
> -			} else {
> -				dev_warn(master->dev,
> -					 "stream %u already in tree\n",
> -					 cur_stream->id);
> -				ret = -EINVAL;
> -				break;
> -			}
> -		}
> -		if (ret)
> +		if (rb_find_add(&new_stream->node, &smmu->streams,
> +				arm_smmu_streams_cmp_node)) {
> +			dev_warn(master->dev, "stream %u already in tree\n",
> +				 sid);
> +			ret = -EINVAL;
>  			break;
> -
> -		rb_link_node(&new_stream->node, parent_node, new_node);
> -		rb_insert_color(&new_stream->node, &smmu->streams);
> +		}
>  	}
>  
>  	if (ret) {
> -- 
> 2.45.2
>