Indexed enums (#3194)

* Indexed enums

* Type generation updates

* Cleanup createType

* Empty bytes default

* output tests

* Explicit entryIndex

packages/api/src/augment/rpc.ts

@@ -333,11 +333,11 @@ declare module '@polkadot/rpc-core/types.jsonrpc' {
       /**
        * Get offchain local storage under given key and prefix
        **/
-      localStorageGet: AugmentedRpc<(kind: StorageKind | '__UNUSED' | 'PERSISTENT' | 'LOCAL' | number | Uint8Array, key: Bytes | string | Uint8Array) => Observable<Option<Bytes>>>;
+      localStorageGet: AugmentedRpc<(kind: StorageKind | 'PERSISTENT' | 'LOCAL' | number | Uint8Array, key: Bytes | string | Uint8Array) => Observable<Option<Bytes>>>;
       /**
        * Set offchain local storage under given key and prefix
        **/
-      localStorageSet: AugmentedRpc<(kind: StorageKind | '__UNUSED' | 'PERSISTENT' | 'LOCAL' | number | Uint8Array, key: Bytes | string | Uint8Array, value: Bytes | string | Uint8Array) => Observable<Null>>;
+      localStorageSet: AugmentedRpc<(kind: StorageKind | 'PERSISTENT' | 'LOCAL' | number | Uint8Array, key: Bytes | string | Uint8Array, value: Bytes | string | Uint8Array) => Observable<Null>>;
     };
     payment: {
       /**

packages/types-known/src/spec/polkadot.ts

@@ -11,7 +11,15 @@ const sharedTypes = {
   Keys: 'SessionKeys6',
   LookupSource: 'MultiAddress',
   ProxyType: {
-    _enum: ['Any', 'NonTransfer', 'Governance', 'Staking', 'DeprecatedSudoBalances', 'IdentityJudgement', 'CancelProxy']
+    _enum: {
+      Any: 0,
+      NonTransfer: 1,
+      Governance: 2,
+      Staking: 3,
+      // SudoBalances: 4,
+      IdentityJudgement: 5,
+      CancelProxy: 6
+    }
   },
   ValidatorPrefs: 'ValidatorPrefsWithBlocked'
 };

packages/types/src/codec/Enum.spec.ts

@@ -359,6 +359,46 @@ describe('Enum', (): void => {
     });
   });
 
+  describe('indexed enum', (): void => {
+    const Test = Enum.with({
+      A: 5,
+      B: 42,
+      C: 69,
+      D: 255
+    });
+
+    it('handles an indexed C-like enum', (): void => {
+      expect(new Test(registry, 'A').toNumber()).toEqual(5);
+      expect(new Test(registry, 'B').toNumber()).toEqual(42);
+      expect(new Test(registry, 'C').toNumber()).toEqual(69);
+      expect(new Test(registry, 69).toNumber()).toEqual(69);
+      expect(new Test(registry, 'D').toNumber()).toEqual(255);
+    });
+
+    it('creates proper raw structure', (): void => {
+      expect(new Test(registry).toRawType()).toEqual(JSON.stringify({
+        _enum: {
+          A: 5,
+          B: 42,
+          C: 69,
+          D: 255
+        }
+      }));
+    });
+
+    it('has the indexes for the enum', (): void => {
+      expect(new Test(registry).defIndexes).toEqual([5, 42, 69, 255]);
+    });
+
+    it('has the correct outputs', (): void => {
+      const test = new Test(registry, 5);
+
+      expect(test.toU8a()).toEqual(new Uint8Array([5]));
+      expect(test.toHex()).toEqual('0x05');
+      expect(test.toJSON()).toEqual('A');
+    });
+  });
+
   describe('toHex', (): void => {
     it('has a proper hex representation & length', (): void => {
       const Test = Enum.with({

packages/types/src/codec/Enum.ts

@@ -15,56 +15,100 @@ export interface EnumConstructor<T = Codec> {
   new(registry: Registry, value?: any, index?: number): T;
 }
 
-type TypesDef = Record<string, Constructor>;
+interface EntryDef {
+  Type: Constructor;
+  index: number;
+}
+
+type TypesDef = Record<string, EntryDef>;
 
 interface Decoded {
   index: number;
   value: Codec;
 }
 
-function extractDef (registry: Registry, _def: Record<string, keyof InterfaceTypes | Constructor> | string[]): { def: TypesDef; isBasic: boolean } {
-  if (!Array.isArray(_def)) {
-    const def = mapToTypeMap(registry, _def);
-    const isBasic = !Object.values(def).some((type): boolean => type !== Null);
+function isRustEnum (def: Record<string, keyof InterfaceTypes | Constructor> | Record<string, number>): def is Record<string, keyof InterfaceTypes | Constructor> {
+  const defValues = Object.values(def);
 
+  if (defValues.some((v) => isNumber(v))) {
+    assert(defValues.every((v) => isNumber(v) && v >= 0 && v <= 255), 'Invalid number-indexed enum definition');
+
+    return false;
+  }
+
+  return true;
+}
+
+function extractDef (registry: Registry, _def: Record<string, keyof InterfaceTypes | Constructor> | Record<string, number> | string[]): { def: TypesDef; isBasic: boolean; isIndexed: boolean } {
+  if (Array.isArray(_def)) {
     return {
-      def,
-      isBasic
+      def: _def.reduce((def: TypesDef, key, index): TypesDef => {
+        def[key] = { Type: Null, index };
+
+        return def;
+      }, {}),
+      isBasic: true,
+      isIndexed: false
     };
   }
 
-  return {
-    def: _def.reduce((def, key): TypesDef => {
-      def[key] = Null;
+  let isBasic: boolean;
+  let isIndexed: boolean;
+  let def: TypesDef;
 
-      return def;
-    }, {} as TypesDef),
-    isBasic: true
+  if (isRustEnum(_def)) {
+    def = Object
+      .entries(mapToTypeMap(registry, _def))
+      .reduce((def: TypesDef, [key, Type], index): TypesDef => {
+        def[key] = { Type, index };
+
+        return def;
+      }, {});
+    isBasic = !Object.values(def).some(({ Type }) => Type !== Null);
+    isIndexed = false;
+  } else {
+    def = Object
+      .entries(_def)
+      .reduce((def: TypesDef, [key, index]): TypesDef => {
+        def[key] = { Type: Null, index };
+
+        return def;
+      }, {});
+    isBasic = true;
+    isIndexed = true;
+  }
+
+  return {
+    def,
+    isBasic,
+    isIndexed
   };
 }
 
 function createFromValue (registry: Registry, def: TypesDef, index = 0, value?: any): Decoded {
-  const Clazz = Object.values(def)[index];
+  const entry = Object.values(def).find((e) => e.index === index);
 
-  assert(!isUndefined(Clazz), `Unable to create Enum via index ${index}, in ${Object.keys(def).join(', ')}`);
+  assert(!isUndefined(entry), `Unable to create Enum via index ${index}, in ${Object.keys(def).join(', ')}`);
 
   return {
     index,
-    value: value instanceof Clazz ? value : new Clazz(registry, value)
+    value: value instanceof entry.Type
+      ? value
+      : new entry.Type(registry, value)
   };
 }
 
 function decodeFromJSON (registry: Registry, def: TypesDef, key: string, value?: any): Decoded {
   // JSON comes in the form of { "<type (lowercased)>": "<value for type>" }, here we
   // additionally force to lower to ensure forward compat
-  const keys = Object.keys(def).map((k): string => k.toLowerCase());
+  const keys = Object.keys(def).map((k) => k.toLowerCase());
   const keyLower = key.toLowerCase();
   const index = keys.indexOf(keyLower);
 
   assert(index !== -1, `Cannot map Enum JSON, unable to find '${key}' in ${keys.join(', ')}`);
 
   try {
-    return createFromValue(registry, def, index, value);
+    return createFromValue(registry, def, Object.values(def)[index].index, value);
   } catch (error) {
     throw new Error(`Enum(${key}):: ${(error as Error).message}`);
   }
@@ -79,7 +123,10 @@ function decodeFromString (registry: Registry, def: TypesDef, value: string): De
 
 function decodeFromValue (registry: Registry, def: TypesDef, value?: any): Decoded {
   if (isU8a(value)) {
-    return createFromValue(registry, def, value[0], value.subarray(1));
+    // nested, we don't want to match isObject below
+    if (value.length) {
+      return createFromValue(registry, def, value[0], value.subarray(1));
+    }
   } else if (isNumber(value)) {
     return createFromValue(registry, def, value);
   } else if (isString(value)) {
@@ -91,7 +138,7 @@ function decodeFromValue (registry: Registry, def: TypesDef, value?: any): Decod
   }
 
   // Worst-case scenario, return the first with default
-  return createFromValue(registry, def, 0);
+  return createFromValue(registry, def, Object.values(def)[0].index);
 }
 
 function decodeEnum (registry: Registry, def: TypesDef, value?: any, index?: number): Decoded {
@@ -104,7 +151,6 @@ function decodeEnum (registry: Registry, def: TypesDef, value?: any, index?: num
     return createFromValue(registry, def, value.index, value.value);
   }
 
-  // Or else, we just look at `value`
   return decodeFromValue(registry, def, value);
 }
 
@@ -122,27 +168,30 @@ export class Enum implements Codec {
 
   readonly #def: TypesDef;
 
-  readonly #index: number;
+  readonly #entryIndex: number;
 
   readonly #indexes: number[];
 
   readonly #isBasic: boolean;
 
+  readonly #isIndexed: boolean;
+
   readonly #raw: Codec;
 
-  constructor (registry: Registry, def: Record<string, keyof InterfaceTypes | Constructor> | string[], value?: unknown, index?: number) {
+  constructor (registry: Registry, def: Record<string, keyof InterfaceTypes | Constructor> | Record<string, number> | string[], value?: unknown, index?: number) {
     const defInfo = extractDef(registry, def);
     const decoded = decodeEnum(registry, defInfo.def, value, index);
 
     this.registry = registry;
     this.#def = defInfo.def;
     this.#isBasic = defInfo.isBasic;
-    this.#indexes = Object.keys(defInfo.def).map((_, index): number => index);
-    this.#index = this.#indexes.indexOf(decoded.index) || 0;
+    this.#isIndexed = defInfo.isIndexed;
+    this.#indexes = Object.values(defInfo.def).map(({ index }) => index);
+    this.#entryIndex = this.#indexes.indexOf(decoded.index) || 0;
     this.#raw = decoded.value;
   }
 
-  public static with (Types: Record<string, keyof InterfaceTypes | Constructor> | string[]): EnumConstructor<Enum> {
+  public static with (Types: Record<string, keyof InterfaceTypes | Constructor> | Record<string, number> | string[]): EnumConstructor<Enum> {
     return class extends Enum {
       constructor (registry: Registry, value?: unknown, index?: number) {
         super(registry, Types, value, index);
@@ -187,10 +236,10 @@ export class Enum implements Codec {
   }
 
   /**
-   * @description The index of the metadata value
+   * @description The index of the enum value
    */
   public get index (): number {
-    return this.#index;
+    return this.#indexes[this.#entryIndex];
   }
 
   /**
@@ -224,8 +273,8 @@ export class Enum implements Codec {
   /**
    * @description The available keys for this enum
    */
-  public get defEntries (): string[] {
-    return Object.keys(this.#def);
+  public get defIndexes (): number[] {
+    return this.#indexes;
   }
 
   /**
@@ -239,7 +288,7 @@ export class Enum implements Codec {
    * @description The name of the type this enum value represents
    */
   public get type (): string {
-    return this.defKeys[this.#index];
+    return this.defKeys[this.#entryIndex];
   }
 
   /**
@@ -301,16 +350,32 @@ export class Enum implements Codec {
    * @description Returns the number representation for the value
    */
   public toNumber (): number {
-    return this.#index;
+    return this.index;
   }
 
   /**
    * @description Returns a raw struct representation of the enum types
    */
-  protected _toRawStruct (): string[] | Record<string, string> {
-    return this.#isBasic
-      ? this.defKeys
-      : Struct.typesToMap(this.registry, this.#def);
+  protected _toRawStruct (): string[] | Record<string, string | number> {
+    if (this.#isBasic) {
+      return this.#isIndexed
+        ? this.defKeys.reduce((out: Record<string, number>, key, index): Record<string, number> => {
+          out[key] = this.#indexes[index];
+
+          return out;
+        }, {})
+        : this.defKeys;
+    }
+
+    const typeMap = Object
+      .entries(this.#def)
+      .reduce((out: Record<string, Constructor>, [key, { Type }]) => {
+        out[key] = Type;
+
+        return out;
+      }, {});
+
+    return Struct.typesToMap(this.registry, typeMap);
   }
 
   /**
@@ -335,7 +400,7 @@ export class Enum implements Codec {
    */
   public toU8a (isBare?: boolean): Uint8Array {
     return u8aConcat(
-      new Uint8Array(isBare ? [] : [this.#indexes[this.#index]]),
+      new Uint8Array(isBare ? [] : [this.index]),
       this.#raw.toU8a(isBare)
     );
   }

packages/types/src/create/createClass.ts

@@ -87,7 +87,19 @@ const infoMapping: Record<TypeDefInfo, (registry: Registry, value: TypeDef) => C
 
   [TypeDefInfo.DoNotConstruct]: (registry: Registry, value: TypeDef): Constructor => DoNotConstruct.with(value.displayName),
 
-  [TypeDefInfo.Enum]: (registry: Registry, value: TypeDef): Constructor => Enum.with(getTypeClassMap(value)),
+  [TypeDefInfo.Enum]: (registry: Registry, value: TypeDef): Constructor => {
+    const subs = getSubDefArray(value);
+
+    return Enum.with(
+      subs.every(({ type }) => type === 'Null')
+        ? subs.reduce((out: Record<string, number>, { index, name }, count): Record<string, number> => {
+          out[name as string] = index || count;
+
+          return out;
+        }, {})
+        : getTypeClassMap(value)
+    );
+  },
 
   [TypeDefInfo.HashMap]: (registry: Registry, value: TypeDef): Constructor => createHashMap(value, HashMap),
 

packages/types/src/create/createType.spec.ts

@@ -57,12 +57,18 @@ describe('createType', (): void => {
     ).toEqual('B');
   });
 
-  it('allows creation of a Enum (parametrised)', (): void => {
+  it('allows creation of a Enum (with types)', (): void => {
     expect(
       createTypeUnsafe(registry, '{"_enum": {"A": null, "B": "u32", "C": null} }', [1]).toJSON()
     ).toEqual({ B: 0 });
   });
 
+  it('allows creation of a Enum (with indexes)', (): void => {
+    expect(
+      createTypeUnsafe(registry, '{"_enum": {"A": 42, "B": 69, "C": 255} }', [69]).toJSON()
+    ).toEqual('B');
+  });
+
   it('allows creation of a Result', (): void => {
     expect(
       createTypeUnsafe(registry, 'Result<u32,Text>', ['0x011064656667']).toJSON()

packages/types/src/create/getTypeDef.ts

@@ -3,7 +3,7 @@
 
 import type { TypeDef } from './types';
 
-import { assert } from '@polkadot/util';
+import { assert, isNumber } from '@polkadot/util';
 
 import { sanitize } from './sanitize';
 import { TypeDefInfo } from './types';
@@ -16,23 +16,47 @@ interface TypeDefOptions {
 
 const MAX_NESTED = 64;
 
+function isRustEnum (details: Record<string, string> | Record<string, number>): details is Record<string, string> {
+  const values = Object.values(details);
+
+  if (values.some((v) => isNumber(v))) {
+    assert(values.every((v) => isNumber(v) && v >= 0 && v <= 255), 'Invalid number-indexed enum definition');
+
+    return false;
+  }
+
+  return true;
+}
+
 // decode an enum of either of the following forms
 //  { _enum: ['A', 'B', 'C'] }
 //  { _enum: { A: AccountId, B: Balance, C: u32 } }
-function _decodeEnum (value: TypeDef, details: string[] | Record<string, string>, count: number): TypeDef {
+//  { _enum: { A: 1, B: 2 } }
+function _decodeEnum (value: TypeDef, details: string[] | Record<string, string> | Record<string, number>, count: number): TypeDef {
   value.info = TypeDefInfo.Enum;
 
   // not as pretty, but remain compatible with oo7 for both struct and Array types
-  value.sub = Array.isArray(details)
-    ? details.map((name): TypeDef => ({
+  if (Array.isArray(details)) {
+    value.sub = details.map((name, index): TypeDef => ({
+      index,
       info: TypeDefInfo.Plain,
       name,
       type: 'Null'
-    }))
-    : Object.entries(details).map(([name, type]): TypeDef =>
+    }));
+  } else if (isRustEnum(details)) {
+    value.sub = Object.entries(details).map(([name, type], index): TypeDef => ({
       // eslint-disable-next-line @typescript-eslint/no-use-before-define
-      getTypeDef(type || 'Null', { name }, count)
-    );
+      ...getTypeDef(type || 'Null', { name }, count),
+      index
+    }));
+  } else {
+    value.sub = Object.entries(details).map(([name, index]): TypeDef => ({
+      index,
+      info: TypeDefInfo.Plain,
+      name,
+      type: 'Null'
+    }));
+  }
 
   return value;
 }

packages/types/src/interfaces/offchain/definitions.ts

@@ -43,7 +43,10 @@ export default {
   },
   types: {
     StorageKind: {
-      _enum: ['__UNUSED', 'PERSISTENT', 'LOCAL']
+      _enum: {
+        PERSISTENT: 1,
+        LOCAL: 2
+      }
     }
   }
 } as Definitions;

packages/types/src/interfaces/offchain/types.ts

@@ -5,7 +5,6 @@ import type { Enum } from '@polkadot/types';
 
 /** @name StorageKind */
 export interface StorageKind extends Enum {
-  readonly isUnused: boolean;
   readonly isPersistent: boolean;
   readonly isLocal: boolean;
 }

packages/types/src/types/registry.ts

@@ -25,7 +25,10 @@ export interface ChainUpgrades {
   versions: ChainUpgradeVersion[];
 }
 
-export type RegistryTypes = Record<string, Constructor | string | Record<string, string> | { _enum: string[] | Record<string, string | null> } | { _set: Record<string, number> }>;
+export type RegistryTypes =
+  Record<string, Constructor | string | Record<string, string> |
+  { _enum: string[] | Record<string, number> | Record<string, string | null> } |
+  { _set: Record<string, number> }>;
 
 export interface RegistryError {
   documentation: string[];