From d60aa5c64a870fee30a93f96dcd22b6e663d420a Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Mon, 1 Jan 2024 11:34:01 -0500 Subject: [PATCH] feat: add buffer truncation --- pkg/btree/buffer.go | 11 +++++++++++ pkg/btree/buffer_test.go | 13 +++++++++++++ 2 files changed, 24 insertions(+) diff --git a/pkg/btree/buffer.go b/pkg/btree/buffer.go index c468a7b1..0fb80a9c 100644 --- a/pkg/btree/buffer.go +++ b/pkg/btree/buffer.go @@ -50,4 +50,15 @@ func (b *seekableBuffer) Read(p []byte) (int, error) { return n, nil } +func (b *seekableBuffer) Truncate(size int64) error { + if size < 0 { + return io.ErrShortBuffer + } + if size > int64(len(b.buf)) { + return io.ErrShortWrite + } + b.buf = b.buf[:size] + return nil +} + var _ io.ReadWriteSeeker = &seekableBuffer{} diff --git a/pkg/btree/buffer_test.go b/pkg/btree/buffer_test.go index ca8f5d40..7552c22b 100644 --- a/pkg/btree/buffer_test.go +++ b/pkg/btree/buffer_test.go @@ -93,4 +93,17 @@ func TestSeekableBuffer(t *testing.T) { t.Fatalf("expected to read 'llo', read %s", string(buf)) } }) + + t.Run("truncate", func(t *testing.T) { + b := newSeekableBuffer() + if _, err := b.Write([]byte("hello")); err != nil { + t.Fatal(err) + } + if err := b.Truncate(3); err != nil { + t.Fatal(err) + } + if string(b.buf) != "hel" { + t.Fatalf("expected to truncate to 'hel', truncated to %s", string(b.buf)) + } + }) }