dirsizer.go

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