codegen/
function.rs

1use std::fmt::{self, Write};
2
3use crate::block::Block;
4use crate::body::Body;
5use crate::bound::Bound;
6use crate::docs::Docs;
7use crate::field::Field;
8use crate::formatter::Formatter;
9use crate::formatter::{fmt_bounds, fmt_generics};
10
11use crate::r#type::Type;
12
13/// Defines a function.
14#[derive(Debug, Clone)]
15pub struct Function {
16    /// Name of the function
17    name: String,
18
19    /// Function documentation
20    docs: Option<Docs>,
21
22    /// A lint attribute used to suppress a warning or error
23    allow: Option<String>,
24
25    /// Function visibility
26    vis: Option<String>,
27
28    /// Function generics
29    generics: Vec<String>,
30
31    /// If the function takes `&self` or `&mut self`
32    arg_self: Option<String>,
33
34    /// Function arguments
35    args: Vec<Field>,
36
37    /// Return type
38    ret: Option<Type>,
39
40    /// Where bounds
41    bounds: Vec<Bound>,
42
43    /// Body contents
44    pub body: Option<Vec<Body>>,
45
46    /// Function attributes, e.g., `#[no_mangle]`.
47    attributes: Vec<String>,
48
49    /// Function `extern` ABI
50    extern_abi: Option<String>,
51
52    /// Whether or not this function is `async` or not
53    r#async: bool,
54}
55
56impl Function {
57    /// Return a new function definition.
58    pub fn new(name: impl Into<String>) -> Self {
59        Function {
60            name: name.into(),
61            docs: None,
62            allow: None,
63            vis: None,
64            generics: vec![],
65            arg_self: None,
66            args: vec![],
67            ret: None,
68            bounds: vec![],
69            body: Some(vec![]),
70            attributes: vec![],
71            extern_abi: None,
72            r#async: false,
73        }
74    }
75
76    /// Set the function documentation.
77    pub fn doc(&mut self, docs: impl Into<String>) -> &mut Self {
78        self.docs = Some(Docs::new(docs));
79        self
80    }
81
82    /// Specify lint attribute to supress a warning or error.
83    pub fn allow(&mut self, allow: &str) -> &mut Self {
84        self.allow = Some(allow.to_string());
85        self
86    }
87
88    /// Set the function visibility.
89    pub fn vis(&mut self, vis: &str) -> &mut Self {
90        self.vis = Some(vis.to_string());
91        self
92    }
93
94    /// Set whether this function is async or not
95    pub fn set_async(&mut self, r#async: bool) -> &mut Self {
96        self.r#async = r#async;
97        self
98    }
99
100    /// Add a generic to the function.
101    pub fn generic(&mut self, name: impl Into<String>) -> &mut Self {
102        self.generics.push(name.into());
103        self
104    }
105
106    /// Add `self` as a function argument.
107    pub fn arg_self(&mut self) -> &mut Self {
108        self.arg_self = Some("self".to_string());
109        self
110    }
111
112    /// Add `&self` as a function argument.
113    pub fn arg_ref_self(&mut self) -> &mut Self {
114        self.arg_self = Some("&self".to_string());
115        self
116    }
117
118    /// Add `&mut self` as a function argument.
119    pub fn arg_mut_self(&mut self) -> &mut Self {
120        self.arg_self = Some("&mut self".to_string());
121        self
122    }
123
124    /// Add a function argument.
125    pub fn arg<T>(&mut self, name: &str, ty: T) -> &mut Self
126    where
127        T: Into<Type>,
128    {
129        self.args.push(Field {
130            name: name.to_string(),
131            ty: ty.into(),
132            // While a `Field` is used here, both `documentation`, `visibility`
133            // and `annotation` does not make sense for function arguments.
134            // Simply use empty strings.
135            documentation: String::new(),
136            annotation: Vec::new(),
137            value: String::new(),
138            visibility: None,
139        });
140
141        self
142    }
143
144    /// Set the function return type.
145    pub fn ret<T>(&mut self, ty: T) -> &mut Self
146    where
147        T: Into<Type>,
148    {
149        self.ret = Some(ty.into());
150        self
151    }
152
153    /// Add a `where` bound to the function.
154    pub fn bound<T>(&mut self, name: &str, ty: T) -> &mut Self
155    where
156        T: Into<Type>,
157    {
158        self.bounds.push(Bound {
159            name: name.to_string(),
160            bound: vec![ty.into()],
161        });
162        self
163    }
164
165    /// Push a line to the function implementation.
166    pub fn line<T>(&mut self, line: T) -> &mut Self
167    where
168        T: ToString,
169    {
170        self.body
171            .get_or_insert(vec![])
172            .push(Body::String(line.to_string()));
173
174        self
175    }
176
177    /// Add an attribute to the function.
178    ///
179    /// ```
180    /// use codegen::Function;
181    ///
182    /// let mut func = Function::new("test");
183    ///
184    /// // add a `#[test]` attribute
185    /// func.attr("test");
186    /// ```
187    pub fn attr(&mut self, attribute: &str) -> &mut Self {
188        self.attributes.push(attribute.to_string());
189        self
190    }
191
192    /// Specify an `extern` ABI for the function.
193    /// ```
194    /// use codegen::Function;
195    ///
196    /// let mut extern_func = Function::new("extern_func");
197    ///
198    /// // use the "C" calling convention
199    /// extern_func.extern_abi("C");
200    /// ```
201    pub fn extern_abi(&mut self, abi: &str) -> &mut Self {
202        self.extern_abi.replace(abi.to_string());
203        self
204    }
205
206    /// Push a block to the function implementation
207    pub fn push_block(&mut self, block: Block) -> &mut Self {
208        self.body.get_or_insert(vec![]).push(Body::Block(block));
209
210        self
211    }
212
213    /// Formats the function using the given formatter.
214    pub fn fmt(&self, is_trait: bool, fmt: &mut Formatter<'_>) -> fmt::Result {
215        if let Some(ref docs) = self.docs {
216            docs.fmt(fmt)?;
217        }
218
219        if let Some(ref allow) = self.allow {
220            write!(fmt, "#[allow({})]\n", allow)?;
221        }
222
223        for attr in self.attributes.iter() {
224            write!(fmt, "#[{}]\n", attr)?;
225        }
226
227        if is_trait {
228            assert!(
229                self.vis.is_none(),
230                "trait fns do not have visibility modifiers"
231            );
232        }
233
234        if let Some(ref vis) = self.vis {
235            write!(fmt, "{} ", vis)?;
236        }
237
238        if let Some(ref extern_abi) = self.extern_abi {
239            write!(fmt, "extern \"{extern_abi}\" ", extern_abi = extern_abi)?;
240        }
241
242        if self.r#async {
243            write!(fmt, "async ")?;
244        }
245
246        write!(fmt, "fn {}", self.name)?;
247        fmt_generics(&self.generics, fmt)?;
248
249        write!(fmt, "(")?;
250
251        if let Some(ref s) = self.arg_self {
252            write!(fmt, "{}", s)?;
253        }
254
255        for (i, arg) in self.args.iter().enumerate() {
256            if i != 0 || self.arg_self.is_some() {
257                write!(fmt, ", ")?;
258            }
259
260            write!(fmt, "{}: ", arg.name)?;
261            arg.ty.fmt(fmt)?;
262        }
263
264        write!(fmt, ")")?;
265
266        if let Some(ref ret) = self.ret {
267            write!(fmt, " -> ")?;
268            ret.fmt(fmt)?;
269        }
270
271        fmt_bounds(&self.bounds, fmt)?;
272
273        match self.body {
274            Some(ref body) => fmt.block(|fmt| {
275                for b in body {
276                    b.fmt(fmt)?;
277                }
278
279                Ok(())
280            }),
281            None => {
282                if !is_trait {
283                    panic!("impl blocks must define fn bodies");
284                }
285
286                write!(fmt, ";\n")
287            }
288        }
289    }
290}
codegen/function.rs

codegen/
function.rs
