ree_exchange_sdk_macro/
lib.rs

1use proc_macro::TokenStream;
2use quote::{ToTokens, format_ident, quote};
3use std::collections::BTreeMap;
4use syn::{Attribute, Ident, ItemMod, parse_macro_input, parse_quote, visit_mut::VisitMut};
5
6#[derive(Clone)]
7struct CanisterVisitor {
8    actions: BTreeMap<String, (String, bool)>,
9    pools: Option<Ident>,
10    hook_present: bool,
11    upgrade_declared: bool,
12    storages: BTreeMap<u8, (proc_macro2::TokenStream, proc_macro2::TokenStream)>,
13}
14
15mod keywords {
16    syn::custom_keyword!(exchange);
17    syn::custom_keyword!(upgrade);
18    syn::custom_keyword!(pools);
19    syn::custom_keyword!(hook);
20    syn::custom_keyword!(storage);
21    syn::custom_keyword!(action);
22    syn::custom_keyword!(memory);
23    syn::custom_keyword!(name);
24}
25
26struct StorageDeclAttr {
27    memory_id: u8,
28}
29
30impl syn::parse::Parse for StorageDeclAttr {
31    fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> {
32        input.parse::<syn::Token![#]>()?;
33        let content;
34        syn::bracketed!(content in input);
35        content.parse::<keywords::storage>()?;
36        let inside;
37        syn::parenthesized!(inside in content);
38        let lookahead = inside.lookahead1();
39        if lookahead.peek(keywords::memory) {
40            let _ = inside.parse::<keywords::memory>()?;
41            let _ = inside.parse::<syn::Token![=]>()?;
42            let lit: syn::LitInt = inside.parse()?;
43            let memory_id = lit.base10_parse::<u8>()?;
44            if memory_id >= 100 {
45                return Err(syn::Error::new_spanned(
46                    lit,
47                    "Memory id must be between 0 and 99",
48                ));
49            }
50            Ok(Self { memory_id })
51        } else {
52            let lit: syn::LitInt = inside.parse()?;
53            let memory_id = lit.base10_parse::<u8>()?;
54            if memory_id >= 100 {
55                return Err(syn::Error::new_spanned(
56                    lit,
57                    "Memory id must be between 0 and 99",
58                ));
59            }
60            Ok(Self { memory_id })
61        }
62    }
63}
64
65enum ActionDeclAttr {
66    Named { value: syn::LitStr },
67    Unnamed,
68}
69
70impl syn::parse::Parse for ActionDeclAttr {
71    fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> {
72        input.parse::<syn::Token![#]>()?;
73        let content;
74        syn::bracketed!(content in input);
75        content.parse::<keywords::action>()?;
76        if content.is_empty() {
77            return Ok(Self::Unnamed);
78        }
79        let inside;
80        syn::parenthesized!(inside in content);
81        let lookahead = inside.lookahead1();
82        if lookahead.peek(keywords::name) {
83            let _ = inside.parse::<keywords::name>()?;
84            let _ = inside.parse::<syn::Token![=]>()?;
85            Ok(Self::Named {
86                value: inside.parse()?,
87            })
88        } else if lookahead.peek(syn::LitStr) {
89            Ok(Self::Named {
90                value: inside.parse()?,
91            })
92        } else {
93            Err(lookahead.error())
94        }
95    }
96}
97
98impl CanisterVisitor {
99    fn new() -> Self {
100        CanisterVisitor {
101            actions: BTreeMap::new(),
102            pools: None,
103            hook_present: false,
104            upgrade_declared: false,
105            storages: BTreeMap::new(),
106        }
107    }
108
109    fn resolve_pools(&mut self, ty: &syn::ItemStruct) {
110        let mark_pools = ty.attrs.iter().find(|a| a.path().is_ident("pools"));
111        if mark_pools.is_none() {
112            return;
113        }
114        if self.pools.is_some() {
115            panic!("Only one struct can have the #[pools] attribute");
116        }
117        self.pools = Some(ty.ident.clone());
118    }
119
120    fn resolve_action(&mut self, attr: &Attribute, func: &syn::ItemFn) {
121        let is_action = attr.path().is_ident("action");
122        if !is_action {
123            return;
124        }
125        let tokens = attr.to_token_stream();
126        let action_decl =
127            syn::parse2::<ActionDeclAttr>(tokens).expect("Failed to parse action attribute");
128        match action_decl {
129            ActionDeclAttr::Unnamed => {
130                self.actions.insert(
131                    func.sig.ident.to_string(),
132                    (func.sig.ident.to_string(), func.sig.asyncness.is_some()),
133                );
134            }
135            ActionDeclAttr::Named { value, .. } => {
136                let action = value.value();
137                self.actions.insert(
138                    action,
139                    (func.sig.ident.to_string(), func.sig.asyncness.is_some()),
140                );
141            }
142        }
143    }
144
145    fn resolve_storage(&mut self, attr: &Attribute, ty: &syn::ItemType) {
146        let is_storage = attr.path().is_ident("storage");
147        if !is_storage {
148            return;
149        }
150        let tokens = attr.to_token_stream();
151        let storage_decl =
152            syn::parse2::<StorageDeclAttr>(tokens).expect("Failed to parse storage attribute");
153        let id = storage_decl.memory_id;
154        let storage_name = to_upper_snake_case(&ty.ident.to_string());
155        let storage_name = format_ident!("__{}", storage_name);
156        let storage_ty = format_ident!("{}", ty.ident);
157        let ic_ty = quote! { <#storage_ty as ::ree_exchange_sdk::store::StorageType>::Type };
158        let decl = quote! {
159            static #storage_name: ::core::cell::RefCell<#ic_ty> = ::core::cell::RefCell::new(
160                <#storage_ty as ::ree_exchange_sdk::store::StorageType>::init(
161                    __MEMORY_MANAGER.with(|m| m.borrow().get(::ic_stable_structures::memory_manager::MemoryId::new(#id))),
162                )
163            );
164        };
165        let access = quote! {
166            impl __CustomStorageAccess<#storage_ty> for #storage_ty {
167                fn with<F, R>(f: F) -> R
168                where
169                    F: FnOnce(&#ic_ty) -> R,
170                {
171                    #storage_name.with(|s| {
172                        let s = s.borrow();
173                        let r = <::std::cell::Ref<'_, #ic_ty> as ::std::ops::Deref>::deref(&s);
174                        f(r)
175                    })
176                }
177
178                fn with_mut<F, R>(f: F) -> R
179                where
180                    F: FnOnce(&mut #ic_ty) -> R,
181                {
182                    #storage_name.with(|s| {
183                        let mut s = s.borrow_mut();
184                        let r = <::std::cell::RefMut<'_, #ic_ty> as ::std::ops::DerefMut>::deref_mut(&mut s);
185                        f(r)
186                    })
187                }
188            }
189        };
190        if let Some(_) = self.storages.insert(id, (decl, access)) {
191            panic!("Memory id {} is already used", id);
192        }
193    }
194}
195
196fn to_upper_snake_case(s: &str) -> String {
197    let mut snake_case = String::new();
198    for (i, ch) in s.chars().enumerate() {
199        if ch.is_uppercase() {
200            if i != 0 {
201                snake_case.push('_');
202            }
203            snake_case.push(ch);
204        } else {
205            snake_case.push(ch.to_ascii_uppercase());
206        }
207    }
208    snake_case
209}
210
211impl VisitMut for CanisterVisitor {
212    fn visit_item_fn_mut(&mut self, item: &mut syn::ItemFn) {
213        for attr in item.attrs.iter() {
214            self.resolve_action(&attr, item);
215        }
216        syn::visit_mut::visit_item_fn_mut(self, item);
217    }
218
219    fn visit_item_struct_mut(&mut self, item: &mut syn::ItemStruct) {
220        self.resolve_pools(item);
221        syn::visit_mut::visit_item_struct_mut(self, item);
222    }
223
224    fn visit_item_impl_mut(&mut self, item: &mut syn::ItemImpl) {
225        if let Some(_attr) = item.attrs.iter().find(|a| a.path().is_ident("hook")) {
226            self.hook_present = true;
227        }
228        if let Some(_attr) = item.attrs.iter().find(|a| a.path().is_ident("upgrade")) {
229            self.upgrade_declared = true;
230        }
231        syn::visit_mut::visit_item_impl_mut(self, item);
232    }
233
234    fn visit_item_type_mut(&mut self, item: &mut syn::ItemType) {
235        for attr in item.attrs.iter() {
236            self.resolve_storage(&attr, item);
237        }
238        syn::visit_mut::visit_item_type_mut(self, item);
239    }
240}
241
242/// REE exchange entrypoint.
243#[proc_macro_attribute]
244pub fn exchange(_attr: TokenStream, item: TokenStream) -> TokenStream {
245    let mut input_mod = parse_macro_input!(item as ItemMod);
246    let mut visitor = CanisterVisitor::new();
247    visitor.visit_item_mod_mut(&mut input_mod);
248    if visitor.pools.is_none() {
249        panic!("#[pools] not found within the exchange mod");
250    }
251    let (storage_decl, storage_access): (
252        Vec<proc_macro2::TokenStream>,
253        Vec<proc_macro2::TokenStream>,
254    ) = visitor.storages.into_values().unzip();
255    let pools = visitor.pools.clone().unwrap();
256    if let Some((_, ref mut items)) = input_mod.content {
257        let branch = visitor
258            .actions
259            .iter()
260            .map(|(action, (func, is_async))| {
261                let call = format_ident!("{}", func);
262                if *is_async {
263                    quote! { #action => #call(&psbt, args).await, }
264                } else {
265                    quote! { #action => #call(&psbt, args), }
266                }
267            })
268            .collect::<Vec<_>>();
269
270        if !visitor.hook_present {
271            items.push(parse_quote! {
272                impl ::ree_exchange_sdk::Hook for #pools {}
273            });
274        }
275
276        items.push(parse_quote! {
277            impl ::ree_exchange_sdk::PoolStorageAccess<#pools> for #pools {
278                fn block_state() -> ::std::option::Option<<#pools as ::ree_exchange_sdk::Pools>::BlockState> {
279                    self::__GLOBAL_STATE.with_borrow(|p| p.last_key_value().map(|(k, v)| v.inner))
280                }
281
282                fn commit(height: u32, block_state: <#pools as ::ree_exchange_sdk::Pools>::BlockState) -> ::std::result::Result<(), ::std::string::String> {
283                    self::__GLOBAL_STATE.with_borrow_mut(|p| {
284                        if p.contains_key(&height) {
285                            ::std::result::Result::Err(format!("Block state for height {} already exists", height))
286                        } else {
287                            p.insert(height, ::ree_exchange_sdk::GlobalStateWrapper { inner: block_state });
288                            ::std::result::Result::Ok(())
289                        }
290                    })
291                }
292
293                fn get(address: &::std::string::String) -> ::std::option::Option<::ree_exchange_sdk::Pool<<#pools as ::ree_exchange_sdk::Pools>::PoolState>> {
294                    self::__CURRENT_POOLS.with_borrow(|p| p.get(address))
295                }
296
297                fn insert(pool: ::ree_exchange_sdk::Pool<<#pools as ::ree_exchange_sdk::Pools>::PoolState>) {
298                    self::__CURRENT_POOLS.with_borrow_mut(|p| {
299                        p.insert(pool.metadata().address.clone(), pool);
300                    });
301                }
302
303                fn remove(address: &::std::string::String) -> ::std::option::Option<::ree_exchange_sdk::Pool<<#pools as ::ree_exchange_sdk::Pools>::PoolState>> {
304                    self::__CURRENT_POOLS.with_borrow_mut(|p| {
305                        p.remove(address)
306                    })
307                }
308
309                fn iter() -> ::ree_exchange_sdk::iter::PoolIterator<#pools> {
310                    let memory = __MEMORY_MANAGER.with(|m| m.borrow().get(::ic_stable_structures::memory_manager::MemoryId::new(
311                        <#pools as ::ree_exchange_sdk::Pools>::POOL_STATE_MEMORY
312                    )));
313                    ::ree_exchange_sdk::iterator::<#pools>(memory)
314                }
315            }
316        });
317
318        items.push(parse_quote! {
319            #[::ic_cdk::update]
320            pub async fn execute_tx(args: ::ree_exchange_sdk::types::exchange_interfaces::ExecuteTxArgs) -> ::core::result::Result<String, String> {
321                ::ree_exchange_sdk::ensure_access::<#pools>()?;
322                let mut psbt = args.psbt()?;
323                let args = <::ree_exchange_sdk::ActionArgs as ::std::convert::From<_>>::from(args);
324                let pool_address = args.intention.pool_address.clone();
325                let _guard = self::__ExecuteTxGuard::new(pool_address.clone())
326                    .ok_or(::ree_exchange_sdk::error::Error::PoolBeingExecuted.to_string())?;
327                let txid = args.txid.clone();
328                let inputs = args.intention.pool_outpoints()
329                    .map_err(|_| ::ree_exchange_sdk::error::Error::IllegalPsbt.to_string())?;
330                let action = args.intention.action.clone();
331                let mut pool = self::__CURRENT_POOLS.with_borrow(|pools| pools.get(&pool_address)).ok_or(::ree_exchange_sdk::error::Error::PoolNotFound.to_string())?;
332                let new_nonce = args.intention.nonce;
333                if args.is_reapply {
334                    pool.truncate(new_nonce)?;
335                    self::__CURRENT_POOLS.with_borrow_mut(|pools| {
336                        pools.insert(pool_address.clone(), pool.clone());
337                    });
338                }
339                let pool_info = pool.get_pool_info();
340                if pool_info.nonce + 1 != new_nonce {
341                    return ::core::result::Result::<String, String>::Err(::ree_exchange_sdk::error::Error::NonceExpired.to_string());
342                }
343                let result: ::ree_exchange_sdk::ActionResult::<<#pools as ::ree_exchange_sdk::Pools>::PoolState> = match action.as_str() {
344                    #(#branch)*
345                    _ => ::ree_exchange_sdk::ActionResult::<<#pools as ::ree_exchange_sdk::Pools>::PoolState>::Err(::ree_exchange_sdk::error::Error::UnknownAction),
346                };
347                match result {
348                    ::ree_exchange_sdk::ActionResult::<<#pools as ::ree_exchange_sdk::Pools>::PoolState>::Ok(mut r) => {
349                        ::ree_exchange_sdk::schnorr::sign_p2tr_inputs(
350                            &mut psbt,
351                            &inputs,
352                            <#pools as ::ree_exchange_sdk::Pools>::network(),
353                            pool.metadata().key_derivation_path.clone(),
354                        ).await?;
355                        r.set_nonce(new_nonce);
356                        pool.states_mut().push(r);
357                        self::__CURRENT_POOLS.with_borrow_mut(|pools| {
358                            pools.insert(pool_address.clone(), pool);
359                        });
360                        self::__TX_RECORDS.with_borrow_mut(|unconfirmed| {
361                            let mut record = unconfirmed.get(&txid).unwrap_or(::ree_exchange_sdk::types::TxRecord {
362                                txid: txid.clone(),
363                                pools: vec![pool_address.clone()],
364                            });
365                            if !record.pools.contains(&pool_address) {
366                                record.pools.push(pool_address.clone());
367                            }
368                            unconfirmed.insert(txid, record);
369                        });
370                        ::core::result::Result::<String, String>::Ok(psbt.serialize_hex())
371                    }
372                    ::ree_exchange_sdk::ActionResult::<<#pools as ::ree_exchange_sdk::Pools>::PoolState>::Err(e) => {
373                        ::core::result::Result::<String, String>::Err(e.to_string())
374                    }
375                }
376            }
377        });
378
379        items.push(parse_quote! {
380            #[::ic_cdk::query]
381            pub fn get_pool_list() -> ::ree_exchange_sdk::types::exchange_interfaces::GetPoolListResponse {
382                self::__CURRENT_POOLS.with_borrow(|pools| {
383                    pools.iter()
384                        .map(|e| e.into_pair())
385                        .map(|(_, p)| p.get_pool_basic())
386                        .collect::<Vec<_>>()
387                })
388            }
389        });
390
391        items.push(parse_quote! {
392            #[::ic_cdk::query]
393            pub fn get_pool_info(
394                args: ::ree_exchange_sdk::types::exchange_interfaces::GetPoolInfoArgs,
395            ) -> ::ree_exchange_sdk::types::exchange_interfaces::GetPoolInfoResponse {
396                self::__CURRENT_POOLS.with_borrow(|pools| {
397                    pools.get(&args.pool_address).map(|p| p.get_pool_info())
398                })
399            }
400        });
401
402        items.push(parse_quote! {
403            #[::ic_cdk::update]
404            pub fn rollback_tx(
405                args: ::ree_exchange_sdk::types::exchange_interfaces::RollbackTxArgs,
406            ) -> ::ree_exchange_sdk::types::exchange_interfaces::RollbackTxResponse {
407                ::ree_exchange_sdk::ensure_access::<#pools>()?;
408                self::__TX_RECORDS.with_borrow_mut(|transactions| {
409                    self::__CURRENT_POOLS.with_borrow_mut(|pools| {
410                        ::ree_exchange_sdk::states::reject_tx::<#pools>(transactions, pools, args)
411                    })
412                })
413            }
414        });
415
416        items.push(parse_quote! {
417            #[::ic_cdk::update]
418            pub fn new_block(
419                args: ::ree_exchange_sdk::types::exchange_interfaces::NewBlockArgs,
420            ) -> ::ree_exchange_sdk::types::exchange_interfaces::NewBlockResponse {
421                ::ree_exchange_sdk::ensure_access::<#pools>()?;
422                let block = self::__TX_RECORDS.with_borrow_mut(|unconfirmed| {
423                    self::__BLOCKS.with_borrow_mut(|blocks| {
424                        self::__GLOBAL_STATE.with_borrow_mut(|state| {
425                            ::ree_exchange_sdk::states::confirm_txs::<#pools>(
426                                state,
427                                blocks,
428                                unconfirmed,
429                                args,
430                            )
431                        })
432                    })
433                })?;
434                if let Some(block) = block {
435                    self::__CURRENT_POOLS.with_borrow_mut(|pools| {
436                        self::__BLOCKS.with_borrow_mut(|blocks| {
437                            self::__GLOBAL_STATE.with_borrow_mut(|state| {
438                                ::ree_exchange_sdk::states::accept_block::<#pools>(
439                                    state,
440                                    blocks,
441                                    pools,
442                                    block.clone(),
443                                )
444                            })
445                        })
446                    })?;
447                    <#pools as ::ree_exchange_sdk::Hook>::on_block_confirmed(block);
448                }
449                Ok(())
450            }
451        });
452
453        items.push(parse_quote! {
454            struct __ExecuteTxGuard(::std::string::String);
455        });
456
457        items.push(parse_quote! {
458            impl __ExecuteTxGuard {
459                pub fn new(pool_address: ::std::string::String) -> ::std::option::Option<Self> {
460                    __GUARDS.with(|guards| {
461                        if guards.borrow().contains(&pool_address) {
462                            return None;
463                        }
464                        guards.borrow_mut().insert(pool_address.clone());
465                        return Some(__ExecuteTxGuard(pool_address));
466                    })
467                }
468            }
469        });
470
471        items.push(parse_quote! {
472            impl ::std::ops::Drop for __ExecuteTxGuard {
473                fn drop(&mut self) {
474                    __GUARDS.with_borrow_mut(|guards| {
475                        guards.remove(&self.0);
476                    });
477                }
478            }
479        });
480
481        items.push(parse_quote! {
482            thread_local! {
483                static __MEMORY_MANAGER: ::core::cell::RefCell<
484                    ::ic_stable_structures::memory_manager::MemoryManager<
485                        ::ic_stable_structures::DefaultMemoryImpl
486                    >
487                > = ::core::cell::RefCell::new(
488                    ::ic_stable_structures::memory_manager::MemoryManager::init(
489                        <::ic_stable_structures::DefaultMemoryImpl as core::default::Default>::default()
490                    )
491                );
492
493                static __GUARDS: ::core::cell::RefCell<::std::collections::HashSet<::std::string::String>> =
494                    ::core::cell::RefCell::new(::std::collections::HashSet::new());
495
496                static __BLOCKS: ::core::cell::RefCell<
497                    ::ic_stable_structures::StableBTreeMap<
498                        u32,
499                        ::ree_exchange_sdk::Block,
500                        ::ic_stable_structures::memory_manager::VirtualMemory<::ic_stable_structures::DefaultMemoryImpl>
501                    >
502                > = ::core::cell::RefCell::new(
503                    ::ic_stable_structures::StableBTreeMap::init(
504                        __MEMORY_MANAGER.with(|m| m.borrow().get(::ic_stable_structures::memory_manager::MemoryId::new(
505                            100
506                        ))),
507                    )
508                );
509                static __TX_RECORDS: ::core::cell::RefCell<
510                    ::ic_stable_structures::StableBTreeMap<
511                        ::ree_exchange_sdk::types::Txid,
512                        ::ree_exchange_sdk::types::TxRecord,
513                        ::ic_stable_structures::memory_manager::VirtualMemory<::ic_stable_structures::DefaultMemoryImpl>
514                    >
515                > = ::core::cell::RefCell::new(
516                    ::ic_stable_structures::StableBTreeMap::init(
517                        __MEMORY_MANAGER.with(|m| m.borrow().get(::ic_stable_structures::memory_manager::MemoryId::new(
518                            101
519                        ))),
520                    )
521                );
522                static __CURRENT_POOLS: ::core::cell::RefCell<
523                    ::ic_stable_structures::StableBTreeMap<
524                        ::std::string::String,
525                        ::ree_exchange_sdk::Pool<
526                            <#pools as ::ree_exchange_sdk::Pools>::PoolState
527                        >,
528                        ::ic_stable_structures::memory_manager::VirtualMemory<::ic_stable_structures::DefaultMemoryImpl>
529                    >
530                > = ::core::cell::RefCell::new(
531                    ::ic_stable_structures::StableBTreeMap::init(
532                        __MEMORY_MANAGER.with(|m| m.borrow().get(::ic_stable_structures::memory_manager::MemoryId::new(
533                            <#pools as ::ree_exchange_sdk::Pools>::POOL_STATE_MEMORY
534                        ))),
535                    )
536                );
537                static __GLOBAL_STATE: ::core::cell::RefCell<
538                    ::ic_stable_structures::StableBTreeMap<
539                        u32,
540                        ::ree_exchange_sdk::GlobalStateWrapper<<#pools as ::ree_exchange_sdk::Pools>::BlockState>,
541                        ::ic_stable_structures::memory_manager::VirtualMemory<::ic_stable_structures::DefaultMemoryImpl>
542                    >
543                > = ::core::cell::RefCell::new(
544                    ::ic_stable_structures::StableBTreeMap::init(
545                        __MEMORY_MANAGER.with(|m| m.borrow().get(::ic_stable_structures::memory_manager::MemoryId::new(
546                            <#pools as ::ree_exchange_sdk::Pools>::BLOCK_STATE_MEMORY
547                        ))),
548                    )
549                );
550                #(#storage_decl)*
551            }
552        });
553
554        if visitor.upgrade_declared {
555            items.push(parse_quote! {
556                impl #pools {
557                    pub fn upgrade() {
558                        let pool_id = <#pools as ::ree_exchange_sdk::Upgrade<#pools>>::POOL_STATE_MEMORY;
559                        if pool_id >= 100 {
560                            panic!("Memory id for pool state upgrade must be between 0 and 99");
561                        }
562                        let block_id = <#pools as ::ree_exchange_sdk::Upgrade<#pools>>::BLOCK_STATE_MEMORY;
563                        if block_id >= 100 {
564                            panic!("Memory id for block state upgrade must be between 0 and 99");
565                        }
566                        let memory_id = ::ic_stable_structures::memory_manager::MemoryId::new(pool_id);
567                        let memory = __MEMORY_MANAGER.with(|m| m.borrow().get(memory_id));
568                        let mut pool_storage = ::ic_stable_structures::StableBTreeMap::<
569                            ::std::string::String,
570                            ::ree_exchange_sdk::Pool<<#pools as ::ree_exchange_sdk::Upgrade<#pools>>::PoolState>,
571                            ::ic_stable_structures::memory_manager::VirtualMemory<::ic_stable_structures::DefaultMemoryImpl>,
572                        >::init(memory);
573                        self::__CURRENT_POOLS.with_borrow_mut(|pools| {
574                            for entry in pool_storage.iter() {
575                                let old_pool = entry.value();
576                                let states = old_pool.states()
577                                    .iter()
578                                    .map(|s| <<#pools as ::ree_exchange_sdk::Upgrade<#pools>>::PoolState as ::std::clone::Clone>::clone(s))
579                                    .map(|s| <<#pools as ::ree_exchange_sdk::Upgrade<#pools>>::PoolState as ::std::convert::Into<<#pools as ::ree_exchange_sdk::Pools>::PoolState>>::into(s))
580                                    .collect::<Vec<<#pools as ::ree_exchange_sdk::Pools>::PoolState>>();
581                                let mut new_pool = ::ree_exchange_sdk::Pool::new(
582                                    old_pool.metadata().clone(),
583                                );
584                                new_pool.states_mut().extend(states);
585                                pools.insert(entry.key().clone(), new_pool);
586                            }
587                        });
588
589                        let memory_id = ::ic_stable_structures::memory_manager::MemoryId::new(block_id);
590                        let memory = __MEMORY_MANAGER.with(|m| m.borrow().get(memory_id));
591                        let mut block_storage = ::ic_stable_structures::StableBTreeMap::<
592                            u32,
593                            ::ree_exchange_sdk::GlobalStateWrapper<<#pools as ::ree_exchange_sdk::Upgrade<#pools>>::BlockState>,
594                            ::ic_stable_structures::memory_manager::VirtualMemory<::ic_stable_structures::DefaultMemoryImpl>,
595                        >::init(memory);
596                        self::__GLOBAL_STATE.with_borrow_mut(|blocks| {
597                            for entry in block_storage.iter() {
598                                let old_block = entry.value().inner;
599                                let height = *entry.key();
600                                let new_block = <<#pools as ::ree_exchange_sdk::Upgrade<#pools>>::BlockState as ::std::convert::Into<<#pools as ::ree_exchange_sdk::Pools>::BlockState>>::into(old_block);
601                                blocks.insert(height, ::ree_exchange_sdk::GlobalStateWrapper { inner: new_block });
602                            }
603                        });
604                        pool_storage.clear_new();
605                        block_storage.clear_new();
606                    }
607                }
608            });
609        }
610
611        items.push(parse_quote! {
612            pub trait __CustomStorageAccess<S: ::ree_exchange_sdk::store::StorageType> {
613                fn with<F, R>(f: F) -> R
614                where
615                    F: FnOnce(&S::Type) -> R;
616                fn with_mut<F, R>(f: F) -> R
617                where
618                    F: FnOnce(&mut S::Type) -> R;
619            }
620        });
621        for access in storage_access {
622            items.push(parse_quote! {
623                #access
624            });
625        }
626    }
627    quote! {
628        #input_mod
629    }
630    .into()
631}
632
633/// Action entrypoint. The macro could be
634/// `#[action(name = "my_action")]` or `#[action("my_action")]` or `#[action]`.
635/// The functions shall have signature `fn(&bitcoin::Psbt, ActionArgs) -> ActionResult<Pools::PoolState>`
636#[proc_macro_attribute]
637pub fn action(_attr: TokenStream, item: TokenStream) -> TokenStream {
638    item
639}
640
641/// Pools definition
642#[proc_macro_attribute]
643pub fn pools(_attr: TokenStream, item: TokenStream) -> TokenStream {
644    item
645}
646
647/// Storage definition
648/// ```rust
649/// #[storage(memory = 3)]
650/// pub type MyStorage = ree_exchange_sdk::store::StableBTreeMap<String, String>;
651/// ```
652#[proc_macro_attribute]
653pub fn storage(_attr: TokenStream, item: TokenStream) -> TokenStream {
654    item
655}
656
657/// Optional hook for `Pools`. It should be marked on the `Hook` impl block of the `Pools` struct.
658/// ```rust
659/// #[hook]
660/// impl Hook for MyPools {
661///     fn on_block_confirmed(block: Block) {
662///         for tx in block.txs {
663///             // for each pool affected by this tx
664///             for addr in tx.pools {
665///                 let pool = DummyPools::get(&addr).unwrap();
666///                 // load the state and do something with it
667///                 let _state = pool.get(tx.txid).unwrap();
668///             }
669///         }
670///         // update block state
671///         MyPools::commit(block.block_height, block.block_height);
672///     }
673/// }
674/// ```
675#[proc_macro_attribute]
676pub fn hook(_attr: TokenStream, item: TokenStream) -> TokenStream {
677    item
678}
679
680/// Upgrade attribute for pool state migration.
681#[proc_macro_attribute]
682pub fn upgrade(_attr: TokenStream, item: TokenStream) -> TokenStream {
683    item
684}
ree_exchange_sdk_macro/lib.rs

ree_exchange_sdk_macro/
lib.rs
