fix(freezeV2): optimize Stake2.0 code

actuator/src/main/java/org/tron/core/actuator/DelegateResourceActuator.java

@@ -148,8 +148,8 @@ public boolean validate() throws ContractValidateException {
 
         long accountNetUsage = ownerCapsule.getNetUsage();
         if (null != this.getTx() && this.getTx().isTransactionCreate()) {
-          accountNetUsage += TransactionUtil.estimateConsumeBandWidthSize(ownerCapsule,
-              chainBaseManager);
+          accountNetUsage += TransactionUtil.estimateConsumeBandWidthSize(
+              ownerCapsule.getBalance());
         }
         long netUsage = (long) (accountNetUsage * TRX_PRECISION * ((double)
             (dynamicStore.getTotalNetWeight()) / dynamicStore.getTotalNetLimit()));

actuator/src/main/java/org/tron/core/actuator/FreezeBalanceV2Actuator.java

@@ -54,8 +54,8 @@ public boolean execute(Object result) throws ContractExeException {
       accountCapsule.initializeOldTronPower();
     }
 
-    long newBalance = accountCapsule.getBalance() - freezeBalanceV2Contract.getFrozenBalance();
     long frozenBalance = freezeBalanceV2Contract.getFrozenBalance();
+    long newBalance = accountCapsule.getBalance() - frozenBalance;
 
     switch (freezeBalanceV2Contract.getResource()) {
       case BANDWIDTH:

actuator/src/main/java/org/tron/core/actuator/UnfreezeBalanceV2Actuator.java

@@ -11,6 +11,8 @@
 import com.google.common.collect.Lists;
 import com.google.protobuf.ByteString;
 import com.google.protobuf.InvalidProtocolBufferException;
+
+import java.util.ArrayList;
 import java.util.Iterator;
 import java.util.List;
 import java.util.Objects;
@@ -365,19 +367,23 @@ private void updateVote(AccountCapsule accountCapsule,
     }
 
     // Update Owner Voting
-    votesCapsule.clearNewVotes();
+    List<Vote> addVotes = new ArrayList<>();
     for (Vote vote : accountCapsule.getVotesList()) {
       long newVoteCount = (long)
           ((double) vote.getVoteCount() / totalVote * ownedTronPower / TRX_PRECISION);
       if (newVoteCount > 0) {
-        votesCapsule.addNewVotes(vote.getVoteAddress(), newVoteCount);
+        Vote newVote = Vote.newBuilder()
+            .setVoteAddress(vote.getVoteAddress())
+            .setVoteCount(newVoteCount)
+            .build();
+        addVotes.add(newVote);
       }
     }
+    votesCapsule.clearNewVotes();
+    votesCapsule.addAllNewVotes(addVotes);
     votesStore.put(ownerAddress, votesCapsule);
 
     accountCapsule.clearVotes();
-    for (Vote vote : votesCapsule.getNewVotes()) {
-      accountCapsule.addVotes(vote.getVoteAddress(), vote.getVoteCount());
-    }
+    accountCapsule.addAllVotes(addVotes);
   }
 }

actuator/src/main/java/org/tron/core/utils/TransactionUtil.java

@@ -253,9 +253,24 @@ public TransactionSignWeight getTransactionSignWeight(Transaction trx) {
     return tswBuilder.build();
   }
 
+  public static long estimateConsumeBandWidthSize(long balance) {
+    DelegateResourceContract.Builder builder = DelegateResourceContract.newBuilder()
+        .setLock(true)
+        .setBalance(balance);
+    long builderSize = builder.build().getSerializedSize();
+
+    DelegateResourceContract.Builder builder2 = DelegateResourceContract.newBuilder()
+        .setBalance(TRX_PRECISION);
+    long builder2Size = builder2.build().getSerializedSize();
+    long addSize = Math.max(builderSize - builder2Size, 0L);
+
+    return DELEGATE_COST_BASE_SIZE + addSize;
+  }
+
+  // only for testing
   public static long consumeBandWidthSize(
-          final TransactionCapsule transactionCapsule,
-          ChainBaseManager chainBaseManager) {
+      final TransactionCapsule transactionCapsule,
+      ChainBaseManager chainBaseManager) {
     long bytesSize;
 
     boolean supportVM = chainBaseManager.getDynamicPropertiesStore().supportVM();
@@ -278,21 +293,21 @@ public static long consumeBandWidthSize(
     return bytesSize;
   }
 
-
-  public static long estimateConsumeBandWidthSize(
-          final AccountCapsule ownerCapsule,
-          ChainBaseManager chainBaseManager) {
+  // only for testing
+  public static long estimateConsumeBandWidthSizeOld(
+      final AccountCapsule ownerCapsule,
+      ChainBaseManager chainBaseManager) {
     DelegateResourceContract.Builder builder = DelegateResourceContract.newBuilder()
-                    .setLock(true)
-                    .setBalance(ownerCapsule.getFrozenV2BalanceForBandwidth());
+        .setLock(true)
+        .setBalance(ownerCapsule.getFrozenV2BalanceForBandwidth());
     TransactionCapsule fakeTransactionCapsule = new TransactionCapsule(builder.build()
-            , ContractType.DelegateResourceContract);
+        , ContractType.DelegateResourceContract);
     long size1 = consumeBandWidthSize(fakeTransactionCapsule, chainBaseManager);
 
     DelegateResourceContract.Builder builder2 = DelegateResourceContract.newBuilder()
-                    .setBalance(TRX_PRECISION);
+        .setBalance(TRX_PRECISION);
     TransactionCapsule fakeTransactionCapsule2 = new TransactionCapsule(builder2.build()
-            , ContractType.DelegateResourceContract);
+        , ContractType.DelegateResourceContract);
     long size2 = consumeBandWidthSize(fakeTransactionCapsule2, chainBaseManager);
     long addSize = Math.max(size1 - size2, 0L);
 

framework/src/main/java/org/tron/core/Wallet.java

@@ -780,6 +780,10 @@ public GrpcAPI.CanWithdrawUnfreezeAmountResponseMessage getCanWithdrawUnfreezeAm
           ByteString ownerAddress, long timestamp) {
     GrpcAPI.CanWithdrawUnfreezeAmountResponseMessage.Builder builder =
             GrpcAPI.CanWithdrawUnfreezeAmountResponseMessage.newBuilder();
+    if (timestamp < 0) {
+      return builder.build();
+    }
+
     long canWithdrawUnfreezeAmount;
 
     AccountStore accountStore = chainBaseManager.getAccountStore();
@@ -842,13 +846,7 @@ public GrpcAPI.GetAvailableUnfreezeCountResponseMessage getAvailableUnfreezeCoun
     }
     long now = dynamicStore.getLatestBlockHeaderTimestamp();
 
-    List<UnFreezeV2> unfrozenV2List = accountCapsule.getInstance().getUnfrozenV2List();
-    long getUsedUnfreezeCount = unfrozenV2List
-            .stream()
-            .filter(unfrozenV2 ->
-                    (unfrozenV2.getUnfreezeAmount() > 0
-                     && unfrozenV2.getUnfreezeExpireTime() > now))
-            .count();
+    long getUsedUnfreezeCount = accountCapsule.getUnfreezingV2Count(now);
     getAvailableUnfreezeCount = UnfreezeBalanceV2Actuator.getUNFREEZE_MAX_TIMES()
             - getUsedUnfreezeCount;
     builder.setCount(getAvailableUnfreezeCount);
@@ -869,7 +867,7 @@ public long calcCanDelegatedBandWidthMaxSize(
 
     long accountNetUsage = ownerCapsule.getNetUsage();
     accountNetUsage += org.tron.core.utils.TransactionUtil.estimateConsumeBandWidthSize(
-            ownerCapsule, chainBaseManager);
+            ownerCapsule.getBalance());
 
     long netUsage = (long) (accountNetUsage * TRX_PRECISION * ((double)
             (dynamicStore.getTotalNetWeight()) / dynamicStore.getTotalNetLimit()));

framework/src/test/java/org/tron/core/actuator/utils/TransactionUtilTest.java

@@ -3,29 +3,40 @@
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertTrue;
 import static org.tron.core.capsule.utils.TransactionUtil.isNumber;
+import static org.tron.core.config.Parameter.ChainConstant.DELEGATE_COST_BASE_SIZE;
+import static org.tron.core.config.Parameter.ChainConstant.TRX_PRECISION;
 
+import com.google.protobuf.ByteString;
 import java.nio.charset.StandardCharsets;
+
 import lombok.extern.slf4j.Slf4j;
-import org.junit.BeforeClass;
+import org.junit.Assert;
 import org.junit.Test;
 import org.tron.common.BaseTest;
+import org.tron.common.utils.ByteArray;
+import org.tron.core.ChainBaseManager;
 import org.tron.core.Constant;
+import org.tron.core.Wallet;
+import org.tron.core.capsule.AccountCapsule;
 import org.tron.core.config.args.Args;
 import org.tron.core.utils.TransactionUtil;
+import org.tron.protos.Protocol;
+import org.tron.protos.contract.BalanceContract;
 
 
 @Slf4j(topic = "capsule")
 public class TransactionUtilTest extends BaseTest {
 
-  /**
-   * Init .
-   */
-  @BeforeClass
-  public static void init() {
-    dbPath = "output_transactionUtil_test";
+  private static final String dbPath = "output_transactionUtil_test";
+  private static final String OWNER_ADDRESS;
+
+  static {
+    OWNER_ADDRESS =
+        Wallet.getAddressPreFixString() + "548794500882809695a8a687866e76d4271a1abc";
     Args.setParam(new String[]{"--output-directory", dbPath}, Constant.TEST_CONF);
   }
 
+
   @Test
   public void validAccountNameCheck() {
     String account = "";
@@ -114,4 +125,185 @@ public void isNumberCheck() {
     assertTrue(isNumber(number.getBytes(StandardCharsets.UTF_8)));
   }
 
+
+
+  @Test
+  public void testEstimateConsumeBandWidthSizeOld() {
+    dbManager.getDynamicPropertiesStore().saveAllowCreationOfContracts(1L);
+    ChainBaseManager chainBaseManager =  dbManager.getChainBaseManager();
+    long balance = 1000_000L;
+
+    AccountCapsule ownerCapsule = new AccountCapsule(ByteString.copyFromUtf8("owner"),
+        ByteString.copyFrom(ByteArray.fromHexString(OWNER_ADDRESS)), Protocol.AccountType.Normal,
+        balance);
+    ownerCapsule.addFrozenBalanceForBandwidthV2(balance);
+    dbManager.getAccountStore().put(ownerCapsule.createDbKey(), ownerCapsule);
+    ownerCapsule = dbManager.getAccountStore().get(ByteArray.fromHexString(OWNER_ADDRESS));
+    long estimateConsumeBandWidthSize1 =  TransactionUtil.estimateConsumeBandWidthSizeOld(
+        ownerCapsule, chainBaseManager);
+    Assert.assertEquals(277, estimateConsumeBandWidthSize1);
+
+    balance = 1000_000_000L;
+    ownerCapsule = new AccountCapsule(ByteString.copyFromUtf8("owner"),
+        ByteString.copyFrom(ByteArray.fromHexString(OWNER_ADDRESS)), Protocol.AccountType.Normal,
+        balance);
+    ownerCapsule.addFrozenBalanceForBandwidthV2(balance);
+    dbManager.getAccountStore().put(ownerCapsule.createDbKey(), ownerCapsule);
+    long estimateConsumeBandWidthSize2 =  TransactionUtil.estimateConsumeBandWidthSizeOld(
+        ownerCapsule, chainBaseManager);
+    Assert.assertEquals(279, estimateConsumeBandWidthSize2);
+
+    balance = 1000_000_000_000L;
+    ownerCapsule = new AccountCapsule(ByteString.copyFromUtf8("owner"),
+        ByteString.copyFrom(ByteArray.fromHexString(OWNER_ADDRESS)), Protocol.AccountType.Normal,
+        balance);
+    ownerCapsule.addFrozenBalanceForBandwidthV2(balance);
+    dbManager.getAccountStore().put(ownerCapsule.createDbKey(), ownerCapsule);
+    long estimateConsumeBandWidthSize3 =  TransactionUtil.estimateConsumeBandWidthSizeOld(
+        ownerCapsule, chainBaseManager);
+    Assert.assertEquals(280, estimateConsumeBandWidthSize3);
+
+    balance = 1000_000_000_000_000L;
+    ownerCapsule = new AccountCapsule(ByteString.copyFromUtf8("owner"),
+        ByteString.copyFrom(ByteArray.fromHexString(OWNER_ADDRESS)), Protocol.AccountType.Normal,
+        balance);
+    ownerCapsule.addFrozenBalanceForBandwidthV2(balance);
+    dbManager.getAccountStore().put(ownerCapsule.createDbKey(), ownerCapsule);
+    long estimateConsumeBandWidthSize4 =  TransactionUtil.estimateConsumeBandWidthSizeOld(
+        ownerCapsule, chainBaseManager);
+    Assert.assertEquals(282, estimateConsumeBandWidthSize4);
+  }
+
+
+  @Test
+  public void testEstimateConsumeBandWidthSizeNew() {
+    long balance = 1000_000L;
+
+    long estimateConsumeBandWidthSize1 =  TransactionUtil.estimateConsumeBandWidthSize(balance);
+    Assert.assertEquals(277, estimateConsumeBandWidthSize1);
+
+    balance = 1000_000_000L;
+    long estimateConsumeBandWidthSize2 =  TransactionUtil.estimateConsumeBandWidthSize(balance);
+    Assert.assertEquals(279, estimateConsumeBandWidthSize2);
+
+    balance = 1000_000_000_000L;
+    long estimateConsumeBandWidthSize3 =  TransactionUtil.estimateConsumeBandWidthSize(balance);
+    Assert.assertEquals(280, estimateConsumeBandWidthSize3);
+
+    balance = 1000_000_000_000_000L;
+    long estimateConsumeBandWidthSize4 =  TransactionUtil.estimateConsumeBandWidthSize(balance);
+    Assert.assertEquals(282, estimateConsumeBandWidthSize4);
+  }
+
+
+
+  @Test
+  public void testEstimateConsumeBandWidthSize() {
+    dbManager.getDynamicPropertiesStore().saveAllowCreationOfContracts(1L);
+    ChainBaseManager chainBaseManager =  dbManager.getChainBaseManager();
+    long balance = 1000_000L;
+
+    AccountCapsule ownerCapsule;
+    long estimateConsumeBandWidthSizeOld;
+    long estimateConsumeBandWidthSizeNew;
+
+    for (int i = 0; i < 100; i++) {
+      // old value is
+      ownerCapsule = new AccountCapsule(ByteString.copyFromUtf8("owner"),
+          ByteString.copyFrom(ByteArray.fromHexString(OWNER_ADDRESS)), Protocol.AccountType.Normal,
+          balance);
+      ownerCapsule.addFrozenBalanceForBandwidthV2(balance);
+      dbManager.getAccountStore().put(ownerCapsule.createDbKey(), ownerCapsule);
+      estimateConsumeBandWidthSizeOld =  TransactionUtil.estimateConsumeBandWidthSizeOld(
+          ownerCapsule, chainBaseManager);
+
+      // new value is
+      estimateConsumeBandWidthSizeNew =  TransactionUtil.estimateConsumeBandWidthSize(balance);
+
+      System.out.println("balance:"
+          + balance
+          + ", estimateConsumeBandWidthSizeOld:"
+          + estimateConsumeBandWidthSizeOld
+          + ", estimateConsumeBandWidthSizeNew:"
+          + estimateConsumeBandWidthSizeNew);
+      // new value assert equal to old value
+      Assert.assertEquals(estimateConsumeBandWidthSizeOld, estimateConsumeBandWidthSizeNew);
+
+      // balance accumulated
+      balance = balance * 10;
+      if (balance < 0) {
+        break;
+      }
+    }
+
+  }
+
+  @Test
+  public void estimateConsumeBandWidthSizePositive() {
+    long balance = 100;
+    BalanceContract.DelegateResourceContract.Builder builder =
+        BalanceContract.DelegateResourceContract.newBuilder()
+        .setLock(true)
+        .setBalance(balance);
+    BalanceContract.DelegateResourceContract.Builder builder2 =
+        BalanceContract.DelegateResourceContract.newBuilder()
+        .setBalance(TRX_PRECISION);
+
+    long expected = DELEGATE_COST_BASE_SIZE + Math.max(
+        builder.build().getSerializedSize() - builder2.build().getSerializedSize(), 0L);
+    long actual = TransactionUtil.estimateConsumeBandWidthSize(balance);
+    Assert.assertEquals(expected, actual);
+  }
+
+  @Test
+  public void estimateConsumeBandWidthSizeBoundary() {
+    long balance = TRX_PRECISION;
+    BalanceContract.DelegateResourceContract.Builder builder =
+        BalanceContract.DelegateResourceContract.newBuilder()
+        .setLock(true)
+        .setBalance(balance);
+    BalanceContract.DelegateResourceContract.Builder builder2 =
+        BalanceContract.DelegateResourceContract.newBuilder()
+        .setBalance(TRX_PRECISION);
+
+    long expected = DELEGATE_COST_BASE_SIZE + Math.max(
+        builder.build().getSerializedSize() - builder2.build().getSerializedSize(), 0L);
+    long actual = TransactionUtil.estimateConsumeBandWidthSize(balance);
+    Assert.assertEquals(expected, actual);
+  }
+
+  @Test
+  public void estimateConsumeBandWidthSizeEdge() {
+    long balance = TRX_PRECISION + 1;
+    BalanceContract.DelegateResourceContract.Builder builder =
+        BalanceContract.DelegateResourceContract.newBuilder()
+        .setLock(true)
+        .setBalance(balance);
+    BalanceContract.DelegateResourceContract.Builder builder2 =
+        BalanceContract.DelegateResourceContract.newBuilder()
+        .setBalance(TRX_PRECISION);
+
+    long expected = DELEGATE_COST_BASE_SIZE + Math.max(
+        builder.build().getSerializedSize() - builder2.build().getSerializedSize(), 0L);
+    long actual = TransactionUtil.estimateConsumeBandWidthSize(balance);
+    Assert.assertEquals(expected, actual);
+  }
+
+  @Test
+  public void estimateConsumeBandWidthSizeCorner() {
+    long balance = Long.MAX_VALUE;
+    BalanceContract.DelegateResourceContract.Builder builder =
+        BalanceContract.DelegateResourceContract.newBuilder()
+        .setLock(true)
+        .setBalance(balance);
+    BalanceContract.DelegateResourceContract.Builder builder2 =
+        BalanceContract.DelegateResourceContract.newBuilder()
+        .setBalance(TRX_PRECISION);
+
+    long expected = DELEGATE_COST_BASE_SIZE + Math.max(
+        builder.build().getSerializedSize() - builder2.build().getSerializedSize(), 0L);
+    long actual = TransactionUtil.estimateConsumeBandWidthSize(balance);
+    Assert.assertEquals(expected, actual);
+  }
+
 }