From f6be02f5c5d22c6d330514f0be04bbaffb9c7541 Mon Sep 17 00:00:00 2001 From: nk_ysg Date: Wed, 8 Jan 2025 09:03:23 +0800 Subject: [PATCH] update merkledistributor contract --- .../modules/MerkleDistributor.move | 105 +++++++++--------- .../src/merkle_distributor_test.rs | 14 +-- vm/framework/move-stdlib/src/lib.rs | 9 ++ 3 files changed, 69 insertions(+), 59 deletions(-) diff --git a/contrib-contracts/modules/MerkleDistributor.move b/contrib-contracts/modules/MerkleDistributor.move index e7d61ad86d..f49ff44ccf 100644 --- a/contrib-contracts/modules/MerkleDistributor.move +++ b/contrib-contracts/modules/MerkleDistributor.move @@ -1,27 +1,29 @@ address StarcoinAssociation { module MerkleDistributorScripts { use StarcoinAssociation::MerkleDistributor; - use StarcoinFramework::Account; - public(script) fun create(signer: signer, merkle_root: vector, token_amounts: u128, leaves: u64) { - let tokens = Account::withdraw(&signer, token_amounts); - MerkleDistributor::create(&signer, merkle_root, tokens, leaves); + use starcoin_framework::coin; + use starcoin_std::signer; + public entry fun create(signer: signer, merkle_root: vector, coin_amounts: u64, leaves: u64) { + let coins = coin::withdraw(&signer, coin_amounts); + MerkleDistributor::create(&signer, merkle_root, coins, leaves); } - public(script) fun claim_for_address(distribution_address: address, index: u64, account: address, amount: u128, merkle_proof: vector>) { - MerkleDistributor::claim_for_address(distribution_address, index, account, amount, merkle_proof); + public entry fun claim_for_address(signer: signer, distribution_address: address, index: u64, account: address, amount: u64, merkle_proof: vector>) { + MerkleDistributor::claim_for_address(&signer, distribution_address, index, account, amount, merkle_proof); } - public(script) fun claim(signer: signer, distribution_address: address, index: u64, amount: u128, merkle_proof: vector>) { - let tokens = MerkleDistributor::claim(&signer, distribution_address, index, amount, merkle_proof); - Account::deposit_to_self(&signer, tokens); + public entry fun claim(signer: signer, distribution_address: address, index: u64, amount: u64, merkle_proof: vector>) { + let coins = MerkleDistributor::claim(&signer, distribution_address, index, amount, merkle_proof); + let account_addr = signer::address_of(&signer); + coin::deposit(account_addr, coins); } } module MerkleProof { - use StarcoinFramework::Hash; - use StarcoinFramework::Vector; - use StarcoinFramework::Compare; + use std::hash; + use std::vector; + use starcoin_std::comparator; - /// verify leaf node with hash of `leaf` with `proof` againest merkle `root`. + /// verify leaf node with hash of `leaf` with `proof` against merkle `root`. public fun verify(proof: &vector>, root: &vector, leaf: vector): bool { let computed_hash = leaf; let i = 0; @@ -29,12 +31,12 @@ module MerkleProof { while(i < proof_length) { let sibling = vector::borrow(proof, i); // computed_hash is left. - if (Compare::cmp_bytes(&computed_hash,sibling) < 2) { + if (!comparator::is_greater_than(&comparator::compare_u8_vector(computed_hash,*sibling))) { let concated = concat(computed_hash, *sibling); - computed_hash = Hash::sha3_256(concated); + computed_hash = hash::sha3_256(concated); } else { let concated = concat(*sibling, computed_hash); - computed_hash = Hash::sha3_256(concated); + computed_hash = hash::sha3_256(concated); }; @@ -53,57 +55,56 @@ module MerkleProof { module MerkleDistributor { - use StarcoinFramework::Token::{Token, Self}; use StarcoinAssociation::MerkleProof; - use StarcoinFramework::Vector; - use StarcoinFramework::BCS; - // use StarcoinFramework::BitOperators; - use StarcoinFramework::Account; - use StarcoinFramework::Signer; - use StarcoinFramework::Errors; - use StarcoinFramework::Hash; + use std::bcs; + use starcoin_framework::coin; + use std::error; + use std::hash; + use std::vector; + use starcoin_std::signer; + struct MerkleDistribution has key { merkle_root: vector, - tokens: Token, - claimed_bitmap: vector, + coins: coin::Coin, + claimed_bitmap: vector, } const INVALID_PROOF: u64 = 1; const ALREADY_CLAIMED: u64 = 2; /// Initialization. - public fun create(signer: &signer, merkle_root: vector, tokens: Token, leaves: u64) { - let bitmap_count = leaves / 128; - if (bitmap_count * 128 < leaves) { + public fun create(signer: &signer, merkle_root: vector, coins: coin::Coin, leaves: u64) { + let bitmap_count = leaves / 64; + if (bitmap_count * 64 < leaves) { bitmap_count = bitmap_count + 1; }; let claimed_bitmap = vector::empty(); let j = 0; while (j < bitmap_count) { - vector::push_back(&mut claimed_bitmap, 0u128); + vector::push_back(&mut claimed_bitmap, 0u64); j = j + 1; }; let distribution = MerkleDistribution{ merkle_root, - tokens, - claimed_bitmap + coins, + claimed_bitmap, }; move_to(signer, distribution); } /// claim for some address. - public fun claim_for_address(distribution_address: address, index: u64, account: address, amount: u128, merkle_proof: vector>) + public fun claim_for_address(signer: &signer, distribution_address: address, index: u64, account: address, amount: u64, merkle_proof: vector>) acquires MerkleDistribution { let distribution = borrow_global_mut>(distribution_address); - let claimed_tokens = internal_claim(distribution, index, account, amount, merkle_proof); - Account::deposit(account, claimed_tokens); + let claimed_coins = internal_claim(signer, distribution, index, account, amount, merkle_proof); + coin::deposit(account, claimed_coins); } /// claim by myself. - public fun claim(signer: &signer, distribution_address: address, index: u64, amount: u128, merkle_proof: vector>): Token + public fun claim(signer: &signer, distribution_address: address, index: u64, amount: u64, merkle_proof: vector>): coin::Coin acquires MerkleDistribution { let distribution = borrow_global_mut>(distribution_address); - internal_claim(distribution, index, Signer::address_of(signer), amount, merkle_proof) + internal_claim(signer, distribution, index, signer::address_of(signer), amount, merkle_proof) } /// Query whether `index` of `distribution_address` has already claimed. @@ -113,41 +114,41 @@ module MerkleDistributor { is_claimed_(distribution, index) } - fun internal_claim(distribution: &mut MerkleDistribution, index: u64, account: address, amount: u128, merkle_proof: vector>): Token { + fun internal_claim(signer: &signer, distribution: &mut MerkleDistribution, index: u64, account: address, amount: u64, merkle_proof: vector>): coin::Coin { let claimed = is_claimed_(distribution, index); - assert!(!claimed, Errors::custom(ALREADY_CLAIMED)); + assert!(!claimed, error::invalid_argument(ALREADY_CLAIMED)); let leaf_data = encode_leaf(&index, &account, &amount); - let verified = MerkleProof::verify(&merkle_proof, &distribution.merkle_root, Hash::sha3_256(leaf_data)); - assert!(verified, Errors::custom(INVALID_PROOF)); + let verified = MerkleProof::verify(&merkle_proof, &distribution.merkle_root, hash::sha3_256(leaf_data)); + assert!(verified, error::invalid_argument(INVALID_PROOF)); set_claimed_(distribution, index); - Token::withdraw(&mut distribution.tokens, amount) + coin::withdraw(signer, amount) } fun is_claimed_(distribution: &MerkleDistribution, index: u64): bool { - let claimed_word_index = index / 128; - let claimed_bit_index = ((index % 128) as u8); + let claimed_word_index = index / 64; + let claimed_bit_index = ((index % 64) as u8); let word = vector::borrow(&distribution.claimed_bitmap, claimed_word_index); - let mask = 1u128 << claimed_bit_index; + let mask = 1u64 << claimed_bit_index; (*word & mask) == mask } fun set_claimed_(distribution: &mut MerkleDistribution, index: u64) { - let claimed_word_index = index / 128; - let claimed_bit_index = ((index % 128) as u8); + let claimed_word_index = index / 64; + let claimed_bit_index = ((index % 64) as u8); let word = vector::borrow_mut(&mut distribution.claimed_bitmap, claimed_word_index); // word | (1 << bit_index) - let mask = 1u128 << claimed_bit_index; + let mask = 1u64 << claimed_bit_index; *word = (*word | mask); } - fun encode_leaf(index: &u64, account: &address, amount: &u128): vector { + fun encode_leaf(index: &u64, account: &address, amount: &u64): vector { let leaf = vector::empty(); - vector::append(&mut leaf, BCS::to_bytes(index)); - vector::append(&mut leaf, BCS::to_bytes(account)); - vector::append(&mut leaf, BCS::to_bytes(amount)); + vector::append(&mut leaf, bcs::to_bytes(index)); + vector::append(&mut leaf, bcs::to_bytes(account)); + vector::append(&mut leaf, bcs::to_bytes(amount)); leaf } } diff --git a/contrib-contracts/src/merkle_distributor_test.rs b/contrib-contracts/src/merkle_distributor_test.rs index 6ab33e153f..4473a6c955 100644 --- a/contrib-contracts/src/merkle_distributor_test.rs +++ b/contrib-contracts/src/merkle_distributor_test.rs @@ -18,17 +18,15 @@ use test_helper::executor::{ struct DataProof { address: AccountAddress, index: u64, - amount: u128, + amount: u64, /// proofs in hex string proof: Vec, } // XXX FIXME YSG next pr -#[ignore] #[stest::test] fn test_merkle_distributor() -> Result<()> { let association = Account::new_association(); - println!("{}", association.address()); let (chain_state, net) = prepare_genesis(); let merkle_data = include_str!("merkle-test.json"); let merkle_data: serde_json::Value = serde_json::from_str(merkle_data)?; @@ -41,7 +39,9 @@ fn test_merkle_distributor() -> Result<()> { let source = include_str!("../modules/MerkleDistributor.move"); let mut dep_libs = starcoin_move_stdlib::move_stdlib_files(); let starcoin_stdlib_files = starcoin_move_stdlib::starcoin_stdlib_files(); + let starcoin_framework_files = starcoin_move_stdlib::starcoin_framework_files(); dep_libs.extend(starcoin_stdlib_files); + dep_libs.extend(starcoin_framework_files); let modules = compile_modules_with_address_ext(association_address(), source, &dep_libs); let package = Package::new(modules, None)?; @@ -56,7 +56,7 @@ fn test_merkle_distributor() -> Result<()> { { let merkle_root = MoveValue::vector_u8(raw_root); - let rewards_total = MoveValue::U128(proofs.iter().map(|p| p.amount).sum()); + let rewards_total = MoveValue::U64(proofs.iter().map(|p| p.amount).sum()); let leaves = MoveValue::U64(proofs.len() as u64); let script_function = EntryFunction::new( @@ -65,7 +65,7 @@ fn test_merkle_distributor() -> Result<()> { Identifier::new("MerkleDistributorScripts").unwrap(), ), Identifier::new("create").unwrap(), - vec![stc_type_tag()], + vec![], vec![ merkle_root.simple_serialize().unwrap(), rewards_total.simple_serialize().unwrap(), @@ -113,7 +113,7 @@ fn test_merkle_distributor() -> Result<()> { let distribution_address = MoveValue::Address(*association.address()); let index = MoveValue::U64(association_proof.index); let account = MoveValue::Address(*association.address()); - let amount = MoveValue::U128(association_proof.amount + 1); + let amount = MoveValue::U64(association_proof.amount + 1); let proofs = MoveValue::Vector( association_proof .proof @@ -158,7 +158,7 @@ fn test_merkle_distributor() -> Result<()> { let distribution_address = MoveValue::Address(*association.address()); let index = MoveValue::U64(association_proof.index); let account = MoveValue::Address(*association.address()); - let amount = MoveValue::U128(association_proof.amount); + let amount = MoveValue::U64(association_proof.amount); let proofs = MoveValue::Vector( association_proof .proof diff --git a/vm/framework/move-stdlib/src/lib.rs b/vm/framework/move-stdlib/src/lib.rs index 60c586bbd4..949b8a3af6 100644 --- a/vm/framework/move-stdlib/src/lib.rs +++ b/vm/framework/move-stdlib/src/lib.rs @@ -23,12 +23,21 @@ pub static STARCOIN_STDLIB_SOURCE_FILES: Lazy = Lazy::new(|| { restore_sources(STARCOIN_STDLIB_SOURCES_DIR, "starcoin-stdlib") .expect("Restore source file error") }); +pub const STARCOIN_FRAMEWORK_SOURCES_DIR: Dir = include_dir!("../starcoin-framework/sources"); +pub static STARCOIN_FRAMEWORK_SOURCE_FILES: Lazy = Lazy::new(|| { + restore_sources(STARCOIN_FRAMEWORK_SOURCES_DIR, "starcoin-framework") + .expect("Restore source file error") +}); pub fn move_stdlib_files() -> Vec { MOVE_STDLIB_SOURCE_FILES.files.clone() } pub fn starcoin_stdlib_files() -> Vec { STARCOIN_STDLIB_SOURCE_FILES.files.clone() } +pub fn starcoin_framework_files() -> Vec { + STARCOIN_FRAMEWORK_SOURCE_FILES.files.clone() +} + //restore the sources files to a tempdir fn restore_sources(dir: Dir, path: &str) -> anyhow::Result { let temp_dir = tempfile::tempdir()?;