1mod parse;
8
9use parse::{Attributes, Field, Struct};
10use proc_macro::TokenStream;
11
12#[proc_macro_derive(Model, attributes(model))]
34pub fn derive_model(input: TokenStream) -> TokenStream {
35 match parse::parse_struct(input).and_then(expand) {
36 Ok(tokens) => tokens,
37 Err(message) => format!("::core::compile_error!({:?});", message).parse().unwrap(),
40 }
41}
42
43fn expand(parsed: Struct) -> Result<TokenStream, String> {
44 let Struct { name, fields, attributes } = parsed;
45
46 let stored: Vec<&Field> = fields.iter().filter(|f| !f.attributes.skip).collect();
47 if stored.is_empty() {
48 return Err(format!("`{name}` has no persisted fields"));
49 }
50
51 let table = attributes.table.clone().unwrap_or_else(|| pluralize(&snake_case(&name)));
52
53 let db = attributes.krate.clone().unwrap_or_else(|| "::rustlavel::db".to_string());
57
58 let primary = stored
59 .iter()
60 .find(|f| f.attributes.primary_key)
61 .or_else(|| stored.iter().find(|f| column_of(f) == "id"))
62 .ok_or_else(|| {
63 format!(
64 "`{name}` has no primary key. Add a field called `id`, or mark one \
65 `#[model(primary_key)]`."
66 )
67 })?;
68 let primary_column = column_of(primary);
69 let primary_field = primary.name.clone();
70 let primary_type = primary.type_text.clone();
71
72 let writable: Vec<&&Field> = stored
74 .iter()
75 .filter(|f| !f.attributes.generated && !f.attributes.primary_key && column_of(f) != "id")
76 .filter(|f| !matches!(column_of(f).as_str(), "created_at" | "updated_at"))
77 .collect();
78
79 let column_list = stored
80 .iter()
81 .map(|f| format!("\"{}\"", column_of(f)))
82 .collect::<Vec<_>>()
83 .join(", ");
84
85 let rest = if stored.len() == fields.len() {
88 String::new()
89 } else {
90 " ..::core::default::Default::default()\n".to_string()
91 };
92
93 let from_row = stored
94 .iter()
95 .map(|field| {
96 format!(
97 " {}: row.get({:?})?,",
98 field.name,
99 column_of(field)
100 )
101 })
102 .collect::<Vec<_>>()
103 .join("\n");
104
105 let to_values = writable
106 .iter()
107 .map(|field| {
108 format!(
109 " ({:?}, {db}::Value::from(self.{}.clone())),",
110 column_of(field),
111 field.name
112 )
113 })
114 .collect::<Vec<_>>()
115 .join("\n");
116
117 let generated = format!(
118 r#"
119impl {db}::model::Model for {name} {{
120 type Key = {primary_type};
121
122 const TABLE: &'static str = {table:?};
123 const PRIMARY_KEY: &'static str = {primary_column:?};
124 const COLUMNS: &'static str = {column_list:?};
125
126 fn from_row(row: &{db}::Row) -> {db}::Result<Self> {{
127 Ok({name} {{
128{from_row}
129{rest} }})
130 }}
131
132 fn key(&self) -> Self::Key {{
133 self.{primary_field}.clone()
134 }}
135
136 fn set_key(&mut self, key: Self::Key) {{
137 self.{primary_field} = key;
138 }}
139
140 fn values(&self) -> ::std::vec::Vec<(&'static str, {db}::Value)> {{
141 ::std::vec![
142{to_values}
143 ]
144 }}
145}}
146"#
147 );
148
149 generated.parse().map_err(|e| format!("generated code did not parse: {e}"))
150}
151
152fn column_of(field: &Field) -> String {
154 field.attributes.column.clone().unwrap_or_else(|| field.name.clone())
155}
156
157fn snake_case(input: &str) -> String {
158 let mut out = String::with_capacity(input.len() + 4);
159 for (index, ch) in input.chars().enumerate() {
160 if ch.is_uppercase() {
161 if index > 0 {
162 out.push('_');
163 }
164 out.extend(ch.to_lowercase());
165 } else {
166 out.push(ch);
167 }
168 }
169 out
170}
171
172fn pluralize(word: &str) -> String {
173 if word.ends_with('s') && !word.ends_with("us") && !word.ends_with("ss") {
174 return word.to_string();
175 }
176 if let Some(stem) = word.strip_suffix('y')
177 && !stem.ends_with(['a', 'e', 'i', 'o', 'u']) {
178 return format!("{stem}ies");
179 }
180 if word.ends_with(['s', 'x', 'z']) || word.ends_with("ch") || word.ends_with("sh") {
181 return format!("{word}es");
182 }
183 format!("{word}s")
184}
185
186const _: Option<&Attributes> = None;
188
189#[proc_macro_attribute]
218pub fn main(_attributes: TokenStream, item: TokenStream) -> TokenStream {
219 runtime_wrapper(item, "main")
220}
221
222#[proc_macro_attribute]
224pub fn test(_attributes: TokenStream, item: TokenStream) -> TokenStream {
225 runtime_wrapper(item, "test")
226}
227
228fn runtime_wrapper(item: TokenStream, which: &str) -> TokenStream {
237 let source = item.to_string();
238
239 let Some(async_at) = source.find("async fn ") else {
240 return compile_error(&format!(
241 "#[rustlavel::{which}] goes on an `async fn`. This one is not async — if it does \
242 not await anything, it does not need the attribute."
243 ));
244 };
245
246 let leading = &source[..async_at];
249 let rest = &source[async_at + "async ".len()..];
250
251 let Some(brace_at) = rest.find('{') else {
252 return compile_error(&format!(
253 "#[rustlavel::{which}] needs a function with a body."
254 ));
255 };
256 let signature = &rest[..brace_at];
257 let body = &rest[brace_at..];
258
259 let test_attribute = if which == "test" { "#[::core::prelude::v1::test]" } else { "" };
260
261 format!(
262 "{test_attribute}{leading}{signature}{{
263 ::rustlavel::tokio::runtime::Builder::new_multi_thread()
264 .enable_all()
265 .build()
266 .expect(\"the async runtime could not be started\")
267 .block_on(async {body})
268 }}"
269 )
270 .parse()
271 .unwrap_or_else(|error| {
272 compile_error(&format!("#[rustlavel::{which}] produced code that will not parse: {error}"))
273 })
274}
275
276fn compile_error(message: &str) -> TokenStream {
277 format!("::core::compile_error!({:?});", message).parse().expect("a literal parses")
278}
279
280#[cfg(test)]
281mod tests {
282 use super::{pluralize, snake_case};
285
286 #[test]
287 fn derives_conventional_table_names() {
288 assert_eq!(pluralize(&snake_case("User")), "users");
289 assert_eq!(pluralize(&snake_case("UserProfile")), "user_profiles");
290 assert_eq!(pluralize(&snake_case("Category")), "categories");
291 assert_eq!(pluralize(&snake_case("Address")), "addresses");
292 }
293}