diff --git a/src/model/overlay.rs b/src/model/overlay.rs index a073ce1..dfdd81e 100644 --- a/src/model/overlay.rs +++ b/src/model/overlay.rs @@ -8,16 +8,173 @@ //! Implementation of applying the Device Tree Overlays (DTBO). +use alloc::borrow::ToOwned; use alloc::collections::BTreeMap; use alloc::format; -use alloc::string::String; +use alloc::string::{String, ToString}; +use alloc::vec::Vec; use crate::error::OverlayError; -use crate::model::DeviceTreeNode; -use crate::overlay::{PHANDLE_PROPS, get_phandle}; +use crate::fdt::Fdt; +use crate::model::{DeviceTree, DeviceTreeNode}; +use crate::overlay::{Fragment, FragmentTarget, NODE_SYMBOLS, Overlay, PHANDLE_PROPS, get_phandle}; use crate::{Node, Property}; +/// A utility struct for applying device tree overlays efficiently. +/// +/// It caches phandle mappings to avoid repeatedly traversing the entire base +/// tree. +#[derive(Debug)] #[allow(unused)] +pub struct OverlayApplier<'a> { + base: &'a mut DeviceTree, + phandles: BTreeMap, + max_phandle: u32, +} + +impl<'a> OverlayApplier<'a> { + /// Creates a new `OverlayApplier` for the given base device tree. + pub fn new(base: &'a mut DeviceTree) -> Self { + let mut phandles = BTreeMap::new(); + let mut max_ph = 0; + let mut current_path = String::new(); + Self::build_cache(&base.root, &mut current_path, &mut phandles, &mut max_ph); + Self { + base, + phandles, + max_phandle: max_ph, + } + } + + fn build_cache( + node: &DeviceTreeNode, + current_path: &mut String, + phandles: &mut BTreeMap, + max_ph: &mut u32, + ) { + if let Some(p) = get_phandle(node) { + *max_ph = (*max_ph).max(p); + let p_str = if current_path.is_empty() { + "/".to_string() + } else { + current_path.clone() + }; + phandles.insert(p, p_str); + } + for child in node.children() { + let old_len = current_path.len(); + if !current_path.ends_with('/') { + current_path.push('/'); + } + current_path.push_str(child.name()); + Self::build_cache(child, current_path, phandles, max_ph); + current_path.truncate(old_len); + } + } + + /// Applies a read-only Device Tree Overlay ([`Fdt`]) to the base + /// [`DeviceTree`]. + /// + /// # Errors + /// + /// Returns an error when applying the overlay fails because of a malformed + /// overlay data. + pub fn apply_overlay(&mut self, overlay: &Fdt<'_>) -> Result<(), OverlayError> { + let overlay_tree = DeviceTree::from_fdt(overlay); + self.apply_overlay_tree(overlay_tree) + } + + /// Applies an in-memory Device Tree Overlay ([`DeviceTree`]) to the base + /// [`DeviceTree`]. + /// + /// # Errors + /// + /// Returns an error when applying the overlay fails because of a malformed + /// overlay data. + pub fn apply_overlay_tree(&mut self, mut overlay: DeviceTree) -> Result<(), OverlayError> { + let overlay_view = Overlay { + node: &overlay.root, + }; + let mut fragment_targets = Vec::new(); + + for frag in overlay_view.fragments() { + let frag_name = frag.name().to_string(); + let target_path = self.resolve_fragment_target_path(frag)?; + fragment_targets.push((frag_name, target_path)); + } + + for (frag_name, target_path) in &fragment_targets { + self.merge_fragment(&mut overlay, frag_name, target_path)?; + } + + Ok(()) + } + + fn resolve_fragment_target_path( + &self, + fragment: Fragment<&DeviceTreeNode>, + ) -> Result { + match fragment.target()? { + FragmentTarget::Phandle(phandle) => { + if let Some(path) = self.phandles.get(&phandle) { + Ok(path.clone()) + } else { + Err(OverlayError::TargetNotFound(format!( + "phandle 0x{phandle:x}" + ))) + } + } + FragmentTarget::Path(path) => { + if path.starts_with('/') { + self.base + .find_node(path) + .map(|_| path.to_string()) + .ok_or_else(|| OverlayError::TargetNotFound(path.to_string())) + } else { + if let Some(sym_node) = self.base.root.child(NODE_SYMBOLS) + && let Some(sym_prop) = sym_node.property(path) + { + let abs_path = sym_prop.as_str()?; + return Ok(abs_path.to_string()); + } + if let Some(aliases_node) = self.base.root.child("aliases") + && let Some(alias_prop) = aliases_node.property(path) + { + let abs_path = alias_prop.as_str()?; + return Ok(abs_path.to_string()); + } + Err(OverlayError::TargetNotFound(path.to_string())) + } + } + } + } + + fn merge_fragment( + &mut self, + overlay: &mut DeviceTree, + frag_name: &str, + target_path: &str, + ) -> Result<(), OverlayError> { + if let Some(mut frag_node) = overlay.root.remove_child(frag_name) + && let Some(overlay_subnode) = frag_node.remove_child("__overlay__") + { + let target_node = self + .base + .find_node_mut(target_path) + .ok_or_else(|| OverlayError::TargetNotFound(target_path.to_owned()))?; + merge_node( + target_node, + overlay_subnode, + target_path, + &mut self.phandles, + &mut self.max_phandle, + ); + } + + Ok(()) + } +} + fn merge_node( target: &mut DeviceTreeNode, overlay: DeviceTreeNode, @@ -46,7 +203,6 @@ fn merge_node( } } -#[allow(unused)] fn update_phandles_recursive( node: &DeviceTreeNode, current_path: &mut String, diff --git a/tests/overlay.rs b/tests/overlay.rs new file mode 100644 index 0000000..ddf89d3 --- /dev/null +++ b/tests/overlay.rs @@ -0,0 +1,57 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 or the MIT license +// , at your +// option. This file may not be copied, modified, or distributed +// except according to those terms. + +#![cfg(feature = "write")] + +use dtoolkit::fdt::Fdt; +use dtoolkit::model::DeviceTree; +use dtoolkit::model::overlay::OverlayApplier; +use dtoolkit::{Node, Property}; + +macro_rules! apply_overlay { + ($base_file:expr, $overlay_file:expr) => {{ + let base_dtb = include_bytes!(concat!("dtb/", $base_file)); + let base_fdt = Fdt::new(base_dtb).unwrap(); + let mut base = DeviceTree::from_fdt(&base_fdt); + + let overlay_dtb = include_bytes!(concat!("dtb/", $overlay_file)); + let overlay_fdt = Fdt::new(overlay_dtb).unwrap(); + + let mut applier = OverlayApplier::new(&mut base); + applier.apply_overlay(&overlay_fdt).unwrap(); + base + }}; +} + +#[test] +fn apply_overlay_target_path() { + let mut base = apply_overlay!( + "overlay_target_path_base.dtb", + "overlay_target_path_overlay.dtb" + ); + + let soc = base.find_node_mut("/soc").unwrap(); + assert_eq!(soc.property("status").unwrap().as_str(), Ok("okay")); + assert_eq!(soc.property("new-prop").unwrap().as_str(), Ok("foo")); + assert!(soc.child("serial@1000").is_some()); + + // Verify round-trip through serialization + let dtb = base.to_dtb(); + let fdt = Fdt::new(&dtb).unwrap(); + let fdt_soc = fdt.find_node("/soc").unwrap(); + assert_eq!(fdt_soc.property("status").unwrap().as_str(), Ok("okay")); +} + +#[test] +fn apply_overlay_dependent_fragments() { + // Just verify it applies without error + let _base = apply_overlay!( + "overlay_dependent_base.dtb", + "overlay_dependent_overlay.dtb" + ); +}