From 4544abab6dbf5d9e9a6b681e77bc2532a642760c Mon Sep 17 00:00:00 2001
From: Blake Bourque <Techplex.Engineer@gmail.com>
Date: Fri, 26 Feb 2021 22:42:41 -0500
Subject: [PATCH] Use stateful regexes to better handle formatting

---
 tools/generate.go | 28 ++++++++++++++++++++++------
 1 file changed, 22 insertions(+), 6 deletions(-)

diff --git a/tools/generate.go b/tools/generate.go
index b1609fd..ca4aabb 100644
--- a/tools/generate.go
+++ b/tools/generate.go
@@ -354,12 +354,28 @@ var SimpleShellOutputLexer = chroma.MustNewLexer(
 	},
 	chroma.Rules{
 		"root": {
-			// first capture group is generic prompt [capture $]
-			// second capture group is text [capture command and newline]
-			// third capture group is output [capture everything else until next $]
-			{`(\$)(.*$\n)([^\$]*)`, chroma.ByGroups(chroma.GenericPrompt, chroma.Text, chroma.GenericOutput), nil}, //BB-gp
-			// everything else is just regular text
-			{`\n|\w`, chroma.Text, nil},
+			// $ or > triggers the start of prompt formatting
+			{`^\$`, chroma.GenericPrompt, chroma.Push("prompt")},
+			{`^>`, chroma.GenericPrompt, chroma.Push("prompt")},
+
+			// empty lines are just text
+			{`^$\n`, chroma.Text, nil},
+
+			// otherwise its all output
+			{`[^\n]+$\n?`, chroma.GenericOutput, nil},
+		},
+		"prompt": {
+			// when we find newline, do output formatting rules
+			{`\n`, chroma.Text, chroma.Push("output")},
+			// otherwise its all text
+			{`[^\n]+$`, chroma.Text, nil},
+		},
+		"output": {
+			// sometimes there isn't output so we go right back to prompt
+			{`^\$`, chroma.GenericPrompt, chroma.Pop(1)},
+			{`^>`, chroma.GenericPrompt, chroma.Pop(1)},
+			// otherwise its all output
+			{`[^\n]+$\n?`, chroma.GenericOutput, nil},
 		},
 	},
 )