cmd/asm: fix handling of nested #if[n]defs

The lexer needs to process all #if[n]defs, even those found when processing is
disabled by a preceding failed conditional, or the first #endif in something
like:

    #ifdef <undefined>
    #ifdef whatever
    #endif
    #endif

terminates the first #ifdef and the second causes an error. And then the
processing of the inner #ifdefs needs to ignore their argument when they are
disabled by an outer failed condition.

Change-Id: Iba259498f1e16042f5b7580b9c000bb0599733d0
Reviewed-on: https://go-review.googlesource.com/14253
Reviewed-by: Rob Pike <r@golang.org>

2 files changed, 115 insertions, 4 deletions

diff --git a/src/cmd/asm/internal/lex/input.go b/src/cmd/asm/internal/lex/input.go
index 7e495b8edf..cd9168064d 100644
--- a/src/cmd/asm/internal/lex/input.go
+++ b/src/cmd/asm/internal/lex/input.go
@@ -136,10 +136,11 @@ func (in *Input) hash() bool {
 		in.expectText("expected identifier after '#'")
 	}
 	if !in.enabled() {
-		// Can only start including again if we are at #else or #endif.
+		// Can only start including again if we are at #else or #endif but also
+		// need to keep track of nested #if[n]defs.
 		// We let #line through because it might affect errors.
 		switch in.Stack.Text() {
-		case "else", "endif", "line":
+		case "else", "endif", "ifdef", "ifndef", "line":
 			// Press on.
 		default:
 			return false
@@ -360,7 +361,9 @@ func (in *Input) collectArgument(macro *Macro) ([]Token, ScanToken) {
 func (in *Input) ifdef(truth bool) {
 	name := in.macroName()
 	in.expectNewline("#if[n]def")
-	if _, defined := in.macros[name]; !defined {
+	if !in.enabled() {
+		truth = false
+	} else if _, defined := in.macros[name]; !defined {
 		truth = !truth
 	}
 	in.ifdefStack = append(in.ifdefStack, truth)
@@ -372,7 +375,9 @@ func (in *Input) else_() {
 	if len(in.ifdefStack) == 0 {
 		in.Error("unmatched #else")
 	}
-	in.ifdefStack[len(in.ifdefStack)-1] = !in.ifdefStack[len(in.ifdefStack)-1]
+	if len(in.ifdefStack) == 1 || in.ifdefStack[len(in.ifdefStack)-2] {
+		in.ifdefStack[len(in.ifdefStack)-1] = !in.ifdefStack[len(in.ifdefStack)-1]
+	}
 }
 
 // #endif processing.
diff --git a/src/cmd/asm/internal/lex/lex_test.go b/src/cmd/asm/internal/lex/lex_test.go
index f034d69c36..32cc13ea66 100644
--- a/src/cmd/asm/internal/lex/lex_test.go
+++ b/src/cmd/asm/internal/lex/lex_test.go
@@ -120,6 +120,112 @@ var lexTests = []lexTest{
 		),
 		"\n.MOVBLZX.(.BP.).(.DX.*.4.).,.R8.\n.\n.MOVBLZX.(.(.8.+.1.).*.4.).(.R12.).,.BX.\n.ADDB.BX.,.DX.\n.MOVB.R8.,.(.8.*.4.).(.R12.).\n.PINSRW.$.0.,.(.BP.).(.R8.*.4.).,.X0.\n",
 	},
+	{
+		"taken #ifdef",
+		lines(
+			"#define A",
+			"#ifdef A",
+			"#define B 1234",
+			"#endif",
+			"B",
+		),
+		"1234.\n",
+	},
+	{
+		"not taken #ifdef",
+		lines(
+			"#ifdef A",
+			"#define B 1234",
+			"#endif",
+			"B",
+		),
+		"B.\n",
+	},
+	{
+		"taken #ifdef with else",
+		lines(
+			"#define A",
+			"#ifdef A",
+			"#define B 1234",
+			"#else",
+			"#define B 5678",
+			"#endif",
+			"B",
+		),
+		"1234.\n",
+	},
+	{
+		"not taken #ifdef with else",
+		lines(
+			"#ifdef A",
+			"#define B 1234",
+			"#else",
+			"#define B 5678",
+			"#endif",
+			"B",
+		),
+		"5678.\n",
+	},
+	{
+		"nested taken/taken #ifdef",
+		lines(
+			"#define A",
+			"#define B",
+			"#ifdef A",
+			"#ifdef B",
+			"#define C 1234",
+			"#else",
+			"#define C 5678",
+			"#endif",
+			"#endif",
+			"C",
+		),
+		"1234.\n",
+	},
+	{
+		"nested taken/not-taken #ifdef",
+		lines(
+			"#define A",
+			"#ifdef A",
+			"#ifdef B",
+			"#define C 1234",
+			"#else",
+			"#define C 5678",
+			"#endif",
+			"#endif",
+			"C",
+		),
+		"5678.\n",
+	},
+	{
+		"nested not-taken/would-be-taken #ifdef",
+		lines(
+			"#define B",
+			"#ifdef A",
+			"#ifdef B",
+			"#define C 1234",
+			"#else",
+			"#define C 5678",
+			"#endif",
+			"#endif",
+			"C",
+		),
+		"C.\n",
+	},
+	{
+		"nested not-taken/not-taken #ifdef",
+		lines(
+			"#ifdef A",
+			"#ifdef B",
+			"#define C 1234",
+			"#else",
+			"#define C 5678",
+			"#endif",
+			"#endif",
+			"C",
+		),
+		"C.\n",
+	},
 }
 
 func TestLex(t *testing.T) {