From b56af9b8b03962ca7fbd00c27f0a18beeb9ed7be Mon Sep 17 00:00:00 2001
From: "Masih H. Derkani" <m@derkani.org>
Date: Thu, 14 Mar 2024 21:03:42 +0000
Subject: [PATCH] Fix go imports at least 2x faster (#11695)

Use native go implementation to sort and adjust imports. Compared to the
previous bash version, this is at least 2X faster.
---
 Makefile                   |  6 +--
 scripts/fiximports         | 22 ---------
 scripts/fiximports/main.go | 92 ++++++++++++++++++++++++++++++++++++++
 3 files changed, 95 insertions(+), 25 deletions(-)
 delete mode 100755 scripts/fiximports
 create mode 100644 scripts/fiximports/main.go

diff --git a/Makefile b/Makefile
index b6e523b12..5724e38fa 100644
--- a/Makefile
+++ b/Makefile
@@ -332,7 +332,7 @@ actors-code-gen:
 	$(GOCC) fmt ./...
 
 actors-gen: actors-code-gen 
-	./scripts/fiximports
+	$(GOCC) run ./scripts/fiximports
 .PHONY: actors-gen
 
 bundle-gen:
@@ -392,10 +392,10 @@ docsgen-openrpc-gateway: docsgen-openrpc-bin
 .PHONY: docsgen docsgen-md-bin docsgen-openrpc-bin
 
 fiximports:
-	./scripts/fiximports
+	$(GOCC) run ./scripts/fiximports
 
 gen: actors-code-gen type-gen cfgdoc-gen docsgen api-gen circleci
-	./scripts/fiximports
+	$(GOCC) run ./scripts/fiximports
 	@echo ">>> IF YOU'VE MODIFIED THE CLI OR CONFIG, REMEMBER TO ALSO RUN 'make docsgen-cli'"
 .PHONY: gen
 
diff --git a/scripts/fiximports b/scripts/fiximports
deleted file mode 100755
index 140cebbce..000000000
--- a/scripts/fiximports
+++ /dev/null
@@ -1,22 +0,0 @@
-#!/usr/bin/env bash
-
-sed_replace='/import (/ {
-  :1
-  $!N
-  s/\n\n/\'$'\n''/
-  /)/!b1
-}'
-
-go_files() {
-    find . -type f -name \*.go -not -name \*_cbor_gen.go | grep -v './extern/filecoin-ffi' | grep -v './extern/test-vectors'
-}
-
-# Because -i works differently on macOS, we need to use a different sed command
-if [[ "$OSTYPE" == "darwin"* ]]; then
-    go_files | xargs -I '{}' sed -i '' -e "$sed_replace" '{}'
-else
-    go_files | xargs -I '{}' sed -i -e "$sed_replace" '{}'
-fi
-
-go_files | xargs -I '{}' goimports -w -local "github.com/filecoin-project" '{}'
-go_files | xargs -I '{}' goimports -w -local "github.com/filecoin-project/lotus" '{}'
diff --git a/scripts/fiximports/main.go b/scripts/fiximports/main.go
new file mode 100644
index 000000000..427975855
--- /dev/null
+++ b/scripts/fiximports/main.go
@@ -0,0 +1,92 @@
+package main
+
+import (
+	"bytes"
+	"fmt"
+	"io"
+	"io/fs"
+	"os"
+	"path/filepath"
+	"regexp"
+	"strings"
+
+	"golang.org/x/tools/imports"
+)
+
+var (
+	// groupByPrefixes is the list of import prefixes that should _each_ be grouped separately.
+	// See: imports.LocalPrefix.
+	groupByPrefixes = []string{
+		"github.com/filecoin-project",
+		"github.com/filecoin-project/lotus",
+	}
+	newline                  = []byte("\n")
+	importBlockRegex         = regexp.MustCompile(`(?s)import\s*\((.*?)\)`)
+	consecutiveNewlinesRegex = regexp.MustCompile(`\n\s*\n`)
+)
+
+func main() {
+	if err := filepath.Walk(".", func(path string, info fs.FileInfo, err error) error {
+		switch {
+		case err != nil:
+			return err
+		case // Skip the entire "./extern/..." directory and its contents.
+			strings.HasPrefix(path, "extern/"):
+			return filepath.SkipDir
+		case // Skip directories, generated cborgen go files and any other non-go files.
+			info.IsDir(),
+			strings.HasSuffix(info.Name(), "_cbor_gen.go"),
+			!strings.HasSuffix(info.Name(), ".go"):
+			return nil
+		}
+		return fixGoImports(path)
+	}); err != nil {
+		fmt.Printf("Error fixing go imports: %v\n", err)
+		os.Exit(1)
+	}
+}
+
+func fixGoImports(path string) error {
+	sourceFile, err := os.OpenFile(path, os.O_RDWR, 0666)
+	if err != nil {
+		return err
+	}
+	defer func() { _ = sourceFile.Close() }()
+
+	source, err := io.ReadAll(sourceFile)
+	if err != nil {
+		return err
+	}
+	formatted := collapseImportNewlines(source)
+	for _, prefix := range groupByPrefixes {
+		imports.LocalPrefix = prefix
+		formatted, err = imports.Process(path, formatted, nil)
+		if err != nil {
+			return err
+		}
+	}
+	if !bytes.Equal(source, formatted) {
+		if err := replaceFileContent(sourceFile, formatted); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+func replaceFileContent(target *os.File, replacement []byte) error {
+	if _, err := target.Seek(0, io.SeekStart); err != nil {
+		return err
+	}
+	written, err := target.Write(replacement)
+	if err != nil {
+		return err
+	}
+	return target.Truncate(int64(written))
+}
+
+func collapseImportNewlines(content []byte) []byte {
+	return importBlockRegex.ReplaceAllFunc(content, func(importBlock []byte) []byte {
+		// Replace consecutive newlines with a single newline within the import block
+		return consecutiveNewlinesRegex.ReplaceAll(importBlock, newline)
+	})
+}