working on passing tests for whitelist contract

src/TokenElection/BaseTokenElection.ts

@@ -30,6 +30,7 @@ export class TokenElection extends SmartContract {
       partial2: PartialBallot.fromBigInts([0n, 0n, 0n, 0n, 0n, 0n, 0n]).packed
     });
     this.actionState.set(Reducer.initialActionState);
+    this.account.delegate.set(this.sender);
   }
 
   @method

src/TokenElection/WhitelistTokenElection.test.ts

@@ -0,0 +1,158 @@
+import { Ballot, PartialBallot } from './BaseTokenElection';
+import { AccountUpdate, Mina, PrivateKey, PublicKey, UInt32 } from 'o1js';
+import { WhitelistTokenElection } from './WhitelistTokenElection';
+
+describe("WhitelistTokenElection", () => {
+  let zkappAddress: PublicKey;
+  let sender: PublicKey;
+  let zkappKey: PrivateKey;
+  let senderKey: PrivateKey;
+  let initialBalance = 10_000_000_000;
+  let zkapp: WhitelistTokenElection;
+
+  beforeAll(async () => {
+    console.time('compile');
+    await WhitelistTokenElection.compile();
+    console.timeEnd('compile');
+  });
+
+  beforeEach(async () => {
+    let Local = Mina.LocalBlockchain({ proofsEnabled: true });
+    zkappKey = PrivateKey.random();
+    zkappAddress = zkappKey.toPublicKey();
+    zkapp = new WhitelistTokenElection(zkappAddress);
+    Mina.setActiveInstance(Local);
+
+    sender = Local.testAccounts[0].publicKey;
+    senderKey = Local.testAccounts[0].privateKey;
+  });
+
+  describe("Votes in the WhitelistTokenElection", () => {
+    beforeEach(async () => {
+      let Local = Mina.LocalBlockchain({ proofsEnabled: true });
+      zkappKey = PrivateKey.random();
+      zkappAddress = zkappKey.toPublicKey();
+      zkapp = new WhitelistTokenElection(zkappAddress);
+      Mina.setActiveInstance(Local);
+
+      sender = Local.testAccounts[0].publicKey;
+      senderKey = Local.testAccounts[0].privateKey;
+
+      let tx = await Mina.transaction(sender, () => {
+        let senderUpdate = AccountUpdate.fundNewAccount(sender);
+        senderUpdate.send({ to: zkappAddress, amount: initialBalance });
+        zkapp.deploy({ zkappKey });
+      });
+      await tx.prove();
+      await tx.sign([senderKey]).send();
+
+      const wlPks = [
+        Local.testAccounts[1].privateKey,
+        Local.testAccounts[2].privateKey,
+        Local.testAccounts[3].privateKey
+      ]
+
+      const wl = [
+        wlPks[0].toPublicKey(),
+        wlPks[1].toPublicKey(),
+        wlPks[2].toPublicKey()
+      ]
+
+      tx = await Mina.transaction(sender, () => {
+        let senderUpdate = AccountUpdate.fundNewAccount(sender);
+        zkapp.addToWhitelist(wl[0]);
+      });
+      await tx.prove();
+      await tx.sign([senderKey]).send();
+
+      tx = await Mina.transaction(sender, () => {
+        let senderUpdate = AccountUpdate.fundNewAccount(sender);
+        zkapp.addToWhitelist(wl[1]);
+      });
+      await tx.prove();
+      await tx.sign([senderKey]).send();
+
+      tx = await Mina.transaction(sender, () => {
+        let senderUpdate = AccountUpdate.fundNewAccount(sender);
+        zkapp.addToWhitelist(wl[2]);
+        zkapp.finalizeWhitelist();
+      });
+      await tx.prove();
+      await tx.sign([senderKey]).send();
+
+      tx = await Mina.transaction(wl[0], () => {
+        const partialBallot1 = PartialBallot.fromBigInts([0n, 0n, 100n, 0n, 0n, 0n, 0n]);
+        const partialBallot2 = PartialBallot.fromBigInts([0n, 0n, 0n, 0n, 0n, 0n, 10n]);
+        const myVote = new Ballot({
+          partial1: partialBallot1.packed,
+          partial2: partialBallot2.packed
+        })
+        zkapp.castVote(myVote, UInt32.from(110));
+      });
+      await tx.prove();
+      await tx.sign([wlPks[0]]).send();
+    });
+
+    it("Updates the State", async () => {
+      let Local = Mina.LocalBlockchain({ proofsEnabled: true });
+      const wlPks = [
+        Local.testAccounts[1].privateKey,
+        Local.testAccounts[2].privateKey,
+        Local.testAccounts[3].privateKey
+      ]
+
+      const wl = [
+        wlPks[0].toPublicKey(),
+        wlPks[1].toPublicKey(),
+        wlPks[2].toPublicKey()
+      ]
+      const tx2 = await Mina.transaction(sender, () => {
+        zkapp.reduceVotes();
+      })
+      await tx2.prove();
+      await tx2.sign([senderKey, wlPks[0]]).send();
+      const zkappState = zkapp.ballot.get();
+      const pb1 = PartialBallot.unpack(zkappState.partial1);
+      const pb2 = PartialBallot.unpack(zkappState.partial2);
+      expect(String(pb1)).toBe(String([0n, 0n, 100n, 0n, 0n, 0n, 0n]));
+      expect(String(pb2)).toBe(String([0n, 0n, 0n, 0n, 0n, 0n, 10n]));
+    });
+
+    it("has the correct number of tokens remaining", async () => {
+      let remainingVoteBalance = await Mina.getBalance(sender, zkapp.token.id);
+      expect(remainingVoteBalance.toString()).toBe(String(50_000))
+      const tx2 = await Mina.transaction(sender, () => {
+        zkapp.reduceVotes();
+      })
+      await tx2.prove();
+      await tx2.sign([senderKey]).send();
+      remainingVoteBalance = await Mina.getBalance(sender, zkapp.token.id);
+      expect(remainingVoteBalance.toString()).toBe(String(50_000 - 110))
+    });
+
+    it("updates a second time", async () => {
+      let tx = await Mina.transaction(sender, () => {
+        const partialBallot1 = PartialBallot.fromBigInts([101n, 0n, 0n, 0n, 0n, 0n, 0n]);
+        const partialBallot2 = PartialBallot.fromBigInts([0n, 0n, 0n, 0n, 20n, 0n, 40n]);
+        const myVote = new Ballot({
+          partial1: partialBallot1.packed,
+          partial2: partialBallot2.packed
+        })
+        zkapp.castVote(myVote, UInt32.from(161));
+      });
+      await tx.prove();
+      await tx.sign([senderKey]).send();
+
+      const tx2 = await Mina.transaction(sender, () => {
+        zkapp.reduceVotes();
+      })
+      await tx2.prove();
+      await tx2.sign([senderKey]).send();
+      const zkappState = zkapp.ballot.get();
+      const pb1 = PartialBallot.unpack(zkappState.partial1);
+      const pb2 = PartialBallot.unpack(zkappState.partial2);
+      expect(String(pb1)).toBe(String([101n, 0n, 100n, 0n, 0n, 0n, 0n]));
+      expect(String(pb2)).toBe(String([0n, 0n, 0n, 0n, 20n, 0n, 50n]));
+    });
+  });
+});

src/TokenElection/WhitelistTokenElection.ts

@@ -1,5 +1,140 @@
-import { TokenElection } from "./BaseTokenElection";
+import { SmartContract, state, Bool, Account, State, method, UInt32, AccountUpdate, UInt64, Reducer, Field, PublicKey, Struct, Provable } from 'o1js';
+import { PackedUInt32Factory, MultiPackedStringFactory, PackedBoolFactory } from 'o1js-pack';
 
-class WhitelistTokenElection extends TokenElection {
+export class Flags extends PackedBoolFactory() { }
+export class IpfsHash extends MultiPackedStringFactory(4) { }
+export class PartialBallot extends PackedUInt32Factory() { }
 
-}
+export class Ballot extends Struct({
+  partial1: Field,
+  partial2: Field
+}) { }
+
+export class VoteAction extends Struct({
+  ballot: Ballot,
+  sender: PublicKey,
+  amount: UInt64
+}) { }
+
+export class WhitelistTokenElection extends SmartContract {
+  @state(IpfsHash) electionDetailsIpfs = State<IpfsHash>();
+  @state(Ballot) ballot = State<Ballot>();
+  @state(Field) actionState = State<Field>();
+  @state(Flags) flags = State<Flags>();
+
+  reducer = Reducer({ actionType: VoteAction });
+
+  init() {
+    super.init();
+    this.electionDetailsIpfs.set(IpfsHash.fromString(''));
+    this.ballot.set({
+      partial1: PartialBallot.fromBigInts([0n, 0n, 0n, 0n, 0n, 0n, 0n]).packed,
+      partial2: PartialBallot.fromBigInts([0n, 0n, 0n, 0n, 0n, 0n, 0n]).packed
+    });
+    this.actionState.set(Reducer.initialActionState);
+    this.account.delegate.set(this.sender);
+  }
+
+  assertWhitelistOpen() {
+    const flags = this.flags.getAndAssertEquals();
+    Flags.unpack(flags.packed)[0].assertEquals(Bool(false));
+  }
+
+  assertAccountUnfunded(address: PublicKey) {
+    const account = Account(address, this.token.id);
+    const tokenBalance = account.balance.getAndAssertEquals();
+    tokenBalance.assertEquals(UInt64.from(0));
+  }
+
+  @method
+  addToWhitelist(address: PublicKey) {
+    this.assertWhitelistOpen();
+    this.assertAccountUnfunded(address);
+    this.token.mint({
+      address: address,
+      amount: 50_000
+    });
+  }
+
+  @method
+  finalizeWhitelist() {
+    const delegate = this.account.delegate.getAndAssertEquals();
+    this.sender.assertEquals(delegate);
+
+    const flags = this.flags.getAndAssertEquals();
+    const unpackedFlags = Flags.unpack(flags.packed)
+    unpackedFlags[0].assertEquals(Bool(false));
+
+    unpackedFlags[0] = Bool(true);
+    this.flags.set(Flags.fromBools(unpackedFlags));
+  }
+
+  @method
+  setElectionDetails(electionDetailsIpfs: IpfsHash) {
+    this.electionDetailsIpfs.getAndAssertEquals();
+    this.electionDetailsIpfs.assertEquals(IpfsHash.fromString(''));
+    this.electionDetailsIpfs.set(electionDetailsIpfs);
+  }
+
+  @method
+  castVote(vote: Ballot, amount: UInt32) {
+    const unpackedVote1 = PartialBallot.unpack(vote.partial1);
+    const unpackedVote2 = PartialBallot.unpack(vote.partial2);
+
+    let voteSum = UInt32.from(0);
+    for (let i = 0; i < PartialBallot.l; i++) {
+      voteSum = voteSum.add(unpackedVote1[i]);
+    }
+    for (let i = 0; i < PartialBallot.l; i++) {
+      voteSum = voteSum.add(unpackedVote2[i]);
+    }
+    voteSum.assertEquals(amount); // sum of votes must equal asserted amount (can vote for multiple options)
+    this.reducer.dispatch({
+      ballot: vote,
+      sender: this.sender,
+      amount: UInt64.from(amount)
+    });
+  }
+
+  @method
+  reduceVotes() {
+    const ballot = this.ballot.getAndAssertEquals();
+    const actionState = this.actionState.getAndAssertEquals();
+
+    let pendingActions = this.reducer.getActions({
+      fromActionState: actionState,
+    }).slice(0, 3); // at most, reduce 3 actions
+
+    let { state: newVotes, actionState: newActionState } =
+      this.reducer.reduce(
+        pendingActions,
+        Ballot,
+        (state: Ballot, _action: VoteAction) => {
+          const unpackedState1 = PartialBallot.unpack(state.partial1);
+          const unpackedState2 = PartialBallot.unpack(state.partial2);
+          const unpackedAction1 = PartialBallot.unpack(_action.ballot.partial1);
+          const unpackedAction2 = PartialBallot.unpack(_action.ballot.partial2);
+          for (let i = 0; i < PartialBallot.l; i++) {
+            unpackedState1[i] = unpackedState1[i].add(unpackedAction1[i])
+          }
+          for (let i = 0; i < PartialBallot.l; i++) {
+            unpackedState2[i] = unpackedState2[i].add(unpackedAction2[i])
+          }
+          this.token.burn({
+            address: _action.sender,
+            amount: UInt64.from(_action.amount)
+          });
+          return {
+            partial1: PartialBallot.fromUInt32s(unpackedState1).packed,
+            partial2: PartialBallot.fromUInt32s(unpackedState2).packed
+          }
+        },
+        {
+          state: ballot, actionState: actionState
+        }
+      );
+
+    this.ballot.set(newVotes);
+    this.actionState.set(newActionState);
+  }
+}