ice_rs/slice/
function.rs

1use crate::slice::types::IceType;
2use quote::{__private::TokenStream, quote};
3
4use super::{function_argument::FunctionArgument, function_return::FunctionReturn, function_throws::FunctionThrows};
5
6
7#[derive(Clone, Debug)]
8pub struct Function {
9    pub id: TokenStream,
10    pub ice_id: String,
11    pub return_type: FunctionReturn,
12    pub arguments: Vec<FunctionArgument>,
13    pub throws: FunctionThrows,
14    idempotent: bool
15}
16
17impl Function {
18    pub fn empty() -> Function {
19        Function {
20            id: TokenStream::new(),
21            ice_id: String::new(),
22            return_type: FunctionReturn::empty(),
23            arguments: Vec::new(),
24            throws: FunctionThrows::empty(),
25            idempotent: false,
26        }
27    }
28
29    pub fn set_idempotent(&mut self) {
30        self.idempotent = true;
31    }
32
33    pub fn add_argument(&mut self, arg: FunctionArgument) {
34        self.arguments.push(arg);
35    }
36
37    pub fn generate_decl(&self) -> Result<TokenStream, Box<dyn std::error::Error>> {       
38        let id_token = &self.id;
39        let return_token = self.return_type.token();
40        let mut arg_tokens = vec![
41            quote! { &mut self }
42        ];
43        arg_tokens.extend(self.arguments.iter().map(|arg| arg.token()).collect::<Vec<_>>());
44        arg_tokens.push(quote! {
45            context: Option<HashMap<String, String>>
46        });
47
48        Ok(quote! {
49            async fn #id_token (#(#arg_tokens),*) -> Result<#return_token, Box<dyn std::error::Error + Send + Sync>>;
50        })
51    }
52
53    pub fn generate_server_decl(&self) -> Result<TokenStream, Box<dyn std::error::Error>> {       
54        let id_token = &self.id;
55        let return_token = self.return_type.token();
56        let mut arg_tokens = vec![
57            quote! { &mut self }
58        ];
59        arg_tokens.extend(self.arguments.iter().map(|arg| arg.token()).collect::<Vec<_>>());
60        arg_tokens.push(quote! {
61            context: Option<HashMap<String, String>>
62        });
63
64        Ok(quote! {
65            async fn #id_token (#(#arg_tokens),*) -> #return_token;
66        })
67    }
68
69    // TODO: return token stream
70    pub fn generate_impl(&self) -> Result<TokenStream, Box<dyn std::error::Error>> {
71        let id_token = &self.id;
72        let return_token = self.return_type.token();
73        let mut arg_tokens = vec![
74            quote! { &mut self }
75        ];
76        arg_tokens.extend(self.arguments.iter().map(|arg| arg.token()).collect::<Vec<_>>());
77        arg_tokens.push(quote! {
78            context: Option<HashMap<String, String>>
79        });
80        let arg_require_result = self.arguments.iter().any(|arg| arg.out);
81        let arg_serialize_input_tokens = self.arguments.iter().map(|arg| arg.serialize_input()).collect::<Vec<_>>();
82        let arg_serialize_output_tokens = self.arguments.iter().map(|arg| arg.serialize_output()).collect::<Vec<_>>();
83        let require_result = arg_require_result || match self.return_type.r#type {
84            IceType::VoidType => false,
85            _ => true
86        };
87        let mut reply_token = None;
88        let mut read_token = None;
89        let throw_token = self.throws.token();
90        if require_result {
91            reply_token = Some(quote! {
92                let reply =    
93            });
94            read_token = Some(quote! {
95                let mut read_bytes: i32 = 0;  
96            })
97        }
98
99        let ice_id_token = &self.ice_id;
100        let mode = if self.idempotent { 1u8 } else { 0u8 };
101        let returned_token = self.return_type.return_token();
102        let bytes_token = if arg_serialize_input_tokens.len() > 0 {
103            quote! { let mut bytes = Vec::new(); }
104        } else {
105            quote! { let bytes = Vec::new(); }
106        };
107
108        Ok(quote! {
109            async fn #id_token (#(#arg_tokens),*) -> Result<#return_token, Box<dyn std::error::Error + Send + Sync>> {
110                #bytes_token
111                #(#arg_serialize_input_tokens)*
112                #reply_token self.proxy.dispatch::<#throw_token>(&String::from(#ice_id_token), #mode, &Encapsulation::from(bytes), context).await?;
113                #read_token
114                #(#arg_serialize_output_tokens)*
115                #returned_token
116            }
117        })
118    }
119
120    fn wrap_result(&self) -> TokenStream {
121        match self.return_type.r#type {
122            IceType::Optional(_, tag) => {
123                quote! { let wrapped_result = OptionalWrapper::new(#tag, result); }
124            },
125            _ => quote! { let wrapped_result = result; }
126        }
127    }
128
129    pub fn generate_server_handler(&self) -> Result<TokenStream, Box<dyn std::error::Error>> {
130        let ice_id_token = self.ice_id.clone();
131        let id_token = &self.id;
132        let wrapped_result = self.wrap_result();
133        
134        let func_call = if self.arguments.len() > 0 {
135            let mut arg_tokens = vec![];
136            arg_tokens.extend(self.arguments.iter().map(|arg| arg.pass_argument()).collect::<Vec<_>>());
137            // TODO: split non optionals and options as non optionals come first and optionals need special handling
138            let decoded_tokens = self.arguments.iter().map(|arg| arg.decode_request()).collect::<Vec<_>>();
139            let decoded_opt_tokens = self.arguments.iter().map(|arg| arg.decode_optional_request()).collect::<Vec<_>>();
140            let encoded_outputs = self.arguments.iter().map(|arg| arg.encode_output()).collect::<Vec<_>>();
141            let has_outputs = self.arguments.iter().any(|arg| arg.out);
142            let mut_token = if has_outputs {
143                quote!{ mut }
144            } else {
145                quote!{ }
146            }; 
147
148            quote!{
149                let mut read_bytes = 0;
150                #(#decoded_tokens)*
151                #(#decoded_opt_tokens)*
152                let result = self.server_impl.#id_token (#(#arg_tokens),*, None).await;
153                #wrapped_result
154                let #mut_token result = wrapped_result.to_bytes()?;
155                #(#encoded_outputs)*                
156            }
157        } else {
158            quote!{
159                let result = self.server_impl.#id_token (None).await;
160                #wrapped_result
161                let result = wrapped_result.to_bytes()?;
162            }
163        };
164
165        Ok(quote! {
166            #ice_id_token => {
167                #func_call
168                Ok(ReplyData {
169                    request_id: request.request_id,
170                    status: 0,
171                    body: Encapsulation::from(result)
172                })
173            },
174        })
175    }
176}