libfdt: Add CompatibleIterator

Add an iterator over FDT nodes sharing a compatible string, to be used
by drivers to locate their devices.

Bug: 237249743
Bug: 255521657
Test: atest vmbase_example.integration_test
Change-Id: I012bb6e93c286189d763685e9be7eba6952022f0
diff --git a/libs/libfdt/src/lib.rs b/libs/libfdt/src/lib.rs
index d15f3c4..01f7b36 100644
--- a/libs/libfdt/src/lib.rs
+++ b/libs/libfdt/src/lib.rs
@@ -269,6 +269,7 @@
 }
 
 /// DT node.
+#[derive(Clone, Copy)]
 pub struct FdtNode<'a> {
     fdt: &'a Fdt,
     offset: c_int,
@@ -351,6 +352,23 @@
         self.fdt
     }
 
+    fn next_compatible(self, compatible: &CStr) -> Result<Option<Self>> {
+        // SAFETY - Accesses (read-only) are constrained to the DT totalsize.
+        let ret = unsafe {
+            libfdt_bindgen::fdt_node_offset_by_compatible(
+                self.fdt.as_ptr(),
+                self.offset,
+                compatible.as_ptr(),
+            )
+        };
+
+        match fdt_err(ret) {
+            Ok(offset) => Ok(Some(Self { fdt: self.fdt, offset })),
+            Err(FdtError::NotFound) => Ok(None),
+            Err(e) => Err(e),
+        }
+    }
+
     fn address_cells(&self) -> Result<AddrCells> {
         // SAFETY - Accesses are constrained to the DT totalsize (validated by ctor).
         unsafe { libfdt_bindgen::fdt_address_cells(self.fdt.as_ptr(), self.offset) }
@@ -366,6 +384,33 @@
     }
 }
 
+/// Iterator over nodes sharing a same compatible string.
+pub struct CompatibleIterator<'a> {
+    node: FdtNode<'a>,
+    compatible: &'a CStr,
+}
+
+impl<'a> CompatibleIterator<'a> {
+    fn new(fdt: &'a Fdt, compatible: &'a CStr) -> Result<Self> {
+        let node = fdt.root()?;
+        Ok(Self { node, compatible })
+    }
+}
+
+impl<'a> Iterator for CompatibleIterator<'a> {
+    type Item = FdtNode<'a>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        let next = self.node.next_compatible(self.compatible).ok()?;
+
+        if let Some(node) = next {
+            self.node = node;
+        }
+
+        next
+    }
+}
+
 /// Wrapper around low-level read-only libfdt functions.
 #[repr(transparent)]
 pub struct Fdt {
@@ -412,12 +457,22 @@
         self.node(CStr::from_bytes_with_nul(b"/chosen\0").unwrap())
     }
 
+    /// Get the root node of the tree.
+    pub fn root(&self) -> Result<FdtNode> {
+        self.node(CStr::from_bytes_with_nul(b"/\0").unwrap())
+    }
+
     /// Find a tree node by its full path.
     pub fn node(&self, path: &CStr) -> Result<FdtNode> {
         let offset = self.path_offset(path)?;
         Ok(FdtNode { fdt: self, offset })
     }
 
+    /// Iterate over nodes with a given compatible string.
+    pub fn compatible_nodes<'a>(&'a self, compatible: &'a CStr) -> Result<CompatibleIterator<'a>> {
+        CompatibleIterator::new(self, compatible)
+    }
+
     fn path_offset(&self, path: &CStr) -> Result<c_int> {
         let len = path.to_bytes().len().try_into().map_err(|_| FdtError::BadPath)?;
         // SAFETY - Accesses are constrained to the DT totalsize (validated by ctor) and the
diff --git a/vmbase/example/src/main.rs b/vmbase/example/src/main.rs
index 3d74168..dcff6e1 100644
--- a/vmbase/example/src/main.rs
+++ b/vmbase/example/src/main.rs
@@ -33,6 +33,7 @@
 };
 use alloc::{vec, vec::Vec};
 use buddy_system_allocator::LockedHeap;
+use core::ffi::CStr;
 use libfdt::Fdt;
 use log::{info, LevelFilter};
 use vmbase::{logger, main, println};
@@ -154,6 +155,13 @@
         info!("memory @ {reg:#x?}");
     }
 
+    let compatible = CStr::from_bytes_with_nul(b"ns16550a\0").unwrap();
+
+    for c in reader.compatible_nodes(compatible).unwrap() {
+        let reg = c.reg().unwrap().next().unwrap();
+        info!("node compatible with '{}' at {reg:?}", compatible.to_str().unwrap());
+    }
+
     info!("FDT checks done.");
 }