From bb174bff5927e57e2d72ba92fd4b22da0e109a37 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Wei=C3=9Fer=20Hase?= <wei3erHase@protonmail.com>
Date: Thu, 20 Mar 2025 16:32:51 -0300
Subject: [PATCH] fix: unconstraining authwit lib method

---
 src/token_contract/src/main.nr                | 36 ++++++-------------
 src/token_contract/src/test/mint_to_public.nr |  6 ++--
 .../src/test/transfer_private_to_public.nr    |  2 +-
 3 files changed, 15 insertions(+), 29 deletions(-)

diff --git a/src/token_contract/src/main.nr b/src/token_contract/src/main.nr
index 76c4cccd..3d6aed3e 100644
--- a/src/token_contract/src/main.nr
+++ b/src/token_contract/src/main.nr
@@ -204,13 +204,7 @@ pub contract Token {
 
     #[public]
     fn transfer_public_to_public(from: AztecAddress, to: AztecAddress, amount: u128, nonce: Field) {
-        // _validate_from_public(from, nonce);
-        // TODO(#34): we cannot call `assert_current_call_valid_authwit_public` from library methods
-        if (!from.eq(context.msg_sender())) {
-            assert_current_call_valid_authwit_public(&mut context, from);
-        } else {
-            assert(nonce == 0, "invalid nonce");
-        }
+        _validate_from_public(&mut context, from, nonce);
         _decrease_public_balance(storage.public_balances, from, amount);
         _increase_public_balance(storage.public_balances, to, amount);
     }
@@ -225,13 +219,7 @@ pub contract Token {
         partial_note: PartialUintNote,
         nonce: Field,
     ) {
-        // _validate_from_public(from, nonce);
-        // TODO(#34): we cannot call `assert_current_call_valid_authwit_public` from library methods
-        if (!from.eq(context.msg_sender())) {
-            assert_current_call_valid_authwit_public(&mut context, from);
-        } else {
-            assert(nonce == 0, "invalid nonce");
-        }
+        _validate_from_public(&mut context, from, nonce);
         _finalize_transfer_public_to_private(
             &mut context,
             storage.public_balances,
@@ -285,8 +273,7 @@ pub contract Token {
      * ===================== UNCONSTRAINED =======================
      * ======================================================== */
 
-    // TODO(#32): unconstrained fns cannot be injected in macros
-    pub(crate) unconstrained fn balance_of_private(owner: AztecAddress) -> pub u128 {
+    unconstrained fn balance_of_private(owner: AztecAddress) -> pub u128 {
         storage.private_balances.at(owner).balance_of()
     }
 
@@ -362,13 +349,7 @@ pub contract Token {
 
     #[public]
     fn burn_public(from: AztecAddress, amount: u128, nonce: Field) {
-        // _validate_from_public(from, nonce);
-        // TODO(#34): we cannot call `assert_current_call_valid_authwit_public` from library methods
-        if (!from.eq(context.msg_sender())) {
-            assert_current_call_valid_authwit_public(&mut context, from);
-        } else {
-            assert(nonce == 0, "invalid nonce");
-        }
+        _validate_from_public(&mut context, from, nonce);
         _decrease_public_balance(storage.public_balances, from, amount);
         _decrease_total_supply(storage.total_supply, amount);
     }
@@ -546,10 +527,13 @@ pub contract Token {
     }
 
     #[contract_library_method]
-    fn _validate_from_public(context: &mut PublicContext, from: AztecAddress, nonce: Field) {
+    unconstrained fn _validate_from_public(
+        context: &mut PublicContext,
+        from: AztecAddress,
+        nonce: Field,
+    ) {
         if (!from.eq(context.msg_sender())) {
-            // TODO(#34): uncomment assert_current_call_valid_authwit_public
-            // assert_current_call_valid_authwit_public(&mut context, from);
+            assert_current_call_valid_authwit_public(context, from);
         } else {
             assert(nonce == 0, "invalid nonce");
         }
diff --git a/src/token_contract/src/test/mint_to_public.nr b/src/token_contract/src/test/mint_to_public.nr
index 70435a5a..d8ec5119 100644
--- a/src/token_contract/src/test/mint_to_public.nr
+++ b/src/token_contract/src/test/mint_to_public.nr
@@ -43,14 +43,16 @@ unconstrained fn mint_to_public_failures() {
 
     Token::at(token_contract_address).mint_to_public(recipient, max_u128).call(&mut env.public());
 
-    let mint_to_public_call_interface = Token::at(token_contract_address).mint_to_public(owner, 1 as u128);
+    let mint_to_public_call_interface =
+        Token::at(token_contract_address).mint_to_public(owner, 1 as u128);
     env.assert_public_call_fails(mint_to_public_call_interface);
 
     utils::check_public_balance(token_contract_address, owner, 0);
     utils::check_total_supply(token_contract_address, max_u128);
 
     // Overflow total supply
-    let mint_to_public_call_interface = Token::at(token_contract_address).mint_to_public(owner, max_u128);
+    let mint_to_public_call_interface =
+        Token::at(token_contract_address).mint_to_public(owner, max_u128);
     env.assert_public_call_fails(mint_to_public_call_interface);
 
     utils::check_public_balance(token_contract_address, owner, 0);
diff --git a/src/token_contract/src/test/transfer_private_to_public.nr b/src/token_contract/src/test/transfer_private_to_public.nr
index e73dc735..a1b75b89 100644
--- a/src/token_contract/src/test/transfer_private_to_public.nr
+++ b/src/token_contract/src/test/transfer_private_to_public.nr
@@ -2,8 +2,8 @@ use crate::test::utils;
 use crate::Token;
 use authwit::cheatcodes as authwit_cheatcodes;
 use aztec::{note::constants::MAX_NOTES_PER_PAGE, oracle::random::random};
-use uint_note::uint_note::UintNote;
 use aztec::protocol_types::traits::ToField;
+use uint_note::uint_note::UintNote;
 
 #[test]
 unconstrained fn transfer_private_to_public_on_behalf_of_self() {