k3_wasm_macros/
lib.rs

1#![cfg_attr(feature = "nightly", feature(proc_macro_span))]
2
3use std::{env::current_dir, fs};
4
5use quote::{format_ident, quote, ToTokens, TokenStreamExt};
6use syn::{parse_macro_input, spanned::Spanned, Item};
7
8/// macro which gives wasm functionality for passing args to http_handler
9#[proc_macro]
10pub fn define_k3_write_inputs(_input: proc_macro::TokenStream) -> proc_macro::TokenStream {
11    let expanded = quote! {
12        #[no_mangle]
13        pub static mut ARG_PTR: *const u8 = std::ptr::null();
14
15        #[no_mangle]
16        pub static mut ARG_LEN: usize = 0;
17
18        #[no_mangle]
19        extern "C" fn __k3_write_inputs(ptr: *const u8, len: usize) -> *const u8 {
20            unsafe {
21                ARG_PTR = ptr;
22                ARG_LEN = len;
23                ARG_PTR
24            }
25        }
26    };
27
28    proc_macro::TokenStream::from(expanded)
29}
30
31#[proc_macro]
32pub fn init(_item: proc_macro::TokenStream) -> proc_macro::TokenStream {
33    let mut tokens = quote! {
34        #[no_mangle]
35        extern "C" fn __k3_alloc(size: u32) -> *mut u8 {
36            unsafe { std::alloc::alloc(std::alloc::Layout::array::<u8>(size as usize).unwrap()) }
37        }
38
39
40        #[no_mangle]
41        extern "C" fn __k3_init_env(ptr: *const u8, len: usize) -> *const u8 {
42            unsafe {
43                k3_wasm_sdk::ENV_VAR_PTR = ptr;
44                k3_wasm_sdk::ENV_VAR_LEN = len;
45                k3_wasm_sdk::ENV_VAR_PTR
46            }
47        }
48
49        #[no_mangle]
50        extern "C" fn __k3_results() -> *const u8 {
51            unsafe {
52                k3_wasm_sdk::RES_PTR
53            }
54        }
55
56        #[no_mangle]
57        extern "C" fn __k3_dealloc(ptr: *mut u8, size: u32) {
58            unsafe {
59                std::alloc::dealloc(ptr, std::alloc::Layout::array::<u8>(size as usize).unwrap())
60            }
61        }
62    };
63
64    let mut added_main = false;
65    let env_path = current_dir().unwrap().join(".env");
66    if env_path.exists() {
67        let env_string = fs::read_to_string(env_path).unwrap();
68        let env_string = syn::LitStr::new(&env_string, tokens.span());
69        tokens.append_all(quote! {
70            #[no_mangle]
71            pub static mut ENV_STRING_PTR: u32 = 0;
72
73            #[no_mangle]
74            pub static mut PROJECT_ID_PTR: u32 = 0;
75
76            fn main() {
77                // .env = PROJECT_ID=<project.id> from build in new version
78                // NOTE - retain ENV_STRING_PTR for writing later in execution
79                k3_wasm_sdk::dotenvy::from_read(#env_string.as_bytes()).unwrap();
80                unsafe {
81                    let mut env_string_buf = #env_string.len().to_le_bytes().to_vec();
82                    env_string_buf.append(&mut #env_string.as_bytes().to_vec());
83                    PROJECT_ID_PTR = env_string_buf.as_ptr() as u32;
84                    std::mem::forget(env_string_buf);
85                    // println!("[__k3_dbg] env string ptr: {}", ENV_STRING_PTR);
86                };
87                // println!("[__k3_dbg] env loaded: {}", #env_string.len());
88            }
89        });
90        added_main = true;
91    }
92
93    if !added_main {
94        tokens.append_all(quote! {fn main() {}});
95    }
96
97    tokens.into()
98}
99
100#[proc_macro_attribute]
101pub fn http_handler(
102    attr: proc_macro::TokenStream,
103    input: proc_macro::TokenStream,
104) -> proc_macro::TokenStream {
105    let input = parse_macro_input!(input as Item);
106
107    let attr_str = attr.to_string();
108
109    let has_env = attr_str.contains("Env");
110
111    if let Item::Fn(mut fn_decl) = input {
112        let name = fn_decl.sig.ident;
113
114
115        #[cfg(feature = "nightly")]
116        let name = {
117            let prefix = if attr.is_empty() || has_env {
118                let file = proc_macro::Span::call_site().source_file().path();
119
120                let file_name = file.file_name().unwrap().to_str().unwrap();
121                let segments = file
122                    .components()
123                    .map(|comp| comp.as_os_str().to_str().unwrap())
124                    .collect::<Vec<_>>();
125                if *segments.last().unwrap() == "main.rs" {
126                    "".to_string()
127                } else {
128                    let mut relevant_segments = segments
129                        .into_iter()
130                        .skip_while(|seg| *seg != "src")
131                        .skip(1)
132                        .map(|s| s.to_string())
133                        .collect::<Vec<_>>();
134                    if file_name == "mod.rs" {
135                        relevant_segments.pop();
136                    } else {
137                        *relevant_segments.last_mut().unwrap() =
138                            file_name[0..file_name.len() - 3].to_string();
139                    }
140                    format!("_{}", relevant_segments.join("_"))
141                }
142            } else {
143                // Allow overriding the prefix even when using nightly
144                syn::parse::<syn::Ident>(attr)
145                    .expect("Expected identifier as argument to http_handler")
146                    .to_string()
147            };
148            format_ident!("__k3_handler_{}_{}", prefix, name.to_string())
149        };
150
151        #[cfg(not(feature = "nightly"))]
152        let name = {
153            let prefix = syn::parse::<syn::Ident>(attr)
154                .expect("Expected identifier as argument to http_handler")
155                .to_string();
156            format_ident!("__k3_handler_{}_{}", prefix, name.to_string())
157        };
158
159        let impl_fn_name = format_ident!("__k3_{}_impl", name);
160        fn_decl.sig.ident = impl_fn_name.clone();
161        let fn_tokens = fn_decl.to_token_stream();
162
163        quote! {
164            #[no_mangle]
165            pub extern "C" fn #name(req_ptr: u32, req_len: u32) -> u32 {
166                #fn_tokens
167                if #has_env {
168                    k3_wasm_sdk::set_env_context();
169                };
170                let req = k3_wasm_sdk::request_from_raw_parts(req_ptr, req_len);
171                let res = #impl_fn_name(req);
172                let r = k3_wasm_sdk::encode_response(res);
173                unsafe {k3_wasm_sdk::RES_PTR = r as *const u8;};
174                r
175
176            }
177        }
178        .into()
179    } else {
180        panic!("Macro can only be used on functions")
181    }
182}
k3_wasm_macros/lib.rs

k3_wasm_macros/
lib.rs
