Skip to main content

leo3_macros_backend/
lean_instance.rs

1//! Implementation of the `#[lean_instance]` macro.
2//!
3//! This macro generates FFI functions that implement Lean type classes for Rust types.
4
5use proc_macro2::TokenStream;
6use quote::{format_ident, quote};
7use syn::{parse::Parse, ImplItem, ItemImpl, Token};
8
9use crate::get_leo3_crate;
10
11/// Options for the `#[lean_instance]` macro.
12pub struct LeanInstanceOptions {
13    /// The type class to implement (BEq, Hashable, Repr, ToString, Ord).
14    pub typeclass: syn::Ident,
15    /// Path to the leo3 crate (for re-exports).
16    pub krate: Option<syn::Path>,
17}
18
19impl Parse for LeanInstanceOptions {
20    fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> {
21        let typeclass: syn::Ident = input.parse()?;
22
23        // Optional: crate path
24        let krate = if input.peek(Token![,]) {
25            input.parse::<Token![,]>()?;
26            if input.peek(syn::Ident) {
27                let ident: syn::Ident = input.parse()?;
28                if ident == "crate" {
29                    input.parse::<Token![=]>()?;
30                    Some(input.parse()?)
31                } else {
32                    return Err(syn::Error::new(ident.span(), "Expected `crate = path`"));
33                }
34            } else {
35                None
36            }
37        } else {
38            None
39        };
40
41        Ok(LeanInstanceOptions { typeclass, krate })
42    }
43}
44
45/// Supported type classes.
46#[derive(Debug, Clone, Copy, PartialEq, Eq)]
47pub enum TypeClass {
48    /// Boolean equality (BEq)
49    BEq,
50    /// Hashing (Hashable)
51    Hashable,
52    /// Representation (Repr)
53    Repr,
54    /// String conversion (ToString)
55    ToString,
56    /// Ordering comparison (Ord)
57    Ord,
58}
59
60impl TypeClass {
61    /// Parse from identifier.
62    pub fn from_ident(ident: &syn::Ident) -> Option<Self> {
63        match ident.to_string().as_str() {
64            "BEq" => Some(TypeClass::BEq),
65            "Hashable" => Some(TypeClass::Hashable),
66            "Repr" => Some(TypeClass::Repr),
67            "ToString" => Some(TypeClass::ToString),
68            "Ord" => Some(TypeClass::Ord),
69            _ => None,
70        }
71    }
72
73    /// Get the required method name.
74    pub fn method_name(&self) -> &'static str {
75        match self {
76            TypeClass::BEq => "beq",
77            TypeClass::Hashable => "hash",
78            TypeClass::Repr => "repr",
79            TypeClass::ToString => "to_string",
80            TypeClass::Ord => "compare",
81        }
82    }
83
84    /// Get the FFI function prefix.
85    pub fn ffi_prefix(&self) -> &'static str {
86        match self {
87            TypeClass::BEq => "__lean_beq",
88            TypeClass::Hashable => "__lean_hash",
89            TypeClass::Repr => "__lean_repr",
90            TypeClass::ToString => "__lean_to_string",
91            TypeClass::Ord => "__lean_compare",
92        }
93    }
94}
95
96/// Build the `#[lean_instance]` macro output.
97pub fn build_lean_instance(
98    item: &mut ItemImpl,
99    options: LeanInstanceOptions,
100) -> syn::Result<TokenStream> {
101    let typeclass = TypeClass::from_ident(&options.typeclass).ok_or_else(|| {
102        syn::Error::new(
103            options.typeclass.span(),
104            format!(
105                "Unsupported type class '{}'. Supported: BEq, Hashable, Repr, ToString, Ord",
106                options.typeclass
107            ),
108        )
109    })?;
110
111    // Extract the struct name from the impl block
112    let struct_name = extract_struct_name(item)?;
113
114    // Find the required method
115    let method = find_method(item, typeclass.method_name())?;
116
117    // Generate the FFI function
118    let leo3_crate = get_leo3_crate(options.krate.as_ref());
119    let ffi_fn = generate_ffi_function(&leo3_crate, &struct_name, typeclass, method)?;
120
121    // Return original impl + generated FFI function
122    Ok(quote! {
123        #item
124
125        #ffi_fn
126    })
127}
128
129/// Extract the struct name from an impl block.
130fn extract_struct_name(item: &ItemImpl) -> syn::Result<syn::Ident> {
131    match &*item.self_ty {
132        syn::Type::Path(type_path) => {
133            if let Some(segment) = type_path.path.segments.last() {
134                Ok(segment.ident.clone())
135            } else {
136                Err(syn::Error::new_spanned(
137                    &item.self_ty,
138                    "Cannot determine struct name from impl block",
139                ))
140            }
141        }
142        _ => Err(syn::Error::new_spanned(
143            &item.self_ty,
144            "Expected a path type in impl block",
145        )),
146    }
147}
148
149/// Find a method in an impl block.
150fn find_method<'a>(item: &'a ItemImpl, name: &str) -> syn::Result<&'a syn::ImplItemFn> {
151    for impl_item in &item.items {
152        if let ImplItem::Fn(method) = impl_item {
153            if method.sig.ident == name {
154                return Ok(method);
155            }
156        }
157    }
158
159    Err(syn::Error::new_spanned(
160        item,
161        format!(
162            "Method '{}' not found in impl block. This method is required for the type class.",
163            name
164        ),
165    ))
166}
167
168/// Generate the FFI function for a type class.
169fn generate_ffi_function(
170    leo3_crate: &TokenStream,
171    struct_name: &syn::Ident,
172    typeclass: TypeClass,
173    _method: &syn::ImplItemFn,
174) -> syn::Result<TokenStream> {
175    let ffi_name = format_ident!("{}_{}", typeclass.ffi_prefix(), struct_name);
176    let try_name = format_ident!("__leo3_try_{}", ffi_name);
177    let method_ident = format_ident!("{}", typeclass.method_name());
178
179    match typeclass {
180        TypeClass::BEq => Ok(quote! {
181            /// Auto-generated BEq instance implementation.
182            #[no_mangle]
183            pub unsafe extern "C" fn #ffi_name(
184                a: *mut #leo3_crate::ffi::lean_object,
185                b: *mut #leo3_crate::ffi::lean_object,
186            ) -> *mut #leo3_crate::ffi::lean_object {
187                #leo3_crate::__private::ffi_panic_boundary(|| #try_name(a, b))
188            }
189
190            #[doc(hidden)]
191            #[allow(non_snake_case)]
192            pub(crate) unsafe fn #try_name(
193                a: *mut #leo3_crate::ffi::lean_object,
194                b: *mut #leo3_crate::ffi::lean_object,
195            ) -> #leo3_crate::LeanResult<*mut #leo3_crate::ffi::lean_object> {
196                let lean = #leo3_crate::Lean::assume_initialized();
197
198                let a_bound: #leo3_crate::LeanBound<'_, #leo3_crate::external::LeanExternalType<#struct_name>> =
199                    #leo3_crate::LeanBound::from_owned_ptr(lean, a);
200                let b_bound: #leo3_crate::LeanBound<'_, #leo3_crate::external::LeanExternalType<#struct_name>> =
201                    #leo3_crate::LeanBound::from_owned_ptr(lean, b);
202
203                let result = a_bound.get_ref().#method_ident(b_bound.get_ref());
204                Ok(#leo3_crate::ffi::inline::lean_box(result as usize))
205            }
206        }),
207
208        TypeClass::Hashable => Ok(quote! {
209            /// Auto-generated Hashable instance implementation.
210            #[no_mangle]
211            pub unsafe extern "C" fn #ffi_name(
212                obj: *mut #leo3_crate::ffi::lean_object,
213            ) -> u64 {
214                #leo3_crate::__private::scalar_u64_ffi_panic_boundary(stringify!(#ffi_name), || #try_name(obj))
215            }
216
217            #[doc(hidden)]
218            #[allow(non_snake_case)]
219            pub(crate) unsafe fn #try_name(
220                obj: *mut #leo3_crate::ffi::lean_object,
221            ) -> #leo3_crate::LeanResult<u64> {
222                let lean = #leo3_crate::Lean::assume_initialized();
223
224                let obj_bound: #leo3_crate::LeanBound<'_, #leo3_crate::external::LeanExternalType<#struct_name>> =
225                    #leo3_crate::LeanBound::from_owned_ptr(lean, obj);
226
227                Ok(obj_bound.get_ref().#method_ident())
228            }
229        }),
230
231        TypeClass::Repr => Ok(quote! {
232            /// Auto-generated Repr instance implementation.
233            #[no_mangle]
234            pub unsafe extern "C" fn #ffi_name(
235                obj: *mut #leo3_crate::ffi::lean_object,
236            ) -> *mut #leo3_crate::ffi::lean_object {
237                #leo3_crate::__private::ffi_panic_boundary(|| #try_name(obj))
238            }
239
240            #[doc(hidden)]
241            #[allow(non_snake_case)]
242            pub(crate) unsafe fn #try_name(
243                obj: *mut #leo3_crate::ffi::lean_object,
244            ) -> #leo3_crate::LeanResult<*mut #leo3_crate::ffi::lean_object> {
245                let lean = #leo3_crate::Lean::assume_initialized();
246
247                let obj_bound: #leo3_crate::LeanBound<'_, #leo3_crate::external::LeanExternalType<#struct_name>> =
248                    #leo3_crate::LeanBound::from_owned_ptr(lean, obj);
249
250                let repr_str: String = obj_bound.get_ref().#method_ident();
251                let s = #leo3_crate::types::LeanString::mk(lean, &repr_str)
252                    .map_err(|e| #leo3_crate::LeanError::Conversion(format!(
253                        "Failed to convert Rust repr result to Lean: {}",
254                        e
255                    )))?;
256                Ok(s.into_ptr())
257            }
258        }),
259
260        TypeClass::ToString => Ok(quote! {
261            /// Auto-generated ToString instance implementation.
262            #[no_mangle]
263            pub unsafe extern "C" fn #ffi_name(
264                obj: *mut #leo3_crate::ffi::lean_object,
265            ) -> *mut #leo3_crate::ffi::lean_object {
266                #leo3_crate::__private::ffi_panic_boundary(|| #try_name(obj))
267            }
268
269            #[doc(hidden)]
270            #[allow(non_snake_case)]
271            pub(crate) unsafe fn #try_name(
272                obj: *mut #leo3_crate::ffi::lean_object,
273            ) -> #leo3_crate::LeanResult<*mut #leo3_crate::ffi::lean_object> {
274                let lean = #leo3_crate::Lean::assume_initialized();
275
276                let obj_bound: #leo3_crate::LeanBound<'_, #leo3_crate::external::LeanExternalType<#struct_name>> =
277                    #leo3_crate::LeanBound::from_owned_ptr(lean, obj);
278
279                let s: String = obj_bound.get_ref().#method_ident();
280                let lean_str = #leo3_crate::types::LeanString::mk(lean, &s)
281                    .map_err(|e| #leo3_crate::LeanError::Conversion(format!(
282                        "Failed to convert Rust string result to Lean: {}",
283                        e
284                    )))?;
285                Ok(lean_str.into_ptr())
286            }
287        }),
288
289        TypeClass::Ord => Ok(quote! {
290            /// Auto-generated Ord instance implementation.
291            #[no_mangle]
292            pub unsafe extern "C" fn #ffi_name(
293                a: *mut #leo3_crate::ffi::lean_object,
294                b: *mut #leo3_crate::ffi::lean_object,
295            ) -> *mut #leo3_crate::ffi::lean_object {
296                #leo3_crate::__private::ffi_panic_boundary(|| #try_name(a, b))
297            }
298
299            #[doc(hidden)]
300            #[allow(non_snake_case)]
301            pub(crate) unsafe fn #try_name(
302                a: *mut #leo3_crate::ffi::lean_object,
303                b: *mut #leo3_crate::ffi::lean_object,
304            ) -> #leo3_crate::LeanResult<*mut #leo3_crate::ffi::lean_object> {
305                let lean = #leo3_crate::Lean::assume_initialized();
306
307                let a_bound: #leo3_crate::LeanBound<'_, #leo3_crate::external::LeanExternalType<#struct_name>> =
308                    #leo3_crate::LeanBound::from_owned_ptr(lean, a);
309                let b_bound: #leo3_crate::LeanBound<'_, #leo3_crate::external::LeanExternalType<#struct_name>> =
310                    #leo3_crate::LeanBound::from_owned_ptr(lean, b);
311
312                let ordering = a_bound.get_ref().#method_ident(b_bound.get_ref());
313                let ord_val = match ordering {
314                    std::cmp::Ordering::Less => 0usize,
315                    std::cmp::Ordering::Equal => 1usize,
316                    std::cmp::Ordering::Greater => 2usize,
317                };
318
319                Ok(#leo3_crate::ffi::inline::lean_box(ord_val))
320            }
321        }),
322    }
323}
324
325#[cfg(test)]
326mod tests {
327    use super::*;
328
329    #[test]
330    fn generated_lean_instance_wrappers_use_boundary_helpers() {
331        let mut item: syn::ItemImpl = syn::parse_quote! {
332            impl Demo {
333                fn repr(&self) -> String {
334                    "demo".to_string()
335                }
336            }
337        };
338
339        let repr_tokens = build_lean_instance(
340            &mut item,
341            LeanInstanceOptions {
342                typeclass: syn::parse_quote!(Repr),
343                krate: None,
344            },
345        )
346        .expect("lean_instance repr expansion should succeed");
347        let rendered_repr = repr_tokens.to_string();
348        assert!(rendered_repr.contains("__private :: ffi_panic_boundary"));
349        assert!(rendered_repr.contains("Failed to convert Rust repr result to Lean"));
350        assert!(!rendered_repr.contains("empty string"));
351        assert!(!rendered_repr.contains(".expect("));
352        assert!(!rendered_repr.contains(". expect ("));
353
354        let mut hash_item: syn::ItemImpl = syn::parse_quote! {
355            impl Demo {
356                fn hash(&self) -> u64 {
357                    7
358                }
359            }
360        };
361
362        let hash_tokens = build_lean_instance(
363            &mut hash_item,
364            LeanInstanceOptions {
365                typeclass: syn::parse_quote!(Hashable),
366                krate: None,
367            },
368        )
369        .expect("lean_instance hash expansion should succeed");
370        let rendered_hash = hash_tokens.to_string();
371        assert!(rendered_hash.contains("__private :: scalar_u64_ffi_panic_boundary"));
372        assert!(!rendered_hash.contains(".expect("));
373        assert!(!rendered_hash.contains(". expect ("));
374    }
375}