refactor: account addresses getter to remove duplicates

Signed-off-by: William Hankins <[email protected]>

Author: whankinsiv

modules/historical_accounts_state/src/immutable_historical_account_store.rs

@@ -269,20 +269,27 @@ impl ImmutableHistoricalAccountStore {
     pub async fn get_addresses(
         &self,
         account: &StakeAddress,
-    ) -> Result<Option<HashSet<ShelleyAddress>>> {
+    ) -> Result<Option<Vec<ShelleyAddress>>> {
         let prefix = account.to_binary();
-        let mut addresses: HashSet<ShelleyAddress> = HashSet::new();
+        let mut addresses = Vec::new();
+        let mut seen = HashSet::new();
 
         for result in self.addresses.prefix(&prefix) {
             let (key, _) = result?;
-            let shelley = ShelleyAddress::from_bytes_key(&key[prefix.len()..])?;
-            addresses.insert(shelley);
+            let shelley = ShelleyAddress::from_bytes_key(&key[prefix.len() + 8..])?;
+            if seen.insert(shelley.clone()) {
+                addresses.push(shelley);
+            }
         }
 
         for block_map in self.pending.lock().await.iter() {
             if let Some(entry) = block_map.get(account) {
                 if let Some(addrs) = &entry.addresses {
-                    addresses.extend(addrs.iter().cloned());
+                    for addr in addrs {
+                        if seen.insert(addr.clone()) {
+                            addresses.push(addr.clone());
+                        }
+                    }
                 }
             }
         }

modules/historical_accounts_state/src/state.rs

@@ -358,25 +358,25 @@ impl State {
         &self,
         account: &StakeAddress,
     ) -> Result<Option<Vec<ShelleyAddress>>> {
-        let immutable = self.immutable.get_addresses(account).await?;
+        let mut addresses = match self.immutable.get_addresses(account).await? {
+            Some(list) => list,
+            None => Vec::new(),
+        };
 
-        let mut volatile = HashSet::new();
+        let mut seen: HashSet<ShelleyAddress> = addresses.iter().cloned().collect();
         for block_map in self.volatile.window.iter() {
             if let Some(entry) = block_map.get(account) {
                 if let Some(pending) = &entry.addresses {
-                    volatile.extend(pending.iter().cloned());
+                    for addr in pending {
+                        if seen.insert(addr.clone()) {
+                            addresses.push(addr.clone());
+                        }
+                    }
                 }
             }
         }
 
-        match immutable {
-            Some(mut addresses) => {
-                addresses.extend(volatile);
-                Ok(Some(addresses.into_iter().collect()))
-            }
-            None if volatile.is_empty() => Ok(None),
-            None => Ok(Some(volatile.into_iter().collect())),
-        }
+        Ok((!addresses.is_empty()).then_some(addresses))
     }
 
     fn handle_stake_registration_change(

modules/stake_delta_filter/src/utils.rs

@@ -321,7 +321,8 @@ pub fn process_message(
     block: &BlockInfo,
     mut tracker: Option<&mut Tracker>,
 ) -> StakeAddressDeltasMessage {
-    let mut grouped: HashMap<StakeAddress, (i64, HashSet<ShelleyAddress>)> = HashMap::new();
+    let mut grouped: HashMap<StakeAddress, (i64, Vec<ShelleyAddress>, HashSet<ShelleyAddress>)> =
+        HashMap::new();
     for d in delta.deltas.iter() {
         // Variants to be processed:
         // 1. Shelley Address delegation is a stake
@@ -388,24 +389,26 @@ pub fn process_message(
             Address::Stake(stake_address) => (stake_address.clone(), None),
         };
 
-        let entry = grouped.entry(stake_address).or_insert((0, HashSet::new()));
+        let entry = grouped.entry(stake_address).or_insert((0, Vec::new(), HashSet::new()));
         entry.0 += d.value.lovelace;
 
         if let Some(shelley) = shelley_opt {
-            entry.1.insert(shelley.clone());
+            if entry.2.insert(shelley.clone()) {
+                entry.1.push(shelley.clone());
+            }
         }
     }
 
     let deltas = grouped
         .into_iter()
         .map(
-            |(stake_address, (delta, shelley_addrs))| StakeAddressDelta {
+            |(stake_address, (delta, shelley_addrs, _))| StakeAddressDelta {
                 stake_address,
-                addresses: shelley_addrs.into_iter().collect(),
+                addresses: shelley_addrs,
                 delta,
             },
         )
-        .collect::<Vec<_>>();
+        .collect();
 
     StakeAddressDeltasMessage { deltas }
 }

processes/omnibus/omnibus.toml

@@ -71,7 +71,7 @@ store-registration-history = false
 store-delegation-history = false
 store-mir-history = false
 store-withdrawal-history = false
-store-addresses = false
+store-addresses = true
 
 [module.drep-state]
 # Enables /governance/dreps/{drep_id} endpoint (Requires store-delegators to be enabled)

processes/omnibus/src/main.rs

@@ -118,7 +118,7 @@ pub async fn main() -> Result<()> {
     SPDDState::register(&mut process);
     DRDDState::register(&mut process);
     Consensus::register(&mut process);
-    ChainStore::register(&mut process);
+    //ChainStore::register(&mut process);
 
     Clock::<Message>::register(&mut process);
     RESTServer::<Message>::register(&mut process);