-
Notifications
You must be signed in to change notification settings - Fork 9
/
zero_copy_sink.go
240 lines (201 loc) · 5.89 KB
/
zero_copy_sink.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
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
/*
* Copyright (C) 2018 The ontology Authors
* This file is part of The ontology library.
*
* The ontology is free software: you can redistribute it and/or modify
* it under the terms of the GNU Lesser General Public License as published by
* the Free Software Foundation, either version 3 of the License, or
* (at your option) any later version.
*
* The ontology is distributed in the hope that it will be useful,
* but WITHOUT ANY WARRANTY; without even the implied warranty of
* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
* GNU Lesser General Public License for more details.
*
* You should have received a copy of the GNU Lesser General Public License
* along with The ontology. If not, see <http://www.gnu.org/licenses/>.
*/
package codec
import (
"bytes"
"encoding/binary"
"errors"
"github.com/laizy/web3"
)
type ZeroCopySink struct {
buf []byte
}
// tryGrowByReslice is a inlineable version of grow for the fast-case where the
// internal buffer only needs to be resliced.
// It returns the index where bytes should be written and whether it succeeded.
func (self *ZeroCopySink) tryGrowByReslice(n int) (int, bool) {
if l := len(self.buf); n <= cap(self.buf)-l {
self.buf = self.buf[:l+n]
return l, true
}
return 0, false
}
const maxInt = int(^uint(0) >> 1)
// grow grows the buffer to guarantee space for n more bytes.
// It returns the index where bytes should be written.
// If the buffer can't grow it will panic with ErrTooLarge.
func (self *ZeroCopySink) grow(n int) int {
// Try to grow by means of a reslice.
if i, ok := self.tryGrowByReslice(n); ok {
return i
}
l := len(self.buf)
c := cap(self.buf)
if c > maxInt-c-n {
panic(ErrTooLarge)
}
// Not enough space anywhere, we need to allocate.
buf := makeSlice(2*c + n)
copy(buf, self.buf)
self.buf = buf[:l+n]
return l
}
func (self *ZeroCopySink) WriteBytes(p []byte) *ZeroCopySink {
data := self.NextBytes(uint64(len(p)))
copy(data, p)
return self
}
func (self *ZeroCopySink) Size() uint64 { return uint64(len(self.buf)) }
func (self *ZeroCopySink) NextBytes(n uint64) (data []byte) {
m, ok := self.tryGrowByReslice(int(n))
if !ok {
m = self.grow(int(n))
}
data = self.buf[m:]
return
}
// Backs up a number of bytes, so that the next call to NextXXX() returns data again
// that was already returned by the last call to NextXXX().
func (self *ZeroCopySink) BackUp(n uint64) {
l := len(self.buf) - int(n)
self.buf = self.buf[:l]
}
func (self *ZeroCopySink) WriteUint8(data uint8) *ZeroCopySink {
buf := self.NextBytes(1)
buf[0] = data
return self
}
func (self *ZeroCopySink) WriteByte(c byte) *ZeroCopySink {
self.WriteUint8(c)
return self
}
func (self *ZeroCopySink) WriteBool(data bool) *ZeroCopySink {
if data {
self.WriteByte(1)
} else {
self.WriteByte(0)
}
return self
}
func (self *ZeroCopySink) WriteUint16(data uint16) *ZeroCopySink {
buf := self.NextBytes(2)
binary.LittleEndian.PutUint16(buf, data)
return self
}
func (self *ZeroCopySink) WriteUint32(data uint32) *ZeroCopySink {
buf := self.NextBytes(4)
binary.LittleEndian.PutUint32(buf, data)
return self
}
func (self *ZeroCopySink) WriteUint64(data uint64) *ZeroCopySink {
buf := self.NextBytes(8)
binary.LittleEndian.PutUint64(buf, data)
return self
}
func (self *ZeroCopySink) WriteUint16BE(data uint16) *ZeroCopySink {
buf := self.NextBytes(2)
binary.BigEndian.PutUint16(buf, data)
return self
}
func (self *ZeroCopySink) WriteUint32BE(data uint32) *ZeroCopySink {
buf := self.NextBytes(4)
binary.BigEndian.PutUint32(buf, data)
return self
}
func (self *ZeroCopySink) WriteUint64BE(data uint64) *ZeroCopySink {
buf := self.NextBytes(8)
binary.BigEndian.PutUint64(buf, data)
return self
}
func (self *ZeroCopySink) WriteInt64(data int64) *ZeroCopySink {
return self.WriteUint64(uint64(data))
}
func (self *ZeroCopySink) WriteInt32(data int32) *ZeroCopySink {
return self.WriteUint32(uint32(data))
}
func (self *ZeroCopySink) WriteInt16(data int16) *ZeroCopySink {
return self.WriteUint16(uint16(data))
}
func (self *ZeroCopySink) WriteInt64BE(data int64) *ZeroCopySink {
return self.WriteUint64BE(uint64(data))
}
func (self *ZeroCopySink) WriteInt32BE(data int32) *ZeroCopySink {
return self.WriteUint32BE(uint32(data))
}
func (self *ZeroCopySink) WriteInt16BE(data int16) *ZeroCopySink {
return self.WriteUint16BE(uint16(data))
}
func (self *ZeroCopySink) WriteVarBytes(data []byte) (size uint64) {
l := uint64(len(data))
size = self.WriteVarUint(l) + l
self.WriteBytes(data)
return
}
func (self *ZeroCopySink) WriteString(data string) (size uint64) {
return self.WriteVarBytes([]byte(data))
}
func (self *ZeroCopySink) WriteAddress(addr web3.Address) *ZeroCopySink {
self.WriteBytes(addr[:])
return self
}
func (self *ZeroCopySink) WriteHash(hash web3.Hash) *ZeroCopySink {
self.WriteBytes(hash[:])
return self
}
func (self *ZeroCopySink) WriteVarUint(data uint64) (size uint64) {
buf := self.NextBytes(9)
if data < 0xFD {
buf[0] = uint8(data)
size = 1
} else if data <= 0xFFFF {
buf[0] = 0xFD
binary.LittleEndian.PutUint16(buf[1:], uint16(data))
size = 3
} else if data <= 0xFFFFFFFF {
buf[0] = 0xFE
binary.LittleEndian.PutUint32(buf[1:], uint32(data))
size = 5
} else {
buf[0] = 0xFF
binary.LittleEndian.PutUint64(buf[1:], uint64(data))
size = 9
}
self.BackUp(9 - size)
return
}
// NewReader returns a new ZeroCopySink reading from b.
func NewZeroCopySink(b []byte) *ZeroCopySink {
if b == nil {
b = make([]byte, 0, 512)
}
return &ZeroCopySink{b}
}
func (self *ZeroCopySink) Bytes() []byte { return self.buf }
func (self *ZeroCopySink) Reset() { self.buf = self.buf[:0] }
var ErrTooLarge = errors.New("bytes.Buffer: too large")
// makeSlice allocates a slice of size n. If the allocation fails, it panics
// with ErrTooLarge.
func makeSlice(n int) []byte {
// If the make fails, give a known error.
defer func() {
if recover() != nil {
panic(bytes.ErrTooLarge)
}
}()
return make([]byte, n)
}