|
1 | 1 | use proc_macro::TokenStream; |
2 | | -use quote::quote; |
| 2 | +use quote::{quote, ToTokens}; |
3 | 3 | use syn::{self}; |
4 | 4 |
|
5 | 5 | /// Automatically implement "Prompt" for all fields in a struct. |
@@ -114,4 +114,264 @@ fn impl_build_image_size(ast: &syn::DeriveInput) -> TokenStream { |
114 | 114 | syntax.into() |
115 | 115 | } |
116 | 116 |
|
117 | | -// TODO add pyclass |
| 117 | +/// Generate a Starlark constructor function definition for this type. |
| 118 | +/// This is used at build time to generate goldboot_dsl.star |
| 119 | +#[proc_macro_derive(StarlarkConstructor, attributes(starlark_doc))] |
| 120 | +pub fn starlark_constructor(input: TokenStream) -> TokenStream { |
| 121 | + let ast = syn::parse(input).unwrap(); |
| 122 | + impl_starlark_constructor(&ast) |
| 123 | +} |
| 124 | + |
| 125 | +fn type_to_starlark_hint(ty: &syn::Type) -> String { |
| 126 | + let ty_string = quote!(#ty).to_string(); |
| 127 | + |
| 128 | + // Handle Option<T> |
| 129 | + if ty_string.starts_with("Option <") { |
| 130 | + // Extract the inner type |
| 131 | + let inner = ty_string.trim_start_matches("Option <").trim_end_matches('>').trim(); |
| 132 | + return format!("{} | None", map_rust_type_to_starlark(inner)); |
| 133 | + } |
| 134 | + |
| 135 | + map_rust_type_to_starlark(&ty_string) |
| 136 | +} |
| 137 | + |
| 138 | +fn map_rust_type_to_starlark(rust_type: &str) -> String { |
| 139 | + match rust_type.trim() { |
| 140 | + "String" | "& str" | "str" => "str".to_string(), |
| 141 | + "bool" => "bool".to_string(), |
| 142 | + "i32" | "i64" | "u32" | "u64" | "usize" => "int".to_string(), |
| 143 | + "f32" | "f64" => "float".to_string(), |
| 144 | + t if t.starts_with("Vec <") => "list".to_string(), |
| 145 | + "Url" | "url :: Url" => "str".to_string(), // URLs are passed as strings in Starlark |
| 146 | + "Size" => "str".to_string(), // Size is passed as string like "20GiB" |
| 147 | + "Arch" => "str".to_string(), // Arch is a wrapper around ImageArch, treat as string |
| 148 | + _ => "dict".to_string(), // Default to dict for custom types |
| 149 | + } |
| 150 | +} |
| 151 | + |
| 152 | +fn impl_starlark_constructor(ast: &syn::DeriveInput) -> TokenStream { |
| 153 | + let name = &ast.ident; |
| 154 | + |
| 155 | + // Extract fields |
| 156 | + let mut fields: Vec<(syn::Ident, syn::Type, bool)> = match &ast.data { |
| 157 | + syn::Data::Struct(data) => match &data.fields { |
| 158 | + syn::Fields::Named(fields) => fields |
| 159 | + .named |
| 160 | + .iter() |
| 161 | + .filter_map(|f| { |
| 162 | + let is_flattened = f.attrs.iter().any(|attr| { |
| 163 | + attr.path().is_ident("serde") |
| 164 | + && attr.meta.to_token_stream().to_string().contains("flatten") |
| 165 | + }); |
| 166 | + |
| 167 | + let has_default = f.attrs.iter().any(|attr| { |
| 168 | + attr.path().is_ident("serde") |
| 169 | + && attr.meta.to_token_stream().to_string().contains("default") |
| 170 | + }); |
| 171 | + |
| 172 | + // Skip fields with defaults (they're optional in Starlark) |
| 173 | + if has_default { |
| 174 | + return None; |
| 175 | + } |
| 176 | + |
| 177 | + // For flattened Hostname, expand to a String field |
| 178 | + if is_flattened { |
| 179 | + let ty_string = quote!(#f.ty).to_string(); |
| 180 | + if ty_string.contains("Hostname") { |
| 181 | + return Some(( |
| 182 | + syn::Ident::new("hostname", f.ident.as_ref().unwrap().span()), |
| 183 | + syn::parse_quote!(String), |
| 184 | + false |
| 185 | + )); |
| 186 | + } |
| 187 | + // Skip other flattened fields |
| 188 | + return None; |
| 189 | + } |
| 190 | + |
| 191 | + Some((f.ident.clone().unwrap(), f.ty.clone(), false)) |
| 192 | + }) |
| 193 | + .collect(), |
| 194 | + _ => panic!("StarlarkConstructor only works on structs with named fields"), |
| 195 | + }, |
| 196 | + _ => panic!("StarlarkConstructor only works on structs"), |
| 197 | + }; |
| 198 | + |
| 199 | + // Sort fields: required first, then optional |
| 200 | + // This ensures Starlark parameter ordering is correct |
| 201 | + let (required_fields, optional_fields): (Vec<_>, Vec<_>) = fields.iter() |
| 202 | + .partition(|(_, ty, _)| !is_option_type(ty)); |
| 203 | + |
| 204 | + // Generate parameter list with type hints |
| 205 | + let mut params: Vec<String> = required_fields.iter().map(|(field, ty, _)| { |
| 206 | + let field_name = field.to_string(); |
| 207 | + let type_hint = type_to_starlark_hint(ty); |
| 208 | + format!("{}: {}", field_name, type_hint) |
| 209 | + }).collect(); |
| 210 | + |
| 211 | + params.extend(optional_fields.iter().map(|(field, ty, _)| { |
| 212 | + let field_name = field.to_string(); |
| 213 | + let type_hint = type_to_starlark_hint(ty); |
| 214 | + format!("{}: {} = None", field_name, type_hint) |
| 215 | + })); |
| 216 | + |
| 217 | + let params_str = params.join(",\n "); |
| 218 | + |
| 219 | + let mut required_entries: Vec<String> = fields.iter() |
| 220 | + .filter(|(_, ty, _)| !is_option_type(ty)) |
| 221 | + .map(|(field, _, _)| format!(r#" "{field}": {field},"#, field = field.to_string())) |
| 222 | + .collect(); |
| 223 | + |
| 224 | + // Check if this is an OS type (needs "os" discriminator) |
| 225 | + // OS types typically have 'iso' field and the name doesn't end with common helper type names |
| 226 | + let field_names: Vec<String> = fields.iter().map(|(f, _, _)| f.to_string()).collect(); |
| 227 | + let name_str = name.to_string(); |
| 228 | + let is_helper_type = name_str == "Iso" |
| 229 | + || name_str.ends_with("Edition") |
| 230 | + || name_str.ends_with("Release") |
| 231 | + || name_str.ends_with("Config") |
| 232 | + || name_str.ends_with("Path"); |
| 233 | + |
| 234 | + let is_os_type = field_names.contains(&"iso".to_string()) && !is_helper_type; |
| 235 | + |
| 236 | + if is_os_type { |
| 237 | + // Add "os" discriminator as the first field |
| 238 | + required_entries.insert(0, format!(r#" "os": "{name}","#, name = name)); |
| 239 | + } |
| 240 | + |
| 241 | + let required_entries_str = required_entries.join("\n"); |
| 242 | + |
| 243 | + let optional_checks: String = fields.iter() |
| 244 | + .filter(|(_, ty, _)| is_option_type(ty)) |
| 245 | + .map(|(field, _, _)| { |
| 246 | + let field_name = field.to_string(); |
| 247 | + format!( |
| 248 | + r#" if {field} != None: |
| 249 | + config["{field}"] = {field}"#, |
| 250 | + field = field_name |
| 251 | + ) |
| 252 | + }) |
| 253 | + .collect::<Vec<_>>() |
| 254 | + .join("\n"); |
| 255 | + |
| 256 | + let starlark_fn = format!( |
| 257 | + r#"def {name}( |
| 258 | + {params} |
| 259 | +) -> dict: |
| 260 | + """Create a {name} configuration.""" |
| 261 | + config = {{ |
| 262 | +{required_entries} |
| 263 | + }} |
| 264 | +{optional_checks} |
| 265 | + return config |
| 266 | +"#, |
| 267 | + name = name, |
| 268 | + params = params_str, |
| 269 | + required_entries = required_entries_str, |
| 270 | + optional_checks = optional_checks |
| 271 | + ); |
| 272 | + |
| 273 | + // Store the function definition as a const |
| 274 | + let const_name = syn::Ident::new(&format!("STARLARK_FN_{}", name.to_string().to_uppercase()), name.span()); |
| 275 | + let registration_name = syn::Ident::new(&format!("STARLARK_FN_{}_REGISTRATION", name.to_string().to_uppercase()), name.span()); |
| 276 | + |
| 277 | + let syntax = quote! { |
| 278 | + #[doc(hidden)] |
| 279 | + pub const #const_name: &str = #starlark_fn; |
| 280 | + |
| 281 | + #[linkme::distributed_slice(crate::config::starlark_dsl::STARLARK_DSL_FUNCTIONS)] |
| 282 | + #[linkme(crate = linkme)] |
| 283 | + static #registration_name: &str = #const_name; |
| 284 | + }; |
| 285 | + |
| 286 | + syntax.into() |
| 287 | +} |
| 288 | + |
| 289 | +/// Generate Starlark constructor functions for each variant of an enum. |
| 290 | +/// Each variant becomes a separate function that returns a dict with the variant name as a key. |
| 291 | +#[proc_macro_derive(StarlarkEnumConstructors)] |
| 292 | +pub fn starlark_enum_constructors(input: TokenStream) -> TokenStream { |
| 293 | + let ast = syn::parse(input).unwrap(); |
| 294 | + impl_starlark_enum_constructors(&ast) |
| 295 | +} |
| 296 | + |
| 297 | +fn impl_starlark_enum_constructors(ast: &syn::DeriveInput) -> TokenStream { |
| 298 | + let enum_name = &ast.ident; |
| 299 | + |
| 300 | + let variants: Vec<_> = match &ast.data { |
| 301 | + syn::Data::Enum(data) => data.variants.iter().collect(), |
| 302 | + _ => panic!("StarlarkEnumConstructors only works on enums"), |
| 303 | + }; |
| 304 | + |
| 305 | + let mut generated_consts = vec![]; |
| 306 | + |
| 307 | + for variant in variants { |
| 308 | + let variant_name = &variant.ident; |
| 309 | + |
| 310 | + // Handle different variant types |
| 311 | + let starlark_fn = match &variant.fields { |
| 312 | + // Tuple variant with one field: Plaintext(String) |
| 313 | + syn::Fields::Unnamed(fields) if fields.unnamed.len() == 1 => { |
| 314 | + let field_type = &fields.unnamed.first().unwrap().ty; |
| 315 | + let type_hint = type_to_starlark_hint(field_type); |
| 316 | + |
| 317 | + format!( |
| 318 | + r#"def {variant_name}(value: {type_hint}) -> dict: |
| 319 | + """Create a {variant_name} {enum_name}.""" |
| 320 | + return {{"{snake_case}": value}} |
| 321 | +"#, |
| 322 | + variant_name = variant_name, |
| 323 | + type_hint = type_hint, |
| 324 | + enum_name = enum_name, |
| 325 | + snake_case = to_snake_case(&variant_name.to_string()) |
| 326 | + ) |
| 327 | + } |
| 328 | + // Unit variant or other types |
| 329 | + _ => { |
| 330 | + format!( |
| 331 | + r#"def {variant_name}() -> dict: |
| 332 | + """Create a {variant_name} {enum_name}.""" |
| 333 | + return {{"{snake_case}": None}} |
| 334 | +"#, |
| 335 | + variant_name = variant_name, |
| 336 | + enum_name = enum_name, |
| 337 | + snake_case = to_snake_case(&variant_name.to_string()) |
| 338 | + ) |
| 339 | + } |
| 340 | + }; |
| 341 | + |
| 342 | + let const_name = syn::Ident::new( |
| 343 | + &format!("STARLARK_FN_{}_{}", enum_name.to_string().to_uppercase(), variant_name.to_string().to_uppercase()), |
| 344 | + variant_name.span() |
| 345 | + ); |
| 346 | + let registration_name = syn::Ident::new( |
| 347 | + &format!("STARLARK_FN_{}_{}_REGISTRATION", enum_name.to_string().to_uppercase(), variant_name.to_string().to_uppercase()), |
| 348 | + variant_name.span() |
| 349 | + ); |
| 350 | + |
| 351 | + generated_consts.push(quote! { |
| 352 | + #[doc(hidden)] |
| 353 | + pub const #const_name: &str = #starlark_fn; |
| 354 | + |
| 355 | + #[linkme::distributed_slice(crate::config::starlark_dsl::STARLARK_DSL_FUNCTIONS)] |
| 356 | + #[linkme(crate = linkme)] |
| 357 | + static #registration_name: &str = #const_name; |
| 358 | + }); |
| 359 | + } |
| 360 | + |
| 361 | + let syntax = quote! { |
| 362 | + #(#generated_consts)* |
| 363 | + }; |
| 364 | + |
| 365 | + syntax.into() |
| 366 | +} |
| 367 | + |
| 368 | +fn to_snake_case(s: &str) -> String { |
| 369 | + let mut result = String::new(); |
| 370 | + for (i, ch) in s.chars().enumerate() { |
| 371 | + if ch.is_uppercase() && i > 0 { |
| 372 | + result.push('_'); |
| 373 | + } |
| 374 | + result.push(ch.to_ascii_lowercase()); |
| 375 | + } |
| 376 | + result |
| 377 | +} |
0 commit comments