HIP-584 Historical: Fix historical NFT/Fungible token balance (#7430)

This PR includes the following fixes:
- get historical fungible token balance from `TokenBalanceRepository`
- get historical NFT balance from `NftRepository`. Historical NFT balance should be obtained from `account.getOwnedNfts()`, as explained in the code comments
- fix `findHistoricalTokenBalanceUpToTimestamp` in edge case where `balance_snapshot` and/or `base` queries do not match anything in the db. Currently reproducible in acceptance tests:
1. create account
2. make fungible token transfer +50
3. make fungible token transfer +50
4. make historical call with timestamp between the token transfers, expected balance is 50, but returns 0

I have observed that the state in db after these operations is the following:
`account_balance` - no entry for this accountId
`token_balance` - no entry for this accountId
`token_transfers` - 2 entries for this accountId, one for first transfer (50) and one for the second transfer (50)

`balance_snapshot` and `base` match nothing because there is no entry for this accountId before the given timestamp.
This causes the statement `tt.consensus_timestamp > s.consensus_timestamp` to be false and cannot match any token transfers, always returning 0 due to the coalesce at the end of the query.
NB. `token_balance` and `account_balance` entries get persisted in ~8 minute intervals so it is not certain that there is entry in the db after the first transfer.

Current solution:
If `balance_snapshot` and/or `base` do not match anything, find token transfers that occured after the account creation and before the given timestamp:
```sql
tt.consensus_timestamp >= coalesce((select consensus_timestamp from base), accountCreatedTimestamp)
```
Changed `>` to `>=` ensures the inclusion of the token transfer that occurred when the account was created.

---------

Signed-off-by: Ivan Ivanov <[email protected]>

hedera-mirror-rest/__tests__/tokens.test.js

@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2020-2023 Hedera Hashgraph, LLC
+ * Copyright (C) 2020-2024 Hedera Hashgraph, LLC
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.

hedera-mirror-rest/tokens.js

@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2020-2023 Hedera Hashgraph, LLC
+ * Copyright (C) 2020-2024 Hedera Hashgraph, LLC
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.

hedera-mirror-test/k6/src/lib/common.js

@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2022-2023 Hedera Hashgraph, LLC
+ * Copyright (C) 2022-2024 Hedera Hashgraph, LLC
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -227,7 +227,9 @@ function markdownReport(data, includeUrlColumn, funcs, scenarios, getUrlFuncs =
   const header = `| Scenario ${
     includeUrlColumn ? '| URL' : ''
   } | VUS | Pass% | RPS | Pass RPS | Avg. Req Duration | Skipped? | Comment |
-|----------${includeUrlColumn ? '|----------' : ''}|-----|-------|-----|----------|-------------------|--------|---------|`;
+|----------${
+    includeUrlColumn ? '|----------' : ''
+  }|-----|-------|-----|----------|-------------------|--------|---------|`;
 
   // collect the metrics
   const {setup_data: availableParams} = data;

hedera-mirror-web3/src/main/java/com/hedera/mirror/web3/evm/store/accessor/AccountDatabaseAccessor.java

@@ -18,6 +18,7 @@
 
 import static com.hedera.mirror.common.domain.entity.EntityType.ACCOUNT;
 import static com.hedera.mirror.common.domain.entity.EntityType.CONTRACT;
+import static com.hedera.mirror.web3.evm.store.accessor.TokenRelationshipDatabaseAccessor.ZERO_BALANCE;
 import static com.hedera.services.utils.EntityIdUtils.idFromEntityId;
 import static com.hedera.services.utils.MiscUtils.asFcKeyUnchecked;
 
@@ -63,7 +64,6 @@ public class AccountDatabaseAccessor extends DatabaseAccessor<Object, Account> {
     private static final BinaryOperator<Long> NO_DUPLICATE_MERGE_FUNCTION = (v1, v2) -> {
         throw new IllegalStateException(String.format("Duplicate key for values %s and %s", v1, v2));
     };
-    private static final Optional<Long> ZERO_BALANCE = Optional.of(0L);
 
     private final EntityDatabaseAccessor entityDatabaseAccessor;
     private final NftAllowanceRepository nftAllowanceRepository;

hedera-mirror-web3/src/main/java/com/hedera/mirror/web3/evm/store/accessor/TokenRelationshipDatabaseAccessor.java

@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2023 Hedera Hashgraph, LLC
+ * Copyright (C) 2023-2024 Hedera Hashgraph, LLC
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -22,7 +22,10 @@
 import com.hedera.mirror.common.domain.token.TokenKycStatusEnum;
 import com.hedera.mirror.web3.evm.store.DatabaseBackedStateFrame.DatabaseAccessIncorrectKeyTypeException;
 import com.hedera.mirror.web3.evm.store.accessor.model.TokenRelationshipKey;
+import com.hedera.mirror.web3.repository.NftRepository;
 import com.hedera.mirror.web3.repository.TokenAccountRepository;
+import com.hedera.mirror.web3.repository.TokenBalanceRepository;
+import com.hedera.node.app.service.evm.store.tokens.TokenType;
 import com.hedera.services.store.models.Account;
 import com.hedera.services.store.models.Token;
 import com.hedera.services.store.models.TokenRelationship;
@@ -39,6 +42,9 @@ public class TokenRelationshipDatabaseAccessor extends DatabaseAccessor<Object,
     private final TokenDatabaseAccessor tokenDatabaseAccessor;
     private final AccountDatabaseAccessor accountDatabaseAccessor;
     private final TokenAccountRepository tokenAccountRepository;
+    private final TokenBalanceRepository tokenBalanceRepository;
+    private final NftRepository nftRepository;
+    static final Optional<Long> ZERO_BALANCE = Optional.of(0L);
 
     @Override
     public @NonNull Optional<TokenRelationship> get(@NonNull Object key, final Optional<Long> timestamp) {
@@ -50,7 +56,7 @@ public class TokenRelationshipDatabaseAccessor extends DatabaseAccessor<Object,
                                     .map(tokenAccount -> new TokenRelationship(
                                             token,
                                             account,
-                                            tokenAccount.getBalance(),
+                                            getBalance(account, token, tokenAccount, timestamp),
                                             TokenFreezeStatusEnum.FROZEN == tokenAccount.getFreezeStatus(),
                                             TokenKycStatusEnum.REVOKED != tokenAccount.getKycStatus(),
                                             false,
@@ -62,6 +68,64 @@ public class TokenRelationshipDatabaseAccessor extends DatabaseAccessor<Object,
                 .formatted(TokenRelationship.class.getTypeName(), key.getClass().getTypeName()));
     }
 
+    /**
+     * Determines fungible or NFT balance based on block context.
+     */
+    private Long getBalance(
+            final Account account, final Token token, final TokenAccount tokenAccount, final Optional<Long> timestamp) {
+        if (token.getType().equals(TokenType.NON_FUNGIBLE_UNIQUE)) {
+            return getNftBalance(tokenAccount, timestamp, account.getCreatedTimestamp());
+        }
+        return getFungibleBalance(tokenAccount, timestamp, account.getCreatedTimestamp());
+    }
+
+    /**
+     * NFT Balance Explanation:
+     * Non-historical Call:
+     * The balance is obtained from `tokenAccount.getBalance()`.
+     * Historical Call:
+     * In historical block queries, as the `token_account` and `token_balance` tables lack historical state for NFT balances,
+     * the NFT balance is retrieved from `NftRepository.nftBalanceByAccountIdTokenIdAndTimestamp`
+     */
+    private Long getNftBalance(
+            final TokenAccount tokenAccount, final Optional<Long> timestamp, long accountCreatedTimestamp) {
+        return timestamp
+                .map(t -> {
+                    if (t >= accountCreatedTimestamp) {
+                        return nftRepository.nftBalanceByAccountIdTokenIdAndTimestamp(
+                                tokenAccount.getAccountId(), tokenAccount.getTokenId(), t);
+                    } else {
+                        return ZERO_BALANCE;
+                    }
+                })
+                .orElseGet(() -> Optional.of(tokenAccount.getBalance()))
+                .orElse(0L);
+    }
+
+    /**
+     * Fungible Token Balance Explanation:
+     * Non-historical Call:
+     * The balance is obtained from `tokenAccount.getBalance()`.
+     * Historical Call:
+     * In historical block queries, since the `token_account` table lacks historical state for fungible balances,
+     * the fungible balance is determined from the `token_balance` table using the `findHistoricalTokenBalanceUpToTimestamp` query.
+     * If the entity creation is after the passed timestamp - return 0L (the entity was not created)
+     */
+    private Long getFungibleBalance(
+            final TokenAccount tokenAccount, final Optional<Long> timestamp, long accountCreatedTimestamp) {
+        return timestamp
+                .map(t -> {
+                    if (t >= accountCreatedTimestamp) {
+                        return tokenBalanceRepository.findHistoricalTokenBalanceUpToTimestamp(
+                                tokenAccount.getTokenId(), tokenAccount.getAccountId(), t);
+                    } else {
+                        return ZERO_BALANCE;
+                    }
+                })
+                .orElseGet(() -> Optional.of(tokenAccount.getBalance()))
+                .orElse(0L);
+    }
+
     private Optional<Account> findAccount(Address address, final Optional<Long> timestamp) {
         return accountDatabaseAccessor.get(address, timestamp);
     }

hedera-mirror-web3/src/main/java/com/hedera/mirror/web3/repository/NftRepository.java

@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2023 Hedera Hashgraph, LLC
+ * Copyright (C) 2023-2024 Hedera Hashgraph, LLC
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -110,7 +110,7 @@ select count(*)
                             select token_id
                             from nft
                             where account_id = :accountId
-                                and lower(timestamp_range) <= :blockTimestamp
+                                and timestamp_range @> :blockTimestamp
                                 and deleted is not true
                         )
                         union all
@@ -127,4 +127,44 @@ and lower(timestamp_range) <= :blockTimestamp
                     """,
             nativeQuery = true)
     long countByAccountIdAndTimestampNotDeleted(long accountId, long blockTimestamp);
+
+    /**
+     * Retrieves the most recent state of nft balance
+     * by accountId and tokenId up to a given block timestamp.
+     * The method considers both the current state of the token account and its historical states
+     * and returns the one that was valid just before or equal to the provided block timestamp.
+     *
+     * @param accountId the ID of the account
+     * @param tokenId the ID of the nft
+     * @param blockTimestamp  the block timestamp used to filter the results.
+     * @return the number of nft serial numbers the accountId owns at the specified timestamp for specific tokenId.
+     */
+    @Query(
+            value =
+                    """
+            select count(*)
+            from (
+                (
+                    select token_id
+                    from nft
+                    where account_id = :accountId
+                        and token_id = :tokenId
+                        and timestamp_range @> :blockTimestamp
+                        and deleted is not true
+                )
+                union all
+                (
+                    select token_id
+                    from nft_history
+                    where account_id = :accountId
+                        and token_id = :tokenId
+                        and timestamp_range @> :blockTimestamp
+                        and deleted is not true
+                )
+                ) as n
+            join entity e on e.id = n.token_id
+            where (e.deleted is not true or lower(e.timestamp_range) > :blockTimestamp)
+            """,
+            nativeQuery = true)
+    Optional<Long> nftBalanceByAccountIdTokenIdAndTimestamp(long accountId, long tokenId, long blockTimestamp);
 }

hedera-mirror-web3/src/main/java/com/hedera/mirror/web3/repository/TokenBalanceRepository.java

@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2023 Hedera Hashgraph, LLC
+ * Copyright (C) 2023-2024 Hedera Hashgraph, LLC
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -54,6 +54,8 @@ public interface TokenBalanceRepository extends CrudRepository<TokenBalance, Tok
      * found at a timestamp less than the given block timestamp. If no token_balance is found for the given token_id,
      * account_id, and consensus timestamp, a balance of 0 will be returned.
      *
+     * For more information please refer to `AccountBalanceRepository.findHistoricalAccountBalanceUpToTimestamp`
+     *
      * @param tokenId         the ID of the token.
      * @param accountId       the ID of the account.
      * @param blockTimestamp  the block timestamp used to filter the results.
@@ -64,27 +66,31 @@ public interface TokenBalanceRepository extends CrudRepository<TokenBalance, Tok
     @Query(
             value =
                     """
-                    with balance_snapshot as (
-                      select consensus_timestamp
-                      from account_balance
-                      where account_id = 2 and consensus_timestamp <= ?3
-                      order by consensus_timestamp desc
-                      limit 1
+                    with balance_timestamp as (
+                    select consensus_timestamp
+                    from account_balance
+                    where account_id = 2 and
+                        consensus_timestamp > ?3 - 2678400000000000 and consensus_timestamp <= ?3
+                    order by consensus_timestamp desc
+                    limit 1
                     ), base as (
-                        select balance
-                        from token_balance as tb, balance_snapshot as s
+                        select tb.balance
+                        from token_balance as tb, balance_timestamp as bt
                         where
-                          tb.consensus_timestamp = s.consensus_timestamp and
-                          token_id = ?1 and
-                          account_id = ?2
+                            token_id = ?1 and
+                            account_id = ?2 and
+                            tb.consensus_timestamp > bt.consensus_timestamp - 2678400000000000 and
+                            tb.consensus_timestamp <= bt.consensus_timestamp
+                        order by tb.consensus_timestamp desc
+                        limit 1
                     ), change as (
                         select sum(amount) as amount
-                        from token_transfer as tt, balance_snapshot as s
+                        from token_transfer as tt
                         where
-                          token_id = ?1 and
-                          account_id = ?2 and
-                          tt.consensus_timestamp > s.consensus_timestamp and
-                          tt.consensus_timestamp <= ?3
+                            token_id = ?1 and
+                            account_id = ?2 and
+                            tt.consensus_timestamp > coalesce((select consensus_timestamp from balance_timestamp), 0) and
+                            tt.consensus_timestamp <= ?3
                     )
                     select coalesce((select balance from base), 0) + coalesce((select amount from change), 0)
                     """,

hedera-mirror-web3/src/test/java/com/hedera/mirror/web3/evm/store/StoreImplTest.java

@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2023 Hedera Hashgraph, LLC
+ * Copyright (C) 2023-2024 Hedera Hashgraph, LLC
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -29,6 +29,7 @@
 import com.hedera.mirror.common.domain.token.Nft;
 import com.hedera.mirror.common.domain.token.Token;
 import com.hedera.mirror.common.domain.token.TokenAccount;
+import com.hedera.mirror.common.domain.token.TokenTypeEnum;
 import com.hedera.mirror.web3.ContextExtension;
 import com.hedera.mirror.web3.evm.store.Store.OnMissing;
 import com.hedera.mirror.web3.evm.store.accessor.AccountDatabaseAccessor;
@@ -46,6 +47,7 @@
 import com.hedera.mirror.web3.repository.NftRepository;
 import com.hedera.mirror.web3.repository.TokenAccountRepository;
 import com.hedera.mirror.web3.repository.TokenAllowanceRepository;
+import com.hedera.mirror.web3.repository.TokenBalanceRepository;
 import com.hedera.mirror.web3.repository.TokenRepository;
 import com.hedera.mirror.web3.repository.projections.TokenAccountAssociationsCount;
 import com.hedera.node.app.service.evm.exceptions.InvalidTransactionException;
@@ -115,6 +117,9 @@ public boolean getIsPositiveBalance() {
     @Mock
     private TokenAccountRepository tokenAccountRepository;
 
+    @Mock
+    private TokenBalanceRepository tokenBalanceRepository;
+
     @Mock
     private CustomFeeDatabaseAccessor customFeeDatabaseAccessor;
 
@@ -163,7 +168,11 @@ void setup() {
         final var tokenDatabaseAccessor = new TokenDatabaseAccessor(
                 tokenRepository, entityDatabaseAccessor, entityRepository, customFeeDatabaseAccessor);
         final var tokenRelationshipDatabaseAccessor = new TokenRelationshipDatabaseAccessor(
-                tokenDatabaseAccessor, accountDatabaseAccessor, tokenAccountRepository);
+                tokenDatabaseAccessor,
+                accountDatabaseAccessor,
+                tokenAccountRepository,
+                tokenBalanceRepository,
+                nftRepository);
         final var uniqueTokenDatabaseAccessor = new UniqueTokenDatabaseAccessor(nftRepository);
         final var entityDatabaseAccessor = new EntityDatabaseAccessor(entityRepository);
         final List<DatabaseAccessor<Object, ?>> accessors = List.of(
@@ -240,6 +249,7 @@ void getTokenRelationshipWithoutThrow() {
         when(accountModel.getType()).thenReturn(EntityType.ACCOUNT);
         when(tokenAccountRepository.findById(any())).thenReturn(Optional.of(tokenAccount));
         when(tokenAccount.getAssociated()).thenReturn(Boolean.TRUE);
+        when(token.getType()).thenReturn(TokenTypeEnum.FUNGIBLE_COMMON);
         final var tokenRelationship = subject.getTokenRelationship(
                 new TokenRelationshipKey(TOKEN_ADDRESS, ACCOUNT_ADDRESS), OnMissing.DONT_THROW);
         assertThat(tokenRelationship.getAccount().getId()).isEqualTo(new Id(0, 0, 12));
@@ -385,6 +395,7 @@ void getHistoricalTimestamp() {
     }
 
     private void setupTokenAndAccount() {
+        when(token.getType()).thenReturn(TokenTypeEnum.FUNGIBLE_COMMON);
         when(entityDatabaseAccessor.get(TOKEN_ADDRESS, Optional.empty())).thenReturn(Optional.of(tokenModel));
         when(tokenModel.getId()).thenReturn(6L);
         when(tokenModel.getNum()).thenReturn(6L);