Fix a bug in preprocessor "busy" logic (#1875)

* Fix a bug in preprocessor "busy" logic

This bug manifested in both incorrect preprocessor output for certain complicated cases, and also (more importantly) a use-after-free issue.

One of the "clever" design choices in the Slang preprocessor implementation is that the set of "busy" macros during expansion is implicitly defined by a linked list of those invocations that are actively being read from as part of the input stack. This logic works very well for checking whether a macro name is busy before triggering expansion, and for computing what macros should be considered busy during expansion of an object-like macro.

The problem case here was with function-like macros where the preprocessor was re-using the same list of busy macros that had been fetched when reading the macro name, but doing so *after* all the macro argument tokens had been read. Because additional tokens had been read from the input stream stack, there was no guarantee that the invocations that had been active before were still live.

The new logic computes the set of busy macros fresh before starting expansion of a function-like macro invocation. A test is included to ensure that the case that showed the use-after-free bug has been fixed.

In addition, the new logic is careful to compute the busy macros only based on where subsequent tokens will come from and not based on any macros that might have contributed to the argument tokens themselves. A test case has been added that relies on getting this detail correct.

* Update slang-preprocessor.cpp

Remove a test typo.

Co-authored-by: jsmall-nvidia <jsmall@nvidia.com>

5 files changed, 103 insertions, 12 deletions

diff --git a/source/slang/slang-preprocessor.cpp b/source/slang/slang-preprocessor.cpp
index d54b35d0b..73d9f7e50 100644
--- a/source/slang/slang-preprocessor.cpp
+++ b/source/slang/slang-preprocessor.cpp
@@ -451,6 +451,8 @@ struct InputStreamStack
     }
 
         /// Get the input stream that the next token would come from
+        ///
+        /// If the input stack is at its end, this will just be the top-most stream.
     InputStream* getNextStream()
     {
         SLANG_ASSERT(m_top);
@@ -673,14 +675,13 @@ struct MacroInvocation : InputStream
     MacroInvocation(
         Preprocessor*       preprocessor,
         MacroDefinition*    macro,
-        MacroInvocation*    nextBusyMacroInvocation,
         SourceLoc           macroInvocationLoc,
         SourceLoc           initiatingMacroInvocationLoc);
 
         /// Prime the input stream
         ///
         /// This operation *must* be called before the first `readToken()` or `peekToken()`
-    void prime();
+    void prime(MacroInvocation* nextBusyMacroInvocation);
 
     // The `readToken()` and `peekToken()` operations for a macro invocation
     // will be implemented by using one token of lookahead, which makes the
@@ -729,7 +730,7 @@ private:
     List<Arg> m_args;
 
         /// Additional macros that should be considered "busy" during this expansion
-    MacroInvocation* m_nextBusyMacroInvocation;
+    MacroInvocation* m_nextBusyMacroInvocation = nullptr;
 
         /// Locatin of the macro invocation that led to this expansion
     SourceLoc m_macroInvocationLoc;
@@ -1088,20 +1089,20 @@ bool MacroInvocation::isBusy(MacroDefinition* macro, MacroInvocation* duringMacr
 MacroInvocation::MacroInvocation(
     Preprocessor*       preprocessor,
     MacroDefinition*    macro,
-    MacroInvocation*    nextBusyMacroInvocation,
     SourceLoc           macroInvocationLoc,
     SourceLoc           initiatingMacroInvocationLoc)
     : Super(preprocessor)
 {
     m_macro = macro;
-    m_nextBusyMacroInvocation = nextBusyMacroInvocation;
     m_firstBusyMacroInvocation = this;
     m_macroInvocationLoc = macroInvocationLoc;
     m_initiatingMacroInvocationLoc = initiatingMacroInvocationLoc;
 }
 
-void MacroInvocation::prime()
+void MacroInvocation::prime(MacroInvocation* nextBusyMacroInvocation)
 {
+    m_nextBusyMacroInvocation = nextBusyMacroInvocation;
+
     _initCurrentOpStream();
     m_lookaheadToken = _readTokenImpl();
 }
@@ -1388,8 +1389,12 @@ void ExpansionInputStream::_maybeBeginMacroInvocation()
                 // prime its input stream, and then push it onto our stack of active
                 // macro invocations.
                 //
-                MacroInvocation* invocation = new MacroInvocation(preprocessor, macro, busyMacros, token.loc, m_initiatingMacroInvocationLoc);
-                invocation->prime();
+                // Note: the macros that should be considered "busy" during the invocation
+                // are all those that were busy at the time we read the name of the macro
+                // to be expanded.
+                //
+                MacroInvocation* invocation = new MacroInvocation(preprocessor, macro, token.loc, m_initiatingMacroInvocationLoc);
+                invocation->prime(busyMacros);
                 _pushMacroInvocation(invocation);
             }
             break;
@@ -1441,7 +1446,7 @@ void ExpansionInputStream::_maybeBeginMacroInvocation()
                 // If we saw an opening `(`, then we know we are starting some kind of
                 // macro invocation, although we don't yet know if it is well-formed.
                 //
-                MacroInvocation* invocation = new MacroInvocation(preprocessor, macro, busyMacros, token.loc, m_initiatingMacroInvocationLoc);
+                MacroInvocation* invocation = new MacroInvocation(preprocessor, macro, token.loc, m_initiatingMacroInvocationLoc);
 
                 // We start by consuming the opening `(` that we checked for above.
                 //
@@ -1483,7 +1488,36 @@ void ExpansionInputStream::_maybeBeginMacroInvocation()
                 // Now that the arguments have been parsed and validated,
                 // we are ready to proceed with expansion of the macro invocation.
                 //
-                invocation->prime();
+                // The main subtle thing we have to figure out is which macros should be considered "busy"
+                // during the expansion of this function-like macro invocation.
+                //
+                // In the case of an object-like macro invocation:
+                //
+                //      1 + M + 2
+                //            ^
+                //
+                // Input will have just read in the `M` token that names the macro
+                // so we needed to consider whatever macro invocations had been in
+                // flight (even if they were at their end) when checking if `M`
+                // was busy.
+                //
+                // In contrast, for a function-like macro invocation:
+                //
+                //      1 + F ( A, B, C ) + 2
+                //                        ^
+                //
+                // We will have just read in the `)` from the argument list, but
+                // we don't actually need/want to worry about any macro invocation
+                // that might have yielded the `)` token, since expanding that macro
+                // again would *not* be able to lead to a recursive case.
+                //
+                // Instead, we really only care about the active stream that the
+                // next token would be read from.
+                //
+                auto nextStream = m_inputStreams.getNextStream();
+                auto busyMacrosForFunctionLikeInvocation = nextStream->getFirstBusyMacroInvocation();
+
+                invocation->prime(busyMacrosForFunctionLikeInvocation);
                 _pushMacroInvocation(invocation);
             }
             break;
@@ -3708,8 +3742,12 @@ Result findMacroValue(
     if(macro->flavor != MacroDefinition::Flavor::ObjectLike)
         return SLANG_FAIL;
 
-    MacroInvocation* invocation = new MacroInvocation(preprocessor, macro, nullptr, SourceLoc(), SourceLoc());
-    invocation->prime();
+    MacroInvocation* invocation = new MacroInvocation(preprocessor, macro, SourceLoc(), SourceLoc());
+
+    // Note: Since we are only expanding the one macro, we should not treat any
+    // other macros as "busy" at the start of expansion.
+    //
+    invocation->prime(/*nextBusyMacroInvocation:*/ nullptr);
 
     String value;
     for(bool first = true;;first = false)
diff --git a/tests/preprocessor/bugs/pp-bug-1.slang b/tests/preprocessor/bugs/pp-bug-1.slang
new file mode 100644
index 000000000..b0f026a6b
--- /dev/null
+++ b/tests/preprocessor/bugs/pp-bug-1.slang
@@ -0,0 +1,20 @@
+// pp-bug-1.slang
+//DIAGNOSTIC_TEST:SIMPLE:-E
+
+// The bug in this case was related to a use-after-free
+// where the list of "busy" macros for a function-like
+// macro invocation was being set based on what was
+// busy when reading the macro name, which could include
+// streams that had been popped by the time arguments
+// had been read.
+
+#define NUM_CASES 2
+#define X1(M) M(0)
+#define X2(M) M(0) M(1)
+#define CONCAT2(a, b) a##b
+#define CONCAT(a, b) CONCAT2(a, b)
+#define FOREACH(M) CONCAT(X, NUM_CASES)(M)
+#define CASE(i) i
+// Should output: 0 1
+// Outputs: 0 CASE ( 1 )
+FOREACH(CASE)
\ No newline at end of file
diff --git a/tests/preprocessor/bugs/pp-bug-1.slang.expected b/tests/preprocessor/bugs/pp-bug-1.slang.expected
new file mode 100644
index 000000000..522fc1b40
--- /dev/null
+++ b/tests/preprocessor/bugs/pp-bug-1.slang.expected
@@ -0,0 +1,6 @@
+result code = 0
+standard error = {
+}
+standard output = {
+0 1 
+}
diff --git a/tests/preprocessor/macros/macro-parens-from-expansion.slang b/tests/preprocessor/macros/macro-parens-from-expansion.slang
new file mode 100644
index 000000000..3886c8db9
--- /dev/null
+++ b/tests/preprocessor/macros/macro-parens-from-expansion.slang
@@ -0,0 +1,21 @@
+// macro-parens-from-expansion.slang
+//DIAGNOSTIC_TEST:SIMPLE:-E
+
+// This test covers a case where the parentheses for
+// a function-like macro invocation argument list are
+// themselves produced by a macro invocation.
+//
+// In the below code, it is important that the macros
+// `LPAREN` and `RPAREN` are not considered busy for
+// the invocation of `INNER`, despite the fact that
+// both of those macros were expanded as part of
+// producing the arguments to the `INNER` invocation.
+
+#define LPAREN (
+#define RPAREN )
+#define INNER(X) LPAREN X RPAREN
+#define M(X) X
+#define OUTER(X) M( INNER LPAREN X RPAREN )
+
+// output: ( 3 )
+OUTER(3)
diff --git a/tests/preprocessor/macros/macro-parens-from-expansion.slang.expected b/tests/preprocessor/macros/macro-parens-from-expansion.slang.expected
new file mode 100644
index 000000000..96558709e
--- /dev/null
+++ b/tests/preprocessor/macros/macro-parens-from-expansion.slang.expected
@@ -0,0 +1,6 @@
+result code = 0
+standard error = {
+}
+standard output = {
+( 3 ) 
+}