Skip to content

Commit c6362fd

Browse files
nerg4lnerg4l
committed
feat: rethink riff writer as io.WriterAt
1 parent f3e1196 commit c6362fd

File tree

2 files changed

+135
-47
lines changed

2 files changed

+135
-47
lines changed

‎internal/riff/riff.go‎

Lines changed: 78 additions & 40 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,6 @@
88
package riff // import "golang.org/x/image/riff"
99

1010
import (
11-
"bytes"
1211
"encoding/binary"
1312
"errors"
1413
"io"
@@ -24,6 +23,7 @@ var (
2423
errShortChunk = errors.New("riff: short chunk")
2524
errStaleReader = errors.New("riff: stale reader")
2625
errStaleWriter = errors.New("riff: stale writer")
26+
errInvalidOffset = errors.New("riff: invalid offset")
2727
)
2828

2929
// FourCC is a four character code.
@@ -126,67 +126,94 @@ func (r *Reader) Next() (FourCC, uint32, *ChunkReader, error) {
126126
return id, l, &cr, nil
127127
}
128128

129-
func NewWriter(w io.Writer, fileType FourCC) (data *Writer, err error) {
130-
if _, err := w.Write([]byte{'R', 'I', 'F', 'F'}); err != nil {
129+
func NewWriter(w io.WriterAt, fileType FourCC) (data *Writer, err error) {
130+
if _, err := w.WriteAt([]byte{'R', 'I', 'F', 'F'}, 0); err != nil {
131131
return nil, err
132132
}
133-
ww, err := NewListWriter(w, fileType)
133+
ww, err := NewListWriter(io.NewOffsetWriter(w, 8), fileType)
134134
if err != nil {
135135
return ww, err
136136
}
137-
ww.root = true
137+
ww.root = w
138138
return ww, err
139139
}
140140

141-
func NewListWriter(chunkData io.Writer, listType FourCC) (listw *Writer, err error) {
142-
w := &Writer{w: chunkData}
143-
_, w.err = w.buf.Write(listType[:])
141+
func NewListWriter(chunkData io.WriterAt, listType FourCC) (listw *Writer, err error) {
142+
w := &Writer{w: chunkData, len: 4}
143+
_, w.err = chunkData.WriteAt(listType[:], 0)
144144
return w, w.err
145145
}
146146

147147
type Writer struct {
148-
w io.Writer
149-
err error
148+
w io.WriterAt
149+
root io.WriterAt
150+
err error
150151

151-
chunk *ChunkWriter
152-
buf bytes.Buffer
153-
root bool
152+
chunk *ChunkWriter
153+
closed bool
154+
len int64
154155
}
155156

156157
type ChunkWriter struct {
157158
listw *Writer
158-
buf bytes.Buffer
159+
len int64
160+
161+
base int64 // the original offset
162+
off int64 // the current offset
159163
}
160164

161165
func (w *ChunkWriter) Write(p []byte) (n int, err error) {
162166
if w.listw.chunk != w {
163167
return 0, errStaleWriter
164168
}
165-
if (w.buf.Len() + len(p)) > math.MaxUint32 {
169+
n, err = w.WriteAt(p, w.off)
170+
w.off += int64(n)
171+
return
172+
}
173+
174+
func (w *ChunkWriter) WriteAt(p []byte, off int64) (n int, err error) {
175+
if w.listw.chunk != w {
176+
return 0, errStaleWriter
177+
}
178+
if off < 0 {
179+
return 0, errInvalidOffset
180+
}
181+
if (w.base + w.len + int64(len(p))) > math.MaxUint32 {
166182
return 0, ErrDataTooLong
167183
}
168-
return w.buf.Write(p)
184+
185+
n, err = w.listw.w.WriteAt(p, off+w.base)
186+
l := off + int64(n)
187+
if l > w.len {
188+
w.len = l
189+
w.listw.len = w.base + w.len
190+
}
191+
return n, err
169192
}
170193

171194
func (w *ChunkWriter) close() error {
172195
var sizebuf [4]byte
173-
l := uint32(w.buf.Len())
196+
l := uint32(w.len)
174197
binary.LittleEndian.PutUint32(sizebuf[:], l)
175-
if _, err := w.listw.buf.Write(sizebuf[:]); err != nil {
198+
if _, err := w.listw.w.WriteAt(sizebuf[:], w.base-4); err != nil {
176199
return err
177200
}
178201
if l&1 == 1 {
179-
if err := w.buf.WriteByte(0); err != nil {
202+
n, err := w.listw.w.WriteAt([]byte{0}, w.base+w.len)
203+
w.len += int64(n)
204+
w.listw.len += int64(n)
205+
if err != nil {
180206
return err
181207
}
182208
}
183-
_, err := w.buf.WriteTo(&w.listw.buf)
184-
if err != nil {
185-
return err
186-
}
187-
return err
209+
return nil
188210
}
189211

212+
var sizePlaceholder = []byte{0, 0, 0, 0}
213+
214+
// Next returns a chunk writer which behaves as an io.OffsetWriter.
215+
//
216+
// Calling Next closes the returned writer.
190217
func (w *Writer) Next(id FourCC) (*ChunkWriter, error) {
191218
if w.err != nil {
192219
return nil, w.err
@@ -199,44 +226,55 @@ func (w *Writer) Next(id FourCC) (*ChunkWriter, error) {
199226
w.chunk = nil
200227
}
201228

202-
if l := w.buf.Len(); l&1 == 1 {
203-
if _, w.err = w.buf.Write([]byte{0}); w.err != nil {
204-
return nil, w.err
205-
}
229+
n, err := w.w.WriteAt(id[:], w.len)
230+
w.len += int64(n)
231+
if err != nil {
232+
w.err = err
233+
return nil, w.err
206234
}
207-
if _, w.err = w.buf.Write(id[:]); w.err != nil {
235+
n, err = w.w.WriteAt(sizePlaceholder, w.len)
236+
w.len += int64(n)
237+
if err != nil {
238+
w.err = err
208239
return nil, w.err
209240
}
210-
w.chunk = &ChunkWriter{listw: w}
241+
w.chunk = &ChunkWriter{listw: w, base: w.len}
211242
return w.chunk, nil
212243
}
213244

245+
// Close closes the list by closing the last writer returned by Next.
214246
func (w *Writer) Close() error {
247+
if w.err != nil {
248+
return w.err
249+
}
250+
if w.closed {
251+
return nil
252+
}
253+
215254
if w.chunk != nil {
216255
if w.err = w.chunk.close(); w.err != nil {
217256
return w.err
218257
}
219258
w.chunk = nil
220259
}
221260

222-
if w.root {
261+
if w.root != nil {
223262
// TODO: think of a better way of writing the size of the root element,
224263
// but not write it for sublists because that occurs in ChunkWriter.
225264
var sizebuf [4]byte
226-
l := uint32(w.buf.Len())
265+
l := uint32(w.len)
227266
binary.LittleEndian.PutUint32(sizebuf[:], l)
228-
if _, err := w.w.Write(sizebuf[:]); err != nil {
229-
return err
267+
if _, w.err = w.root.WriteAt(sizebuf[:], 4); w.err != nil {
268+
return w.err
230269
}
231270
if l&1 == 1 {
232-
if err := w.buf.WriteByte(0); err != nil {
271+
n, err := w.w.WriteAt([]byte{0}, w.len)
272+
w.len += int64(n)
273+
if err != nil {
233274
return err
234275
}
235276
}
236277
}
237-
_, err := w.buf.WriteTo(w.w)
238-
if err != nil {
239-
return err
240-
}
241-
return err
278+
w.closed = true
279+
return nil
242280
}

‎internal/riff/riff_test.go‎

Lines changed: 57 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -7,8 +7,8 @@ import (
77
)
88

99
func Test(t *testing.T) {
10-
var buf bytes.Buffer
11-
lw, err := NewWriter(&buf, FourCC{'T', 'E', 'S', 'T'})
10+
var writeAt bytesWriterAt
11+
lw, err := NewWriter(&writeAt, FourCC{'T', 'E', 'S', 'T'})
1212
if err != nil {
1313
t.Fatal(err)
1414
}
@@ -30,12 +30,21 @@ func Test(t *testing.T) {
3030
t.Fatal(err)
3131
}
3232
}
33+
{ // ck02
34+
w, err := lw.Next(FourCC{'c', 'k', '0', '2'})
35+
if err != nil {
36+
t.Fatal(err)
37+
}
38+
if _, err := w.Write([]byte("Hello world")); err != nil {
39+
t.Fatal(err)
40+
}
41+
}
3342
if err := lw.Close(); err != nil {
3443
t.Fatal(err)
3544
}
3645
}
37-
{ // ck02
38-
w, err := lw.Next(FourCC{'c', 'k', '0', '2'})
46+
{ // ck03
47+
w, err := lw.Next(FourCC{'c', 'k', '0', '3'})
3948
if err != nil {
4049
t.Fatal(err)
4150
}
@@ -47,7 +56,7 @@ func Test(t *testing.T) {
4756
t.Fatal(err)
4857
}
4958

50-
id, lr, err := NewReader(&buf)
59+
id, lr, err := NewReader(bytes.NewReader(writeAt.buf))
5160
if err != nil {
5261
t.Fatal(err)
5362
}
@@ -85,16 +94,32 @@ func Test(t *testing.T) {
8594
t.Errorf("Read() got %s, want %s", got, want)
8695
}
8796
}
97+
{ // ck02
98+
id, l, r, err := lr.Next()
99+
if err != nil {
100+
t.Fatal(err)
101+
}
102+
if got, want := id, (FourCC{'c', 'k', '0', '2'}); got != want {
103+
t.Errorf("Next() got %v, want %v", got, want)
104+
}
105+
b := make([]byte, l)
106+
if _, err := r.Read(b); err != nil {
107+
t.Fatal(err)
108+
}
109+
if got, want := b, []byte("Hello world"); !bytes.Equal(got, want) {
110+
t.Errorf("Read() got %s, want %s", got, want)
111+
}
112+
}
88113
if _, _, _, err := lr.Next(); err != io.EOF {
89114
t.Errorf("Next() got %v, want EOF", err)
90115
}
91116
}
92-
{ // ck02
117+
{ // ck03
93118
id, l, r, err := lr.Next()
94119
if err != nil {
95120
t.Fatal(err)
96121
}
97-
if got, want := id, (FourCC{'c', 'k', '0', '2'}); got != want {
122+
if got, want := id, (FourCC{'c', 'k', '0', '3'}); got != want {
98123
t.Errorf("Next() got %v, want %v", got, want)
99124
}
100125
b := make([]byte, l)
@@ -109,3 +134,28 @@ func Test(t *testing.T) {
109134
t.Errorf("Next() got %v, want EOF", err)
110135
}
111136
}
137+
138+
type bytesWriterAt struct {
139+
buf []byte
140+
}
141+
142+
func (w *bytesWriterAt) grow(n int64) {
143+
if n > int64(cap(w.buf)) {
144+
m := max(int64(cap(w.buf))*2, 8)
145+
for m < n {
146+
m *= 2
147+
}
148+
tmp := make([]byte, m)
149+
copy(tmp, w.buf)
150+
w.buf = tmp[:n]
151+
}
152+
if n > int64(len(w.buf)) {
153+
w.buf = w.buf[:n]
154+
}
155+
}
156+
157+
func (b *bytesWriterAt) WriteAt(p []byte, off int64) (n int, err error) {
158+
b.grow(off + int64(len(p)))
159+
copy(b.buf[off:off+int64(len(p))], p)
160+
return len(p), nil
161+
}

0 commit comments

Comments
 (0)