### Description

Resolves issue #163

This adds `NetworkId` into the `BlockInfo` struct, and introduces this identifier into the account state, drep state, epochs state and rest_blockfrost as these all rely on querying the `StakeAddressMap`.

Author: lowhung

common/src/address.rs

@@ -4,12 +4,10 @@
 
 use crate::cip19::{VarIntDecoder, VarIntEncoder};
 use crate::types::{KeyHash, ScriptHash};
-use crate::Credential;
+use crate::{Credential, NetworkId};
 use anyhow::{anyhow, bail, Result};
 use serde_with::{hex::Hex, serde_as};
-use std::borrow::Borrow;
 use std::fmt::{Display, Formatter};
-use std::hash::{Hash, Hasher};
 
 /// a Byron-era address
 #[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
@@ -28,6 +26,15 @@ pub enum AddressNetwork {
     Test,
 }
 
+impl From<NetworkId> for AddressNetwork {
+    fn from(network: NetworkId) -> Self {
+        match network {
+            NetworkId::Mainnet => Self::Main,
+            NetworkId::Testnet => Self::Test,
+        }
+    }
+}
+
 impl Default for AddressNetwork {
     fn default() -> Self {
         Self::Main
@@ -201,7 +208,7 @@ impl StakeAddressPayload {
 }
 
 /// A stake address
-#[derive(Debug, Clone, Eq, serde::Serialize, serde::Deserialize)]
+#[derive(Debug, Clone, Eq, PartialEq, Hash, serde::Serialize, serde::Deserialize)]
 pub struct StakeAddress {
     /// Network id
     pub network: AddressNetwork,
@@ -211,7 +218,7 @@ pub struct StakeAddress {
 }
 
 impl StakeAddress {
-    pub fn new(network: AddressNetwork, payload: StakeAddressPayload) -> Self {
+    pub fn new(payload: StakeAddressPayload, network: AddressNetwork) -> Self {
         StakeAddress { network, payload }
     }
 
@@ -305,26 +312,6 @@ impl Display for StakeAddress {
     }
 }
 
-impl Hash for StakeAddress {
-    fn hash<H: Hasher>(&self, state: &mut H) {
-        self.network.hash(state);
-        std::mem::discriminant(&self.payload).hash(state);
-        self.get_hash().hash(state);
-    }
-}
-
-impl Borrow<[u8]> for StakeAddress {
-    fn borrow(&self) -> &[u8] {
-        self.get_hash()
-    }
-}
-
-impl PartialEq for StakeAddress {
-    fn eq(&self, other: &Self) -> bool {
-        self.network == other.network && self.payload == other.payload
-    }
-}
-
 impl<C> minicbor::Encode<C> for StakeAddress {
     fn encode<W: minicbor::encode::Write>(
         &self,

common/src/stake_addresses.rs

@@ -528,10 +528,10 @@ mod tests {
 
     fn create_stake_address(hash: &[u8]) -> StakeAddress {
         StakeAddress::new(
-            AddressNetwork::Main,
             StakeAddressPayload::StakeKeyHash(
                 hash.to_vec().try_into().expect("Invalid hash length"),
             ),
+            AddressNetwork::Main,
         )
     }
 

common/src/types.rs

@@ -26,6 +26,16 @@ pub enum NetworkId {
     Mainnet,
 }
 
+impl From<String> for NetworkId {
+    fn from(s: String) -> Self {
+        match s.as_str() {
+            "testnet" => NetworkId::Testnet,
+            "mainnet" => NetworkId::Mainnet,
+            _ => NetworkId::Mainnet,
+        }
+    }
+}
+
 /// Protocol era
 #[derive(
     Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, serde::Serialize, serde::Deserialize,
@@ -120,6 +130,9 @@ pub struct BlockInfo {
     #[serde(default)]
     pub timestamp: u64,
 
+    /// Network ID
+    pub network_id: NetworkId,
+
     /// Protocol era
     pub era: Era,
 }
@@ -614,7 +627,7 @@ impl Credential {
 pub type StakeCredential = Credential;
 
 impl StakeCredential {
-    pub fn to_stake_address(&self, network: Option<AddressNetwork>) -> StakeAddress {
+    pub fn to_stake_address(&self, network: AddressNetwork) -> StakeAddress {
         let payload = match self {
             StakeCredential::AddrKeyHash(hash) => StakeAddressPayload::StakeKeyHash(
                 hash.clone().try_into().expect("Invalid hash length"),
@@ -623,7 +636,7 @@ impl StakeCredential {
                 hash.clone().try_into().expect("Invalid hash length"),
             ),
         };
-        StakeAddress::new(network.unwrap_or(AddressNetwork::Main), payload)
+        StakeAddress::new(payload, network)
     }
 }
 

modules/accounts_state/src/snapshot.rs

@@ -165,12 +165,12 @@ impl Snapshot {
             return 0;
         };
 
-        let addr_set: std::collections::HashSet<_> = addresses.iter().collect();
+        let address_set: std::collections::HashSet<_> = addresses.iter().collect();
         snapshot_spo
             .delegators
             .iter()
-            .filter_map(|(addr, amount)| {
-                if addr_set.contains(&addr.get_credential()) {
+            .filter_map(|(address, amount)| {
+                if address_set.contains(&address.get_credential()) {
                     Some(*amount)
                 } else {
                     None

modules/accounts_state/src/state.rs

@@ -14,9 +14,9 @@ use acropolis_common::{
     protocol_params::ProtocolParams,
     stake_addresses::{StakeAddressMap, StakeAddressState},
     BlockInfo, DRepChoice, DRepCredential, DelegatedStake, InstantaneousRewardSource,
-    InstantaneousRewardTarget, KeyHash, Lovelace, MoveInstantaneousReward, PoolLiveStakeInfo,
-    PoolRegistration, Pot, SPORewards, StakeAddress, StakeCredential, StakeRewardDelta,
-    TxCertificate,
+    InstantaneousRewardTarget, KeyHash, Lovelace, MoveInstantaneousReward, NetworkId,
+    PoolLiveStakeInfo, PoolRegistration, Pot, SPORewards, StakeAddress, StakeCredential,
+    StakeRewardDelta, TxCertificate,
 };
 use anyhow::Result;
 use imbl::OrdMap;
@@ -95,6 +95,9 @@ pub struct State {
     /// List of SPOs (by operator ID) retiring in the current epoch
     retiring_spos: Vec<KeyHash>,
 
+    /// Network ID
+    network_id: NetworkId,
+
     /// Map of staking address values
     /// Wrapped in an Arc so it doesn't get cloned in full by StateHistory
     stake_addresses: Arc<Mutex<StakeAddressMap>>,
@@ -550,7 +553,8 @@ impl State {
                     // Transfer to (in theory also from) stake addresses from (to) a pot
                     let mut total_value: u64 = 0;
                     for (credential, value) in deltas.iter() {
-                        let stake_address = credential.to_stake_address(None); // Need to convert credential to address
+                        let stake_address =
+                            credential.to_stake_address(self.network_id.clone().into());
 
                         // Get old stake address state, or create one
                         let mut stake_addresses = self.stake_addresses.lock().unwrap();
@@ -792,8 +796,8 @@ impl State {
 
     /// Register a stake address, with a specified deposit if known
     fn register_stake_address(&mut self, credential: &StakeCredential, deposit: Option<Lovelace>) {
-        // TODO: Handle network
-        let stake_address = credential.to_stake_address(None);
+
+        let stake_address = credential.to_stake_address(self.network_id.clone().into());
 
         // Stake addresses can be registered after being used in UTXOs
         let mut stake_addresses = self.stake_addresses.lock().unwrap();
@@ -823,8 +827,8 @@ impl State {
 
     /// Deregister a stake address, with specified refund if known
     fn deregister_stake_address(&mut self, credential: &StakeCredential, refund: Option<Lovelace>) {
-        // TODO: Handle network
-        let stake_address = credential.to_stake_address(None);
+
+        let stake_address = credential.to_stake_address(self.network_id.clone().into());
 
         // Check if it existed
         let mut stake_addresses = self.stake_addresses.lock().unwrap();
@@ -860,8 +864,8 @@ impl State {
 
     /// Record a stake delegation
     fn record_stake_delegation(&mut self, credential: &StakeCredential, spo: &KeyHash) {
-        // TODO: Handle network
-        let stake_address = credential.to_stake_address(None);
+
+        let stake_address = credential.to_stake_address(self.network_id.clone().into());
         let mut stake_addresses = self.stake_addresses.lock().unwrap();
         stake_addresses.record_stake_delegation(&stake_address, spo);
     }
@@ -874,8 +878,8 @@ impl State {
 
     /// record a drep delegation
     fn record_drep_delegation(&mut self, credential: &StakeCredential, drep: &DRepChoice) {
-        // TODO: Handle network
-        let stake_address = credential.to_stake_address(None);
+
+        let stake_address = credential.to_stake_address(self.network_id.clone().into());
         let mut stake_addresses = self.stake_addresses.lock().unwrap();
         stake_addresses.record_drep_delegation(&stake_address, drep);
     }

modules/drep_state/src/state.rs

@@ -4,7 +4,7 @@ use acropolis_common::{messages::{Message, StateQuery, StateQueryResponse}, quer
     accounts::{AccountsStateQuery, AccountsStateQueryResponse, DEFAULT_ACCOUNTS_QUERY_TOPIC},
     get_query_topic,
     governance::{DRepActionUpdate, DRepUpdateEvent, VoteRecord},
-}, Anchor, Credential, DRepChoice, DRepCredential, KeyHash, Lovelace, StakeAddress, StakeCredential, TxCertificate, TxHash, Voter, VotingProcedures};
+}, Anchor, Credential, DRepChoice, DRepCredential, KeyHash, Lovelace, NetworkId, StakeAddress, StakeCredential, TxCertificate, TxHash, Voter, VotingProcedures};
 use anyhow::{anyhow, Result};
 use caryatid_sdk::Context;
 use serde_with::serde_as;
@@ -89,6 +89,7 @@ impl DRepStorageConfig {
 pub struct State {
     pub config: DRepStorageConfig,
     pub dreps: HashMap<DRepCredential, DRepRecord>,
+    pub network_id: NetworkId,
     pub historical_dreps: Option<HashMap<DRepCredential, HistoricalDRepState>>,
 }
 
@@ -97,6 +98,7 @@ impl State {
         Self {
             config,
             dreps: HashMap::new(),
+            network_id: NetworkId::default(),
             historical_dreps: if config.any_enabled() {
                 Some(HashMap::new())
             } else {
@@ -470,8 +472,7 @@ impl State {
         delegators: Vec<(&StakeCredential, &DRepChoice)>,
     ) -> Result<()> {
         let stake_keys: Vec<KeyHash> = delegators.iter().map(|(sc, _)| sc.get_hash()).collect();
-        // TODO: USE NETWORK ID
-        let stake_addresses: Vec<StakeAddress> = delegators.iter().map(|(k, _) | k.to_stake_address(None) ).collect();
+        let stake_addresses: Vec<StakeAddress> = delegators.iter().map(|(k, _) | k.to_stake_address(self.network_id.clone().into()) ).collect();
         let stake_key_to_input: HashMap<KeyHash, _> = delegators
             .iter()
             .zip(&stake_keys)

modules/epochs_state/src/epochs_history.rs

@@ -78,7 +78,7 @@ impl EpochsHistoryState {
 #[cfg(test)]
 mod tests {
     use super::*;
-    use acropolis_common::{BlockHash, BlockStatus, Era};
+    use acropolis_common::{BlockHash, BlockStatus, Era, NetworkId};
 
     fn make_block(epoch: u64) -> BlockInfo {
         BlockInfo {
@@ -89,6 +89,7 @@ mod tests {
             epoch,
             epoch_slot: 99,
             new_epoch: false,
+            network_id: NetworkId::default(),
             timestamp: 99999,
             era: Era::Conway,
         }

modules/epochs_state/src/state.rs

@@ -269,7 +269,7 @@ mod tests {
         crypto::keyhash_224,
         protocol_params::{Nonce, NonceHash},
         state_history::{StateHistory, StateHistoryStore},
-        BlockHash, BlockInfo, BlockStatus, Era,
+        BlockHash, BlockInfo, BlockStatus, Era, NetworkId,
     };
     use tokio::sync::Mutex;
 
@@ -281,6 +281,7 @@ mod tests {
             hash: BlockHash::default(),
             epoch,
             epoch_slot: 99,
+            network_id: NetworkId::default(),
             new_epoch: false,
             timestamp: 99999,
             era: Era::Shelley,
@@ -295,6 +296,7 @@ mod tests {
             hash: BlockHash::default(),
             epoch,
             epoch_slot: 99,
+            network_id: NetworkId::default(),
             new_epoch: true,
             timestamp: 99999,
             era: Era::Shelley,
@@ -309,6 +311,7 @@ mod tests {
             hash: BlockHash::default(),
             epoch,
             epoch_slot: 99,
+            network_id: NetworkId::default(),
             new_epoch: false,
             timestamp: 99999,
             era: Era::Conway,

modules/genesis_bootstrapper/src/genesis_bootstrapper.rs

@@ -7,8 +7,8 @@ use acropolis_common::{
         CardanoMessage, GenesisCompleteMessage, GenesisUTxOsMessage, Message, PotDeltasMessage,
         UTXODeltasMessage,
     },
-    Address, BlockHash, BlockInfo, BlockStatus, ByronAddress, Era, Lovelace, LovelaceDelta, Pot,
-    PotDelta, TxIdentifier, TxOutRef, TxOutput, UTXODelta, UTxOIdentifier, Value,
+    Address, BlockHash, BlockInfo, BlockStatus, ByronAddress, Era, Lovelace, LovelaceDelta,
+    NetworkId, Pot, PotDelta, TxIdentifier, TxOutRef, TxOutput, UTXODelta, UTxOIdentifier, Value,
 };
 use anyhow::Result;
 use blake2::{digest::consts::U32, Blake2b, Digest};
@@ -128,6 +128,7 @@ impl GenesisBootstrapper {
                     hash: BlockHash::default(),
                     epoch: 0,
                     epoch_slot: 0,
+                    network_id: NetworkId::from(network_name),
                     new_epoch: false,
                     timestamp: byron_genesis.start_time,
                     era: Era::Byron,

modules/governance_state/src/alonzo_babbage_voting.rs

@@ -113,7 +113,7 @@ mod tests {
     use crate::alonzo_babbage_voting::AlonzoBabbageVoting;
     use acropolis_common::{
         rational_number::rational_number_from_f32, AlonzoBabbageUpdateProposal,
-        AlonzoBabbageVotingOutcome, BlockHash, BlockInfo, BlockStatus, GenesisKeyhash,
+        AlonzoBabbageVotingOutcome, BlockHash, BlockInfo, BlockStatus, GenesisKeyhash, NetworkId,
         ProtocolParamUpdate,
     };
     use anyhow::Result;
@@ -154,6 +154,7 @@ mod tests {
                 epoch,
                 epoch_slot: 0,
                 era: era.try_into()?,
+                network_id: NetworkId::default(),
                 new_epoch: new_epoch != 0,
                 timestamp: 0,
                 hash: BlockHash::default(),