Merge pull request #95 from sentrysoftware/feature/94-exit-nn-in-begi…

…n-statement-doesnt-properly-return-the-nn-exit-code

Fixed issue #94 with `exit NN` not working in BEGIN and normal rules

src/main/java/org/sentrysoftware/jawk/backend/AVM.java

@@ -194,12 +194,23 @@ public AVM(final AwkSettings parameters, final Map<String, JawkExtension> extens
 	private int oldseed;
 
 	private Address exit_address = null;
+
 	/**
 	 * <code>true</code> if execution position is within an END block;
 	 * <code>false</code> otherwise.
 	 */
 	private boolean within_end_blocks = false;
 
+	/**
+	 * Exit code set by the <code>exit NN</code> command (0 by default)
+	 */
+	private int exit_code = 0;
+
+	/**
+	 * Whether <code>exit</code> has been called and we should throw ExitException
+	 */
+	private boolean throw_exit_exception = false;
+
 	/**
 	 * Maps global variable names to their global array offsets.
 	 * It is useful when passing variable assignments from the file-list
@@ -1782,18 +1793,23 @@ public void interpret(AwkTuples tuples)
 						position.next();
 						break;
 					}
+					case AwkTuples._EXIT_WITHOUT_CODE_:
 					case AwkTuples._EXIT_WITH_CODE_: {
-						// stack[0] = exit code
-						final int exit_code = (int) JRT.toDouble(pop());
+						if (opcode == AwkTuples._EXIT_WITH_CODE_) {
+							// stack[0] = exit code
+							exit_code = (int) JRT.toDouble(pop());
+						}
+						throw_exit_exception = true;
+
+						// If in BEGIN or in a rule, jump to the END section
 						if (!within_end_blocks) {
-							assert exit_address != null;
 							// clear runtime stack
 							runtime_stack.popAllFrames();
 							// clear operand stack
 							operand_stack.clear();
 							position.jump(exit_address);
 						} else {
-							// break;
+							// Exit immediately with ExitException
 							jrt.jrtCloseAll();
 							// clear operand stack
 							operand_stack.clear();
@@ -1957,8 +1973,10 @@ public void interpret(AwkTuples tuples)
 						throw new Error("invalid opcode: " + AwkTuples.toOpcodeString(position.opcode()));
 				}
 			}
+
+			// End of the instructions
 			jrt.jrtCloseAll();
-			assert operand_stack.size() == 0 : "operand stack is NOT empty upon script termination. operand_stack (size=" + operand_stack.size() + ") = " + operand_stack;
+
 		} catch (RuntimeException re) {
 			LOG.error("", re);
 			LOG.error("operand_stack = {}", operand_stack);
@@ -1985,12 +2003,13 @@ public void interpret(AwkTuples tuples)
 				LOG.error("{ could not report on line number", t);
 			}
 			throw ae;
-		} finally {
-//			assert operand_stack.size() == 0 : "operand stack is NOT empty upon script termination. operand_stack (size="+operand_stack.size()+") = "+operand_stack;
-//			if (operand_stack.size() != 0) {
-//				throw new Error("operand stack is NOT empty upon script termination. operand_stack (size=" + operand_stack.size() + ") = " + operand_stack);
-//			}
 		}
+
+		// If <code>exit</code> was called, throw an ExitException
+		if (throw_exit_exception) {
+			throw new ExitException(exit_code, "The AWK script requested an exit");
+		}
+
 	}
 
 	/**

src/main/java/org/sentrysoftware/jawk/frontend/AwkParser.java

@@ -4833,10 +4833,10 @@ public int populateTuples(AwkTuples tuples) {
 			if (ast1 != null) {
 				int ast1_result = ast1.populateTuples(tuples);
 				assert ast1_result == 1;
+				tuples.exitWithCode();
 			} else {
-				tuples.push(0);
+				tuples.exitWithoutCode();
 			}
-			tuples.exitWithCode();
 			popSourceLineNumber(tuples);
 			return 0;
 		}

src/main/java/org/sentrysoftware/jawk/intermediate/AwkTuples.java

@@ -1781,6 +1781,14 @@ private HasFunctionAddress getHasFuncAddr() {
 	 */
 	public static final int _UNARY_PLUS_ = 385;	// x -> -x
 
+	/**
+	 * Terminates execution without specifying an exit code.
+	 * <p>
+	 * Stack before: N/A <br>
+	 * Stack after: N/A
+	 */
+	public static final int _EXIT_WITHOUT_CODE_ = 386;	// 0 -> 0
+
 	/**
 	 * Override add() to populate the line number for each tuple,
 	 * rather than polluting all the constructors with this assignment.
@@ -2887,6 +2895,13 @@ public void exitWithCode() {
 		queue.add(new Tuple(_EXIT_WITH_CODE_));
 	}
 
+	/**
+	 * <p>exitWithCode.</p>
+	 */
+	public void exitWithoutCode() {
+		queue.add(new Tuple(_EXIT_WITHOUT_CODE_));
+	}
+
 	/**
 	 * <p>regexp.</p>
 	 *

src/test/java/org/sentrysoftware/jawk/AwkTest.java

@@ -2,7 +2,9 @@
 
 import static org.junit.Assert.assertArrayEquals;
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertThrows;
 import static org.sentrysoftware.jawk.AwkTestHelper.evalAwk;
+import static org.sentrysoftware.jawk.AwkTestHelper.runAwk;
 
 import java.io.File;
 import java.io.FileNotFoundException;
@@ -209,4 +211,20 @@ public void testNot() throws Exception {
 		assertEquals("!\"a\" must return 0", "0", evalAwk("!\"a\""));
 		assertEquals("!uninitialized must return true", "1", evalAwk("!uninitialized"));
 	}
+
+	@Test
+	public void testExit() throws Exception {
+		assertThrows("exit NN must throw ExitException", ExitException.class, () -> runAwk("BEGIN { exit 17 }", null));
+		assertEquals("exit in BEGIN prevents any rules execution", "", runAwk("BEGIN { exit 0 }\n{ print $0 }", "failure"));
+		assertEquals("exit in BEGIN jumps to END", "success", runAwk("BEGIN { exit 0 ; printf \"failure\" }\nEND { printf \"success\" }", null));
+		assertEquals("exit in END stops immediately", "success", runAwk("END { printf \"success\" ; exit 0 ; printf \"failure\" }", null));
+		assertEquals("exit without code works", "", runAwk("BEGIN { exit }\n{ print $0 }", "failure"));
+		int code = 0;
+		try {
+			runAwk("BEGIN { exit 2 }\nEND { exit }", null);
+		} catch (ExitException e) {
+			code = e.getCode();
+		}
+		assertEquals("exit without code must not alter previous exit with code", 2, code);
+	}
 }

src/test/java/org/sentrysoftware/jawk/AwkTestHelper.java

@@ -61,7 +61,13 @@ static String runAwk(File scriptFile, List<String> inputFileList) throws IOExcep
 
 		// Execute the awk script against the specified input
 		Awk awk = new Awk();
-		awk.invoke(settings);
+		try {
+			awk.invoke(settings);
+		} catch (ExitException e) {
+			if (e.getCode() != 0) {
+				throw e;
+			}
+		}
 
 		// Return the result as a string
 		return resultBytesStream.toString("UTF-8");
@@ -101,7 +107,13 @@ static String runAwk(String script, String input) throws IOException, ExitExcept
 
 		// Execute the awk script against the specified input
 		Awk awk = new Awk();
-		awk.invoke(settings);
+		try {
+			awk.invoke(settings);
+		} catch (ExitException e) {
+			if (e.getCode() != 0) {
+				throw e;
+			}
+		}
 
 		// Return the result as a string
 		return resultBytesStream.toString("UTF-8");