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.

include/asm/fstack.hpp

@@ -41,6 +41,9 @@ struct FileStackNode {
 	    : type(type_), data(data_){};
 
 	std::string const &dump(uint32_t curLineNo) const;
+
+	// If true, entering this context generates a new unique ID.
+	bool generatesUniqueID() const { return type == NODE_REPT || type == NODE_MACRO; }
 };
 
 #define DEFAULT_MAX_DEPTH 64
@@ -50,6 +53,7 @@ struct MacroArgs;
 
 void fstk_DumpCurrent();
 std::shared_ptr<FileStackNode> fstk_GetFileStack();
+std::shared_ptr<std::string> fstk_GetUniqueIDStr();
 
 void fstk_AddIncludePath(std::string const &path);
 void fstk_SetPreIncludeFile(std::string const &path);

include/asm/macro.hpp

@@ -24,11 +24,6 @@ void macro_UseNewArgs(MacroArgs *args);
 char const *macro_GetArg(uint32_t i);
 char const *macro_GetAllArgs();
 
-uint32_t macro_GetUniqueID();
-char const *macro_GetUniqueIDStr();
-void macro_SetUniqueID(uint32_t id);
-uint32_t macro_UseNewUniqueID();
-uint32_t macro_UndefUniqueID();
 void macro_ShiftCurrentArgs(int32_t count);
 uint32_t macro_NbArgs();
 

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();

src/asm/lexer.cpp

@@ -591,7 +591,8 @@ static char const *readMacroArg(char name) {
 	char const *str = nullptr;
 
 	if (name == '@') {
-		str = macro_GetUniqueIDStr();
+		auto maybeStr = fstk_GetUniqueIDStr();
+		str = maybeStr ? maybeStr->c_str() : nullptr;
 	} else if (name == '#') {
 		str = macro_GetAllArgs();
 	} else if (name == '<') {

src/asm/macro.cpp

@@ -11,10 +11,6 @@
 #define MAXMACROARGS 99999
 
 static MacroArgs *macroArgs = nullptr;
-static uint32_t uniqueID = 0;
-static uint32_t maxUniqueID = 0;
-static char uniqueIDBuf[sizeof("_u4294967295")] = {}; // UINT32_MAX
-static char *uniqueIDPtr = nullptr;
 
 void MacroArgs::append(std::string s) {
 	if (s.empty())
@@ -77,42 +73,6 @@ char const *macro_GetAllArgs() {
 	return str;
 }
 
-uint32_t macro_GetUniqueID() {
-	return uniqueID;
-}
-
-char const *macro_GetUniqueIDStr() {
-	// Generate a new unique ID on the first use of `\@`
-	if (uniqueID == 0)
-		macro_SetUniqueID(++maxUniqueID);
-
-	return uniqueIDPtr;
-}
-
-void macro_SetUniqueID(uint32_t id) {
-	uniqueID = id;
-	if (id == 0 || id == (uint32_t)-1) {
-		uniqueIDPtr = nullptr;
-	} else {
-		// The buffer is guaranteed to be the correct size
-		// This is a valid label fragment, but not a valid numeric
-		sprintf(uniqueIDBuf, "_u%" PRIu32, id);
-		uniqueIDPtr = uniqueIDBuf;
-	}
-}
-
-uint32_t macro_UseNewUniqueID() {
-	// A new ID will be generated on the first use of `\@`
-	macro_SetUniqueID(0);
-	return uniqueID;
-}
-
-uint32_t macro_UndefUniqueID() {
-	// No ID will be generated; use of `\@` is an error
-	macro_SetUniqueID((uint32_t)-1);
-	return uniqueID;
-}
-
 void macro_ShiftCurrentArgs(int32_t count) {
 	if (!macroArgs) {
 		error("Cannot shift macro arguments outside of a macro\n");