From b631e7d0ee04af7d7f57a13ae292d28e47180adc Mon Sep 17 00:00:00 2001
From: Nat Quayle Nelson <natquaylenelson@gmail.com>
Date: Sun, 26 Oct 2025 08:27:36 -0500
Subject: [PATCH] Support gather at a choice

---
 examples/gather-at-choice/main.hank  | 11 +++++++++++
 examples/gather-at-choice/test1.hlog | 10 ++++++++++
 hank/HankAST.hx                      |  9 +++++++--
 hank/Story.hx                        |  5 +++--
 4 files changed, 31 insertions(+), 4 deletions(-)
 create mode 100644 examples/gather-at-choice/main.hank
 create mode 100644 examples/gather-at-choice/test1.hlog

diff --git a/examples/gather-at-choice/main.hank b/examples/gather-at-choice/main.hank
new file mode 100644
index 0000000..d02a374
--- /dev/null
+++ b/examples/gather-at-choice/main.hank
@@ -0,0 +1,11 @@
+What a beautiful morning!
+
+* Drink water
+
+- * Brush teeth
+
+- * Egg breakfast
+* Toast breakfast
+* Easy breakfast
+
+- Done
\ No newline at end of file
diff --git a/examples/gather-at-choice/test1.hlog b/examples/gather-at-choice/test1.hlog
new file mode 100644
index 0000000..afffeeb
--- /dev/null
+++ b/examples/gather-at-choice/test1.hlog
@@ -0,0 +1,10 @@
+What a beautiful morning!
+* Drink water
+> 1: Drink water
+* Brush teeth
+> 1: Brush teeth
+* Egg breakfast
+* Toast breakfast
+* Easy breakfast
+> 2: Toast breakfast
+Done
\ No newline at end of file
diff --git a/hank/HankAST.hx b/hank/HankAST.hx
index 886de30..a0f79de 100644
--- a/hank/HankAST.hx
+++ b/hank/HankAST.hx
@@ -80,7 +80,12 @@ class ASTExtension {
 				case EChoice(choice):
 					if (tryAdd(choice, [])) lastChoiceIndex = i;
 				case ETagged(EChoice(choice), tags):
-					if (tryAdd(choice, tags)) lastChoiceIndex = 1;
+					if (tryAdd(choice, tags)) lastChoiceIndex = i;
+				// Choice immediatly after gather at current depth:
+				case EGather(_, d, EChoice(choice)) if(i == startingIndex && d == depth):
+					if (tryAdd(choice, [])) lastChoiceIndex = i;
+				case EGather(_, d, ETagged(EChoice(choice), tags)) if(i == startingIndex && d == depth):
+					if (tryAdd(choice, tags)) lastChoiceIndex = i;
 				// Stop at the next gather of this depth
 				case EGather(_, d, _) if (d == depth):
 					break;
@@ -111,7 +116,7 @@ class ASTExtension {
 		for (i in 0...ast.length) {
 			var expr = ast[i].expr;
 			switch (expr) {
-				case EChoice(c) | ETagged(EChoice(c), _):
+				case EChoice(c) | ETagged(EChoice(c), _) | EGather(_, _, EChoice(c)):
 					if (c.id == id)
 						return i;
 				default:
diff --git a/hank/Story.hx b/hank/Story.hx
index f7029e8..8ec9bec 100644
--- a/hank/Story.hx
+++ b/hank/Story.hx
@@ -383,7 +383,9 @@ class Story {
 			}
 		}
 
-		if (exprIndex < ast.length && ast[exprIndex].expr.match(EChoice(_))) {
+		if (exprIndex < ast.length &&
+			 	(ast[exprIndex].expr.match(EChoice(_))
+					|| ast[exprIndex].expr.match(EGather(_, _, EChoice(_))))) {
 			var allChoiceInfo = ast.collectChoices(exprIndex, weaveDepth).choices;
 			for (choiceInfo in allChoiceInfo) {
 				if (choicesTaken.indexOf(choiceInfo.choice.id) == -1 || !choiceInfo.choice.onceOnly) {
@@ -404,7 +406,6 @@ class Story {
 		// trace('final:');
 		// traceChoiceArray(choices);
 
-		// traceChoiceArray(choices);
 		return choices;
 	}