Fix sorting of Option values and error on duplicates from BTreeSet/BTreeMap.

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

@@ -6,7 +6,7 @@
 import type { CodecClass, ITuple } from '@polkadot/types-codec/types';
 
 import { TypeRegistry } from '@polkadot/types';
-import { BTreeMap, Enum, I32, Struct, Text, Tuple, U32 } from '@polkadot/types-codec';
+import { BTreeMap, Enum, I32, Option, Struct, Text, Tuple, U32 } from '@polkadot/types-codec';
 import { stringToU8a } from '@polkadot/util';
 
 const registry = new TypeRegistry();
@@ -20,15 +20,21 @@ class MockEnum extends Enum.with({
   Key2: MockStruct,
   Key3: U32TextTuple
 }) {}
+class MockOptionEnum extends Option.with(MockEnum) {}
 
 const mockU32TextMap = new Map<Text, U32>();
+const mockU32DuplicateTextMap = 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>();
+const mockU32OptionEnumMap = new Map<MockOptionEnum, U32>();
 
 mockU32TextMap.set(new Text(registry, 'bazzing'), new U32(registry, 69));
 
+mockU32DuplicateTextMap.set(new Text(registry, 'bazzing'), new U32(registry, 42));
+mockU32DuplicateTextMap.set(new Text(registry, 'bazzing'), new U32(registry, 43));
+
 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));
@@ -50,6 +56,13 @@ mockU32EnumMap.set(new MockEnum(registry, { Key2: new MockStruct(registry, { int
 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));
 
+mockU32OptionEnumMap.set(new Option(registry, MockEnum, null), new U32(registry, 13));
+mockU32OptionEnumMap.set(new Option(registry, MockEnum, { Key3: new U32TextTuple(registry, [2, 'ba']) }), new U32(registry, 13));
+mockU32OptionEnumMap.set(new Option(registry, MockEnum, { Key3: new U32TextTuple(registry, [2, 'b']) }), new U32(registry, 42));
+mockU32OptionEnumMap.set(new Option(registry, MockEnum, { Key2: new MockStruct(registry, { int: -1, text: 'b' }) }), new U32(registry, 7));
+mockU32OptionEnumMap.set(new Option(registry, MockEnum, { Key1: new MockStruct(registry, { int: 1, text: 'b' }) }), new U32(registry, 25));
+mockU32OptionEnumMap.set(new Option(registry, MockEnum, { Key1: new MockStruct(registry, { int: -1, text: 'b' }) }), new U32(registry, 69));
+
 describe('BTreeMap', (): void => {
   it('decodes null', (): void => {
     expect(
@@ -78,6 +91,12 @@ describe('BTreeMap', (): void => {
     expect(s.toString()).toBe('{"placeholder":0,"value":{"bazzing":69}}');
   });
 
+  it('throws on duplicate keys', (): void => {
+    expect(
+      () => new (BTreeMap.with(Text, U32))(registry, mockU32DuplicateTextMap)
+    ).toThrow(/Duplicate value in BTreeMap/);
+  });
+
   it('throws when it cannot decode', (): void => {
     expect(
       (): BTreeMap<Text, U32> => new (
@@ -116,6 +135,19 @@ describe('BTreeMap', (): void => {
     ]);
   });
 
+  it('correctly sorts Option(Enum) keys', (): void => {
+    expect(
+      Array.from(new (BTreeMap.with(MockOptionEnum, U32))(registry, mockU32OptionEnumMap).keys()).map((k) => k.value.toJSON())
+    ).toEqual([
+      null,
+      { key1: { int: -1, text: 'b' } },
+      { key1: { int: 1, text: 'b' } },
+      { key2: { int: -1, text: 'b' } },
+      { key3: [2, 'b'] },
+      { key3: [2, 'ba'] }
+    ]);
+  });
+
   it('correctly sorts enum keys', (): void => {
     expect(
       Array.from(new (BTreeMap.with(MockEnum, U32))(registry, mockU32EnumMap).keys()).map((k) => k.toJSON())

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

@@ -7,7 +7,7 @@ import type { CodecClass, CodecTo } from '@polkadot/types-codec/types';
 import type { ITuple } from '../types/interfaces.js';
 
 import { TypeRegistry } from '@polkadot/types';
-import { BTreeSet, Enum, I32, Struct, Text, Tuple, U32 } from '@polkadot/types-codec';
+import { BTreeSet, Enum, I32, Option, Struct, Text, Tuple, U32 } from '@polkadot/types-codec';
 
 const registry = new TypeRegistry();
 
@@ -20,6 +20,7 @@ class MockEnum extends Enum.with({
   Key2: MockStruct,
   Key3: U32TextTuple
 }) {}
+class MockOptionEnum extends Option.with(MockEnum) {}
 
 const mockU32Set = new Set<U32>();
 
@@ -61,6 +62,22 @@ const mockEnumSetObj = [
   new MockEnum(registry, { Key1: new MockStruct(registry, { int: -1, text: 'b' }) })
 ];
 
+const mockOptionEnumSetObj = [
+  new Option(registry, MockEnum, null),
+  new Option(registry, MockEnum, { Key3: new U32TextTuple(registry, [2, 'ba']) }),
+  new Option(registry, MockEnum, { Key3: new U32TextTuple(registry, [2, 'b']) }),
+  new Option(registry, MockEnum, { Key2: new MockStruct(registry, { int: -1, text: 'b' }) }),
+  new Option(registry, MockEnum, { Key1: new MockStruct(registry, { int: 1, text: 'b' }) }),
+  new Option(registry, MockEnum, { Key1: new MockStruct(registry, { int: -1, text: 'b' }) })
+];
+const mockDuplicateTextSetObj = [
+  new Text(registry, 'baz'),
+  new Text(registry, 'bb'),
+  new Text(registry, 'bb'), // duplicate.
+  new Text(registry, 'ba'),
+  new Text(registry, 'c')
+];
+
 describe('BTreeSet', (): void => {
   describe('decoding', (): void => {
     const testDecode = (type: string, input: unknown, output: string): void =>
@@ -171,6 +188,12 @@ describe('BTreeSet', (): void => {
       ).toEqual(['b', 'ba', 'baz', 'bb', 'c']);
     });
 
+    it('Reject duplicate values', (): void => {
+      expect(
+        () => new (BTreeSet.with(Text))(registry, mockDuplicateTextSetObj)
+      ).toThrow(/Duplicate value in BTreeSet/);
+    });
+
     it('correctly sorts complex tuple values', (): void => {
       expect(
         Array.from(new (BTreeSet.with(U32TextTuple))(registry, mockTupleSetObj)).map((k) => k.toJSON())
@@ -188,6 +211,19 @@ describe('BTreeSet', (): void => {
       ]);
     });
 
+    it('correctly sorts complex Option(enum) values', (): void => {
+      expect(
+        Array.from(new (BTreeSet.with(MockOptionEnum))(registry, mockOptionEnumSetObj)).map((k) => k.value.toJSON())
+      ).toEqual([
+        null,
+        { key1: { int: -1, text: 'b' } },
+        { key1: { int: 1, text: 'b' } },
+        { key2: { int: -1, text: 'b' } },
+        { key3: [2, 'b'] },
+        { key3: [2, 'ba'] }
+      ]);
+    });
+
     it('correctly sorts complex enum values', (): void => {
       expect(
         Array.from(new (BTreeSet.with(MockEnum))(registry, mockEnumSetObj)).map((k) => k.toJSON())

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

@@ -2,10 +2,12 @@
 // SPDX-License-Identifier: Apache-2.0
 
 import type { BN } from '@polkadot/util';
+import type { HexString } from '@polkadot/util/types';
 import type { Enum } from '../base/Enum.js';
+import type { Option } from '../base/Option.js';
 import type { Codec } from '../types/index.js';
 
-import { bnToBn, isBigInt, isBn, isCodec, isNumber, stringify } from '@polkadot/util';
+import { bnToBn, isBigInt, isBn, isBoolean, isCodec, isNumber, stringify } from '@polkadot/util';
 
 type SortArg = Codec | Codec[] | number[] | BN | bigint | number | Uint8Array;
 
@@ -19,6 +21,11 @@ function isEnum (arg: SortArg): arg is Enum {
   return isCodec<Codec>(arg) && isNumber((arg as Enum).index) && isCodec((arg as Enum).value);
 }
 
+/** @internal **/
+function isOption (arg: SortArg): arg is Option<Codec> {
+  return isCodec<Codec>(arg) && isBoolean((arg as Option<Codec>).isSome) && isCodec((arg as Option<Codec>).value);
+}
+
 /** @internal */
 function isNumberLike (arg: SortArg): arg is BN | bigint | number {
   return isNumber(arg) || isBn(arg) || isBigInt(arg);
@@ -41,6 +48,24 @@ function sortArray (a: Uint8Array | Codec[] | number[], b: Uint8Array | Codec[]
   return a.length - b.length;
 }
 
+/** @internal */
+function checkForDuplicates (container: string, seen: Set<HexString>, arg: SortArg): boolean {
+  // Convert the value to hex.
+  if (isCodec<Codec>(arg)) {
+    const hex = arg.toHex();
+
+    // Check if we have seen the value.
+    if (seen.has(hex)) {
+      // Duplicates are not allowed.
+      throw new Error(`Duplicate value in ${container}: ${stringify(arg)}`);
+    }
+
+    seen.add(hex);
+  }
+
+  return true;
+}
+
 /**
 * Sort keys/values of BTreeSet/BTreeMap in ascending order for encoding compatibility with Rust's BTreeSet/BTreeMap
 * (https://doc.rust-lang.org/stable/std/collections/struct.BTreeSet.html)
@@ -53,6 +78,8 @@ export function sortAsc<V extends SortArg = Codec> (a: V, b: V): number {
     return sortAsc(Array.from(a.values()), Array.from(b.values()));
   } else if (isEnum(a) && isEnum(b)) {
     return sortAsc(a.index, b.index) || sortAsc(a.value, b.value);
+  } else if (isOption(a) && isOption(b)) {
+    return sortAsc(a.isNone ? 0 : 1, b.isNone ? 0 : 1) || sortAsc(a.value, b.value);
   } else if (isArrayLike(a) && isArrayLike(b)) {
     return sortArray(a, b);
   } else if (isCodec<Codec>(a) && isCodec<Codec>(b)) {
@@ -64,9 +91,13 @@ export function sortAsc<V extends SortArg = Codec> (a: V, b: V): number {
 }
 
 export function sortSet<V extends Codec = Codec> (set: Set<V>): Set<V> {
-  return new Set(Array.from(set).sort(sortAsc));
+  const seen = new Set<HexString>();
+
+  return new Set(Array.from(set).filter((value) => checkForDuplicates('BTreeSet', seen, value)).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)));
+  const seen = new Set<HexString>();
+
+  return new Map(Array.from(map.entries()).filter(([key]) => checkForDuplicates('BTreeMap', seen, key)).sort(([keyA], [keyB]) => sortAsc(keyA, keyB)));
 }