codegen/
trait.rs

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

codegen/
trait.rs
