Simplify \@ handling by using std::shared_ptr<std::string>

This has been relocated from macro.cpp to fstack.cpp, since both
MACRO and REPT/FOR nodes have their own unique `\@` values.

src/asm/fstack.cpp

@@ -7,6 +7,7 @@
 #include <assert.h>
 #include <errno.h>
 #include <inttypes.h>
+#include <memory>
 #include <stack>
 #include <stdio.h>
 #include <stdlib.h>
@@ -23,10 +24,17 @@
 #include "asm/symbol.hpp"
 #include "asm/warning.hpp"
 
+using namespace std::literals;
+
 struct Context {
 	std::shared_ptr<FileStackNode> fileInfo;
 	LexerState lexerState{};
-	uint32_t uniqueID = 0;
+	// If the shared_ptr is empty, `\@` is not permitted for this context.
+	// Otherwise, if the pointee string is empty, it means that a unique ID has not been requested
+	// for this context yet, and it should be generated.
+	// Note that several contexts can share the same unique ID (since `INCLUDE` preserves its
+	// parent's, and likewise "back-propagates" a unique ID if requested), hence using `shared_ptr`.
+	std::shared_ptr<std::string> uniqueIDStr;
 	MacroArgs *macroArgs = nullptr; // Macro args are *saved* here
 	uint32_t nbReptIters = 0;
 	bool isForLoop = false;
@@ -99,6 +107,18 @@ std::shared_ptr<FileStackNode> fstk_GetFileStack() {
 	return contextStack.empty() ? nullptr : contextStack.top().fileInfo;
 }
 
+std::shared_ptr<std::string> fstk_GetUniqueIDStr() {
+	static uint64_t nextUniqueID = 1;
+
+	std::shared_ptr<std::string> &str = contextStack.top().uniqueIDStr;
+
+	// If a unique ID is allowed but has not been generated yet, generate one now.
+	if (str && str->empty())
+		*str = "_u"s + std::to_string(nextUniqueID++);
+
+	return str;
+}
+
 void fstk_AddIncludePath(std::string const &path) {
 	if (path.empty())
 		return;
@@ -187,7 +207,7 @@ bool yywrap() {
 		// If this wasn't the last iteration, wrap instead of popping
 		if (fileInfoIters.front() <= context.nbReptIters) {
 			lexer_RestartRept(context.fileInfo->lineNo);
-			context.uniqueID = macro_UseNewUniqueID();
+			context.uniqueIDStr->clear(); // Invalidate the current unique ID (if any).
 			return false;
 		}
 	} else if (contextStack.size() == 1) {
@@ -202,11 +222,6 @@ bool yywrap() {
 	// 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);
 
 	lexer_SetState(&contextStack.top().lexerState);
 
@@ -220,13 +235,15 @@ static Context &newContext(std::shared_ptr<FileStackNode> fileInfo) {
 	if (contextStack.size() > maxRecursionDepth)
 		fatalerror("Recursion limit (%zu) exceeded\n", maxRecursionDepth);
 
-	// Save the current `\@` value, to be restored when this context ends
-	contextStack.top().uniqueID = macro_GetUniqueID();
-
 	fileInfo->parent = contextStack.top().fileInfo;
 	fileInfo->lineNo = lexer_GetLineNo();
 
-	return contextStack.emplace(Context{.fileInfo = fileInfo});
+	return contextStack.emplace(Context{
+	    .fileInfo = fileInfo,
+	    .uniqueIDStr = fileInfo->generatesUniqueID()
+	                       ? std::make_shared<std::string>() // Create a new, not-yet-generated ID.
+	                       : contextStack.top().uniqueIDStr, // Make a copy.
+	});
 }
 
 void fstk_RunInclude(std::string const &path) {
@@ -242,17 +259,11 @@ void fstk_RunInclude(std::string const &path) {
 		return;
 	}
 
-	uint32_t uniqueID = contextStack.top().uniqueID;
-
 	auto fileInfo = std::make_shared<FileStackNode>(NODE_FILE, *fullPath);
 	Context &context = newContext(fileInfo);
 	if (!lexer_OpenFile(context.lexerState, fileInfo->name()))
 		fatalerror("Failed to set up lexer for file include\n");
 	lexer_SetStateAtEOL(&context.lexerState);
-	// We're back at top-level, so most things are reset,
-	// but not the unique ID, since INCLUDE may be inside a
-	// MACRO or REPT/FOR loop
-	context.uniqueID = uniqueID;
 }
 
 // Similar to `fstk_RunInclude`, but not subject to `-MG`, and
@@ -272,8 +283,6 @@ static void runPreIncludeFile() {
 	if (!lexer_OpenFile(context.lexerState, fileInfo->name()))
 		fatalerror("Failed to set up lexer for file include\n");
 	lexer_SetState(&context.lexerState);
-	// We're back at top-level, so most things are reset
-	context.uniqueID = macro_UndefUniqueID();
 }
 
 void fstk_RunMacro(std::string const &macroName, MacroArgs &args) {
@@ -316,7 +325,6 @@ void fstk_RunMacro(std::string const &macroName, MacroArgs &args) {
 	    context.lexerState, "MACRO", macroView->data(), macroView->size(), macro->fileLine
 	);
 	lexer_SetStateAtEOL(&context.lexerState);
-	context.uniqueID = macro_UseNewUniqueID();
 	macro_UseNewArgs(&args);
 }
 
@@ -336,7 +344,6 @@ static bool newReptContext(int32_t reptLineNo, char const *body, size_t size) {
 	context.fileInfo->lineNo = reptLineNo;
 	lexer_OpenFileView(context.lexerState, "REPT", body, size, reptLineNo);
 	lexer_SetStateAtEOL(&context.lexerState);
-	context.uniqueID = macro_UseNewUniqueID();
 
 	return true;
 }
@@ -414,7 +421,10 @@ void fstk_NewRecursionDepth(size_t newDepth) {
 }
 
 void fstk_Init(std::string const &mainPath, size_t maxDepth) {
-	Context &context = contextStack.emplace();
+	Context &context = contextStack.emplace(Context{
+	    .fileInfo = nullptr,    // We're going to init it just below.
+	    .uniqueIDStr = nullptr, // `\@` is not allowed at top level.
+	});
 	if (!lexer_OpenFile(context.lexerState, mainPath))
 		fatalerror("Failed to open main file\n");
 	lexer_SetState(&context.lexerState);
@@ -424,8 +434,6 @@ void fstk_Init(std::string const &mainPath, size_t maxDepth) {
 	context.fileInfo->parent = nullptr;
 	context.fileInfo->lineNo = 0; // This still gets written to the object file, so init it
 
-	context.uniqueID = macro_UndefUniqueID();
-
 	maxRecursionDepth = maxDepth;
 
 	runPreIncludeFile();