1#[doc(hidden)]
95pub mod schnorr;
96#[doc(hidden)]
97pub mod states;
98pub mod store;
99pub mod prelude {
100 pub use crate::*;
101 pub use ree_exchange_sdk_macro::*;
102}
103
104use crate::types::{
105 CoinBalance, Intention, IntentionSet, Pubkey, TxRecord, Txid, Utxo, exchange_interfaces::*,
106};
107use candid::CandidType;
108use ic_stable_structures::{
109 BTreeMap, DefaultMemoryImpl, Storable, memory_manager::VirtualMemory, storable::Bound,
110};
111use serde::{Deserialize, Serialize};
112
113pub use ree_types as types;
115
116#[doc(hidden)]
117pub type BlockStateStorage<S> = BTreeMap<u32, GlobalStateWrapper<S>, Memory>;
118#[doc(hidden)]
119pub type Memory = VirtualMemory<DefaultMemoryImpl>;
120#[doc(hidden)]
121pub type BlockStorage = BTreeMap<u32, Block, Memory>;
122#[doc(hidden)]
123pub type UnconfirmedTxStorage = BTreeMap<Txid, TxRecord, Memory>;
124#[doc(hidden)]
125pub type PoolStorage<S> = BTreeMap<String, Pool<S>, Memory>;
126
127#[derive(CandidType, Clone, Debug, Deserialize, Serialize, PartialEq, Eq, Copy)]
129pub enum Network {
130 Bitcoin,
131 Testnet4,
132}
133
134impl Into<crate::types::bitcoin::Network> for Network {
135 fn into(self) -> crate::types::bitcoin::Network {
136 match self {
137 Network::Bitcoin => crate::types::bitcoin::Network::Bitcoin,
138 Network::Testnet4 => crate::types::bitcoin::Network::Testnet4,
139 }
140 }
141}
142
143#[doc(hidden)]
144pub fn ensure_access<P: Pools>() -> Result<(), String> {
145 match P::network() {
146 Network::Bitcoin => crate::types::orchestrator_interfaces::ensure_orchestrator(),
147 Network::Testnet4 => crate::types::orchestrator_interfaces::ensure_testnet4_orchestrator(),
148 }
149}
150
151#[derive(CandidType, Clone, Debug, Deserialize, Serialize, PartialEq, Eq)]
153pub struct Block {
154 pub block_height: u32,
156 pub block_hash: String,
158 pub block_timestamp: u64,
160 pub txs: Vec<TxRecord>,
162}
163
164impl Storable for Block {
165 fn to_bytes(&self) -> std::borrow::Cow<'_, [u8]> {
166 let bytes = bincode::serialize(self).unwrap();
167 std::borrow::Cow::Owned(bytes)
168 }
169
170 fn into_bytes(self) -> Vec<u8> {
171 bincode::serialize(&self).unwrap()
172 }
173
174 fn from_bytes(bytes: std::borrow::Cow<'_, [u8]>) -> Self {
175 bincode::deserialize(bytes.as_ref()).unwrap()
176 }
177
178 const BOUND: Bound = Bound::Unbounded;
179}
180
181#[derive(CandidType, Clone, Debug, Deserialize, Serialize, PartialEq, Eq)]
184pub struct Metadata {
185 pub key: Pubkey,
186 pub key_derivation_path: Vec<Vec<u8>>,
187 pub name: String,
188 pub address: String,
189}
190
191impl Metadata {
192 pub async fn new<P: Pools>(name: String) -> Result<Self, String> {
194 let key_derivation_path: Vec<Vec<u8>> = vec![name.clone().into_bytes()];
195 let (key, _, address) =
196 crate::schnorr::request_p2tr_address(key_derivation_path.clone(), P::network())
197 .await
198 .map_err(|e| format!("Failed to generate pool address: {}", e))?;
199 Ok(Self {
200 key,
201 key_derivation_path,
202 name,
203 address: address.to_string(),
204 })
205 }
206}
207
208#[derive(CandidType, Clone, Debug, Deserialize, Serialize, PartialEq, Eq, Default)]
210pub struct StateInfo {
211 pub nonce: u64,
212 pub txid: Txid,
213 pub coin_reserved: Vec<CoinBalance>,
214 pub btc_reserved: u64,
215 pub utxos: Vec<Utxo>,
216 pub attributes: String,
217}
218
219#[derive(CandidType, Clone, Debug, Deserialize, Serialize, PartialEq, Eq)]
221pub struct ActionArgs {
222 pub txid: Txid,
223 pub initiator_address: String,
224 pub intention: Intention,
225 pub other_intentions: Vec<Intention>,
226 pub unconfirmed_tx_count: usize,
227}
228
229impl From<ExecuteTxArgs> for ActionArgs {
230 fn from(args: ExecuteTxArgs) -> Self {
231 let ExecuteTxArgs {
232 psbt_hex: _,
233 txid,
234 intention_set,
235 intention_index,
236 zero_confirmed_tx_queue_length,
237 } = args;
238 let IntentionSet {
239 mut intentions,
240 initiator_address,
241 tx_fee_in_sats: _,
242 } = intention_set;
243 let intention = intentions.swap_remove(intention_index as usize);
244 Self {
245 txid,
246 initiator_address,
247 intention,
248 other_intentions: intentions,
249 unconfirmed_tx_count: zero_confirmed_tx_queue_length as usize,
250 }
251 }
252}
253
254pub type ActionResult<S> = Result<S, String>;
256
257pub trait StateView {
259 fn inspect_state(&self) -> StateInfo;
260}
261
262#[derive(Clone, Debug, Deserialize, Serialize)]
265pub struct Pool<S> {
266 metadata: Metadata,
267 states: Vec<S>,
268}
269
270impl<S> Storable for Pool<S>
271where
272 S: Serialize + for<'de> Deserialize<'de>,
273{
274 const BOUND: Bound = Bound::Unbounded;
275
276 fn to_bytes(&self) -> std::borrow::Cow<'_, [u8]> {
277 let bytes = bincode::serialize(self).unwrap();
278 std::borrow::Cow::Owned(bytes)
279 }
280
281 fn into_bytes(self) -> Vec<u8> {
282 bincode::serialize(&self).unwrap()
283 }
284
285 fn from_bytes(bytes: std::borrow::Cow<'_, [u8]>) -> Self {
286 bincode::deserialize(bytes.as_ref()).unwrap()
287 }
288}
289
290impl<S> Pool<S>
291where
292 S: StateView,
293{
294 pub fn new(metadata: Metadata) -> Self {
296 Self {
297 metadata,
298 states: Vec::new(),
299 }
300 }
301
302 pub fn metadata(&self) -> &Metadata {
304 &self.metadata
305 }
306
307 pub fn last_state(&self) -> Option<&S> {
309 self.states.last()
310 }
311
312 pub fn states(&self) -> &Vec<S> {
314 &self.states
315 }
316
317 pub fn get(&self, txid: Txid) -> Option<&S> {
319 self.states
320 .iter()
321 .find(|state| state.inspect_state().txid == txid)
322 }
323
324 pub fn states_mut(&mut self) -> &mut Vec<S> {
326 &mut self.states
327 }
328}
329
330#[doc(hidden)]
331pub trait ReePool<S> {
332 fn get_pool_info(&self) -> PoolInfo;
333
334 fn get_pool_basic(&self) -> PoolBasic;
335
336 fn rollback(&mut self, txid: Txid) -> Result<Vec<S>, String>;
337
338 fn finalize(&mut self, txid: Txid) -> Result<(), String>;
339}
340
341#[doc(hidden)]
342impl<S> ReePool<S> for Pool<S>
343where
344 S: StateView,
345{
346 fn get_pool_basic(&self) -> PoolBasic {
347 PoolBasic {
348 name: self.metadata.name.clone(),
349 address: self.metadata.address.clone(),
350 }
351 }
352
353 fn get_pool_info(&self) -> PoolInfo {
354 let metadata: Metadata = self.metadata.clone();
355 let Metadata {
356 key,
357 key_derivation_path,
358 name,
359 address,
360 } = metadata;
361 let state = self
362 .states
363 .last()
364 .map(|s| s.inspect_state())
365 .unwrap_or_default();
366 let StateInfo {
367 txid: _,
368 nonce,
369 coin_reserved,
370 btc_reserved,
371 utxos,
372 attributes,
373 } = state;
374 PoolInfo {
375 key,
376 key_derivation_path,
377 name,
378 address,
379 nonce,
380 coin_reserved,
381 btc_reserved,
382 utxos,
383 attributes,
384 }
385 }
386
387 fn rollback(&mut self, txid: Txid) -> Result<Vec<S>, String> {
388 let idx = self
389 .states
390 .iter()
391 .position(|state| state.inspect_state().txid == txid)
392 .ok_or("txid not found".to_string())?;
393
394 let mut rollbacked_states = vec![];
395 while self.states.len() > idx {
396 rollbacked_states.push(self.states.pop().unwrap());
397 }
398
399 Ok(rollbacked_states)
400 }
401
402 fn finalize(&mut self, txid: Txid) -> Result<(), String> {
403 let idx = self
404 .states
405 .iter()
406 .position(|state| state.inspect_state().txid == txid)
407 .ok_or("txid not found".to_string())?;
408 if idx == 0 {
409 return Ok(());
410 }
411 self.states.rotate_left(idx);
412 self.states.truncate(self.states.len() - idx);
413 Ok(())
414 }
415}
416
417pub trait Pools {
419 type PoolState: StateView + Serialize + for<'de> Deserialize<'de>;
421
422 type BlockState: Serialize + for<'de> Deserialize<'de>;
424
425 const BLOCK_STATE_MEMORY: u8;
427
428 const POOL_STATE_MEMORY: u8;
430
431 fn network() -> Network;
433
434 fn finalize_threshold() -> u32 {
437 60
438 }
439}
440
441pub trait Hook: Pools {
444 fn on_tx_rollbacked(
446 _address: String,
447 _txid: Txid,
448 _reason: String,
449 _rollbacked_states: Vec<Self::PoolState>,
450 ) {
451 }
452
453 fn on_tx_confirmed(_address: String, _txid: Txid, _block: Block) {}
455
456 fn on_block_confirmed(_block: Block) {}
458
459 fn pre_block_confirmed(_height: u32) {}
461}
462
463pub trait PoolStorageAccess<P: Pools> {
466 fn block_state() -> Option<P::BlockState>;
467
468 fn commit(height: u32, block_state: P::BlockState) -> Result<(), String>;
469
470 fn get(address: &String) -> Option<Pool<P::PoolState>>;
471
472 fn insert(pool: Pool<P::PoolState>);
473
474 fn remove(address: &String) -> Option<Pool<P::PoolState>>;
475
476 fn iter() -> iter::PoolIterator<P>;
477}
478
479#[doc(hidden)]
480#[derive(Clone, Debug, Deserialize, Serialize)]
481pub struct GlobalStateWrapper<S> {
482 pub inner: S,
483}
484
485#[doc(hidden)]
486impl<S> GlobalStateWrapper<S> {
487 pub fn new(s: S) -> Self {
488 Self { inner: s }
489 }
490}
491
492#[doc(hidden)]
493impl<S> Storable for GlobalStateWrapper<S>
494where
495 S: Serialize + for<'de> Deserialize<'de>,
496{
497 const BOUND: Bound = Bound::Unbounded;
498
499 fn to_bytes(&self) -> std::borrow::Cow<'_, [u8]> {
500 let bytes = bincode::serialize(self).unwrap();
501 std::borrow::Cow::Owned(bytes)
502 }
503
504 fn into_bytes(self) -> Vec<u8> {
505 bincode::serialize(&self).unwrap()
506 }
507
508 fn from_bytes(bytes: std::borrow::Cow<'_, [u8]>) -> Self {
509 bincode::deserialize(bytes.as_ref()).unwrap()
510 }
511}
512
513pub trait Upgrade<P: Pools> {
600 type PoolState: Into<P::PoolState> + for<'de> Deserialize<'de> + Clone;
602
603 type BlockState: Into<P::BlockState> + for<'de> Deserialize<'de> + Clone;
605
606 const POOL_STATE_MEMORY: u8;
608
609 const BLOCK_STATE_MEMORY: u8;
611}
612
613#[doc(hidden)]
614pub fn iterator<P>(memory: Memory) -> iter::PoolIterator<P>
615where
616 P: Pools,
617{
618 let inner = PoolStorage::<P::PoolState>::init(memory);
619 let keys = inner.keys().collect::<Vec<_>>();
620 iter::PoolIterator {
621 inner,
622 cursor: 0,
623 keys,
624 }
625}
626
627#[doc(hidden)]
628pub mod iter {
629 pub struct PoolIterator<P: super::Pools> {
630 pub(crate) inner: super::PoolStorage<P::PoolState>,
631 pub(crate) cursor: usize,
632 pub(crate) keys: Vec<String>,
633 }
634
635 impl<P> std::iter::Iterator for PoolIterator<P>
636 where
637 P: super::Pools,
638 {
639 type Item = (String, super::Pool<P::PoolState>);
640
641 fn next(&mut self) -> Option<Self::Item> {
642 if self.cursor < self.keys.len() {
643 let key = self.keys[self.cursor].clone();
644 self.cursor += 1;
645 self.inner.get(&key).map(|v| (key.clone(), v))
646 } else {
647 None
648 }
649 }
650 }
651}
652
653#[cfg(test)]
654pub mod test {
655 use std::str::FromStr;
656
657 use super::*;
658
659 #[derive(CandidType, Clone, Debug, Deserialize, Serialize, PartialEq, Eq)]
660 struct DummyPoolState {
661 nonce: u64,
662 txid: Txid,
663 coin_reserved: Vec<CoinBalance>,
664 btc_reserved: u64,
665 utxos: Vec<Utxo>,
666 attributes: String,
667 }
668
669 impl StateView for DummyPoolState {
670 fn inspect_state(&self) -> StateInfo {
671 StateInfo {
672 txid: self.txid.clone(),
673 nonce: self.nonce,
674 coin_reserved: self.coin_reserved.clone(),
675 btc_reserved: self.btc_reserved,
676 utxos: self.utxos.clone(),
677 attributes: self.attributes.clone(),
678 }
679 }
680 }
681
682 #[test]
683 pub fn test_candid_and_bincode_serialize() {
684 let state = DummyPoolState {
685 nonce: 1,
686 txid: Txid::default(),
687 coin_reserved: vec![],
688 btc_reserved: 0,
689 utxos: vec![],
690 attributes: "{}".to_string(),
691 };
692 let pool = Pool::<DummyPoolState> {
693 metadata: Metadata {
694 key: Pubkey::from_raw(vec![2u8; 33]).unwrap(),
695 key_derivation_path: vec![vec![0; 32]],
696 name: "Test Pool".to_string(),
697 address: "test-address".to_string(),
698 },
699 states: vec![state.clone()],
700 };
701 let bincode_serialized = pool.to_bytes();
702 Pool::<DummyPoolState>::from_bytes(bincode_serialized);
703 assert_eq!(pool.metadata.name, "Test Pool");
704
705 let mut candid_ser = candid::ser::IDLBuilder::new();
706 candid_ser.arg(&state).unwrap();
707 let candid_serialized = candid_ser.serialize_to_vec();
708 assert!(candid_serialized.is_ok());
709 let candid_serialized = candid_serialized.unwrap();
710 let mut candid_de = candid::de::IDLDeserialize::new(&candid_serialized).unwrap();
711 let candid_deserialized = candid_de.get_value::<DummyPoolState>();
712 assert!(candid_deserialized.is_ok());
713 }
714
715 #[test]
716 fn test_pool_rollback() {
717 let mut pool = Pool::<DummyPoolState> {
718 metadata: Metadata {
719 key: Pubkey::from_raw(vec![2u8; 33]).unwrap(),
720 key_derivation_path: vec![vec![0; 32]],
721 name: "Test Pool".to_string(),
722 address: "test-address".to_string(),
723 },
724 states: vec![],
725 };
726 let push_random_state_by_txid = |txid: &str, pool: &mut Pool<DummyPoolState>| {
727 let txid = Txid::from_str(txid).unwrap();
728 let nonce = pool.states.len() as u64;
729 let state = DummyPoolState {
730 nonce,
731 txid,
732 coin_reserved: vec![],
733 btc_reserved: 0,
734 utxos: vec![],
735 attributes: "{}".to_string(),
736 };
737 pool.states.push(state);
738 };
739
740 let txs = [
741 "51230fe70deae44a92f8f44a600585e3e57b8c8720a0b67c4c422f579d9ace2a",
742 "51230fe70deae44a92f8f44a600585e3e57b8c8720a0b67c4c422f579d9ace2b",
743 "51230fe70deae44a92f8f44a600585e3e57b8c8720a0b67c4c422f579d9ace2c",
744 ];
745
746 let init_pool_state = |pool: &mut Pool<DummyPoolState>| {
747 pool.states.clear();
748 for txid in txs.iter() {
749 push_random_state_by_txid(txid, pool);
750 }
751 };
752
753 init_pool_state(&mut pool);
755 assert_eq!(pool.states.len(), 3);
756 let before_rollback_states = pool.states.clone();
757 let rollbacked_states = pool.rollback(Txid::from_str(txs[0]).unwrap()).unwrap();
758 assert_eq!(rollbacked_states.len(), 3);
759 assert_eq!(pool.states.len(), 0);
760 assert_eq!(rollbacked_states[0], before_rollback_states[2]);
761 assert_eq!(rollbacked_states[1], before_rollback_states[1]);
762 assert_eq!(rollbacked_states[2], before_rollback_states[0]);
763
764 init_pool_state(&mut pool);
766 assert_eq!(pool.states.len(), 3);
767 let before_rollback_states = pool.states.clone();
768 let rollbacked_states = pool.rollback(Txid::from_str(txs[1]).unwrap()).unwrap();
769 assert_eq!(rollbacked_states.len(), 2);
770 assert_eq!(pool.states.len(), 1);
771 assert_eq!(pool.states[0], before_rollback_states[0]);
772 assert_eq!(rollbacked_states[0], before_rollback_states[2]);
773 assert_eq!(rollbacked_states[1], before_rollback_states[1]);
774
775 init_pool_state(&mut pool);
777 assert_eq!(pool.states.len(), 3);
778 let before_rollback_states = pool.states.clone();
779 let rollbacked_states = pool.rollback(Txid::from_str(txs[2]).unwrap()).unwrap();
780 assert_eq!(rollbacked_states.len(), 1);
781 assert_eq!(pool.states.len(), 2);
782 assert_eq!(pool.states[0], before_rollback_states[0]);
783 assert_eq!(pool.states[1], before_rollback_states[1]);
784 assert_eq!(rollbacked_states[0], before_rollback_states[2]);
785 }
786}