aztec_pxe/stores/
contract_store.rs

1//! Contract artifact and instance storage.
2
3use std::sync::Arc;
4
5use aztec_core::abi::ContractArtifact;
6use aztec_core::error::Error;
7use aztec_core::types::{AztecAddress, ContractInstanceWithAddress, Fr};
8
9use super::kv::KvStore;
10
11/// Stores contract artifacts, instances, and class registrations.
12pub struct ContractStore {
13    kv: Arc<dyn KvStore>,
14}
15
16impl ContractStore {
17    pub fn new(kv: Arc<dyn KvStore>) -> Self {
18        Self { kv }
19    }
20
21    // --- Contract Instances ---
22
23    /// Store a contract instance by its address.
24    pub async fn add_instance(&self, instance: &ContractInstanceWithAddress) -> Result<(), Error> {
25        let key = instance_key(&instance.address);
26        let value = serde_json::to_vec(instance)?;
27        self.kv.put(&key, &value).await
28    }
29
30    /// Get a contract instance by address.
31    pub async fn get_instance(
32        &self,
33        address: &AztecAddress,
34    ) -> Result<Option<ContractInstanceWithAddress>, Error> {
35        let key = instance_key(address);
36        match self.kv.get(&key).await? {
37            Some(bytes) => Ok(Some(serde_json::from_slice(&bytes)?)),
38            None => Ok(None),
39        }
40    }
41
42    /// List all registered contract addresses.
43    pub async fn get_contract_addresses(&self) -> Result<Vec<AztecAddress>, Error> {
44        let entries = self.kv.list_prefix(b"contract:instance:").await?;
45        entries
46            .into_iter()
47            .map(|(_, v)| {
48                let inst: ContractInstanceWithAddress = serde_json::from_slice(&v)?;
49                Ok(inst.address)
50            })
51            .collect()
52    }
53
54    // --- Contract Artifacts ---
55
56    /// Store a contract artifact by class ID.
57    pub async fn add_artifact(
58        &self,
59        class_id: &Fr,
60        artifact: &ContractArtifact,
61    ) -> Result<(), Error> {
62        let key = artifact_key(class_id);
63        let value = serde_json::to_vec(artifact)?;
64        self.kv.put(&key, &value).await
65    }
66
67    /// Get a contract artifact by class ID.
68    pub async fn get_artifact(&self, class_id: &Fr) -> Result<Option<ContractArtifact>, Error> {
69        let key = artifact_key(class_id);
70        match self.kv.get(&key).await? {
71            Some(bytes) => Ok(Some(serde_json::from_slice(&bytes)?)),
72            None => Ok(None),
73        }
74    }
75
76    /// Update a contract's artifact (by address — looks up the class ID).
77    pub async fn update_artifact(
78        &self,
79        address: &AztecAddress,
80        artifact: &ContractArtifact,
81    ) -> Result<(), Error> {
82        let instance = self
83            .get_instance(address)
84            .await?
85            .ok_or_else(|| Error::InvalidData(format!("contract not found at {address}")))?;
86        self.add_artifact(&instance.inner.current_contract_class_id, artifact)
87            .await
88    }
89
90    // --- Contract Classes ---
91
92    /// Register a contract class (stores the artifact keyed by computed class ID).
93    pub async fn add_class(&self, artifact: &ContractArtifact) -> Result<Fr, Error> {
94        let class_id = aztec_core::hash::compute_contract_class_id_from_artifact(artifact)?;
95        self.add_artifact(&class_id, artifact).await?;
96        Ok(class_id)
97    }
98}
99
100fn instance_key(address: &AztecAddress) -> Vec<u8> {
101    format!("contract:instance:{address}").into_bytes()
102}
103
104fn artifact_key(class_id: &Fr) -> Vec<u8> {
105    format!("contract:artifact:{class_id}").into_bytes()
106}
107
108#[cfg(test)]
109mod tests {
110    use super::*;
111    use crate::stores::InMemoryKvStore;
112    use aztec_core::types::ContractInstance;
113
114    fn test_instance() -> ContractInstanceWithAddress {
115        ContractInstanceWithAddress {
116            address: AztecAddress::from(42u64),
117            inner: ContractInstance {
118                version: 1,
119                salt: Fr::from(1u64),
120                deployer: AztecAddress::zero(),
121                current_contract_class_id: Fr::from(100u64),
122                original_contract_class_id: Fr::from(100u64),
123                initialization_hash: Fr::zero(),
124                public_keys: Default::default(),
125            },
126        }
127    }
128
129    #[tokio::test]
130    async fn store_and_retrieve_instance() {
131        let kv = Arc::new(InMemoryKvStore::new());
132        let store = ContractStore::new(kv);
133        let inst = test_instance();
134
135        store.add_instance(&inst).await.unwrap();
136        let retrieved = store.get_instance(&inst.address).await.unwrap().unwrap();
137        assert_eq!(retrieved.address, inst.address);
138    }
139
140    #[tokio::test]
141    async fn list_contracts() {
142        let kv = Arc::new(InMemoryKvStore::new());
143        let store = ContractStore::new(kv);
144
145        assert!(store.get_contract_addresses().await.unwrap().is_empty());
146
147        store.add_instance(&test_instance()).await.unwrap();
148        let addrs = store.get_contract_addresses().await.unwrap();
149        assert_eq!(addrs.len(), 1);
150    }
151}
aztec_pxe/stores/contract_store.rs

aztec_pxe/stores/
contract_store.rs
