From a5937e0dc71065dcfa5113fc3f2a112a3d0a1be3 Mon Sep 17 00:00:00 2001
From: Ned Palacios <npdoesmc@gmail.com>
Date: Sat, 3 Feb 2024 21:06:13 +0800
Subject: [PATCH] refactor: remove contextData param in OutputGen, add hooks

---
 errgoengine.go                           | 46 +++++++++++++++++++++++-
 error_templates/test_utils/test_utils.go |  2 +-
 output_gen.go                            | 39 +++++---------------
 3 files changed, 55 insertions(+), 32 deletions(-)

diff --git a/errgoengine.go b/errgoengine.go
index 7ced40b..fc13e1b 100644
--- a/errgoengine.go
+++ b/errgoengine.go
@@ -13,6 +13,7 @@ type ErrgoEngine struct {
 	ErrorTemplates ErrorTemplates
 	FS             *MultiReadFileFS
 	OutputGen      *OutputGenerator
+	IsTesting      bool
 }
 
 func New() *ErrgoEngine {
@@ -125,7 +126,50 @@ func (e *ErrgoEngine) Translate(template *CompiledErrorTemplate, contextData *Co
 		template.OnGenBugFixFn(contextData, fixGen)
 	}
 
-	output := e.OutputGen.Generate(contextData, expGen, fixGen)
+	if e.IsTesting {
+		// add a code snippet that points to the error
+		e.OutputGen.GenAfterExplain = func(gen *OutputGenerator) {
+			err := contextData.MainError
+			if err == nil {
+				return
+			}
+
+			doc := err.Document
+			if doc == nil || err.Nearest.IsNull() {
+				return
+			}
+
+			startLineNr := err.Nearest.StartPosition().Line
+			startLines := doc.LinesAt(max(startLineNr-1, 0), startLineNr)
+			endLines := doc.LinesAt(min(startLineNr+1, doc.TotalLines()), min(startLineNr+2, doc.TotalLines()))
+			arrowLength := int(err.Nearest.EndByte() - err.Nearest.StartByte())
+			if arrowLength == 0 {
+				arrowLength = 1
+			}
+
+			startArrowPos := err.Nearest.StartPosition().Column
+			gen.Writeln("```")
+			gen.WriteLines(startLines...)
+
+			for i := 0; i < startArrowPos; i++ {
+				if startLines[len(startLines)-1][i] == '\t' {
+					gen.Builder.WriteString("    ")
+				} else {
+					gen.Builder.WriteByte(' ')
+				}
+			}
+
+			for i := 0; i < arrowLength; i++ {
+				gen.Builder.WriteByte('^')
+			}
+
+			gen._break()
+			gen.WriteLines(endLines...)
+			gen.Writeln("```")
+		}
+	}
+
+	output := e.OutputGen.Generate(expGen, fixGen)
 	defer e.OutputGen.Reset()
 
 	return expGen.Builder.String(), output
diff --git a/error_templates/test_utils/test_utils.go b/error_templates/test_utils/test_utils.go
index 704a3b5..378891b 100644
--- a/error_templates/test_utils/test_utils.go
+++ b/error_templates/test_utils/test_utils.go
@@ -57,7 +57,7 @@ func SetupTest(tb testing.TB, cfg SetupTestConfig) TestCases {
 
 	// load error templates
 	engine := lib.New()
-	engine.OutputGen.IsTesting = true
+	engine.IsTesting = true
 
 	cfg.TemplateLoader(&engine.ErrorTemplates)
 
diff --git a/output_gen.go b/output_gen.go
index fddb50b..caff2eb 100644
--- a/output_gen.go
+++ b/output_gen.go
@@ -6,8 +6,8 @@ import (
 )
 
 type OutputGenerator struct {
-	IsTesting bool
-	Builder   *strings.Builder
+	GenAfterExplain func(*OutputGenerator)
+	Builder         *strings.Builder
 }
 
 func (gen *OutputGenerator) Heading(level int, text string) {
@@ -70,7 +70,7 @@ func (gen *OutputGenerator) WriteLines(lines ...string) {
 	}
 }
 
-func (gen *OutputGenerator) Generate(cd *ContextData, explain *ExplainGenerator, bugFix *BugFixGenerator) string {
+func (gen *OutputGenerator) Generate(explain *ExplainGenerator, bugFix *BugFixGenerator) string {
 	if gen.Builder == nil {
 		gen.Builder = &strings.Builder{}
 	}
@@ -80,33 +80,8 @@ func (gen *OutputGenerator) Generate(cd *ContextData, explain *ExplainGenerator,
 	}
 
 	gen.ExpGen(1, explain)
-	doc := cd.MainError.Document
-
-	if doc != nil && gen.IsTesting && !cd.MainError.Nearest.IsNull() {
-		startLineNr := cd.MainError.Nearest.StartPosition().Line
-		startLines := doc.LinesAt(max(startLineNr-1, 0), startLineNr)
-		endLines := doc.LinesAt(min(startLineNr+1, doc.TotalLines()), min(startLineNr+2, doc.TotalLines()))
-		arrowLength := int(cd.MainError.Nearest.EndByte() - cd.MainError.Nearest.StartByte())
-		if arrowLength == 0 {
-			arrowLength = 1
-		}
-
-		startArrowPos := cd.MainError.Nearest.StartPosition().Column
-		gen.Writeln("```")
-		gen.WriteLines(startLines...)
-		for i := 0; i < startArrowPos; i++ {
-			if startLines[len(startLines)-1][i] == '\t' {
-				gen.Builder.WriteString("    ")
-			} else {
-				gen.Builder.WriteByte(' ')
-			}
-		}
-		for i := 0; i < arrowLength; i++ {
-			gen.Builder.WriteByte('^')
-		}
-		gen._break()
-		gen.WriteLines(endLines...)
-		gen.Writeln("```")
+	if gen.GenAfterExplain != nil {
+		gen.GenAfterExplain(gen)
 	}
 
 	gen.Heading(2, "Steps to fix")
@@ -131,6 +106,7 @@ func (gen *OutputGenerator) Generate(cd *ContextData, explain *ExplainGenerator,
 				}
 
 				if len(step.Fixes) != 0 {
+					doc := step.Doc.Document
 					descriptionBuilder := &strings.Builder{}
 
 					// get the start and end line after applying the diff
@@ -225,5 +201,8 @@ func (gen *OutputGenerator) Generate(cd *ContextData, explain *ExplainGenerator,
 }
 
 func (gen *OutputGenerator) Reset() {
+	if gen.GenAfterExplain != nil {
+		gen.GenAfterExplain = nil
+	}
 	gen.Builder.Reset()
 }