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

codegen/
trait.rs
