hdp_primitives/task/datalake/block_sampled/
collection.rs

1use std::{fmt::Display, str::FromStr};
2
3use alloy::primitives::{Address, StorageKey};
4use anyhow::{bail, Result};
5use serde::{Deserialize, Serialize};
6
7use crate::task::datalake::{DatalakeCollection, DatalakeField};
8
9use super::rlp_fields::{AccountField, HeaderField};
10
11#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
12#[serde(try_from = "String")]
13pub enum BlockSampledCollection {
14    Header(HeaderField),
15    Account(Address, AccountField),
16    Storage(Address, StorageKey),
17}
18
19pub enum BlockSampledCollectionType {
20    Header,
21    Account,
22    Storage,
23}
24
25impl BlockSampledCollectionType {
26    pub fn variants() -> Vec<String> {
27        vec![
28            "HEADER".to_string(),
29            "ACCOUNT".to_string(),
30            "STORAGE".to_string(),
31        ]
32    }
33}
34
35impl FromStr for BlockSampledCollectionType {
36    type Err = anyhow::Error;
37
38    fn from_str(s: &str) -> Result<Self> {
39        match s.to_uppercase().as_str() {
40            "HEADER" => Ok(BlockSampledCollectionType::Header),
41            "ACCOUNT" => Ok(BlockSampledCollectionType::Account),
42            "STORAGE" => Ok(BlockSampledCollectionType::Storage),
43            _ => bail!("Unknown block sampled collection type"),
44        }
45    }
46}
47
48impl DatalakeCollection for BlockSampledCollection {
49    fn to_index(&self) -> u8 {
50        match self {
51            BlockSampledCollection::Header(_) => 1,
52            BlockSampledCollection::Account(..) => 2,
53            BlockSampledCollection::Storage(..) => 3,
54        }
55    }
56
57    fn serialize(&self) -> Result<Vec<u8>> {
58        let mut serialized = Vec::new();
59        match self {
60            BlockSampledCollection::Header(field) => {
61                serialized.push(1);
62                serialized.push(field.to_index());
63            }
64            BlockSampledCollection::Account(address, field) => {
65                serialized.push(2);
66                serialized.extend_from_slice(address.as_slice());
67                serialized.push(field.to_index());
68            }
69            BlockSampledCollection::Storage(address, slot) => {
70                serialized.push(3);
71                serialized.extend_from_slice(address.as_slice());
72                serialized.extend_from_slice(slot.as_ref());
73            }
74        }
75
76        Ok(serialized)
77    }
78
79    fn deserialize(serialized: &[u8]) -> Result<Self> {
80        if serialized.is_empty() {
81            bail!("Invalid block sampled collection");
82        }
83
84        match serialized[0] {
85            1 => {
86                if serialized.len() != 2 {
87                    bail!("Invalid header property");
88                }
89                Ok(BlockSampledCollection::Header(HeaderField::from_index(
90                    serialized[1],
91                )?))
92            }
93            2 => {
94                if serialized.len() != 22 {
95                    bail!("Invalid account property");
96                }
97                let address = Address::from_slice(&serialized[1..21]);
98                Ok(BlockSampledCollection::Account(
99                    address,
100                    AccountField::from_index(serialized[21])?,
101                ))
102            }
103            3 => {
104                if serialized.len() != 53 {
105                    bail!("Invalid storage property");
106                }
107                let address = Address::from_slice(&serialized[1..21]);
108                let slot = StorageKey::from_slice(&serialized[21..53]);
109                Ok(BlockSampledCollection::Storage(address, slot))
110            }
111            _ => bail!("Unknown block sampled collection"),
112        }
113    }
114}
115
116impl FromStr for BlockSampledCollection {
117    type Err = anyhow::Error;
118
119    fn from_str(s: &str) -> Result<Self> {
120        // Split into two parts by '.'
121        let parts: Vec<&str> = s.split('.').collect();
122        if !(parts.len() == 2 || parts.len() == 3) {
123            bail!("Invalid block sampled collection format");
124        }
125
126        match parts[0].to_uppercase().as_str() {
127            "HEADER" => Ok(BlockSampledCollection::Header(HeaderField::from_str(
128                parts[1].to_uppercase().as_str(),
129            )?)),
130            "ACCOUNT" => {
131                let address = Address::from_str(parts[1])?;
132                let field = AccountField::from_str(parts[2].to_uppercase().as_str())?;
133                Ok(BlockSampledCollection::Account(address, field))
134            }
135            "STORAGE" => {
136                let address = Address::from_str(parts[1])?;
137                let slot = StorageKey::from_str(parts[2])?;
138                Ok(BlockSampledCollection::Storage(address, slot))
139            }
140            _ => bail!("Unknown block sampled collection"),
141        }
142    }
143}
144
145impl TryFrom<String> for BlockSampledCollection {
146    type Error = anyhow::Error;
147
148    fn try_from(value: String) -> Result<Self> {
149        BlockSampledCollection::from_str(&value)
150    }
151}
152
153impl Display for BlockSampledCollection {
154    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
155        match self {
156            BlockSampledCollection::Header(field) => write!(f, "header.{}", field),
157            BlockSampledCollection::Account(address, field) => {
158                write!(f, "account.{}.{}", address, field)
159            }
160            BlockSampledCollection::Storage(address, slot) => {
161                write!(f, "storage.{}.{}", address, slot)
162            }
163        }
164    }
165}
hdp_primitives/task/datalake/block_sampled/collection.rs

hdp_primitives/task/datalake/block_sampled/
collection.rs
