p-token: Update amount on self-unwrap lamports (#113)

Update amount on self unwrap

Author: febo

pinocchio/program/src/processor/unwrap_lamports.rs

@@ -60,30 +60,30 @@ pub fn process_unwrap_lamports(accounts: &[AccountInfo], instruction_data: &[u8]
         (source_account.amount(), 0)
     };
 
-    // Comparing whether the AccountInfo's "point" to the same account or
-    // not - this is a faster comparison since it just checks the internal
-    // raw pointer.
-    let self_transfer = source_account_info == destination_account_info;
-
-    if unlikely(self_transfer || amount == 0) {
+    if unlikely(amount == 0) {
         // Validates the token account owner since we are not writing
         // to the account.
         check_account_owner(source_account_info)
     } else {
         source_account.set_amount(remaining_amount);
 
-        // SAFETY: single mutable borrow to `source_account_info` lamports.
-        let source_lamports = unsafe { source_account_info.borrow_mut_lamports_unchecked() };
-        // Note: The amount of a source token account is already validated and the
-        // `lamports` on the account is always greater than `amount`.
-        *source_lamports -= amount;
+        // Comparing whether the AccountInfo's "point" to the same account or
+        // not - this is a faster comparison since it just checks the internal
+        // raw pointer.
+        if source_account_info != destination_account_info {
+            // SAFETY: single mutable borrow to `source_account_info` lamports.
+            let source_lamports = unsafe { source_account_info.borrow_mut_lamports_unchecked() };
+            // Note: The amount of a source token account is already validated and the
+            // `lamports` on the account is always greater than `amount`.
+            *source_lamports -= amount;
 
-        // SAFETY: single mutable borrow to `destination_account_info` lamports; the
-        // account is already validated to be different from `source_account_info`.
-        let destination_lamports =
-            unsafe { destination_account_info.borrow_mut_lamports_unchecked() };
-        // Note: The total lamports supply is bound to `u64::MAX`.
-        *destination_lamports += amount;
+            // SAFETY: single mutable borrow to `destination_account_info` lamports; the
+            // account is already validated to be different from `source_account_info`.
+            let destination_lamports =
+                unsafe { destination_account_info.borrow_mut_lamports_unchecked() };
+            // Note: The total lamports supply is bound to `u64::MAX`.
+            *destination_lamports += amount;
+        }
 
         Ok(())
     }

pinocchio/program/tests/unwrap_lamports.rs

@@ -409,7 +409,8 @@ fn unwrap_lamports_with_self_transfer() {
     .unwrap();
 
     // It should succeed to unwrap lamports with the same source and destination
-    // accounts.
+    // accounts. The amount should be deducted from the source account but the
+    // lamports should remain the same.
 
     let result = mollusk().process_and_validate_instruction(
         &instruction,
@@ -432,7 +433,7 @@ fn unwrap_lamports_with_self_transfer() {
 
     let account = account.unwrap();
     let token_account = spl_token_interface::state::Account::unpack(&account.data).unwrap();
-    assert_eq!(token_account.amount, 2_000_000_000);
+    assert_eq!(token_account.amount, 1_000_000_000);
 }
 
 #[test]