better wrapped mina example

examples/zkapps/11-advanced-account-updates/src/WrappedMina.ts

@@ -1,155 +1,125 @@
 import {
-  Bool,
-  DeployArgs,
   Int64,
   method,
-  Mina,
   AccountUpdate,
-  Permissions,
   PublicKey,
   UInt64,
   State,
   state,
   TokenContract,
   AccountUpdateForest,
+  AccountUpdateTree,
+  assert,
+  Bool,
+  Reducer,
+  TokenId,
+  Provable,
+  Types,
+  Permissions,
 } from 'o1js';
 
-export class WrappedMina extends TokenContract {
-  async deploy(args?: DeployArgs) {
-    await super.deploy(args);
-    this.account.permissions.set({
-      ...Permissions.default(),
-      send: Permissions.proof(),
-    });
-  }
+export { WrappedMina };
 
-  @state(UInt64) priorMina = State<UInt64>();
+const $MINA = TokenId.default;
 
-  // ----------------------------------------------------------------------
+class WrappedMina extends TokenContract {
+  @state(UInt64) totalSupply = State<UInt64>();
 
-  @method async init() {
-    super.init();
+  // actions are total supply changes triggered by minting or burning
+  totalSupplyReducer = Reducer({ actionType: Int64 });
 
-    let receiver = this.internal.mint({
-      address: this.address,
-      amount: UInt64.from(0),
-    });
-    // require that the receiving account is new, so this can be only done once
-    receiver.account.isNew.requireEquals(Bool(true));
-    // pay fees for opened account
-    this.balance.subInPlace(Mina.getNetworkConstants().accountCreationFee);
-    this.priorMina.set(UInt64.from(0));
+  @method async init() {
+    super.init(); // totalSupply provable starts at 0
   }
 
-  // ----------------------------------------------------------------------
-  async approveBase(forest: AccountUpdateForest) {
-    this.checkZeroBalanceChange(forest);
+  get $wMINA() {
+    return this.deriveTokenId();
   }
 
-  // ----------------------------------------------------------------------
-  @method async mintWrappedMina(amount: UInt64, destination: PublicKey) {
-    const priorMina = this.priorMina.get();
-    this.priorMina.requireEquals(this.priorMina.get());
-
-    const newMina = amount.add(priorMina);
-
-    // TODO is there a way to directly get the balance change for this transaction?
-    this.account.balance.requireBetween(newMina, UInt64.MAXINT());
+  // approve any transaction which leaves the token supply unchanged
+  @method async approveBase(forest: AccountUpdateForest) {
+    let sum = Int64.from(0);
 
-    this.internal.mint({ address: destination, amount });
+    this.forEachUpdate(forest, (update, usesToken) => {
+      sum = Provable.if(usesToken, sum.add(update.balanceChange), sum);
+      checkPermissionsUpdate(update);
+    });
 
-    this.priorMina.set(newMina);
+    sum.assertEquals(Int64.zero);
   }
 
-  // ----------------------------------------------------------------------
-
-  @method async redeemWrappedMinaApprove(
-    burnWMINA: AccountUpdate,
-    amount: UInt64
-  ) {
-    // check that the burn account update has our token id
-    burnWMINA.body.tokenId.assertEquals(this.tokenId);
+  @method async wrap(sender: AccountUpdateTree) {
+    // get sender update and ensure there are no other updates
+    assert(sender.children.isEmpty());
+    let senderUpdate = sender.accountUpdate.unhash();
+    this.approve(sender);
 
-    // approve burn with at most 2 child account updates, which don't get token permissions
-    this.approve(burnWMINA);
+    // ensure sender is giving away a positive amount of MINA
+    senderUpdate.tokenId.assertEquals($MINA);
+    let amount = senderUpdate.balanceChange.neg();
+    assert(amount.isPositive());
 
-    // check that the account update burns the specified amount
-    let balanceChange = Int64.fromObject(burnWMINA.body.balanceChange);
-    balanceChange.assertEquals(Int64.from(amount).neg());
+    // move MINA from sender to this contract
+    this.balance.addInPlace(amount);
 
-    // in return for burn, decrease our MINA balance (can be picked up as balance increase anywhere it suits the caller)
-    this.balance.subInPlace(amount);
+    // mint same amount of wrapped MINA to sender
+    let senderTokenUpdate = this.internal.mint({
+      address: senderUpdate.publicKey,
+      amount: amount.magnitude,
+    });
+    // allow minting to pay for account creation if the account doesn't exist
+    senderTokenUpdate.body.implicitAccountCreationFee = Bool(true);
 
-    // update priorMina
-    const priorMina = this.priorMina.get();
-    this.priorMina.requireEquals(this.priorMina.get());
-    const newMina = priorMina.sub(amount);
-    this.priorMina.set(newMina);
+    // increase total wMINA supply
+    this.totalSupplyReducer.dispatch(amount);
   }
 
-  // ----------------------------------------------------------------------
-
-  @method async redeemWrappedMinaWithoutApprove(
-    source: PublicKey,
-    destination: PublicKey,
-    amount: UInt64
-  ) {
-    this.internal.burn({ address: source, amount });
-
-    const priorMina = this.priorMina.get();
-    this.priorMina.requireEquals(this.priorMina.get());
-
-    const newMina = priorMina.sub(amount);
-
-    this.send({ to: destination, amount });
-
-    this.priorMina.set(newMina);
-  }
+  @method async unwrap(sender: AccountUpdateTree) {
+    /// get sender update and ensure there are no other updates
+    assert(sender.children.isEmpty());
+    let senderTokenUpdate = sender.accountUpdate.unhash();
+    checkPermissionsUpdate(senderTokenUpdate);
+    this.approve(sender);
+
+    // ensure sender is burning a positive amount of wrapped MINA
+    senderTokenUpdate.tokenId.assertEquals(this.$wMINA);
+    let amount = senderTokenUpdate.balanceChange.neg();
+    assert(amount.isPositive());
+
+    // release same amount of MINA in return for burning
+    let senderUpdate = this.send({
+      to: senderTokenUpdate.publicKey,
+      amount: amount.magnitude,
+    });
+    // allow sending to pay for account creation if the account doesn't exist
+    senderUpdate.body.implicitAccountCreationFee = Bool(true);
 
-  // ----------------------------------------------------------------------
-
-  // let a zkapp send tokens to someone, provided the token supply stays constant
-  @method async approveUpdateAndSend(
-    zkappUpdate: AccountUpdate,
-    to: PublicKey,
-    amount: UInt64
-  ) {
-    this.approve(zkappUpdate); // TODO is this secretly approving other changes?
-
-    // see if balance change cancels the amount sent
-    let balanceChange = Int64.fromObject(zkappUpdate.body.balanceChange);
-    balanceChange.assertEquals(Int64.from(amount).neg());
-    // add same amount of tokens to the receiving address
-    this.internal.mint({ address: to, amount });
+    // decrease total wMINA supply
+    this.totalSupplyReducer.dispatch(amount.neg());
   }
 
-  // ----------------------------------------------------------------------
-
-  // let a zkapp do anything, provided the token supply stays constant
-  @method async approveUpdate(zkappUpdate: AccountUpdate) {
-    this.approve(zkappUpdate); // TODO is this secretly approving other changes?
-
-    // see if balance change is zero
-    let balanceChange = Int64.fromObject(zkappUpdate.body.balanceChange);
-    balanceChange.assertEquals(Int64.from(0));
+  @method.returns(UInt64) async getBalance(publicKey: PublicKey) {
+    let accountUpdate = AccountUpdate.create(publicKey, this.$wMINA);
+    return accountUpdate.account.balance.getAndRequireEquals();
   }
+}
 
-  // ----------------------------------------------------------------------
+function checkPermissionsUpdate(update: AccountUpdate) {
+  let permissions = update.update.permissions;
 
-  @method async transfer(from: PublicKey, to: PublicKey, value: UInt64) {
-    this.internal.send({ from, to, amount: value });
-  }
+  // account must not change its permissions in a way that prevents sending to it
+  let { access, receive } = permissions.value;
+  let accessIsNone = permissionEquals(access, Permissions.none());
+  let receiveIsNone = permissionEquals(receive, Permissions.none());
+  let updateAllowed = accessIsNone.and(receiveIsNone);
 
-  // ----------------------------------------------------------------------
-
-  @method async getBalance(publicKey: PublicKey): UInt64 {
-    let accountUpdate = AccountUpdate.create(publicKey, this.tokenId);
-    let balance = accountUpdate.account.balance.get();
-    accountUpdate.account.balance.requireEquals(
-      accountUpdate.account.balance.get()
-    );
-    return balance;
-  }
+  // either do an allowed update, or don't change permissions
+  assert(updateAllowed.or(permissions.isSome.not()));
+}
 
-  // ----------------------------------------------------------------------
+function permissionEquals(p1: Types.AuthRequired, p2: Types.AuthRequired) {
+  return p1.constant
+    .equals(p2.constant)
+    .and(p1.signatureNecessary.equals(p2.signatureNecessary))
+    .and(p1.signatureSufficient.equals(p2.signatureSufficient));
 }