1// Copyright 2022 The Go Authors. All rights reserved.
2// Use of this source code is governed by a BSD-style
3// license that can be found in the LICENSE file.
4
5package slicewriter
6
7import (
8	"fmt"
9	"io"
10)
11
12// WriteSeeker is a helper object that implements the io.WriteSeeker
13// interface. Clients can create a WriteSeeker, make a series of Write
14// calls to add data to it (and possibly Seek calls to update
15// previously written portions), then finally invoke BytesWritten() to
16// get a pointer to the constructed byte slice.
17type WriteSeeker struct {
18	payload []byte
19	off     int64
20}
21
22func (sws *WriteSeeker) Write(p []byte) (n int, err error) {
23	amt := len(p)
24	towrite := sws.payload[sws.off:]
25	if len(towrite) < amt {
26		sws.payload = append(sws.payload, make([]byte, amt-len(towrite))...)
27		towrite = sws.payload[sws.off:]
28	}
29	copy(towrite, p)
30	sws.off += int64(amt)
31	return amt, nil
32}
33
34// Seek repositions the read/write position of the WriteSeeker within
35// its internally maintained slice. Note that it is not possible to
36// expand the size of the slice using SEEK_SET; trying to seek outside
37// the slice will result in an error.
38func (sws *WriteSeeker) Seek(offset int64, whence int) (int64, error) {
39	switch whence {
40	case io.SeekStart:
41		if sws.off != offset && (offset < 0 || offset > int64(len(sws.payload))) {
42			return 0, fmt.Errorf("invalid seek: new offset %d (out of range [0 %d]", offset, len(sws.payload))
43		}
44		sws.off = offset
45		return offset, nil
46	case io.SeekCurrent:
47		newoff := sws.off + offset
48		if newoff != sws.off && (newoff < 0 || newoff > int64(len(sws.payload))) {
49			return 0, fmt.Errorf("invalid seek: new offset %d (out of range [0 %d]", newoff, len(sws.payload))
50		}
51		sws.off += offset
52		return sws.off, nil
53	case io.SeekEnd:
54		newoff := int64(len(sws.payload)) + offset
55		if newoff != sws.off && (newoff < 0 || newoff > int64(len(sws.payload))) {
56			return 0, fmt.Errorf("invalid seek: new offset %d (out of range [0 %d]", newoff, len(sws.payload))
57		}
58		sws.off = newoff
59		return sws.off, nil
60	}
61	// other modes not supported
62	return 0, fmt.Errorf("unsupported seek mode %d", whence)
63}
64
65// BytesWritten returns the underlying byte slice for the WriteSeeker,
66// containing the data written to it via Write/Seek calls.
67func (sws *WriteSeeker) BytesWritten() []byte {
68	return sws.payload
69}
70
71func (sws *WriteSeeker) Read(p []byte) (n int, err error) {
72	amt := len(p)
73	toread := sws.payload[sws.off:]
74	if len(toread) < amt {
75		amt = len(toread)
76	}
77	copy(p, toread)
78	sws.off += int64(amt)
79	return amt, nil
80}
81