massbit_sol/generator/
handler.rs

1use crate::generator::graphql::MAPPING_RUST_TYPES_TO_DB;
2use crate::generator::Generator;
3use crate::schema::{Schema, Variant, VariantArray};
4use inflector::Inflector;
5use std::fmt::Write;
6
7const MODULES: &str = r#"
8use crate::generated::instruction::*;
9use crate::STORE;
10use massbit_solana_sdk::entity::{Attribute, Entity, Value};
11use massbit_solana_sdk::types::SolanaBlock;
12use serde_json;
13use solana_program::pubkey::Pubkey;
14use solana_transaction_status::TransactionWithStatusMeta;
15use std::collections::HashMap;
16use uuid::Uuid;
17"#;
18const ENTITY_SAVE: &str = r#"
19pub trait EntityExt {
20    fn save(&self, entity_name: &str);
21}
22impl EntityExt for Entity {
23    fn save(&self, entity_name: &str) {
24        unsafe {
25            STORE
26                .as_mut()
27                .unwrap()
28                .save(String::from(entity_name), self.clone());
29        }
30    }
31}
32"#;
33
34impl<'a> Generator<'a> {
35    pub fn generate_handler(&self, schema: &Schema) -> String {
36        let mut out = String::new();
37        let _ = writeln!(out, "{}", MODULES);
38        let _ = writeln!(out, "{}", ENTITY_SAVE);
39        if schema.name.is_some() && schema.variants.is_some() {
40            let name = schema.get_pascal_name(schema.name.as_ref().unwrap());
41            let patterns = self.expand_handler_patterns(&name, schema.variants.as_ref().unwrap());
42            let handler_functions =
43                self.expand_handler_functions(schema.variants.as_ref().unwrap());
44            let _ = write!(
45                &mut out,
46                r#"pub struct Handler {{}}
47                    impl Handler {{
48                        pub fn process(
49                            &self, 
50                            block: &SolanaBlock, 
51                            transaction: &TransactionWithStatusMeta, 
52                            program_id: &Pubkey, 
53                            accounts: &Vec<Pubkey>, 
54                            input: &[u8],
55                        ) {{
56                            println!("Process block {{}} with input {{:?}}", block.block_number, input);
57                            if let Some(instruction) = {name}::unpack(input) {{
58                                match instruction {{
59                                    {patterns}
60                                }}
61                            }}
62                        }}
63                        {handler_functions}
64                    }}"#,
65                name = name,
66                patterns = patterns.join(",\n"),
67                handler_functions = handler_functions.join("\n")
68            );
69        }
70        out
71    }
72    pub fn expand_handler_patterns(
73        &self,
74        enum_name: &String,
75        variants: &VariantArray,
76    ) -> Vec<String> {
77        variants
78            .iter()
79            .map(|variant| {
80                let method_name = format!("process_{}", &variant.name.to_snake_case());
81                let args = match &variant.inner_type {
82                    None => (String::default(), String::default()),
83                    Some(_) => (String::from("(arg)"), String::from(", arg")),
84                };
85                format!(
86                    r#"{enum_name}::{var_name}{var_inner} => {{
87                        self.{method_name}(block,transaction,program_id, accounts{arg});
88                    }}"#,
89                    enum_name = enum_name,
90                    var_name = &variant.name,
91                    method_name = method_name,
92                    var_inner = &args.0,
93                    arg = &args.1
94                )
95            })
96            .collect::<Vec<String>>()
97    }
98    pub fn expand_handler_functions(&self, variants: &VariantArray) -> Vec<String> {
99        variants
100            .iter()
101            .map(|variant| {
102                let function_name = format!("process_{}", &variant.name.to_snake_case());
103                let function_body = self.expand_function_body(variant);
104                let mut inner_arg = String::default();
105                if let Some(inner_type) = &variant.inner_type {
106                    let _ =  write!(&mut inner_arg, "arg: {}", inner_type);
107                };
108                let log = if let Some(_inner_type) = &variant.inner_type {
109                    format!(r#"println!("call function {} for handle incoming block {{}} with argument {{:?}}", block.block_number, &arg);"#, function_name)
110                } else {
111                    format!(r#"println!("call function {} for handle incoming block {{}}", block.block_number);"#, function_name)
112                };
113                format!(
114                    r#"pub fn {function_name}(
115                                &self,
116                                block: &SolanaBlock,
117                                transaction: &TransactionWithStatusMeta,
118                                program_id: &Pubkey,
119                                accounts: &Vec<Pubkey>,
120                                {inner_arg}
121                            ) -> Result<(), anyhow::Error> {{
122                                {log}
123                                {function_body}
124                            }}"#,
125                    function_name = function_name,
126                    log = log,
127                    function_body = function_body,
128                    inner_arg = inner_arg
129                )
130            })
131            .collect::<Vec<String>>()
132    }
133    pub fn expand_function_body(&self, variant: &Variant) -> String {
134        let mut assignments: Vec<String> = Vec::default();
135        //Account assigment
136        if let Some(accounts) = &variant.accounts {
137            for account in accounts {
138                assignments.push(format!(
139                    r#"map.insert("{}".to_string(), Value::from(
140                        accounts.get({})
141                            .and_then(|pubkey| Some(pubkey.to_string()))
142                            .unwrap_or_default()));"#,
143                    account.name, account.index
144                ));
145            }
146        }
147        // Write table if there is inner_type
148        if let Some(inner_type) = &variant.inner_type {
149            // Get definitions
150            if let Some(inner_schema) = self.definitions.get(inner_type.as_str()) {
151                // get a table corresponding to inner_schema
152                self.expand_entity_assignment(&mut assignments, inner_schema);
153            } else if MAPPING_RUST_TYPES_TO_DB.contains_key(inner_type.as_str()) {
154                //Inner type is primitive
155                self.expand_single_assignment(&mut assignments, "value", "arg");
156            }
157        }
158        format!(
159            r#"
160                let mut map : HashMap<Attribute, Value> = HashMap::default();
161                map.insert("id".to_string(), Value::from(Uuid::new_v4().to_simple().to_string()));
162                {assignments}
163                Entity::from(map).save("{entity_name}");
164                Ok(())
165            "#,
166            assignments = assignments.join("\n"),
167            entity_name = &variant.name
168        )
169    }
170    fn expand_single_assignment(
171        &self,
172        assignments: &mut Vec<String>,
173        field_name: &str,
174        field_value: &str,
175    ) {
176        assignments.push(format!(
177            r#"map.insert("{}".to_string(), Value::from({}));"#,
178            field_name, field_value
179        ));
180    }
181    pub fn expand_entity_assignment(&self, assignments: &mut Vec<String>, inner_schema: &Schema) {
182        //If inner schema is a struct
183        if let Some(properties) = &inner_schema.properties {
184            for property in properties {
185                let db_type = MAPPING_RUST_TYPES_TO_DB.get(property.data_type.as_str());
186                //.unwrap_or(&*DEFAULT_TYPE_DB);
187                // If data_type is not primitive (e.g. Enum, Struct)
188                match db_type {
189                    // If data_type is primitive (e.g. Enum, Struct)
190                    Some(db_type) => {
191                        let elm_value = if property.data_type.starts_with("NonZero") {
192                            format!("{}.get()", property.name)
193                        } else {
194                            format!("{}", property.name)
195                        };
196                        let property_value = match property.array_length {
197                            Some(_) => {
198                                format!(
199                                    r#"arg.{property_name}.iter().map(|&{property_name}| Value::from({elm_value})).collect::<Vec<Value>>()"#,
200                                    property_name = &property.name,
201                                    elm_value = elm_value
202                                )
203                            }
204                            None => {
205                                format!("arg.{}", elm_value)
206                            }
207                        };
208                        assignments.push(format!(
209                            r#"map.insert("{}".to_string(), Value::from({}));"#,
210                            &property.name, &property_value
211                        ));
212                    }
213                    None => {
214                        assignments.push(format!(
215                            r#"map.insert("{name}".to_string(), Value::from(serde_json::to_string(&arg.{name}).unwrap_or(Default::default())));"#,
216                            name=&property.name
217                        ));
218                    }
219                }
220            }
221        }
222    }
223}