Merge pull request #5 from Levana-Protocol/perp-4042/wallet-count

perp-4042 | include wallet count

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)?;

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<OutcomeId>,
     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)]

contract/src/tests.rs

@@ -168,14 +168,18 @@ impl Predict {
         )
     }
 
+    fn query<T: serde::de::DeserializeOwned>(&self, msg: &QueryMsg) -> StdResult<T> {
+        self.app
+            .borrow()
+            .wrap()
+            .query_wasm_smart(&self.contract, msg)
+    }
+
     fn query_tokens(&self, better: &Addr, outcome: u8) -> StdResult<Token> {
-        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<AppResponse> {
         self.execute(addr, &ExecuteMsg::Collect { id: self.id }, None)
     }
+
+    fn query_wallet_count(&self) -> StdResult<(u32, Vec<u32>)> {
+        let resp = self.query::<MarketResp>(&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();