wakflo_sdk_derive/
lib.rs

1extern crate proc_macro;
2use proc_macro::TokenStream;
3#[allow(unused)]
4use wakflo_common::TaskResult;
5#[allow(unused)]
6use wakflo_core::TaskContext;
7
8use quote::quote;
9use syn::{parse_macro_input, FnArg, ItemFn, ReturnType, Type};
10
11/// wakflo_plugin [procedural macro attribute][1].
12///
13/// [1]: https://doc.rust-lang.org/reference/procedural-macros.html#attribute-macros
14#[proc_macro_attribute]
15pub fn wakflo_plugin(_args: TokenStream, input: TokenStream) -> TokenStream {
16    let mut function = parse_macro_input!(input as ItemFn);
17    if matches!(function.vis, syn::Visibility::Public(..)) {
18        panic!("wakflo_sdk::wakflo_plugin expects a private function");
19    }
20
21    let name = &function.sig.ident;
22    let inputs = &mut function.sig.inputs;
23    let output = &mut function.sig.output;
24
25    if name != "execute" {
26        panic!("wakflo_sdk::wakflo_plugin expects the name of the function to be `execute`")
27    }
28    if inputs.is_empty() || inputs.len() > 1 {
29        panic!(
30            "wakflo_sdk::wakflo_plugin expects a function with one argument of type `PluginContext`"
31        );
32    }
33
34    let context_arg = inputs.first().expect("missing input");
35    match context_arg {
36        FnArg::Receiver(_) => {
37            panic!(
38                "wakflo_sdk::wakflo_plugin expects a function that returns wakflo_sdk::PluginContext"
39            );
40        }
41        FnArg::Typed(t) => {
42            if let Type::Path(p) = t.ty.as_ref() {
43                if let Some(t) = p.path.segments.last() {
44                    if t.ident != "TaskContext" {
45                        panic!("wakflo_sdk::wakflo_plugin expects a function that returns wakflo_sdk::PluginContext");
46                    }
47                } else {
48                    panic!("wakflo_sdk::wakflo_plugin expects a function that returns wakflo_sdk::PluginContext");
49                }
50            }
51        }
52    };
53
54    match output {
55        ReturnType::Default => {
56            panic!("wakflo_sdk::wakflo_plugin expects a return value of type `PluginResult`")
57        }
58        ReturnType::Type(_, t) => {
59            if let Type::Path(p) = t.as_ref() {
60                if let Some(t) = p.path.segments.last() {
61                    if t.ident != "TaskResult" {
62                        panic!("wakflo_sdk::wakflo_plugin expects a function that returns wakflo_sdk::PluginResult");
63                    }
64                } else {
65                    panic!("wakflo_sdk::wakflo_connector expects a function that returns wakflo_sdk::PluginResult");
66                }
67            }
68        }
69    }
70
71    let fn_name = &function.sig.ident;
72    let tokens = quote! {
73        #[repr(C)]
74        struct Slice {
75            len: usize,
76            ptr: *const u8,
77        }
78
79        #[no_mangle]
80        fn main(input_ptr: u32) -> Box<Slice> {
81            // Read object from memory
82            let ctx = unsafe {
83                let len = *(input_ptr as *const u32);
84                let bytes = (input_ptr + 4) as *const u8;
85                let slice = core::slice::from_raw_parts(bytes, len as usize);
86                let d = rmp_serde::from_slice::<JsonValue>(slice.clone());
87                wakflo_sdk::prelude::rmp_serde::from_slice::<TaskContext>(slice).unwrap()
88            };
89
90            let output = #fn_name(ctx).unwrap();
91            let output_bytes = wakflo_sdk::prelude::rmp_serde::to_vec_named(&output).unwrap().leak();
92
93            // return output to host
94            Box::new(Slice {
95                len: output_bytes.len(),
96                ptr: output_bytes.as_ptr(),
97            })
98        }
99
100        #function
101    };
102    tokens.into()
103}
wakflo_sdk_derive/lib.rs

wakflo_sdk_derive/
lib.rs
