Replace ballots with indexed map

packages/voting-contract/src/ballots.rs

@@ -0,0 +1,74 @@
+use schemars::JsonSchema;
+use serde::{Deserialize, Serialize};
+
+use cosmwasm_std::{Addr, Storage};
+use cw_storage_plus::{Index, IndexList, IndexedMap, MultiIndex};
+use tg3::Vote;
+
+use crate::ContractError;
+
+// we cast a ballot with our chosen vote and a given points
+#[derive(Serialize, Deserialize, Clone, PartialEq, JsonSchema, Debug)]
+pub struct Ballot {
+    pub voter: Addr,
+    pub proposal_id: u64,
+    pub points: u64,
+    pub vote: Vote,
+}
+
+pub struct BallotIndexes<'a> {
+    // This PrimaryKey allows quering over all proposal ids for given voter address
+    pub voter: MultiIndex<'a, Addr, Ballot, (u64, Addr)>,
+}
+
+impl<'a> IndexList<Ballot> for BallotIndexes<'a> {
+    fn get_indexes(&'_ self) -> Box<dyn Iterator<Item = &'_ dyn Index<Ballot>> + '_> {
+        let v: Vec<&dyn Index<Ballot>> = vec![&self.voter];
+        Box::new(v.into_iter())
+    }
+}
+
+pub struct Ballots<'a> {
+    pub ballots: IndexedMap<'a, (u64, &'a Addr), Ballot, BallotIndexes<'a>>,
+}
+
+impl<'a> Ballots<'a> {
+    pub fn new(storage_key: &'a str, release_subkey: &'a str) -> Self {
+        let indexes = BallotIndexes {
+            voter: MultiIndex::new(|ballot| ballot.voter.clone(), storage_key, release_subkey),
+        };
+        let ballots = IndexedMap::new(storage_key, indexes);
+
+        Self { ballots }
+    }
+
+    pub fn create_ballot(
+        &self,
+        storage: &mut dyn Storage,
+        addr: &Addr,
+        proposal_id: u64,
+        points: u64,
+        vote: Vote,
+    ) -> Result<(), ContractError> {
+        self.ballots.update(
+            storage,
+            (proposal_id, addr),
+            move |ballot| -> Result<_, ContractError> {
+                match ballot {
+                    Some(_) => Err(ContractError::AlreadyVoted {}),
+                    None => Ok(Ballot {
+                        voter: addr.clone(),
+                        proposal_id,
+                        points,
+                        vote,
+                    }),
+                }
+            },
+        )?;
+        Ok(())
+    }
+}
+
+pub fn ballots() -> Ballots<'static> {
+    Ballots::new("ballots", "ballots__proposal_id")
+}

packages/voting-contract/src/lib.rs

@@ -1,3 +1,4 @@
+mod ballots;
 mod error;
 pub mod msg;
 #[cfg(test)]
@@ -7,11 +8,11 @@ pub mod state;
 use serde::de::DeserializeOwned;
 use serde::Serialize;
 
+use ballots::ballots;
 pub use error::ContractError;
 use state::{
-    next_id, proposals, Ballot, Config, Proposal, ProposalListResponse, ProposalResponse,
-    TextProposalListResponse, Votes, VotingRules, BALLOTS, BALLOTS_BY_VOTER, CONFIG,
-    TEXT_PROPOSALS,
+    next_id, proposals, Config, Proposal, ProposalListResponse, ProposalResponse,
+    TextProposalListResponse, Votes, VotingRules, CONFIG, TEXT_PROPOSALS,
 };
 
 use cosmwasm_std::{
@@ -51,17 +52,6 @@ pub fn instantiate<Q: CustomQuery>(
     Ok(Response::default())
 }
 
-fn save_ballot(
-    storage: &mut dyn Storage,
-    proposal_id: u64,
-    sender: &Addr,
-    ballot: Ballot,
-) -> Result<(), ContractError> {
-    BALLOTS.save(storage, (proposal_id, sender), &ballot)?;
-    BALLOTS_BY_VOTER.save(storage, (sender, proposal_id), &ballot)?;
-    Ok(())
-}
-
 pub fn propose<P, Q: CustomQuery>(
     deps: DepsMut<Q>,
     env: Env,
@@ -103,15 +93,7 @@ where
     proposals().save(deps.storage, id, &prop)?;
 
     // add the first yes vote from voter
-    save_ballot(
-        deps.storage,
-        id,
-        &info.sender,
-        Ballot {
-            points: vote_power,
-            vote: Vote::Yes,
-        },
-    )?;
+    ballots().create_ballot(deps.storage, &info.sender, id, vote_power, Vote::Yes)?;
 
     let resp = msg::ProposalCreationResponse { proposal_id: id };
 
@@ -148,21 +130,7 @@ where
             .was_voting_member(&deps.querier, &info.sender, prop.start_height)?;
 
     // cast vote if no vote previously cast
-    if BALLOTS
-        .may_load(deps.storage, (proposal_id, &info.sender))?
-        .is_some()
-    {
-        return Err(ContractError::AlreadyVoted {});
-    }
-    save_ballot(
-        deps.storage,
-        proposal_id,
-        &info.sender,
-        Ballot {
-            points: vote_power,
-            vote,
-        },
-    )?;
+    ballots().create_ballot(deps.storage, &info.sender, proposal_id, vote_power, vote)?;
 
     // update vote tally
     prop.votes.add_vote(vote, vote_power);
@@ -362,7 +330,9 @@ pub fn query_vote<Q: CustomQuery>(
     voter: String,
 ) -> StdResult<VoteResponse> {
     let voter_addr = deps.api.addr_validate(&voter)?;
-    let prop = BALLOTS.may_load(deps.storage, (proposal_id, &voter_addr))?;
+    let prop = ballots()
+        .ballots
+        .may_load(deps.storage, (proposal_id, &voter_addr))?;
     let vote = prop.map(|b| VoteInfo {
         proposal_id,
         voter,
@@ -381,7 +351,8 @@ pub fn list_votes<Q: CustomQuery>(
     let addr = maybe_addr(deps.api, start_after)?;
     let start = addr.as_ref().map(Bound::exclusive);
 
-    let votes: StdResult<Vec<_>> = BALLOTS
+    let votes: StdResult<Vec<_>> = ballots()
+        .ballots
         .prefix(proposal_id)
         .range(deps.storage, start, None, Order::Ascending)
         .take(limit)
@@ -405,18 +376,22 @@ pub fn list_votes_by_voter<Q: CustomQuery>(
     start_after: Option<u64>,
     limit: usize,
 ) -> StdResult<VoteListResponse> {
-    let start = start_after.map(Bound::exclusive);
     let voter_addr = deps.api.addr_validate(&voter)?;
-
-    let votes: StdResult<Vec<_>> = BALLOTS_BY_VOTER
-        .prefix(&voter_addr)
+    // PrimaryKey of that IndexMap is (proposal_id, voter_address) -> (u64, Addr)
+    let start = start_after.map(|m| Bound::exclusive((m, voter_addr.clone())));
+
+    let votes: StdResult<Vec<_>> = ballots()
+        .ballots
+        .idx
+        .voter
+        .prefix(voter_addr)
         .range(deps.storage, start, None, Order::Ascending)
         .take(limit)
         .map(|item| {
-            let (proposal_id, ballot) = item?;
+            let (_, ballot) = item?;
             Ok(VoteInfo {
-                proposal_id,
-                voter: voter.clone(),
+                proposal_id: ballot.proposal_id,
+                voter: ballot.voter.into(),
                 vote: ballot.vote,
                 points: ballot.points,
             })

packages/voting-contract/src/state.rs

@@ -1,7 +1,7 @@
 use schemars::JsonSchema;
 use serde::{Deserialize, Serialize};
 
-use cosmwasm_std::{Addr, BlockInfo, Decimal, StdResult, Storage, Uint128};
+use cosmwasm_std::{BlockInfo, Decimal, StdResult, Storage, Uint128};
 use cw_storage_plus::{Item, Map};
 use tg3::{Status, Vote};
 use tg4::Tg4Contract;
@@ -253,21 +253,10 @@ fn votes_needed(points: u64, percentage: Decimal) -> u64 {
     ((applied.u128() + PRECISION_FACTOR - 1) / PRECISION_FACTOR) as u64
 }
 
-// we cast a ballot with our chosen vote and a given points
-// stored under the key that voted
-#[derive(Serialize, Deserialize, Clone, PartialEq, JsonSchema, Debug)]
-pub struct Ballot {
-    pub points: u64,
-    pub vote: Vote,
-}
-
 // unique items
 pub const CONFIG: Item<Config> = Item::new("voting_config");
 pub const PROPOSAL_COUNT: Item<u64> = Item::new("proposal_count");
 
-// multiple-item map
-pub const BALLOTS: Map<(u64, &Addr), Ballot> = Map::new("votes");
-pub const BALLOTS_BY_VOTER: Map<(&Addr, u64), Ballot> = Map::new("votes_by_voter");
 pub fn proposals<'m, P>() -> Map<'m, u64, Proposal<P>> {
     Map::new("proposals")
 }