• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1package main
2
3import (
4	"flag"
5	"fmt"
6	"io/ioutil"
7	"os"
8	"path"
9	"path/filepath"
10	"sort"
11	"strings"
12	"sync"
13
14	"google.golang.org/protobuf/proto"
15)
16
17var root string
18var repo string
19var outDir string
20var channel chan SourceFile
21var waiter sync.WaitGroup
22var sourceTree SourceTree
23
24func normalizeOutDir(outDir, root string) string {
25	if len(outDir) == 0 {
26		return ""
27	}
28	if filepath.IsAbs(outDir) {
29		if strings.HasPrefix(outDir, root) {
30			// Absolute path inside root, use it
31			return outDir
32		} else {
33			// Not inside root, we don't care about it
34			return ""
35		}
36	} else {
37		// Relative path inside root, make it absolute
38		return root + outDir
39	}
40}
41
42func walkDir(dir string) {
43	defer waiter.Done()
44
45	visit := func(path string, info os.FileInfo, err error) error {
46		name := info.Name()
47
48		// Repo git projects are symlinks.  A real directory called .git counts as checked in
49		// (and is very likely to be wasted space)
50		if info.Mode().Type()&os.ModeSymlink != 0 && name == ".git" {
51			return nil
52		}
53
54		// Skip .repo and out
55		if info.IsDir() && (path == repo || path == outDir) {
56			return filepath.SkipDir
57		}
58
59		if info.IsDir() && path != dir {
60			waiter.Add(1)
61			go walkDir(path)
62			return filepath.SkipDir
63		}
64
65		if !info.IsDir() {
66			sourcePath := strings.TrimPrefix(path, root)
67			file := SourceFile{
68				Path:      proto.String(sourcePath),
69				SizeBytes: proto.Int32(42),
70			}
71			channel <- file
72		}
73		return nil
74
75	}
76	filepath.Walk(dir, visit)
77}
78
79func main() {
80	var outputFile string
81	flag.StringVar(&outputFile, "o", "", "The file to write")
82	flag.StringVar(&outDir, "out_dir", "out", "The out directory")
83	flag.Parse()
84
85	if outputFile == "" {
86		fmt.Fprintf(os.Stderr, "source_tree_size: Missing argument: -o\n")
87		os.Exit(1)
88	}
89
90	root, _ = os.Getwd()
91	if root[len(root)-1] != '/' {
92		root += "/"
93	}
94
95	outDir = normalizeOutDir(outDir, root)
96	repo = path.Join(root, ".repo")
97
98	// The parallel scanning reduces the run time by about a minute
99	channel = make(chan SourceFile)
100	waiter.Add(1)
101	go walkDir(root)
102	go func() {
103		waiter.Wait()
104		close(channel)
105	}()
106	for sourceFile := range channel {
107		sourceTree.Files = append(sourceTree.Files, &sourceFile)
108	}
109
110	// Sort the results, for a stable output
111	sort.Slice(sourceTree.Files, func(i, j int) bool {
112		return *sourceTree.Files[i].Path < *sourceTree.Files[j].Path
113	})
114
115	// Flatten and write
116	buf, err := proto.Marshal(&sourceTree)
117	if err != nil {
118		fmt.Fprintf(os.Stderr, "Couldn't marshal protobuf\n")
119		os.Exit(1)
120	}
121	err = ioutil.WriteFile(outputFile, buf, 0644)
122	if err != nil {
123		fmt.Fprintf(os.Stderr, "Error writing file: %v\n", err)
124		os.Exit(1)
125	}
126}
127