graphql_query/visit/
compose.rs

1use super::visitor::*;
2use crate::ast::*;
3
4/// This structure implements the `Visitor` trait and runs two child Visitors in parallel,
5/// executing and calling callbacks on them both.
6///
7/// During traversal the Composed Visitor will keep track of the visitor's [`VisitFlow`] signals and
8/// will avoid calling callbacks on them appropriately, while letting the other visitor continue as
9/// usual. In short, this visitor aims to minimize the work it does while preserving expected
10/// behavior.
11///
12/// Visitors may be composed indefinitely since a Composed Visitor can be passed into another
13/// Composed Visitor, as long as all visitors accept the same `Context` type.
14pub struct ComposedVisitor<'a, Context, A: Visitor<'a, Context>, B: Visitor<'a, Context>> {
15    _marker: std::marker::PhantomData<&'a Context>,
16    depth: usize,
17    skip_a: usize,
18    skip_b: usize,
19    pub a: A,
20    pub b: B,
21}
22
23impl<'a, C, A: Visitor<'a, C>, B: Visitor<'a, C>> ComposedVisitor<'a, C, A, B> {
24    /// Composes two input visitors into one Composed Visitor.
25    #[inline]
26    pub fn new(a: A, b: B) -> ComposedVisitor<'a, C, A, B> {
27        ComposedVisitor {
28            _marker: std::marker::PhantomData,
29            depth: 1,
30            skip_a: usize::MAX,
31            skip_b: usize::MAX,
32            a,
33            b,
34        }
35    }
36
37    #[inline(always)]
38    fn compose_flow_enter<'b, Node: 'a>(
39        &mut self,
40        fn_a: fn(_self: &mut A, ctx: &mut C, node: &'a Node, info: &VisitInfo) -> VisitFlow,
41        fn_b: fn(_self: &mut B, ctx: &mut C, node: &'a Node, info: &VisitInfo) -> VisitFlow,
42        node: &'a Node,
43        info: &VisitInfo,
44        ctx: &mut C,
45    ) -> VisitFlow {
46        let mut all_skip = true;
47        let mut all_break = true;
48
49        if self.skip_a == usize::MAX || self.skip_a == self.depth {
50            self.skip_a = usize::MAX;
51            let flow = fn_a(&mut self.a, ctx, node, info);
52            if flow == VisitFlow::Break {
53                self.skip_a = 0;
54                all_skip = false;
55            } else if flow == VisitFlow::Skip {
56                self.skip_a = self.depth;
57                all_break = false;
58            } else {
59                all_break = false;
60                all_skip = false;
61            }
62        } else if self.skip_b == 0 {
63            all_skip = false;
64        } else {
65            all_break = false;
66        }
67
68        if self.skip_b == usize::MAX || self.skip_b == self.depth {
69            self.skip_b = usize::MAX;
70            let flow = fn_b(&mut self.b, ctx, node, info);
71            if flow == VisitFlow::Break {
72                self.skip_b = 0;
73                all_skip = false;
74            } else if flow == VisitFlow::Skip {
75                self.skip_b = self.depth;
76                all_break = false;
77            } else {
78                all_break = false;
79                all_skip = false;
80            }
81        } else if self.skip_b == 0 {
82            all_skip = false;
83        } else {
84            all_break = false;
85        };
86
87        if all_break {
88            VisitFlow::Break
89        } else if all_skip {
90            if self.skip_a == self.depth {
91                self.skip_a = usize::MAX;
92            }
93            if self.skip_b == self.depth {
94                self.skip_b = usize::MAX;
95            }
96            VisitFlow::Skip
97        } else {
98            self.depth += 1;
99            VisitFlow::Next
100        }
101    }
102
103    #[inline(always)]
104    fn compose_flow_leave<'b, Node: 'a>(
105        &mut self,
106        fn_a: fn(_self: &mut A, ctx: &mut C, node: &'a Node, info: &VisitInfo) -> VisitFlow,
107        fn_b: fn(_self: &mut B, ctx: &mut C, node: &'a Node, info: &VisitInfo) -> VisitFlow,
108        node: &'a Node,
109        info: &VisitInfo,
110        ctx: &'b mut C,
111    ) -> VisitFlow {
112        self.depth -= 1;
113        let mut all_break = true;
114
115        if self.skip_a == usize::MAX {
116            let flow = fn_a(&mut self.a, ctx, node, info);
117            if flow == VisitFlow::Break {
118                self.skip_a = 0;
119            } else {
120                all_break = false;
121            }
122        } else if self.skip_a == self.depth {
123            self.skip_a = usize::MAX;
124            all_break = false;
125        } else if self.skip_a != 0 {
126            all_break = false;
127        }
128
129        if self.skip_b == usize::MAX {
130            let flow = fn_b(&mut self.b, ctx, node, info);
131            if flow == VisitFlow::Break {
132                self.skip_b = 0;
133            } else {
134                all_break = false;
135            }
136        } else if self.skip_b == self.depth {
137            self.skip_b = usize::MAX;
138            all_break = false;
139        } else if self.skip_b != 0 {
140            all_break = false;
141        }
142
143        if all_break {
144            VisitFlow::Break
145        } else {
146            VisitFlow::Next
147        }
148    }
149}
150
151impl<'a, C, A: Visitor<'a, C>, B: Visitor<'a, C>> Visitor<'a, C> for ComposedVisitor<'a, C, A, B> {
152    #[inline]
153    fn enter_document(
154        &mut self,
155        ctx: &mut C,
156        document: &'a Document<'a>,
157        info: &VisitInfo,
158    ) -> VisitFlow {
159        self.compose_flow_enter(A::enter_document, B::enter_document, document, info, ctx)
160    }
161
162    #[inline]
163    fn leave_document(
164        &mut self,
165        ctx: &mut C,
166        document: &'a Document<'a>,
167        info: &VisitInfo,
168    ) -> VisitFlow {
169        self.compose_flow_leave(A::leave_document, B::leave_document, document, info, ctx)
170    }
171
172    #[inline]
173    fn enter_operation(
174        &mut self,
175        ctx: &mut C,
176        operation: &'a OperationDefinition<'a>,
177        info: &VisitInfo,
178    ) -> VisitFlow {
179        self.compose_flow_enter(A::enter_operation, B::enter_operation, operation, info, ctx)
180    }
181
182    #[inline]
183    fn leave_operation(
184        &mut self,
185        ctx: &mut C,
186        operation: &'a OperationDefinition<'a>,
187        info: &VisitInfo,
188    ) -> VisitFlow {
189        self.compose_flow_leave(A::leave_operation, B::leave_operation, operation, info, ctx)
190    }
191
192    #[inline]
193    fn enter_fragment(
194        &mut self,
195        ctx: &mut C,
196        fragment: &'a FragmentDefinition<'a>,
197        info: &VisitInfo,
198    ) -> VisitFlow {
199        self.compose_flow_enter(A::enter_fragment, B::enter_fragment, fragment, info, ctx)
200    }
201
202    #[inline]
203    fn leave_fragment(
204        &mut self,
205        ctx: &mut C,
206        fragment: &'a FragmentDefinition<'a>,
207        info: &VisitInfo,
208    ) -> VisitFlow {
209        self.compose_flow_leave(A::leave_fragment, B::leave_fragment, fragment, info, ctx)
210    }
211
212    #[inline]
213    fn enter_variable_definition(
214        &mut self,
215        ctx: &mut C,
216        var_def: &'a VariableDefinition<'a>,
217        info: &VisitInfo,
218    ) -> VisitFlow {
219        self.compose_flow_enter(
220            A::enter_variable_definition,
221            B::enter_variable_definition,
222            var_def,
223            info,
224            ctx,
225        )
226    }
227
228    #[inline]
229    fn leave_variable_definition(
230        &mut self,
231        ctx: &mut C,
232        var_def: &'a VariableDefinition<'a>,
233        info: &VisitInfo,
234    ) -> VisitFlow {
235        self.compose_flow_leave(
236            A::leave_variable_definition,
237            B::leave_variable_definition,
238            var_def,
239            info,
240            ctx,
241        )
242    }
243
244    #[inline]
245    fn enter_selection_set(
246        &mut self,
247        ctx: &mut C,
248        selection_set: &'a SelectionSet<'a>,
249        info: &VisitInfo,
250    ) -> VisitFlow {
251        self.compose_flow_enter(
252            A::enter_selection_set,
253            B::enter_selection_set,
254            selection_set,
255            info,
256            ctx,
257        )
258    }
259
260    #[inline]
261    fn leave_selection_set(
262        &mut self,
263        ctx: &mut C,
264        selection_set: &'a SelectionSet<'a>,
265        info: &VisitInfo,
266    ) -> VisitFlow {
267        self.compose_flow_leave(
268            A::leave_selection_set,
269            B::leave_selection_set,
270            selection_set,
271            info,
272            ctx,
273        )
274    }
275
276    #[inline]
277    fn enter_fragment_spread(
278        &mut self,
279        ctx: &mut C,
280        fragment_spread: &'a FragmentSpread<'a>,
281        info: &VisitInfo,
282    ) -> VisitFlow {
283        self.compose_flow_enter(
284            A::enter_fragment_spread,
285            B::enter_fragment_spread,
286            fragment_spread,
287            info,
288            ctx,
289        )
290    }
291
292    #[inline]
293    fn leave_fragment_spread(
294        &mut self,
295        ctx: &mut C,
296        fragment_spread: &'a FragmentSpread<'a>,
297        info: &VisitInfo,
298    ) -> VisitFlow {
299        self.compose_flow_leave(
300            A::leave_fragment_spread,
301            B::leave_fragment_spread,
302            fragment_spread,
303            info,
304            ctx,
305        )
306    }
307
308    #[inline]
309    fn enter_inline_fragment(
310        &mut self,
311        ctx: &mut C,
312        inline_fragment: &'a InlineFragment<'a>,
313        info: &VisitInfo,
314    ) -> VisitFlow {
315        self.compose_flow_enter(
316            A::enter_inline_fragment,
317            B::enter_inline_fragment,
318            inline_fragment,
319            info,
320            ctx,
321        )
322    }
323
324    #[inline]
325    fn leave_inline_fragment(
326        &mut self,
327        ctx: &mut C,
328        inline_fragment: &'a InlineFragment<'a>,
329        info: &VisitInfo,
330    ) -> VisitFlow {
331        self.compose_flow_leave(
332            A::leave_inline_fragment,
333            B::leave_inline_fragment,
334            inline_fragment,
335            info,
336            ctx,
337        )
338    }
339
340    #[inline]
341    fn enter_field(&mut self, ctx: &mut C, field: &'a Field<'a>, info: &VisitInfo) -> VisitFlow {
342        self.compose_flow_enter(A::enter_field, B::enter_field, field, info, ctx)
343    }
344
345    #[inline]
346    fn leave_field(&mut self, ctx: &mut C, field: &'a Field<'a>, info: &VisitInfo) -> VisitFlow {
347        self.compose_flow_leave(A::leave_field, B::leave_field, field, info, ctx)
348    }
349
350    #[inline]
351    fn enter_directive(
352        &mut self,
353        ctx: &mut C,
354        directive: &'a Directive<'a>,
355        info: &VisitInfo,
356    ) -> VisitFlow {
357        self.compose_flow_enter(A::enter_directive, B::enter_directive, directive, info, ctx)
358    }
359
360    #[inline]
361    fn leave_directive(
362        &mut self,
363        ctx: &mut C,
364        directive: &'a Directive<'a>,
365        info: &VisitInfo,
366    ) -> VisitFlow {
367        self.compose_flow_leave(A::leave_directive, B::leave_directive, directive, info, ctx)
368    }
369
370    #[inline]
371    fn enter_argument(
372        &mut self,
373        ctx: &mut C,
374        argument: &'a Argument<'a>,
375        info: &VisitInfo,
376    ) -> VisitFlow {
377        self.compose_flow_enter(A::enter_argument, B::enter_argument, argument, info, ctx)
378    }
379
380    #[inline]
381    fn leave_argument(
382        &mut self,
383        ctx: &mut C,
384        argument: &'a Argument<'a>,
385        info: &VisitInfo,
386    ) -> VisitFlow {
387        self.compose_flow_leave(A::leave_argument, B::leave_argument, argument, info, ctx)
388    }
389}
390
391#[cfg(test)]
392mod tests {
393    use super::super::visitor::tests::*;
394    use super::*;
395
396    #[derive(Debug, PartialEq, Default)]
397    struct SkipVisitor {
398        pub(crate) enter: usize,
399        pub(crate) leave: usize,
400        pub(crate) enter_field: usize,
401        pub(crate) leave_field: usize,
402        pub(crate) enter_directive: usize,
403        pub(crate) leave_directive: usize,
404    }
405
406    impl<'a> Visitor<'a, ()> for SkipVisitor {
407        fn enter_selection_set(
408            &mut self,
409            _: &mut (),
410            _selection_set: &SelectionSet,
411            _info: &VisitInfo,
412        ) -> VisitFlow {
413            self.enter += 1;
414            VisitFlow::Next
415        }
416        fn leave_selection_set(
417            &mut self,
418            _: &mut (),
419            _selection_set: &SelectionSet,
420            _info: &VisitInfo,
421        ) -> VisitFlow {
422            self.leave += 1;
423            VisitFlow::Next
424        }
425
426        fn enter_field(&mut self, _: &mut (), _field: &Field, _info: &VisitInfo) -> VisitFlow {
427            self.enter_field += 1;
428            VisitFlow::Skip
429        }
430        fn leave_field(&mut self, _: &mut (), _field: &Field, _info: &VisitInfo) -> VisitFlow {
431            self.leave_field += 1;
432            VisitFlow::Skip
433        }
434
435        fn enter_directive(
436            &mut self,
437            _: &mut (),
438            _field: &Directive,
439            _info: &VisitInfo,
440        ) -> VisitFlow {
441            self.enter_directive += 1;
442            VisitFlow::Skip
443        }
444        fn leave_directive(
445            &mut self,
446            _: &mut (),
447            _field: &Directive,
448            _info: &VisitInfo,
449        ) -> VisitFlow {
450            self.leave_directive += 1;
451            VisitFlow::Skip
452        }
453    }
454
455    #[test]
456    fn kitchen_sink() {
457        let ctx = ASTContext::new();
458        let query = include_str!("../../fixture/kitchen_sink.graphql");
459        let ast = Document::parse(&ctx, query).unwrap();
460        let mut visitor = ComposedVisitor::new(CountVisitor::default(), SkipVisitor::default());
461        ast.visit(&mut (), &mut visitor);
462
463        assert_eq!(
464            visitor.a,
465            CountVisitor {
466                in_document: 1,
467                out_document: 1,
468                in_operation: 5,
469                out_operation: 5,
470                in_fragment: 1,
471                out_fragment: 1,
472                in_variable_definition: 3,
473                out_variable_definition: 3,
474                in_selection_set: 30,
475                out_selection_set: 30,
476                in_fragment_spread: 1,
477                out_fragment_spread: 1,
478                in_inline_fragment: 3,
479                out_inline_fragment: 3,
480                in_field: 21,
481                out_field: 21,
482                in_directive: 10,
483                out_directive: 10,
484                in_argument: 13,
485                out_argument: 13,
486            }
487        );
488
489        assert_eq!(
490            visitor.b,
491            SkipVisitor {
492                enter: 6,
493                leave: 6,
494                enter_field: 7,
495                leave_field: 0,
496                enter_directive: 4,
497                leave_directive: 0,
498            }
499        );
500    }
501}
graphql_query/visit/compose.rs

graphql_query/visit/
compose.rs
