X-Git-Url: http://www.git.stargrave.org/?a=blobdiff_plain;f=dirsizer.go;fp=dirsizer.go;h=b4f5b4e51c46905addae4a0e5a061a445d23f1a3d6e0e6ce0967c2f36d221a93;hb=d5b8c235a1f3088c6c1e7261df3d1b565d042db2ba2ad1bbd1018782b9178e1f;hp=0000000000000000000000000000000000000000000000000000000000000000;hpb=411a031ec7cc707b8269acc3dfe28bc8db1bab5a9a91781c26809ae9853c6f6a;p=glocate.git

diff --git a/dirsizer.go b/dirsizer.go
new file mode 100644
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
+}