dirsizer.go

   1 package main
   2
   3 import (
   4         "bufio"
   5         "encoding/binary"
   6         "io"
   7         "log"
   8         "os"
   9
  10         "github.com/klauspost/compress/zstd"
  11 )
  12
  13 func dirSizer(dirSizes *[]int64, depth int, sinkBack, sinkIn, sinkOut chan Ent) (curSize int64) {
  14         var ent Ent
  15         var opened bool
  16         var dirIdx int
  17         for {
  18                 select {
  19                 case ent = <-sinkBack:
  20                         goto Got
  21                 default:
  22                 }
  23                 ent, opened = <-sinkIn
  24                 if !opened {
  25                         break
  26                 }
  27         Got:
  28                 if len(ent.name) < depth {
  29                         sinkBack <- ent
  30                         return
  31                 }
  32                 sinkOut <- ent
  33                 curSize += ent.size
  34                 if !ent.IsDir() {
  35                         continue
  36                 }
  37                 dirIdx = len(*dirSizes)
  38                 (*dirSizes) = append(*dirSizes, 0)
  39                 dirSize := dirSizer(dirSizes, depth+1, sinkBack, sinkIn, sinkOut)
  40                 (*dirSizes)[dirIdx] = dirSize
  41                 curSize += dirSize
  42         }
  43         return
  44 }
  45
  46 func applyDirSizes(src *os.File, dirSizes []int64) *os.File {
  47         _, err := src.Seek(0, io.SeekStart)
  48         if err != nil {
  49                 log.Fatalln(err)
  50         }
  51         tmp, err := os.CreateTemp("", "glocate-idx")
  52         if err != nil {
  53                 log.Fatalln(err)
  54         }
  55
  56         compR, err := zstd.NewReader(src)
  57         if err != nil {
  58                 log.Fatalln(err)
  59         }
  60         br := bufio.NewReaderSize(compR, 1<<17)
  61
  62         compW, err := zstd.NewWriter(tmp,
  63                 zstd.WithEncoderLevel(zstd.SpeedBestCompression))
  64         if err != nil {
  65                 log.Fatalln(err)
  66         }
  67         bw := bufio.NewWriterSize(compW, 1<<17)
  68
  69         num := make([]byte, 8)
  70         var nameLen int
  71         name := make([]byte, 0, 1<<16)
  72         var dirIdx int
  73         for {
  74                 if _, err = io.ReadFull(br, num[:2]); err != nil {
  75                         if err == io.EOF {
  76                                 break
  77                         }
  78                         log.Fatalln(err)
  79                 }
  80                 mustWrite(bw, num[:2])
  81                 nameLen = int(binary.BigEndian.Uint16(num[:2]))
  82                 name = name[:nameLen]
  83                 mustReadFull(br, name)
  84                 mustWrite(bw, name)
  85                 if _, err = io.CopyN(bw, br, 1+8); err != nil {
  86                         log.Fatalln(err)
  87                 }
  88                 if name[len(name)-1] == byte('/') {
  89                         if _, err = br.Discard(8); err != nil {
  90                                 log.Fatalln(err)
  91                         }
  92                         binary.BigEndian.PutUint64(num, uint64(dirSizes[dirIdx]))
  93                         mustWrite(bw, num)
  94                         dirIdx++
  95                 } else {
  96                         if _, err = io.CopyN(bw, br, 8); err != nil {
  97                                 log.Fatalln(err)
  98                         }
  99                 }
 100         }
 101         if err = bw.Flush(); err != nil {
 102                 log.Fatalln(err)
 103         }
 104         if err = compW.Close(); err != nil {
 105                 log.Fatalln(err)
 106         }
 107         compR.Close()
 108         return tmp
 109 }