rust_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 /// The name of the function.
17 name: String,
18 /// The function's documentation.
19 docs: Option<Docs>,
20 /// A lint attribute used to suppress a warning or error.
21 allow: Option<String>,
22 /// Function visibility.
23 vis: Option<String>,
24 /// Function generics.
25 generics: Vec<String>,
26 /// If the function takes `&self` or `&mut self`.
27 arg_self: Option<String>,
28 /// Function arguments.
29 args: Vec<Field>,
30 /// Return type.
31 ret: Option<Type>,
32 /// Where bounds.
33 bounds: Vec<Bound>,
34 /// Body contents.
35 pub body: Option<Vec<Body>>,
36 /// Function attributes, e.g., `#[no_mangle]`.
37 attributes: Vec<String>,
38 /// Function `extern` ABI.
39 extern_abi: Option<String>,
40 /// Whether or not this function is `async` or not.
41 r#async: bool,
42}
43
44impl Function {
45 /// Return a new function definition.
46 ///
47 /// # Arguments
48 ///
49 /// * `name` - The name of the function.
50 ///
51 /// # Examples
52 ///
53 /// ```
54 /// use rust_codegen::Function;
55 ///
56 /// let foo_fn = Function::new("foo_fn");
57 /// ```
58 pub fn new(name: &str) -> Self {
59 Function {
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 ///
78 /// # Arguments
79 ///
80 /// * `docs` - The docs to set for the function.
81 ///
82 /// # Examples
83 ///
84 /// ```
85 /// use rust_codegen::Function;
86 ///
87 /// let mut foo_fn = Function::new("foo_fn");
88 /// foo_fn.doc("Sample Foo function documentation.");
89 /// ```
90 pub fn doc(&mut self, docs: &str) -> &mut Self {
91 self.docs = Some(Docs::new(docs));
92 self
93 }
94
95 /// Specify lint attribute to supress a warning or error.
96 ///
97 /// # Arguments
98 ///
99 /// * `allow` - The lint attribute to add.
100 ///
101 /// # Examples
102 ///
103 /// ```
104 /// use rust_codegen::Function;
105 ///
106 /// let mut foo_fn = Function::new("foo_fn");
107 /// foo_fn.allow("dead_code");
108 /// ```
109 pub fn allow(&mut self, allow: &str) -> &mut Self {
110 self.allow = Some(allow.to_string());
111 self
112 }
113
114 /// Set the function visibility.
115 ///
116 /// # Arguments
117 ///
118 /// * `vis` - The visiblity to set for the function.
119 ///
120 /// # Examples
121 ///
122 /// ```
123 /// use rust_codegen::Function;
124 ///
125 /// let mut foo_fn = Function::new("foo_fn");
126 /// foo_fn.vis("pub");
127 /// ```
128 pub fn vis(&mut self, vis: &str) -> &mut Self {
129 self.vis = Some(vis.to_string());
130 self
131 }
132
133 /// Set whether this function is async or not.
134 ///
135 /// # Arguments
136 ///
137 /// * `async` - Indicates whether this function is async or not.
138 ///
139 /// # Examples
140 ///
141 /// ```
142 /// use rust_codegen::Function;
143 ///
144 /// let mut foo_fn = Function::new("foo_fn");
145 /// foo_fn.set_async(true);
146 /// ```
147 pub fn set_async(&mut self, r#async: bool) -> &mut Self {
148 self.r#async = r#async;
149 self
150 }
151
152 /// Add a generic to the function.
153 ///
154 /// # Arguments
155 ///
156 /// * `name` - The name of the generic to add.
157 ///
158 /// # Examples
159 ///
160 /// ```
161 /// use rust_codegen::Function;
162 ///
163 /// let mut foo_fn = Function::new("foo_fn");
164 /// foo_fn.generic("T");
165 /// ```
166 pub fn generic(&mut self, name: &str) -> &mut Self {
167 self.generics.push(name.to_string());
168 self
169 }
170
171 /// Add `self` as a function argument.
172 ///
173 /// # Examples
174 ///
175 /// ```
176 /// use rust_codegen::Function;
177 ///
178 /// let mut foo_fn = Function::new("foo_fn");
179 /// foo_fn.arg_self();
180 /// ```
181 pub fn arg_self(&mut self) -> &mut Self {
182 self.arg_self = Some("self".to_string());
183 self
184 }
185
186 /// Add `&self` as a function argument.
187 ///
188 /// # Examples
189 ///
190 /// ```
191 /// use rust_codegen::Function;
192 ///
193 /// let mut foo_fn = Function::new("foo_fn");
194 /// foo_fn.arg_ref_self();
195 /// ```
196 pub fn arg_ref_self(&mut self) -> &mut Self {
197 self.arg_self = Some("&self".to_string());
198 self
199 }
200
201 /// Add `&mut self` as a function argument.
202 ///
203 /// # Examples
204 ///
205 /// ```
206 /// use rust_codegen::Function;
207 ///
208 /// let mut foo_fn = Function::new("foo_fn");
209 /// foo_fn.arg_mut_self();
210 /// ```
211 pub fn arg_mut_self(&mut self) -> &mut Self {
212 self.arg_self = Some("&mut self".to_string());
213 self
214 }
215
216 /// Add a function argument.
217 ///
218 /// # Arguments
219 ///
220 /// * `name` - The name of the argument.
221 /// * `ty` - The type of the argument.
222 ///
223 /// # Examples
224 ///
225 /// ```
226 /// use rust_codegen::Function;
227 ///
228 /// let mut foo_fn = Function::new("foo_fn");
229 /// foo_fn.arg("name", "&str");
230 /// ```
231 pub fn arg<T>(&mut self, name: &str, ty: T) -> &mut Self
232 where
233 T: Into<Type>,
234 {
235 self.args.push(Field {
236 name: name.to_string(),
237 ty: ty.into(),
238 // While a `Field` is used here, both `documentation`
239 // and `annotation` does not make sense for function arguments.
240 // Simply use empty strings.
241 documentation: Vec::new(),
242 annotation: Vec::new(),
243 });
244
245 self
246 }
247
248 /// Set the function return type.
249 ///
250 /// # Arguments
251 ///
252 /// * `ty` - The return type.
253 ///
254 /// # Examples
255 ///
256 /// ```
257 /// use rust_codegen::Function;
258 ///
259 /// let mut foo_fn = Function::new("foo_fn");
260 /// foo_fn.ret("String");
261 /// ```
262 pub fn ret<T>(&mut self, ty: T) -> &mut Self
263 where
264 T: Into<Type>,
265 {
266 self.ret = Some(ty.into());
267 self
268 }
269
270 /// Add a `where` bound to the function.
271 ///
272 /// # Arguments
273 ///
274 /// * `name ` - The name of the bound.
275 /// * `ty` - The type of the bound.
276 ///
277 /// # Examples
278 ///
279 /// ```
280 /// use rust_codegen::Function;
281 ///
282 /// let mut foo_fn = Function::new("foo_fn");
283 /// foo_fn.bound("A", "TraitA");
284 /// ```
285 pub fn bound<T>(&mut self, name: &str, ty: T) -> &mut Self
286 where
287 T: Into<Type>,
288 {
289 self.bounds.push(Bound {
290 name: name.to_string(),
291 bound: vec![ty.into()],
292 });
293 self
294 }
295
296 /// Push a line to the function implementation.
297 ///
298 /// # Arguments
299 ///
300 /// * `line` - The line to add to the function.
301 ///
302 /// # Examples
303 ///
304 /// ```
305 /// use rust_codegen::Function;
306 ///
307 /// let mut foo_fn = Function::new("foo_fn");
308 /// foo_fn.line("println!(\"Hello, world!\")");
309 /// ```
310 pub fn line<T>(&mut self, line: T) -> &mut Self
311 where
312 T: ToString,
313 {
314 self.body
315 .get_or_insert(vec![])
316 .push(Body::String(line.to_string()));
317
318 self
319 }
320
321 /// Add an attribute to the function.
322 ///
323 /// # Arguments
324 ///
325 /// * `attribute` - The attribute to add.
326 ///
327 /// # Examples
328 ///
329 /// ```
330 /// use rust_codegen::Function;
331 ///
332 /// let mut foo_fn = Function::new("foo_fn");
333 /// foo_fn.attr("test");
334 /// ```
335 pub fn attr(&mut self, attribute: &str) -> &mut Self {
336 self.attributes.push(attribute.to_string());
337 self
338 }
339
340 /// Specify an `extern` ABI for the function.
341 ///
342 /// # Arguments
343 ///
344 /// * `abi` - The extern ABI to add.
345 ///
346 /// # Examples
347 ///
348 /// ```
349 /// use rust_codegen::Function;
350 ///
351 /// let mut foo_fn = Function::new("foo_fn");
352 /// foo_fn.extern_abi("C");
353 /// ```
354 pub fn extern_abi(&mut self, abi: &str) -> &mut Self {
355 self.extern_abi.replace(abi.to_string());
356 self
357 }
358
359 /// Push a block to the function implementation.
360 ///
361 /// # Arguments
362 ///
363 /// * `block` - The block to push to the function.
364 ///
365 /// # Examples
366 ///
367 /// ```
368 /// use rust_codegen::*;
369 ///
370 /// let mut foo_fn = Function::new("foo_fn");
371 ///
372 /// let mut block = Block::new("");
373 /// block.line("println!(\"Hello, world!\");");
374 ///
375 /// foo_fn.push_block(block);
376 /// ```
377 pub fn push_block(&mut self, block: Block) -> &mut Self {
378 self.body.get_or_insert(vec![]).push(Body::Block(block));
379
380 self
381 }
382
383 /// Formats the function using the given formatter.
384 ///
385 /// # Arguments
386 ///
387 /// * `is_trait` - Indicates whether it is a trait or not.
388 /// * `fmt` - The formatter to use.
389 ///
390 /// # Examples
391 ///
392 /// ```
393 /// use rust_codegen::*;
394 ///
395 /// let mut dest = String::new();
396 /// let mut fmt = Formatter::new(&mut dest);
397 ///
398 /// let mut foo_fn = Function::new("foo_fn");
399 /// foo_fn.fmt(false, &mut fmt);
400 /// ```
401 pub fn fmt(&self, is_trait: bool, fmt: &mut Formatter<'_>) -> fmt::Result {
402 if let Some(ref docs) = self.docs {
403 docs.fmt(fmt)?;
404 }
405
406 if let Some(ref allow) = self.allow {
407 write!(fmt, "#[allow({})]\n", allow)?;
408 }
409
410 for attr in self.attributes.iter() {
411 write!(fmt, "#[{}]\n", attr)?;
412 }
413
414 if is_trait {
415 assert!(
416 self.vis.is_none(),
417 "trait fns do not have visibility modifiers"
418 );
419 }
420
421 if let Some(ref vis) = self.vis {
422 write!(fmt, "{} ", vis)?;
423 }
424
425 if let Some(ref extern_abi) = self.extern_abi {
426 write!(fmt, "extern \"{extern_abi}\" ", extern_abi = extern_abi)?;
427 }
428
429 if self.r#async {
430 write!(fmt, "async ")?;
431 }
432
433 write!(fmt, "fn {}", self.name)?;
434 fmt_generics(&self.generics, fmt)?;
435
436 write!(fmt, "(")?;
437
438 if let Some(ref s) = self.arg_self {
439 write!(fmt, "{}", s)?;
440 }
441
442 for (i, arg) in self.args.iter().enumerate() {
443 if i != 0 || self.arg_self.is_some() {
444 write!(fmt, ", ")?;
445 }
446
447 write!(fmt, "{}: ", arg.name)?;
448 arg.ty.fmt(fmt)?;
449 }
450
451 write!(fmt, ")")?;
452
453 if let Some(ref ret) = self.ret {
454 write!(fmt, " -> ")?;
455 ret.fmt(fmt)?;
456 }
457
458 fmt_bounds(&self.bounds, fmt)?;
459
460 match self.body {
461 Some(ref body) => fmt.block(|fmt| {
462 for b in body {
463 b.fmt(fmt)?;
464 }
465
466 Ok(())
467 }),
468 None => {
469 if !is_trait {
470 panic!("impl blocks must define fn bodies");
471 }
472
473 write!(fmt, ";\n")
474 }
475 }
476 }
477}