add split teams and revote functionality. Still need tests

gate approval with nft

combine create_shortlist and split_shortlist

update project list collection and add tests

Author: zihehuang

contracts/voting/Move.lock

@@ -21,7 +21,7 @@ dependencies = [
 ]
 
 [move.toolchain-version]
-compiler-version = "1.33.0"
+compiler-version = "1.33.1"
 edition = "2024.beta"
 flavor = "sui"
 

contracts/voting/sources/approval.move

@@ -0,0 +1,82 @@
+module voting::approval;
+
+use voting::voting::{Votes, Project, AdminCap, new_votes, vote_internal,
+assert_token_in_whitelist, assert_valid_project_ids, assert_voting_is_active};
+use sui::random::{Random};
+use sui::table;
+use std::debug;
+
+public struct ProjectWithAddress has drop {
+  project: Project,
+  address: address,
+}
+
+public struct TeamOrca has key { id: UID }
+public struct TeamPolarBear has key { id: UID }
+
+fun init(_: &mut TxContext) {}
+
+// AdminCap is required to create a shortlist
+entry fun create_shortlist(_: &AdminCap, votes: &Votes, project_ids: vector<u64>, addresses: vector<address>, r: &Random, ctx: &mut TxContext) {
+  let mut projects = project_ids.zip_map!(addresses, |i, a| {
+    let p = votes.project_list(i);
+    ProjectWithAddress {
+      project: p,
+      address: a,
+    }
+  });
+
+  // Split shortlist into 2 teams
+  let mut rng = r.new_generator(ctx);
+  rng.shuffle(&mut projects);
+
+  let mut team_orca = new_votes<TeamOrca>(
+    0,
+    table::new(ctx),
+    table::new(ctx),
+    false,
+    ctx
+  );
+  let mut team_polar_bear = new_votes<TeamPolarBear>(
+    0,
+    table::new(ctx),
+    table::new(ctx),
+    false,
+    ctx
+  );
+
+  // divide the projects into 2 teams
+  // if the number of projects is odd, the first team will have 1 more project
+  let n = projects.length();
+  debug::print(&b"project length".to_string());
+  debug::print(&n);
+  let n1 = n / 2;
+
+  vector::tabulate!(n, |i| {
+    let p = projects.pop_back();
+    if (i <= n1) {
+      team_orca.append_project_list(p.project);
+      transfer::transfer(TeamPolarBear{id: object::new(ctx)}, p.address);
+      debug::print(&b"TeamPolarBear".to_string());
+      debug::print(&p.address);
+    } else {
+      team_polar_bear.append_project_list(p.project);
+      transfer::transfer(TeamOrca{id: object::new(ctx)}, p.address);
+      debug::print(&b"TeamOrca".to_string());
+      debug::print(&p.address);
+    };
+    true
+  });
+
+  team_orca.share_votes();
+  team_polar_bear.share_votes();
+  projects.destroy_empty();
+}
+
+public fun approve<T>(nft: &T, ballot: vector<u64>, v: &mut Votes, ctx: &TxContext) {
+  assert_voting_is_active(v);
+  assert_token_in_whitelist(nft, v);
+  assert_valid_project_ids(ballot, v);
+  
+  vote_internal(v, ballot, ctx);
+}

contracts/voting/sources/voting.move

@@ -6,22 +6,64 @@ module voting::voting {
   use sui::vec_map;
   use sui::url;
   use sui::zklogin_verified_issuer::check_zklogin_issuer;
+  use sui::vec_set::{Self,VecSet};
+  use std::type_name::{Self, TypeName};
 
   const EInvalidProof: u64 = 1;
   const EUserAlreadyVoted: u64 = 2;
-  const ETooManyVotes: u64 = 3;
   const EInvalidProjectId: u64 = 4;
   const EVotingInactive: u64 = 5;
+  const ENotInWhitelist: u64 = 6;
 
   public struct Votes has key {
     id: UID, 
     total_votes: u64, 
-    project_list: vector<Project>,
-    votes: table::Table<address, vector<u64>>,
-    voting_active: bool
+    project_list: table::Table<u64, Project>,
+    ballots: table::Table<address, vector<u64>>,
+    voting_active: bool,
+    whitelist_tokens: VecSet<TypeName>,
   }
 
-  public struct Project has store {
+  public(package) fun project_list(self: &Votes, i: u64): Project {
+    return self.project_list[i]
+  }
+
+  public(package) fun append_project_list(self: &mut Votes, p: Project) {
+    self.project_list.add(p.id, p);
+  }
+
+  public(package) fun share_votes(self: Votes) {
+    transfer::share_object(self);
+  }
+
+  public fun total_votes(self: &Votes): u64 {
+    self.total_votes
+  }
+
+  public fun ballots(self: &Votes, user: address): vector<u64> {
+    *self.ballots.borrow(user)
+  }
+
+  public(package) fun new_votes<T>(
+    total_votes: u64,
+    project_list: table::Table<u64, Project>,
+    ballots: table::Table<address, vector<u64>>,
+    voting_active: bool,
+    ctx: &mut TxContext
+  ): Votes {
+    let tn = type_name::get<T>();
+    return Votes {
+      id: object::new(ctx),
+      total_votes,
+      project_list,
+      // project_ids: vec_set::from_keys(project_list.map!(|p| p.id)),
+      ballots,
+      voting_active: voting_active,
+      whitelist_tokens: vec_set::singleton(tn),
+    }
+  }
+
+  public struct Project has store, copy, drop {
     id: u64,
     name: string::String, 
     description: string::String, 
@@ -31,6 +73,10 @@ module voting::voting {
     votes: u64
   }
 
+  public fun project_votes(self: &Project): u64 {
+    self.votes
+  }
+
   public struct AdminCap has key, store {
     id: UID
   }
@@ -520,12 +566,12 @@ module voting::voting {
       ],
     ];
 
-    let mut project_list = vector[];
+    let mut project_list = table::new(ctx);
 
     let mut index = 0;
 
     while (index < projects.length()) {
-      project_list.push_back(Project {
+      project_list.add(index, Project {
         id: index, 
         votes: 0, 
         name: projects[index][0].to_string(), 
@@ -542,8 +588,9 @@ module voting::voting {
       id: object::new(ctx),
       total_votes: 0, 
       project_list,
-      votes: table::new(ctx),
-      voting_active: false
+      ballots: table::new(ctx),
+      voting_active: false,
+      whitelist_tokens: vec_set::empty()
     };
     transfer::share_object(votes);
 
@@ -554,9 +601,8 @@ module voting::voting {
       ctx.sender()
     );
   }
-
+  
   public fun vote(project_ids: vector<u64>, votes: &mut Votes, address_seed: u256, ctx: &TxContext) {
-
     let voter = ctx.sender();
 
     assert_user_has_not_voted(voter, votes);
@@ -565,58 +611,72 @@ module voting::voting {
     assert_voting_is_active(votes);
 
     // Update project's vote
-    let mut curr_index = 0;
-    while (curr_index < project_ids.length()) {
-      let project = &mut votes.project_list[project_ids[curr_index]];
-      project.votes = project.votes + 1;
-
-      // Increment total votes
-      votes.total_votes = votes.total_votes + 1;
-
-      curr_index = curr_index + 1;
-    };
+    vote_internal(votes, project_ids, ctx);
 
     // Record user's vote
     table::add(
-      &mut votes.votes, 
+      &mut votes.ballots, 
       voter, 
       project_ids
     );
   }
 
+  public(package) fun vote_internal(votes: &mut Votes, ballot: vector<u64>, ctx: &TxContext) {
+    let voter = ctx.sender();
+    let already_voted = votes.ballots.contains(voter);
+    // Clean up old ballot
+    if (already_voted) {
+      let og_ballot = votes.ballots[voter];
+      og_ballot.do!(|v| {
+        let p = &mut votes.project_list[v];
+        p.votes = p.votes - 1;
+        votes.total_votes = votes.total_votes - 1;
+      });
+      votes.ballots.remove(voter);
+    };
+    // add new ballot
+    votes.ballots.add(voter, ballot);
+    ballot.do!(|v| {
+      let p = &mut votes.project_list[v];
+      p.votes = p.votes + 1;
+      votes.total_votes = votes.total_votes + 1;
+    });
+  }
+
   public entry fun toggle_voting(_: &AdminCap, can_vote: bool, votes: &mut Votes) {
     votes.voting_active = can_vote;
   }
 
+  public(package) fun assert_token_in_whitelist<T>(_: &T, votes: &Votes) {
+    let tn = type_name::get<T>();
+    assert!(
+      votes.whitelist_tokens.contains(&tn), 
+      ENotInWhitelist
+    );
+  }
+
   fun assert_user_has_not_voted(user: address, votes: &Votes) {
     assert!(
       table::contains(
-        &votes.votes, 
+        &votes.ballots, 
         user
       ) == false, 
       EUserAlreadyVoted
     );
   }
 
-  fun assert_valid_project_ids(project_ids: vector<u64>, votes: &Votes) {
-    // assert!(
-    //   project_ids.length() <= 3, 
-    //   ETooManyVotes
-    // );
-    
-    let mut curr_index = 0;
+  public(package) fun assert_valid_project_ids(project_ids: vector<u64>, votes: &Votes) {
     let mut ids = vec_map::empty();
-    while (curr_index < project_ids.length()) {
+    project_ids.do!(|id| {
       assert!(
-        project_ids[curr_index] < votes.project_list.length(),
-        EInvalidProjectId
+          votes.project_list.contains(id),
+          EInvalidProjectId
       );
-      vec_map::insert(&mut ids, project_ids[curr_index], 0); // this will abort if there is a dup
-      curr_index = curr_index + 1;
-    };
+      vec_map::insert(&mut ids, id, 0); // this will abort if there is a dup
+    });
   }
 
-  fun assert_voting_is_active(votes: &Votes) {
+  public(package) fun assert_voting_is_active(votes: &Votes) {
     assert!(
       votes.voting_active, 
       EVotingInactive
@@ -628,4 +688,9 @@ module voting::voting {
     let issuer = string::utf8(b"https://accounts.google.com");
     assert!(check_zklogin_issuer(sender, address_seed, &issuer), EInvalidProof);
   }
+
+  #[test_only]
+  public fun init_for_test(ctx: &mut TxContext) {
+    init(ctx);
+  }
 }

contracts/voting/tests/voting_tests.move

@@ -1,19 +1,247 @@
-/*
 #[test_only]
-module voting::voting_tests {
-    // uncomment this line to import the module
-    // use voting::voting;
+module voting::voting_tests;
+// uncomment this line to import the module
+use voting::voting::{Self, AdminCap, Votes, toggle_voting, ENotInWhitelist, EInvalidProjectId, EVotingInactive};
+use voting::approval::{create_shortlist, TeamOrca, approve};
+use sui::test_scenario::{Self as ts, Scenario};
+use sui::random::{Self,Random};
+use std::debug;
 
-    const ENotImplemented: u64 = 0;
+const ORGANIZER: address = @0xAAA;
+const Voter1: address = @0x111;
+const Voter2: address = @0x222;
+const Voter3: address = @0x333;
+const Voter4: address = @0x444;
+const Voter5: address = @0x555;
 
-    #[test]
-    fun test_voting() {
-        // pass
-    }
+#[test]
+fun initialize(): Scenario {
+    let mut scenario = ts::begin(ORGANIZER);
+    voting::init_for_test(scenario.ctx());
 
-    #[test, expected_failure(abort_code = ::voting::voting_tests::ENotImplemented)]
-    fun test_voting_fail() {
-        abort ENotImplemented
-    }
+    ts::next_tx(&mut scenario, @0x0);
+    {
+        random::create_for_testing(scenario.ctx());
+    };
+
+    ts::next_tx(&mut scenario, @0x0);
+    let mut random_state: Random = scenario.take_shared();
+    {
+        random_state.update_randomness_state_for_testing(
+            0,
+            x"1F1F1F1F1F1F1F1F1F1F1F1F1F1F1F1F1F1F1F1F1F1F1F1F1F1F1F1F1F1FFF",
+            scenario.ctx(),
+        );
+        ts::return_shared(random_state);
+    };
+    ts::next_tx(&mut scenario, ORGANIZER);
+    {
+        let random_state: Random = scenario.take_shared<Random>();
+        let cap = scenario.take_from_sender<AdminCap>();
+        let og_votes = scenario.take_shared<Votes>();
+        let shortlisted_projects = vector[1, 2, 3, 4, 5];
+        let whitelist_addresses = vector[Voter1, Voter2, Voter3, Voter4, Voter5];
+        // Guaranteed to split the team the same way due to hardcoded randomness
+        // TeamPolarBear: @0x222, @0x555, @0x444
+        // TeamOrca: @0x111, @0x333
+        create_shortlist(&cap, &og_votes, shortlisted_projects, whitelist_addresses, &random_state, scenario.ctx());
+        ts::return_shared(og_votes);
+        ts::return_shared(random_state);
+        ts::return_to_sender(&scenario, cap);
+    };
+    scenario
+}
+
+#[test]
+fun happy_path(): Scenario {
+    let mut scenario = initialize();
+
+    let effects = ts::next_tx(&mut scenario, ORGANIZER);
+    let approval_votes = ts::shared(&effects);
+    {
+        let mut team_polar_bear_projects = scenario.take_shared_by_id<Votes>(approval_votes[0]);
+        let mut team_orca_projects = scenario.take_shared_by_id<Votes>(approval_votes[1]);
+        let cap = scenario.take_from_sender<AdminCap>();
+        toggle_voting(&cap, true, &mut team_polar_bear_projects);
+        toggle_voting(&cap, true, &mut team_orca_projects);
+        ts::return_shared(team_polar_bear_projects);
+        ts::return_shared(team_orca_projects);
+        ts::return_to_sender(&scenario, cap);
+    };
+
+
+    // Team Orca member votes for project 4 and 5
+    ts::next_tx(&mut scenario, Voter1);
+    {
+        let mut votes = scenario.take_shared_by_id<Votes>(approval_votes[0]);
+        let orca = scenario.take_from_sender<TeamOrca>();
+        debug::print(&votes);
+
+        approve<TeamOrca>(&orca, vector[4, 5], &mut votes, scenario.ctx());
+
+        ts::return_shared(votes);
+        ts::return_to_sender(&scenario, orca);
+    };
+    ts::next_tx(&mut scenario, Voter3);
+    {
+        let mut votes = scenario.take_shared_by_id<Votes>(approval_votes[0]);
+        let orca = scenario.take_from_sender<TeamOrca>();
+        debug::print(&votes);
+
+        approve<TeamOrca>(&orca, vector[4], &mut votes, scenario.ctx());
+
+        ts::return_shared(votes);
+        ts::return_to_sender(&scenario, orca);
+    };
+    // Revote #1
+    ts::next_tx(&mut scenario, Voter1);
+    {
+        let mut votes = scenario.take_shared_by_id<Votes>(approval_votes[0]);
+        let orca = scenario.take_from_sender<TeamOrca>();
+        debug::print(&votes);
+
+        approve<TeamOrca>(&orca, vector[5, 4], &mut votes, scenario.ctx());
+
+        ts::return_shared(votes);
+        ts::return_to_sender(&scenario, orca);
+    };
+
+    // Assert first vote
+    ts::next_tx(&mut scenario, Voter1);
+    {
+        let votes = scenario.take_shared_by_id<Votes>(approval_votes[0]);
+        debug::print(&votes);
+        assert!(votes.project_list(4).project_votes() == 2);
+        assert!(votes.project_list(5).project_votes() == 1);
+        assert!(votes.total_votes() == 3);
+        assert!(votes.ballots(Voter1) == vector[5, 4]);
+        assert!(votes.ballots(Voter3) == vector[4]);
+        ts::return_shared(votes);
+    };
+
+    // Revote #2
+    ts::next_tx(&mut scenario, Voter1);
+    {
+        let mut votes = scenario.take_shared_by_id<Votes>(approval_votes[0]);
+        let orca = scenario.take_from_sender<TeamOrca>();
+        debug::print(&votes);
+
+        // Team Orca member votes for project 4 and 5
+        approve<TeamOrca>(&orca, vector[2], &mut votes, scenario.ctx());
+
+        ts::return_shared(votes);
+        ts::return_to_sender(&scenario, orca);
+    };
+    // Assert revote
+    ts::next_tx(&mut scenario, Voter1);
+    {
+        let votes = scenario.take_shared_by_id<Votes>(approval_votes[0]);
+        debug::print(&votes);
+        assert!(votes.project_list(2).project_votes() == 1);
+        assert!(votes.project_list(4).project_votes() == 1);
+        assert!(votes.project_list(5).project_votes() == 0);
+        assert!(votes.total_votes() == 2);
+        assert!(votes.ballots(Voter1) == vector[2]);
+        assert!(votes.ballots(Voter3) == vector[4]);
+        ts::return_shared(votes);
+    };
+
+    scenario
+}
+
+#[test]
+#[expected_failure(abort_code = EVotingInactive)]
+fun voting_not_active_error(): Scenario {
+    let mut scenario = initialize();
+
+    let effects = ts::next_tx(&mut scenario, ORGANIZER);
+    let approval_votes = ts::shared(&effects);
+    let team_polar_bear_id = approval_votes[0];
+
+    // Team Orca member votes for project 4 and 5
+    ts::next_tx(&mut scenario, Voter1);
+    {
+        let mut votes = scenario.take_shared_by_id<Votes>(team_polar_bear_id);
+        let orca = scenario.take_from_sender<TeamOrca>();
+        debug::print(&votes);
+
+        approve<TeamOrca>(&orca, vector[4, 5], &mut votes, scenario.ctx());
+
+        ts::return_shared(votes);
+        ts::return_to_sender(&scenario, orca);
+    };
+
+    scenario
+}
+
+#[test]
+#[expected_failure(abort_code = ENotInWhitelist)]
+fun not_in_whitelist_error(): Scenario {
+    let mut scenario = initialize();
+
+    let effects = ts::next_tx(&mut scenario, ORGANIZER);
+    let approval_votes = ts::shared(&effects);
+    let team_polar_bear_id = approval_votes[0];
+    let team_orca_id = approval_votes[1];
+    {
+        let mut team_polar_bear_projects = scenario.take_shared_by_id<Votes>(team_polar_bear_id);
+        let mut team_orca_projects = scenario.take_shared_by_id<Votes>(team_orca_id);
+        let cap = scenario.take_from_sender<AdminCap>();
+        toggle_voting(&cap, true, &mut team_polar_bear_projects);
+        toggle_voting(&cap, true, &mut team_orca_projects);
+        ts::return_shared(team_polar_bear_projects);
+        ts::return_shared(team_orca_projects);
+        ts::return_to_sender(&scenario, cap);
+    };
+
+    // Team Orca member votes for Team Orca projects
+    ts::next_tx(&mut scenario, Voter1);
+    {
+        let mut votes = scenario.take_shared_by_id<Votes>(team_orca_id);
+        let orca = scenario.take_from_sender<TeamOrca>();
+        debug::print(&votes);
+
+        approve<TeamOrca>(&orca, vector[3], &mut votes, scenario.ctx());
+
+        ts::return_shared(votes);
+        ts::return_to_sender(&scenario, orca);
+    };
+
+    scenario
+}
+
+#[test]
+#[expected_failure(abort_code = EInvalidProjectId)]
+fun invalid_project_id_error(): Scenario {
+    let mut scenario = initialize();
+
+    let effects = ts::next_tx(&mut scenario, ORGANIZER);
+    let approval_votes = ts::shared(&effects);
+    let team_polar_bear_id = approval_votes[0];
+    let team_orca_id = approval_votes[1];
+    {
+        let mut team_polar_bear_projects = scenario.take_shared_by_id<Votes>(team_polar_bear_id);
+        let mut team_orca_projects = scenario.take_shared_by_id<Votes>(team_orca_id);
+        let cap = scenario.take_from_sender<AdminCap>();
+        toggle_voting(&cap, true, &mut team_polar_bear_projects);
+        toggle_voting(&cap, true, &mut team_orca_projects);
+        ts::return_shared(team_polar_bear_projects);
+        ts::return_shared(team_orca_projects);
+        ts::return_to_sender(&scenario, cap);
+    };
+
+    // Team Orca member tries to vote for project 1
+    ts::next_tx(&mut scenario, Voter1);
+    {
+        let mut votes = scenario.take_shared_by_id<Votes>(team_polar_bear_id);
+        let orca = scenario.take_from_sender<TeamOrca>();
+        debug::print(&votes);
+
+        approve<TeamOrca>(&orca, vector[1], &mut votes, scenario.ctx());
+
+        ts::return_shared(votes);
+        ts::return_to_sender(&scenario, orca);
+    };
+
+    scenario
 }
-*/