kodept_ast/graph/
tree.rs

1pub use main::{SyntaxTree, SyntaxTreeBuilder, SyntaxTreeMutView};
2pub use utils::{ChildScope, DfsIter};
3
4use crate::graph::nodes::Inaccessible;
5use crate::graph::tags::ChildTag;
6
7type Graph<T = Inaccessible, E = ChildTag> = slotgraph::DiGraph<T, E>;
8
9mod stage {
10    use crate::graph::PermTkn;
11
12    #[derive(Debug)]
13    pub struct FullAccess(pub(super) PermTkn);
14
15    #[derive(Debug)]
16    pub struct ViewingAccess<'a>(pub(super) &'a PermTkn);
17
18    pub struct ModificationAccess<'a>(pub(super) &'a mut PermTkn);
19
20    #[derive(Default, Debug)]
21    pub struct NoAccess;
22
23    pub(super) trait CanAccess {
24        fn tkn(&self) -> &PermTkn;
25    }
26
27    impl CanAccess for FullAccess {
28        fn tkn(&self) -> &PermTkn {
29            &self.0
30        }
31    }
32
33    impl CanAccess for ViewingAccess<'_> {
34        fn tkn(&self) -> &PermTkn {
35            self.0
36        }
37    }
38
39    impl CanAccess for ModificationAccess<'_> {
40        fn tkn(&self) -> &PermTkn {
41            self.0
42        }
43    }
44}
45
46mod main {
47    use std::fmt::{Display, Formatter};
48
49    use kodept_core::{ConvertibleToMut, ConvertibleToRef};
50    use slotgraph::export::{Config, Direction, Dot};
51    use slotgraph::SubDiGraph;
52
53    use crate::graph::{AnyNode, Change, ChangeSet, Identifiable, NodeId, PermTkn};
54    use crate::graph::nodes::Inaccessible;
55    use crate::graph::tags::{ChildTag, TAGS_DESC};
56    use crate::graph::tree::{ChildScope, DfsIter, Graph};
57    use crate::graph::tree::stage::{
58        CanAccess, FullAccess, ModificationAccess, NoAccess, ViewingAccess,
59    };
60    use crate::graph::utils::OptVec;
61    use crate::node_properties::{Node, NodeWithParent};
62    use crate::Uninit;
63
64    #[derive(Debug)]
65    pub struct SyntaxTree<Permission = NoAccess> {
66        pub(super) graph: Graph,
67        pub(super) stage: Permission,
68    }
69
70    pub type SyntaxTreeBuilder = SyntaxTree<FullAccess>;
71    pub type SyntaxTreeMutView<'token> = SyntaxTree<ModificationAccess<'token>>;
72
73    impl Default for SyntaxTree<FullAccess> {
74        fn default() -> Self {
75            // SAFE: While tree is building, token should be owned by it
76            Self {
77                graph: Default::default(),
78                stage: FullAccess(PermTkn::new()),
79            }
80        }
81    }
82
83    impl SyntaxTree<FullAccess> {
84        pub fn new() -> SyntaxTree<FullAccess> {
85            SyntaxTree::default()
86        }
87
88        pub fn export_dot<'a>(&'a self, config: &'a [Config]) -> impl Display + 'a {
89            struct Wrapper<'c>(SubDiGraph<String, &'static str>, &'c [Config]);
90            impl Display for Wrapper<'_> {
91                fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
92                    let dot = Dot::with_config(&self.0, self.1);
93                    write!(f, "{dot}")
94                }
95            }
96
97            let mapping = self.graph.map(
98                |id, node| format!("{} [{}]", node.ro(&self.stage.0).name(), id),
99                |_, &edge| TAGS_DESC[edge as usize],
100            );
101            Wrapper(mapping, config)
102        }
103    }
104
105    #[allow(private_bounds)]
106    impl<U: CanAccess> SyntaxTree<U> {
107        pub fn add_node<T>(&mut self, node: Uninit<T>) -> ChildScope<'_, T, U>
108        where
109            T: Identifiable + Into<AnyNode>,
110        {
111            let id = self
112                .graph
113                .add_node_with_key(|id| Inaccessible::new(node.unwrap(id.into())));
114            ChildScope::new(self, id.into())
115        }
116
117        pub fn build(self) -> SyntaxTree<NoAccess> {
118            SyntaxTree {
119                graph: self.graph,
120                stage: NoAccess,
121            }
122        }
123    }
124
125    impl SyntaxTree<ViewingAccess<'_>> {
126        fn apply_change(&mut self, change: Change) {
127            match change {
128                Change::Delete { child_id, .. } => {
129                    self.graph.remove_node(child_id.into());
130                }
131                Change::Add {
132                    parent_id,
133                    child,
134                    tag,
135                } => {
136                    let child_id = self
137                        .graph
138                        .add_node_with_key(|id| Inaccessible::new(child.unwrap(id.into())));
139                    self.graph.add_edge(parent_id.into(), child_id, tag);
140                }
141                Change::Replace { from_id, to } => {
142                    match self.graph.node_weight_mut(from_id.into()) {
143                        None => {}
144                        Some(x) => *x = Inaccessible::new(to.unwrap(from_id)),
145                    };
146                }
147                Change::DeleteSelf { node_id } => {
148                    self.graph.remove_node(node_id.into());
149                }
150            };
151        }
152    }
153
154    impl<U> SyntaxTree<U> {
155        pub(crate) fn children_of_raw<T>(
156            &self,
157            id: NodeId<T>,
158            tag: ChildTag,
159        ) -> OptVec<&Inaccessible> {
160            self.graph
161                .children(id.into())
162                .filter(|it| self.graph[it.0].data == tag)
163                .map(|it| &self.graph[it.1])
164                .collect()
165        }
166
167        pub fn dfs(&self) -> DfsIter {
168            let mut roots = self.graph.externals(Direction::Incoming);
169            let (Some(start), None) = (roots.next(), roots.next()) else {
170                panic!("Syntax tree should have a root")
171            };
172
173            DfsIter::new(&self.graph, start)
174        }
175    }
176
177    impl SyntaxTree {
178        pub fn children_of<'b, T, U>(
179            &'b self,
180            id: NodeId<T>,
181            token: &'b PermTkn,
182            tag: ChildTag,
183        ) -> OptVec<&U>
184        where
185            AnyNode: ConvertibleToRef<U>,
186        {
187            self.graph
188                .children(id.into())
189                .filter(|it| self.graph[it.0].data == tag)
190                .filter_map(|it| self.graph[it.1].ro(token).try_as_ref())
191                .collect()
192        }
193
194        pub fn get<'b, T>(&'b self, id: NodeId<T>, token: &'b PermTkn) -> Option<&T>
195        where
196            AnyNode: ConvertibleToRef<T>,
197        {
198            let node_ref = self.graph.node_weight(id.into())?;
199            node_ref.ro(token).try_as_ref()
200        }
201
202        pub fn get_mut<'b, T>(&'b self, id: NodeId<T>, token: &'b mut PermTkn) -> Option<&mut T>
203        where
204            AnyNode: ConvertibleToMut<T>,
205        {
206            let node_ref = self.graph.node_weight(id.into())?;
207            node_ref.rw(token).try_as_mut()
208        }
209
210        pub fn parent_of<'a, T>(&'a self, id: NodeId<T>, token: &'a PermTkn) -> Option<&AnyNode> {
211            let mut parents = self.graph.parents(id.into());
212            if let (Some((_, parent_id)), None) = (parents.next(), parents.next()) {
213                Some(self.graph[parent_id].ro(token))
214            } else {
215                None
216            }
217        }
218
219        pub fn parent_of_mut<'a, T>(
220            &'a self,
221            id: NodeId<T>,
222            token: &'a mut PermTkn,
223        ) -> &mut T::Parent
224        where
225            T: NodeWithParent + Node,
226            AnyNode: ConvertibleToMut<T::Parent>,
227        {
228            let mut parents = self.graph.parents(id.into());
229            let (Some((_, parent_id)), None) = (parents.next(), parents.next()) else {
230                panic!(
231                    "NodeWithParent contract violated: such kind of nodes should always have a parent"
232                )
233            };
234            let parent_ref = self.graph[parent_id].rw(token);
235            parent_ref.try_as_mut().expect("Node has wrong type")
236        }
237
238        pub fn apply_changes(self, changes: ChangeSet, token: &PermTkn) -> Self {
239            let mut this = self.modify(token);
240            for change in changes {
241                this.apply_change(change);
242            }
243            this.build()
244        }
245
246        fn modify(self, token: &PermTkn) -> SyntaxTree<ViewingAccess> {
247            SyntaxTree {
248                graph: self.graph,
249                stage: ViewingAccess(token),
250            }
251        }
252    }
253}
254
255mod utils {
256    use std::collections::VecDeque;
257    use std::iter::FusedIterator;
258
259    use kodept_core::structure::span::CodeHolder;
260    use slotgraph::{Key, NodeKey};
261    use slotgraph::export::NodeCount;
262
263    use crate::graph::{AnyNode, HasChildrenMarker, NodeId, SyntaxTree};
264    use crate::graph::nodes::Inaccessible;
265    use crate::graph::tags::ChildTag;
266    use crate::graph::tree::Graph;
267    use crate::graph::tree::stage::{CanAccess, FullAccess};
268    use crate::rlt_accessor::RLTFamily;
269    use crate::traits::{Linker, PopulateTree};
270    use crate::visit_side::VisitSide;
271
272    pub struct ChildScope<'arena, T, Stage = FullAccess> {
273        tree: &'arena mut SyntaxTree<Stage>,
274        id: Key<T>,
275    }
276
277    pub enum TraverseState {
278        DescendDeeper,
279        Exit,
280    }
281
282    pub struct DfsIter<'a> {
283        stack: VecDeque<(NodeKey, TraverseState)>,
284        edges_buffer: Vec<NodeKey>,
285        graph: &'a Graph,
286    }
287
288    impl<'arena> DfsIter<'arena> {
289        pub(super) fn new(graph: &'arena Graph, start: NodeKey) -> Self {
290            let mut stack = VecDeque::with_capacity(graph.node_count());
291            stack.push_back((start, TraverseState::DescendDeeper));
292
293            Self {
294                stack,
295                edges_buffer: vec![],
296                graph,
297            }
298        }
299    }
300
301    impl<'arena> Iterator for DfsIter<'arena> {
302        type Item = (&'arena Inaccessible, VisitSide);
303
304        fn next(&mut self) -> Option<Self::Item> {
305            let (next, descend) = self.stack.pop_back()?;
306            let current = self.graph.node_weight(next)?;
307            if matches!(descend, TraverseState::Exit) {
308                return Some((current, VisitSide::Exiting));
309            }
310
311            self.edges_buffer.clear();
312            self.edges_buffer
313                .extend(self.graph.children(next).map(|it| it.1));
314            self.edges_buffer.reverse();
315            let edges_iter = self.edges_buffer.iter();
316            if edges_iter.len() != 0 {
317                self.stack.push_back((next, TraverseState::Exit));
318                for child in edges_iter {
319                    self.stack.push_back((*child, TraverseState::DescendDeeper));
320                }
321                Some((current, VisitSide::Entering))
322            } else {
323                Some((current, VisitSide::Leaf))
324            }
325        }
326
327        fn size_hint(&self) -> (usize, Option<usize>) {
328            (self.stack.len(), Some(self.graph.node_count() * 2))
329        }
330    }
331
332    impl FusedIterator for DfsIter<'_> {}
333
334    impl<'arena, T, S> ChildScope<'arena, T, S> {
335        pub(super) fn new(tree: &'arena mut SyntaxTree<S>, node_id: Key<T>) -> Self {
336            Self { tree, id: node_id }
337        }
338
339        fn add_child_by_ref<U, const TAG: ChildTag>(&mut self, child_id: NodeKey)
340        where
341            U: Into<AnyNode>,
342            T: HasChildrenMarker<U, TAG>,
343        {
344            self.tree.graph.add_edge(self.id.into(), child_id, TAG);
345        }
346
347        #[allow(private_bounds)]
348        pub fn with_rlt<U>(self, context: &mut impl Linker, rlt_node: &U) -> Self
349        where
350            U: Into<RLTFamily> + Clone,
351            T: Into<AnyNode>,
352            S: CanAccess,
353        {
354            let element = &self.tree.graph[NodeKey::from(self.id)];
355            let node = element.ro(self.tree.stage.tkn());
356
357            context.link(node, rlt_node);
358            self
359        }
360
361        pub fn id(&self) -> NodeId<T> {
362            self.id.into()
363        }
364    }
365
366    impl<T> ChildScope<'_, T, FullAccess> {
367        #[allow(private_bounds)]
368        pub fn with_children_from<'b, const TAG: ChildTag, U>(
369            mut self,
370            iter: impl IntoIterator<Item = &'b U>,
371            context: &mut (impl Linker + CodeHolder),
372        ) -> Self
373        where
374            T: HasChildrenMarker<U::Output, TAG>,
375            U: PopulateTree + 'b,
376        {
377            for item in iter {
378                let child_id = item.convert(self.tree, context);
379                self.add_child_by_ref(child_id.into());
380            }
381            self
382        }
383    }
384}
385
386#[cfg(test)]
387mod tests {
388    use crate::FileDecl;
389    use crate::graph::SyntaxTreeBuilder;
390
391    #[test]
392    fn test_tree_creation() {
393        let mut builder = SyntaxTreeBuilder::new();
394
395        let id = builder.add_node(FileDecl::uninit()).id();
396
397        let tree = builder.build();
398        let children = tree.children_of_raw(id, 0);
399
400        assert!(children.is_empty());
401    }
402}
kodept_ast/graph/tree.rs

kodept_ast/graph/
tree.rs
