Skip to content

Commit 60fdfdb

Browse files
klauspostklauspost
authored
Add binary map key and shimming (#398)
* Allows for binary map keys * Allows for non-string map keys to be shimmed as well as cast from string type aliases. * Allows automatic conversion of some map key types to/from string. Adds directive `maps` * `//msgp:maps binkeys|shim|autoshim` Not adding the directive will retain current behaviour where maps without string keys are ignored. Adding `binkeys` will allow for binary marshaling of map keys. Types are inferred with standard rules. Binary encoding can be shimmed and can return any standard type. See `_generated/map_bin_key.go` for examples. Adding `shim` will require you to add shims to custom map key types. Shims must be with string type. See `_generated/map_shim_key.go` for examples. Adding `autoshim` will add automatic shimming to map keys that are `uint, uint8, uint16, uint32, uint64, int, int8, int16, int32, int64, bool, float32, float64, byte` - so these are written as a string. Standard `strconv.Parse/Format` functions are used. Aliased types are not supported, but `msgp:replace MyInt with:int` directive can be used and this will be picked up. Adds `witherr:true|false` to `//msgp:shim` directives, so shim de-serializing can return errors. Updates/Resolves: #331 #345 #305 #257 #232 Quirks: * Sizing is not perfect and can sometimes be collapsed from a range to constant. * These maps are not currently readable by "Interface" functions.
1 parent 17a251f commit 60fdfdb

18 files changed

+1037
-31
lines changed

‎_generated/def.go‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -342,3 +342,8 @@ type Flob struct {
342342
}
343343

344344
type Numberwang int8
345+
346+
//msgp:ignore ExternalString
347+
type ExternalString string
348+
type ExternalArr [4]byte
349+
type ExternalInt int

‎_generated/map_autoshim.go‎

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,35 @@
1+
package _generated
2+
3+
//go:generate msgp -v
4+
5+
//msgp:maps autoshim
6+
7+
//msgp:replace ExternalString with:string
8+
//msgp:replace ExternalInt with:int
9+
10+
type MyMapKeyStruct3 struct {
11+
MapString map[string]int `msg:",allownil"`
12+
MapString3 map[ExternalString]int `msg:",allownil"`
13+
MapString2 map[ExternalInt]int `msg:",allownil"`
14+
MapFloat32 map[float32]int `msg:",allownil"`
15+
MapFloat64 map[float64]int `msg:",allownil"`
16+
MapUint map[uint]int `msg:",allownil"`
17+
MapUint8 map[uint8]int `msg:",allownil"`
18+
MapUint16 map[uint16]int `msg:",allownil"`
19+
MapUint32 map[uint32]int `msg:",allownil"`
20+
MapUint64 map[uint64]int `msg:",allownil"`
21+
MapByte map[byte]int `msg:",allownil"`
22+
MapInt map[int]int `msg:",allownil"`
23+
MapInt8 map[int8]int `msg:",allownil"`
24+
MapInt16 map[int16]int `msg:",allownil"`
25+
MapInt32 map[int32]int `msg:",allownil"`
26+
MapInt64 map[int64]int `msg:",allownil"`
27+
MapBool map[bool]int `msg:",allownil"`
28+
MapMapInt map[int]map[int]int `msg:",allownil"`
29+
30+
// Maps with unsupported base types.
31+
MapArray map[[4]byte]int `msg:",allownil"` // Ignored
32+
MapArray4 map[[4]uint32]int `msg:",allownil"` // Ignored
33+
MapComplex64 map[complex64]int `msg:",allownil"` // Ignored
34+
MapComplex128 map[complex128]int `msg:",allownil"` // Ignored
35+
}

‎_generated/map_autoshim_test.go‎

Lines changed: 141 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,141 @@
1+
package _generated
2+
3+
import (
4+
"bytes"
5+
"math"
6+
"math/rand/v2"
7+
"reflect"
8+
"strconv"
9+
"testing"
10+
11+
"github.com/tinylib/msgp/msgp"
12+
)
13+
14+
func TestAutoShimKeys(t *testing.T) {
15+
rng := rand.New(rand.NewPCG(0, 0))
16+
17+
// Generate a bunch of random maps
18+
for i := range 500 {
19+
var test MyMapKeyStruct3
20+
if i != 0 {
21+
// Don't add anything to the first object
22+
test.MapString = make(map[string]int)
23+
test.MapString2 = make(map[ExternalInt]int)
24+
test.MapString3 = make(map[ExternalString]int)
25+
test.MapFloat32 = make(map[float32]int)
26+
test.MapFloat64 = make(map[float64]int)
27+
test.MapUint = make(map[uint]int)
28+
test.MapUint8 = make(map[uint8]int)
29+
test.MapUint16 = make(map[uint16]int)
30+
test.MapUint32 = make(map[uint32]int)
31+
test.MapUint64 = make(map[uint64]int)
32+
test.MapByte = make(map[byte]int)
33+
test.MapInt = make(map[int]int)
34+
test.MapInt8 = make(map[int8]int)
35+
test.MapInt16 = make(map[int16]int)
36+
test.MapInt32 = make(map[int32]int)
37+
test.MapInt64 = make(map[int64]int)
38+
test.MapBool = make(map[bool]int)
39+
test.MapMapInt = make(map[int]map[int]int)
40+
41+
for range rng.IntN(50) {
42+
test.MapString[string(strconv.Itoa(rng.IntN(math.MaxInt32)))] = rng.IntN(100)
43+
}
44+
for range rng.IntN(50) {
45+
test.MapString3[ExternalString(strconv.Itoa(rng.IntN(math.MaxInt32)))] = rng.IntN(100)
46+
}
47+
for range rng.IntN(50) {
48+
test.MapString2[ExternalInt(rng.IntN(math.MaxInt32))] = rng.IntN(100)
49+
}
50+
for range rng.IntN(50) {
51+
test.MapFloat32[float32(rng.Float32())] = rng.IntN(100)
52+
}
53+
for range rng.IntN(50) {
54+
test.MapFloat64[rng.Float64()] = rng.IntN(100)
55+
}
56+
for range rng.IntN(50) {
57+
test.MapUint[uint(rng.Uint64())] = rng.IntN(100)
58+
}
59+
for range rng.IntN(50) {
60+
test.MapUint8[uint8(rng.Uint64())] = rng.IntN(100)
61+
}
62+
for range rng.IntN(50) {
63+
test.MapUint16[uint16(rng.Uint64())] = rng.IntN(100)
64+
}
65+
for range rng.IntN(50) {
66+
test.MapUint32[rng.Uint32()] = rng.IntN(100)
67+
}
68+
for range rng.IntN(50) {
69+
test.MapUint64[rng.Uint64()] = rng.IntN(100)
70+
}
71+
for range rng.IntN(50) {
72+
test.MapByte[byte(uint8(rng.Uint64()))] = rng.IntN(100)
73+
}
74+
for range rng.IntN(50) {
75+
test.MapInt[rng.IntN(math.MaxInt32)] = rng.IntN(100)
76+
}
77+
for range rng.IntN(50) {
78+
test.MapInt8[int8(rng.IntN(int(math.MaxInt8)+1))] = rng.IntN(100)
79+
}
80+
for range rng.IntN(50) {
81+
test.MapInt16[int16(rng.IntN(int(math.MaxInt16)+1))] = rng.IntN(100)
82+
}
83+
for range rng.IntN(50) {
84+
test.MapInt32[int32(rng.IntN(int(math.MaxInt32)))] = rng.IntN(100)
85+
}
86+
for range rng.IntN(50) {
87+
// Use only non-negative values to stay within IntN capabilities
88+
test.MapInt64[int64(rng.IntN(int(math.MaxInt32)))] = rng.IntN(100)
89+
}
90+
if rng.IntN(2) == 0 {
91+
test.MapBool[true] = rng.IntN(100)
92+
}
93+
if rng.IntN(2) == 0 {
94+
test.MapBool[false] = rng.IntN(100)
95+
}
96+
97+
for range rng.IntN(50) {
98+
dst := make(map[int]int, 50)
99+
test.MapMapInt[rng.Int()] = dst
100+
for range rng.IntN(50) {
101+
dst[rng.Int()] = rng.IntN(100)
102+
}
103+
}
104+
}
105+
var encoded [][]byte
106+
b, err := test.MarshalMsg(nil)
107+
if err != nil {
108+
t.Fatal(err)
109+
}
110+
encoded = append(encoded, b)
111+
var buf bytes.Buffer
112+
en := msgp.NewWriter(&buf)
113+
err = test.EncodeMsg(en)
114+
if err != nil {
115+
t.Fatal(err)
116+
}
117+
err = en.Flush()
118+
if err != nil {
119+
t.Fatal(err)
120+
}
121+
encoded = append(encoded, buf.Bytes())
122+
for _, enc := range encoded {
123+
var decoded, decoded2 MyMapKeyStruct3
124+
_, err = decoded.UnmarshalMsg(enc)
125+
if err != nil {
126+
t.Fatal(err)
127+
}
128+
if !reflect.DeepEqual(&decoded, &test) {
129+
t.Errorf("decoded != test")
130+
}
131+
dec := msgp.NewReader(bytes.NewReader(enc))
132+
err = decoded2.DecodeMsg(dec)
133+
if err != nil {
134+
t.Fatal(err)
135+
}
136+
if !reflect.DeepEqual(&decoded2, &test) {
137+
t.Errorf("decoded2 != test")
138+
}
139+
}
140+
}
141+
}

‎_generated/map_bin_key.go‎

Lines changed: 81 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,81 @@
1+
package _generated
2+
3+
import (
4+
"encoding/hex"
5+
"time"
6+
)
7+
8+
//go:generate msgp -unexported -v
9+
10+
//msgp:maps binkeys
11+
12+
type ArrayMapKey [4]byte
13+
14+
//msgp:replace ExternalArr with:[4]byte
15+
16+
//msgp:replace ExternalString with:string
17+
18+
type mapKeyBytes2 [8]byte
19+
20+
//msgp:shim mapKeyBytes2 as:string using:hexEncode2/hexDecode2 witherr:false
21+
22+
type mapKeyShimmed time.Duration
23+
24+
//msgp:shim mapKeyShimmed as:[]byte using:durEncode/durDecode witherr:true
25+
26+
type MyStringType string
27+
28+
type MyMapKeyStruct2 struct {
29+
MapString map[string]int `msg:",allownil"`
30+
MapString2 map[MyStringType]int `msg:",allownil"`
31+
MapString3 map[ExternalString]int `msg:",allownil"`
32+
MapString4 map[mapKeyBytes2]int `msg:",allownil"`
33+
MapFloat32 map[float32]int `msg:",allownil"`
34+
MapFloat64 map[float64]int `msg:",allownil"`
35+
MapComplex64 map[complex64]int `msg:",allownil"`
36+
MapComplex128 map[complex128]int `msg:",allownil"`
37+
MapUint map[uint]int `msg:",allownil"`
38+
MapUint8 map[uint8]int `msg:",allownil"`
39+
MapUint16 map[uint16]int `msg:",allownil"`
40+
MapUint32 map[uint32]int `msg:",allownil"`
41+
MapUint64 map[uint64]int `msg:",allownil"`
42+
MapByte map[byte]int `msg:",allownil"`
43+
MapInt map[int]int `msg:",allownil"`
44+
MapInt8 map[int8]int `msg:",allownil"`
45+
MapInt16 map[int16]int `msg:",allownil"`
46+
MapInt32 map[int32]int `msg:",allownil"`
47+
MapInt64 map[int64]int `msg:",allownil"`
48+
MapBool map[bool]int `msg:",allownil"`
49+
MapMapInt map[int]map[int]int `msg:",allownil"`
50+
51+
// Maps with array keys
52+
MapArray map[[4]byte]int `msg:",allownil"`
53+
MapArray2 map[ArrayMapKey]int `msg:",allownil"`
54+
MapArray3 map[ExternalArr]int `msg:",allownil"`
55+
MapArray4 map[[4]uint32]int `msg:",allownil"`
56+
57+
// Maps with shimmed types
58+
MapDuration map[mapKeyShimmed]int `msg:",allownil"`
59+
}
60+
61+
func hexEncode2(b mapKeyBytes2) string {
62+
return hex.EncodeToString(b[:])
63+
}
64+
65+
func hexDecode2(s string) mapKeyBytes2 {
66+
var b [8]byte
67+
_, err := hex.Decode(b[:], []byte(s))
68+
if err != nil {
69+
panic(err)
70+
}
71+
return b
72+
}
73+
74+
func durEncode(v mapKeyShimmed) []byte {
75+
return []byte(time.Duration(v).String())
76+
}
77+
78+
func durDecode(b []byte) (mapKeyShimmed, error) {
79+
v, err := time.ParseDuration(string(b))
80+
return mapKeyShimmed(v), err
81+
}

0 commit comments

Comments
 (0)