BTreeSet and BTreeMap - encoding and serialization fixes (#3789)

* BTreeSet and BTreeMap - encoding and serialization fixes

* Fix linter issues

* BTreeSet/BTreeMap: Support for Struct and Enum sorting

* Cleanup sorting function

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

@@ -2,15 +2,52 @@
 // SPDX-License-Identifier: Apache-2.0
 
 import { TypeRegistry } from '../create';
-import { Text, U32 } from '../primitive';
-import { BTreeMap, Struct } from '.';
+import { I32, Text, U32 } from '../primitive';
+import { Constructor } from '../types';
+import { ITuple } from '../types/interfaces';
+import { BTreeMap, Enum, Struct, Tuple } from '.';
 
 const registry = new TypeRegistry();
 
+class U32TextTuple extends (Tuple.with([U32, Text]) as unknown as Constructor<ITuple<[U32, Text]>>) {}
+// Reason: We purposefully want `text` to be the first key of the struct and take priority during sorting
+// eslint-disable-next-line sort-keys
+class MockStruct extends Struct.with({ text: Text, int: I32 }) {}
+class MockEnum extends Enum.with({
+  Key1: MockStruct,
+  Key2: MockStruct,
+  Key3: U32TextTuple
+}) {}
+
 const mockU32TextMap = new Map<Text, U32>();
+const mockU32TupleMap = new Map<ITuple<[U32, Text]>, U32>();
+const mockU32I32Map = new Map<I32, U32>();
+const mockU32StructMap = new Map<MockStruct, U32>();
+const mockU32EnumMap = new Map<MockEnum, U32>();
 
 mockU32TextMap.set(new Text(registry, 'bazzing'), new U32(registry, 69));
 
+mockU32TupleMap.set((new U32TextTuple(registry, [2, 'ba'])), new U32(registry, 42));
+mockU32TupleMap.set((new U32TextTuple(registry, [2, 'b'])), new U32(registry, 7));
+mockU32TupleMap.set((new U32TextTuple(registry, [1, 'baz'])), new U32(registry, 13));
+
+mockU32I32Map.set(new I32(registry, 255), new U32(registry, 69));
+mockU32I32Map.set(new I32(registry, -255), new U32(registry, 42));
+mockU32I32Map.set(new I32(registry, 1000), new U32(registry, 7));
+mockU32I32Map.set(new I32(registry, -1000), new U32(registry, 25));
+mockU32I32Map.set(new I32(registry, 0), new U32(registry, 13));
+
+mockU32StructMap.set(new MockStruct(registry, { int: 1, text: 'b' }), new U32(registry, 42));
+mockU32StructMap.set(new MockStruct(registry, { int: -1, text: 'b' }), new U32(registry, 7));
+mockU32StructMap.set(new MockStruct(registry, { int: -1, text: 'ba' }), new U32(registry, 25));
+mockU32StructMap.set(new MockStruct(registry, { int: -2, text: 'baz' }), new U32(registry, 13));
+
+mockU32EnumMap.set(new MockEnum(registry, { Key3: new U32TextTuple(registry, [2, 'ba']) }), new U32(registry, 13));
+mockU32EnumMap.set(new MockEnum(registry, { Key3: new U32TextTuple(registry, [2, 'b']) }), new U32(registry, 42));
+mockU32EnumMap.set(new MockEnum(registry, { Key2: new MockStruct(registry, { int: -1, text: 'b' }) }), new U32(registry, 7));
+mockU32EnumMap.set(new MockEnum(registry, { Key1: new MockStruct(registry, { int: 1, text: 'b' }) }), new U32(registry, 25));
+mockU32EnumMap.set(new MockEnum(registry, { Key1: new MockStruct(registry, { int: -1, text: 'b' }) }), new U32(registry, 69));
+
 describe('BTreeMap', (): void => {
   it('decodes null', (): void => {
     expect(
@@ -54,6 +91,97 @@ describe('BTreeMap', (): void => {
     ).toEqual(13);
   });
 
+  it('correctly sorts simple keys', (): void => {
+    expect(
+      Array.from(new (BTreeMap.with(I32, U32))(registry, mockU32I32Map).keys()).map((k) => k.toNumber())
+    ).toEqual([-1000, -255, 0, 255, 1000]);
+  });
+
+  it('correctly sorts tuple keys', (): void => {
+    expect(
+      Array.from(new (BTreeMap.with(U32TextTuple, U32))(registry, mockU32TupleMap).keys()).map((k) => k.toJSON())
+    ).toEqual([[1, 'baz'], [2, 'b'], [2, 'ba']]);
+  });
+
+  it('correctly sorts struct keys', (): void => {
+    expect(
+      Array.from(new (BTreeMap.with(MockStruct, U32))(registry, mockU32StructMap).keys()).map((k) => k.toJSON())
+    ).toEqual([
+      { int: -1, text: 'b' },
+      { int: 1, text: 'b' },
+      { int: -1, text: 'ba' },
+      { int: -2, text: 'baz' }
+    ]);
+  });
+
+  it('correctly sorts enum keys', (): void => {
+    expect(
+      Array.from(new (BTreeMap.with(MockEnum, U32))(registry, mockU32EnumMap).keys()).map((k) => k.toJSON())
+    ).toEqual([
+      { key1: { int: -1, text: 'b' } },
+      { key1: { int: 1, text: 'b' } },
+      { key2: { int: -1, text: 'b' } },
+      { key3: [2, 'b'] },
+      { key3: [2, 'ba'] }
+    ]);
+  });
+
+  it('correctly serializes/deserializes to/from json with numeric keys', (): void => {
+    expect(
+      new (BTreeMap.with(I32, U32))(
+        registry,
+        new (BTreeMap.with(I32, U32))(registry, mockU32I32Map).toJSON()
+      ).toJSON()
+    ).toEqual({ '-1000': 25, '-255': 42, 0: 13, 1000: 7, 255: 69 });
+  });
+
+  it('correctly serializes/deserializes to/from json with text keys', (): void => {
+    expect(
+      new (BTreeMap.with(Text, U32))(
+        registry,
+        new (BTreeMap.with(Text, U32))(registry, mockU32TextMap).toJSON()
+      ).toJSON()
+    ).toEqual({ bazzing: 69 });
+  });
+
+  it('correctly serializes/deserializes to/from json with tuple keys', (): void => {
+    expect(
+      new (BTreeMap.with(U32TextTuple, U32))(
+        registry,
+        new (BTreeMap.with(U32TextTuple, U32))(registry, mockU32TupleMap).toJSON()
+      ).toJSON()
+    ).toEqual({ '[1,"baz"]': 13, '[2,"b"]': 7, '[2,"ba"]': 42 });
+  });
+
+  it('correctly serializes/deserializes to/from json with struct keys', (): void => {
+    expect(
+      new (BTreeMap.with(MockStruct, U32))(
+        registry,
+        new (BTreeMap.with(MockStruct, U32))(registry, mockU32StructMap).toJSON()
+      ).toJSON()
+    ).toEqual({
+      '{"text":"b","int":-1}': 7,
+      '{"text":"b","int":1}': 42,
+      '{"text":"ba","int":-1}': 25,
+      '{"text":"baz","int":-2}': 13
+    });
+  });
+
+  it('correctly serializes/deserializes to/from json with enum keys', (): void => {
+    expect(
+      new (BTreeMap.with(MockEnum, U32))(
+        registry,
+        new (BTreeMap.with(MockEnum, U32))(registry, mockU32EnumMap).toJSON()
+      ).toJSON()
+    ).toEqual({
+      '{"key1":{"text":"b","int":-1}}': 69,
+      '{"key1":{"text":"b","int":1}}': 25,
+      '{"key2":{"text":"b","int":-1}}': 7,
+      '{"key3":[2,"b"]}': 42,
+      '{"key3":[2,"ba"]}': 13
+    });
+  });
+
   it('generates sane toRawTypes', (): void => {
     expect(new (BTreeMap.with(Text, U32))(registry).toRawType()).toBe('BTreeMap<Text,u32>');
     expect(new (BTreeMap.with(Text, Text))(registry).toRawType()).toBe('BTreeMap<Text,Text>');

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

@@ -4,11 +4,24 @@
 import type { CodecTo } from '../types';
 
 import { TypeRegistry } from '../create';
-import { Text, U32 } from '../primitive';
-import { BTreeSet, Struct } from '.';
+import { I32, Text, U32 } from '../primitive';
+import { Constructor } from '../types';
+import { ITuple } from '../types/interfaces';
+import { Enum } from './Enum';
+import { BTreeSet, Struct, Tuple } from '.';
 
 const registry = new TypeRegistry();
 
+class U32TextTuple extends (Tuple.with([U32, Text]) as unknown as Constructor<ITuple<[U32, Text]>>) {}
+// Reason: We purposefully want `text` to be the first key of the struct and take priority during sorting
+// eslint-disable-next-line sort-keys
+class MockStruct extends Struct.with({ text: Text, int: I32 }) {}
+class MockEnum extends Enum.with({
+  Key1: MockStruct,
+  Key2: MockStruct,
+  Key3: U32TextTuple
+}) {}
+
 const mockU32Set = new Set<U32>();
 
 mockU32Set.add(new U32(registry, 2));
@@ -21,6 +34,34 @@ const mockU32SetObject = [2, 24, 30, 80];
 const mockU32SetHexString = '0x1002000000180000001e00000050000000';
 const mockU32SetUint8Array = Uint8Array.from([16, 2, 0, 0, 0, 24, 0, 0, 0, 30, 0, 0, 0, 80, 0, 0, 0]);
 
+const mockI32SetObj = [1000, 0, 255, -255, -1000];
+const mockTextSetObj = [
+  new Text(registry, 'baz'),
+  new Text(registry, 'b'),
+  new Text(registry, 'bb'),
+  new Text(registry, 'ba'),
+  new Text(registry, 'c')
+];
+const mockTupleSetObj = [
+  new U32TextTuple(registry, [2, 'ba']),
+  new U32TextTuple(registry, [2, 'bb']),
+  new U32TextTuple(registry, [2, 'b']),
+  new U32TextTuple(registry, [1, 'baz'])
+];
+const mockStructSetObj = [
+  new MockStruct(registry, { int: 1, text: 'b' }),
+  new MockStruct(registry, { int: -1, text: 'b' }),
+  new MockStruct(registry, { int: -1, text: 'ba' }),
+  new MockStruct(registry, { int: -2, text: 'baz' })
+];
+const mockEnumSetObj = [
+  new MockEnum(registry, { Key3: new U32TextTuple(registry, [2, 'ba']) }),
+  new MockEnum(registry, { Key3: new U32TextTuple(registry, [2, 'b']) }),
+  new MockEnum(registry, { Key2: new MockStruct(registry, { int: -1, text: 'b' }) }),
+  new MockEnum(registry, { Key1: new MockStruct(registry, { int: 1, text: 'b' }) }),
+  new MockEnum(registry, { Key1: new MockStruct(registry, { int: -1, text: 'b' }) })
+];
+
 describe('BTreeSet', (): void => {
   describe('decoding', (): void => {
     const testDecode = (type: string, input: any, output: string): void =>
@@ -94,6 +135,47 @@ describe('BTreeSet', (): void => {
     ).toEqual(17);
   });
 
+  it('correctly sorts numeric values', (): void => {
+    expect(
+      Array.from(new (BTreeSet.with(I32))(registry, mockI32SetObj)).map((k) => k.toNumber())
+    ).toEqual([-1000, -255, 0, 255, 1000]);
+  });
+
+  it('correctly sorts text values', (): void => {
+    expect(
+      Array.from(new (BTreeSet.with(Text))(registry, mockTextSetObj)).map((k) => k.toString())
+    ).toEqual(['b', 'ba', 'baz', 'bb', 'c']);
+  });
+
+  it('correctly sorts complex tuple values', (): void => {
+    expect(
+      Array.from(new (BTreeSet.with(U32TextTuple))(registry, mockTupleSetObj)).map((k) => k.toJSON())
+    ).toEqual([[1, 'baz'], [2, 'b'], [2, 'ba'], [2, 'bb']]);
+  });
+
+  it('correctly sorts complex struct values', (): void => {
+    expect(
+      Array.from(new (BTreeSet.with(MockStruct))(registry, mockStructSetObj)).map((k) => k.toJSON())
+    ).toEqual([
+      { int: -1, text: 'b' },
+      { int: 1, text: 'b' },
+      { int: -1, text: 'ba' },
+      { int: -2, text: 'baz' }
+    ]);
+  });
+
+  it('correctly sorts complex enum values', (): void => {
+    expect(
+      Array.from(new (BTreeSet.with(MockEnum))(registry, mockEnumSetObj)).map((k) => k.toJSON())
+    ).toEqual([
+      { key1: { int: -1, text: 'b' } },
+      { key1: { int: 1, text: 'b' } },
+      { key2: { int: -1, text: 'b' } },
+      { key3: [2, 'b'] },
+      { key3: [2, 'ba'] }
+    ]);
+  });
+
   it('generates sane toRawTypes', (): void => {
     expect(new (BTreeSet.with(U32))(registry).toRawType()).toBe('BTreeSet<u32>');
     expect(new (BTreeSet.with(Text))(registry).toRawType()).toBe('BTreeSet<Text>');

packages/types/src/codec/BTreeSet.ts

@@ -6,7 +6,7 @@ import type { AnyJson, Codec, Constructor, InterfaceTypes, Registry } from '../t
 
 import { compactFromU8a, compactToU8a, isHex, isU8a, logger, stringify, u8aConcat, u8aToHex, u8aToU8a } from '@polkadot/util';
 
-import { compareSet, decodeU8a, typeToConstructor } from './utils';
+import { compareSet, decodeU8a, sortSet, typeToConstructor } from './utils';
 
 const l = logger('BTreeSet');
 
@@ -84,7 +84,7 @@ export class BTreeSet<V extends Codec = Codec> extends Set<V> implements Codec {
   readonly #ValClass: Constructor<V>;
 
   constructor (registry: Registry, valType: Constructor<V> | keyof InterfaceTypes, rawValue?: Uint8Array | string | string[] | Set<any>) {
-    super(decodeSet(registry, valType, rawValue));
+    super(sortSet(decodeSet(registry, valType, rawValue)));
 
     this.registry = registry;
     this.#ValClass = typeToConstructor(registry, valType);

packages/types/src/codec/Map.ts

@@ -6,7 +6,10 @@ import type { AnyJson, Codec, Constructor, InterfaceTypes, Registry } from '../t
 
 import { compactFromU8a, compactToU8a, isHex, isObject, isU8a, logger, stringify, u8aConcat, u8aToHex, u8aToU8a } from '@polkadot/util';
 
-import { compareMap, decodeU8a, typeToConstructor } from './utils';
+import { AbstractArray } from './AbstractArray';
+import { Enum } from './Enum';
+import { Struct } from './Struct';
+import { compareMap, decodeU8a, sortMap, typeToConstructor } from './utils';
 
 const l = logger('Map');
 
@@ -34,11 +37,15 @@ function decodeMapFromMap<K extends Codec = Codec, V extends Codec = Codec> (reg
   const output = new Map<K, V>();
 
   value.forEach((val: any, key: any) => {
+    const isComplex = KeyClass.prototype instanceof AbstractArray ||
+      KeyClass.prototype instanceof Struct ||
+      KeyClass.prototype instanceof Enum;
+
     try {
       output.set(
         key instanceof KeyClass
           ? key
-          : new KeyClass(registry, key),
+          : new KeyClass(registry, isComplex ? JSON.parse(key) : key),
         val instanceof ValClass
           ? val
           : new ValClass(registry, val)
@@ -97,7 +104,9 @@ export class CodecMap<K extends Codec = Codec, V extends Codec = Codec> extends
   readonly #type: string;
 
   constructor (registry: Registry, keyType: Constructor<K> | keyof InterfaceTypes, valType: Constructor<V> | keyof InterfaceTypes, rawValue: Uint8Array | string | Map<any, any> | undefined, type: 'BTreeMap' | 'HashMap' = 'HashMap') {
-    super(decodeMap(registry, keyType, valType, rawValue));
+    const decoded = decodeMap(registry, keyType, valType, rawValue);
+
+    super(type === 'BTreeMap' ? sortMap(decoded) : decoded);
 
     this.registry = registry;
     this.#KeyClass = typeToConstructor(registry, keyType);

packages/types/src/codec/utils/index.ts

@@ -7,3 +7,4 @@ export { compareSet } from './compareSet';
 export { decodeU8a } from './decodeU8a';
 export { mapToTypeMap } from './mapToTypeMap';
 export { typeToConstructor } from './typeToConstructor';
+export { sortAsc, sortSet, sortMap } from './sortValues';

packages/types/src/codec/utils/sortValues.ts

@@ -0,0 +1,64 @@
+// Copyright 2017-2021 @polkadot/types authors & contributors
+// SPDX-License-Identifier: Apache-2.0
+
+import type { Codec } from '../../types';
+
+import { AbstractArray } from '../AbstractArray';
+import { AbstractInt } from '../AbstractInt';
+import { Enum, Struct } from '..';
+
+type SortArg = Codec | Codec[] | number[] | number | Uint8Array
+
+/** @internal **/
+function isArrayLike (arg: SortArg): arg is AbstractArray<Codec> | Uint8Array | Codec[] | number[] {
+  return arg instanceof AbstractArray || arg instanceof Uint8Array || Array.isArray(arg);
+}
+
+/** @internal **/
+function isCodec (arg: SortArg): arg is Codec {
+  return !!(arg && typeof (arg as Codec).toU8a === 'function');
+}
+
+/**
+* Sort keys/values of BTreeSet/BTreeMap in acending order for encoding compatibility with Rust's BTreeSet/BTreeMap
+* (https://doc.rust-lang.org/stable/std/collections/struct.BTreeSet.html)
+* (https://doc.rust-lang.org/stable/std/collections/struct.BTreeMap.html)
+*/
+export function sortAsc<V extends SortArg = Codec> (a: V, b: V): number {
+  if (typeof a === 'number' && typeof b === 'number') {
+    return a - b;
+  } else if (a instanceof Struct && b instanceof Struct) {
+    return sortAsc(Array.from(a.values()), Array.from(b.values()));
+  } else if (a instanceof Enum && b instanceof Enum) {
+    return sortAsc(a.index, b.index) || sortAsc(a.value, b.value);
+  } else if (isArrayLike(a) && isArrayLike(b)) {
+    // Vec, Tuple, Bytes etc.
+    let sortRes = 0; const lenA = a.length; const lenB = b.length;
+    const minLen = Math.min(lenA, lenB);
+
+    for (let i = 0; i < minLen; ++i) {
+      sortRes = sortAsc(a[i], b[i]);
+
+      if (sortRes !== 0) {
+        return sortRes;
+      }
+    }
+
+    return lenA - lenB;
+  } else if (a instanceof AbstractInt && b instanceof AbstractInt) {
+    return a.cmp(b);
+  } else if (isCodec(a) && isCodec(b)) {
+    // Text, Bool etc.
+    return sortAsc(a.toU8a(true), b.toU8a(true));
+  } else {
+    throw new Error(`Attempting to sort unrecognized value: ${JSON.stringify(a)}`);
+  }
+}
+
+export function sortSet<V extends Codec = Codec> (set: Set<V>): Set<V> {
+  return new Set(Array.from(set).sort(sortAsc));
+}
+
+export function sortMap<K extends Codec = Codec, V extends Codec = Codec> (map: Map<K, V>): Map<K, V> {
+  return new Map(Array.from(map.entries()).sort(([keyA], [keyB]) => sortAsc(keyA, keyB)));
+}