[PATCH RFC v2 1/4] iommu/arm-smmu-v3: Convert streams from RB tree to XArray

["Peng Fan (OSS)" <peng.fan@oss.nxp.com>]

From: Peng Fan <peng.fan@nxp.com>

Replace the smmu->streams RB tree with an XArray for SID ->
arm_smmu_stream lookups. The existing streams_mutex serialises all
accesses (both xa_store/xa_erase and xa_load), protecting the lifetime
of returned pointers against concurrent arm_smmu_remove_master()
without requiring RCU grace periods. A mutex (rather than xa_lock) is
needed because several paths sleep while the lock is held:
dmam_alloc_coherent(GFP_KERNEL) in arm_smmu_init_sid_strtab(), and
down_read() inside iommu_report_device_fault().

This removes the RB tree comparators, the rb_node from
arm_smmu_stream, and simplifies duplicate-SID handling for bridged PCI
devices.

No behavioural change intended; preparation for shared-SID support.

Assisted-by: Claude:claude-opus-4-6
Signed-off-by: Peng Fan <peng.fan@nxp.com>
---
 drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3.c | 83 ++++++++++++++---------------
 drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3.h | 11 +++-
 2 files changed, 48 insertions(+), 46 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 5732f3ba0122d..65e448a69a019 100644
--- a/drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3.c
+++ b/drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3.c
@@ -2024,37 +2024,17 @@ 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 = rb_find(&sid, &smmu->streams, arm_smmu_streams_cmp_key);
-	if (!node)
+	stream = xa_load(&smmu->streams, sid);
+	if (!stream)
 		return NULL;
-	return rb_entry(node, struct arm_smmu_stream, node)->master;
+	return stream->master;
 }
 
 /* IRQ and event handlers */
@@ -4123,38 +4103,47 @@ static int arm_smmu_insert_master(struct arm_smmu_device *smmu,
 	     NULL);
 
 	mutex_lock(&smmu->streams_mutex);
-	for (i = 0; i < fwspec->num_ids; i++) {
+	for (i = 0; i < master->num_streams; i++) {
 		struct arm_smmu_stream *new_stream = &master->streams[i];
-		struct rb_node *existing;
+		struct arm_smmu_stream *existing;
 		u32 sid = new_stream->id;
 
 		ret = arm_smmu_init_sid_strtab(smmu, sid);
 		if (ret)
 			break;
 
-		/* Insert into SID tree */
-		existing = rb_find_add(&new_stream->node, &smmu->streams,
-				       arm_smmu_streams_cmp_node);
-		if (existing) {
-			struct arm_smmu_master *existing_master =
-				rb_entry(existing, struct arm_smmu_stream, node)
-					->master;
-
-			/* Bridged PCI devices may end up with duplicated IDs */
-			if (existing_master == master)
-				continue;
+		/* Bridged PCI devices may end up with duplicated IDs */
+		if (i > 0 && master->streams[i - 1].id == sid)
+			continue;
 
+		existing = xa_load(&smmu->streams, sid);
+		if (existing) {
 			dev_warn(master->dev,
 				 "Aliasing StreamID 0x%x (from %s) unsupported, expect DMA to be broken\n",
-				 sid, dev_name(existing_master->dev));
+				 sid, dev_name(existing->master->dev));
 			ret = -ENODEV;
 			break;
 		}
+
+		/*
+		 * xa_store() returns the old entry (void *) on success
+		 * or an ERR_PTR on allocation failure.  Use xa_err() to
+		 * convert to a standard errno.
+		 */
+		ret = xa_err(xa_store(&smmu->streams, sid, new_stream,
+				      GFP_KERNEL));
+		if (ret)
+			break;
 	}
 
 	if (ret) {
-		for (i--; i >= 0; i--)
-			rb_erase(&master->streams[i].node, &smmu->streams);
+		for (i--; i >= 0; i--) {
+			u32 sid = master->streams[i].id;
+
+			if (i > 0 && master->streams[i - 1].id == sid)
+				continue;
+			xa_erase(&smmu->streams, sid);
+		}
 		kfree(master->streams);
 		kfree(master->build_invs);
 	}
@@ -4167,14 +4156,19 @@ static void arm_smmu_remove_master(struct arm_smmu_master *master)
 {
 	int i;
 	struct arm_smmu_device *smmu = master->smmu;
-	struct iommu_fwspec *fwspec = dev_iommu_fwspec_get(master->dev);
 
 	if (!smmu || !master->streams)
 		return;
 
 	mutex_lock(&smmu->streams_mutex);
-	for (i = 0; i < fwspec->num_ids; i++)
-		rb_erase(&master->streams[i].node, &smmu->streams);
+	for (i = 0; i < master->num_streams; i++) {
+		u32 sid = master->streams[i].id;
+
+		/* Skip intra-master duplicate SIDs */
+		if (i > 0 && master->streams[i - 1].id == sid)
+			continue;
+		xa_erase(&smmu->streams, sid);
+	}
 	mutex_unlock(&smmu->streams_mutex);
 
 	kfree(master->streams);
@@ -4602,7 +4596,7 @@ static int arm_smmu_init_structures(struct arm_smmu_device *smmu)
 	int ret;
 
 	mutex_init(&smmu->streams_mutex);
-	smmu->streams = RB_ROOT;
+	xa_init(&smmu->streams);
 
 	ret = arm_smmu_init_queues(smmu);
 	if (ret)
@@ -5627,6 +5621,7 @@ static void arm_smmu_device_remove(struct platform_device *pdev)
 
 	iommu_device_unregister(&smmu->iommu);
 	iommu_device_sysfs_remove(&smmu->iommu);
+	xa_destroy(&smmu->streams);
 }
 
 static void arm_smmu_device_shutdown(struct platform_device *pdev)
diff --git a/drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3.h b/drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3.h
index dd2fee2f560e6..97dc97ac704d9 100644
--- a/drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3.h
+++ b/drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3.h
@@ -14,6 +14,7 @@
 #include <linux/kernel.h>
 #include <linux/mmzone.h>
 #include <linux/sizes.h>
+#include <linux/xarray.h>
 
 struct arm_smmu_device;
 struct arm_vsmmu;
@@ -961,14 +962,20 @@ struct arm_smmu_device {
 	/* IOMMU core code handle */
 	struct iommu_device		iommu;
 
-	struct rb_root			streams;
+	/*
+	 * XArray of arm_smmu_stream, indexed by SID.
+	 * All accesses (reads and writes) are serialised by streams_mutex.
+	 * The mutex is held across xa_load() and all subsequent uses of the
+	 * returned pointer to prevent use-after-free from concurrent
+	 * arm_smmu_remove_master().
+	 */
+	struct xarray			streams;
 	struct mutex			streams_mutex;
 };
 
 struct arm_smmu_stream {
 	u32				id;
 	struct arm_smmu_master		*master;
-	struct rb_node			node;
 };
 
 struct arm_smmu_vmaster {

-- 
2.34.1