indextree_macros/
lib.rs

1use either::Either;
2use itertools::Itertools;
3use proc_macro2::TokenStream;
4use quote::{quote, ToTokens};
5use strum::EnumDiscriminants;
6use syn::{
7    braced,
8    parse::{Parse, ParseStream},
9    parse_macro_input,
10    punctuated::Punctuated,
11    Expr, Token,
12};
13
14#[derive(Clone, Debug)]
15struct IndexNode {
16    node: Expr,
17    children: Punctuated<Self, Token![,]>,
18}
19
20impl Parse for IndexNode {
21    fn parse(input: ParseStream) -> syn::Result<Self> {
22        let node = input.parse::<Expr>()?;
23
24        if input.parse::<Token![=>]>().is_err() {
25            return Ok(IndexNode {
26                node,
27                children: Punctuated::new(),
28            });
29        }
30
31        let children_stream;
32        braced!(children_stream in input);
33        let children = children_stream.parse_terminated(Self::parse, Token![,])?;
34
35        Ok(IndexNode { node, children })
36    }
37}
38
39#[derive(Clone, Debug)]
40struct IndexTree {
41    arena: Expr,
42    root_node: Expr,
43    nodes: Punctuated<IndexNode, Token![,]>,
44}
45
46impl Parse for IndexTree {
47    fn parse(input: ParseStream) -> syn::Result<Self> {
48        let arena = input.parse::<Expr>()?;
49
50        input.parse::<Token![,]>()?;
51
52        let root_node = input.parse::<Expr>()?;
53
54        let nodes = if input.parse::<Token![=>]>().is_ok() {
55            let braced_nodes;
56            braced!(braced_nodes in input);
57            braced_nodes.parse_terminated(IndexNode::parse, Token![,])?
58        } else {
59            Punctuated::new()
60        };
61
62        let _ = input.parse::<Token![,]>();
63
64        Ok(IndexTree {
65            arena,
66            root_node,
67            nodes,
68        })
69    }
70}
71
72#[derive(Clone, EnumDiscriminants, Debug)]
73#[strum_discriminants(name(ActionKind))]
74enum Action {
75    Append(Expr),
76    Parent,
77    Nest,
78}
79
80impl ToTokens for Action {
81    fn to_tokens(&self, tokens: &mut TokenStream) {
82        tokens.extend(self.to_stream())
83    }
84}
85
86impl Action {
87    fn to_stream(&self) -> TokenStream {
88        match self {
89            Action::Append(expr) => quote! {
90                __last = __node.append_value(#expr, __arena);
91            },
92            Action::Parent => quote! {
93                let __temp = ::indextree::Arena::get(__arena, __node);
94                let __temp = ::core::option::Option::unwrap(__temp);
95                let __temp = ::indextree::Node::parent(__temp);
96                let __temp = ::core::option::Option::unwrap(__temp);
97                __node = __temp;
98            },
99            Action::Nest => quote! {
100                __node = __last;
101            },
102        }
103    }
104}
105
106#[derive(Clone, Debug)]
107struct NestingLevelMarker;
108
109#[derive(Clone, Debug)]
110struct ActionStream {
111    count: usize,
112    kind: ActionKind,
113    stream: TokenStream,
114}
115
116impl ToTokens for ActionStream {
117    fn to_tokens(&self, tokens: &mut TokenStream) {
118        tokens.extend(self.stream.clone());
119    }
120}
121
122/// Construct a tree for a given arena.
123///
124/// This macro creates a tree in an [`Arena`] with a pre-defined layout. If the root node is of
125/// type [`NodeId`], then that [`NodeId`] is used for the root node, but if it's any other type,
126/// then it creates a new root node on-the-fly. The macro returns [`NodeId`] of the root node.
127///
128/// # Examples
129///
130/// ```
131/// # use indextree::{Arena, macros::tree};
132/// # let mut arena = Arena::new();
133/// let root_node = arena.new_node("root node");
134/// tree!(
135///     &mut arena,
136///     root_node => {
137///         "1",
138///         "2" => {
139///             "2_1" => { "2_1_1" },
140///             "2_2",
141///         },
142///         "3",
143///     }
144/// );
145///
146/// let automagical_root_node = tree!(
147///     &mut arena,
148///     "root node, but automagically created" => {
149///         "1",
150///         "2" => {
151///             "2_1" => { "2_1_1" },
152///             "2_2",
153///         },
154///         "3",
155///     }
156/// );
157/// ```
158///
159/// Note that you can anchor the root node in the macro to any node at any nesting. So you can take
160/// an already existing node of a tree and attach another tree to it:
161/// ```
162/// # use indextree::{Arena, macros::tree};
163/// # let mut arena = Arena::new();
164/// let root_node = tree!(
165///     &mut arena,
166///     "root node" => {
167///         "1",
168///         "2",
169///         "3",
170///     }
171/// );
172///
173/// let node_1 = arena.get(root_node).unwrap().first_child().unwrap();
174/// let node_2 = arena.get(node_1).unwrap().next_sibling().unwrap();
175/// tree!(
176///     &mut arena,
177///     node_2 => {
178///         "2_1" => { "2_1_1" },
179///         "2_2",
180///     }
181/// );
182/// ```
183///
184/// It is also possible to create an empty root_node, although, I'm not sure why you'd want to do
185/// that.
186/// ```
187/// # use indextree::{Arena, macros::tree};
188/// # let mut arena = Arena::new();
189/// let root_node = tree!(
190///     &mut arena,
191///     "my root node",
192/// );
193/// ```
194/// Empty nodes can also be defined as `=> {}`
195/// ```
196/// # use indextree::{Arena, macros::tree};
197/// # let mut arena = Arena::new();
198/// let root_node = tree!(
199///     &mut arena,
200///     "my root node" => {},
201/// );
202/// ```
203///
204/// [`Arena`]: https://docs.rs/indextree/latest/indextree/struct.Arena.html
205/// [`NodeId`]: https://docs.rs/indextree/latest/indextree/struct.NodeId.html
206#[proc_macro]
207pub fn tree(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
208    let IndexTree {
209        arena,
210        root_node,
211        nodes,
212    } = parse_macro_input!(input as IndexTree);
213
214    let mut stack: Vec<Either<_, NestingLevelMarker>> =
215        nodes.into_iter().map(Either::Left).rev().collect();
216
217    let mut action_buffer: Vec<Action> = Vec::new();
218
219    while let Some(item) = stack.pop() {
220        let Either::Left(IndexNode { node, children }) = item else {
221            action_buffer.push(Action::Parent);
222            continue;
223        };
224
225        action_buffer.push(Action::Append(node));
226
227        if children.is_empty() {
228            continue;
229        }
230
231        // going one level deeper
232        stack.push(Either::Right(NestingLevelMarker));
233        action_buffer.push(Action::Nest);
234        stack.extend(children.into_iter().map(Either::Left).rev());
235    }
236
237    let mut actions: Vec<ActionStream> = action_buffer
238        .into_iter()
239        .map(|action| ActionStream {
240            count: 1,
241            kind: ActionKind::from(&action),
242            stream: action.to_stream(),
243        })
244        .coalesce(|action1, action2| {
245            if action1.kind != action2.kind {
246                return Err((action1, action2));
247            }
248
249            let count = action1.count + action2.count;
250            let kind = action1.kind;
251            let mut stream = action1.stream;
252            stream.extend(action2.stream);
253            Ok(ActionStream {
254                count,
255                kind,
256                stream,
257            })
258        })
259        .collect();
260
261    let is_last_action_useless = actions
262        .last()
263        .map(|last| last.kind == ActionKind::Parent)
264        .unwrap_or(false);
265    if is_last_action_useless {
266        actions.pop();
267    }
268
269    // HACK(alexmozaidze): Due to the fact that specialization is unstable, we must resort to
270    // autoref specialization trick.
271    // https://github.com/dtolnay/case-studies/blob/master/autoref-specialization/README.md
272    quote! {{
273        let mut __arena: &mut ::indextree::Arena<_> = #arena;
274
275        #[repr(transparent)]
276        struct __Wrapping<__T>(::core::mem::ManuallyDrop<__T>);
277
278        trait __ToNodeId<__T> {
279            fn __to_node_id(&mut self, __arena: &mut ::indextree::Arena<__T>) -> ::indextree::NodeId;
280        }
281
282        trait __NodeIdToNodeId<__T> {
283            fn __to_node_id(&mut self, __arena: &mut ::indextree::Arena<__T>) -> ::indextree::NodeId;
284        }
285
286        impl<__T> __NodeIdToNodeId<__T> for __Wrapping<::indextree::NodeId> {
287            fn __to_node_id(&mut self, __arena: &mut ::indextree::Arena<__T>) -> ::indextree::NodeId {
288                unsafe { ::core::mem::ManuallyDrop::take(&mut self.0) }
289            }
290        }
291
292        impl<__T> __ToNodeId<__T> for &mut __Wrapping<__T> {
293            fn __to_node_id(&mut self, __arena: &mut ::indextree::Arena<__T>) -> ::indextree::NodeId {
294                ::indextree::Arena::new_node(__arena, unsafe { ::core::mem::ManuallyDrop::take(&mut self.0) })
295            }
296        }
297
298        let __root_node: ::indextree::NodeId = {
299            let mut __root_node = __Wrapping(::core::mem::ManuallyDrop::new(#root_node));
300            (&mut __root_node).__to_node_id(__arena)
301        };
302        let mut __node: ::indextree::NodeId = __root_node;
303        let mut __last: ::indextree::NodeId;
304
305        #(#actions)*
306
307        __root_node
308    }}.into()
309}
indextree_macros/lib.rs

indextree_macros/
lib.rs
