fdt_edit/node/
iter.rs

1use core::{
2    ops::{Deref, DerefMut},
3    ptr::NonNull,
4    slice::Iter,
5};
6
7use alloc::vec::Vec;
8
9use crate::{
10    Context, Node, NodeKind, NodeRefClock, NodeRefInterruptController, NodeRefMemory, NodeRefPci,
11    node::gerneric::{NodeMutGen, NodeRefGen},
12};
13
14#[derive(Clone)]
15pub enum NodeRef<'a> {
16    Gerneric(NodeRefGen<'a>),
17    Pci(NodeRefPci<'a>),
18    Clock(NodeRefClock<'a>),
19    InterruptController(NodeRefInterruptController<'a>),
20    Memory(NodeRefMemory<'a>),
21}
22
23impl<'a> NodeRef<'a> {
24    pub fn new(node: &'a Node, ctx: Context<'a>) -> Self {
25        let mut g = NodeRefGen { node, ctx };
26
27        // 先尝试 PCI
28        g = match NodeRefPci::try_from(g) {
29            Ok(pci) => return Self::Pci(pci),
30            Err(v) => v,
31        };
32
33        // 再尝试 Clock
34        g = match NodeRefClock::try_from(g) {
35            Ok(clock) => return Self::Clock(clock),
36            Err(v) => v,
37        };
38
39        // 然后尝试 InterruptController
40        g = match NodeRefInterruptController::try_from(g) {
41            Ok(ic) => return Self::InterruptController(ic),
42            Err(v) => v,
43        };
44
45        // 最后尝试 Memory
46        g = match NodeRefMemory::try_from(g) {
47            Ok(mem) => return Self::Memory(mem),
48            Err(v) => v,
49        };
50
51        Self::Gerneric(g)
52    }
53
54    /// 获取节点的具体类型用于模式匹配
55    pub fn as_ref(&self) -> NodeKind<'a> {
56        match self {
57            NodeRef::Clock(clock) => NodeKind::Clock(clock.clone()),
58            NodeRef::Pci(pci) => NodeKind::Pci(pci.clone()),
59            NodeRef::InterruptController(ic) => NodeKind::InterruptController(ic.clone()),
60            NodeRef::Memory(mem) => NodeKind::Memory(mem.clone()),
61            NodeRef::Gerneric(generic) => NodeKind::Generic(generic.clone()),
62        }
63    }
64}
65
66impl<'a> Deref for NodeRef<'a> {
67    type Target = NodeRefGen<'a>;
68
69    fn deref(&self) -> &Self::Target {
70        match self {
71            NodeRef::Gerneric(n) => n,
72            NodeRef::Pci(n) => &n.node,
73            NodeRef::Clock(n) => &n.node,
74            NodeRef::InterruptController(n) => &n.node,
75            NodeRef::Memory(n) => &n.node,
76        }
77    }
78}
79
80pub enum NodeMut<'a> {
81    Gerneric(NodeMutGen<'a>),
82}
83
84impl<'a> NodeMut<'a> {
85    pub fn new(node: &'a mut Node, ctx: Context<'a>) -> Self {
86        Self::Gerneric(NodeMutGen { node, ctx })
87    }
88}
89
90impl<'a> Deref for NodeMut<'a> {
91    type Target = NodeMutGen<'a>;
92
93    fn deref(&self) -> &Self::Target {
94        match self {
95            NodeMut::Gerneric(n) => n,
96        }
97    }
98}
99
100impl<'a> DerefMut for NodeMut<'a> {
101    fn deref_mut(&mut self) -> &mut Self::Target {
102        match self {
103            NodeMut::Gerneric(n) => n,
104        }
105    }
106}
107
108pub struct NodeIter<'a> {
109    ctx: Context<'a>,
110    node: Option<&'a Node>,
111    stack: Vec<Iter<'a, Node>>,
112}
113
114impl<'a> NodeIter<'a> {
115    pub fn new(root: &'a Node) -> Self {
116        let mut ctx = Context::new();
117        // 预先构建整棵树的 phandle_map
118        // 这样在遍历任何节点时都能通过 phandle 找到其他节点
119        Context::build_phandle_map_from_node(root, &mut ctx.phandle_map);
120
121        Self {
122            ctx,
123            node: Some(root),
124            stack: vec![],
125        }
126    }
127}
128
129impl<'a> Iterator for NodeIter<'a> {
130    type Item = NodeRef<'a>;
131
132    fn next(&mut self) -> Option<Self::Item> {
133        if let Some(n) = self.node.take() {
134            // 返回当前节点，并将其子节点压入栈中
135            let ctx = self.ctx.clone();
136            self.ctx.push(n);
137            self.stack.push(n.children.iter());
138            return Some(NodeRef::new(n, ctx));
139        }
140
141        let iter = self.stack.last_mut()?;
142
143        if let Some(child) = iter.next() {
144            // 返回子节点，并将其子节点压入栈中
145            let ctx = self.ctx.clone();
146            self.ctx.push(child);
147            self.stack.push(child.children.iter());
148            return Some(NodeRef::new(child, ctx));
149        }
150
151        // 当前迭代器耗尽，弹出栈顶
152        self.stack.pop();
153        self.ctx.parents.pop();
154        self.next()
155    }
156}
157
158pub struct NodeIterMut<'a> {
159    ctx: Context<'a>,
160    node: Option<NonNull<Node>>,
161    stack: Vec<RawChildIter>,
162    _marker: core::marker::PhantomData<&'a mut Node>,
163}
164
165/// 原始指针子节点迭代器
166struct RawChildIter {
167    ptr: *mut Node,
168    end: *mut Node,
169}
170
171impl RawChildIter {
172    fn new(children: &mut Vec<Node>) -> Self {
173        let ptr = children.as_mut_ptr();
174        let end = unsafe { ptr.add(children.len()) };
175        Self { ptr, end }
176    }
177
178    fn next(&mut self) -> Option<NonNull<Node>> {
179        if self.ptr < self.end {
180            let current = self.ptr;
181            self.ptr = unsafe { self.ptr.add(1) };
182            NonNull::new(current)
183        } else {
184            None
185        }
186    }
187}
188
189impl<'a> NodeIterMut<'a> {
190    pub fn new(root: &'a mut Node) -> Self {
191        let mut ctx = Context::new();
192        // 预先构建整棵树的 phandle_map
193        // 使用原始指针来避免借用冲突
194        let root_ptr = root as *mut Node;
195        unsafe {
196            // 用不可变引用构建 phandle_map
197            Context::build_phandle_map_from_node(&*root_ptr, &mut ctx.phandle_map);
198        }
199
200        Self {
201            ctx,
202            node: NonNull::new(root_ptr),
203            stack: vec![],
204            _marker: core::marker::PhantomData,
205        }
206    }
207}
208
209impl<'a> Iterator for NodeIterMut<'a> {
210    type Item = NodeMut<'a>;
211
212    fn next(&mut self) -> Option<Self::Item> {
213        if let Some(node_ptr) = self.node.take() {
214            // 返回当前节点，并将其子节点压入栈中
215            let ctx = self.ctx.clone();
216            unsafe {
217                let node_ref = node_ptr.as_ref();
218                self.ctx.push(node_ref);
219                let node_mut = &mut *node_ptr.as_ptr();
220                self.stack.push(RawChildIter::new(&mut node_mut.children));
221                return Some(NodeMut::new(node_mut, ctx));
222            }
223        }
224
225        let iter = self.stack.last_mut()?;
226
227        if let Some(child_ptr) = iter.next() {
228            // 返回子节点，并将其子节点压入栈中
229            let ctx = self.ctx.clone();
230            unsafe {
231                let child_ref = child_ptr.as_ref();
232                self.ctx.push(child_ref);
233                let child_mut = &mut *child_ptr.as_ptr();
234                self.stack.push(RawChildIter::new(&mut child_mut.children));
235                return Some(NodeMut::new(child_mut, ctx));
236            }
237        }
238
239        // 当前迭代器耗尽，弹出栈顶
240        self.stack.pop();
241        self.ctx.parents.pop();
242        self.next()
243    }
244}
fdt_edit/node/iter.rs

fdt_edit/node/
iter.rs
