Browse Source

Merge pull request #5 from arya/master

Return the correct length after flushing a buffered write.
Pierre Curto 10 years ago
parent
commit
222ab1fc79
2 changed files with 31 additions and 1 deletions
  1. 17 0
      lz4_test.go
  2. 14 1
      writer.go

+ 17 - 0
lz4_test.go

@@ -565,3 +565,20 @@ func TestSkippable(t *testing.T) {
 	}
 
 }
+
+func TestWrittenCountAfterBufferedWrite(t *testing.T) {
+	w := lz4.NewWriter(bytes.NewBuffer(nil))
+	w.Header.BlockDependency = true
+
+	if n, _ := w.Write([]byte{1}); n != 1 {
+		t.Errorf("expected to write 1 byte, wrote %d", n)
+		t.FailNow()
+	}
+
+	forcesWrite := make([]byte, 1<<16)
+
+	if n, _ := w.Write(forcesWrite); n != len(forcesWrite) {
+		t.Errorf("expected to write %d bytes, wrote %d", len(forcesWrite), n)
+		t.FailNow()
+	}
+}

+ 14 - 1
writer.go

@@ -125,7 +125,9 @@ func (z *Writer) Write(buf []byte) (n int, err error) {
 
 	// with block dependency, require at least 64Kb of data to work with
 	// not having 64Kb only matters initially to setup the first window
+	bl := 0
 	if z.BlockDependency && len(z.window) == 0 {
+		bl = len(z.data)
 		z.data = append(z.data, buf...)
 		if len(z.data) < winSize {
 			z.wg.Wait()
@@ -184,7 +186,18 @@ func (z *Writer) Write(buf []byte) (n int, err error) {
 	for zi, zb := range zblocks {
 		_, err = z.writeBlock(&zb)
 
-		n += len(zb.data)
+		written := len(zb.data)
+		if bl > 0 {
+			if written >= bl {
+				written -= bl
+				bl = 0
+			} else {
+				bl -= written
+				written = 0
+			}
+		}
+
+		n += written
 		// remove the window in zb.data
 		if z.BlockDependency {
 			if zi == 0 {