cherry_rpc_call/
lib.rs

1use alloy_multicall::Multicall;
2use alloy_primitives::{Address, U256};
3use alloy_provider::ProviderBuilder;
4use alloy_sol_types::{sol, JsonAbiExt};
5use anyhow::{Context, Result};
6use arrow::{
7    array::{Array, FixedSizeBinaryBuilder, StringArray, UInt8Array},
8    datatypes::{DataType, Field, Schema},
9    record_batch::RecordBatch,
10};
11use std::{str::FromStr, sync::Arc};
12
13sol! {
14    #[derive(Debug)]
15    #[sol(abi)]
16    function decimals()
17        public
18        view
19        virtual
20        override
21        returns (uint8);
22
23    #[derive(Debug)]
24    #[sol(abi)]
25    function symbol()
26        public
27        view
28        virtual
29        override
30        returns (string memory);
31
32    #[derive(Debug)]
33    #[sol(abi)]
34    function name()
35        public
36        view
37        virtual
38        override
39        returns (string memory);
40
41    #[derive(Debug)]
42    #[sol(abi)]
43    function totalSupply()
44        public
45        view
46        virtual
47        override
48        returns (uint256);
49}
50
51#[derive(Debug)]
52pub struct TokenMetadata {
53    pub address: Option<Address>,
54    pub decimals: Option<u8>,
55    pub symbol: Option<String>,
56    pub name: Option<String>,
57    pub total_supply: Option<U256>,
58}
59
60#[derive(Debug)]
61pub struct TokenMetadataSelector {
62    pub decimals: bool,
63    pub symbol: bool,
64    pub name: bool,
65    pub total_supply: bool,
66}
67
68impl Default for TokenMetadataSelector {
69    fn default() -> Self {
70        Self {
71            decimals: true,
72            symbol: true,
73            name: true,
74            total_supply: false,
75        }
76    }
77}
78
79#[cfg(feature = "pyo3")]
80impl<'py> pyo3::FromPyObject<'py> for TokenMetadataSelector {
81    fn extract_bound(ob: &pyo3::Bound<'py, pyo3::PyAny>) -> pyo3::PyResult<Self> {
82        use pyo3::types::PyAnyMethods;
83        use pyo3::types::PyDict;
84        // Get the dictionary
85        let dict = ob.downcast::<PyDict>()?;
86
87        let decimals = dict.get_item("decimals").unwrap();
88        let symbol = dict.get_item("symbol").unwrap();
89        let name = dict.get_item("name").unwrap();
90        let total_supply = dict.get_item("total_supply").unwrap();
91
92        Ok(TokenMetadataSelector {
93            decimals: decimals.extract::<bool>()?,
94            symbol: symbol.extract::<bool>()?,
95            name: name.extract::<bool>()?,
96            total_supply: total_supply.extract::<bool>()?,
97        })
98    }
99}
100
101pub async fn get_token_metadata(
102    rpc_url: &str,
103    addresses: Vec<String>,
104    selector: &TokenMetadataSelector,
105) -> Result<Vec<TokenMetadata>> {
106    let provider = ProviderBuilder::new().on_http(rpc_url.parse().context("invalid rpc url")?);
107    let mut multicall = Multicall::with_provider_chain_id(&provider)
108        .await
109        .context("failed to create multicall")?;
110
111    let decimals = decimalsCall::abi();
112    let symbol = symbolCall::abi();
113    let name = nameCall::abi();
114    let total_supply = totalSupplyCall::abi();
115
116    let addresses: Vec<Option<Address>> = addresses
117        .into_iter()
118        .map(|addr| Address::from_str(&addr).ok())
119        .collect();
120    for address in addresses.iter().flatten() {
121        if selector.decimals {
122            multicall.add_call(*address, &decimals, &[], true);
123        }
124        if selector.symbol {
125            multicall.add_call(*address, &symbol, &[], true);
126        }
127        if selector.name {
128            multicall.add_call(*address, &name, &[], true);
129        }
130        if selector.total_supply {
131            multicall.add_call(*address, &total_supply, &[], true);
132        }
133    }
134
135    let results = multicall.call().await.context("failed to call multicall")?;
136    let mut token_metadata: Vec<TokenMetadata> = Vec::new();
137
138    // Process results in chunks (decimals, symbol, name, total_supply)
139    let mut i = 0;
140    let chuck_size = selector.decimals as usize
141        + selector.symbol as usize
142        + selector.name as usize
143        + selector.total_supply as usize;
144    for address in addresses.iter() {
145        if let Some(address) = address {
146            let mut base_idx = i * chuck_size;
147            let decimals: Option<u8> = if selector.decimals {
148                results
149                    .get(base_idx)
150                    .and_then(|result| result.as_ref().ok())
151                    .and_then(|v| v.as_uint())
152                    .map(|uint| uint.0.as_limbs()[0] as u8)
153            } else {
154                None
155            };
156            let symbol: Option<String> = if selector.symbol {
157                base_idx += 1;
158                results
159                    .get(base_idx)
160                    .and_then(|result| result.as_ref().ok())
161                    .and_then(|v| v.as_str())
162                    .map(|s| s.to_string())
163            } else {
164                None
165            };
166            let name: Option<String> = if selector.name {
167                base_idx += 1;
168                results
169                    .get(base_idx)
170                    .and_then(|result| result.as_ref().ok())
171                    .and_then(|v| v.as_str())
172                    .map(|s| s.to_string())
173            } else {
174                None
175            };
176            let total_supply: Option<U256> = if selector.total_supply {
177                base_idx += 1;
178                results
179                    .get(base_idx)
180                    .and_then(|result| result.as_ref().ok())
181                    .and_then(|v| v.as_uint())
182                    .map(|uint| uint.0)
183            } else {
184                None
185            };
186
187            token_metadata.push(TokenMetadata {
188                address: Some(*address),
189                decimals,
190                symbol,
191                name,
192                total_supply,
193            });
194            i += 1;
195        } else {
196            token_metadata.push(TokenMetadata {
197                address: None,
198                decimals: None,
199                symbol: None,
200                name: None,
201                total_supply: None,
202            });
203        }
204    }
205
206    Ok(token_metadata)
207}
208
209pub fn token_metadata_to_table(
210    token_metadata: Vec<TokenMetadata>,
211    selector: &TokenMetadataSelector,
212) -> Result<RecordBatch> {
213    let mut fields = Vec::new();
214    fields.push(Field::new("address", DataType::FixedSizeBinary(20), true));
215    if selector.decimals {
216        fields.push(Field::new("decimals", DataType::UInt8, true));
217    }
218    if selector.symbol {
219        fields.push(Field::new("symbol", DataType::Utf8, true));
220    }
221    if selector.name {
222        fields.push(Field::new("name", DataType::Utf8, true));
223    }
224    if selector.total_supply {
225        fields.push(Field::new(
226            "total_supply",
227            DataType::FixedSizeBinary(32),
228            true,
229        ));
230    }
231
232    let schema = Schema::new(fields);
233
234    let array_len = token_metadata.len();
235    let mut address_builder = FixedSizeBinaryBuilder::with_capacity(array_len, 20);
236    let mut decimals_values: Vec<Option<u8>> = Vec::with_capacity(array_len);
237    let mut symbol_values: Vec<Option<String>> = Vec::with_capacity(array_len);
238    let mut name_values: Vec<Option<String>> = Vec::with_capacity(array_len);
239    let mut total_supply_builder = FixedSizeBinaryBuilder::with_capacity(array_len, 32);
240
241    for token in token_metadata {
242        let address_bytes: Option<[u8; 20]> = token
243            .address
244            .and_then(|addr| addr.as_slice().try_into().ok());
245        match address_bytes {
246            Some(address_bytes) => {
247                let _ = address_builder.append_value(address_bytes);
248            }
249            None => address_builder.append_null(),
250        }
251
252        if selector.decimals {
253            decimals_values.push(token.decimals);
254        }
255        if selector.symbol {
256            symbol_values.push(token.symbol);
257        }
258        if selector.name {
259            name_values.push(token.name);
260        }
261        if selector.total_supply {
262            match token.total_supply {
263                Some(supply) => {
264                    // Add explicit type annotation to fix the error
265                    let bytes: [u8; 32] = supply.to_be_bytes();
266                    let _ = total_supply_builder.append_value(bytes);
267                }
268                None => total_supply_builder.append_null(),
269            }
270        }
271    }
272
273    let mut arrays = Vec::new();
274    arrays.push(Arc::new(address_builder.finish()) as Arc<dyn Array>);
275    if selector.decimals {
276        arrays.push(Arc::new(UInt8Array::from(decimals_values)) as Arc<dyn Array>);
277    }
278    if selector.symbol {
279        arrays.push(Arc::new(StringArray::from(symbol_values)) as Arc<dyn Array>);
280    }
281    if selector.name {
282        arrays.push(Arc::new(StringArray::from(name_values)) as Arc<dyn Array>);
283    }
284    if selector.total_supply {
285        arrays.push(Arc::new(total_supply_builder.finish()) as Arc<dyn Array>);
286    }
287    let batch = RecordBatch::try_new(Arc::new(schema), arrays)?;
288
289    Ok(batch)
290}
291
292#[tokio::test]
293#[ignore]
294async fn test_get_token_metadata() {
295    let selector = TokenMetadataSelector {
296        decimals: true,
297        symbol: false,
298        name: true,
299        total_supply: false,
300    };
301    let token_metadata = get_token_metadata(
302        "https://ethereum-rpc.publicnode.com",
303        vec![
304            "Invalid address".to_string(),
305            "0x0000000000000000000000000000000000000000".to_string(),
306            "0x6B175474E89094C44Da98b954EedeAC495271d0F".to_string(),
307            "0xA0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48".to_string(),
308            "0xae7ab96520DE3A18E5e111B5EaAb095312D7fE84".to_string(),
309        ],
310        &selector,
311    )
312    .await;
313
314    let table = token_metadata_to_table(token_metadata.unwrap(), &selector).unwrap();
315
316    println!("{:?}", table);
317}
cherry_rpc_call/lib.rs

cherry_rpc_call/
lib.rs
