Merge pull request #122 from input-output-hk/whankinsiv/drep-state-expansion

feat: Expanded drep_state with historical state and rollback support

Author: sandtreader

common/src/messages.rs

@@ -129,7 +129,7 @@ pub struct GovernanceProceduresMessage {
     pub proposal_procedures: Vec<ProposalProcedure>,
 
     /// Voting
-    pub voting_procedures: Vec<(DataHash, VotingProcedures)>,
+    pub voting_procedures: Vec<([u8; 32], VotingProcedures)>,
 
     /// Alonzo-compatible (from Shelley) and Babbage updates
     pub alonzo_babbage_updates: Vec<AlonzoBabbageUpdateProposal>,

common/src/queries/accounts.rs

@@ -1,5 +1,10 @@
+use std::collections::HashMap;
+
 use crate::{DRepChoice, KeyHash};
 
+pub const DEFAULT_ACCOUNTS_QUERY_TOPIC: (&str, &str) =
+    ("accounts-state-query-topic", "cardano.query.accounts");
+
 #[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
 pub enum AccountsStateQuery {
     GetAccountInfo { stake_key: Vec<u8> },
@@ -16,6 +21,9 @@ pub enum AccountsStateQuery {
 
     // Pools related queries
     GetPoolsLiveStakes { pools_operators: Vec<Vec<u8>> },
+
+    // Dreps related queries
+    GetAccountsDrepDelegationsMap { stake_keys: Vec<Vec<u8>> },
 }
 
 #[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
@@ -35,6 +43,9 @@ pub enum AccountsStateQueryResponse {
     // Pools related responses
     PoolsLiveStakes(PoolsLiveStakes),
 
+    // DReps related responses
+    AccountsDrepDelegationsMap(HashMap<Vec<u8>, Option<DRepChoice>>),
+
     NotFound,
     Error(String),
 }

common/src/queries/mod.rs

@@ -1,3 +1,7 @@
+use crate::messages::Message;
+use caryatid_sdk::Context;
+use std::sync::Arc;
+
 pub mod accounts;
 pub mod addresses;
 pub mod assets;
@@ -12,3 +16,7 @@ pub mod pools;
 pub mod scripts;
 pub mod transactions;
 pub mod utils;
+
+pub fn get_query_topic(context: Arc<Context<Message>>, topic: (&str, &str)) -> String {
+    context.config.get_string(topic.0).unwrap_or_else(|_| topic.1.to_string())
+}

common/src/state_history.rs

@@ -89,3 +89,17 @@ impl<S: Clone + Default> StateHistory<S> {
         });
     }
 }
+
+/// Helper that lets callers initialize the first state with custom config.
+impl<S: Clone> StateHistory<S> {
+    pub fn get_or_init_with<F>(&mut self, init: F) -> S
+    where
+        F: FnOnce() -> S,
+    {
+        if let Some(current) = self.history.back() {
+            current.state.clone()
+        } else {
+            init()
+        }
+    }
+}

common/src/types.rs

@@ -658,6 +658,13 @@ pub struct DRepRegistration {
     pub anchor: Option<Anchor>,
 }
 
+#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
+pub struct DRepRegistrationWithPos {
+    pub reg: DRepRegistration,
+    pub tx_hash: [u8; 32],
+    pub cert_index: u64,
+}
+
 /// DRep Deregistration = unreg_drep_cert
 #[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
 pub struct DRepDeregistration {
@@ -668,6 +675,13 @@ pub struct DRepDeregistration {
     pub refund: Lovelace,
 }
 
+#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
+pub struct DRepDeregistrationWithPos {
+    pub reg: DRepDeregistration,
+    pub tx_hash: [u8; 32],
+    pub cert_index: u64,
+}
+
 /// DRep Update = update_drep_cert
 #[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
 pub struct DRepUpdate {
@@ -678,6 +692,13 @@ pub struct DRepUpdate {
     pub anchor: Option<Anchor>,
 }
 
+#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
+pub struct DRepUpdateWithPos {
+    pub reg: DRepUpdate,
+    pub tx_hash: [u8; 32],
+    pub cert_index: u64,
+}
+
 pub type CommitteeCredential = Credential;
 
 /// Authorise a committee hot credential
@@ -1209,6 +1230,7 @@ pub enum Vote {
 pub struct VotingProcedure {
     pub vote: Vote,
     pub anchor: Option<Anchor>,
+    pub vote_index: u32,
 }
 
 #[serde_as]
@@ -1375,13 +1397,13 @@ pub enum TxCertificate {
     ResignCommitteeCold(ResignCommitteeCold),
 
     /// DRep registration
-    DRepRegistration(DRepRegistration),
+    DRepRegistration(DRepRegistrationWithPos),
 
     /// DRep deregistration
-    DRepDeregistration(DRepDeregistration),
+    DRepDeregistration(DRepDeregistrationWithPos),
 
     /// DRep update
-    DRepUpdate(DRepUpdate),
+    DRepUpdate(DRepUpdateWithPos),
 }
 
 #[cfg(test)]
@@ -1440,6 +1462,7 @@ mod tests {
             VotingProcedure {
                 anchor: None,
                 vote: Vote::Abstain,
+                vote_index: 0,
             },
         );
         voting.votes.insert(

modules/accounts_state/src/accounts_state.rs

@@ -439,6 +439,15 @@ impl AccountsState {
                         })
                     }
 
+                    AccountsStateQuery::GetAccountsDrepDelegationsMap { stake_keys } => match state
+                        .get_drep_delegations_map(stake_keys)
+                    {
+                        Some(map) => AccountsStateQueryResponse::AccountsDrepDelegationsMap(map),
+                        None => AccountsStateQueryResponse::Error(
+                            "Error retrieving DRep delegations map".to_string(),
+                        ),
+                    },
+
                     _ => AccountsStateQueryResponse::Error(format!(
                         "Unimplemented query variant: {:?}",
                         query

modules/accounts_state/src/state.rs

@@ -133,6 +133,24 @@ impl State {
             .collect()
     }
 
+    /// Map stake_keys to their delegated DRep
+    pub fn get_drep_delegations_map(
+        &self,
+        stake_keys: &[Vec<u8>],
+    ) -> Option<HashMap<Vec<u8>, Option<DRepChoice>>> {
+        let accounts = self.stake_addresses.lock().ok()?; // If lock fails, return None
+
+        let mut map = HashMap::new();
+
+        for stake_key in stake_keys {
+            let account = accounts.get(stake_key)?;
+            let maybe_drep = account.delegated_drep.clone();
+            map.insert(stake_key.clone(), maybe_drep);
+        }
+
+        Some(map)
+    }
+
     /// Log statistics
     fn log_stats(&self) {
         info!(num_stake_addresses = self.stake_addresses.lock().unwrap().keys().len(),);