fix bug from unwrap_unchecked usage

AztecProtocol · Jul 31, 2023 · 6f4d187 · 6f4d187
1 parent 1530d57
commit 6f4d187
Show file tree

Hide file tree

Showing 4 changed files with 16 additions and 23 deletions.
diff --git a/yarn-project/noir-contracts/src/artifacts/pokeable_token_contract.json b/yarn-project/noir-contracts/src/artifacts/pokeable_token_contract.json
diff --git a/yarn-project/noir-contracts/src/contracts/pokeable_token_contract/src/main.nr b/yarn-project/noir-contracts/src/contracts/pokeable_token_contract/src/main.nr
@@ -6,7 +6,7 @@ contract PokeableToken {
     use dep::value_note::{
         balance_utils,
         utils::{send_note, spend_notes},
-        value_note::{VALUE_NOTE_LEN, ValueNoteInterface},
+        value_note::{VALUE_NOTE_LEN, ValueNoteInterface, ValueNote},
         filter::get_2_notes,
     };
     use dep::aztec::abi;
@@ -96,16 +96,20 @@ contract PokeableToken {
         let (mut new_context, maybe_notes) = sender_balance.get_notes(context, options);
         context = new_context;
 
-        // Note: dangerous to unwrap_unchecked. See zk_token_contract example for a safer implementation.
-        // TODO: any reason this doesn't use the 'spend_notes' util function?
-        let note1 = maybe_notes[0].unwrap_unchecked();
-        let note2 = maybe_notes[1].unwrap_unchecked();
+        let note0 = maybe_notes[0].unwrap_or(ValueNote::dummy());
+        let note1 = maybe_notes[1].unwrap_or(ValueNote::dummy());
 
-        let note_sum = note1.value + note2.value;
+        let note_sum = note0.value + note1.value;
 
         // Removes the 2 notes from the sender's set of notes.
-        context = sender_balance.remove(context, note1);
-        context = sender_balance.remove(context, note2);
+        if maybe_notes[0].is_some() {
+          assert(sender.address == note0.owner);
+          context = sender_balance.remove(context, note0);
+        }
+        if maybe_notes[1].is_some() {
+          assert(sender.address == note1.owner);
+          context = sender_balance.remove(context, note1);
+        }
 
         // Create new note for the recipient.
         let recipient_balance = storage.balances.at(recipient.address);

diff --git a/yarn-project/noir-contracts/src/contracts/pokeable_token_contract/src/storage.nr b/yarn-project/noir-contracts/src/contracts/pokeable_token_contract/src/storage.nr
@@ -25,7 +25,7 @@ impl Storage {
         Storage {
             sender: ImmutableSingleton::new(1, AddressNoteInterface),
             recipient: ImmutableSingleton::new(2, AddressNoteInterface),
-            balances: Map::new(1, |s| Set::new(s, ValueNoteInterface)),
+            balances: Map::new(3, |s| Set::new(s, ValueNoteInterface)),
         }
     }
 }
diff --git a/yarn-project/noir-libs/noir-aztec/src/types/option.nr b/yarn-project/noir-libs/noir-aztec/src/types/option.nr
@@ -45,22 +45,11 @@ impl<T> Option<T> {
         self._value
     }
 
-    // TODO: this would be cool.
-    // fn unwrap_mut_ref_unchecked(&mut self) -> &mut T {
-    //     &mut self._value
-    // }
-
     fn unwrap(self) -> T {
         assert(self._is_some);
         self._value
     }
 
-    // TODO: this would be cool.
-    // fn unwrap_mut_ref(&mut self) -> &mut T {
-    //     assert(self._is_some);
-    //     self._value
-    // }
-
     fn unwrap_or(self, default: T) -> T {
         if self._is_some {
             self._value