345 add mapkeymanager (#346)

Co-authored-by: BenRey <[email protected]>

assembly/helpers/__tests__/mapKeyManager.spec.ts

@@ -0,0 +1,94 @@
+import { Address } from '../../std';
+import { resetStorage } from '../../vm-mock';
+import { KeyIncrementer } from '../keyIncrementer';
+import { MapManager } from '../mapKeyManager';
+import { Args } from '@massalabs/as-types';
+
+beforeEach(() => {
+  resetStorage();
+});
+
+describe('MapManager - use cases', () => {
+  test('executes a basic scenario - simple', () => {
+    const addr1 = new Address(
+      'AU12UBnqTHDQALpocVBnkPNy7y5CndUJQTLutaVDDFgMJcq5kQiKq',
+    );
+    const addr2 = new Address(
+      'AU12UBnqTHDQALpocVBnkPNy7y5CndUJQTLutaVDDFgMJcq5kQiKr',
+    );
+    const balance = new MapManager<Address, u64>();
+    balance.set(addr1, 100);
+    balance.set(addr2, 200);
+
+    expect(balance.mustValue(addr1)).toBe(100);
+    expect(balance.mustValue(addr2)).toBe(200);
+    expect<u64>(
+      balance
+        .value(
+          new Address('AU12UBnqTHDQALpocVBnkPNy7y5CndUJQTLutaVDDFgMJcq5kQiKs'),
+        )
+        .unwrapOrDefault(),
+    ).toBe(0);
+  });
+
+  test('executes a basic scenario - complex', () => {
+    const addr1 = new Address(
+      'AU12UBnqTHDQALpocVBnkPNy7y5CndUJQTLutaVDDFgMJcq5kQiKq',
+    );
+    const addr2 = new Address(
+      'AU12UBnqTHDQALpocVBnkPNy7y5CndUJQTLutaVDDFgMJcq5kQiKr',
+    );
+
+    const keyer = new KeyIncrementer<u8>();
+    const balance = new MapManager<Address, u64>(keyer);
+    const allowance = new MapManager<StaticArray<u8>, u64>(keyer);
+
+    balance.set(addr1, 100);
+    balance.set(addr2, 200);
+
+    const allowanceKey = new Args().add(addr1).add(addr2).serialize();
+    allowance.set(allowanceKey, 100);
+
+    expect(balance.mustValue(addr1)).toBe(100);
+    expect(balance.mustValue(addr2)).toBe(200);
+    expect(allowance.mustValue(allowanceKey)).toBe(100);
+    expect<u64>(
+      allowance
+        .value(new Args().add(addr1).add(addr1).serialize())
+        .unwrapOrDefault(),
+    ).toBe(0);
+  });
+});
+
+describe('ConstantManager - unit tests', () => {
+  test('mustValue - key not found', () => {
+    expect(() => {
+      const m = new MapManager<u64, u64>();
+      m.mustValue(1);
+    }).toThrow('Key not found');
+  });
+
+  test('tryValue', () => {
+    const m = new MapManager<Array<u64>, Array<u64>>();
+    m.set([0], [1, 2, 3]);
+    expect<u64[]>(m.value([0]).unwrap()).toStrictEqual([1, 2, 3]);
+    expect(m.value([1]).isErr()).toBe(true);
+    expect(m.value([1]).error).toBe('Key not found');
+    expect<u64[]>(m.value([1]).unwrapOrDefault()).toStrictEqual([]);
+  });
+
+  test('get/set - array of serializable', () => {
+    const addr1 = new Address(
+      'AU12UBnqTHDQALpocVBnkPNy7y5CndUJQTLutaVDDFgMJcq5kQiKq',
+    );
+    const addr2 = new Address(
+      'AU12UBnqTHDQALpocVBnkPNy7y5CndUJQTLutaVDDFgMJcq5kQiKr',
+    );
+
+    const addresses = new MapManager<u64, Array<Address>, u8, Address>();
+    addresses.set(1, [addr1, addr2]);
+
+    expect(addresses.mustValue(1)).toStrictEqual([addr1, addr2]);
+    expect(addresses.value(2).unwrapOrDefault()).toStrictEqual([]);
+  });
+});

assembly/helpers/mapKeyManager.ts

@@ -0,0 +1,68 @@
+import { Args, Result } from '@massalabs/as-types';
+import { KeyIncrementer, KeySequenceManager } from './keyIncrementer';
+import { Storage } from '../std';
+
+/**
+ * Manages a Map in storage.
+ *
+ * @typeParam TValue - The type of the value stored.
+ * @typeParam TKey - The type of the key used to store the value.
+ * @typeParam TKeySequenceManager - The type of the key sequence manager.
+ * @typeParam TArray - When value is an array of serializable, the underlying serializable type.
+ */
+export class MapManager<TKey, TValue, TKeySequenceManager = u8, TArray = void> {
+  public keyPrefix: StaticArray<u8>;
+
+  constructor(
+    manager: KeySequenceManager = new KeyIncrementer<TKeySequenceManager>(0),
+  ) {
+    this.keyPrefix = manager.nextKey();
+  }
+
+  @inline
+  private storageKey(key: TKey): StaticArray<u8> {
+    return new Args(this.keyPrefix).add(key).serialize();
+  }
+
+  /**
+   * Retrieves the value from storage and panics in case of failure.
+   *
+   * @returns the value stored.
+   * @throws if the key is not found in storage.
+   * @throws if the value is not found in storage.
+   */
+  public mustValue(key: TKey): TValue {
+    return new Args(Storage.get(this.storageKey(key)))
+      .next<TValue, TArray>()
+      .unwrap();
+  }
+
+  /**
+   * Retrieves the value from storage and returns a result.
+   *
+   * @returns the value stored wrapped in a result.
+   */
+  public value(key: TKey): Result<TValue> {
+    const storageKey = this.storageKey(key);
+    if (!Storage.has(storageKey)) {
+      if (isManaged<TValue>()) {
+        return new Result<TValue>(instantiate<TValue>(), 'Key not found');
+      }
+      return new Result<TValue>(0, 'Key not found');
+    }
+
+    return new Args(Storage.get(storageKey)).next<TValue, TArray>();
+  }
+
+  /**
+   * Sets the value in storage.
+   *
+   * @param value - The value to store. Must be an A
+   */
+  public set(key: TKey, value: TValue): void {
+    Storage.set(
+      this.storageKey(key),
+      new Args().add<TValue, TArray>(value).serialize(),
+    );
+  }
+}

package.json

@@ -38,7 +38,7 @@
   "devDependencies": {
     "@as-pect/cli": "^8.1.0",
     "@massalabs/as-transformer": "^0.3.2",
-    "@massalabs/as-types": "^2.0.1-dev.20240524080933",
+    "@massalabs/as-types": "^2.0.1-dev",
     "@massalabs/eslint-config": "^0.0.8",
     "@massalabs/prettier-config-as": "^0.0.2",
     "as-bignum": "^0.2.40",