Cw4 group list members by weight

contracts/cw4-group/schema/query_msg.json

@@ -57,6 +57,47 @@
       },
       "additionalProperties": false
     },
+    {
+      "description": "Returns MembersListResponse, sorted by weight descending",
+      "type": "object",
+      "required": [
+        "list_members_by_weight"
+      ],
+      "properties": {
+        "list_members_by_weight": {
+          "type": "object",
+          "properties": {
+            "limit": {
+              "type": [
+                "integer",
+                "null"
+              ],
+              "format": "uint32",
+              "minimum": 0.0
+            },
+            "start_after": {
+              "type": [
+                "array",
+                "null"
+              ],
+              "items": [
+                {
+                  "type": "integer",
+                  "format": "uint64",
+                  "minimum": 0.0
+                },
+                {
+                  "type": "string"
+                }
+              ],
+              "maxItems": 2,
+              "minItems": 2
+            }
+          }
+        }
+      },
+      "additionalProperties": false
+    },
     {
       "description": "Returns MemberResponse",
       "type": "object",

contracts/cw4-group/src/contract.rs

@@ -9,11 +9,11 @@ use cw4::{
     Member, MemberChangedHookMsg, MemberDiff, MemberListResponse, MemberResponse,
     TotalWeightResponse,
 };
-use cw_storage_plus::Bound;
+use cw_storage_plus::{Bound, PrimaryKey, U64Key};
 
 use crate::error::ContractError;
 use crate::msg::{ExecuteMsg, InstantiateMsg, QueryMsg};
-use crate::state::{ADMIN, HOOKS, MEMBERS, TOTAL};
+use crate::state::{members, ADMIN, HOOKS, TOTAL};
 
 // version info for migration info
 const CONTRACT_NAME: &str = "crates.io:cw4-group";
@@ -38,7 +38,7 @@ pub fn instantiate(
 pub fn create(
     mut deps: DepsMut,
     admin: Option<String>,
-    members: Vec<Member>,
+    members_list: Vec<Member>,
     height: u64,
 ) -> Result<(), ContractError> {
     let admin_addr = admin
@@ -47,10 +47,10 @@ pub fn create(
     ADMIN.set(deps.branch(), admin_addr)?;
 
     let mut total = 0u64;
-    for member in members.into_iter() {
+    for member in members_list.into_iter() {
         total += member.weight;
         let member_addr = deps.api.addr_validate(&member.addr)?;
-        MEMBERS.save(deps.storage, &member_addr, &member.weight, height)?;
+        members().save(deps.storage, &member_addr, &member.weight, height)?;
     }
     TOTAL.save(deps.storage, &total)?;
 
@@ -126,7 +126,7 @@ pub fn update_members(
     // add all new members and update total
     for add in to_add.into_iter() {
         let add_addr = deps.api.addr_validate(&add.addr)?;
-        MEMBERS.update(deps.storage, &add_addr, height, |old| -> StdResult<_> {
+        members().update(deps.storage, &add_addr, height, |old| -> StdResult<_> {
             total -= old.unwrap_or_default();
             total += add.weight;
             diffs.push(MemberDiff::new(add.addr, old, Some(add.weight)));
@@ -136,12 +136,12 @@ pub fn update_members(
 
     for remove in to_remove.into_iter() {
         let remove_addr = deps.api.addr_validate(&remove)?;
-        let old = MEMBERS.may_load(deps.storage, &remove_addr)?;
+        let old = members().may_load(deps.storage, &remove_addr)?;
         // Only process this if they were actually in the list before
         if let Some(weight) = old {
             diffs.push(MemberDiff::new(remove, Some(weight), None));
             total -= weight;
-            MEMBERS.remove(deps.storage, &remove_addr, height)?;
+            members().remove(deps.storage, &remove_addr, height)?;
         }
     }
 
@@ -159,6 +159,9 @@ pub fn query(deps: Deps, _env: Env, msg: QueryMsg) -> StdResult<Binary> {
         QueryMsg::ListMembers { start_after, limit } => {
             to_binary(&list_members(deps, start_after, limit)?)
         }
+        QueryMsg::ListMembersByWeight { start_after, limit } => {
+            to_binary(&list_members_by_weight(deps, start_after, limit)?)
+        }
         QueryMsg::TotalWeight {} => to_binary(&query_total_weight(deps)?),
         QueryMsg::Admin {} => to_binary(&ADMIN.query_admin(deps)?),
         QueryMsg::Hooks {} => to_binary(&HOOKS.query_hooks(deps)?),
@@ -173,8 +176,8 @@ fn query_total_weight(deps: Deps) -> StdResult<TotalWeightResponse> {
 fn query_member(deps: Deps, addr: String, height: Option<u64>) -> StdResult<MemberResponse> {
     let addr = deps.api.addr_validate(&addr)?;
     let weight = match height {
-        Some(h) => MEMBERS.may_load_at_height(deps.storage, &addr, h),
-        None => MEMBERS.may_load(deps.storage, &addr),
+        Some(h) => members().may_load_at_height(deps.storage, &addr, h),
+        None => members().may_load(deps.storage, &addr),
     }?;
     Ok(MemberResponse { weight })
 }
@@ -190,9 +193,9 @@ fn list_members(
 ) -> StdResult<MemberListResponse> {
     let limit = limit.unwrap_or(DEFAULT_LIMIT).min(MAX_LIMIT) as usize;
     let addr = maybe_addr(deps.api, start_after)?;
-    let start = addr.map(|addr| Bound::exclusive(addr.to_string()));
+    let start = addr.map(|addr| Bound::exclusive(addr.as_ref()));
 
-    let members: StdResult<Vec<_>> = MEMBERS
+    let members: StdResult<Vec<_>> = members()
         .range(deps.storage, start, None, Order::Ascending)
         .take(limit)
         .map(|item| {
@@ -207,6 +210,32 @@ fn list_members(
     Ok(MemberListResponse { members: members? })
 }
 
+fn list_members_by_weight(
+    deps: Deps,
+    start_after: Option<(u64, String)>,
+    limit: Option<u32>,
+) -> StdResult<MemberListResponse> {
+    let limit = limit.unwrap_or(DEFAULT_LIMIT).min(MAX_LIMIT) as usize;
+    let start =
+        start_after.map(|(w, a)| Bound::exclusive((U64Key::from(w), a.as_str()).joined_key()));
+
+    let members: StdResult<Vec<_>> = members()
+        .idx
+        .weight
+        .range(deps.storage, None, start, Order::Descending)
+        .take(limit)
+        .map(|item| {
+            let (key, weight) = item?;
+            Ok(Member {
+                addr: String::from_utf8(key)?,
+                weight,
+            })
+        })
+        .collect();
+
+    Ok(MemberListResponse { members: members? })
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -265,9 +294,117 @@ mod tests {
         let member3 = query_member(deps.as_ref(), USER3.into(), None).unwrap();
         assert_eq!(member3.weight, None);
 
-        let members = list_members(deps.as_ref(), None, None).unwrap();
-        assert_eq!(members.members.len(), 2);
-        // TODO: assert the set is proper
+        let members = list_members(deps.as_ref(), None, None).unwrap().members;
+        assert_eq!(members.len(), 2);
+        // Assert the set is proper
+        assert_eq!(
+            members,
+            vec![
+                Member {
+                    addr: USER2.into(),
+                    weight: 6
+                },
+                Member {
+                    addr: USER1.into(),
+                    weight: 11
+                },
+            ]
+        );
+
+        // Test pagination / limits
+        let members = list_members(deps.as_ref(), None, Some(1)).unwrap().members;
+        assert_eq!(members.len(), 1);
+        // Assert the set is proper
+        assert_eq!(
+            members,
+            vec![Member {
+                addr: USER2.into(),
+                weight: 6
+            },]
+        );
+
+        // Next page
+        let start_after = Some(members[0].addr.clone());
+        let members = list_members(deps.as_ref(), start_after, Some(1))
+            .unwrap()
+            .members;
+        assert_eq!(members.len(), 1);
+        // Assert the set is proper
+        assert_eq!(
+            members,
+            vec![Member {
+                addr: USER1.into(),
+                weight: 11
+            },]
+        );
+
+        // Assert there's no more
+        let start_after = Some(members[0].addr.clone());
+        let members = list_members(deps.as_ref(), start_after, Some(1))
+            .unwrap()
+            .members;
+        assert_eq!(members.len(), 0);
+    }
+
+    #[test]
+    fn try_list_members_by_weight() {
+        let mut deps = mock_dependencies(&[]);
+        do_instantiate(deps.as_mut());
+
+        let members = list_members_by_weight(deps.as_ref(), None, None)
+            .unwrap()
+            .members;
+        assert_eq!(members.len(), 2);
+        // Assert the set is sorted by (descending) weight
+        assert_eq!(
+            members,
+            vec![
+                Member {
+                    addr: USER1.into(),
+                    weight: 11
+                },
+                Member {
+                    addr: USER2.into(),
+                    weight: 6
+                }
+            ]
+        );
+
+        // Test pagination / limits
+        let members = list_members_by_weight(deps.as_ref(), None, Some(1))
+            .unwrap()
+            .members;
+        assert_eq!(members.len(), 1);
+        // Assert the set is proper
+        assert_eq!(
+            members,
+            vec![Member {
+                addr: USER1.into(),
+                weight: 11
+            },]
+        );
+
+        // Next page
+        let start_after = Some((members[0].weight, members[0].addr.clone()));
+        let members = list_members_by_weight(deps.as_ref(), start_after, None)
+            .unwrap()
+            .members;
+        assert_eq!(members.len(), 1);
+        // Assert the set is proper
+        assert_eq!(
+            members,
+            vec![Member {
+                addr: USER2.into(),
+                weight: 6
+            },]
+        );
+
+        // Assert there's no more
+        let start_after = Some((members[0].weight, members[0].addr.clone()));
+        let members = list_members_by_weight(deps.as_ref(), start_after, Some(1))
+            .unwrap()
+            .members;
+        assert_eq!(members.len(), 0);
     }
 
     fn assert_users<S: Storage, A: Api, Q: Querier>(

contracts/cw4-group/src/msg.rs

@@ -41,6 +41,11 @@ pub enum QueryMsg {
         start_after: Option<String>,
         limit: Option<u32>,
     },
+    /// Returns MembersListResponse, sorted by weight descending
+    ListMembersByWeight {
+        start_after: Option<(u64, String)>,
+        limit: Option<u32>,
+    },
     /// Returns MemberResponse
     Member {
         addr: String,

contracts/cw4-group/src/state.rs

@@ -1,16 +1,40 @@
 use cosmwasm_std::Addr;
 use cw4::TOTAL_KEY;
 use cw_controllers::{Admin, Hooks};
-use cw_storage_plus::{Item, SnapshotMap, Strategy};
+use cw_storage_plus::{
+    Index, IndexList, IndexedSnapshotMap, Item, MultiIndex, PkOwned, Strategy, U64Key,
+};
 
 pub const ADMIN: Admin = Admin::new("admin");
 pub const HOOKS: Hooks = Hooks::new("cw4-hooks");
 
 pub const TOTAL: Item<u64> = Item::new(TOTAL_KEY);
 
-pub const MEMBERS: SnapshotMap<&Addr, u64> = SnapshotMap::new(
-    cw4::MEMBERS_KEY,
-    cw4::MEMBERS_CHECKPOINTS,
-    cw4::MEMBERS_CHANGELOG,
-    Strategy::EveryBlock,
-);
+pub struct MemberIndexes<'a> {
+    // pk goes to second tuple element
+    pub weight: MultiIndex<'a, (U64Key, PkOwned), u64>,
+}
+
+impl<'a> IndexList<u64> for MemberIndexes<'a> {
+    fn get_indexes(&'_ self) -> Box<dyn Iterator<Item = &'_ dyn Index<u64>> + '_> {
+        let v: Vec<&dyn Index<u64>> = vec![&self.weight];
+        Box::new(v.into_iter())
+    }
+}
+
+pub fn members<'a>() -> IndexedSnapshotMap<'a, &'a Addr, u64, MemberIndexes<'a>> {
+    let indexes = MemberIndexes {
+        weight: MultiIndex::new(
+            |&w, k| (U64Key::new(w), PkOwned(k)),
+            cw4::MEMBERS_KEY,
+            "members__weight",
+        ),
+    };
+    IndexedSnapshotMap::new(
+        cw4::MEMBERS_KEY,
+        cw4::MEMBERS_CHECKPOINTS,
+        cw4::MEMBERS_CHANGELOG,
+        Strategy::EveryBlock,
+        indexes,
+    )
+}