forked from ava-labs/avalanchego
-
Notifications
You must be signed in to change notification settings - Fork 4
/
codec.go
151 lines (129 loc) · 3.98 KB
/
codec.go
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
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
// Copyright (C) 2019-2024, Ava Labs, Inc. All rights reserved.
// See the file LICENSE for licensing terms.
package hierarchycodec
import (
"fmt"
"reflect"
"sync"
"time"
"github.com/MetalBlockchain/metalgo/codec"
"github.com/MetalBlockchain/metalgo/codec/reflectcodec"
"github.com/MetalBlockchain/metalgo/utils/bimap"
"github.com/MetalBlockchain/metalgo/utils/wrappers"
)
const (
// default max length of a slice being marshalled by Marshal(). Should be <= math.MaxUint32.
defaultMaxSliceLength = 256 * 1024
)
var (
_ Codec = (*hierarchyCodec)(nil)
_ codec.Codec = (*hierarchyCodec)(nil)
_ codec.Registry = (*hierarchyCodec)(nil)
_ codec.GeneralCodec = (*hierarchyCodec)(nil)
)
// Codec marshals and unmarshals
type Codec interface {
codec.Registry
codec.Codec
SkipRegistrations(int)
NextGroup()
}
type typeID struct {
groupID uint16
typeID uint16
}
// Codec handles marshaling and unmarshaling of structs
type hierarchyCodec struct {
codec.Codec
lock sync.RWMutex
currentGroupID uint16
nextTypeID uint16
registeredTypes *bimap.BiMap[typeID, reflect.Type]
}
// New returns a new, concurrency-safe codec
func New(durangoTime time.Time, tagNames []string, maxSliceLen uint32) Codec {
hCodec := &hierarchyCodec{
currentGroupID: 0,
nextTypeID: 0,
registeredTypes: bimap.New[typeID, reflect.Type](),
}
hCodec.Codec = reflectcodec.New(hCodec, tagNames, durangoTime, maxSliceLen)
return hCodec
}
// NewDefault returns a new codec with reasonable default values
func NewDefault(durangoTime time.Time) Codec {
return New(durangoTime, []string{reflectcodec.DefaultTagName}, defaultMaxSliceLength)
}
// SkipRegistrations some number of type IDs
func (c *hierarchyCodec) SkipRegistrations(num int) {
c.lock.Lock()
c.nextTypeID += uint16(num)
c.lock.Unlock()
}
// NextGroup moves to the next group registry
func (c *hierarchyCodec) NextGroup() {
c.lock.Lock()
c.currentGroupID++
c.nextTypeID = 0
c.lock.Unlock()
}
// RegisterType is used to register types that may be unmarshaled into an interface
// [val] is a value of the type being registered
func (c *hierarchyCodec) RegisterType(val interface{}) error {
c.lock.Lock()
defer c.lock.Unlock()
valType := reflect.TypeOf(val)
if c.registeredTypes.HasValue(valType) {
return fmt.Errorf("%w: %v", codec.ErrDuplicateType, valType)
}
valTypeID := typeID{
groupID: c.currentGroupID,
typeID: c.nextTypeID,
}
c.nextTypeID++
c.registeredTypes.Put(valTypeID, valType)
return nil
}
func (*hierarchyCodec) PrefixSize(reflect.Type) int {
// see PackPrefix implementation
return wrappers.ShortLen + wrappers.ShortLen
}
func (c *hierarchyCodec) PackPrefix(p *wrappers.Packer, valueType reflect.Type) error {
c.lock.RLock()
defer c.lock.RUnlock()
typeID, ok := c.registeredTypes.GetKey(valueType) // Get the type ID of the value being marshaled
if !ok {
return fmt.Errorf("can't marshal unregistered type %q", valueType)
}
// Pack type ID so we know what to unmarshal this into
p.PackShort(typeID.groupID)
p.PackShort(typeID.typeID)
return p.Err
}
func (c *hierarchyCodec) UnpackPrefix(p *wrappers.Packer, valueType reflect.Type) (reflect.Value, error) {
c.lock.RLock()
defer c.lock.RUnlock()
groupID := p.UnpackShort() // Get the group ID
typeIDShort := p.UnpackShort() // Get the type ID
if p.Err != nil {
return reflect.Value{}, fmt.Errorf("couldn't unmarshal interface: %w", p.Err)
}
t := typeID{
groupID: groupID,
typeID: typeIDShort,
}
// Get a type that implements the interface
implementingType, ok := c.registeredTypes.GetValue(t)
if !ok {
return reflect.Value{}, fmt.Errorf("couldn't unmarshal interface: unknown type ID %+v", t)
}
// Ensure type actually does implement the interface
if !implementingType.Implements(valueType) {
return reflect.Value{}, fmt.Errorf("couldn't unmarshal interface: %s %w %s",
implementingType,
codec.ErrDoesNotImplementInterface,
valueType,
)
}
return reflect.New(implementingType).Elem(), nil // instance of the proper type
}