codegen_rs/
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: &str) -> Self {
59        Self {
60            name: name.to_string(),
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: &str) -> &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: &str) -> &mut Self {
102        self.generics.push(name.to_string());
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`
133            // and `annotation` does not make sense for function arguments.
134            // Simply use empty strings.
135            documentation: Vec::new(),
136            annotation: Vec::new(),
137        });
138
139        self
140    }
141
142    /// Set the function return type.
143    pub fn ret<T>(&mut self, ty: T) -> &mut Self
144    where
145        T: Into<Type>,
146    {
147        self.ret = Some(ty.into());
148        self
149    }
150
151    /// Add a `where` bound to the function.
152    pub fn bound<T>(&mut self, name: &str, ty: T) -> &mut Self
153    where
154        T: Into<Type>,
155    {
156        self.bounds.push(Bound {
157            name: name.to_string(),
158            bound: vec![ty.into()],
159        });
160        self
161    }
162
163    /// Push a line to the function implementation.
164    pub fn line<T>(&mut self, line: T) -> &mut Self
165    where
166        T: ToString,
167    {
168        self.body
169            .get_or_insert(vec![])
170            .push(Body::String(line.to_string()));
171
172        self
173    }
174
175    /// Add an attribute to the function.
176    ///
177    /// ```
178    /// use codegen_rs::Function;
179    ///
180    /// let mut func = Function::new("test");
181    ///
182    /// // add a `#[test]` attribute
183    /// func.attr("test");
184    /// ```
185    pub fn attr(&mut self, attribute: &str) -> &mut Self {
186        self.attributes.push(attribute.to_string());
187        self
188    }
189
190    /// Specify an `extern` ABI for the function.
191    /// ```
192    /// use codegen_rs::Function;
193    ///
194    /// let mut extern_func = Function::new("extern_func");
195    ///
196    /// // use the "C" calling convention
197    /// extern_func.extern_abi("C");
198    /// ```
199    pub fn extern_abi(&mut self, abi: &str) -> &mut Self {
200        self.extern_abi.replace(abi.to_string());
201        self
202    }
203
204    /// Push a block to the function implementation
205    pub fn push_block(&mut self, block: Block) -> &mut Self {
206        self.body.get_or_insert(vec![]).push(Body::Block(block));
207
208        self
209    }
210
211    /// Formats the function using the given formatter.
212    pub fn fmt(&self, is_trait: bool, fmt: &mut Formatter<'_>) -> fmt::Result {
213        if let Some(docs) = &self.docs {
214            docs.fmt(fmt)?;
215        }
216
217        if let Some(allow) = &self.allow {
218            writeln!(fmt, "#[allow({})]", allow)?;
219        }
220
221        for attr in self.attributes.iter() {
222            writeln!(fmt, "#[{}]", attr)?;
223        }
224
225        if is_trait {
226            assert!(
227                self.vis.is_none(),
228                "trait fns do not have visibility modifiers"
229            );
230        }
231
232        if let Some(vis) = &self.vis {
233            write!(fmt, "{} ", vis)?;
234        }
235
236        if let Some(extern_abi) = &self.extern_abi {
237            write!(fmt, "extern \"{extern_abi}\" ", extern_abi = extern_abi)?;
238        }
239
240        if self.r#async {
241            write!(fmt, "async ")?;
242        }
243
244        write!(fmt, "fn {}", self.name)?;
245        fmt_generics(&self.generics, fmt)?;
246
247        write!(fmt, "(")?;
248
249        if let Some(s) = &self.arg_self {
250            write!(fmt, "{}", s)?;
251        }
252
253        for (i, arg) in self.args.iter().enumerate() {
254            if i != 0 || self.arg_self.is_some() {
255                write!(fmt, ", ")?;
256            }
257
258            write!(fmt, "{}: ", arg.name)?;
259            arg.ty.fmt(fmt)?;
260        }
261
262        write!(fmt, ")")?;
263
264        if let Some(ret) = &self.ret {
265            write!(fmt, " -> ")?;
266            ret.fmt(fmt)?;
267        }
268
269        fmt_bounds(&self.bounds, fmt)?;
270
271        match &self.body {
272            Some(body) => fmt.block(|fmt| {
273                for b in body {
274                    b.fmt(fmt)?;
275                }
276
277                Ok(())
278            }),
279            None => {
280                if !is_trait {
281                    panic!("impl blocks must define fn bodies");
282                }
283
284                writeln!(fmt, ";")
285            }
286        }
287    }
288}
codegen_rs/function.rs

codegen_rs/
function.rs
