-
Notifications
You must be signed in to change notification settings - Fork 4
/
union.ts
105 lines (99 loc) · 3.86 KB
/
union.ts
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
import { AnyCodec, Input, Output } from "../common/codec.ts"
import { Codec, createCodec, Expand, metadata, Narrow, ScaleAssertError, ScaleDecodeError } from "../common/mod.ts"
import { constant } from "./constant.ts"
import { field, InputObject, object, ObjectMembers, OutputObject } from "./object.ts"
export class Variant<T extends string, I, O> {
constructor(readonly tag: T, readonly codec: Codec<I, O>) {}
}
export type AnyVariant = Variant<any, never, unknown>
export function variant<T extends string, E extends AnyCodec[]>(
tag: T,
...members: ObjectMembers<E>
): Variant<T, InputObject<E>, OutputObject<E>> {
return new Variant(tag, object(...members))
}
export type InputTaggedUnion<
K extends keyof any,
M extends Record<number, AnyVariant>,
> = {
[I in keyof M]: Expand<
& Readonly<Record<K, Extract<M[I], AnyVariant>["tag"]>>
& Input<Extract<M[I], AnyVariant>["codec"]>
>
}[keyof M & number]
export type OutputTaggedUnion<
K extends keyof any,
M extends Record<number, AnyVariant>,
> = {
[I in keyof M]: Expand<
& Record<K, Extract<M[I], AnyVariant>["tag"]>
& Output<Extract<M[I], AnyVariant>["codec"]>
>
}[keyof M & number]
export function taggedUnion<
K extends keyof any,
M extends [] | Record<number, Variant<any, never, unknown>>,
>(tagKey: K, members: M): Codec<InputTaggedUnion<K, M>, OutputTaggedUnion<K, M>> {
const tagToDiscriminant: Record<string, number> = Object.create(null)
const discriminantToMember: Record<number, AnyCodec> = Object.create(null)
for (const _discriminant in members) {
const discriminant = +_discriminant
if (isNaN(discriminant)) continue
const { tag, codec } = (members as M)[discriminant]!
tagToDiscriminant[tag] = discriminant
discriminantToMember[discriminant] = object(field(tagKey, constant(tag)) as any, codec)
}
return createCodec({
_metadata: metadata("$.taggedUnion", taggedUnion, tagKey, members),
_staticSize: 1 + Math.max(...Object.values(discriminantToMember).map((x) => x._staticSize)),
_encode(buffer, value) {
const discriminant = tagToDiscriminant[value[tagKey]]!
const $member = discriminantToMember[discriminant]!
buffer.array[buffer.index++] = discriminant
$member._encode(buffer, value as never)
},
_decode(buffer) {
const discriminant = buffer.array[buffer.index++]!
const $member = discriminantToMember[discriminant]
if (!$member) {
throw new ScaleDecodeError(this, buffer, `No such member codec matching the discriminant \`${discriminant}\``)
}
return $member._decode(buffer) as any
},
_assert(assert) {
const assertTag = assert.key(this, tagKey)
assertTag.typeof(this, "string")
if (!((assertTag.value as string) in tagToDiscriminant)) {
throw new ScaleAssertError(this, assertTag.value, `${assertTag.path}: invalid tag`)
}
discriminantToMember[tagToDiscriminant[assertTag.value as string]!]!._assert(assert)
},
})
}
export function literalUnion<T extends Narrow>(members: Record<number, T>): Codec<T> {
const keyToDiscriminant: Map<T, number> = new Map()
for (const _discriminant in members) {
const discriminant = +_discriminant
if (isNaN(discriminant)) continue
const key = members[discriminant] as T
keyToDiscriminant.set(key, discriminant)
}
return createCodec({
_metadata: metadata("$.literalUnion", literalUnion, members),
_staticSize: 1,
_encode(buffer, value) {
const discriminant = keyToDiscriminant.get(value)!
buffer.array[buffer.index++] = discriminant
},
_decode(buffer) {
const discriminant = buffer.array[buffer.index++]!
return members[discriminant]!
},
_assert(assert) {
assert.typeof(this, "string")
if (!keyToDiscriminant.has(assert.value as T)) {
throw new ScaleAssertError(this, assert.value, `${assert.path} invalid value`)
}
},
})
}