moa_fdt 0.1.4

// Based on fdt 0.1.5 (MPL-2.0) by repnop, with modifications:
// - Single lifetime 'a (no header back-reference)
// - FDT_END = 0x9 (per DTSpec)
// - Non-recursive skip_current_node
// - MemoryRegion.address is usize, not *const u8

use crate::parsing::{CStr, FdtData, skip_4_aligned};

pub(crate) const FDT_BEGIN_NODE: u32 = 0x1;
pub(crate) const FDT_END_NODE: u32 = 0x2;
pub(crate) const FDT_PROP: u32 = 0x3;
pub(crate) const FDT_NOP: u32 = 0x4;
pub(crate) const FDT_END: u32 = 0x9;

const MAX_DEPTH: usize = 12;

/// 设备树节点（零拷贝引用）
#[derive(Debug, Clone, Copy)]
pub struct FdtNode<'a> {
    /// 节点名（如 `"memory@80000000"`，根节点为 `"/"`）
    pub name: &'a str,
    props: &'a [u8],
    parent_props: Option<&'a [u8]>,
    strings: &'a [u8],
}

impl<'a> FdtNode<'a> {
    #[inline(always)]
    pub(crate) fn new(
        name: &'a str,
        props: &'a [u8],
        parent_props: Option<&'a [u8]>,
        strings: &'a [u8],
    ) -> Self {
        Self { name, props, parent_props, strings }
    }

    /// 遍历节点的所有属性
    #[inline(always)]
    pub fn properties(self) -> impl Iterator<Item = NodeProperty<'a>> {
        let strings = self.strings;
        let mut stream = FdtData::new(self.props);
        let mut done = false;

        core::iter::from_fn(move || {
            if stream.is_empty() || done {
                return None;
            }

            while stream.peek_u32()? == FDT_NOP {
                stream.skip(4);
            }

            if stream.peek_u32()? == FDT_PROP {
                NodeProperty::parse(&mut stream, strings)
            } else {
                done = true;
                None
            }
        })
    }

    /// 按名称查找属性
    #[inline(always)]
    pub fn property(self, name: &str) -> Option<NodeProperty<'a>> {
        self.properties().find(|p| p.name == name)
    }

    /// 遍历直接子节点
    #[moa_sec_macros::init]
    pub fn children(self) -> impl Iterator<Item = FdtNode<'a>> {
        let strings = self.strings;
        let mut stream = FdtData::new(self.props);

        while stream.peek_u32() == Some(FDT_NOP) {
            stream.skip(4);
        }

        while stream.peek_u32() == Some(FDT_PROP) {
            if NodeProperty::parse(&mut stream, strings).is_none() {
                break;
            }
        }

        let parent_props = self.props;
        let mut done = false;

        core::iter::from_fn(move || {
            if stream.is_empty() || done {
                return None;
            }

            while stream.peek_u32()? == FDT_NOP {
                stream.skip(4);
            }

            if stream.peek_u32()? == FDT_BEGIN_NODE {
                let origin = stream.remaining();
                let ret = {
                    stream.skip(4);
                    let unit_name = CStr::new(stream.remaining())?.as_str()?;
                    let full_name_len = unit_name.len() + 1;
                    skip_4_aligned(&mut stream, full_name_len);

                    Some(FdtNode::new(unit_name, stream.remaining(), Some(parent_props), strings))
                };

                stream = FdtData::new(origin);
                skip_current_node(&mut stream);

                ret
            } else {
                done = true;
                None
            }
        })
    }

    /// 此节点声明的 cell sizes（用于解析子节点的 reg）
    #[moa_sec_macros::init]
    pub fn cell_sizes(self) -> CellSizes {
        let mut cell_sizes = CellSizes::default();
        let mut found = 0u8;
        for property in self.properties() {
            match property.name {
                "#address-cells" => {
                    cell_sizes.address_cells = u32::from_be_bytes(
                        property.value.get(..4).and_then(|b| b.try_into().ok()).unwrap_or([0; 4]),
                    ) as usize;
                    found |= 1;
                },
                "#size-cells" => {
                    cell_sizes.size_cells = u32::from_be_bytes(
                        property.value.get(..4).and_then(|b| b.try_into().ok()).unwrap_or([0; 4]),
                    ) as usize;
                    found |= 2;
                },
                _ => {},
            }
            if found == 3 {
                break;
            }
        }
        cell_sizes
    }

    /// 父节点声明的 cell sizes（用于解析本节点的 reg）
    #[inline(always)]
    pub fn parent_cell_sizes(self) -> CellSizes {
        match self.parent_props {
            Some(parent) => FdtNode::new("", parent, None, self.strings).cell_sizes(),
            None => CellSizes::default(),
        }
    }

    /// `reg` 属性
    #[moa_sec_macros::init]
    pub fn reg(self) -> Option<RegIter<'a>> {
        let sizes = self.parent_cell_sizes();
        if sizes.address_cells > 2 || sizes.size_cells > 2 {
            return None;
        }
        let prop = self.property("reg")?;
        Some(RegIter { stream: FdtData::new(prop.value), sizes })
    }

    /// `compatible` 属性
    #[inline(always)]
    pub fn compatible(self) -> Option<Compatible<'a>> {
        self.property("compatible").map(|p| Compatible { data: p.value })
    }
}

/// `#address-cells` / `#size-cells` 配置
#[derive(Debug, Clone, Copy)]
pub struct CellSizes {
    /// 地址单元数
    pub address_cells: usize,
    /// 大小单元数
    pub size_cells: usize,
}

impl Default for CellSizes {
    #[inline(always)]
    fn default() -> Self {
        Self { address_cells: 2, size_cells: 1 }
    }
}

/// 内存区域
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct MemoryRegion {
    /// 起始物理地址
    pub address: usize,
    /// 区域大小
    pub size: Option<usize>,
}

/// `reg` 属性迭代器
pub struct RegIter<'a> {
    stream: FdtData<'a>,
    sizes: CellSizes,
}

#[allow(clippy::cast_possible_truncation)]
impl Iterator for RegIter<'_> {
    type Item = MemoryRegion;

    #[moa_sec_macros::init]
    fn next(&mut self) -> Option<Self::Item> {
        let address = match self.sizes.address_cells {
            1 => self.stream.u32()? as usize,
            2 => self.stream.u64()? as usize,
            _ => return None,
        };
        let size = match self.sizes.size_cells {
            0 => None,
            1 => Some(self.stream.u32()? as usize),
            2 => Some(self.stream.u64()? as usize),
            _ => return None,
        };
        Some(MemoryRegion { address, size })
    }
}

/// `compatible` 属性
#[derive(Clone, Copy)]
pub struct Compatible<'a> {
    pub(crate) data: &'a [u8],
}

impl<'a> Compatible<'a> {
    /// 第一个 compatible 字符串
    #[inline(always)]
    pub fn first(self) -> Option<&'a str> {
        CStr::new(self.data)?.as_str()
    }

    /// 遍历所有 compatible 字符串
    #[inline(always)]
    pub fn all(self) -> impl Iterator<Item = &'a str> {
        let mut data = self.data;
        core::iter::from_fn(move || {
            if data.is_empty() {
                return None;
            }
            if let Some(idx) = data.iter().position(|&b| b == 0) {
                let ret = core::str::from_utf8(data.get(..idx)?).ok()?;
                data = data.get(idx + 1..).unwrap_or(&[]);
                Some(ret)
            } else {
                let ret = core::str::from_utf8(data).ok()?;
                data = &[];
                Some(ret)
            }
        })
    }

    /// 是否包含指定的 compatible 字符串
    #[inline(always)]
    pub fn contains(self, s: &str) -> bool {
        self.all().any(|c| c == s)
    }
}

/// 节点属性
#[derive(Debug, Clone, Copy)]
pub struct NodeProperty<'a> {
    /// 属性名
    pub name: &'a str,
    /// 属性原始值
    pub value: &'a [u8],
}

impl<'a> NodeProperty<'a> {
    /// 解析为 usize
    #[moa_sec_macros::init]
    #[allow(clippy::cast_possible_truncation)]
    pub fn as_usize(self) -> Option<usize> {
        match self.value.len() {
            4 => Some(u32::from_be_bytes(self.value.get(..4)?.try_into().ok()?) as usize),
            8 => Some(u64::from_be_bytes(self.value.get(..8)?.try_into().ok()?) as usize),
            _ => None,
        }
    }

    /// 解析为字符串
    #[inline(always)]
    pub fn as_str(self) -> Option<&'a str> {
        core::str::from_utf8(self.value).map(|s| s.trim_end_matches('\0')).ok()
    }

    #[moa_sec_macros::init]
    fn parse(stream: &mut FdtData<'a>, strings: &'a [u8]) -> Option<Self> {
        if stream.u32()? != FDT_PROP {
            return None;
        }
        let len = stream.u32()? as usize;
        let name_offset = stream.u32()? as usize;
        let data = stream.remaining().get(..len)?;
        skip_4_aligned(stream, len);

        let name = CStr::new(strings.get(name_offset..).unwrap_or(&[]))
            .and_then(|c| c.as_str())
            .unwrap_or("");

        Some(NodeProperty { name, value: data })
    }
}

#[moa_sec_macros::init]
pub(crate) fn find_node<'a>(
    stream: &mut FdtData<'a>,
    name: &str,
    strings: &'a [u8],
    parent_props: Option<&'a [u8]>,
) -> Option<FdtNode<'a>> {
    let mut parts = name.splitn(2, '/');
    let looking_for = parts.next()?;

    stream.skip_nops();
    let curr_data = stream.remaining();

    if stream.u32()? != FDT_BEGIN_NODE {
        return None;
    }

    let unit_name = CStr::new(stream.remaining())?.as_str()?;
    let full_name_len = unit_name.len() + 1;
    skip_4_aligned(stream, full_name_len);

    let looking_contains_addr = looking_for.contains('@');
    let addr_name_same = unit_name == looking_for;
    let base_name_same = unit_name.split('@').next()? == looking_for;

    if (looking_contains_addr && !addr_name_same) || (!looking_contains_addr && !base_name_same) {
        *stream = FdtData::new(curr_data);
        skip_current_node(stream);
        return None;
    }

    let next_part = match parts.next() {
        None | Some("") => {
            return Some(FdtNode::new(unit_name, stream.remaining(), parent_props, strings));
        },
        Some(part) => part,
    };

    stream.skip_nops();
    let new_parent_props = Some(stream.remaining());

    while stream.peek_u32()? == FDT_PROP {
        NodeProperty::parse(stream, strings)?;
    }

    while stream.peek_u32()? == FDT_BEGIN_NODE {
        if let Some(p) = find_node(stream, next_part, strings, new_parent_props) {
            return Some(p);
        }
    }

    stream.skip_nops();

    if stream.u32()? != FDT_END_NODE {
        return None;
    }

    None
}

/// 深度优先全节点迭代器
pub struct AllNodes<'a> {
    stream: FdtData<'a>,
    strings: &'a [u8],
    done: bool,
    parents: [&'a [u8]; MAX_DEPTH],
    parent_index: usize,
}

impl<'a> AllNodes<'a> {
    #[inline(always)]
    pub(crate) fn new(structs: &'a [u8], strings: &'a [u8]) -> Self {
        Self {
            stream: FdtData::new(structs),
            strings,
            done: false,
            parents: [&[]; MAX_DEPTH],
            parent_index: 0,
        }
    }
}

impl<'a> Iterator for AllNodes<'a> {
    type Item = FdtNode<'a>;

    #[moa_sec_macros::init]
    fn next(&mut self) -> Option<Self::Item> {
        if self.stream.is_empty() || self.done {
            return None;
        }

        while self.stream.peek_u32()? == FDT_END_NODE {
            self.parent_index = self.parent_index.checked_sub(1)?;
            self.stream.skip(4);
        }

        if self.stream.peek_u32()? == FDT_END {
            self.done = true;
            return None;
        }

        while self.stream.peek_u32()? == FDT_NOP {
            self.stream.skip(4);
        }

        if self.stream.u32()? != FDT_BEGIN_NODE {
            return None;
        }

        let unit_name = CStr::new(self.stream.remaining())?.as_str()?;
        let full_name_len = unit_name.len() + 1;
        skip_4_aligned(&mut self.stream, full_name_len);

        let curr_node = self.stream.remaining();

        self.parent_index += 1;
        if self.parent_index < MAX_DEPTH {
            self.parents[self.parent_index] = curr_node;
        }

        while self.stream.peek_u32()? == FDT_NOP {
            self.stream.skip(4);
        }

        while self.stream.peek_u32()? == FDT_PROP {
            NodeProperty::parse(&mut self.stream, self.strings)?;
        }

        Some(FdtNode {
            name: if unit_name.is_empty() { "/" } else { unit_name },
            parent_props: match self.parent_index {
                1 => None,
                i if i < MAX_DEPTH => Some(self.parents[i - 1]),
                _ => None,
            },
            props: curr_node,
            strings: self.strings,
        })
    }
}

/// 创建深度优先全节点迭代器
#[inline(always)]
pub(crate) fn all_nodes<'a>(structs: &'a [u8], strings: &'a [u8]) -> AllNodes<'a> {
    AllNodes::new(structs, strings)
}

/// 跳过当前节点（非递归，深度计数）
#[moa_sec_macros::init]
pub(crate) fn skip_current_node(stream: &mut FdtData<'_>) {
    if stream.u32() != Some(FDT_BEGIN_NODE) {
        return;
    }

    let name_len = CStr::new(stream.remaining()).map_or(1, |c| c.len() + 1);
    skip_4_aligned(stream, name_len);

    let mut depth: u32 = 1;
    while depth > 0 {
        match stream.u32() {
            Some(FDT_BEGIN_NODE) => {
                let len = CStr::new(stream.remaining()).map_or(1, |c| c.len() + 1);
                skip_4_aligned(stream, len);
                depth += 1;
            },
            Some(FDT_END_NODE) => depth -= 1,
            Some(FDT_PROP) => {
                let len = stream.u32().unwrap_or(0) as usize;
                stream.u32();
                skip_4_aligned(stream, len);
            },
            Some(FDT_NOP) => {},
            _ => break,
        }
    }
}