Skip to content

Commit

Permalink
refactor(countrw): move to iout
Browse files Browse the repository at this point in the history
  • Loading branch information
sentriz committed Apr 12, 2022
1 parent d7655cb commit fd211d7
Show file tree
Hide file tree
Showing 6 changed files with 132 additions and 27 deletions.
26 changes: 0 additions & 26 deletions countrw/cr.go

This file was deleted.

23 changes: 23 additions & 0 deletions iout/copy_range.go
@@ -0,0 +1,23 @@
package iout

import (
"errors"
"fmt"
"io"
)

func CopyRange(w io.Writer, r io.Reader, start, length int64) error {
if _, err := io.CopyN(io.Discard, r, start); err != nil && !errors.Is(err, io.EOF) {
return fmt.Errorf("discard %d: %w", start, err)
}
if length == 0 {
if _, err := io.Copy(w, r); err != nil && !errors.Is(err, io.EOF) {
return fmt.Errorf("direct copy: %w", err)
}
return nil
}
if _, err := io.CopyN(w, io.MultiReader(r, NewNullReader()), length); err != nil && !errors.Is(err, io.EOF) {
return fmt.Errorf("copy %d: %w", length, err)
}
return nil
}
42 changes: 42 additions & 0 deletions iout/copy_range_test.go
@@ -0,0 +1,42 @@
package iout_test

import (
"bytes"
"testing"

"github.com/matryer/is"
"go.senan.xyz/gonic/iout"
)

func TestCopyRange(t *testing.T) {
t.Parallel()
is := is.New(t)

realLength := 50
cr := func(start, length int64) []byte {
is.Helper()
var data []byte
for i := 0; i < realLength; i++ {
data = append(data, byte(i%10))
}
var buff bytes.Buffer
is.NoErr(iout.CopyRange(&buff, bytes.NewReader(data), start, length))
return buff.Bytes()
}

// range
is.Equal(len(cr(0, 50)), 50)
is.Equal(len(cr(10, 10)), 10)
is.Equal(cr(10, 10)[0], byte(0))
is.Equal(cr(10, 10)[5], byte(5))
is.Equal(cr(25, 35)[0], byte(5))
is.Equal(cr(25, 35)[5], byte(0))

// 0 padding
is.Equal(len(cr(0, 5000)), 5000)
is.Equal(cr(0, 5000)[50:], make([]byte, 5000-50))

// no bound
is.Equal(len(cr(0, 0)), 50)
is.Equal(len(cr(50, 0)), 0)
}
22 changes: 21 additions & 1 deletion countrw/cw.go → iout/countrw.go
@@ -1,10 +1,30 @@
package countrw
package iout

import (
"io"
"sync/atomic"
)

type CountReader struct {
r io.Reader
c *uint64
}

func NewCountReader(r io.Reader) *CountReader {
return &CountReader{r: r, c: new(uint64)}
}

func (c *CountReader) Reset() { atomic.StoreUint64(c.c, 0) }
func (c *CountReader) Count() uint64 { return atomic.LoadUint64(c.c) }

func (c *CountReader) Read(p []byte) (int, error) {
n, err := c.r.Read(p)
atomic.AddUint64(c.c, uint64(n))
return n, err
}

var _ io.Reader = (*CountReader)(nil)

type CountWriter struct {
r io.Writer
c *uint64
Expand Down
18 changes: 18 additions & 0 deletions iout/null_reader.go
@@ -0,0 +1,18 @@
package iout

import "io"

type nullReader struct{}

func NewNullReader() io.Reader {
return &nullReader{}
}

func (*nullReader) Read(p []byte) (n int, err error) {
for b := range p {
p[b] = 0
}
return len(p), nil
}

var _ io.Reader = (*nullReader)(nil)
28 changes: 28 additions & 0 deletions iout/tee_closer.go
@@ -0,0 +1,28 @@
package iout

import "io"

type teeCloser struct {
r io.ReadCloser
w io.WriteCloser
}

func NewTeeCloser(r io.ReadCloser, w io.WriteCloser) io.ReadCloser {
return &teeCloser{r, w}
}

func (t *teeCloser) Read(p []byte) (int, error) {
n, err := t.r.Read(p)
if n > 0 {
if n, err := t.w.Write(p[:n]); err != nil {
return n, err
}
}
return n, err
}

func (t *teeCloser) Close() error {
t.r.Close()
t.w.Close()
return nil
}

0 comments on commit fd211d7

Please sign in to comment.