Improved version

[glocate.git] / dirsizer.go

diff --git a/dirsizer.go b/dirsizer.go
new file mode 100644 (file)
index 0000000..b4f5b4e
--- /dev/null
+++ b/dirsizer.go
@@ -0,0 +1,94 @@
+package main
+
+import (
+       "bufio"
+       "encoding/binary"
+       "io"
+       "log"
+       "os"
+)
+
+func dirSizer(dirSizes *[]int64, depth int, sinkBack, sinkIn, sinkOut chan Ent) (curSize int64) {
+       var ent Ent
+       var opened bool
+       var dirIdx int
+       for {
+               select {
+               case ent = <-sinkBack:
+                       goto Got
+               default:
+               }
+               ent, opened = <-sinkIn
+               if !opened {
+                       break
+               }
+       Got:
+               if len(ent.name) < depth {
+                       sinkBack <- ent
+                       return
+               }
+               sinkOut <- ent
+               curSize += ent.size
+               if !ent.IsDir() {
+                       continue
+               }
+               dirIdx = len(*dirSizes)
+               (*dirSizes) = append(*dirSizes, 0)
+               dirSize := dirSizer(dirSizes, depth+1, sinkBack, sinkIn, sinkOut)
+               (*dirSizes)[dirIdx] = dirSize
+               curSize += dirSize
+       }
+       return
+}
+
+func applyDirSizes(src *os.File, dirSizes []int64) *os.File {
+       _, err := src.Seek(0, io.SeekStart)
+       if err != nil {
+               log.Fatalln(err)
+       }
+       tmp, err := os.CreateTemp("", "glocate-idx")
+       if err != nil {
+               log.Fatalln(err)
+       }
+
+       br := bufio.NewReaderSize(src, 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 {
+                       if err == io.EOF {
+                               break
+                       }
+                       log.Fatalln(err)
+               }
+               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)
+               }
+               mustWrite(bw, name)
+               if _, err = io.CopyN(bw, br, 1+8); err != nil {
+                       log.Fatalln(err)
+               }
+               if name[len(name)-1] == byte('/') {
+                       if _, err = br.Discard(8); err != nil {
+                               log.Fatalln(err)
+                       }
+                       binary.BigEndian.PutUint64(num, uint64(dirSizes[dirIdx]))
+                       mustWrite(bw, num)
+                       dirIdx++
+               } else {
+                       if _, err = io.CopyN(bw, br, 8); err != nil {
+                               log.Fatalln(err)
+                       }
+               }
+       }
+       if err = bw.Flush(); err != nil {
+               log.Fatalln(err)
+       }
+       return tmp
+}