From 3e0c70bfe1cfe608b3ebd102badb3739ce04a04b Mon Sep 17 00:00:00 2001
From: xonix <xonixx@gmail.com>
Date: Thu, 18 Aug 2022 01:28:06 +0300
Subject: [PATCH] Work on annotating multiple files (multiple -f provided).

---
 cover/cover.go | 27 ++++++++++++++++-----------
 goawk.go       |  6 +-----
 2 files changed, 17 insertions(+), 16 deletions(-)

diff --git a/cover/cover.go b/cover/cover.go
index 98641dd4..156de5c8 100644
--- a/cover/cover.go
+++ b/cover/cover.go
@@ -18,17 +18,18 @@ type annotator struct {
 	program         *parser.Program
 }
 
-func NewAnnotator(covermode string) *annotator {
+func NewAnnotator(covermode string, parserConfig *parser.ParserConfig) *annotator {
 	return &annotator{covermode, "", 0,
-		map[int]ast.Boundary{}, map[int]string{}, parseProg("") /*TODO apply the parser config from caller*/}
+		map[int]ast.Boundary{}, map[int]string{}, parseProg("", parserConfig)}
+}
+
+var onlyParseParserConfig = &parser.ParserConfig{
+	DebugWriter:    os.Stderr,
+	OnlyParseToAST: true,
 }
 
 func (annotator *annotator) AddFile(filename string, code []byte) {
-	parserConfig := &parser.ParserConfig{
-		DebugWriter:    os.Stderr,
-		OnlyParseToAST: true,
-	}
-	prog, err := parser.ParseProgram(code, parserConfig)
+	prog, err := parser.ParseProgram(code, onlyParseParserConfig)
 	if err != nil {
 		panic(err) // at this point the code should be already valid
 	}
@@ -48,6 +49,10 @@ func (annotator *annotator) GetResultProgram() *parser.Program {
 	annotator.addCoverageEnd()
 	program := annotator.program
 	program.Resolve()
+	err := program.Compile()
+	if err != nil {
+		panic(err)
+	}
 	return program
 }
 
@@ -144,11 +149,11 @@ func (annotator *annotator) trackStatement(statements []ast.Stmt) ast.Stmt {
 		Start: statements[0].(ast.SimpleStmt).GetBoundary().Start,
 		End:   statements[len(statements)-1].(ast.SimpleStmt).GetBoundary().End,
 	}
-	return parseProg(fmt.Sprintf(`BEGIN { __COVER[%d]%s }`, annotator.annotationIdx, op)).Begin[0][0]
+	return parseProg(fmt.Sprintf(`BEGIN { __COVER[%d]%s }`, annotator.annotationIdx, op), nil).Begin[0][0]
 }
 
-func parseProg(code string) *parser.Program {
-	prog, err := parser.ParseProgram([]byte(code), nil)
+func parseProg(code string, parserConfig *parser.ParserConfig) *parser.Program {
+	prog, err := parser.ParseProgram([]byte(code), parserConfig)
 	if err != nil {
 		panic(err)
 	}
@@ -163,7 +168,7 @@ func (annotator *annotator) addCoverageEnd() {
 			i, annotator.fileNames[i], renderCoverBoundary(annotator.boundaries[i])))
 	}
 	code.WriteString("}")
-	annotator.program.End = append(annotator.program.End, parseProg(code.String()).End...)
+	annotator.program.End = append(annotator.program.End, parseProg(code.String(), nil).End...)
 }
 
 func renderCoverBoundary(boundary ast.Boundary) string {
diff --git a/goawk.go b/goawk.go
index add61f71..061b8e36 100644
--- a/goawk.go
+++ b/goawk.go
@@ -267,7 +267,7 @@ argsLoop:
 	//fmt.Println("after parse")
 
 	if covermode != "" {
-		annotator := cover.NewAnnotator(covermode)
+		annotator := cover.NewAnnotator(covermode, parserConfig)
 		// Read source: the concatenation of all source files specified
 		for _, progFile := range progFiles {
 			var f *os.File
@@ -292,10 +292,6 @@ argsLoop:
 			fmt.Fprintln(os.Stdout, prog)
 			os.Exit(0)
 		}
-		//err := prog.Compile() // recompile for annotations to take an effect
-		if err != nil {
-			errorExitf("%s", err)
-		}
 	}
 
 	if debug {