light_macros/
lib.rs

1extern crate proc_macro;
2
3use proc_macro::TokenStream;
4use quote::quote;
5use syn::{parse_macro_input, parse_quote, ItemFn};
6
7mod pubkey;
8
9/// Converts a base58 encoded public key into a byte array.
10#[proc_macro]
11pub fn pubkey(input: TokenStream) -> TokenStream {
12    let args = parse_macro_input!(input as pubkey::PubkeyArgs);
13    pubkey::pubkey(args)
14        .unwrap_or_else(|err| err.to_compile_error())
15        .into()
16}
17
18#[proc_macro_attribute]
19pub fn heap_neutral(_: TokenStream, input: TokenStream) -> TokenStream {
20    let mut function = parse_macro_input!(input as ItemFn);
21
22    // Insert memory management code at the beginning of the function
23    let init_code: syn::Stmt = parse_quote! {
24        #[cfg(target_os = "solana")]
25        let pos = light_heap::GLOBAL_ALLOCATOR.get_heap_pos();
26    };
27    let msg = format!("pre: {}", function.sig.ident);
28    let log_pre: syn::Stmt = parse_quote! {
29        #[cfg(all(target_os = "solana", feature = "mem-profiling"))]
30        light_heap::GLOBAL_ALLOCATOR.log_total_heap(#msg);
31    };
32    function.block.stmts.insert(0, init_code);
33    function.block.stmts.insert(1, log_pre);
34
35    // Insert memory management code at the end of the function
36    let msg = format!("post: {}", function.sig.ident);
37    let log_post: syn::Stmt = parse_quote! {
38        #[cfg(all(target_os = "solana", feature = "mem-profiling"))]
39        light_heap::GLOBAL_ALLOCATOR.log_total_heap(#msg);
40    };
41    let cleanup_code: syn::Stmt = parse_quote! {
42        #[cfg(target_os = "solana")]
43        light_heap::GLOBAL_ALLOCATOR.free_heap(pos)?;
44    };
45    let len = function.block.stmts.len();
46    function.block.stmts.insert(len - 1, log_post);
47    function.block.stmts.insert(len - 1, cleanup_code);
48    TokenStream::from(quote! { #function })
49}