Skip to content

Commit e30dc06

Browse files
orensvfacundomedica
orensv
and
facundomedica
authored
feat(collections): quad collection (#21090)
Co-authored-by: Facundo Medica <14063057+facundomedica@users.noreply.github.com>
1 parent 4e97c49 commit e30dc06

2 files changed

Lines changed: 439 additions & 0 deletions

File tree

‎collections/quad.go‎

Lines changed: 378 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,378 @@
1+
package collections
2+
3+
import (
4+
"encoding/json"
5+
"fmt"
6+
"strings"
7+
8+
"cosmossdk.io/collections/codec"
9+
)
10+
11+
// Quad defines a multipart key composed of four keys.
12+
type Quad[K1, K2, K3, K4 any] struct {
13+
k1 *K1
14+
k2 *K2
15+
k3 *K3
16+
k4 *K4
17+
}
18+
19+
// Join4 instantiates a new Quad instance composed of the four provided keys, in order.
20+
func Join4[K1, K2, K3, K4 any](k1 K1, k2 K2, k3 K3, k4 K4) Quad[K1, K2, K3, K4] {
21+
return Quad[K1, K2, K3, K4]{&k1, &k2, &k3, &k4}
22+
}
23+
24+
// K1 returns the first part of the key. If nil, the zero value is returned.
25+
func (t Quad[K1, K2, K3, K4]) K1() (x K1) {
26+
if t.k1 != nil {
27+
return *t.k1
28+
}
29+
return x
30+
}
31+
32+
// K2 returns the second part of the key. If nil, the zero value is returned.
33+
func (t Quad[K1, K2, K3, K4]) K2() (x K2) {
34+
if t.k2 != nil {
35+
return *t.k2
36+
}
37+
return x
38+
}
39+
40+
// K3 returns the third part of the key. If nil, the zero value is returned.
41+
func (t Quad[K1, K2, K3, K4]) K3() (x K3) {
42+
if t.k3 != nil {
43+
return *t.k3
44+
}
45+
return x
46+
}
47+
48+
// K4 returns the fourth part of the key. If nil, the zero value is returned.
49+
func (t Quad[K1, K2, K3, K4]) K4() (x K4) {
50+
if t.k4 != nil {
51+
return *t.k4
52+
}
53+
return x
54+
}
55+
56+
// QuadPrefix creates a new Quad instance composed only of the first part of the key.
57+
func QuadPrefix[K1, K2, K3, K4 any](k1 K1) Quad[K1, K2, K3, K4] {
58+
return Quad[K1, K2, K3, K4]{k1: &k1}
59+
}
60+
61+
// QuadSuperPrefix creates a new Quad instance composed only of the first two parts of the key.
62+
func QuadSuperPrefix[K1, K2, K3, K4 any](k1 K1, k2 K2) Quad[K1, K2, K3, K4] {
63+
return Quad[K1, K2, K3, K4]{k1: &k1, k2: &k2}
64+
}
65+
66+
// QuadSuperPrefix3 creates a new Quad instance composed only of the first three parts of the key.
67+
func QuadSuperPrefix3[K1, K2, K3, K4 any](k1 K1, k2 K2, k3 K3) Quad[K1, K2, K3, K4] {
68+
return Quad[K1, K2, K3, K4]{k1: &k1, k2: &k2, k3: &k3}
69+
}
70+
71+
// QuadKeyCodec instantiates a new KeyCodec instance that can encode the Quad, given
72+
// the KeyCodecs of the four parts of the key, in order.
73+
func QuadKeyCodec[K1, K2, K3, K4 any](keyCodec1 codec.KeyCodec[K1], keyCodec2 codec.KeyCodec[K2], keyCodec3 codec.KeyCodec[K3], keyCodec4 codec.KeyCodec[K4]) codec.KeyCodec[Quad[K1, K2, K3, K4]] {
74+
return quadKeyCodec[K1, K2, K3, K4]{
75+
keyCodec1: keyCodec1,
76+
keyCodec2: keyCodec2,
77+
keyCodec3: keyCodec3,
78+
keyCodec4: keyCodec4,
79+
}
80+
}
81+
82+
type quadKeyCodec[K1, K2, K3, K4 any] struct {
83+
keyCodec1 codec.KeyCodec[K1]
84+
keyCodec2 codec.KeyCodec[K2]
85+
keyCodec3 codec.KeyCodec[K3]
86+
keyCodec4 codec.KeyCodec[K4]
87+
}
88+
89+
type jsonQuadKey [4]json.RawMessage
90+
91+
func (t quadKeyCodec[K1, K2, K3, K4]) EncodeJSON(value Quad[K1, K2, K3, K4]) ([]byte, error) {
92+
json1, err := t.keyCodec1.EncodeJSON(*value.k1)
93+
if err != nil {
94+
return nil, err
95+
}
96+
97+
json2, err := t.keyCodec2.EncodeJSON(*value.k2)
98+
if err != nil {
99+
return nil, err
100+
}
101+
102+
json3, err := t.keyCodec3.EncodeJSON(*value.k3)
103+
if err != nil {
104+
return nil, err
105+
}
106+
107+
json4, err := t.keyCodec4.EncodeJSON(*value.k4)
108+
if err != nil {
109+
return nil, err
110+
}
111+
112+
return json.Marshal(jsonQuadKey{json1, json2, json3, json4})
113+
}
114+
115+
func (t quadKeyCodec[K1, K2, K3, K4]) DecodeJSON(b []byte) (Quad[K1, K2, K3, K4], error) {
116+
var jsonKey jsonQuadKey
117+
err := json.Unmarshal(b, &jsonKey)
118+
if err != nil {
119+
return Quad[K1, K2, K3, K4]{}, err
120+
}
121+
122+
key1, err := t.keyCodec1.DecodeJSON(jsonKey[0])
123+
if err != nil {
124+
return Quad[K1, K2, K3, K4]{}, err
125+
}
126+
127+
key2, err := t.keyCodec2.DecodeJSON(jsonKey[1])
128+
if err != nil {
129+
return Quad[K1, K2, K3, K4]{}, err
130+
}
131+
132+
key3, err := t.keyCodec3.DecodeJSON(jsonKey[2])
133+
if err != nil {
134+
return Quad[K1, K2, K3, K4]{}, err
135+
}
136+
137+
key4, err := t.keyCodec4.DecodeJSON(jsonKey[3])
138+
if err != nil {
139+
return Quad[K1, K2, K3, K4]{}, err
140+
}
141+
142+
return Join4(key1, key2, key3, key4), nil
143+
}
144+
145+
func (t quadKeyCodec[K1, K2, K3, K4]) Stringify(key Quad[K1, K2, K3, K4]) string {
146+
b := new(strings.Builder)
147+
b.WriteByte('(')
148+
if key.k1 != nil {
149+
b.WriteByte('"')
150+
b.WriteString(t.keyCodec1.Stringify(*key.k1))
151+
b.WriteByte('"')
152+
} else {
153+
b.WriteString("<nil>")
154+
}
155+
156+
b.WriteString(", ")
157+
if key.k2 != nil {
158+
b.WriteByte('"')
159+
b.WriteString(t.keyCodec2.Stringify(*key.k2))
160+
b.WriteByte('"')
161+
} else {
162+
b.WriteString("<nil>")
163+
}
164+
165+
b.WriteString(", ")
166+
if key.k3 != nil {
167+
b.WriteByte('"')
168+
b.WriteString(t.keyCodec3.Stringify(*key.k3))
169+
b.WriteByte('"')
170+
} else {
171+
b.WriteString("<nil>")
172+
}
173+
174+
b.WriteString(", ")
175+
if key.k4 != nil {
176+
b.WriteByte('"')
177+
b.WriteString(t.keyCodec4.Stringify(*key.k4))
178+
b.WriteByte('"')
179+
} else {
180+
b.WriteString("<nil>")
181+
}
182+
183+
b.WriteByte(')')
184+
return b.String()
185+
}
186+
187+
func (t quadKeyCodec[K1, K2, K3, K4]) KeyType() string {
188+
return fmt.Sprintf("Quad[%s,%s,%s,%s]", t.keyCodec1.KeyType(), t.keyCodec2.KeyType(), t.keyCodec3.KeyType(), t.keyCodec4.KeyType())
189+
}
190+
191+
func (t quadKeyCodec[K1, K2, K3, K4]) Encode(buffer []byte, key Quad[K1, K2, K3, K4]) (int, error) {
192+
writtenTotal := 0
193+
if key.k1 != nil {
194+
written, err := t.keyCodec1.EncodeNonTerminal(buffer, *key.k1)
195+
if err != nil {
196+
return 0, err
197+
}
198+
writtenTotal += written
199+
}
200+
if key.k2 != nil {
201+
written, err := t.keyCodec2.EncodeNonTerminal(buffer[writtenTotal:], *key.k2)
202+
if err != nil {
203+
return 0, err
204+
}
205+
writtenTotal += written
206+
}
207+
if key.k3 != nil {
208+
written, err := t.keyCodec3.EncodeNonTerminal(buffer[writtenTotal:], *key.k3)
209+
if err != nil {
210+
return 0, err
211+
}
212+
writtenTotal += written
213+
}
214+
if key.k4 != nil {
215+
written, err := t.keyCodec4.Encode(buffer[writtenTotal:], *key.k4)
216+
if err != nil {
217+
return 0, err
218+
}
219+
writtenTotal += written
220+
}
221+
return writtenTotal, nil
222+
}
223+
224+
func (t quadKeyCodec[K1, K2, K3, K4]) Decode(buffer []byte) (int, Quad[K1, K2, K3, K4], error) {
225+
readTotal := 0
226+
read, key1, err := t.keyCodec1.DecodeNonTerminal(buffer)
227+
if err != nil {
228+
return 0, Quad[K1, K2, K3, K4]{}, err
229+
}
230+
readTotal += read
231+
read, key2, err := t.keyCodec2.DecodeNonTerminal(buffer[readTotal:])
232+
if err != nil {
233+
return 0, Quad[K1, K2, K3, K4]{}, err
234+
}
235+
readTotal += read
236+
read, key3, err := t.keyCodec3.DecodeNonTerminal(buffer[readTotal:])
237+
if err != nil {
238+
return 0, Quad[K1, K2, K3, K4]{}, err
239+
}
240+
readTotal += read
241+
read, key4, err := t.keyCodec4.Decode(buffer[readTotal:])
242+
if err != nil {
243+
return 0, Quad[K1, K2, K3, K4]{}, err
244+
}
245+
readTotal += read
246+
return readTotal, Join4(key1, key2, key3, key4), nil
247+
}
248+
249+
func (t quadKeyCodec[K1, K2, K3, K4]) Size(key Quad[K1, K2, K3, K4]) int {
250+
size := 0
251+
if key.k1 != nil {
252+
size += t.keyCodec1.SizeNonTerminal(*key.k1)
253+
}
254+
if key.k2 != nil {
255+
size += t.keyCodec2.SizeNonTerminal(*key.k2)
256+
}
257+
if key.k3 != nil {
258+
size += t.keyCodec3.SizeNonTerminal(*key.k3)
259+
}
260+
if key.k4 != nil {
261+
size += t.keyCodec4.Size(*key.k4)
262+
}
263+
return size
264+
}
265+
266+
func (t quadKeyCodec[K1, K2, K3, K4]) EncodeNonTerminal(buffer []byte, key Quad[K1, K2, K3, K4]) (int, error) {
267+
writtenTotal := 0
268+
if key.k1 != nil {
269+
written, err := t.keyCodec1.EncodeNonTerminal(buffer, *key.k1)
270+
if err != nil {
271+
return 0, err
272+
}
273+
writtenTotal += written
274+
}
275+
if key.k2 != nil {
276+
written, err := t.keyCodec2.EncodeNonTerminal(buffer[writtenTotal:], *key.k2)
277+
if err != nil {
278+
return 0, err
279+
}
280+
writtenTotal += written
281+
}
282+
if key.k3 != nil {
283+
written, err := t.keyCodec3.EncodeNonTerminal(buffer[writtenTotal:], *key.k3)
284+
if err != nil {
285+
return 0, err
286+
}
287+
writtenTotal += written
288+
}
289+
if key.k4 != nil {
290+
written, err := t.keyCodec4.EncodeNonTerminal(buffer[writtenTotal:], *key.k4)
291+
if err != nil {
292+
return 0, err
293+
}
294+
writtenTotal += written
295+
}
296+
return writtenTotal, nil
297+
}
298+
299+
func (t quadKeyCodec[K1, K2, K3, K4]) DecodeNonTerminal(buffer []byte) (int, Quad[K1, K2, K3, K4], error) {
300+
readTotal := 0
301+
read, key1, err := t.keyCodec1.DecodeNonTerminal(buffer)
302+
if err != nil {
303+
return 0, Quad[K1, K2, K3, K4]{}, err
304+
}
305+
readTotal += read
306+
read, key2, err := t.keyCodec2.DecodeNonTerminal(buffer[readTotal:])
307+
if err != nil {
308+
return 0, Quad[K1, K2, K3, K4]{}, err
309+
}
310+
readTotal += read
311+
read, key3, err := t.keyCodec3.DecodeNonTerminal(buffer[readTotal:])
312+
if err != nil {
313+
return 0, Quad[K1, K2, K3, K4]{}, err
314+
}
315+
readTotal += read
316+
read, key4, err := t.keyCodec4.DecodeNonTerminal(buffer[readTotal:])
317+
if err != nil {
318+
return 0, Quad[K1, K2, K3, K4]{}, err
319+
}
320+
readTotal += read
321+
return readTotal, Join4(key1, key2, key3, key4), nil
322+
}
323+
324+
func (t quadKeyCodec[K1, K2, K3, K4]) SizeNonTerminal(key Quad[K1, K2, K3, K4]) int {
325+
size := 0
326+
if key.k1 != nil {
327+
size += t.keyCodec1.SizeNonTerminal(*key.k1)
328+
}
329+
if key.k2 != nil {
330+
size += t.keyCodec2.SizeNonTerminal(*key.k2)
331+
}
332+
if key.k3 != nil {
333+
size += t.keyCodec3.SizeNonTerminal(*key.k3)
334+
}
335+
if key.k4 != nil {
336+
size += t.keyCodec4.SizeNonTerminal(*key.k4)
337+
}
338+
return size
339+
}
340+
341+
// NewPrefixUntilQuadRange defines a collection query which ranges until the provided Quad prefix.
342+
// Unstable: this API might change in the future.
343+
func NewPrefixUntilQuadRange[K1, K2, K3, K4 any](k1 K1) Ranger[Quad[K1, K2, K3, K4]] {
344+
key := QuadPrefix[K1, K2, K3, K4](k1)
345+
return &Range[Quad[K1, K2, K3, K4]]{
346+
end: RangeKeyPrefixEnd(key),
347+
}
348+
}
349+
350+
// NewPrefixedQuadRange provides a Range for all keys prefixed with the given
351+
// first part of the Quad key.
352+
func NewPrefixedQuadRange[K1, K2, K3, K4 any](k1 K1) Ranger[Quad[K1, K2, K3, K4]] {
353+
key := QuadPrefix[K1, K2, K3, K4](k1)
354+
return &Range[Quad[K1, K2, K3, K4]]{
355+
start: RangeKeyExact(key),
356+
end: RangeKeyPrefixEnd(key),
357+
}
358+
}
359+
360+
// NewSuperPrefixedQuadRange provides a Range for all keys prefixed with the given
361+
// first and second parts of the Quad key.
362+
func NewSuperPrefixedQuadRange[K1, K2, K3, K4 any](k1 K1, k2 K2) Ranger[Quad[K1, K2, K3, K4]] {
363+
key := QuadSuperPrefix[K1, K2, K3, K4](k1, k2)
364+
return &Range[Quad[K1, K2, K3, K4]]{
365+
start: RangeKeyExact(key),
366+
end: RangeKeyPrefixEnd(key),
367+
}
368+
}
369+
370+
// NewSuperPrefixedQuadRange provides a Range for all keys prefixed with the given
371+
// first, second and third parts of the Quad key.
372+
func NewSuperPrefixedQuadRange3[K1, K2, K3, K4 any](k1 K1, k2 K2, k3 K3) Ranger[Quad[K1, K2, K3, K4]] {
373+
key := QuadSuperPrefix3[K1, K2, K3, K4](k1, k2, k3)
374+
return &Range[Quad[K1, K2, K3, K4]]{
375+
start: RangeKeyExact(key),
376+
end: RangeKeyPrefixEnd(key),
377+
}
378+
}

0 commit comments

Comments
 (0)
⚡