Skip to main content

shape_macros/
lib.rs

1//! Procedural macros for Shape introspection
2//!
3//! This crate provides:
4//! - `#[shape_builtin]` attribute macro for function metadata extraction
5//! - `#[derive(ShapeType)]` derive macro for type/struct property metadata
6
7use proc_macro::TokenStream;
8use proc_macro2::TokenStream as TokenStream2;
9use quote::{format_ident, quote};
10use syn::{
11    Attribute, DeriveInput, Field, ItemFn, Lit, Meta, MetaNameValue, Token, Type,
12    parse_macro_input, punctuated::Punctuated,
13};
14
15/// Attribute macro for Shape builtin functions.
16///
17/// Extracts metadata from doc comments and generates a `METADATA_<NAME>` constant.
18///
19/// # Usage
20///
21/// ```ignore
22/// /// Calculate Simple Moving Average
23/// ///
24/// /// # Parameters
25/// /// * `series: Series` - Input price series
26/// /// * `period: Number` - Lookback period
27/// ///
28/// /// # Returns
29/// /// `Series` - Smoothed series
30/// ///
31/// /// # Example
32/// /// ```shape
33/// /// sma(series("close"), 20)
34/// /// ```
35/// #[shape_builtin(category = "Indicator")]
36/// pub fn eval_sma(args: Vec<Value>, ctx: &mut ExecutionContext) -> Result<Value> {
37///     // implementation
38/// }
39/// ```
40#[proc_macro_attribute]
41pub fn shape_builtin(attr: TokenStream, item: TokenStream) -> TokenStream {
42    let args = parse_macro_input!(attr with Punctuated::<Meta, Token![,]>::parse_terminated);
43    let input = parse_macro_input!(item as ItemFn);
44
45    let expanded = impl_shape_builtin(&args, &input);
46
47    TokenStream::from(expanded)
48}
49
50fn impl_shape_builtin(args: &Punctuated<Meta, Token![,]>, input: &ItemFn) -> TokenStream2 {
51    // Extract category from attribute args
52    let category = extract_category(args).unwrap_or_else(|| "Utility".to_string());
53
54    // Extract function name (remove eval_ prefix if present)
55    let fn_name = input.sig.ident.to_string();
56    let builtin_name = fn_name
57        .strip_prefix("eval_")
58        .or_else(|| fn_name.strip_prefix("intrinsic_"))
59        .unwrap_or(&fn_name)
60        .to_string();
61
62    // Parse doc comments
63    let doc_info = parse_doc_comments(&input.attrs);
64
65    // Generate metadata constant name
66    let metadata_ident = format_ident!("METADATA_{}", builtin_name.to_uppercase());
67
68    // Generate parameter array
69    let params = generate_params(&doc_info.parameters);
70
71    // Build signature string
72    let signature = build_signature(&builtin_name, &doc_info.parameters, &doc_info.return_type);
73
74    // Extract string values from doc_info
75    let description = &doc_info.description;
76    let return_type = &doc_info.return_type;
77
78    // Generate example option
79    let example_tokens = match &doc_info.example {
80        Some(ex) => quote! { Some(#ex) },
81        None => quote! { None },
82    };
83
84    // Generate the metadata constant and preserve the original function
85    let vis = &input.vis;
86    let sig = &input.sig;
87    let block = &input.block;
88    let attrs = &input.attrs;
89
90    quote! {
91        /// Metadata for the builtin function (auto-generated)
92        pub const #metadata_ident: crate::builtin_metadata::BuiltinMetadata = crate::builtin_metadata::BuiltinMetadata {
93            name: #builtin_name,
94            signature: #signature,
95            description: #description,
96            category: #category,
97            parameters: &[#params],
98            return_type: #return_type,
99            example: #example_tokens,
100        };
101
102        #(#attrs)*
103        #vis #sig #block
104    }
105}
106
107fn extract_category(args: &Punctuated<Meta, Token![,]>) -> Option<String> {
108    for meta in args {
109        if let Meta::NameValue(MetaNameValue {
110            path,
111            value: syn::Expr::Lit(expr_lit),
112            ..
113        }) = meta
114        {
115            if path.is_ident("category") {
116                if let Lit::Str(lit_str) = &expr_lit.lit {
117                    return Some(lit_str.value());
118                }
119            }
120        }
121    }
122    None
123}
124
125#[derive(Default)]
126struct DocInfo {
127    description: String,
128    parameters: Vec<ParamInfo>,
129    return_type: String,
130    example: Option<String>,
131}
132
133struct ParamInfo {
134    name: String,
135    param_type: String,
136    optional: bool,
137    description: String,
138}
139
140fn parse_doc_comments(attrs: &[Attribute]) -> DocInfo {
141    let mut info = DocInfo::default();
142    let mut current_section = Section::Description;
143    let mut example_lines: Vec<String> = Vec::new();
144    let mut in_code_block = false;
145
146    for attr in attrs {
147        if attr.path().is_ident("doc") {
148            if let Meta::NameValue(meta) = &attr.meta {
149                if let syn::Expr::Lit(expr_lit) = &meta.value {
150                    if let Lit::Str(lit_str) = &expr_lit.lit {
151                        let line = lit_str.value();
152                        let trimmed = line.trim();
153
154                        // Check for section headers
155                        if trimmed == "# Parameters" {
156                            current_section = Section::Parameters;
157                            continue;
158                        } else if trimmed == "# Returns" {
159                            current_section = Section::Returns;
160                            continue;
161                        } else if trimmed == "# Example" || trimmed == "# Examples" {
162                            current_section = Section::Example;
163                            continue;
164                        }
165
166                        // Check for code block markers
167                        if trimmed.starts_with("```") {
168                            in_code_block = !in_code_block;
169                            if current_section == Section::Example && !in_code_block {
170                                // End of example code block
171                                info.example = Some(example_lines.join("\n"));
172                                example_lines.clear();
173                            }
174                            continue;
175                        }
176
177                        match current_section {
178                            Section::Description => {
179                                if !trimmed.is_empty() {
180                                    if !info.description.is_empty() {
181                                        info.description.push(' ');
182                                    }
183                                    info.description.push_str(trimmed);
184                                }
185                            }
186                            Section::Parameters => {
187                                if let Some(param) = parse_param_line(trimmed) {
188                                    info.parameters.push(param);
189                                }
190                            }
191                            Section::Returns => {
192                                if let Some(ret) = parse_returns_line(trimmed) {
193                                    info.return_type = ret;
194                                }
195                            }
196                            Section::Example => {
197                                if in_code_block {
198                                    example_lines.push(line.to_string());
199                                }
200                            }
201                        }
202                    }
203                }
204            }
205        }
206    }
207
208    // Default return type if not specified
209    if info.return_type.is_empty() {
210        info.return_type = "Any".to_string();
211    }
212
213    info
214}
215
216#[derive(PartialEq)]
217enum Section {
218    Description,
219    Parameters,
220    Returns,
221    Example,
222}
223
224fn parse_param_line(line: &str) -> Option<ParamInfo> {
225    // Parse: * `name: Type` - Description
226    // or:   * `name?: Type` - Description (optional)
227    let line = line.trim_start_matches('*').trim();
228    if !line.starts_with('`') {
229        return None;
230    }
231
232    let line = line.trim_start_matches('`');
233    let end_tick = line.find('`')?;
234    let param_spec = &line[..end_tick];
235    let description = line[end_tick + 1..]
236        .trim_start_matches(" - ")
237        .trim()
238        .to_string();
239
240    // Parse name: Type or name?: Type
241    let (name, param_type, optional) = if let Some(colon_pos) = param_spec.find(':') {
242        let name_part = &param_spec[..colon_pos];
243        let type_part = param_spec[colon_pos + 1..].trim();
244
245        let (name, optional) = if name_part.ends_with('?') {
246            (name_part.trim_end_matches('?').to_string(), true)
247        } else {
248            (name_part.to_string(), false)
249        };
250
251        (name, type_part.to_string(), optional)
252    } else {
253        (param_spec.to_string(), "Any".to_string(), false)
254    };
255
256    Some(ParamInfo {
257        name,
258        param_type,
259        optional,
260        description,
261    })
262}
263
264fn parse_returns_line(line: &str) -> Option<String> {
265    // Parse: `Type` - Description
266    let line = line.trim();
267    if !line.starts_with('`') {
268        return None;
269    }
270
271    let line = line.trim_start_matches('`');
272    let end_tick = line.find('`')?;
273    Some(line[..end_tick].to_string())
274}
275
276fn generate_params(params: &[ParamInfo]) -> TokenStream2 {
277    let param_tokens: Vec<TokenStream2> = params
278        .iter()
279        .map(|p| {
280            let name = &p.name;
281            let param_type = &p.param_type;
282            let optional = p.optional;
283            let description = &p.description;
284
285            quote! {
286                crate::builtin_metadata::BuiltinParam {
287                    name: #name,
288                    param_type: #param_type,
289                    optional: #optional,
290                    description: #description,
291                }
292            }
293        })
294        .collect();
295
296    quote! { #(#param_tokens),* }
297}
298
299fn build_signature(name: &str, params: &[ParamInfo], return_type: &str) -> String {
300    let param_strs: Vec<String> = params
301        .iter()
302        .map(|p| {
303            if p.optional {
304                format!("{}?: {}", p.name, p.param_type)
305            } else {
306                format!("{}: {}", p.name, p.param_type)
307            }
308        })
309        .collect();
310
311    format!("{}({}) -> {}", name, param_strs.join(", "), return_type)
312}
313
314/// Derive macro for Shape type metadata.
315///
316/// Extracts property metadata from struct fields and generates a `TYPE_METADATA_<NAME>` constant.
317///
318/// # Usage
319///
320/// ```ignore
321/// /// A single price bar (OHLCV data)
322/// #[derive(ShapeType)]
323/// #[shape(name = "Candle")]
324/// pub struct CandleValue {
325///     /// Opening price
326///     pub open: f64,
327///     /// Highest price
328///     pub high: f64,
329///     /// Closing price
330///     pub close: f64,
331/// }
332/// ```
333#[proc_macro_derive(ShapeType, attributes(shape))]
334pub fn derive_shape_type(input: TokenStream) -> TokenStream {
335    let input = parse_macro_input!(input as DeriveInput);
336    let expanded = impl_shape_type(&input);
337    TokenStream::from(expanded)
338}
339
340fn impl_shape_type(input: &DeriveInput) -> TokenStream2 {
341    // Extract type name from attribute or use struct name
342    let type_name = extract_type_name(&input.attrs).unwrap_or_else(|| input.ident.to_string());
343
344    // Extract description from doc comments
345    let description = extract_struct_description(&input.attrs);
346
347    // Generate metadata constant name
348    let metadata_ident = format_ident!("TYPE_METADATA_{}", type_name.to_uppercase());
349
350    // Extract fields and generate property metadata
351    let properties = match &input.data {
352        syn::Data::Struct(data) => match &data.fields {
353            syn::Fields::Named(fields) => generate_property_metadata(&fields.named),
354            _ => quote! {},
355        },
356        _ => {
357            return syn::Error::new_spanned(input, "ShapeType can only be derived for structs")
358                .to_compile_error();
359        }
360    };
361
362    quote! {
363        /// Type metadata for LSP introspection (auto-generated)
364        pub const #metadata_ident: crate::builtin_metadata::TypeMetadata = crate::builtin_metadata::TypeMetadata {
365            name: #type_name,
366            description: #description,
367            properties: &[#properties],
368        };
369    }
370}
371
372fn extract_type_name(attrs: &[Attribute]) -> Option<String> {
373    for attr in attrs {
374        if attr.path().is_ident("shape") {
375            if let Ok(nested) =
376                attr.parse_args_with(Punctuated::<Meta, Token![,]>::parse_terminated)
377            {
378                for meta in nested {
379                    if let Meta::NameValue(MetaNameValue {
380                        path,
381                        value: syn::Expr::Lit(expr_lit),
382                        ..
383                    }) = meta
384                    {
385                        if path.is_ident("name") {
386                            if let Lit::Str(lit_str) = &expr_lit.lit {
387                                return Some(lit_str.value());
388                            }
389                        }
390                    }
391                }
392            }
393        }
394    }
395    None
396}
397
398fn extract_struct_description(attrs: &[Attribute]) -> String {
399    let mut description = String::new();
400
401    for attr in attrs {
402        if attr.path().is_ident("doc") {
403            if let Meta::NameValue(meta) = &attr.meta {
404                if let syn::Expr::Lit(expr_lit) = &meta.value {
405                    if let Lit::Str(lit_str) = &expr_lit.lit {
406                        let line = lit_str.value();
407                        let trimmed = line.trim();
408                        if !trimmed.is_empty() {
409                            if !description.is_empty() {
410                                description.push(' ');
411                            }
412                            description.push_str(trimmed);
413                        }
414                    }
415                }
416            }
417        }
418    }
419
420    description
421}
422
423fn generate_property_metadata(fields: &Punctuated<Field, Token![,]>) -> TokenStream2 {
424    let props: Vec<TokenStream2> = fields
425        .iter()
426        .filter_map(|field| {
427            // Skip fields with #[shape(skip)]
428            if has_shape_skip(&field.attrs) {
429                return None;
430            }
431
432            let name = field.ident.as_ref()?.to_string();
433            let prop_type = extract_field_type(&field.attrs, &field.ty);
434            let description = extract_field_description(&field.attrs);
435
436            Some(quote! {
437                crate::builtin_metadata::PropertyMetadata {
438                    name: #name,
439                    prop_type: #prop_type,
440                    description: #description,
441                }
442            })
443        })
444        .collect();
445
446    quote! { #(#props),* }
447}
448
449fn has_shape_skip(attrs: &[Attribute]) -> bool {
450    for attr in attrs {
451        if attr.path().is_ident("shape") {
452            if let Ok(nested) =
453                attr.parse_args_with(Punctuated::<Meta, Token![,]>::parse_terminated)
454            {
455                for meta in nested {
456                    if let Meta::Path(path) = meta {
457                        if path.is_ident("skip") {
458                            return true;
459                        }
460                    }
461                }
462            }
463        }
464    }
465    false
466}
467
468fn extract_field_type(attrs: &[Attribute], ty: &Type) -> String {
469    // First check for explicit #[shape(type = "...")] override
470    for attr in attrs {
471        if attr.path().is_ident("shape") {
472            if let Ok(nested) =
473                attr.parse_args_with(Punctuated::<Meta, Token![,]>::parse_terminated)
474            {
475                for meta in nested {
476                    if let Meta::NameValue(MetaNameValue {
477                        path,
478                        value: syn::Expr::Lit(expr_lit),
479                        ..
480                    }) = meta
481                    {
482                        if path.is_ident("type") {
483                            if let Lit::Str(lit_str) = &expr_lit.lit {
484                                return lit_str.value();
485                            }
486                        }
487                    }
488                }
489            }
490        }
491    }
492
493    // Otherwise, infer from Rust type
494    rust_type_to_shape(ty)
495}
496
497fn rust_type_to_shape(ty: &Type) -> String {
498    match ty {
499        Type::Path(type_path) => {
500            let segments: Vec<_> = type_path
501                .path
502                .segments
503                .iter()
504                .map(|s| s.ident.to_string())
505                .collect();
506            let type_str = segments.last().map(|s| s.as_str()).unwrap_or("Any");
507
508            match type_str {
509                "f64" | "f32" => "Number".to_string(),
510                "i64" | "i32" | "i16" | "i8" | "u64" | "u32" | "u16" | "u8" | "usize" | "isize" => {
511                    "Number".to_string()
512                }
513                "String" => "String".to_string(),
514                "bool" => "Boolean".to_string(),
515                "DateTime" => "DateTime".to_string(),
516                "Series" => "Series".to_string(),
517                "Vec" => {
518                    // Try to extract inner type
519                    if let Some(seg) = type_path.path.segments.last() {
520                        if let syn::PathArguments::AngleBracketed(args) = &seg.arguments {
521                            if let Some(syn::GenericArgument::Type(inner)) = args.args.first() {
522                                let inner_type = rust_type_to_shape(inner);
523                                return format!("Array<{}>", inner_type);
524                            }
525                        }
526                    }
527                    "Array".to_string()
528                }
529                "Option" => {
530                    if let Some(seg) = type_path.path.segments.last() {
531                        if let syn::PathArguments::AngleBracketed(args) = &seg.arguments {
532                            if let Some(syn::GenericArgument::Type(inner)) = args.args.first() {
533                                let inner_type = rust_type_to_shape(inner);
534                                return format!("{}?", inner_type);
535                            }
536                        }
537                    }
538                    "Any?".to_string()
539                }
540                "HashMap" | "BTreeMap" => "Object".to_string(),
541                _ => type_str.to_string(),
542            }
543        }
544        _ => "Any".to_string(),
545    }
546}
547
548fn extract_field_description(attrs: &[Attribute]) -> String {
549    let mut description = String::new();
550
551    for attr in attrs {
552        if attr.path().is_ident("doc") {
553            if let Meta::NameValue(meta) = &attr.meta {
554                if let syn::Expr::Lit(expr_lit) = &meta.value {
555                    if let Lit::Str(lit_str) = &expr_lit.lit {
556                        let line = lit_str.value();
557                        let trimmed = line.trim();
558                        if !trimmed.is_empty() {
559                            if !description.is_empty() {
560                                description.push(' ');
561                            }
562                            description.push_str(trimmed);
563                        }
564                    }
565                }
566            }
567        }
568    }
569
570    description
571}
572
573/// Attribute macro for Shape data providers.
574///
575/// Extracts metadata from doc comments and generates a `PROVIDER_METADATA_<NAME>` constant.
576///
577/// # Usage
578///
579/// ```ignore
580/// /// Market data provider with DuckDB backend
581/// ///
582/// /// # Parameters
583/// /// * `symbol: String` - Stock symbol (required)
584/// /// * `timeframe: String` - Time period (required)
585/// ///
586/// /// # Example
587/// /// ```shape
588/// /// data('market_data', {symbol: 'ES', timeframe: '1h'})
589/// /// ```
590/// #[shape_provider(category = "Market Data")]
591/// pub fn market_data_provider(args: Vec<Value>, ctx: &mut ExecutionContext) -> Result<Value> {
592///     // implementation
593/// }
594/// ```
595#[proc_macro_attribute]
596pub fn shape_provider(attr: TokenStream, item: TokenStream) -> TokenStream {
597    let args = parse_macro_input!(attr with Punctuated::<Meta, Token![,]>::parse_terminated);
598    let input = parse_macro_input!(item as ItemFn);
599
600    let expanded = impl_shape_provider(&args, &input);
601
602    TokenStream::from(expanded)
603}
604
605fn impl_shape_provider(args: &Punctuated<Meta, Token![,]>, input: &ItemFn) -> TokenStream2 {
606    // Extract category from attribute args
607    let category = extract_category(args).unwrap_or_else(|| "Data Provider".to_string());
608
609    // Extract provider name (remove _provider suffix if present, or eval_ prefix)
610    let fn_name = input.sig.ident.to_string();
611    let provider_name = fn_name
612        .strip_suffix("_provider")
613        .or_else(|| fn_name.strip_prefix("eval_"))
614        .unwrap_or(&fn_name)
615        .to_string();
616
617    // Parse doc comments
618    let doc_info = parse_doc_comments(&input.attrs);
619
620    // Generate metadata constant name
621    let metadata_ident = format_ident!("PROVIDER_METADATA_{}", provider_name.to_uppercase());
622
623    // Generate parameter array
624    let params = generate_provider_params(&doc_info.parameters);
625
626    // Extract string values from doc_info
627    let description = &doc_info.description;
628
629    // Generate example option
630    let example_tokens = match &doc_info.example {
631        Some(ex) => quote! { Some(#ex) },
632        None => quote! { None },
633    };
634
635    // Generate the metadata constant and preserve the original function
636    let vis = &input.vis;
637    let sig = &input.sig;
638    let block = &input.block;
639    let attrs = &input.attrs;
640
641    quote! {
642        /// Provider metadata for LSP introspection (auto-generated)
643        pub const #metadata_ident: crate::data::provider_metadata::ProviderMetadata = crate::data::provider_metadata::ProviderMetadata {
644            name: #provider_name,
645            description: #description,
646            category: #category,
647            parameters: &[#params],
648            example: #example_tokens,
649        };
650
651        #(#attrs)*
652        #vis #sig #block
653    }
654}
655
656fn generate_provider_params(params: &[ParamInfo]) -> TokenStream2 {
657    let param_tokens: Vec<TokenStream2> = params
658        .iter()
659        .map(|p| {
660            let name = &p.name;
661            let param_type = &p.param_type;
662            let required = !p.optional; // Invert: optional=false means required=true
663            let description = &p.description;
664
665            quote! {
666                crate::data::provider_metadata::ProviderParam {
667                    name: #name,
668                    param_type: #param_type,
669                    required: #required,
670                    description: #description,
671                    default: None,
672                }
673            }
674        })
675        .collect();
676
677    quote! { #(#param_tokens),* }
678}