Merge pull request #1 from mymiracle0118/test

complete main structure and logic

Cargo.toml

@@ -42,6 +42,7 @@ cosmwasm-std = { version = "1.5.0", features = [
 ] }
 cw-storage-plus = "1.1.0"
 cw2 = "1.1.1"
+cw721 = "0.18.0"
 schemars = "0.8.15"
 serde = { version = "1.0.189", default-features = false, features = ["derive"] }
 thiserror = { version = "1.0.49" }

src/contract.rs

@@ -1,11 +1,12 @@
 #[cfg(not(feature = "library"))]
 use cosmwasm_std::entry_point;
-use cosmwasm_std::{to_json_binary, Binary, Deps, DepsMut, Env, MessageInfo, Response, StdResult};
+use cosmwasm_std::{to_json_binary, Binary, Deps, DepsMut, Env, MessageInfo, Response, StdResult, Addr, CosmosMsg, WasmMsg, SubMsg};
 use cw2::set_contract_version;
+use cw721::Cw721ExecuteMsg;
 
 use crate::error::ContractError;
-use crate::msg::{ExecuteMsg, GetCountResponse, InstantiateMsg, QueryMsg};
-use crate::state::{State, STATE};
+use crate::msg::{ExecuteMsg, GetCountResponse, InstantiateMsg, QueryMsg, AllNftsResponse, NftContractAddrResponse};
+use crate::state::{State, STATE, NFTS, NFT_CONTRACT_ADDR};
 
 // version info for migration info
 const CONTRACT_NAME: &str = "crates.io:jarvis-airdrop";
@@ -34,17 +35,22 @@ pub fn instantiate(
 #[cfg_attr(not(feature = "library"), entry_point)]
 pub fn execute(
     deps: DepsMut,
-    _env: Env,
+    env: Env,
     info: MessageInfo,
     msg: ExecuteMsg,
 ) -> Result<Response, ContractError> {
     match msg {
         ExecuteMsg::Increment {} => execute::increment(deps),
         ExecuteMsg::Reset { count } => execute::reset(deps, info, count),
+        ExecuteMsg::SetNftContractAddr { addr } => execute::set_nft_contract_addr(deps, env, info, addr),
+        ExecuteMsg::ReceiveNft { sender, token_id, msg } => execute::receive_nft(deps, env, info, token_id),
+        ExecuteMsg::SendNfts { allocations } => execute::send_nfts(deps, env, info, allocations), 
     }
 }
 
 pub mod execute {
+    use std::ops::Sub;
+
     use super::*;
 
     pub fn increment(deps: DepsMut) -> Result<Response, ContractError> {
@@ -66,12 +72,76 @@ pub mod execute {
         })?;
         Ok(Response::new().add_attribute("action", "reset"))
     }
+
+    pub fn set_nft_contract_addr(
+        deps: DepsMut,
+        _env: Env,
+        info: MessageInfo,
+        addr: String,
+    ) -> Result<Response, ContractError> {
+        // Optionally, add authorization checks here to ensure only specific addresses can update this
+        let nft_contract_addr = deps.api.addr_validate(&addr)?;
+        NFT_CONTRACT_ADDR.save(deps.storage, &nft_contract_addr)?;
+        Ok(Response::new().add_attribute("action", "set_nft_contract_addr").add_attribute("address", addr))
+    }
+
+    pub fn receive_nft(
+        deps: DepsMut,
+        _env: Env,
+        _info: MessageInfo,
+        token_id: String,
+    ) -> Result<Response, ContractError> {
+        let mut nfts = NFTS.load(deps.storage).unwrap_or_default();
+        nfts.push(token_id);
+        NFTS.save(deps.storage, &nfts)?;
+
+        Ok(Response::new().add_attribute("action", "receive_nft"))
+    }
+
+
+    pub fn send_nfts(
+        deps: DepsMut,
+        _env: Env,
+        info: MessageInfo,
+        allocations: Vec<(Addr, u32)>,
+    ) -> Result<Response, ContractError> {
+        let nft_contract_addr = NFT_CONTRACT_ADDR.load(deps.storage)?;
+        let mut nfts = NFTS.load(deps.storage)?;
+        let mut response = Response::new().add_attribute("action", "send_nfts");
+
+        for (recipient, amount) in allocations {
+            for _ in 0..amount {
+                if let Some(token_id) = nfts.pop() {
+                    // Create a transfer message for the cw721 NFT
+                    let transfer_msg = Cw721ExecuteMsg::TransferNft {
+                        recipient: recipient.to_string(),
+                        token_id: token_id,
+                    };
+
+                    let msg = CosmosMsg::Wasm(WasmMsg::Execute {
+                        contract_addr: nft_contract_addr.clone().to_string(),
+                        msg: to_json_binary(&transfer_msg)?,
+                        funds: vec![],
+                    });
+
+                    response.messages.push(SubMsg::new(msg));
+                } else {
+                    return Err(ContractError::InsufficientNFTs {});
+                }
+            }
+        }
+
+        NFTS.save(deps.storage, &nfts)?;
+        Ok(response)
+    }
 }
 
 #[cfg_attr(not(feature = "library"), entry_point)]
 pub fn query(deps: Deps, _env: Env, msg: QueryMsg) -> StdResult<Binary> {
     match msg {
         QueryMsg::GetCount {} => to_json_binary(&query::count(deps)?),
+        QueryMsg::GetAllNfts {  } => to_json_binary(&query::all_nfts(deps)?),
+        QueryMsg::GetNftContractAddr {  } => to_json_binary(&query::nft_contract_addr(deps)?),
     }
 }
 
@@ -82,6 +152,16 @@ pub mod query {
         let state = STATE.load(deps.storage)?;
         Ok(GetCountResponse { count: state.count })
     }
+
+    pub fn all_nfts(deps: Deps) -> StdResult<AllNftsResponse> {
+        let nfts = NFTS.load(deps.storage)?;
+        Ok(AllNftsResponse { nfts })
+    }
+
+    pub fn nft_contract_addr(deps: Deps) -> StdResult<NftContractAddrResponse> {
+        let nft_contract_addr = NFT_CONTRACT_ADDR.load(deps.storage)?;
+        Ok(NftContractAddrResponse { nft_contract_addr })
+    }
 }
 
 #[cfg(test)]

src/error.rs

@@ -8,6 +8,9 @@ pub enum ContractError {
 
     #[error("Unauthorized")]
     Unauthorized {},
+
+    #[error("Insufficient NFTs")]
+    InsufficientNFTs {},
     // Add any other custom errors you like here.
     // Look at https://docs.rs/thiserror/1.0.21/thiserror/ for details.
 }

src/msg.rs

@@ -1,4 +1,5 @@
 use cosmwasm_schema::{cw_serde, QueryResponses};
+use cosmwasm_std::{Addr, Binary};
 
 #[cw_serde]
 pub struct InstantiateMsg {
@@ -9,6 +10,15 @@ pub struct InstantiateMsg {
 pub enum ExecuteMsg {
     Increment {},
     Reset { count: i32 },
+    ReceiveNft {
+        sender: String,
+        token_id: String,
+        msg: Binary,
+    },
+    SendNfts {
+        allocations: Vec<(Addr, u32)>, // Each tuple contains an address and the number of NFTs to send
+    },
+    SetNftContractAddr { addr: String },
 }
 
 #[cw_serde]
@@ -17,10 +27,26 @@ pub enum QueryMsg {
     // GetCount returns the current count as a json-encoded number
     #[returns(GetCountResponse)]
     GetCount {},
+
+    #[returns(AllNftsResponse)]
+    GetAllNfts {},
+
+    #[returns(NftContractAddrResponse)]
+    GetNftContractAddr {},
 }
 
 // We define a custom struct for each query response
 #[cw_serde]
 pub struct GetCountResponse {
-    pub count: i32,
+    pub count: i32, 
+}
+
+#[cw_serde]
+pub struct AllNftsResponse {
+    pub nfts: Vec<String>,
 }
+
+#[cw_serde]
+pub struct NftContractAddrResponse {
+    pub nft_contract_addr: Addr,
+}

src/state.rs

@@ -11,3 +11,5 @@ pub struct State {
 }
 
 pub const STATE: Item<State> = Item::new("state");
+pub const NFTS: Item<Vec<String>> = Item::new("nfts");
+pub const NFT_CONTRACT_ADDR: Item<Addr> = Item::new("nft_contract_addr");