Return zstd back

[glocate.git] / dirsizer.go

diff --git a/dirsizer.go b/dirsizer.go
index b4f5b4e51c46905addae4a0e5a061a445d23f1a3d6e0e6ce0967c2f36d221a93..fda37f4fe46098cf687690e5e03984aba77f952ed5814359e94deefcccbbbdb3 100644 (file)
--- a/dirsizer.go
+++ b/dirsizer.go
@@ -6,6 +6,8 @@ import (
        "io"
        "log"
        "os"
+
+       "github.com/klauspost/compress/zstd"
 )
 
 func dirSizer(dirSizes *[]int64, depth int, sinkBack, sinkIn, sinkOut chan Ent) (curSize int64) {
@@ -51,11 +53,22 @@ func applyDirSizes(src *os.File, dirSizes []int64) *os.File {
                log.Fatalln(err)
        }
 
-       br := bufio.NewReaderSize(src, 1<<17)
+       compR, err := zstd.NewReader(src)
+       if err != nil {
+               log.Fatalln(err)
+       }
+       br := bufio.NewReaderSize(compR, 1<<17)
+
+       compW, err := zstd.NewWriter(tmp,
+               zstd.WithEncoderLevel(zstd.SpeedBestCompression))
+       if err != nil {
+               log.Fatalln(err)
+       }
+       bw := bufio.NewWriterSize(compW, 1<<17)
+
        num := make([]byte, 8)
        var nameLen int
        name := make([]byte, 0, 1<<16)
-       bw := bufio.NewWriterSize(tmp, 1<<17)
        var dirIdx int
        for {
                if _, err = io.ReadFull(br, num[:2]); err != nil {
@@ -67,9 +80,7 @@ func applyDirSizes(src *os.File, dirSizes []int64) *os.File {
                mustWrite(bw, num[:2])
                nameLen = int(binary.BigEndian.Uint16(num[:2]))
                name = name[:nameLen]
-               if _, err = io.ReadFull(br, name); err != nil {
-                       log.Fatalln(err)
-               }
+               mustReadFull(br, name)
                mustWrite(bw, name)
                if _, err = io.CopyN(bw, br, 1+8); err != nil {
                        log.Fatalln(err)
@@ -90,5 +101,9 @@ func applyDirSizes(src *os.File, dirSizes []int64) *os.File {
        if err = bw.Flush(); err != nil {
                log.Fatalln(err)
        }
+       if err = compW.Close(); err != nil {
+               log.Fatalln(err)
+       }
+       compR.Close()
        return tmp
 }