Merge pull request #2 from PotatoesFall/memory

Reduce memory usage

Author: EwenQuim

Dockerfile

@@ -4,9 +4,12 @@ ARG TARGETARCH
 
 WORKDIR /go/src
 
-COPY . .
+COPY go.mod go.sum .
 
 RUN go mod download
+
+COPY . .
+
 RUN mkdir /data
 
 RUN GOOS=linux GOARCH=${TARGETARCH} go build -ldflags "-s -w" -o entropy .

main.go

@@ -35,6 +35,38 @@ type Entropy struct {
 	Line    string  // Line with high entropy
 }
 
+// Entropies should be created with a size n using make()
+// it should not be written to manually, instead use Entropies.Add
+type Entropies struct {
+	sync.Mutex
+	Entropies []Entropy
+}
+
+// Add assumes that es contains an ordered set of entropies.
+// It preserves ordering, and inserts an additional value e, if it has high enough entropy.
+// In that case, the entry with lowest entropy is rejected.
+func (es *Entropies) Add(e Entropy) {
+	es.Lock()
+	defer es.Unlock()
+
+	if es.Entropies[len(es.Entropies)-1].Entropy >= e.Entropy {
+		return
+	}
+
+	i, _ := slices.BinarySearchFunc(es.Entropies, e, func(a, b Entropy) int {
+		if b.Entropy > a.Entropy {
+			return 1
+		}
+		if a.Entropy > b.Entropy {
+			return -1
+		}
+		return 0
+	})
+
+	copy(es.Entropies[i+1:], es.Entropies[i:])
+	es.Entropies[i] = e
+}
+
 func main() {
 	minCharactersFlag := flag.Int("min", minCharactersDefault, "Minimum number of characters in the line to consider computing entropy")
 	resultCountFlag := flag.Int("top", resultCountDefault, "Number of results to display")
@@ -64,17 +96,16 @@ func main() {
 		fmt.Println("No files provided, defaults to current folder.")
 		fileNames = []string{"."}
 	}
-	entropies := make([]Entropy, 0, 10*len(fileNames))
+	entropies := &Entropies{
+		Entropies: make([]Entropy, resultCount),
+	}
 	for _, fileName := range fileNames {
-		fileEntropies, err := readFile(fileName)
+		err := readFile(entropies, fileName)
 		if err != nil {
 			fmt.Fprintf(os.Stderr, "Error reading file %s: %v\n", fileName, err)
 		}
-		entropies = append(entropies, fileEntropies...)
 	}
 
-	entropies = sortAndCutTop(entropies)
-
 	redMark := "\033[31m"
 	resetMark := "\033[0m"
 	if !term.IsTerminal(int(os.Stdout.Fd())) {
@@ -83,59 +114,54 @@ func main() {
 		resetMark = ""
 	}
 
-	for _, entropy := range entropies {
+	for _, entropy := range entropies.Entropies {
+		if entropy == (Entropy{}) {
+			return
+		}
 		fmt.Printf("%.2f: %s%s:%d%s %s\n", entropy.Entropy, redMark, entropy.File, entropy.LineNum, resetMark, entropy.Line)
 	}
 }
 
-func readFile(fileName string) ([]Entropy, error) {
+func readFile(entropies *Entropies, fileName string) error {
 	// If file is a folder, walk inside the folder
 	fileInfo, err := os.Stat(fileName)
 	if err != nil {
-		return nil, err
+		return err
 	}
 
 	if isFileHidden(fileInfo.Name()) && !exploreHidden {
-		return nil, nil
+		return nil
 	}
 
-	entropies := make([]Entropy, 0, 10)
 	if fileInfo.IsDir() {
 		// Walk through the folder and read all files
 		dir, err := os.ReadDir(fileName)
 		if err != nil {
-			return nil, err
+			return err
 		}
 
-		entropiies := make([][]Entropy, len(dir))
-
 		var wg sync.WaitGroup
 		for i, file := range dir {
 			wg.Add(1)
 			go func(i int, file os.DirEntry) {
 				defer wg.Done()
-				fileEntropies, err := readFile(fileName + "/" + file.Name())
+				err := readFile(entropies, fileName+"/"+file.Name())
 				if err != nil {
 					fmt.Fprintf(os.Stderr, "Error reading file %s: %v\n", file.Name(), err)
 				}
-				entropiies[i] = fileEntropies
 			}(i, file)
 		}
 
 		wg.Wait()
-
-		for _, fileEntropies := range entropiies {
-			entropies = append(entropies, fileEntropies...)
-		}
 	}
 
 	if !isFileIncluded(fileInfo.Name()) {
-		return sortAndCutTop(entropies), nil
+		return nil
 	}
 
 	file, err := os.Open(fileName)
 	if err != nil {
-		return nil, err
+		return err
 	}
 	defer file.Close()
 
@@ -150,7 +176,7 @@ func readFile(fileName string) ([]Entropy, error) {
 				continue
 			}
 
-			entropies = append(entropies, Entropy{
+			entropies.Add(Entropy{
 				Entropy: entropy(field),
 				File:    fileName,
 				LineNum: i,
@@ -159,7 +185,7 @@ func readFile(fileName string) ([]Entropy, error) {
 		}
 	}
 
-	return sortAndCutTop(entropies), nil
+	return nil
 }
 
 func entropy(text string) float64 {
@@ -210,15 +236,3 @@ func isFileIncluded(filename string) bool {
 
 	return false
 }
-
-func sortAndCutTop(entropies []Entropy) []Entropy {
-	slices.SortFunc(entropies, func(a, b Entropy) int {
-		return int((b.Entropy - a.Entropy) * 10000)
-	})
-
-	if len(entropies) > resultCount {
-		return entropies[:resultCount]
-	}
-
-	return entropies
-}

main_test.go

@@ -5,8 +5,9 @@ import (
 )
 
 func BenchmarkFile(b *testing.B) {
+	entropies := &Entropies{Entropies: make([]Entropy, 10)}
 	for range b.N {
-		readFile("testdata")
+		_ = readFile(entropies, "testdata")
 	}
 }
 
@@ -46,82 +47,38 @@ func TestEntropy(t *testing.T) {
 
 func TestReadFile(t *testing.T) {
 	t.Run("random.js", func(t *testing.T) {
-		res, err := readFile("testdata/random.js")
+		res := &Entropies{Entropies: make([]Entropy, 10)}
+		err := readFile(res, "testdata/random.js")
 		if err != nil {
 			t.Errorf("expected nil, got %v", err)
 		}
 
-		Expect(t, len(res), 10)
-		ExpectFloat(t, res[0].Entropy, 5.53614242151549)
-		Expect(t, res[0].LineNum, 7) // The token is hidden here
-		ExpectFloat(t, res[4].Entropy, 3.321928094887362)
+		ExpectFloat(t, res.Entropies[0].Entropy, 5.53614242151549)
+		Expect(t, res.Entropies[0].LineNum, 7) // The token is hidden here
+		ExpectFloat(t, res.Entropies[4].Entropy, 3.321928094887362)
 	})
 
 	t.Run("testdata/folder", func(t *testing.T) {
-		res, err := readFile("testdata/folder")
+		res := &Entropies{Entropies: make([]Entropy, 10)}
+		err := readFile(res, "testdata/folder")
 		if err != nil {
 			t.Errorf("expected nil, got %v", err)
 		}
 
-		Expect(t, len(res), 10)
-		ExpectFloat(t, res[0].Entropy, 3.7667029194153567)
-		Expect(t, res[0].LineNum, 7) // The token is hidden here
-		ExpectFloat(t, res[6].Entropy, 2.8553885422075336)
+		Expect(t, len(res.Entropies), 10)
+		ExpectFloat(t, res.Entropies[0].Entropy, 3.7667029194153567)
+		Expect(t, res.Entropies[0].LineNum, 7) // The token is hidden here
+		ExpectFloat(t, res.Entropies[6].Entropy, 2.8553885422075336)
 	})
 
 	t.Run("dangling symlink in testdata folder", func(t *testing.T) {
-		res, err := readFile("testdata")
+		res := &Entropies{Entropies: make([]Entropy, 10)}
+		err := readFile(res, "testdata")
 		if err != nil {
 			t.Errorf("expected nil, got %v", err)
 		}
 
-		Expect(t, len(res), 10)
-	})
-}
-
-func TestSortAndCutTop(t *testing.T) {
-	resultCount = 5
-
-	t.Run("nil", func(t *testing.T) {
-		res := sortAndCutTop(nil)
-		if len(res) != 0 {
-			t.Errorf("expected 0, got %d", len(res))
-		}
-	})
-
-	t.Run("empty", func(t *testing.T) {
-		res := sortAndCutTop([]Entropy{})
-		if len(res) != 0 {
-			t.Errorf("expected 0, got %d", len(res))
-		}
-	})
-
-	t.Run("less than resultCount", func(t *testing.T) {
-		res := sortAndCutTop([]Entropy{
-			{Entropy: 0.1},
-			{Entropy: 0.6},
-			{Entropy: 0.3},
-		})
-
-		Expect(t, len(res), 3)
-		Expect(t, res[0].Entropy, 0.6)
-		Expect(t, res[2].Entropy, 0.1)
-	})
-
-	t.Run("more than resultCount", func(t *testing.T) {
-		res := sortAndCutTop([]Entropy{
-			{Entropy: 0.1},
-			{Entropy: 0.6},
-			{Entropy: 0.3},
-			{Entropy: 0.7},
-			{Entropy: 0.4},
-			{Entropy: 0.5},
-			{Entropy: 0.2},
-		})
-
-		Expect(t, len(res), 5)
-		Expect(t, res[0].Entropy, 0.7)
-		Expect(t, res[4].Entropy, 0.3)
+		Expect(t, len(res.Entropies), 10)
 	})
 }
 
@@ -164,6 +121,19 @@ func TestIsFileHidden(t *testing.T) {
 	Expect(t, isFileHidden(".env"), true)
 }
 
+func TestEntropies(t *testing.T) {
+	res := &Entropies{Entropies: make([]Entropy, 5)}
+	for _, i := range []float64{1, 3, 5, 7, 2, 4, 6, 8} {
+		res.Add(Entropy{Entropy: i})
+	}
+
+	Expect(t, res.Entropies[0].Entropy, 8)
+	Expect(t, res.Entropies[1].Entropy, 7)
+	Expect(t, res.Entropies[2].Entropy, 6)
+	Expect(t, res.Entropies[3].Entropy, 5)
+	Expect(t, res.Entropies[4].Entropy, 4)
+}
+
 func Expect[T comparable](t *testing.T, got, expected T) {
 	t.Helper()
 	if got != expected {