[aptos-framework] Remove balance checks in DFA

aptos-move/framework/aptos-framework/doc/dispatchable_fungible_asset.md

@@ -228,7 +228,6 @@ The semantics of deposit will be governed by the function specified in DispatchF
             <a href="../../aptos-stdlib/../move-stdlib/doc/features.md#0x1_features_dispatchable_fungible_asset_enabled">features::dispatchable_fungible_asset_enabled</a>(),
             <a href="../../aptos-stdlib/../move-stdlib/doc/error.md#0x1_error_aborted">error::aborted</a>(<a href="dispatchable_fungible_asset.md#0x1_dispatchable_fungible_asset_ENOT_ACTIVATED">ENOT_ACTIVATED</a>)
         );
-        <b>let</b> start_balance = <a href="fungible_asset.md#0x1_fungible_asset_balance">fungible_asset::balance</a>(store);
         <b>let</b> func = <a href="../../aptos-stdlib/../move-stdlib/doc/option.md#0x1_option_borrow">option::borrow</a>(&func_opt);
         <a href="function_info.md#0x1_function_info_load_module_from_function">function_info::load_module_from_function</a>(func);
         <b>let</b> fa = <a href="dispatchable_fungible_asset.md#0x1_dispatchable_fungible_asset_dispatchable_withdraw">dispatchable_withdraw</a>(
@@ -237,8 +236,6 @@ The semantics of deposit will be governed by the function specified in DispatchF
             <a href="dispatchable_fungible_asset.md#0x1_dispatchable_fungible_asset_borrow_transfer_ref">borrow_transfer_ref</a>(store),
             func,
         );
-        <b>let</b> end_balance = <a href="fungible_asset.md#0x1_fungible_asset_balance">fungible_asset::balance</a>(store);
-        <b>assert</b>!(amount &lt;= start_balance - end_balance, <a href="../../aptos-stdlib/../move-stdlib/doc/error.md#0x1_error_aborted">error::aborted</a>(<a href="dispatchable_fungible_asset.md#0x1_dispatchable_fungible_asset_EAMOUNT_MISMATCH">EAMOUNT_MISMATCH</a>));
         fa
     } <b>else</b> {
         <a href="fungible_asset.md#0x1_fungible_asset_unchecked_withdraw">fungible_asset::unchecked_withdraw</a>(<a href="object.md#0x1_object_object_address">object::object_address</a>(&store), amount)

aptos-move/framework/aptos-framework/sources/dispatchable_fungible_asset.move

@@ -84,7 +84,6 @@ module aptos_framework::dispatchable_fungible_asset {
                 features::dispatchable_fungible_asset_enabled(),
                 error::aborted(ENOT_ACTIVATED)
             );
-            let start_balance = fungible_asset::balance(store);
             let func = option::borrow(&func_opt);
             function_info::load_module_from_function(func);
             let fa = dispatchable_withdraw(
@@ -93,8 +92,6 @@ module aptos_framework::dispatchable_fungible_asset {
                 borrow_transfer_ref(store),
                 func,
             );
-            let end_balance = fungible_asset::balance(store);
-            assert!(amount <= start_balance - end_balance, error::aborted(EAMOUNT_MISMATCH));
             fa
         } else {
             fungible_asset::unchecked_withdraw(object::object_address(&store), amount)

aptos-move/framework/aptos-framework/tests/clamped_token.move

@@ -0,0 +1,92 @@
+#[test_only]
+module 0xcafe::clamped_token {
+    // Create a token with max amount one can withdraw on each withdraw call.
+
+    use aptos_framework::fungible_asset::{Self, FungibleAsset, RawBalanceRef, RawSupplyRef, TransferRef};
+    use aptos_framework::dispatchable_fungible_asset;
+    use aptos_framework::object::{ConstructorRef, Object};
+    use aptos_framework::function_info;
+
+    use std::option;
+    use std::option::Option;
+    use std::signer;
+    use std::string;
+
+    struct BalanceStore has key {
+        balance_ref: RawBalanceRef,
+        supply_ref: RawSupplyRef,
+    }
+
+    public fun initialize(account: &signer, constructor_ref: &ConstructorRef) {
+        assert!(signer::address_of(account) == @0xcafe, 1);
+        let balance_ref = fungible_asset::generate_raw_balance_ref(constructor_ref);
+        let supply_ref = fungible_asset::generate_raw_supply_ref(constructor_ref);
+        move_to<BalanceStore>(account, BalanceStore { balance_ref, supply_ref });
+
+        let balance_value = function_info::new_function_info(
+            account,
+            string::utf8(b"clamped_token"),
+            string::utf8(b"derived_balance"),
+        );
+        let supply_value = function_info::new_function_info(
+            account,
+            string::utf8(b"clamped_token"),
+            string::utf8(b"derived_supply"),
+        );
+
+        let withdraw = function_info::new_function_info(
+            account,
+            string::utf8(b"clamped_token"),
+            string::utf8(b"withdraw"),
+        );
+
+        let deposit = function_info::new_function_info(
+            account,
+            string::utf8(b"clamped_token"),
+            string::utf8(b"deposit"),
+        );
+
+        dispatchable_fungible_asset::register_dispatch_functions(
+            constructor_ref,
+            option::some(withdraw),
+            option::some(deposit),
+            option::some(balance_value)
+        );
+        dispatchable_fungible_asset::register_derive_supply_dispatch_function(
+            constructor_ref,
+            option::some(supply_value)
+        );
+    }
+
+    public fun derived_balance<T: key>(store: Object<T>): u64 acquires BalanceStore {
+        fungible_asset::balance_with_ref(
+            &borrow_global<BalanceStore>(@0xcafe).balance_ref,
+            store
+        )
+    }
+
+    public fun derived_supply<T: key>(metadata: Object<T>): Option<u128> acquires BalanceStore {
+        option::some(option::extract(&mut fungible_asset::supply_with_ref(
+            &borrow_global<BalanceStore>(@0xcafe).supply_ref,
+            metadata
+        )))
+    }
+
+    public fun withdraw<T: key>(
+        store: Object<T>,
+        amount: u64,
+        transfer_ref: &TransferRef,
+    ): FungibleAsset {
+        // Clamp the max amount of asset to withdraw: at most 10 can be withdrawn each call.
+        assert!(amount <= 10, 0);
+        fungible_asset::withdraw_with_ref(transfer_ref, store, amount)
+    }
+
+    public fun deposit<T: key>(
+        store: Object<T>,
+        fa: FungibleAsset,
+        transfer_ref: &TransferRef,
+    ) {
+        fungible_asset::deposit_with_ref(transfer_ref, store, fa)
+    }
+}

aptos-move/framework/aptos-framework/tests/clamped_token_tests.move

@@ -0,0 +1,55 @@
+#[test_only]
+module aptos_framework::clamped_token_tests {
+    use aptos_framework::fungible_asset::{Self, Metadata, TestToken};
+    use aptos_framework::dispatchable_fungible_asset;
+    use aptos_framework::object;
+    use 0xcafe::clamped_token;
+    use std::option;
+
+    #[test(creator = @0xcafe)]
+    fun test_clamped(
+        creator: &signer,
+    ) {
+        let (creator_ref, token_object) = fungible_asset::create_test_token(creator);
+        let (mint, _, _, _) = fungible_asset::init_test_metadata(&creator_ref);
+        let metadata = object::convert<TestToken, Metadata>(token_object);
+
+        let creator_store = fungible_asset::create_test_store(creator, metadata);
+
+        clamped_token::initialize(creator, &creator_ref);
+
+        assert!(dispatchable_fungible_asset::derived_supply(metadata) == option::some(0), 2);
+        // Mint
+        let fa = fungible_asset::mint(&mint, 100);
+        dispatchable_fungible_asset::deposit(creator_store, fa);
+
+        assert!(dispatchable_fungible_asset::derived_balance(creator_store) == 100, 4);
+        assert!(dispatchable_fungible_asset::derived_supply(metadata) == option::some(100), 5);
+
+        let fa = dispatchable_fungible_asset::withdraw(creator, creator_store, 5);
+        dispatchable_fungible_asset::deposit(creator_store, fa);
+    }
+
+    #[test(creator = @0xcafe)]
+    #[expected_failure(abort_code = 0, location = 0xcafe::clamped_token)]
+    fun test_clamped_aborted(
+        creator: &signer,
+    ) {
+        let (creator_ref, token_object) = fungible_asset::create_test_token(creator);
+        let (mint, _, _, _) = fungible_asset::init_test_metadata(&creator_ref);
+        let metadata = object::convert<TestToken, Metadata>(token_object);
+
+        let creator_store = fungible_asset::create_test_store(creator, metadata);
+
+        clamped_token::initialize(creator, &creator_ref);
+
+        assert!(dispatchable_fungible_asset::derived_supply(metadata) == option::some(0), 2);
+        // Mint
+        let fa = fungible_asset::mint(&mint, 100);
+        dispatchable_fungible_asset::deposit(creator_store, fa);
+
+        // Failed to withdraw as it exceeds the withdraw limit.
+        let fa = dispatchable_fungible_asset::withdraw(creator, creator_store, 20);
+        dispatchable_fungible_asset::deposit(creator_store, fa);
+    }
+}

aptos-move/framework/aptos-framework/tests/nil_op_token_tests.move

@@ -7,7 +7,6 @@ module aptos_framework::nil_op_token_tests {
     use std::option;
 
     #[test(creator = @0xcafe)]
-    #[expected_failure(abort_code=0x70002, location=aptos_framework::dispatchable_fungible_asset)]
     fun test_nil_op_token(
         creator: &signer,
     ) {
@@ -26,7 +25,6 @@ module aptos_framework::nil_op_token_tests {
         // Deposit will cause an re-entrant call into dispatchable_fungible_asset
         dispatchable_fungible_asset::deposit(creator_store, fa);
 
-        // Withdraw will fail because it's not drawing the basic amount.
         let fa = dispatchable_fungible_asset::withdraw(creator, creator_store, 10);
         dispatchable_fungible_asset::deposit(creator_store, fa);
     }

aptos-move/framework/aptos-framework/tests/ten_x_token.move

@@ -1,6 +1,6 @@
 #[test_only]
 module 0xcafe::ten_x_token {
-    use aptos_framework::fungible_asset::{Self, RawBalanceRef, RawSupplyRef};
+    use aptos_framework::fungible_asset::{Self, FungibleAsset, RawBalanceRef, RawSupplyRef, TransferRef};
     use aptos_framework::dispatchable_fungible_asset;
     use aptos_framework::object::{ConstructorRef, Object};
     use aptos_framework::function_info;
@@ -31,10 +31,23 @@ module 0xcafe::ten_x_token {
             string::utf8(b"ten_x_token"),
             string::utf8(b"derived_supply"),
         );
+
+        let withdraw = function_info::new_function_info(
+            account,
+            string::utf8(b"ten_x_token"),
+            string::utf8(b"withdraw"),
+        );
+
+        let deposit = function_info::new_function_info(
+            account,
+            string::utf8(b"ten_x_token"),
+            string::utf8(b"deposit"),
+        );
+
         dispatchable_fungible_asset::register_dispatch_functions(
             constructor_ref,
-            option::none(),
-            option::none(),
+            option::some(withdraw),
+            option::some(deposit),
             option::some(balance_value)
         );
         dispatchable_fungible_asset::register_derive_supply_dispatch_function(
@@ -58,4 +71,20 @@ module 0xcafe::ten_x_token {
             metadata
         )) * 10)
     }
+
+    public fun withdraw<T: key>(
+        store: Object<T>,
+        amount: u64,
+        transfer_ref: &TransferRef,
+    ): FungibleAsset {
+        fungible_asset::withdraw_with_ref(transfer_ref, store, amount)
+    }
+
+    public fun deposit<T: key>(
+        store: Object<T>,
+        fa: FungibleAsset,
+        transfer_ref: &TransferRef,
+    ) {
+        fungible_asset::deposit_with_ref(transfer_ref, store, fa)
+    }
 }