preserves/implementations/javascript/packages/core/src/dictionary.ts

import { Encoder, canonicalString } from "./encoder";
import { Tag } from "./constants";
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 { Writer, PreserveWritable } from "./writer";
import { annotations, Annotated } from "./annotated";

export type DictionaryType = 'Dictionary' | 'Set';
export const DictionaryType = Symbol.for('DictionaryType');

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 CompoundKey<T>, V, T = GenericEmbedded>(
        x: any,
    ): x is KeyedDictionary<K, V, T> {
        return x?.[DictionaryType] === 'Dictionary';
    }

    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> {
        const result = new KeyedDictionary<S, W, R>();
        for (let oldEntry of this.entries()) {
            const newEntry = f(oldEntry);
            result.set(newEntry[0], newEntry[1])
        }
        return result;
    }

    clone(): KeyedDictionary<K, V, T> {
        return new KeyedDictionary(this);
    }

    get [Symbol.toStringTag]() { return 'Dictionary'; }
}

export class Dictionary<T = GenericEmbedded, V = Value<T>> extends KeyedDictionary<Value<T>, V, T> {
    static isDictionary<T = GenericEmbedded, V = Value<T>>(x: any): x is Dictionary<T, V> {
        return x?.[DictionaryType] === 'Dictionary';
    }

    static __from_preserve__<T>(v: Value<T>): undefined | Dictionary<T> {
        return Dictionary.isDictionary<T>(v) ? v : void 0;
    }
}

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 CompoundKey<T>, T = GenericEmbedded>(
        x: any,
    ): x is KeyedSet<K, T> {
        return x?.[DictionaryType] === 'Set';
    }

    constructor(items?: Iterable<K>) {
        super(k => k, items);
    }

    map<S extends Value<R>, R = GenericEmbedded>(f: (value: K) => S): KeyedSet<S, R> {
        return new KeyedSet(_iterMap(this[Symbol.iterator](), f));
    }

    filter(f: (value: K) => boolean): KeyedSet<K, T> {
        const result = new KeyedSet<K, T>();
        for (let k of this) if (f(k)) result.add(k);
        return result;
    }

    clone(): KeyedSet<K, T> {
        return new KeyedSet(this);
    }

    get [Symbol.toStringTag]() { return 'Set'; }
}

export class Set<T = GenericEmbedded> extends KeyedSet<Value<T>, T> {
    static isSet<T = GenericEmbedded>(x: any): x is Set<T> {
        return x?.[DictionaryType] === 'Set';
    }

    static __from_preserve__<T>(v: Value<T>): undefined | Set<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));
}