From f792580816ff8728fb0abf5d97efc7c7204ec48f Mon Sep 17 00:00:00 2001
From: Sylvie <35663410+Rangi42@users.noreply.github.com>
Date: Thu, 21 Mar 2024 19:53:49 -0400
Subject: [PATCH] Only restore parent context's `\@` value if it had one
 defined (#1366)

This way, if a child context initializes `\@`, the parent won't
reset it. And if the child context did not initialize `\@`,
then resetting it would be redundant.
---
 include/asm/fstack.hpp         |  2 --
 src/asm/fstack.cpp             | 10 +++++++---
 test/asm/unique-id-include.asm | 27 +++++++++++++++++++++++++++
 test/asm/unique-id-include.inc |  1 +
 test/asm/unique-id-include.out | 19 +++++++++++++++++++
 test/asm/unique-id-nested.asm  | 33 +++++++++++++++++++++++++++++++++
 test/asm/unique-id-nested.out  | 21 +++++++++++++++++++++
 7 files changed, 108 insertions(+), 5 deletions(-)
 create mode 100644 test/asm/unique-id-include.asm
 create mode 100644 test/asm/unique-id-include.inc
 create mode 100644 test/asm/unique-id-include.out
 create mode 100644 test/asm/unique-id-nested.asm
 create mode 100644 test/asm/unique-id-nested.out

diff --git a/include/asm/fstack.hpp b/include/asm/fstack.hpp
index 2a9de798..96b0b1b9 100644
--- a/include/asm/fstack.hpp
+++ b/include/asm/fstack.hpp
@@ -14,8 +14,6 @@
 
 #include "linkdefs.hpp"
 
-#include "asm/lexer.hpp"
-
 struct FileStackNode {
 	FileStackNodeType type;
 	std::variant<
diff --git a/src/asm/fstack.cpp b/src/asm/fstack.cpp
index cf1e525f..64199324 100644
--- a/src/asm/fstack.cpp
+++ b/src/asm/fstack.cpp
@@ -9,6 +9,7 @@
 #include <inttypes.h>
 #include <new>
 #include <stack>
+#include <stdio.h>
 #include <stdlib.h>
 #include <string_view>
 
@@ -207,20 +208,23 @@ bool yywrap() {
 		return true;
 	}
 
-	Context oldContext = contextStack.top();
-
+	Context oldContext = std::move(contextStack.top());
 	contextStack.pop();
 
 	lexer_CleanupState(oldContext.lexerState);
 	// Restore args if a macro (not REPT) saved them
 	if (oldContext.fileInfo->type == NODE_MACRO)
 		macro_UseNewArgs(contextStack.top().macroArgs);
+	// Restore the `\@` value if the old context was supposed to define one
+	// (i.e. it either is a macro or REPT, or it initialized `\@` for its parent)
+	if (oldContext.fileInfo->type == NODE_REPT || oldContext.fileInfo->type == NODE_MACRO
+	    || contextStack.top().uniqueID > 0)
+		macro_SetUniqueID(contextStack.top().uniqueID);
 	// Free the file stack node
 	if (!oldContext.fileInfo->referenced)
 		delete oldContext.fileInfo;
 
 	lexer_SetState(&contextStack.top().lexerState);
-	macro_SetUniqueID(contextStack.top().uniqueID);
 
 	return false;
 }
diff --git a/test/asm/unique-id-include.asm b/test/asm/unique-id-include.asm
new file mode 100644
index 00000000..9b0f9582
--- /dev/null
+++ b/test/asm/unique-id-include.asm
@@ -0,0 +1,27 @@
+REPT 2
+	println "Within REPT"
+	INCLUDE "unique-id-include.inc"
+	println "Outside INCLUDE: \@"
+ENDR
+
+MACRO m1
+	println "Within MACRO"
+	INCLUDE "unique-id-include.inc"
+	println "Outside INCLUDE: \@"
+ENDM
+	m1
+
+println
+
+REPT 2
+	println "Within REPT: \@"
+	INCLUDE "unique-id-include.inc"
+	println "Outside INCLUDE: \@"
+ENDR
+
+MACRO m2
+	println "Within MACRO: \@"
+	INCLUDE "unique-id-include.inc"
+	println "Outside INCLUDE: \@"
+ENDM
+	m2
diff --git a/test/asm/unique-id-include.inc b/test/asm/unique-id-include.inc
new file mode 100644
index 00000000..d9c65540
--- /dev/null
+++ b/test/asm/unique-id-include.inc
@@ -0,0 +1 @@
+println "Within INCLUDE: \@"
diff --git a/test/asm/unique-id-include.out b/test/asm/unique-id-include.out
new file mode 100644
index 00000000..9b4705bd
--- /dev/null
+++ b/test/asm/unique-id-include.out
@@ -0,0 +1,19 @@
+Within REPT
+Within INCLUDE: _u1
+Outside INCLUDE: _u1
+Within REPT
+Within INCLUDE: _u2
+Outside INCLUDE: _u2
+Within MACRO
+Within INCLUDE: _u3
+Outside INCLUDE: _u3
+
+Within REPT: _u4
+Within INCLUDE: _u4
+Outside INCLUDE: _u4
+Within REPT: _u5
+Within INCLUDE: _u5
+Outside INCLUDE: _u5
+Within MACRO: _u6
+Within INCLUDE: _u6
+Outside INCLUDE: _u6
diff --git a/test/asm/unique-id-nested.asm b/test/asm/unique-id-nested.asm
new file mode 100644
index 00000000..94c55bf4
--- /dev/null
+++ b/test/asm/unique-id-nested.asm
@@ -0,0 +1,33 @@
+MACRO m1
+	PRINTLN "Begin MACRO"
+	DEF nested EQUS """MACRO mm
+		PRINTLN "Within nested MACRO: \\@"
+		\n ENDM
+		mm"""
+	nested
+	PURGE nested, mm
+	PRINTLN "Within MACRO: \@"
+ENDM
+REPT 2
+	PRINTLN "Begin REPT"
+	m1
+	PRINTLN "Within REPT: \@"
+ENDR
+
+PRINTLN
+
+MACRO m2
+	PRINTLN "Begin MACRO: \@"
+	DEF nested EQUS """MACRO mm
+		PRINTLN "Within nested MACRO: \\@"
+		\n ENDM
+		mm"""
+	nested
+	PURGE nested, mm
+	PRINTLN "Within MACRO: \@"
+ENDM
+REPT 2
+	PRINTLN "Begin REPT: \@"
+	m2
+	PRINTLN "Within REPT: \@"
+ENDR
diff --git a/test/asm/unique-id-nested.out b/test/asm/unique-id-nested.out
new file mode 100644
index 00000000..04eadb07
--- /dev/null
+++ b/test/asm/unique-id-nested.out
@@ -0,0 +1,21 @@
+Begin REPT
+Begin MACRO
+Within nested MACRO: _u1
+Within MACRO: _u2
+Within REPT: _u3
+Begin REPT
+Begin MACRO
+Within nested MACRO: _u4
+Within MACRO: _u5
+Within REPT: _u6
+
+Begin REPT: _u7
+Begin MACRO: _u8
+Within nested MACRO: _u9
+Within MACRO: _u8
+Within REPT: _u7
+Begin REPT: _u10
+Begin MACRO: _u11
+Within nested MACRO: _u12
+Within MACRO: _u11
+Within REPT: _u10