From cddfd69257594e8353f929c9f434424b7c083e12 Mon Sep 17 00:00:00 2001
From: pagran <67878280+pagran@users.noreply.github.com>
Date: Sun, 23 Aug 2020 16:24:52 +0300
Subject: [PATCH] Limit fake line numbers

---
 line_obfuscator.go | 5 ++---
 main.go            | 2 +-
 2 files changed, 3 insertions(+), 4 deletions(-)

diff --git a/line_obfuscator.go b/line_obfuscator.go
index c1f472c..61024bb 100644
--- a/line_obfuscator.go
+++ b/line_obfuscator.go
@@ -3,7 +3,6 @@ package main
 import (
 	"fmt"
 	"go/ast"
-	"go/token"
 	mathrand "math/rand"
 	"strings"
 
@@ -80,12 +79,12 @@ func findBuildTags(commentGroups []*ast.CommentGroup) (buildTags []string) {
 	return buildTags
 }
 
-func transformLineInfo(file *ast.File, fset *token.FileSet) ([]string, *ast.File) {
+func transformLineInfo(file *ast.File) ([]string, *ast.File) {
 	// Save build tags and add file name leak protection
 	extraComments := append(findBuildTags(file.Comments), "", "//line :1")
 	file.Comments = nil
 
-	newLines := mathrand.Perm(fset.File(file.Package).LineCount())
+	newLines := mathrand.Perm(len(file.Decls))
 
 	funcCounter := 0
 	pre := func(cursor *astutil.Cursor) bool {
diff --git a/main.go b/main.go
index 8fe3fe7..0e05e77 100644
--- a/main.go
+++ b/main.go
@@ -526,7 +526,7 @@ func transformCompile(args []string) ([]string, error) {
 			// messy.
 			name = "_cgo_" + name
 		default:
-			extraComments, file = transformLineInfo(file, fset)
+			extraComments, file = transformLineInfo(file)
 			file = transformGo(file, info, blacklist)
 
 			// Uncomment for some quick debugging. Do not delete.