diff --git a/contract/TODO.md b/contract/TODO.md deleted file mode 100644 index c0ded34..0000000 --- a/contract/TODO.md +++ /dev/null @@ -1,2 +0,0 @@ -* Keep track of wallet count per outcome -* Add tests diff --git a/contract/src/execute.rs b/contract/src/execute.rs index 903dc33..98a9e0f 100644 --- a/contract/src/execute.rs +++ b/contract/src/execute.rs @@ -66,6 +66,7 @@ pub(crate) fn to_stored_outcome( label, pool_tokens, total_tokens: pool_tokens, + wallets: 0, }) }, ) @@ -133,6 +134,7 @@ fn add_market( withdrawal_stop_date, winner: None, house: deps.api.addr_validate(&house)?, + total_wallets: 0, }, )?; @@ -164,15 +166,29 @@ fn deposit( market.add_liquidity(fee); let funds = deposit_amount.checked_sub(fee)?; let tokens = market.buy(outcome, funds)?; - MARKETS.save(deps.storage, id, &market)?; let mut share_info = ShareInfo::load(deps.storage, &market, &info.sender)? .unwrap_or_else(|| ShareInfo::new(market.outcomes.len())); assert_eq!(share_info.outcomes.len(), market.outcomes.len()); + let had_any_tokens = share_info.has_tokens(); + let outcome_tokens = share_info.get_outcome_mut(id, outcome).unwrap(); + let had_these_tokens = !outcome_tokens.is_zero(); + + if !had_any_tokens { + market.total_wallets += 1; + } + + if !had_these_tokens { + market.get_outcome_mut(outcome)?.wallets += 1; + } + *outcome_tokens += tokens; share_info.save(deps.storage, &market, &info.sender)?; + + MARKETS.save(deps.storage, id, &market)?; + Ok(Response::new().add_event( Event::new("deposit") .add_attribute("market-id", id.to_string()) @@ -220,6 +236,13 @@ fn withdraw( *user_tokens -= tokens; + if user_tokens.is_zero() { + market.get_outcome_mut(outcome)?.wallets -= 1; + if !share_info.has_tokens() { + market.total_wallets -= 1; + } + } + share_info.save(deps.storage, &market, &info.sender)?; let funds = market.sell(outcome, tokens)?; diff --git a/contract/src/state.rs b/contract/src/state.rs index 934928c..93f9721 100644 --- a/contract/src/state.rs +++ b/contract/src/state.rs @@ -63,6 +63,10 @@ impl ShareInfo { outcome, }) } + + pub(crate) fn has_tokens(&self) -> bool { + self.outcomes.iter().any(|token| !token.is_zero()) + } } #[derive(Serialize, Deserialize)] @@ -80,6 +84,7 @@ pub struct StoredMarket { pub withdrawal_stop_date: Timestamp, pub winner: Option, pub house: Addr, + pub total_wallets: u32, } impl StoredMarket { @@ -117,6 +122,8 @@ pub struct StoredOutcome { pub label: String, pub pool_tokens: Token, pub total_tokens: Token, + /// Count of wallets holding tokens + pub wallets: u32, } #[derive(Serialize, Deserialize)] diff --git a/contract/src/tests.rs b/contract/src/tests.rs index 623b7dd..adf739a 100644 --- a/contract/src/tests.rs +++ b/contract/src/tests.rs @@ -168,14 +168,18 @@ impl Predict { ) } + fn query(&self, msg: &QueryMsg) -> StdResult { + self.app + .borrow() + .wrap() + .query_wasm_smart(&self.contract, msg) + } + fn query_tokens(&self, better: &Addr, outcome: u8) -> StdResult { - let PositionsResp { outcomes } = self.app.borrow().wrap().query_wasm_smart( - &self.contract, - &QueryMsg::Positions { - id: self.id, - addr: better.to_string(), - }, - )?; + let PositionsResp { outcomes } = self.query(&QueryMsg::Positions { + id: self.id, + addr: better.to_string(), + })?; outcomes .get(usize::from(outcome)) .copied() @@ -199,6 +203,14 @@ impl Predict { fn collect(&self, addr: &Addr) -> AnyResult { self.execute(addr, &ExecuteMsg::Collect { id: self.id }, None) } + + fn query_wallet_count(&self) -> StdResult<(u32, Vec)> { + let resp = self.query::(&QueryMsg::Market { id: self.id })?; + Ok(( + resp.total_wallets, + resp.outcomes.iter().map(|o| o.wallets).collect(), + )) + } } #[test] @@ -343,6 +355,34 @@ fn invalid_outcome_ids() { app.set_winner(&app.arbitrator, 0).unwrap(); } +#[test] +fn wallet_count() { + let app = Predict::new(); + + assert_eq!(app.query_wallet_count().unwrap(), (0, vec![0, 0])); + + app.place_bet(&app.better, 0, 1_000).unwrap(); + assert_eq!(app.query_wallet_count().unwrap(), (1, vec![1, 0])); + + app.place_bet(&app.better, 1, 1_000).unwrap(); + assert_eq!(app.query_wallet_count().unwrap(), (1, vec![1, 1])); + + app.place_bet(&app.admin, 0, 1_000).unwrap(); + assert_eq!(app.query_wallet_count().unwrap(), (2, vec![2, 1])); + + let tokens0 = app.query_tokens(&app.better, 0).unwrap(); + app.withdraw(&app.better, 0, tokens0).unwrap(); + assert_eq!(app.query_wallet_count().unwrap(), (2, vec![1, 1])); + + let tokens1 = app.query_tokens(&app.better, 1).unwrap(); + app.withdraw(&app.better, 1, tokens1).unwrap(); + assert_eq!(app.query_wallet_count().unwrap(), (1, vec![1, 0])); + + let tokens0 = app.query_tokens(&app.admin, 0).unwrap(); + app.withdraw(&app.admin, 0, tokens0).unwrap(); + assert_eq!(app.query_wallet_count().unwrap(), (0, vec![0, 0])); +} + proptest! { #[test] fn test_cpmm_buy_sell(pool_one in 1..1000u32, pool_two in 1..1000u32, buy in 2..50u32) { @@ -383,6 +423,7 @@ fn test_cpmm_buy_sell(pool_one in 1..1000u32, pool_two in 1..1000u32, buy in 2.. withdrawal_stop_date: ts.plus_days(1), winner: None, house: Addr::unchecked("house"), + total_wallets: 0 }; let yes_id = OutcomeId::from(0); let yes_tokens = stored.buy(yes_id, buy).unwrap();