Skip to content

Commit

Permalink
update merkledistributor contract
Browse files Browse the repository at this point in the history
  • Loading branch information
@nkysg
nkysg committed Jan 15, 2025
1 parent 89ed235 commit f6be02f
Show file tree
Hide file tree
Showing 3 changed files with 69 additions and 59 deletions.
105 changes: 53 additions & 52 deletions contrib-contracts/modules/MerkleDistributor.move
View file
Original file line number Diff line number Diff line change
@@ -1,40 +1,42 @@
address StarcoinAssociation {
module MerkleDistributorScripts {
use StarcoinAssociation::MerkleDistributor;
use StarcoinFramework::Account;
public(script) fun create<T: store>(signer: signer, merkle_root: vector<u8>, token_amounts: u128, leaves: u64) {
let tokens = Account::withdraw<T>(&signer, token_amounts);
MerkleDistributor::create<T>(&signer, merkle_root, tokens, leaves);
use starcoin_framework::coin;
use starcoin_std::signer;
public entry fun create<T: store>(signer: signer, merkle_root: vector<u8>, coin_amounts: u64, leaves: u64) {
let coins = coin::withdraw<T>(&signer, coin_amounts);
MerkleDistributor::create<T>(&signer, merkle_root, coins, leaves);
}

public(script) fun claim_for_address<T: store>(distribution_address: address, index: u64, account: address, amount: u128, merkle_proof: vector<vector<u8>>) {
MerkleDistributor::claim_for_address<T>(distribution_address, index, account, amount, merkle_proof);
public entry fun claim_for_address<T: store>(signer: signer, distribution_address: address, index: u64, account: address, amount: u64, merkle_proof: vector<vector<u8>>) {
MerkleDistributor::claim_for_address<T>(&signer, distribution_address, index, account, amount, merkle_proof);
}
public(script) fun claim<T: store>(signer: signer, distribution_address: address, index: u64, amount: u128, merkle_proof: vector<vector<u8>>) {
let tokens = MerkleDistributor::claim<T>(&signer, distribution_address, index, amount, merkle_proof);
Account::deposit_to_self<T>(&signer, tokens);
public entry fun claim<T: store>(signer: signer, distribution_address: address, index: u64, amount: u64, merkle_proof: vector<vector<u8>>) {
let coins = MerkleDistributor::claim<T>(&signer, distribution_address, index, amount, merkle_proof);
let account_addr = signer::address_of(&signer);
coin::deposit<T>(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<vector<u8>>, root: &vector<u8>, leaf: vector<u8>): bool {
let computed_hash = leaf;
let i = 0;
let proof_length = vector::length(proof);
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);

};

Expand All @@ -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<T: store> has key {
merkle_root: vector<u8>,
tokens: Token<T>,
claimed_bitmap: vector<u128>,
coins: coin::Coin<T>,
claimed_bitmap: vector<u64>,
}
const INVALID_PROOF: u64 = 1;
const ALREADY_CLAIMED: u64 = 2;

/// Initialization.
public fun create<T: store>(signer: &signer, merkle_root: vector<u8>, tokens: Token<T>, leaves: u64) {
let bitmap_count = leaves / 128;
if (bitmap_count * 128 < leaves) {
public fun create<T: store>(signer: &signer, merkle_root: vector<u8>, coins: coin::Coin<T>, 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<T: store>(distribution_address: address, index: u64, account: address, amount: u128, merkle_proof: vector<vector<u8>>)
public fun claim_for_address<T: store>(signer: &signer, distribution_address: address, index: u64, account: address, amount: u64, merkle_proof: vector<vector<u8>>)
acquires MerkleDistribution {
let distribution = borrow_global_mut<MerkleDistribution<T>>(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<T: store>(signer: &signer, distribution_address: address, index: u64, amount: u128, merkle_proof: vector<vector<u8>>): Token<T>
public fun claim<T: store>(signer: &signer, distribution_address: address, index: u64, amount: u64, merkle_proof: vector<vector<u8>>): coin::Coin<T>
acquires MerkleDistribution {
let distribution = borrow_global_mut<MerkleDistribution<T>>(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.
Expand All @@ -113,41 +114,41 @@ module MerkleDistributor {
is_claimed_(distribution, index)
}

fun internal_claim<T: store>(distribution: &mut MerkleDistribution<T>, index: u64, account: address, amount: u128, merkle_proof: vector<vector<u8>>): Token<T> {
fun internal_claim<T: store>(signer: &signer, distribution: &mut MerkleDistribution<T>, index: u64, account: address, amount: u64, merkle_proof: vector<vector<u8>>): coin::Coin<T> {
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_<T: store>(distribution: &MerkleDistribution<T>, 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_<T: store>(distribution: &mut MerkleDistribution<T>, 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<u8> {
fun encode_leaf(index: &u64, account: &address, amount: &u64): vector<u8> {
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
}
}
Expand Down
14 changes: 7 additions & 7 deletions contrib-contracts/src/merkle_distributor_test.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,17 +18,15 @@ use test_helper::executor::{
struct DataProof {
address: AccountAddress,
index: u64,
amount: u128,
amount: u64,
/// proofs in hex string
proof: Vec<String>,
}

// 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)?;
Expand All @@ -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)?;
Expand All @@ -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(
Expand All @@ -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(),
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down
9 changes: 9 additions & 0 deletions vm/framework/move-stdlib/src/lib.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -23,12 +23,21 @@ pub static STARCOIN_STDLIB_SOURCE_FILES: Lazy<SourceFiles> = 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<SourceFiles> = Lazy::new(|| {
restore_sources(STARCOIN_FRAMEWORK_SOURCES_DIR, "starcoin-framework")
.expect("Restore source file error")
});
pub fn move_stdlib_files() -> Vec<String> {
MOVE_STDLIB_SOURCE_FILES.files.clone()
}
pub fn starcoin_stdlib_files() -> Vec<String> {
STARCOIN_STDLIB_SOURCE_FILES.files.clone()
}
pub fn starcoin_framework_files() -> Vec<String> {
STARCOIN_FRAMEWORK_SOURCE_FILES.files.clone()
}

//restore the sources files to a tempdir
fn restore_sources(dir: Dir, path: &str) -> anyhow::Result<SourceFiles> {
let temp_dir = tempfile::tempdir()?;
Expand Down

0 comments on commit f6be02f

Please sign in to comment.