Add liquidity sources to the contract (#32)

bisenzone-cw-mvp/contracts/cw-tee-mtcs/src/contract.rs

@@ -15,7 +15,10 @@ use crate::{
         execute::{SubmitObligationMsg, SubmitObligationsMsg, SubmitSetoffsMsg},
         ExecuteMsg, InstantiateMsg, QueryMsg,
     },
-    state::{current_epoch_key, ObligationsItem, State, OBLIGATIONS_KEY, STATE},
+    state::{
+        current_epoch_key, LiquiditySourcesItem, ObligationsItem, State, LIQUIDITY_SOURCES_KEY,
+        OBLIGATIONS_KEY, STATE,
+    },
 };
 
 // version info for migration info
@@ -43,6 +46,9 @@ pub fn instantiate(
     ObligationsItem::new(&current_epoch_key(OBLIGATIONS_KEY, deps.storage)?)
         .save(deps.storage, &Default::default())?;
 
+    LiquiditySourcesItem::new(&current_epoch_key(LIQUIDITY_SOURCES_KEY, deps.storage)?)
+        .save(deps.storage, &Default::default())?;
+
     // store token info using cw20-base format
     let data = TokenInfo {
         name: "USD".to_string(),
@@ -103,10 +109,14 @@ pub fn execute(
         ExecuteMsg::SubmitObligation(SubmitObligationMsg { ciphertext, digest }) => {
             execute::submit_obligation(deps, ciphertext, digest)
         }
-        ExecuteMsg::SubmitObligations(SubmitObligationsMsg(obligations)) => {
+        ExecuteMsg::SubmitObligations(SubmitObligationsMsg {
+            obligations,
+            liquidity_sources,
+        }) => {
             for o in obligations {
                 execute::submit_obligation(deps.branch(), o.ciphertext, o.digest)?;
             }
+            execute::append_liquidity_sources(deps, liquidity_sources)?;
             Ok(Response::new())
         }
         ExecuteMsg::SubmitSetoffs(SubmitSetoffsMsg { setoffs_enc }) => {
@@ -125,8 +135,8 @@ pub mod execute {
 
     use crate::{
         state::{
-            current_epoch_key, previous_epoch_key, ObligationsItem, RawHash, SetoffsItem,
-            SettleOff, OBLIGATIONS_KEY, SETOFFS_KEY,
+            current_epoch_key, previous_epoch_key, LiquiditySourcesItem, ObligationsItem, RawHash,
+            SetoffsItem, SettleOff, LIQUIDITY_SOURCES_KEY, OBLIGATIONS_KEY, SETOFFS_KEY,
         },
         ContractError,
     };
@@ -155,6 +165,24 @@ pub mod execute {
             .add_attribute("ciphertext", ciphertext.to_string()))
     }
 
+    pub fn append_liquidity_sources(
+        deps: DepsMut,
+        liquidity_sources: Vec<String>,
+    ) -> Result<(), ContractError> {
+        liquidity_sources
+            .iter()
+            .try_for_each(|ls| deps.api.addr_validate(ls).map(|_| ()))?;
+
+        // store the liquidity sources
+        LiquiditySourcesItem::new(&current_epoch_key(LIQUIDITY_SOURCES_KEY, deps.storage)?)
+            .update(deps.storage, |mut ls| {
+                ls.extend(liquidity_sources);
+                Ok::<_, ContractError>(ls)
+            })?;
+
+        Ok(())
+    }
+
     pub fn submit_setoffs(
         mut deps: DepsMut,
         env: Env,
@@ -206,6 +234,9 @@ pub mod execute {
 pub fn query(deps: Deps, _env: Env, msg: QueryMsg) -> StdResult<Binary> {
     match msg {
         QueryMsg::GetAllSetoffs => to_json_binary(&query::get_all_setoffs(deps)?),
+        QueryMsg::GetLiquiditySources { epoch } => {
+            to_json_binary(&query::get_liquidity_sources(deps, epoch)?)
+        }
         QueryMsg::Balance { address } => to_json_binary(&cw20_query_balance(deps, address)?),
     }
 }
@@ -214,8 +245,11 @@ pub mod query {
     use cosmwasm_std::{Deps, StdResult};
 
     use crate::{
-        msg::GetAllSetoffsResponse,
-        state::{previous_epoch_key, SetoffsItem, SETOFFS_KEY},
+        msg::{GetAllSetoffsResponse, GetLiquiditySourcesResponse},
+        state::{
+            current_epoch_key, epoch_key, previous_epoch_key, LiquiditySourcesItem, SetoffsItem,
+            LIQUIDITY_SOURCES_KEY, SETOFFS_KEY,
+        },
     };
 
     pub fn get_all_setoffs(deps: Deps) -> StdResult<GetAllSetoffsResponse> {
@@ -225,4 +259,20 @@ pub mod query {
             .collect();
         Ok(GetAllSetoffsResponse { setoffs })
     }
+
+    pub fn get_liquidity_sources(
+        deps: Deps,
+        epoch: Option<usize>,
+    ) -> StdResult<GetLiquiditySourcesResponse> {
+        let epoch_key = match epoch {
+            None => current_epoch_key(LIQUIDITY_SOURCES_KEY, deps.storage)?,
+            Some(e) => epoch_key(LIQUIDITY_SOURCES_KEY, e)?,
+        };
+
+        let liquidity_sources = LiquiditySourcesItem::new(&epoch_key)
+            .load(deps.storage)?
+            .into_iter()
+            .collect();
+        Ok(GetLiquiditySourcesResponse { liquidity_sources })
+    }
 }

bisenzone-cw-mvp/contracts/cw-tee-mtcs/src/msg.rs

@@ -31,8 +31,10 @@ pub mod execute {
     }
 
     #[cw_serde]
-    #[serde(transparent)]
-    pub struct SubmitObligationsMsg(pub Vec<SubmitObligationMsg>);
+    pub struct SubmitObligationsMsg {
+        pub obligations: Vec<SubmitObligationMsg>,
+        pub liquidity_sources: Vec<String>,
+    }
 
     #[cw_serde]
     pub struct SubmitTenderMsg {
@@ -47,11 +49,14 @@ pub mod execute {
         // pub proof: π,
     }
 }
+
 #[cw_serde]
 #[derive(QueryResponses)]
 pub enum QueryMsg {
     #[returns(GetAllSetoffsResponse)]
     GetAllSetoffs,
+    #[returns(GetLiquiditySourcesResponse)]
+    GetLiquiditySources { epoch: Option<usize> }, // `None` means latest
     #[returns(cw20::BalanceResponse)]
     Balance { address: String },
 }
@@ -62,6 +67,11 @@ pub struct GetAllSetoffsResponse {
     pub setoffs: Vec<(HexBinary, SettleOff)>,
 }
 
+#[cw_serde]
+pub struct GetLiquiditySourcesResponse {
+    pub liquidity_sources: Vec<String>,
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

bisenzone-cw-mvp/contracts/cw-tee-mtcs/src/state.rs

@@ -1,4 +1,4 @@
-use std::collections::BTreeMap;
+use std::collections::{BTreeMap, BTreeSet};
 
 use cosmwasm_schema::cw_serde;
 use cosmwasm_std::{HexBinary, StdError, Storage};
@@ -10,6 +10,7 @@ pub type RawCipherText = HexBinary;
 
 pub type ObligationsItem<'a> = Item<'a, BTreeMap<RawHash, RawCipherText>>;
 pub type SetoffsItem<'a> = Item<'a, BTreeMap<RawHash, SettleOff>>;
+pub type LiquiditySourcesItem<'a> = Item<'a, BTreeSet<String>>;
 
 #[cw_serde]
 pub struct State {
@@ -33,11 +34,16 @@ pub enum SettleOff {
 pub const STATE: Item<State> = Item::new("state");
 pub const OBLIGATIONS_KEY: &str = "obligations";
 pub const SETOFFS_KEY: &str = "setoffs";
+pub const LIQUIDITY_SOURCES_KEY: &str = "liquidity_sources";
 
 pub fn current_epoch_key(key: &str, storage: &dyn Storage) -> Result<String, StdError> {
-    Ok(format!("{}/{key}", EPOCH_COUNTER.load(storage)?))
+    epoch_key(key, EPOCH_COUNTER.load(storage)?)
 }
 
 pub fn previous_epoch_key(key: &str, storage: &dyn Storage) -> Result<String, StdError> {
-    Ok(format!("{}/{key}", EPOCH_COUNTER.load(storage)? - 1))
+    epoch_key(key, EPOCH_COUNTER.load(storage)? - 1)
+}
+
+pub fn epoch_key(key: &str, epoch: usize) -> Result<String, StdError> {
+    Ok(format!("{}/{key}", epoch))
 }