gmsol/store/
utils.rs

1use std::{
2    collections::HashMap,
3    iter::{Peekable, Zip},
4    slice::Iter,
5};
6
7use anchor_client::solana_sdk::{instruction::AccountMeta, pubkey::Pubkey};
8use gmsol_store::states::{common::TokensWithFeed, PriceProviderKind};
9
10use crate::{pyth::find_pyth_feed_account, utils::builder::FeedAddressMap};
11
12type Parser = Box<dyn Fn(Pubkey) -> crate::Result<AccountMeta>>;
13
14/// Feeds parser.
15pub struct FeedsParser {
16    parsers: HashMap<PriceProviderKind, Parser>,
17}
18
19impl Default for FeedsParser {
20    fn default() -> Self {
21        Self {
22            parsers: HashMap::from([(
23                PriceProviderKind::Pyth,
24                Box::new(|feed: Pubkey| {
25                    let pubkey = find_pyth_feed_account(0, feed.to_bytes()).0;
26                    Ok(AccountMeta {
27                        pubkey,
28                        is_signer: false,
29                        is_writable: false,
30                    })
31                }) as Parser,
32            )]),
33        }
34    }
35}
36
37impl FeedsParser {
38    /// Parse a [`TokensWithFeed`]
39    pub fn parse<'a>(
40        &'a self,
41        tokens_with_feed: &'a TokensWithFeed,
42    ) -> impl Iterator<Item = crate::Result<AccountMeta>> + 'a {
43        Feeds::new(tokens_with_feed).map(|res| {
44            res.and_then(|FeedConfig { provider, feed, .. }| self.dispatch(&provider, &feed))
45        })
46    }
47
48    /// Parse and sort by tokens.
49    pub fn parse_and_sort_by_tokens(
50        &self,
51        tokens_with_feed: &TokensWithFeed,
52    ) -> crate::Result<Vec<AccountMeta>> {
53        let accounts = self
54            .parse(tokens_with_feed)
55            .collect::<crate::Result<Vec<_>>>()?;
56
57        let mut combined = tokens_with_feed
58            .tokens
59            .iter()
60            .zip(accounts)
61            .collect::<Vec<_>>();
62
63        combined.sort_by_key(|(key, _)| *key);
64
65        Ok(combined.into_iter().map(|(_, account)| account).collect())
66    }
67
68    fn dispatch(&self, provider: &PriceProviderKind, feed: &Pubkey) -> crate::Result<AccountMeta> {
69        let Some(parser) = self.parsers.get(provider) else {
70            return Ok(AccountMeta {
71                pubkey: *feed,
72                is_signer: false,
73                is_writable: false,
74            });
75        };
76        (parser)(*feed)
77    }
78
79    /// Insert a pull oracle feed parser.
80    pub fn insert_pull_oracle_feed_parser(
81        &mut self,
82        provider: PriceProviderKind,
83        map: FeedAddressMap,
84    ) -> &mut Self {
85        self.parsers.insert(
86            provider,
87            Box::new(move |feed_id| {
88                let price_update = map.get(&feed_id).ok_or_else(|| {
89                    crate::Error::invalid_argument(format!(
90                        "feed account for {feed_id} not provided"
91                    ))
92                })?;
93
94                Ok(AccountMeta {
95                    pubkey: *price_update,
96                    is_signer: false,
97                    is_writable: false,
98                })
99            }),
100        );
101        self
102    }
103}
104
105#[cfg(feature = "pyth-pull-oracle")]
106mod pyth_pull_oracle {
107    use pyth_sdk::Identifier;
108
109    use super::*;
110    use crate::pyth::pull_oracle::Prices;
111
112    impl FeedsParser {
113        /// Parse Pyth feeds with price updates map.
114        pub fn with_pyth_price_updates(&mut self, price_updates: Prices) -> &mut Self {
115            self.parsers.insert(
116                PriceProviderKind::Pyth,
117                Box::new(move |feed| {
118                    let feed_id = Identifier::new(feed.to_bytes());
119                    let price_update = price_updates.get(&feed_id).ok_or_else(|| {
120                        crate::Error::invalid_argument(format!(
121                            "price update account for {feed_id}"
122                        ))
123                    })?;
124
125                    Ok(AccountMeta {
126                        pubkey: *price_update,
127                        is_signer: false,
128                        is_writable: false,
129                    })
130                }),
131            );
132            self
133        }
134    }
135}
136
137/// Feed account metas.
138pub struct Feeds<'a> {
139    provider_with_lengths: Peekable<Zip<Iter<'a, u8>, Iter<'a, u16>>>,
140    tokens: Iter<'a, Pubkey>,
141    feeds: Iter<'a, Pubkey>,
142    current: usize,
143    failed: bool,
144}
145
146impl<'a> Feeds<'a> {
147    /// Create from [`TokensWithFeed`].
148    pub fn new(token_with_feeds: &'a TokensWithFeed) -> Self {
149        let providers = token_with_feeds.providers.iter();
150        let nums = token_with_feeds.nums.iter();
151        let provider_with_lengths = providers.zip(nums).peekable();
152        let tokens = token_with_feeds.tokens.iter();
153        let feeds = token_with_feeds.feeds.iter();
154        Self {
155            provider_with_lengths,
156            tokens,
157            feeds,
158            current: 0,
159            failed: false,
160        }
161    }
162}
163
164/// A feed config.
165#[derive(Debug, Clone)]
166pub struct FeedConfig {
167    /// Token.
168    pub token: Pubkey,
169    /// Provider Kind.
170    pub provider: PriceProviderKind,
171    /// Feed.
172    pub feed: Pubkey,
173}
174
175impl Iterator for Feeds<'_> {
176    type Item = crate::Result<FeedConfig>;
177
178    fn next(&mut self) -> Option<Self::Item> {
179        if self.failed {
180            return None;
181        }
182        loop {
183            let (provider, length) = self.provider_with_lengths.peek()?;
184            if self.current == (**length as usize) {
185                self.provider_with_lengths.next();
186                self.current = 0;
187                continue;
188            }
189            let Ok(provider) = PriceProviderKind::try_from(**provider) else {
190                self.failed = true;
191                return Some(Err(crate::Error::invalid_argument(
192                    "invalid provider index",
193                )));
194            };
195            let Some(feed) = self.feeds.next() else {
196                return Some(Err(crate::Error::invalid_argument("not enough feeds")));
197            };
198            let Some(token) = self.tokens.next() else {
199                return Some(Err(crate::Error::invalid_argument("not enough tokens")));
200            };
201            self.current += 1;
202            return Some(Ok(FeedConfig {
203                token: *token,
204                provider,
205                feed: *feed,
206            }));
207        }
208    }
209}