macro_compose/
context.rs

1use proc_macro2::TokenStream;
2use quote::ToTokens;
3use std::ops::Deref;
4use syn::{parse, parse::Parse, parse2, Error};
5
6use crate::{Expand, Lint};
7
8/// Collector collects the results and errors of a macro expansion
9pub struct Collector {
10    err_count: usize,
11    output: TokenStream,
12}
13
14impl Collector {
15    /// create a new collector
16    pub fn new() -> Self {
17        Collector {
18            err_count: 0,
19            output: TokenStream::new(),
20        }
21    }
22
23    /// report an error
24    ///
25    /// once an error has been reported to an collector, `Expand`s will no longer be run
26    pub fn error(&mut self, e: Error) {
27        let error: TokenStream = e.to_compile_error();
28        self.output.extend(error);
29        self.err_count += 1;
30    }
31
32    /// checks if any errors have been reported yet
33    pub fn has_errors(&self) -> bool {
34        self.err_count != 0
35    }
36
37    /// finish the expansion and return the result
38    pub fn finish(self) -> TokenStream {
39        self.output
40    }
41}
42
43impl Default for Collector {
44    fn default() -> Self {
45        Self::new()
46    }
47}
48
49enum Data<'a, T> {
50    Owned(T),
51    Borrowed(&'a T),
52}
53
54impl<T> Deref for Data<'_, T> {
55    type Target = T;
56
57    fn deref(&self) -> &Self::Target {
58        match self {
59            Data::Owned(data) => &data,
60            Data::Borrowed(data) => data,
61        }
62    }
63}
64
65/// used to lint, expand and capture
66///
67/// a context might not have data in it (eg. if parsing the data failed), so calling the lint and expand functions does not guarantee the `Lint`s and `Expand`s actually run
68pub struct Context<'a, T> {
69    collector: &'a mut Collector,
70    data: Option<Data<'a, T>>,
71}
72
73impl<'a, T> Context<'a, T> {
74    /// create a new context with the data
75    pub fn new(collector: &'a mut Collector, data: T) -> Self {
76        Context {
77            collector,
78            data: Some(Data::Owned(data)),
79        }
80    }
81
82    /// create a new context with the data passed by reference
83    pub fn new_by_ref(collector: &'a mut Collector, data: &'a T) -> Self {
84        Context {
85            collector,
86            data: Some(Data::Borrowed(data)),
87        }
88    }
89
90    /// create a new context without any data in it
91    pub fn new_empty(collector: &'a mut Collector) -> Self {
92        Context {
93            collector,
94            data: None,
95        }
96    }
97
98    /// try to parse the data from a [`proc_macro::TokenStream`]
99    ///
100    /// if parsing the data fails the error is reported to the collector
101    pub fn new_parse(collector: &'a mut Collector, data: proc_macro::TokenStream) -> Self
102    where
103        T: Parse,
104    {
105        match parse::<T>(data) {
106            Ok(data) => Self::new(collector, data),
107            Err(e) => {
108                collector.error(e);
109                Self {
110                    collector,
111                    data: None,
112                }
113            }
114        }
115    }
116
117    /// try to parse the data from a [`proc_macro2::TokenStream`]
118    ///
119    /// if parsing the data fails the error is reported to the collector
120    pub fn new_parse2(collector: &'a mut Collector, data: TokenStream) -> Self
121    where
122        T: Parse,
123    {
124        match parse2::<T>(data) {
125            Ok(data) => Self::new(collector, data),
126            Err(e) => {
127                collector.error(e);
128                Self {
129                    collector,
130                    data: None,
131                }
132            }
133        }
134    }
135
136    /// lint the macro input
137    ///
138    /// returns true if the lint ran without reporting an error
139    pub fn lint<L: Lint<T>>(&mut self, lint: &L) -> bool {
140        if let Some(data) = self.data.take() {
141            let start = self.collector.err_count;
142            lint.lint(&*data, &mut self.collector);
143            self.data = Some(data);
144            start == self.collector.err_count
145        } else {
146            false
147        }
148    }
149
150    /// expand the macro and add the result to the collector
151    pub fn expand(&mut self, expand: &impl Expand<T>) {
152        if let Some(res) = self.capture(expand) {
153            let tts: TokenStream = res.to_token_stream();
154            self.collector.output.extend(tts);
155        }
156    }
157
158    /// expand the macro and return the output
159    pub fn capture<E: Expand<T>>(&mut self, expand: &E) -> Option<E::Output> {
160        if self.collector.has_errors() {
161            return None;
162        }
163        if let Some(data) = self.data.as_ref() {
164            expand.expand(&*data, &mut self.collector)
165        } else {
166            None
167        }
168    }
169}