//! A small reader for the **device tree** (a "DTB", flattened device tree).
//!
//! Arm boards have no PCI-style "ask the hardware what's there", so the
//! firmware (here QEMU) hands the kernel a description: a tree of nodes
//! like `pl011@9000000` with properties like `compatible = "arm,pl011"`
//! and `reg = <address size>`. The format is simple (devicetree.org spec,
//! chapter 5): a header, then a list of tokens (begin node, property, end
//! node), then a table of property names. Every number in it is big-endian.
//!
//! This reader only does what the bring-up needs: walk every node and hand
//! over its name, `compatible`, `reg`, `interrupts` and `device_type`. It
//! assumes `#address-cells = 2` and `#size-cells = 2` (64-bit addresses
//! and sizes), which is what QEMU's virt uses for every node we read.

const MAGIC: u32 = 0xd00d_feed;
const BEGIN_NODE: u32 = 1;
const END_NODE: u32 = 2;
const PROP: u32 = 3;
const NOP: u32 = 4;
#[allow(dead_code)]
const END: u32 = 9; // (the walk stops at anything it doesn't know, END included)
const MAX_DEPTH: usize = 16;

pub struct Fdt {
    data: &'static [u8],
    structs: usize,
    strings: usize,
}

/// One node, as `for_each_node` hands it over.
#[derive(Clone, Copy, Default)]
pub struct Node<'a> {
    pub name: &'a [u8],
    pub compatible: &'a [u8],
    pub reg: &'a [u8],
    pub interrupts: &'a [u8],
    pub device_type: &'a [u8],
    pub depth: usize,
}

fn be32(d: &[u8], at: usize) -> Option<u32> {
    let b = d.get(at..at + 4)?;
    Some(u32::from_be_bytes([b[0], b[1], b[2], b[3]]))
}

fn cstr(d: &[u8], at: usize) -> &[u8] {
    let rest = d.get(at..).unwrap_or(&[]);
    let n = rest.iter().position(|&b| b == 0).unwrap_or(rest.len());
    &rest[..n]
}

impl Fdt {
    /// The device tree at physical address `addr`, if there is one there.
    ///
    /// # Safety
    /// `addr` must be readable memory (with the MMU on: mapped).
    pub unsafe fn at(addr: usize) -> Option<Fdt> {
        let head = unsafe { core::slice::from_raw_parts(addr as *const u8, 40) };
        if be32(head, 0)? != MAGIC {
            return None;
        }
        let size = be32(head, 4)? as usize;
        if !(40..=(2 << 20)).contains(&size) {
            return None;
        }
        let data = unsafe { core::slice::from_raw_parts(addr as *const u8, size) };
        Some(Fdt { data, structs: be32(data, 8)? as usize, strings: be32(data, 12)? as usize })
    }

    pub fn size(&self) -> usize {
        self.data.len()
    }

    /// Call `f` for every node (children before their parents finish).
    pub fn for_each_node(&self, mut f: impl FnMut(&Node)) {
        let d = self.data;
        let mut stack = [Node::default(); MAX_DEPTH];
        let mut depth = 0usize;
        let mut at = self.structs;
        while let Some(token) = be32(d, at) {
            at += 4;
            match token {
                BEGIN_NODE => {
                    let name = cstr(d, at);
                    at = (at + name.len() + 1 + 3) & !3;
                    if depth < MAX_DEPTH {
                        stack[depth] = Node { name, depth, ..Node::default() };
                    }
                    depth += 1;
                }
                END_NODE => {
                    if depth == 0 {
                        return;
                    }
                    depth -= 1;
                    if depth < MAX_DEPTH {
                        f(&stack[depth]);
                    }
                }
                PROP => {
                    let (Some(len), Some(nameoff)) = (be32(d, at), be32(d, at + 4)) else { return };
                    let value = d.get(at + 8..at + 8 + len as usize).unwrap_or(&[]);
                    at = (at + 8 + len as usize + 3) & !3;
                    if depth == 0 || depth > MAX_DEPTH {
                        continue;
                    }
                    let node = &mut stack[depth - 1];
                    match cstr(d, self.strings + nameoff as usize) {
                        b"compatible" => node.compatible = value,
                        b"reg" => node.reg = value,
                        b"interrupts" => node.interrupts = value,
                        b"device_type" => node.device_type = value,
                        _ => {}
                    }
                }
                NOP => {}
                _ => return, // END (9), or something we don't understand
            }
        }
    }
}

impl Node<'_> {
    /// Is one of the strings in `compatible` exactly `what`?
    pub fn is_compatible(&self, what: &str) -> bool {
        self.compatible.split(|&b| b == 0).any(|s| s == what.as_bytes())
    }

    /// The `i`th (address, size) pair in `reg` (2 + 2 cells).
    pub fn reg(&self, i: usize) -> Option<(u64, u64)> {
        let at = i * 16;
        let cell = |k: usize| be32(self.reg, at + 4 * k).map(|v| v as u64);
        Some(((cell(0)? << 32) | cell(1)?, (cell(2)? << 32) | cell(3)?))
    }

    /// The `i`th 32-bit cell of `interrupts`.
    pub fn interrupt_cell(&self, i: usize) -> Option<u32> {
        be32(self.interrupts, 4 * i)
    }

    pub fn name(&self) -> &str {
        core::str::from_utf8(self.name).unwrap_or("?")
    }
}