From 9420cc723659b8427c145eb9a03c5db1f9e41b12 Mon Sep 17 00:00:00 2001
From: Tony Garnock-Jones <tonyg@leastfixedpoint.com>
Date: Sat, 16 Dec 2023 22:13:17 +1300
Subject: [PATCH] New EncodableDictionary/EncodableSet as intermediate steps
 between Flex- and Keyed-

---
 .../packages/core/src/dictionary.ts           | 207 ++++++++++++------
 .../javascript/packages/core/src/encoder.ts   |   8 +-
 2 files changed, 148 insertions(+), 67 deletions(-)

diff --git a/implementations/javascript/packages/core/src/dictionary.ts b/implementations/javascript/packages/core/src/dictionary.ts
index 570614b..358cc8b 100644
--- a/implementations/javascript/packages/core/src/dictionary.ts
+++ b/implementations/javascript/packages/core/src/dictionary.ts
@@ -1,31 +1,61 @@
-import { Encodable, Encoder, canonicalEncode, canonicalString } from "./encoder";
+import { Encoder, canonicalString } from "./encoder";
 import { Tag } from "./constants";
-import { FlexMap, FlexSet, _iterMap } from "./flex";
+import { FlexMap, FlexSet, _iterMap, IdentitySet } from "./flex";
 import { Value } from "./values";
 import { Bytes } from './bytes';
 import { GenericEmbedded } from "./embedded";
 import type { Preservable } from "./encoder";
-import type { Writable, Writer, PreserveWritable } from "./writer";
+import type { Writer, PreserveWritable } from "./writer";
 import { annotations, Annotated } from "./annotated";
 
 export type DictionaryType = 'Dictionary' | 'Set';
 export const DictionaryType = Symbol.for('DictionaryType');
 
-export class KeyedDictionary<K extends (Encodable<T> & Writable<T>), V, T = GenericEmbedded> extends FlexMap<K, V>
+export type CompoundKey<T> = Value<T> | (Preservable<T> & PreserveWritable<T>);
+
+export class EncodableDictionary<K, V, T = GenericEmbedded> extends FlexMap<K, V>
     implements Preservable<T>, PreserveWritable<T>
+{
+    constructor(
+        public readonly encodeK: (k: K) => CompoundKey<T>,
+        public readonly encodeV: (v: V) => CompoundKey<T>,
+        items?: readonly [K, V][] | Iterable<readonly [K, V]>
+    ) {
+        super((k: K) => canonicalString(encodeK(k)), items);
+    }
+
+    __preserve_on__(encoder: Encoder<T>) {
+        encodeDictionaryOn(this,
+                           encoder,
+                           (k, e) => e.push(this.encodeK(k)),
+                           (v, e) => e.push(this.encodeV(v)));
+    }
+
+    __preserve_text_on__(w: Writer<T>) {
+        writeDictionaryOn(this,
+                          w,
+                          (k, w) => w.push(this.encodeK(k)),
+                          (v, w) => w.push(this.encodeV(v)));
+    }
+}
+
+export class KeyedDictionary<K extends CompoundKey<T>, V, T = GenericEmbedded>
+    extends EncodableDictionary<K, V, T>
 {
     get [DictionaryType](): DictionaryType {
         return 'Dictionary';
     }
 
-    static isKeyedDictionary<K extends Value<T>, V, T = GenericEmbedded>(x: any): x is KeyedDictionary<K, V, T> {
+    static isKeyedDictionary<K extends CompoundKey<T>, V, T = GenericEmbedded>(
+        x: any,
+    ): x is KeyedDictionary<K, V, T> {
         return x?.[DictionaryType] === 'Dictionary';
     }
 
-    constructor(items?: readonly [K, V][]);
-    constructor(items?: Iterable<readonly [K, V]>);
-    constructor(items?: Iterable<readonly [K, V]>) {
-        super(canonicalString, items);
+    constructor(items?: readonly [K, V][] | Iterable<readonly [K, V]>) {
+        // The cast in encodeV is suuuuuuuper unsound, since V may not in fact be Encodable and Writable.
+        // Don't try to encode/write dictionaries holding non-encodable/non-writable values.
+        super(k => k, v => v as CompoundKey<T>, items);
     }
 
     mapEntries<W, S extends Value<R>, R = GenericEmbedded>(f: (entry: [K, V]) => [S, W]): KeyedDictionary<S, W, R> {
@@ -42,44 +72,6 @@ export class KeyedDictionary<K extends (Encodable<T> & Writable<T>), V, T = Gene
     }
 
     get [Symbol.toStringTag]() { return 'Dictionary'; }
-
-    __preserve_on__(encoder: Encoder<T>) {
-        if (encoder.canonical) {
-            const entries = Array.from(this);
-            const pieces = entries.map<[Bytes, number]>(([k, _v], i) => [canonicalEncode(k), i]);
-            pieces.sort((a, b) => Bytes.compare(a[0], b[0]));
-            encoder.state.emitbyte(Tag.Dictionary);
-            pieces.forEach(([_encodedKey, i]) => {
-                const [k, v] = entries[i];
-                encoder.push(k);
-                encoder.push(v as unknown as Value<T>); // Suuuuuuuper unsound
-            });
-            encoder.state.emitbyte(Tag.End);
-        } else {
-            encoder.state.emitbyte(Tag.Dictionary);
-            this.forEach((v, k) => {
-                encoder.push(k);
-                encoder.push(v as unknown as Value<T>); // Suuuuuuuper unsound
-            });
-            encoder.state.emitbyte(Tag.End);
-        }
-    }
-
-    __preserve_text_on__(w: Writer<T>) {
-        w.state.writeSeq('{', '}', this.entries(), ([k, v]) => {
-            w.push(k);
-            if (Annotated.isAnnotated<T>(v) && (annotations(v).length > 1) && w.state.isIndenting) {
-                w.state.pieces.push(':');
-                w.state.indentCount++;
-                w.state.writeIndent();
-                w.push(v);
-                w.state.indentCount--;
-            } else {
-                w.state.pieces.push(': ');
-                w.push(v as unknown as Value<T>); // Suuuuuuuper unsound
-            }
-        });
-    }
 }
 
 export class Dictionary<T = GenericEmbedded, V = Value<T>> extends KeyedDictionary<Value<T>, V, T> {
@@ -92,19 +84,91 @@ export class Dictionary<T = GenericEmbedded, V = Value<T>> extends KeyedDictiona
     }
 }
 
-export class KeyedSet<K extends Value<T>, T = GenericEmbedded> extends FlexSet<K>
+export function encodeDictionaryOn<K, V, T>(
+    dict: Map<K, V>,
+    encoder: Encoder<T>,
+    encodeK: (k: K, encoder: Encoder<T>) => void,
+    encodeV: (v: V, encoder: Encoder<T>) => void,
+) {
+    if (encoder.canonical) {
+        const entries = Array.from(dict);
+        const canonicalEncoder = new Encoder<T>({
+            canonical: true,
+            embeddedEncode: encoder.embeddedEncode,
+        });
+        const pieces = entries.map<[Bytes, number]>(([k, _v], i) => {
+            encodeK(k, canonicalEncoder);
+            return [canonicalEncoder.contents(), i];
+        });
+        pieces.sort((a, b) => Bytes.compare(a[0], b[0]));
+        encoder.grouped(Tag.Dictionary, () => pieces.forEach(([_encodedKey, i]) => {
+            const [k, v] = entries[i];
+            encodeK(k, encoder);
+            encodeV(v, encoder);
+        }));
+    } else {
+        encoder.grouped(Tag.Dictionary, () => dict.forEach((v, k) => {
+            encodeK(k, encoder);
+            encodeV(v, encoder);
+        }));
+    }
+}
+
+export function writeDictionaryOn<K, V, T>(
+    dict: Map<K, V>,
+    w: Writer<T>,
+    writeK: (k: K, w: Writer<T>) => void,
+    writeV: (v: V, w: Writer<T>) => void,
+) {
+    w.state.writeSeq('{', '}', dict.entries(), ([k, v]) => {
+        writeK(k, w);
+        if (Annotated.isAnnotated<T>(v) && (annotations(v).length > 1) && w.state.isIndenting) {
+            w.state.pieces.push(':');
+            w.state.indentCount++;
+            w.state.writeIndent();
+            writeV(v, w);
+            w.state.indentCount--;
+        } else {
+            w.state.pieces.push(': ');
+            writeV(v, w);
+        }
+    });
+}
+
+export class EncodableSet<V, T = GenericEmbedded> extends FlexSet<V>
     implements Preservable<T>, PreserveWritable<T>
+{
+    constructor(
+        public readonly encodeV: (v: V) => CompoundKey<T>,
+        items?: Iterable<V>,
+    ) {
+        super((v: V) => canonicalString(encodeV(v)), items);
+    }
+
+    __preserve_on__(encoder: Encoder<T>) {
+        encodeSetOn(this, encoder, (v, e) => e.push(this.encodeV(v)));
+    }
+
+    __preserve_text_on__(w: Writer<T>) {
+        writeSetOn(this, w, (v, w) => w.push(this.encodeV(v)));
+    }
+}
+
+export class KeyedSet<K extends CompoundKey<T>, T = GenericEmbedded>
+    extends EncodableSet<K, T>
 {
     get [DictionaryType](): DictionaryType {
         return 'Set';
     }
 
-    static isKeyedSet<K extends Value<T>, T = GenericEmbedded>(x: any): x is KeyedSet<K, T> {
+    static isKeyedSet<K extends CompoundKey<T>, T = GenericEmbedded>(
+        x: any,
+    ): x is KeyedSet<K, T> {
         return x?.[DictionaryType] === 'Set';
     }
 
     constructor(items?: Iterable<K>) {
-        super(canonicalString, items);
+        super(k => k, items);
     }
 
     map<S extends Value<R>, R = GenericEmbedded>(f: (value: K) => S): KeyedSet<S, R> {
@@ -122,20 +186,6 @@ export class KeyedSet<K extends Value<T>, T = GenericEmbedded> extends FlexSet<K
     }
 
     get [Symbol.toStringTag]() { return 'Set'; }
-
-    __preserve_on__(encoder: Encoder<T>) {
-        if (encoder.canonical) {
-            const pieces = Array.from(this).map<[Bytes, K]>(k => [canonicalEncode(k), k]);
-            pieces.sort((a, b) => Bytes.compare(a[0], b[0]));
-            encoder.encodevalues(Tag.Set, pieces.map(e => e[1]));
-        } else {
-            encoder.encodevalues(Tag.Set, this);
-        }
-    }
-
-    __preserve_text_on__(w: Writer<T>) {
-        w.state.writeSeq('#{', '}', this, vv => w.push(vv));
-    }
 }
 
 export class Set<T = GenericEmbedded> extends KeyedSet<Value<T>, T> {
@@ -147,3 +197,32 @@ export class Set<T = GenericEmbedded> extends KeyedSet<Value<T>, T> {
         return Set.isSet<T>(v) ? v : void 0;
     }
 }
+
+export function encodeSetOn<V, T>(
+    s: IdentitySet<V>,
+    encoder: Encoder<T>,
+    encodeV: (v: V, encoder: Encoder<T>) => void,
+) {
+    if (encoder.canonical) {
+        const canonicalEncoder = new Encoder<T>({
+            canonical: true,
+            embeddedEncode: encoder.embeddedEncode,
+        });
+        const pieces = Array.from(s).map<[Bytes, V]>(v => {
+            encodeV(v, canonicalEncoder);
+            return [canonicalEncoder.contents(), v];
+        });
+        pieces.sort((a, b) => Bytes.compare(a[0], b[0]));
+        encoder.grouped(Tag.Set, () => pieces.forEach(([_e, v]) => encodeV(v, encoder)));
+    } else {
+        encoder.grouped(Tag.Set, () => s.forEach(v => encodeV(v, encoder)));
+    }
+}
+
+export function writeSetOn<V, T>(
+    s: IdentitySet<V>,
+    w: Writer<T>,
+    writeV: (v: V, w: Writer<T>) => void,
+) {
+    w.state.writeSeq('#{', '}', s, vv => writeV(vv, w));
+}
diff --git a/implementations/javascript/packages/core/src/encoder.ts b/implementations/javascript/packages/core/src/encoder.ts
index a283c69..d8b9cd7 100644
--- a/implementations/javascript/packages/core/src/encoder.ts
+++ b/implementations/javascript/packages/core/src/encoder.ts
@@ -242,9 +242,9 @@ export class Encoder<T = object> {
         return this.state.contentsString();
     }
 
-    encodevalues(tag: Tag, items: Iterable<Value<T>>) {
+    grouped(tag: Tag, f: () => void) {
         this.state.emitbyte(tag);
-        for (let i of items) { this.push(i); }
+        f();
         this.state.emitbyte(Tag.End);
     }
 
@@ -284,7 +284,9 @@ export class Encoder<T = object> {
             this.state.emitbyte(Tag.End);
         }
         else if (isIterable<Value<T>>(v)) {
-            this.encodevalues(Tag.Sequence, v);
+            this.grouped(Tag.Sequence, () => {
+                for (let i of v) this.push(i);
+            });
         }
         else {
             ((v: Embedded<T>) => {