constrainer/
lib.rs

1use proc_macro2::{
2    Ident,
3    TokenStream,
4    TokenTree,
5    Delimiter,
6    Span
7};
8use quote::{
9    TokenStreamExt,
10    quote
11};
12
13use indexmap::IndexMap;
14use std::collections::{
15    BTreeMap,
16    BTreeSet,
17};
18
19#[proc_macro]
20pub fn create_constrainer(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
21    let input = TokenStream::from(input);
22    let mut trees = input.into_iter();
23    let name = if let TokenTree::Ident(name) = trees.next().unwrap() {
24        name
25    } else {
26        panic!("Environment needs a name!");
27    };
28
29    let data = if let TokenTree::Group(data) = trees.next().unwrap() {
30        data
31    } else {
32        panic!("Need data");
33    };
34    // println!("{:#?}", data);
35
36    let mut identifiers: IndexMap<Ident, Identifier> = IndexMap::new();
37
38    let mut dynamic_fields = TokenStream::new();
39    let mut constrained_fields = TokenStream::new();
40    let mut external_fields = TokenStream::new();
41    let mut deliminated_dynamics = TokenStream::new();
42    let mut deliminated_constraineds = TokenStream::new();
43    let mut init_constraineds = TokenStream::new();
44    let mut ops: TokenStream = TokenStream::new();
45
46    let mut parse_state = ParseState::Key;
47    for token in data.stream() {
48        match parse_state {
49            ParseState::Key => match token {
50                TokenTree::Ident(ident) if ident.to_string().as_str() == "dynamic" => parse_state = ParseState::DynamicName,
51                TokenTree::Ident(ident) if ident.to_string().as_str() == "constrained" => parse_state = ParseState::ConstrainedName,
52                TokenTree::Ident(ident) if ident.to_string().as_str() == "external" => parse_state = ParseState::ExternalName,
53                TokenTree::Ident(ident) if ident.to_string().as_str() == "listener" => parse_state = ParseState::ListenerName,
54                TokenTree::Ident(ident) if ident.to_string().as_str() == "opgenset" => parse_state = ParseState::OpGenSet,
55                _ => panic!("Unexpected token: {}", token)
56            },
57            ParseState::DynamicName => match token {
58                TokenTree::Ident(name) => parse_state = ParseState::DynamicType(name),
59                _ => panic!("Unexpected token: {}", token)
60            },
61            ParseState::DynamicType(name) => match token {
62                TokenTree::Ident(ty) => {
63                    let get_fn_name = Ident::new(&format!("get_{}", name), Span::call_site());
64                    ops.append_all(quote! {
65                        pub fn #get_fn_name(&self) -> &#ty {
66                            &self.#name
67                        }
68                    });
69                    dynamic_fields.append_all(quote! {
70                        #name: #ty,
71                    });
72                    deliminated_dynamics.append_all(quote! {
73                        #name,
74                    });
75                    identifiers.insert(name, Identifier::Dynamic(Dynamic {
76                        ty,
77                        dependents: BTreeSet::new(),
78                    }));
79                    parse_state = ParseState::Key;
80                },
81                _ => panic!("Unexpected token: {}", token)
82            },
83            ParseState::ConstrainedName => match token {
84                TokenTree::Ident(name) => parse_state = ParseState::ConstrainedType(name),
85                _ => panic!("Unexpected token: {}", token)
86            },
87            ParseState::ConstrainedType(name) => match token {
88                TokenTree::Ident(ty) => {
89                    parse_state = ParseState::ConstrainedParams(name, ty);
90                },
91                _ => panic!("Unexpected token: {}", token)
92            },
93            ParseState::ConstrainedParams(name, ty) => match token {
94                TokenTree::Group(group) if group.delimiter() == Delimiter::Parenthesis => {
95                    let mut params = Vec::new();
96                    for token in group.stream() {
97                        match token {
98                            TokenTree::Punct(punct) if punct.as_char() == ',' => {}, // TODO: Remove >1 comma, no comma, and leading comma
99                            TokenTree::Ident(param) => params.push(param),
100                            _ => panic!("Unexpected token: {}", token)
101                        }
102                    }
103                    parse_state = ParseState::ConstrainedBlock(name, ty, params);
104                },
105                _ => panic!("Unexpected token: {}", token)
106            },
107            ParseState::ConstrainedBlock(name, ty, params) => match token {
108                TokenTree::Group(group) if group.delimiter() == Delimiter::Brace => {
109                    let index = identifiers.len();
110                    let mut compute_args = TokenStream::new();
111                    let mut init_args = TokenStream::new();
112                    for param in &params {
113                        let identifier = identifiers.get_mut(param).unwrap();
114                        let param_ty = match identifier {
115                            Identifier::Dynamic(dynamic) => {
116                                dynamic.dependents.insert(index);
117                                &dynamic.ty
118                            },
119                            Identifier::Constrained(constrained) => {
120                                constrained.dependents.insert(index);
121                                &constrained.ty
122                            },
123                            Identifier::External(external) => {
124                                external.dependents.insert(index);
125                                &external.ty
126                            },
127                            Identifier::Listener(_) => {
128                                panic!("A constrained cannot depend on a listener.");
129                            },
130                        };
131                        compute_args.append_all(quote! {
132                            #param: #param_ty,
133                        });
134                        init_args.append_all(quote! {
135                            #param,
136                        });
137                    }
138
139                    let get_fn_name = Ident::new(&format!("get_{}", name), Span::call_site());
140                    ops.append_all(quote! {
141                        pub fn #get_fn_name(&self) -> &#ty {
142                            &self.#name
143                        }
144                    });
145                    constrained_fields.append_all(quote! {
146                        #name: #ty,
147                    });
148                    deliminated_constraineds.append_all(quote! {
149                        #name,
150                    });
151                    let compute_fn_name = Ident::new(&format!("compute_{}", name), Span::call_site());
152                    init_constraineds.append_all(quote! {
153                        let #name = Self::#compute_fn_name(#init_args);
154                    });
155                    let block = group.stream();
156                    ops.append_all(quote! { fn #compute_fn_name (#compute_args) -> #ty { #block }});
157                    identifiers.insert(name, Identifier::Constrained(Constrained {
158                        ty,
159                        params,
160                        block,
161                        compute_fn_name,
162                        dependents: BTreeSet::new(),
163                    }));
164                    parse_state = ParseState::Key;
165                },
166                _ => panic!("Unexpected token: {}", token)
167            },
168            ParseState::ExternalName => match token {
169                TokenTree::Ident(name) => parse_state = ParseState::ExternalType(name),
170                _ => panic!("Unexpected token: {}", token)
171            },
172            ParseState::ExternalType(name) => match token {
173                TokenTree::Ident(ty) => {
174                    external_fields.append_all(quote! {
175                        #name: #ty,
176                    });
177                    identifiers.insert(name, Identifier::External(External {
178                        ty,
179                        dependents: BTreeSet::new(),
180                    }));
181                    parse_state = ParseState::Key;
182                },
183                _ => panic!("Unexpected token: {}", token)
184            },
185            ParseState::ListenerName => match token {
186                TokenTree::Ident(name) => parse_state = ParseState::ListenerParams(name),
187                _ => panic!("Unexpected token: {}", token)
188            },
189            ParseState::ListenerParams(name) => match token {
190                TokenTree::Group(group) if group.delimiter() == Delimiter::Parenthesis => {
191                    let mut params = Vec::new();
192                    for token in group.stream() {
193                        match token {
194                            TokenTree::Punct(punct) if punct.as_char() == ',' => {}, // TODO: Remove >1 comma, no comma, and leading comma
195                            TokenTree::Ident(param) => params.push(param),
196                            _ => panic!("Unexpected token: {}", token)
197                        }
198                    }
199                    parse_state = ParseState::ListenerBlock(name, params);
200                },
201                _ => panic!("Unexpected token: {}", token)
202            },
203            ParseState::ListenerBlock(listener_fn_name, params) => match token {
204                TokenTree::Group(group) if group.delimiter() == Delimiter::Brace => {
205                    let index = identifiers.len();
206                    let mut listener_args = TokenStream::new();
207                    let mut init_args = TokenStream::new();
208                    for param in &params {
209                        let identifier = identifiers.get_mut(param).unwrap();
210                        let param_ty = match identifier {
211                            Identifier::Dynamic(dynamic) => {
212                                dynamic.dependents.insert(index);
213                                &dynamic.ty
214                            },
215                            Identifier::Constrained(constrained) => {
216                                constrained.dependents.insert(index);
217                                &constrained.ty
218                            },
219                            Identifier::External(external) => {
220                                external.dependents.insert(index);
221                                &external.ty
222                            },
223                            Identifier::Listener(_) => {
224                                panic!("A listener cannot depend on a listener.");
225                            },
226                        };
227                        listener_args.append_all(quote! {
228                            #param: #param_ty,
229                        });
230                        init_args.append_all(quote! {
231                            #param,
232                        });
233                    }
234
235                    init_constraineds.append_all(quote! {
236                        Self::#listener_fn_name(#init_args);
237                    });
238                    let block = group.stream();
239                    ops.append_all(quote! { fn #listener_fn_name (#listener_args) { #block }});
240                    identifiers.insert(listener_fn_name, Identifier::Listener(Listener {
241                        params,
242                        block,
243                    }));
244                    parse_state = ParseState::Key;
245                },
246                _ => panic!("Unexpected token: {}", token)
247            },
248            ParseState::OpGenSet => match token {
249                TokenTree::Group(group) if group.delimiter() == Delimiter::Parenthesis => {
250                    let mut set_dynamics = BTreeMap::new();
251                    for token in group.stream() {
252                        match token {
253                            TokenTree::Punct(punct) if punct.as_char() == ',' => {}, // TODO: Remove >1 comma, no comma, and leading comma
254                            TokenTree::Ident(ident) => {
255                                let index = identifiers.get_index_of(&ident).unwrap();
256                                let key_val = identifiers.get_index(index).unwrap();
257                                let name = key_val.0;
258                                let dynamic = if let Identifier::Dynamic(dynamic) = key_val.1 {
259                                    dynamic
260                                } else {
261                                    panic!("OpGenSet can only take dynamics");
262                                };
263                                set_dynamics.insert(index, (name, dynamic));
264                            },
265                            _ => panic!("Unexpected token: {}", token)
266                        }
267                    }
268                    if set_dynamics.len() == 0 {
269                        todo!("opgenset needs at least one dynamic");
270                    }
271
272                    let mut set_dynamics_iter = set_dynamics.iter();
273                    let mut fn_name = format!("set_{}", set_dynamics_iter.next().unwrap().1.0);
274                    for (_, (name, _)) in set_dynamics_iter{
275                        fn_name.push_str(&format!("_{}", name));
276                    }
277                    let fn_name = Ident::new(&fn_name, Span::call_site());
278
279                    let mut set_fn_args = TokenStream::new();
280                    let mut fn_block = TokenStream::new();
281                    set_fn_args.append_all(quote! { &mut self, });
282                    for (_, (name, dynamic)) in &set_dynamics {
283                        let ty = &dynamic.ty;
284                        set_fn_args.append_all(quote! {
285                            #name: #ty,
286                        });
287                        fn_block.append_all(quote! {
288                            self.#name = #name;
289                        });
290                    }
291                    
292                    let mut to_update = BTreeMap::new();
293                    let mut to_call = BTreeMap::new();
294                    let mut found_new_matches = false;
295                    for (_, (_, dynamic)) in &set_dynamics {
296                        for dependent in &dynamic.dependents {
297                            match identifiers.get_index(*dependent).unwrap() {
298                                (name, Identifier::Constrained(constrained)) => {
299                                    to_update.insert(*dependent, (name, constrained));
300                                    
301                                },
302                                (name, Identifier::Listener(listener)) => {
303                                    to_call.insert(*dependent, (name, listener));
304                                },
305                                _ => unreachable!()
306                            }
307                            found_new_matches = true;
308                        }
309                    }
310
311                    if found_new_matches {
312                        let mut last_updates_and_call_size = to_update.len()+to_call.len();
313                        let mut new_updates = BTreeMap::new();
314                        let mut new_calls = BTreeMap::new();
315                        for (_, (_, queued)) in &to_update {
316                            for dependent in &queued.dependents {
317                                match identifiers.get_index(*dependent).unwrap() {
318                                    (name, Identifier::Constrained(constrained)) => {
319                                        new_updates.insert(*dependent, (name, constrained));
320                                    },
321                                    (name, Identifier::Listener(listener)) => {
322                                        new_calls.insert(*dependent, (name, listener));
323                                    },
324                                    _ => unreachable!()
325                                }
326                            }
327                        }
328                        to_update.append(&mut new_updates);
329                        to_call.append(&mut new_calls);
330                        
331                        while last_updates_and_call_size != to_update.len()+to_call.len() {
332                            let mut new_updates = BTreeMap::new();
333                            for (_, (_, queued)) in &to_update {
334                                for dependent in &queued.dependents {
335                                    match identifiers.get_index(*dependent).unwrap() {
336                                        (name, Identifier::Constrained(constrained)) => {
337                                            new_updates.insert(*dependent, (name, constrained));
338                                        },
339                                        (name, Identifier::Listener(listener)) => {
340                                            new_calls.insert(*dependent, (name, listener));
341                                        },
342                                        _ => unreachable!()
343                                    }
344                                }
345                            }
346                            last_updates_and_call_size = to_update.len()+to_call.len();
347                            to_update.append(&mut new_updates);
348                            to_call.append(&mut new_calls);
349                        }
350                    }
351
352                    let mut set_fn_external_args = BTreeMap::new();
353
354                    for (_, (name, constrained)) in to_update {
355                        let compute_fn_name = &constrained.compute_fn_name;
356                        let mut compute_fn_args = TokenStream::new();
357                        for param in &constrained.params {
358                            let param_index = identifiers.get_index_of(param).unwrap();
359                            match identifiers.get(param).unwrap() {
360                                Identifier::Dynamic(_) | Identifier::Constrained(_) => {
361                                    compute_fn_args.append_all(quote! {
362                                        self.#param,
363                                    });
364                                },
365                                Identifier::External(External {
366                                    ty,
367                                    ..
368                                }) => {
369                                    set_fn_external_args.insert(param_index, quote! {
370                                        #param: #ty,
371                                    });
372                                    compute_fn_args.append_all(quote! {
373                                        #param,
374                                    });
375                                },
376                                Identifier::Listener(_) => unreachable!()
377                            }
378                        }
379                        fn_block.append_all(quote! {
380                            self.#name = Self::#compute_fn_name(#compute_fn_args);
381                        });
382                    }
383                    for (_, (listener_fn_name, listener)) in to_call {
384                        let mut listener_fn_args = TokenStream::new();
385                        for param in &listener.params {
386                            let param_index = identifiers.get_index_of(param).unwrap();
387                            match identifiers.get(param).unwrap() {
388                                Identifier::Dynamic(_) | Identifier::Constrained(_) => {
389                                    listener_fn_args.append_all(quote! {
390                                        #param,
391                                    });
392                                },
393                                Identifier::External(External {
394                                    ty,
395                                    ..
396                                }) => {
397                                    set_fn_external_args.insert(param_index, quote! {
398                                        #param: #ty,
399                                    });
400                                    listener_fn_args.append_all(quote! {
401                                        #param,
402                                    });
403                                },
404                                Identifier::Listener(_) => unreachable!()
405                            }
406                        }
407                        fn_block.append_all(quote! {
408                            Self::#listener_fn_name(#listener_fn_args);
409                        });
410                    }
411
412                    for (_, set_fn_external_arg) in set_fn_external_args {
413                        set_fn_args.append_all(set_fn_external_arg);
414                    }
415
416                    ops.append_all(quote! {
417                        pub fn #fn_name(#set_fn_args) {
418                            #fn_block
419                        }
420                    });
421
422                    parse_state = ParseState::Key;
423                },
424                _ => panic!("Unexpected token: {}", token)
425            },
426        }
427    }
428
429    let mut out = TokenStream::new();
430
431    out.append_all(quote! {
432        #[derive(Debug)]
433        struct #name {
434            #dynamic_fields
435            #constrained_fields
436        }
437
438        impl #name {
439            pub fn new ( #dynamic_fields #external_fields ) -> Self {
440                #init_constraineds
441
442                Self {
443                    #deliminated_dynamics
444                    #deliminated_constraineds
445                }
446            }
447
448            #ops
449        }
450    });
451
452    // println!("{:#}", out);
453
454    out.into()
455}
456
457#[derive(Debug)]
458enum ParseState {
459    Key,
460    DynamicName,
461    DynamicType(Ident),
462    ConstrainedName,
463    ConstrainedType(Ident),
464    ConstrainedParams(Ident, Ident),
465    ConstrainedBlock(Ident, Ident, Vec<Ident>),
466    ExternalName,
467    ExternalType(Ident),
468    ListenerName,
469    ListenerParams(Ident),
470    ListenerBlock(Ident, Vec<Ident>),
471    OpGenSet,
472}
473
474#[derive(Debug)]
475enum Identifier {
476    Dynamic(Dynamic),
477    Constrained(Constrained),
478    External(External),
479    Listener(Listener),
480}
481
482#[derive(Debug)]
483struct Dynamic {
484    ty: Ident,
485    dependents: BTreeSet<usize>,
486}
487
488#[derive(Debug)]
489struct Constrained {
490    ty: Ident,
491    params: Vec<Ident>,
492    block: TokenStream,
493    compute_fn_name: Ident,
494    dependents: BTreeSet<usize>,
495}
496
497#[derive(Debug)]
498struct External {
499    ty: Ident,
500    dependents: BTreeSet<usize>,
501}
502
503#[derive(Debug)]
504struct Listener {
505    params: Vec<Ident>,
506    block: TokenStream,
507}