WIP: new lazy versions of storage values and read+write functionality

core/Cargo.toml

@@ -90,6 +90,7 @@ tendermint-proto-abcipp = {package = "tendermint-proto", git = "https://github.c
 thiserror = "1.0.30"
 tracing = "0.1.30"
 zeroize = {version = "1.5.5", features = ["zeroize_derive"]}
+regex = "1"
 
 [dev-dependencies]
 assert_matches = "1.5.0"

core/src/types/storage.rs

@@ -12,6 +12,7 @@ use borsh::{BorshDeserialize, BorshSchema, BorshSerialize};
 use data_encoding::BASE32HEX_NOPAD;
 use serde::{Deserialize, Serialize};
 use thiserror::Error;
+use regex::Regex;
 
 use crate::bytes::ByteBuf;
 use crate::types::address::{self, Address};
@@ -771,6 +772,31 @@ impl_int_key_seg!(u32, i32, 4);
 impl_int_key_seg!(u64, i64, 8);
 impl_int_key_seg!(u128, i128, 16);
 
+impl KeySeg for (Epoch, Epoch) {
+    fn parse(string: String) -> Result<Self>
+    where
+        Self: Sized,
+    {
+        let re = Regex::new(r"\((\d{1}),(\d{1})\)").unwrap();
+        let caps = re.captures(string.as_str()).unwrap();
+        let first = caps.get(1).map(|m| m.as_str().to_owned()).unwrap();
+        let second = caps.get(2).map(|m| m.as_str().to_owned()).unwrap();
+
+        let first = u64::parse(first)?;
+        let second = u64::parse(second)?;
+
+        Ok((Epoch(first), Epoch(second)))
+    }
+
+    fn raw(&self) -> String {
+        format!("({},{})",self.0,self.1)
+    }
+
+    fn to_db_key(&self) -> DbKeySeg {
+        DbKeySeg::StringSeg(self.raw())
+    }
+}
+
 impl KeySeg for Epoch {
     fn parse(string: String) -> Result<Self>
     where

proof_of_stake/src/epoched_new.rs

@@ -185,7 +185,7 @@ where
                 }
             }
             let key = self.get_last_update_storage_key();
-            storage.write(&key, expected_epoch)?;
+            storage.write(&key, current_epoch)?;
         }
         Ok(())
     }
@@ -425,7 +425,7 @@ where
                 }
             }
             let key = self.get_last_update_storage_key();
-            storage.write(&key, expected_oldest_epoch)?;
+            storage.write(&key, current_epoch)?;
         }
         Ok(())
     }

proof_of_stake/src/lib.rs

@@ -29,6 +29,7 @@ use epoched::{
     DynEpochOffset, EpochOffset, Epoched, EpochedDelta, OffsetPipelineLen,
 };
 use namada_core::ledger::storage_api::{self, StorageRead, StorageWrite};
+use namada_core::ledger::storage_api::collections::{LazyCollection, LazyMap};
 use namada_core::types::address::{self, Address, InternalAddress};
 use namada_core::types::{key::common, token, storage::Epoch};
 use parameters::PosParams;
@@ -39,7 +40,8 @@ use types::{
     Slash, SlashType, Slashes, TotalDeltas, Unbond, Unbonds,
     ValidatorConsensusKeys, ValidatorConsensusKeys_NEW, ValidatorSet,
     ValidatorSetUpdate, ValidatorSets, ValidatorState, ValidatorStates,
-    ValidatorDeltas
+    ValidatorDeltas, ValidatorStates_NEW,
+    ValidatorDeltas_NEW, ValidatorSets_NEW, BondId_NEW
 };
 
 use crate::btree_set::BTreeSetShims;
@@ -1640,6 +1642,10 @@ fn withdraw_unbonds(
     })
 }
 
+// ------------------------------------------------------------------------------------------
+// ------------------------------------------------------------------------------------------
+// ------------------------------------------------------------------------------------------
+// ------------------------------------------------------------------------------------------
 impl From<BecomeValidatorError> for storage_api::Error {
     fn from(err: BecomeValidatorError) -> Self {
         Self::new(err)
@@ -1670,6 +1676,12 @@ impl From<CommissionRateChangeError> for storage_api::Error {
     }
 }
 
+/// Get the storage handle to the Validator sets
+pub fn validator_sets_handle() -> ValidatorSets_NEW {
+    let key = storage::validator_set_key();
+    crate::epoched_new::Epoched::open(key)
+}
+
 /// Get the storage handle to a PoS validator's consensus key (used for
 /// signing block votes).
 pub fn validator_consensus_key_handle(
@@ -1679,6 +1691,33 @@ pub fn validator_consensus_key_handle(
     crate::epoched_new::Epoched::open(key)
 }
 
+/// Get the storage handle to a PoS validator's state
+pub fn validator_state_handle(
+    validator: &Address
+) -> ValidatorStates_NEW {
+    let key = storage::validator_state_key(&validator);
+    crate::epoched_new::Epoched::open(key)
+}
+
+/// Get the storage handle to a PoS validator's deltas
+pub fn validator_deltas_handle(
+    validator: &Address
+) -> ValidatorDeltas_NEW {
+    let key = storage::validator_total_deltas_key(&validator);
+    crate::epoched_new::EpochedDelta::open(key)
+}
+
+/// Get the storage handle to a bonds
+pub fn bond_handle(
+    source: &Address,
+    validator: &Address
+) -> LazyMap<Epoch, token::Amount> {
+    let bond_id = BondId_NEW {source: source.clone(), validator: validator.clone()};
+    let key = storage::bond_key(&bond_id);
+    LazyMap::open(key)
+}
+
+/// new init genesis
 pub fn init_genesis_NEW<S>(
     storage: &mut S,
     params: &PosParams,
@@ -1688,6 +1727,7 @@ pub fn init_genesis_NEW<S>(
 where
     S: for<'iter> StorageRead<'iter> + StorageWrite,
 {
+    // validator_sets_handle().init_at_genesis(storage, value, current_epoch)
     for GenesisValidator {
         address,
         tokens,
@@ -1699,6 +1739,17 @@ where
             consensus_key,
             current_epoch,
         )?;
+        validator_state_handle(&address).init_at_genesis(
+            storage,
+            ValidatorState::Candidate,
+            current_epoch
+        )?;
+        let delta = token::Change::from(tokens);
+        validator_deltas_handle(&address).init_at_genesis(
+            storage,
+            delta,
+            current_epoch
+        )?;
     }
 
     Ok(())
@@ -1733,3 +1784,112 @@ where
     let offset = OffsetPipelineLen::value(params);
     handle.set(storage, consensus_key, current_epoch, offset)
 }
+
+/// Read PoS validator's delta value.
+pub fn read_validator_delta<S>(
+    storage: &S,
+    params: &PosParams,
+    validator: &Address,
+    epoch: namada_core::types::storage::Epoch,
+) -> storage_api::Result<Option<token::Change>>
+where
+    S: for<'iter> StorageRead<'iter>,
+{
+    let handle = validator_deltas_handle(&validator);
+    handle.get_delta_val(storage, epoch, params)
+}
+
+/// Read PoS validator's stake (sum of deltas).
+pub fn read_validator_stake<S>(
+    storage: &S,
+    params: &PosParams,
+    validator: &Address,
+    epoch: namada_core::types::storage::Epoch,
+) -> storage_api::Result<Option<token::Change>>
+where
+    S: for<'iter> StorageRead<'iter>,
+{
+    let handle = validator_deltas_handle(&validator);
+    handle.get_sum(storage, epoch, params)
+}
+
+/// Write PoS validator's consensus key (used for signing block votes).
+/// Note: for EpochedDelta, write the value to change storage by
+pub fn write_validator_deltas<S>(
+    storage: &mut S,
+    params: &PosParams,
+    validator: &Address,
+    delta: token::Change,
+    current_epoch: namada_core::types::storage::Epoch,
+) -> storage_api::Result<()>
+where
+    S: for<'iter> StorageRead<'iter> + StorageWrite,
+{
+    let handle = validator_deltas_handle(&validator);
+    let offset = OffsetPipelineLen::value(params);
+
+    // TODO: either use read_validator_deltas here to update the value properly
+    // or use a new or updated method to update the Data val for EpochedDelta (set currently just sets the val like discrete Epoched)
+    handle.set(storage, delta, current_epoch, offset)
+}
+/// Read PoS validator's state.
+pub fn read_validator_state<S>(
+    storage: &S,
+    params: &PosParams,
+    validator: &Address,
+    epoch: namada_core::types::storage::Epoch,
+) -> storage_api::Result<Option<ValidatorState>>
+where
+    S: for<'iter> StorageRead<'iter>,
+{
+    let handle = validator_state_handle(&validator);
+    handle.get(storage, epoch, params)
+}
+
+/// Write PoS validator's consensus key (used for signing block votes).
+pub fn write_validator_state<S>(
+    storage: &mut S,
+    params: &PosParams,
+    validator: &Address,
+    state: ValidatorState,
+    current_epoch: namada_core::types::storage::Epoch,
+) -> storage_api::Result<()>
+where
+    S: for<'iter> StorageRead<'iter> + StorageWrite,
+{
+    let handle = validator_state_handle(&validator);
+    let offset = OffsetPipelineLen::value(params);
+    handle.set(storage, state, current_epoch, offset)
+}
+
+/// Read PoS validator's bonds
+pub fn read_bonds<S>(
+    storage: &S,
+    params: &PosParams,
+    source: &Address,
+    validator: &Address,
+    epoch: namada_core::types::storage::Epoch,
+) -> storage_api::Result<Option<LazyMap<(Epoch, Epoch), token::Amount>>>
+where
+    S: for<'iter> StorageRead<'iter>,
+{
+    // What exactly should we have this read? Should we look up by source, validator?
+    // Should it be the sum of all bonds?
+    todo!()
+}
+
+/// Write PoS validator's consensus key (used for signing block votes).
+pub fn write_bond<S>(
+    storage: &mut S,
+    params: &PosParams,
+    validator: &Address,
+    state: ValidatorState,
+    current_epoch: namada_core::types::storage::Epoch,
+) -> storage_api::Result<()>
+where
+    S: for<'iter> StorageRead<'iter> + StorageWrite,
+{
+    let handle = validator_state_handle(&validator);
+    let offset = OffsetPipelineLen::value(params);
+    handle.set(storage, state, current_epoch, offset)
+}

proof_of_stake/src/types.rs

@@ -9,6 +9,8 @@ use std::ops::{Add};
 
 use borsh::{BorshDeserialize, BorshSchema, BorshSerialize};
 use namada_core::types::address::Address;
+use namada_core::ledger::storage_api::collections::lazy_map::LazyMap;
+use namada_core::ledger::storage_api::collections::LazyCollection;
 use namada_core::types::key::common;
 use namada_core::types::storage::Epoch;
 use namada_core::types::token;
@@ -25,6 +27,37 @@ pub type ValidatorConsensusKeys_NEW = crate::epoched_new::Epoched<
     crate::epoched_new::OffsetPipelineLen,
     0,
 >;
+
+/// Epoched validator's state.
+pub type ValidatorStates_NEW = crate::epoched_new::Epoched<
+    ValidatorState,
+    crate::epoched_new::OffsetPipelineLen,
+    0
+>;
+
+/// Epoched validator sets.
+pub type ValidatorSets_NEW = crate::epoched_new::Epoched<
+    ValidatorSet_NEW,
+    crate::epoched_new::OffsetPipelineLen,
+    0
+>;
+
+/// Epoched validator's deltas.
+pub type ValidatorDeltas_NEW = crate::epoched_new::EpochedDelta<
+    token::Change,
+    // TODO: check the offsets
+    crate::epoched_new::OffsetUnbondingLen,
+    21,
+>;
+
+/// Epoched validator's bonds
+pub type Bonds_NEW = crate::epoched_new::Epoched<
+    Bond_NEW,
+    // TODO: check the offsets
+    crate::epoched_new::OffsetUnbondingLen,
+    21,
+>;
+
 /// Epoched validator's consensus key.
 pub type ValidatorConsensusKeys = Epoched<common::PublicKey, OffsetPipelineLen>;
 /// Epoched validator's state.
@@ -180,6 +213,8 @@ pub struct ValidatorSet {
     pub inactive: BTreeSet<WeightedValidator>,
 }
 
+pub type ValidatorSet_NEW = ValidatorSet<Address>;
+
 /// Validator's state.
 #[derive(
     Debug,
@@ -222,6 +257,8 @@ pub struct Bond {
     pub neg_deltas: token::Amount,
 }
 
+type Bond_NEW = LazyMap<Address, LazyMap<Address, LazyMap<(Epoch, Epoch), token::Amount>>>;
+
 /// An unbond contains unbonded tokens from a validator's self-bond or a
 /// delegation from a regular account to a validator.
 #[derive(