feat(state): implements sizeOfIndex on WalletRepository (#479)

* Add sizeOfIndex * Add initial implementation * Wallet repository clone * Implement on wallet repo clone * Add missing attribute * FIx validator resignation
ArkEcosystem · Mar 13, 2024 · 221d08a · 221d08a
1 parent 79eafc0
commit 221d08a
Show file tree

Hide file tree

Showing 6 changed files with 37 additions and 6 deletions.
diff --git a/packages/contracts/source/contracts/state/wallets.ts b/packages/contracts/source/contracts/state/wallets.ts
@@ -105,6 +105,7 @@ export interface WalletRepository {
 	getIndex(name: string): WalletIndex;
 	setOnIndex(index: string, key: string, wallet: Wallet): void;
 	forgetOnIndex(index: string, key: string): void;
+	sizeOfIndex(index: string): number;
 
 	setDirtyWallet(wallet: Wallet): void;
 	getDirtyWallets(): IterableIterator<Wallet>;

diff --git a/packages/crypto-transaction-validator-resignation/source/handlers/validator-resignation.ts b/packages/crypto-transaction-validator-resignation/source/handlers/validator-resignation.ts
@@ -37,11 +37,10 @@ export class ValidatorResignationTransactionHandler extends Handlers.Transaction
 			throw new Exceptions.WalletAlreadyResignedError();
 		}
 
-		// TODO: use validator count relative to proposed block height
 		const requiredValidatorsCount: number = this.configuration.getMilestone().activeValidators;
-		const currentValidatorsCount: number = walletRepository
-			.allValidators()
-			.filter((w) => w.hasAttribute("validatorResigned") === false).length;
+		const currentValidatorsCount: number =
+			walletRepository.sizeOfIndex(Contracts.State.WalletIndexes.Validators) -
+			walletRepository.sizeOfIndex(Contracts.State.WalletIndexes.Resignations);
 
 		if (currentValidatorsCount - 1 < requiredValidatorsCount) {
 			throw new Exceptions.NotEnoughValidatorsError();

diff --git a/packages/state/source/service-provider.ts b/packages/state/source/service-provider.ts
@@ -50,6 +50,7 @@ export class ServiceProvider extends Providers.ServiceProvider {
 		walletAttributeRepository.set("validatorForgedTotal", Contracts.State.AttributeType.BigNumber);
 		walletAttributeRepository.set("validatorProducedBlocks", Contracts.State.AttributeType.Number);
 		walletAttributeRepository.set("validatorApproval", Contracts.State.AttributeType.Number);
+		walletAttributeRepository.set("validatorResigned", Contracts.State.AttributeType.Boolean);
 
 		this.app
 			.bind(Identifiers.State.Wallet.Factory)

diff --git a/packages/state/source/wallets/wallet-repository-by-sender.ts b/packages/state/source/wallets/wallet-repository-by-sender.ts
@@ -70,8 +70,13 @@ export class WalletRepositoryBySender extends WalletRepository {
 		return this.#blockchainWalletRepository.hasByIndex(index, key);
 	}
 
+	// TODO: set, forget & sizeOfIndex methods are not implemented fully
 	public setOnIndex(index: string, key: string, wallet: Contracts.State.Wallet): void {}
 
+	public sizeOfIndex(index: string): number {
+		return this.#blockchainWalletRepository.sizeOfIndex(index);
+	}
+
 	#cloneWallet(address: string): Contracts.State.Wallet {
 		return this.#blockchainWalletRepository.findByAddress(address).clone(this);
 	}

diff --git a/packages/state/source/wallets/wallet-repository-clone.ts b/packages/state/source/wallets/wallet-repository-clone.ts
@@ -8,6 +8,7 @@ export class WalletRepositoryClone extends WalletRepository implements Contracts
 	#originalWalletRepository!: WalletRepository;
 
 	readonly #forgetIndexes: Record<string, Set<string>> = {};
+	readonly #indexOffset: Record<string, number> = {};
 	readonly #dirtyWallets = new Set<Contracts.State.Wallet>();
 
 	@postConstruct()
@@ -16,6 +17,7 @@ export class WalletRepositoryClone extends WalletRepository implements Contracts
 
 		for (const name of this.indexSet.all()) {
 			this.#forgetIndexes[name] = new Set();
+			this.#indexOffset[name] = 0;
 		}
 	}
 
@@ -83,17 +85,36 @@ export class WalletRepositoryClone extends WalletRepository implements Contracts
 	}
 
 	public setOnIndex(index: string, key: string, wallet: Contracts.State.Wallet): void {
+		// New key registration
+		if (!this.#originalWalletRepository.hasByIndex(index, key) && !this.getIndex(index).has(key)) {
+			this.#indexOffset[index]++;
+		}
+
+		// Remove from forget indexes
+		if (this.#getForgetSet(index).has(key)) {
+			this.#getForgetSet(index).delete(key);
+			this.#indexOffset[index]++;
+		}
+
 		this.getIndex(index).set(key, wallet);
-		this.#getForgetSet(index).delete(key);
 	}
 
 	public forgetOnIndex(index: string, key: string): void {
-		if (this.getIndex(index).has(key) || this.#originalWalletRepository.getIndex(index).has(key)) {
+		if (
+			!this.#getForgetSet(index).has(key) &&
+			(this.getIndex(index).has(key) || this.#originalWalletRepository.getIndex(index).has(key))
+		) {
 			this.getIndex(index).forget(key);
 			this.#getForgetSet(index).add(key);
+
+			this.#indexOffset[index]--;
 		}
 	}
 
+	public sizeOfIndex(index: string): number {
+		return this.#originalWalletRepository.sizeOfIndex(index) + this.#indexOffset[index];
+	}
+
 	public getDirtyWallets(): IterableIterator<Contracts.State.Wallet> {
 		return this.#dirtyWallets.values();
 	}

diff --git a/packages/state/source/wallets/wallet-repository.ts b/packages/state/source/wallets/wallet-repository.ts
@@ -95,6 +95,10 @@ export class WalletRepository implements Contracts.State.WalletRepository {
 		this.getIndex(index).forget(key);
 	}
 
+	public sizeOfIndex(index: string): number {
+		return this.getIndex(index).size();
+	}
+
 	public setDirtyWallet(wallet: Contracts.State.Wallet): void {}
 
 	public getDirtyWallets(): IterableIterator<Contracts.State.Wallet> {