Return zstd back

[glocate.git] / writer.go

diff --git a/writer.go b/writer.go
index e12a19c0d7f2103ae9c88e4c4f3c237bb3d03f62819ea98eb33c93210f221256..75b82e57abc5e88fff375e6c1a4a69ae1335b4e2ce944ad401c940ee08e3c8e1 100644 (file)
--- a/writer.go
+++ b/writer.go
@@ -1,9 +1,12 @@
 package main
 
 import (
+       "bufio"
        "encoding/binary"
        "io"
        "log"
+
+       "github.com/klauspost/compress/zstd"
 )
 
 func mustWrite(w io.Writer, buf []byte) {
@@ -12,7 +15,12 @@ func mustWrite(w io.Writer, buf []byte) {
        }
 }
 
-func writer(w io.Writer, sink chan Ent) {
+func writer(dst io.Writer, sink chan Ent) {
+       comp, err := zstd.NewWriter(dst)
+       if err != nil {
+               log.Fatalln(err)
+       }
+       bw := bufio.NewWriterSize(comp, 1<<17)
        num := make([]byte, 8)
        var name string
        for ent := range sink {
@@ -21,15 +29,21 @@ func writer(w io.Writer, sink chan Ent) {
                        panic("too long")
                }
                binary.BigEndian.PutUint16(num[:2], uint16(len(name)))
-               mustWrite(w, num[:2])
-               mustWrite(w, []byte(name))
+               mustWrite(bw, num[:2])
+               mustWrite(bw, []byte(name))
                if len(ent.name) >= 1<<8 {
                        panic("too deep")
                }
-               mustWrite(w, []byte{byte(len(ent.name) - 1)})
+               mustWrite(bw, []byte{byte(len(ent.name) - 1)})
                binary.BigEndian.PutUint64(num, uint64(ent.mtime))
-               mustWrite(w, num)
+               mustWrite(bw, num)
                binary.BigEndian.PutUint64(num, uint64(ent.size))
-               mustWrite(w, num)
+               mustWrite(bw, num)
+       }
+       if err = bw.Flush(); err != nil {
+               log.Fatalln(err)
+       }
+       if err = comp.Close(); err != nil {
+               log.Fatalln(err)
        }
 }