Merge changes I8557b04c,I31610911 into main am: c8e471f5d2

Original change: https://android-review.googlesource.com/c/platform/packages/modules/Virtualization/+/2761473

Change-Id: Iecaaa2b6ee2f896b4077c72646ee1fe8e11b58a6
Signed-off-by: Automerger Merge Worker <android-build-automerger-merge-worker@system.gserviceaccount.com>
diff --git a/libs/libfdt/src/iterators.rs b/libs/libfdt/src/iterators.rs
index 280257b..7bffd8c 100644
--- a/libs/libfdt/src/iterators.rs
+++ b/libs/libfdt/src/iterators.rs
@@ -14,11 +14,41 @@
 
 //! Iterators over cells, and various layers on top of them.
 
+use crate::Fdt;
 use crate::FdtError;
+use crate::FdtNode;
 use crate::{AddrCells, SizeCells};
+use core::ffi::CStr;
 use core::marker::PhantomData;
 use core::{mem::size_of, ops::Range, slice::ChunksExact};
 
+/// Iterator over nodes sharing a same compatible string.
+pub struct CompatibleIterator<'a> {
+    node: FdtNode<'a>,
+    compatible: &'a CStr,
+}
+
+impl<'a> CompatibleIterator<'a> {
+    pub(crate) fn new(fdt: &'a Fdt, compatible: &'a CStr) -> Result<Self, FdtError> {
+        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
+    }
+}
+
 /// Iterator over cells of a DT property.
 #[derive(Debug)]
 pub struct CellIterator<'a> {
@@ -265,3 +295,29 @@
         }
     }
 }
+
+/// Iterator over subnodes
+#[derive(Debug)]
+pub struct SubnodeIterator<'a> {
+    subnode: Option<FdtNode<'a>>,
+}
+
+impl<'a> SubnodeIterator<'a> {
+    pub(crate) fn new(node: &'a FdtNode) -> Result<Self, FdtError> {
+        let subnode = node.first_subnode()?;
+
+        Ok(Self { subnode })
+    }
+}
+
+impl<'a> Iterator for SubnodeIterator<'a> {
+    type Item = FdtNode<'a>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        let res = self.subnode;
+
+        self.subnode = self.subnode.and_then(|node| node.next_subnode().ok()?);
+
+        res
+    }
+}
diff --git a/libs/libfdt/src/lib.rs b/libs/libfdt/src/lib.rs
index 19ce0f7..758df2a 100644
--- a/libs/libfdt/src/lib.rs
+++ b/libs/libfdt/src/lib.rs
@@ -19,7 +19,10 @@
 
 mod iterators;
 
-pub use iterators::{AddressRange, CellIterator, MemRegIterator, RangesIterator, Reg, RegIterator};
+pub use iterators::{
+    AddressRange, CellIterator, CompatibleIterator, MemRegIterator, RangesIterator, Reg,
+    RegIterator, SubnodeIterator,
+};
 
 use core::cmp::max;
 use core::ffi::{c_int, c_void, CStr};
@@ -381,6 +384,25 @@
             .try_into()
             .map_err(|_| FdtError::Internal)
     }
+
+    /// Returns an iterator of subnodes
+    pub fn subnodes(&'a self) -> Result<SubnodeIterator<'a>> {
+        SubnodeIterator::new(self)
+    }
+
+    fn first_subnode(&self) -> Result<Option<Self>> {
+        // SAFETY: Accesses (read-only) are constrained to the DT totalsize.
+        let ret = unsafe { libfdt_bindgen::fdt_first_subnode(self.fdt.as_ptr(), self.offset) };
+
+        Ok(fdt_err_or_option(ret)?.map(|offset| FdtNode { fdt: self.fdt, offset }))
+    }
+
+    fn next_subnode(&self) -> Result<Option<Self>> {
+        // SAFETY: Accesses (read-only) are constrained to the DT totalsize.
+        let ret = unsafe { libfdt_bindgen::fdt_next_subnode(self.fdt.as_ptr(), self.offset) };
+
+        Ok(fdt_err_or_option(ret)?.map(|offset| FdtNode { fdt: self.fdt, offset }))
+    }
 }
 
 /// Mutable FDT node.
@@ -586,33 +608,6 @@
     }
 }
 
-/// 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 libfdt functions.
 #[derive(Debug)]
 #[repr(transparent)]
diff --git a/libs/libfdt/tests/api_test.rs b/libs/libfdt/tests/api_test.rs
index 866a2bb..513ffd8 100644
--- a/libs/libfdt/tests/api_test.rs
+++ b/libs/libfdt/tests/api_test.rs
@@ -87,3 +87,15 @@
     let nested_node = fdt.node(nested_node_path).unwrap().unwrap();
     assert_eq!(nested_node.name().unwrap().to_str().unwrap(), "PowerPC,970@0");
 }
+
+#[test]
+fn node_subnodes() {
+    let data = fs::read(TEST_TREE_WITH_NO_MEMORY_NODE_PATH).unwrap();
+    let fdt = Fdt::from_slice(&data).unwrap();
+    let root = fdt.root().unwrap();
+    let expected: Vec<&str> = vec!["cpus", "randomnode", "chosen"];
+
+    for (node, name) in root.subnodes().unwrap().zip(expected) {
+        assert_eq!(node.name().unwrap().to_str().unwrap(), name);
+    }
+}