Rewrite comment cleanup

Fix method shuffling
Added compatibility with marked //go directives methods

line_obfuscator.go

@@ -3,8 +3,10 @@ package main
 import (
 	"fmt"
 	"go/ast"
-	"golang.org/x/tools/go/ast/astutil"
 	mathrand "math/rand"
+	"strings"
+
+	"golang.org/x/tools/go/ast/astutil"
 )
 
 const (
@@ -17,26 +19,71 @@ const (
 	PosMin = 1
 )
 
+func prependComment(group *ast.CommentGroup, comment *ast.Comment) *ast.CommentGroup {
+	if group == nil {
+		return &ast.CommentGroup{List: []*ast.Comment{comment}}
+	}
+
+	group.List = append([]*ast.Comment{comment}, group.List...)
+	return group
+}
+
+func clearCommentGroup(group *ast.CommentGroup) *ast.CommentGroup {
+	if group == nil {
+		return nil
+	}
+
+	var comments []*ast.Comment
+
+	for _, comment := range group.List {
+		if strings.HasPrefix(comment.Text, "//go:") {
+			comments = append(comments, &ast.Comment{Text: comment.Text})
+		}
+	}
+	if len(comments) == 0 {
+		return nil
+	}
+	return &ast.CommentGroup{List: comments}
+}
+
+func clearNodeComments(node ast.Node) {
+	switch n := node.(type) {
+	case *ast.Field:
+		n.Doc = clearCommentGroup(n.Doc)
+	case *ast.ImportSpec:
+		n.Doc = clearCommentGroup(n.Doc)
+	case *ast.ValueSpec:
+		n.Doc = clearCommentGroup(n.Doc)
+	case *ast.TypeSpec:
+		n.Doc = clearCommentGroup(n.Doc)
+	case *ast.GenDecl:
+		n.Doc = clearCommentGroup(n.Doc)
+	case *ast.FuncDecl:
+		n.Doc = clearCommentGroup(n.Doc)
+	case *ast.File:
+		n.Doc = clearCommentGroup(n.Doc)
+	}
+}
+
 func transformLineInfo(fileIndex int, file *ast.File) *ast.File {
+	file.Comments = nil
 	pre := func(cursor *astutil.Cursor) bool {
-		funcDecl, ok := cursor.Node().(*ast.FuncDecl)
-		if !ok {
-			return true
-		}
+		node := cursor.Node()
+		clearNodeComments(node)
 
-		// Ignore functions with //go: directives
-		if funcDecl.Doc != nil && len(funcDecl.Doc.List) != 0 {
+		funcDecl, ok := node.(*ast.FuncDecl)
+		if !ok {
 			return true
 		}
 
 		if envGarbleTiny {
-			funcDecl.Doc = &ast.CommentGroup{List: []*ast.Comment{{Text: "//line :1"}}}
+			funcDecl.Doc = prependComment(funcDecl.Doc, &ast.Comment{Text: "//line :1"})
 			return true
 		}
 
 		linePos := hashWithAsUint64(buildInfo.buildID, fmt.Sprintf("%d:%s", fileIndex, funcDecl.Name), PosMin, PosMax)
 		comment := &ast.Comment{Text: fmt.Sprintf("//line %c.go:%d", nameCharset[mathrand.Intn(len(nameCharset))], linePos)}
-		funcDecl.Doc = &ast.CommentGroup{List: []*ast.Comment{comment}}
+		funcDecl.Doc = prependComment(funcDecl.Doc, comment)
 		return true
 	}
 

main.go

@@ -739,38 +739,19 @@ func buildBlacklist(files []*ast.File, info *types.Info, pkg *types.Package) map
 
 // transformGo garbles the provided Go syntax node.
 func transformGo(file *ast.File, info *types.Info, blacklist map[types.Object]struct{}) *ast.File {
-	// Remove all comments, minus the "//go:" compiler directives.
-	// The final binary should still not contain comment text, but removing
-	// it helps ensure that (and makes position info less predictable).
-	origComments := file.Comments
-	file.Comments = nil
-	for _, commentGroup := range origComments {
-		for _, comment := range commentGroup.List {
-			if strings.HasPrefix(comment.Text, "//go:") {
-				file.Comments = append(file.Comments, &ast.CommentGroup{
-					List: []*ast.Comment{comment},
-				})
-			}
+	// Shuffle top level declarations
+	mathrand.Shuffle(len(file.Decls), func(i, j int) {
+		decl1 := file.Decls[i]
+		decl2 := file.Decls[j]
+
+		// Import declarations must remain at the top of the file.
+		gd1, ok1 := decl1.(*ast.GenDecl)
+		gd2, ok2 := decl2.(*ast.GenDecl)
+		if (ok1 && gd1.Tok == token.IMPORT) || (ok2 && gd2.Tok == token.IMPORT) {
+			return
 		}
-	}
-
-	// Shuffle top level declarations if there are no remaining compiler
-	// directives.
-	if len(file.Comments) == 0 {
-		// TODO: Also allow files with compiler directives.
-		mathrand.Shuffle(len(file.Decls), func(i, j int) {
-			decl1 := file.Decls[i]
-			decl2 := file.Decls[j]
-
-			// Import declarations must remain at the top of the file.
-			gd1, ok1 := decl1.(*ast.GenDecl)
-			gd2, ok2 := decl2.(*ast.GenDecl)
-			if (ok1 && gd1.Tok == token.IMPORT) || (ok2 && gd2.Tok == token.IMPORT) {
-				return
-			}
-			file.Decls[i], file.Decls[j] = decl2, decl1
-		})
-	}
+		file.Decls[i], file.Decls[j] = decl2, decl1
+	})
 
 	pre := func(cursor *astutil.Cursor) bool {
 		node, ok := cursor.Node().(*ast.Ident)