codegen_rs/
trait.rs

1use std::fmt::{self, Write};
2
3use crate::associated_type::AssociatedType;
4use crate::bound::Bound;
5use crate::formatter::{fmt_bound_rhs, Formatter};
6use crate::function::Function;
7use crate::type_def::TypeDef;
8
9use crate::r#type::Type;
10
11/// Define a trait.
12#[derive(Debug, Clone)]
13pub struct Trait {
14    type_def: TypeDef,
15    parents: Vec<Type>,
16    associated_tys: Vec<AssociatedType>,
17    fns: Vec<Function>,
18    macros: Vec<String>,
19}
20
21impl Trait {
22    /// Return a trait definition with the provided name
23    pub fn new(name: &str) -> Self {
24        Self {
25            type_def: TypeDef::new(name),
26            parents: vec![],
27            associated_tys: vec![],
28            fns: vec![],
29            macros: vec![],
30        }
31    }
32
33    /// Returns a reference to the type
34    pub const fn ty(&self) -> &Type {
35        &self.type_def.ty
36    }
37
38    /// Set the trait visibility.
39    pub fn vis(&mut self, vis: &str) -> &mut Self {
40        self.type_def.vis(vis);
41        self
42    }
43
44    /// Add a generic to the trait
45    pub fn generic(&mut self, name: &str) -> &mut Self {
46        self.type_def.ty.generic(name);
47        self
48    }
49
50    /// Add a `where` bound to the trait.
51    pub fn bound<T>(&mut self, name: &str, ty: T) -> &mut Self
52    where
53        T: Into<Type>,
54    {
55        self.type_def.bound(name, ty);
56        self
57    }
58
59    /// Add a macro to the trait def (e.g. `"#[async_trait]"`)
60    pub fn r#macro(&mut self, r#macro: &str) -> &mut Self {
61        self.type_def.r#macro(r#macro);
62        self
63    }
64
65    /// Add a parent trait.
66    pub fn parent<T>(&mut self, ty: T) -> &mut Self
67    where
68        T: Into<Type>,
69    {
70        self.parents.push(ty.into());
71        self
72    }
73
74    /// Set the trait documentation.
75    pub fn doc(&mut self, docs: &str) -> &mut Self {
76        self.type_def.doc(docs);
77        self
78    }
79
80    /// Add an associated type. Returns a mutable reference to the new
81    /// associated type for futher configuration.
82    pub fn associated_type(&mut self, name: &str) -> &mut AssociatedType {
83        self.associated_tys.push(AssociatedType(Bound {
84            name: name.to_string(),
85            bound: vec![],
86        }));
87
88        self.associated_tys.last_mut().unwrap()
89    }
90
91    /// Push a new function definition, returning a mutable reference to it.
92    pub fn new_fn(&mut self, name: &str) -> &mut Function {
93        let mut func = Function::new(name);
94        func.body = None;
95
96        self.push_fn(func);
97        self.fns.last_mut().unwrap()
98    }
99
100    /// Push a function definition.
101    pub fn push_fn(&mut self, item: Function) -> &mut Self {
102        self.fns.push(item);
103        self
104    }
105
106    /// Formats the scope using the given formatter.
107    pub fn fmt(&self, fmt: &mut Formatter<'_>) -> fmt::Result {
108        self.type_def.fmt_head("trait", &self.parents, fmt)?;
109
110        fmt.block(|fmt| {
111            let assoc = &self.associated_tys;
112
113            // format associated types
114            if !assoc.is_empty() {
115                for ty in assoc {
116                    let ty = &ty.0;
117
118                    write!(fmt, "type {}", ty.name)?;
119
120                    if !ty.bound.is_empty() {
121                        write!(fmt, ": ")?;
122                        fmt_bound_rhs(&ty.bound, fmt)?;
123                    }
124
125                    writeln!(fmt, ";")?;
126                }
127            }
128
129            for (i, func) in self.fns.iter().enumerate() {
130                if i != 0 || !assoc.is_empty() {
131                    writeln!(fmt)?;
132                }
133
134                func.fmt(true, fmt)?;
135            }
136
137            Ok(())
138        })
139    }
140}
codegen_rs/trait.rs

codegen_rs/
trait.rs
