pr review fixes

DA0-DA0 · Oct 31, 2024 · 195da14 · 195da14
1 parent ed90960
commit 195da14
Show file tree

Hide file tree

Showing 9 changed files with 191 additions and 112 deletions.
diff --git a/contracts/voting/dao-voting-cw721-staked/schema/dao-voting-cw721-staked.json b/contracts/voting/dao-voting-cw721-staked/schema/dao-voting-cw721-staked.json
@@ -265,23 +265,20 @@
         "additionalProperties": false
       },
       {
-        "description": "Claim NFTs that have been unstaked for the specified duration. If none are provided, it attempts to claim all legacy claims. If token IDs are provided, only those are claimed. If an empty vector is provided, it attempts to claim all non-legacy claims.",
+        "description": "Claim NFTs that have been unstaked for the specified duration.",
         "type": "object",
         "required": [
           "claim_nfts"
         ],
         "properties": {
           "claim_nfts": {
             "type": "object",
+            "required": [
+              "type"
+            ],
             "properties": {
-              "token_ids": {
-                "type": [
-                  "array",
-                  "null"
-                ],
-                "items": {
-                  "type": "string"
-                }
+              "type": {
+                "$ref": "#/definitions/ClaimType"
               }
             },
             "additionalProperties": false
@@ -440,6 +437,40 @@
         "description": "Binary is a wrapper around Vec<u8> to add base64 de/serialization with serde. It also adds some helper methods to help encode inline.\n\nThis is only needed as serde-json-{core,wasm} has a horrible encoding for Vec<u8>. See also <https://github.com/CosmWasm/cosmwasm/blob/main/docs/MESSAGE_TYPES.md>.",
         "type": "string"
       },
+      "ClaimType": {
+        "oneOf": [
+          {
+            "description": "Claims all legacy claims.",
+            "type": "string",
+            "enum": [
+              "legacy"
+            ]
+          },
+          {
+            "description": "Claims all non-legacy claims.",
+            "type": "string",
+            "enum": [
+              "all"
+            ]
+          },
+          {
+            "description": "Claims specific non-legacy NFTs.",
+            "type": "object",
+            "required": [
+              "specific"
+            ],
+            "properties": {
+              "specific": {
+                "type": "array",
+                "items": {
+                  "type": "string"
+                }
+              }
+            },
+            "additionalProperties": false
+          }
+        ]
+      },
       "Cw721ReceiveMsg": {
         "description": "Cw721ReceiveMsg should be de/serialized under `Receive()` variant in a ExecuteMsg",
         "type": "object",

diff --git a/contracts/voting/dao-voting-cw721-staked/src/contract.rs b/contracts/voting/dao-voting-cw721-staked/src/contract.rs
@@ -17,7 +17,7 @@ use dao_voting::threshold::{
     ActiveThresholdResponse,
 };
 
-use crate::msg::{ExecuteMsg, InstantiateMsg, MigrateMsg, NftContract, QueryMsg};
+use crate::msg::{ClaimType, ExecuteMsg, InstantiateMsg, MigrateMsg, NftContract, QueryMsg};
 use crate::state::{
     register_staked_nft, register_unstaked_nfts, Config, ACTIVE_THRESHOLD, CONFIG, DAO, HOOKS,
     INITIAL_NFTS, LEGACY_NFT_CLAIMS, NFT_BALANCES, NFT_CLAIMS, STAKED_NFTS_PER_OWNER,
@@ -205,7 +205,7 @@ pub fn execute(
     match msg {
         ExecuteMsg::ReceiveNft(msg) => execute_stake(deps, env, info, msg),
         ExecuteMsg::Unstake { token_ids } => execute_unstake(deps, env, info, token_ids),
-        ExecuteMsg::ClaimNfts { token_ids } => execute_claim_nfts(deps, env, info, token_ids),
+        ExecuteMsg::ClaimNfts { r#type } => execute_claim_nfts(deps, env, info, r#type),
         ExecuteMsg::UpdateConfig { duration } => execute_update_config(info, deps, duration),
         ExecuteMsg::AddHook { addr } => execute_add_hook(deps, info, addr),
         ExecuteMsg::RemoveHook { addr } => execute_remove_hook(deps, info, addr),
@@ -337,47 +337,43 @@ pub fn execute_claim_nfts(
     deps: DepsMut,
     env: Env,
     info: MessageInfo,
-    token_ids: Option<Vec<String>>,
+    claim_type: ClaimType,
 ) -> Result<Response, ContractError> {
-    let token_ids = match token_ids {
+    let token_ids = match claim_type {
         // attempt to claim all legacy NFTs
-        None => {
-            // claim all legacy NFTs
-            let legacy_nfts =
-                LEGACY_NFT_CLAIMS.claim_nfts(deps.storage, &info.sender, &env.block)?;
+        ClaimType::Legacy => {
+            LEGACY_NFT_CLAIMS.claim_nfts(deps.storage, &info.sender, &env.block)?
+        }
+        // attempt to claim all non-legacy NFTs
+        ClaimType::All => {
+            let token_ids = NFT_CLAIMS
+                .query_claims(deps.as_ref(), &info.sender, None, None)?
+                .nft_claims
+                .into_iter()
+                .filter(|nft| nft.release_at.is_expired(&env.block))
+                .map(|nft| nft.token_id)
+                .collect::<Vec<_>>();
 
-            if legacy_nfts.is_empty() {
-                return Err(ContractError::NothingToClaim {});
+            if !token_ids.is_empty() {
+                NFT_CLAIMS.claim_nfts(deps.storage, &info.sender, &token_ids, &env.block)?;
             }
 
-            legacy_nfts
+            token_ids
         }
-        // attempt to claim non-legacy NFTs
-        Some(token_ids) => {
-            let token_ids = if token_ids.is_empty() {
-                // query all NFT claims if none are provided
-                NFT_CLAIMS
-                    .query_claims(deps.as_ref(), &info.sender, None, None)?
-                    .nft_claims
-                    .into_iter()
-                    .filter(|nft| nft.release_at.is_expired(&env.block))
-                    .map(|nft| nft.token_id)
-                    .collect::<Vec<_>>()
-            } else {
-                // use provided NFTs if any
-                token_ids
-            };
-
-            if token_ids.is_empty() {
-                return Err(ContractError::NothingToClaim {});
+        // attempt to claim specific non-legacy NFTs
+        ClaimType::Specific(token_ids) => {
+            if !token_ids.is_empty() {
+                NFT_CLAIMS.claim_nfts(deps.storage, &info.sender, &token_ids, &env.block)?;
             }
 
-            NFT_CLAIMS.claim_nfts(deps.storage, &info.sender, &token_ids, &env.block)?;
-
             token_ids
         }
     };
 
+    if token_ids.is_empty() {
+        return Err(ContractError::NothingToClaim {});
+    }
+
     let config = CONFIG.load(deps.storage)?;
 
     let msgs = token_ids

diff --git a/contracts/voting/dao-voting-cw721-staked/src/msg.rs b/contracts/voting/dao-voting-cw721-staked/src/msg.rs
@@ -54,11 +54,8 @@ pub enum ExecuteMsg {
     /// sender. token_ids must have unique values and have non-zero
     /// length.
     Unstake { token_ids: Vec<String> },
-    /// Claim NFTs that have been unstaked for the specified duration. If none
-    /// are provided, it attempts to claim all legacy claims. If token IDs are
-    /// provided, only those are claimed. If an empty vector is provided, it
-    /// attempts to claim all non-legacy claims.
-    ClaimNfts { token_ids: Option<Vec<String>> },
+    /// Claim NFTs that have been unstaked for the specified duration.
+    ClaimNfts { r#type: ClaimType },
     /// Updates the contract configuration, namely unstaking duration.
     /// Only callable by the DAO that initialized this voting contract.
     UpdateConfig { duration: Option<Duration> },
@@ -75,6 +72,16 @@ pub enum ExecuteMsg {
     },
 }
 
+#[cw_serde]
+pub enum ClaimType {
+    /// Claims all legacy claims.
+    Legacy,
+    /// Claims all non-legacy claims.
+    All,
+    /// Claims specific non-legacy NFTs.
+    Specific(Vec<String>),
+}
+
 #[active_query]
 #[voting_module_query]
 #[cw_serde]

diff --git a/contracts/voting/dao-voting-cw721-staked/src/testing/execute.rs b/contracts/voting/dao-voting-cw721-staked/src/testing/execute.rs
@@ -5,7 +5,7 @@ use cw_multi_test::{App, AppResponse, Executor};
 use anyhow::Result as AnyResult;
 use cw_utils::Duration;
 
-use crate::msg::ExecuteMsg;
+use crate::msg::{ClaimType, ExecuteMsg};
 
 // Shorthand for an unchecked address.
 macro_rules! addr {
@@ -111,7 +111,7 @@ pub fn claim_nfts(app: &mut App, module: &Addr, sender: &str) -> AnyResult<AppRe
         addr!(sender),
         module.clone(),
         &ExecuteMsg::ClaimNfts {
-            token_ids: Some(vec![]),
+            r#type: ClaimType::All,
         },
         &[],
     )
@@ -127,7 +127,7 @@ pub fn claim_specific_nfts(
         addr!(sender),
         module.clone(),
         &ExecuteMsg::ClaimNfts {
-            token_ids: Some(token_ids.to_vec()),
+            r#type: ClaimType::Specific(token_ids.to_vec()),
         },
         &[],
     )
@@ -137,7 +137,9 @@ pub fn claim_legacy_nfts(app: &mut App, module: &Addr, sender: &str) -> AnyResul
     app.execute_contract(
         addr!(sender),
         module.clone(),
-        &ExecuteMsg::ClaimNfts { token_ids: None },
+        &ExecuteMsg::ClaimNfts {
+            r#type: ClaimType::Legacy,
+        },
         &[],
     )
 }

diff --git a/contracts/voting/dao-voting-onft-staked/schema/dao-voting-onft-staked.json b/contracts/voting/dao-voting-onft-staked/schema/dao-voting-onft-staked.json
@@ -274,23 +274,20 @@
         "additionalProperties": false
       },
       {
-        "description": "Claim NFTs that have been unstaked for the specified duration. If none are provided, it attempts to claim all legacy claims. If token IDs are provided, only those are claimed. If an empty vector is provided, it attempts to claim all non-legacy claims.",
+        "description": "Claim NFTs that have been unstaked for the specified duration.",
         "type": "object",
         "required": [
           "claim_nfts"
         ],
         "properties": {
           "claim_nfts": {
             "type": "object",
+            "required": [
+              "type"
+            ],
             "properties": {
-              "token_ids": {
-                "type": [
-                  "array",
-                  "null"
-                ],
-                "items": {
-                  "type": "string"
-                }
+              "type": {
+                "$ref": "#/definitions/ClaimType"
               }
             },
             "additionalProperties": false
@@ -445,6 +442,40 @@
           }
         ]
       },
+      "ClaimType": {
+        "oneOf": [
+          {
+            "description": "Claims all legacy claims.",
+            "type": "string",
+            "enum": [
+              "legacy"
+            ]
+          },
+          {
+            "description": "Claims all non-legacy claims.",
+            "type": "string",
+            "enum": [
+              "all"
+            ]
+          },
+          {
+            "description": "Claims specific non-legacy NFTs.",
+            "type": "object",
+            "required": [
+              "specific"
+            ],
+            "properties": {
+              "specific": {
+                "type": "array",
+                "items": {
+                  "type": "string"
+                }
+              }
+            },
+            "additionalProperties": false
+          }
+        ]
+      },
       "Decimal": {
         "description": "A fixed-point decimal value with 18 fractional digits, i.e. Decimal(1_000_000_000_000_000_000) == 1.0\n\nThe greatest possible value that can be represented is 340282366920938463463.374607431768211455 (which is (2^128 - 1) / 10^18)",
         "type": "string"

diff --git a/contracts/voting/dao-voting-onft-staked/src/contract.rs b/contracts/voting/dao-voting-onft-staked/src/contract.rs
@@ -15,7 +15,7 @@ use dao_voting::threshold::{
     ActiveThresholdResponse,
 };
 
-use crate::msg::{ExecuteMsg, InstantiateMsg, MigrateMsg, OnftCollection, QueryMsg};
+use crate::msg::{ClaimType, ExecuteMsg, InstantiateMsg, MigrateMsg, OnftCollection, QueryMsg};
 use crate::omniflix::{get_onft_transfer_msg, query_onft_owner, query_onft_supply};
 use crate::state::{
     register_staked_nfts, register_unstaked_nfts, Config, ACTIVE_THRESHOLD, CONFIG, DAO, HOOKS,
@@ -98,7 +98,7 @@ pub fn execute(
             recipient,
         } => execute_cancel_stake(deps, env, info, token_ids, recipient),
         ExecuteMsg::Unstake { token_ids } => execute_unstake(deps, env, info, token_ids),
-        ExecuteMsg::ClaimNfts { token_ids } => execute_claim_nfts(deps, env, info, token_ids),
+        ExecuteMsg::ClaimNfts { r#type } => execute_claim_nfts(deps, env, info, r#type),
         ExecuteMsg::UpdateConfig { duration } => execute_update_config(info, deps, duration),
         ExecuteMsg::AddHook { addr } => execute_add_hook(deps, info, addr),
         ExecuteMsg::RemoveHook { addr } => execute_remove_hook(deps, info, addr),
@@ -394,47 +394,43 @@ pub fn execute_claim_nfts(
     deps: DepsMut,
     env: Env,
     info: MessageInfo,
-    token_ids: Option<Vec<String>>,
+    claim_type: ClaimType,
 ) -> Result<Response, ContractError> {
-    let token_ids = match token_ids {
+    let token_ids = match claim_type {
         // attempt to claim all legacy NFTs
-        None => {
-            // claim all legacy NFTs
-            let legacy_nfts =
-                LEGACY_NFT_CLAIMS.claim_nfts(deps.storage, &info.sender, &env.block)?;
+        ClaimType::Legacy => {
+            LEGACY_NFT_CLAIMS.claim_nfts(deps.storage, &info.sender, &env.block)?
+        }
+        // attempt to claim all non-legacy NFTs
+        ClaimType::All => {
+            let token_ids = NFT_CLAIMS
+                .query_claims(deps.as_ref(), &info.sender, None, None)?
+                .nft_claims
+                .into_iter()
+                .filter(|nft| nft.release_at.is_expired(&env.block))
+                .map(|nft| nft.token_id)
+                .collect::<Vec<_>>();
 
-            if legacy_nfts.is_empty() {
-                return Err(ContractError::NothingToClaim {});
+            if !token_ids.is_empty() {
+                NFT_CLAIMS.claim_nfts(deps.storage, &info.sender, &token_ids, &env.block)?;
             }
 
-            legacy_nfts
+            token_ids
         }
-        // attempt to claim non-legacy NFTs
-        Some(token_ids) => {
-            let token_ids = if token_ids.is_empty() {
-                // query all NFT claims if none are provided
-                NFT_CLAIMS
-                    .query_claims(deps.as_ref(), &info.sender, None, None)?
-                    .nft_claims
-                    .into_iter()
-                    .filter(|nft| nft.release_at.is_expired(&env.block))
-                    .map(|nft| nft.token_id)
-                    .collect::<Vec<_>>()
-            } else {
-                // use provided NFTs if any
-                token_ids
-            };
-
-            if token_ids.is_empty() {
-                return Err(ContractError::NothingToClaim {});
+        // attempt to claim specific non-legacy NFTs
+        ClaimType::Specific(token_ids) => {
+            if !token_ids.is_empty() {
+                NFT_CLAIMS.claim_nfts(deps.storage, &info.sender, &token_ids, &env.block)?;
             }
 
-            NFT_CLAIMS.claim_nfts(deps.storage, &info.sender, &token_ids, &env.block)?;
-
             token_ids
         }
     };
 
+    if token_ids.is_empty() {
+        return Err(ContractError::NothingToClaim {});
+    }
+
     let config = CONFIG.load(deps.storage)?;
 
     let msgs = token_ids