Add support for multiple tranches

contracts/atom_wars/src/contract.rs

@@ -16,8 +16,7 @@ use crate::error::ContractError;
 use crate::msg::{ExecuteMsg, InstantiateMsg};
 use crate::query::{QueryMsg, RoundProposalsResponse, UserLockupsResponse};
 use crate::state::{
-    Constants, LockEntry, Proposal, Round, Vote, CONSTANTS, LOCKS_MAP, LOCK_ID, PROPOSAL_MAP,
-    PROPS_BY_SCORE, PROP_ID, ROUND_ID, ROUND_MAP, TOTAL_POWER_VOTING, VOTE_MAP,
+    Constants, LockEntry, Proposal, Round, Vote, Tranche, CONSTANTS, LOCKS_MAP, LOCK_ID, PROPOSAL_MAP, PROPS_BY_SCORE, PROP_ID, ROUND_ID, ROUND_MAP, TOTAL_POWER_VOTING, TRANCHE_MAP, VOTE_MAP
 };
 
 pub const ONE_MONTH_IN_NANO_SECONDS: u64 = 2629746000000000; // 365 days / 12
@@ -55,7 +54,12 @@ pub fn instantiate(
             round_end: env.block.time.plus_nanos(msg.round_length),
         },
     )?;
-    TOTAL_POWER_VOTING.save(deps.storage, round_id, &Uint128::zero())?;
+
+    // For each tranche, create a tranche in the TRANCHE_MAP and set the total power to 0
+    for tranche in msg.tranches {
+        TRANCHE_MAP.save(deps.storage, tranche.tranche_id, &tranche)?;
+        TOTAL_POWER_VOTING.save(deps.storage, (round_id, tranche.tranche_id), &Uint128::zero())?;
+    }
 
     Ok(Response::new()
         .add_attribute("action", "initialisation")
@@ -73,8 +77,8 @@ pub fn execute(
     match msg {
         ExecuteMsg::LockTokens { lock_duration } => lock_tokens(deps, env, info, lock_duration),
         ExecuteMsg::UnlockTokens {} => unlock_tokens(deps, env, info),
-        ExecuteMsg::CreateProposal { covenant_params } => create_proposal(deps, covenant_params),
-        ExecuteMsg::Vote { proposal_id } => vote(deps, info, proposal_id),
+        ExecuteMsg::CreateProposal { tranche_id, covenant_params } => create_proposal(deps, tranche_id, covenant_params),
+        ExecuteMsg::Vote { tranche_id, proposal_id } => vote(deps, info, tranche_id, proposal_id),
         ExecuteMsg::EndRound {} => end_round(deps, env, info),
         // ExecuteMsg::ExecuteProposal { proposal_id } => {
         //     execute_proposal(deps, env, info, proposal_id)
@@ -186,23 +190,24 @@ fn validate_covenant_params(_covenant_params: String) -> Result<(), ContractErro
 //     Validate covenant_params
 //     Hold tribute in contract's account
 //     Create in PropMap
-fn create_proposal(deps: DepsMut, covenant_params: String) -> Result<Response, ContractError> {
+fn create_proposal(deps: DepsMut, tranche_id: u64, covenant_params: String) -> Result<Response, ContractError> {
     validate_covenant_params(covenant_params.clone())?;
 
     let round_id = ROUND_ID.load(deps.storage)?;
 
     // Create proposal in PropMap
     let proposal = Proposal {
-        covenant_params,
         round_id,
+        tranche_id,
+        covenant_params,
         executed: false,
         power: Uint128::zero(),
         percentage: Uint128::zero(),
     };
 
     let prop_id = PROP_ID.load(deps.storage)?;
     PROP_ID.save(deps.storage, &(prop_id + 1))?;
-    PROPOSAL_MAP.save(deps.storage, (round_id, prop_id), &proposal)?;
+    PROPOSAL_MAP.save(deps.storage, (round_id, tranche_id, prop_id), &proposal)?;
 
     Ok(Response::new().add_attribute("action", "create_proposal"))
 }
@@ -230,7 +235,7 @@ fn scale_lockup_power(lockup_time: u64, raw_power: Uint128) -> Uint128 {
     scaled_power
 }
 
-fn vote(deps: DepsMut, info: MessageInfo, proposal_id: u64) -> Result<Response, ContractError> {
+fn vote(deps: DepsMut, info: MessageInfo, tranche_id: u64, proposal_id: u64) -> Result<Response, ContractError> {
     // This voting system is designed to allow for an unlimited number of proposals and an unlimited number of votes
     // to be created, without being vulnerable to DOS. A naive implementation, where all votes or all proposals were iterated
     // at the end of the round could be DOSed by creating a large number of votes or proposals. This is not a problem
@@ -252,36 +257,36 @@ fn vote(deps: DepsMut, info: MessageInfo, proposal_id: u64) -> Result<Response,
     // Get any existing vote for this sender and reverse it- this may be a vote for a different proposal (if they are switching their vote),
     // or it may be a vote for the same proposal (if they have increased their power by locking more and want to update their vote).
     // TODO: this could be made more gas-efficient by using a separate path with fewer writes if the vote is for the same proposal
-    let vote = VOTE_MAP.load(deps.storage, (round_id, info.sender.clone()));
+    let vote = VOTE_MAP.load(deps.storage, (round_id, tranche_id, info.sender.clone()));
     if let Ok(vote) = vote {
         // Load the proposal in the vote
-        let mut proposal = PROPOSAL_MAP.load(deps.storage, (round_id, vote.prop_id))?;
+        let mut proposal = PROPOSAL_MAP.load(deps.storage, (round_id, tranche_id, vote.prop_id))?;
 
         // Remove proposal's old power in PROPS_BY_SCORE
         PROPS_BY_SCORE.remove(
             deps.storage,
-            (round_id, proposal.power.into(), vote.prop_id),
+            ((round_id, proposal.tranche_id), proposal.power.into(), vote.prop_id),
         );
 
         // Decrement proposal's power
         proposal.power -= vote.power;
 
         // Save the proposal
-        PROPOSAL_MAP.save(deps.storage, (round_id, vote.prop_id), &proposal)?;
+        PROPOSAL_MAP.save(deps.storage, (round_id, tranche_id, vote.prop_id), &proposal)?;
 
         // Add proposal's new power in PROPS_BY_SCORE
         PROPS_BY_SCORE.save(
             deps.storage,
-            (round_id, proposal.power.into(), vote.prop_id),
+            ((round_id, proposal.tranche_id), proposal.power.into(), vote.prop_id),
             &vote.prop_id,
         )?;
 
         // Decrement total power voting
-        let total_power_voting = TOTAL_POWER_VOTING.load(deps.storage, round_id)?;
-        TOTAL_POWER_VOTING.save(deps.storage, round_id, &(total_power_voting - vote.power))?;
+        let total_power_voting = TOTAL_POWER_VOTING.load(deps.storage, (round_id, tranche_id))?;
+        TOTAL_POWER_VOTING.save(deps.storage, (round_id, tranche_id), &(total_power_voting - vote.power))?;
 
         // Delete vote
-        VOTE_MAP.remove(deps.storage, (round_id, info.sender.clone()));
+        VOTE_MAP.remove(deps.storage, (round_id, tranche_id, info.sender.clone()));
     }
 
     // Get sender's total locked power
@@ -305,34 +310,34 @@ fn vote(deps: DepsMut, info: MessageInfo, proposal_id: u64) -> Result<Response,
     }
 
     // Load the proposal being voted on
-    let mut proposal = PROPOSAL_MAP.load(deps.storage, (round_id, proposal_id))?;
+    let mut proposal = PROPOSAL_MAP.load(deps.storage, (round_id, tranche_id, proposal_id))?;
 
     // Delete the proposal's old power in PROPS_BY_SCORE
-    PROPS_BY_SCORE.remove(deps.storage, (round_id, proposal.power.into(), proposal_id));
+    PROPS_BY_SCORE.remove(deps.storage, ((round_id, tranche_id), proposal.power.into(), proposal_id));
 
     // Update proposal's power
     proposal.power += power;
 
     // Save the proposal
-    PROPOSAL_MAP.save(deps.storage, (round_id, proposal_id), &proposal)?;
+    PROPOSAL_MAP.save(deps.storage, (round_id, tranche_id, proposal_id), &proposal)?;
 
     // Save the proposal's new power in PROPS_BY_SCORE
     PROPS_BY_SCORE.save(
         deps.storage,
-        (round_id, proposal.power.into(), proposal_id),
+        ((round_id, tranche_id), proposal.power.into(), proposal_id),
         &proposal_id,
     )?;
 
     // Increment total power voting
-    let total_power_voting = TOTAL_POWER_VOTING.load(deps.storage, round_id)?;
-    TOTAL_POWER_VOTING.save(deps.storage, round_id, &(total_power_voting + power))?;
+    let total_power_voting = TOTAL_POWER_VOTING.load(deps.storage, (round_id, tranche_id))?;
+    TOTAL_POWER_VOTING.save(deps.storage, (round_id, tranche_id), &(total_power_voting + power))?;
 
     // Create vote in Votemap
     let vote = Vote {
         prop_id: proposal_id,
         power,
     };
-    VOTE_MAP.save(deps.storage, (round_id, info.sender), &vote)?;
+    VOTE_MAP.save(deps.storage, (round_id, tranche_id, info.sender), &vote)?;
 
     Ok(Response::new().add_attribute("action", "vote"))
 }
@@ -366,8 +371,17 @@ fn end_round(deps: DepsMut, env: Env, _info: MessageInfo) -> Result<Response, Co
             round_id,
         },
     )?;
-    // Initialize total voting power for new round
-    TOTAL_POWER_VOTING.save(deps.storage, round_id, &Uint128::zero())?;
+
+    let tranches = TRANCHE_MAP
+        .range(deps.storage, None, None, Order::Ascending)
+        .map(|t| t.unwrap().1)
+        .collect::<Vec<_>>();
+
+    // Iterate through each tranche
+    for tranche in tranches {
+        // Initialize total voting power for new round
+        TOTAL_POWER_VOTING.save(deps.storage, (round_id, tranche.tranche_id), &Uint128::zero())?;
+    }
 
     Ok(Response::new().add_attribute("action", "tally"))
 }
@@ -391,16 +405,18 @@ pub fn query(deps: Deps, _env: Env, msg: QueryMsg) -> StdResult<Binary> {
         }
         QueryMsg::Proposal {
             round_id,
+            tranche_id, 
             proposal_id,
-        } => to_json_binary(&query_proposal(deps, round_id, proposal_id)?),
-        QueryMsg::RoundProposals { round_id } => {
-            to_json_binary(&query_round_proposals(deps, round_id)?)
+        } => to_json_binary(&query_proposal(deps, round_id, tranche_id, proposal_id)?),
+        QueryMsg::RoundProposals { round_id , tranche_id} => {
+            to_json_binary(&query_round_tranche_proposals(deps, round_id, tranche_id)?)
         }
         QueryMsg::CurrentRound {} => to_json_binary(&query_current_round(deps)?),
         QueryMsg::TopNProposals {
             round_id,
+            tranche_id,
             number_of_proposals,
-        } => to_json_binary(&query_top_n_proposals(deps, round_id, number_of_proposals)?),
+        } => to_json_binary(&query_top_n_proposals(deps, round_id, tranche_id, number_of_proposals)?),
     }
 }
 
@@ -425,18 +441,18 @@ pub fn query_all_user_lockups(deps: Deps, address: String) -> StdResult<UserLock
     })
 }
 
-pub fn query_proposal(deps: Deps, round_id: u64, proposal_id: u64) -> StdResult<Proposal> {
-    Ok(PROPOSAL_MAP.load(deps.storage, (round_id, proposal_id))?)
+pub fn query_proposal(deps: Deps, round_id: u64, tranche_id: u64, proposal_id: u64) -> StdResult<Proposal> {
+    Ok(PROPOSAL_MAP.load(deps.storage, (round_id, tranche_id, proposal_id))?)
 }
 
-pub fn query_round_proposals(deps: Deps, round_id: u64) -> StdResult<RoundProposalsResponse> {
+pub fn query_round_tranche_proposals(deps: Deps, round_id: u64, tranche_id: u64) -> StdResult<RoundProposalsResponse> {
     // check if the round exists so that we can make distinction between non-existing round and round without proposals
     if let Err(_) = ROUND_MAP.may_load(deps.storage, round_id) {
         return Err(StdError::generic_err("Round does not exist"));
     }
 
     let props = PROPOSAL_MAP
-        .prefix(round_id)
+        .prefix((round_id, tranche_id))
         .range(deps.storage, None, None, Order::Ascending);
 
     let mut proposals = vec![];
@@ -452,15 +468,15 @@ pub fn query_current_round(deps: Deps) -> StdResult<Round> {
     Ok(ROUND_MAP.load(deps.storage, ROUND_ID.load(deps.storage)?)?)
 }
 
-pub fn query_top_n_proposals(deps: Deps, round_id: u64, num: usize) -> StdResult<Vec<Proposal>> {
+pub fn query_top_n_proposals(deps: Deps, round_id: u64, tranche_id: u64, num: usize) -> StdResult<Vec<Proposal>> {
     // check if the round exists
     if let Err(_) = ROUND_MAP.may_load(deps.storage, round_id) {
         return Err(StdError::generic_err("Round does not exist"));
     }
 
     // Iterate through PROPS_BY_SCORE to find the top ten props
     let top_prop_ids: Vec<u64> = PROPS_BY_SCORE
-        .sub_prefix(round_id)
+        .sub_prefix((round_id, tranche_id))
         .range(deps.storage, None, None, Order::Descending)
         .take(num)
         .map(|x| match x {
@@ -472,7 +488,7 @@ pub fn query_top_n_proposals(deps: Deps, round_id: u64, num: usize) -> StdResult
     let mut top_props = vec![];
 
     for prop_id in top_prop_ids {
-        let prop = PROPOSAL_MAP.load(deps.storage, (round_id, prop_id))?;
+        let prop = PROPOSAL_MAP.load(deps.storage, (round_id, tranche_id, prop_id))?;
         top_props.push(prop);
     }
 
@@ -491,3 +507,12 @@ pub fn query_top_n_proposals(deps: Deps, round_id: u64, num: usize) -> StdResult
         })
         .collect());
 }
+
+pub fn query_tranches (deps: Deps) -> StdResult<Vec<Tranche>> {
+    let tranches = TRANCHE_MAP
+        .range(deps.storage, None, None, Order::Ascending)
+        .map(|t| t.unwrap().1)
+        .collect::<Vec<_>>();
+
+    Ok(tranches)
+}

contracts/atom_wars/src/msg.rs

@@ -2,20 +2,24 @@ use cosmwasm_std::Uint128;
 use schemars::JsonSchema;
 use serde::{Deserialize, Serialize};
 
+use crate::state::Tranche;
+
+
 #[derive(Serialize, Deserialize, Clone, Debug, PartialEq, JsonSchema)]
 pub struct InstantiateMsg {
     pub denom: String,
     pub round_length: u64,
     pub total_pool: Uint128,
+    pub tranches: Vec<Tranche>,
 }
 
 #[derive(Serialize, Deserialize, Clone, Debug, PartialEq, JsonSchema)]
 #[serde(rename_all = "snake_case")]
 pub enum ExecuteMsg {
     LockTokens { lock_duration: u64 },
     UnlockTokens {},
-    CreateProposal { covenant_params: String },
-    Vote { proposal_id: u64 },
+    CreateProposal { tranche_id: u64, covenant_params: String },
+    Vote { tranche_id: u64, proposal_id: u64 },
     EndRound {},
     // ExecuteProposal { proposal_id: u64 },
 }

contracts/atom_wars/src/query.rs

@@ -14,13 +14,16 @@ pub enum QueryMsg {
     // Round { round_id } ?
     RoundProposals {
         round_id: u64,
+        tranche_id: u64
     },
     Proposal {
         round_id: u64,
+        tranche_id: u64,
         proposal_id: u64,
     },
     TopNProposals {
         round_id: u64,
+        tranche_id: u64,
         number_of_proposals: usize,
     },
 }

contracts/atom_wars/src/state.rs

@@ -30,28 +30,30 @@ pub struct LockEntry {
     pub lock_end: Timestamp,
 }
 
-// PROP_MAP: key(round_id, prop_id) -> Proposal {
+// PROP_MAP: key(round_id, tranche_id, prop_id) -> Proposal {
 //     round_id: u64,
+//     tranche_id: u64,
 //     covenant_params: String,
 //     executed: bool,
 //     power: Uint128
 // }
-pub const PROPOSAL_MAP: Map<(u64, u64), Proposal> = Map::new("prop_map");
+pub const PROPOSAL_MAP: Map<(u64, u64, u64), Proposal> = Map::new("prop_map");
 #[cw_serde]
 pub struct Proposal {
     pub round_id: u64,
+    pub tranche_id: u64,
     pub covenant_params: String,
     pub executed: bool, // TODO: maybe remove in the future
     pub power: Uint128,
     pub percentage: Uint128,
 }
 
-// VOTE_MAP: key(round_id, sender_addr) -> Vote {
+// VOTE_MAP: key(round_id, tranche_id, sender_addr) -> Vote {
 //     prop_id: u64,
 //     power: Uint128,
 //     tribute_claimed: bool
 // }
-pub const VOTE_MAP: Map<(u64, Addr), Vote> = Map::new("vote_map");
+pub const VOTE_MAP: Map<(u64, u64, Addr), Vote> = Map::new("vote_map");
 #[cw_serde]
 pub struct Vote {
     pub prop_id: u64,
@@ -69,8 +71,19 @@ pub struct Round {
     pub round_end: Timestamp,
 }
 
-// PROPS_BY_SCORE: key(round_id, score, prop_id) -> prop_id
-pub const PROPS_BY_SCORE: Map<(u64, u128, u64), u64> = Map::new("props_by_score");
+// PROPS_BY_SCORE: key((round_id, tranche_id), score, prop_id) -> prop_id
+pub const PROPS_BY_SCORE: Map<((u64, u64), u128, u64), u64> = Map::new("props_by_score");
 
-// TOTAL_POWER_VOTING: key(round_id) -> Uint128
-pub const TOTAL_POWER_VOTING: Map<u64, Uint128> = Map::new("total_power_voting");
+// TOTAL_POWER_VOTING: key(round_id, tranche_id) -> Uint128
+pub const TOTAL_POWER_VOTING: Map<(u64, u64), Uint128> = Map::new("total_power_voting");
+
+// TRANCHE_MAP: key(tranche_id) -> Tranche {
+//     tranche_id: u64,
+//     metadata: String
+// }
+pub const TRANCHE_MAP: Map<u64, Tranche> = Map::new("tranche_map");
+#[cw_serde]
+pub struct Tranche {
+    pub tranche_id: u64,
+    pub metadata: String,
+}