Optimize chained string concatenation with counted CONCAT

bertysentry · bertysentry · commit bf22b460bfc9 · 2026-06-02T19:30:52.000+02:00
diff --git a/src/jmh/java/io/jawk/backend/AVMExpressionBenchmark.java b/src/jmh/java/io/jawk/backend/AVMExpressionBenchmark.java
@@ -61,6 +61,9 @@ public class AVMExpressionBenchmark {
 	private AwkExpression fieldConcatenation;
 	private AwkExpression fieldRegexMatch;
 	private AwkExpression multiStringConcatenation;
+	private AwkExpression constantStringConcatenation;
+	private AwkExpression stringConstantStringConstantConcatenation;
+	private AwkExpression fourStringConcatenation;
 	private AwkExpression mixedExpression;
 
 	/**
@@ -78,6 +81,9 @@ public void setup() throws IOException {
 		this.fieldConcatenation = awk.compileExpression("$1 \" test\"");
 		this.fieldRegexMatch = awk.compileExpression("$1 ~ /test/");
 		this.multiStringConcatenation = awk.compileExpression("$1 \" test1\" \" test2\" \" test3\"");
+		this.constantStringConcatenation = awk.compileExpression("\"constant\" \"constant\" \"constant\" \"constant\"");
+		this.stringConstantStringConstantConcatenation = awk.compileExpression("$1 \"constant\" $2 \"constant\"");
+		this.fourStringConcatenation = awk.compileExpression("$1 $2 $3 $4");
 		this.mixedExpression = awk.compileExpression("($1 + $2) \":\" ($3 ~ /test/) \":\" $4");
 		this.avm = new AVM(new AwkSettings(), Collections.emptyMap());
 		this.avm.prepareForEval("42 3.14 test-value suffix");
@@ -159,6 +165,39 @@ public Object multiStringConcatenation() throws IOException {
 		return this.avm.eval(this.multiStringConcatenation);
 	}
 
+	/**
+	 * Measures concatenation of four constant string operands.
+	 *
+	 * @return expression result
+	 * @throws IOException if input preparation or evaluation fails
+	 */
+	@Benchmark
+	public Object constantStringConcatenation() throws IOException {
+		return this.avm.eval(this.constantStringConcatenation);
+	}
+
+	/**
+	 * Measures alternating field and constant string concatenation.
+	 *
+	 * @return expression result
+	 * @throws IOException if input preparation or evaluation fails
+	 */
+	@Benchmark
+	public Object stringConstantStringConstantConcatenation() throws IOException {
+		return this.avm.eval(this.stringConstantStringConstantConcatenation);
+	}
+
+	/**
+	 * Measures concatenation of four field string operands.
+	 *
+	 * @return expression result
+	 * @throws IOException if input preparation or evaluation fails
+	 */
+	@Benchmark
+	public Object fourStringConcatenation() throws IOException {
+		return this.avm.eval(this.fourStringConcatenation);
+	}
+
 	/**
 	 * Measures mixed numeric, string, field, and regular expression operations.
 	 *
diff --git a/src/main/java/io/jawk/backend/AVM.java b/src/main/java/io/jawk/backend/AVM.java
@@ -1168,12 +1168,25 @@ private void executeTuples(PositionTracker position)
 					break;
 				}
 				case CONCAT: {
-					// stack[0] = string1
-					// stack[1] = string2
-					String s2 = jrt.toAwkString(pop());
-					String s1 = jrt.toAwkString(pop());
-					String resultString = s1 + s2;
-					push(resultString);
+					// arg[0] = number of stack items to concatenate
+					// stack[0] = last concatenation operand
+					CountTuple countTuple = (CountTuple) tuple;
+					int count = (int) countTuple.getCount();
+					// Store String references so appends run left-to-right. Converting
+					// operands to char[] would copy them once before StringBuilder
+					// copies them again, and front-inserting would shift existing
+					// content on each operand.
+					String[] values = new String[count];
+					int resultLength = 0;
+					for (int i = count - 1; i >= 0; i--) {
+						values[i] = jrt.toAwkString(pop());
+						resultLength += values[i].length();
+					}
+					StringBuilder resultString = new StringBuilder(resultLength);
+					for (String value : values) {
+						resultString.append(value);
+					}
+					push(resultString.toString());
 					position.next();
 					break;
 				}
diff --git a/src/main/java/io/jawk/frontend/AwkParser.java b/src/main/java/io/jawk/frontend/AwkParser.java
@@ -3628,12 +3628,25 @@ private ConcatExpressionAst(AST lhs, AST rhs) {
 		@Override
 		public int populateTuples(AwkTuples tuples) {
 			pushSourceLineNumber(tuples);
-			getAst1().populateTuples(tuples);
-			getAst2().populateTuples(tuples);
-			tuples.concat();
+			List<AST> operands = new ArrayList<>();
+			collectConcatOperands(this, operands);
+			for (AST operand : operands) {
+				operand.populateTuples(tuples);
+			}
+			tuples.concat(operands.size());
 			popSourceLineNumber(tuples);
 			return 1;
 		}
+
+		private void collectConcatOperands(AST ast, List<AST> operands) {
+			if (ast instanceof ConcatExpressionAst) {
+				ConcatExpressionAst concat = (ConcatExpressionAst) ast;
+				collectConcatOperands(concat.getAst1(), operands);
+				collectConcatOperands(concat.getAst2(), operands);
+			} else {
+				operands.add(ast);
+			}
+		}
 	}
 
 	private final class NegativeExpressionAst extends ScalarExpressionAst {
diff --git a/src/main/java/io/jawk/intermediate/AwkTuples.java b/src/main/java/io/jawk/intermediate/AwkTuples.java
@@ -32,6 +32,7 @@
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.IdentityHashMap;
+import java.util.List;
 import java.util.Map;
 import java.util.Set;
 import java.util.function.Supplier;
@@ -69,7 +70,7 @@ public class AwkTuples implements Serializable {
 	 * can be serialized and patched efficiently. A linked list would make every
 	 * lookup O(n) and complicate address reassignment.
 	 */
-	private java.util.List<Tuple> queue = new ArrayList<Tuple>(100) {
+	private List<Tuple> queue = new ArrayList<Tuple>(100) {
 		private static final long serialVersionUID = -6334362156408598578L;
 
 		@Override
@@ -296,11 +297,25 @@ public void length(int numExprs) {
 
 	/**
 	 * <p>
-	 * concat.
+	 * Concatenates two stack items.
 	 * </p>
 	 */
 	public void concat() {
-		queue.add(new Tuple.NoOperandTuple(Opcode.CONCAT));
+		concat(2);
+	}
+
+	/**
+	 * <p>
+	 * Concatenates the requested number of stack items.
+	 * </p>
+	 *
+	 * @param count the number of stack items to concatenate
+	 */
+	public void concat(int count) {
+		if (count < 2) {
+			throw new IllegalArgumentException("CONCAT requires at least two stack items");
+		}
+		queue.add(new Tuple.CountTuple(Opcode.CONCAT, count));
 	}
 
 	/**
@@ -1878,10 +1893,10 @@ private boolean peepholeOptimizePass() {
 			return false;
 		}
 
-		java.util.List<Tuple> original = new ArrayList<Tuple>(queue);
+		List<Tuple> original = new ArrayList<Tuple>(queue);
 		int[] indexMapping = new int[originalSize];
 		Arrays.fill(indexMapping, -1);
-		java.util.List<Tuple> optimizedQueue = new ArrayList<Tuple>(originalSize);
+		List<Tuple> optimizedQueue = new ArrayList<Tuple>(originalSize);
 		boolean[] isAddressTarget = addressTargets(original, originalSize);
 
 		boolean modified = false;
@@ -1929,6 +1944,20 @@ private boolean peepholeOptimizePass() {
 						continue;
 					}
 				}
+				Tuple countedStringConcat = foldCountedStringConcat(original, oldIndex);
+				if (countedStringConcat != null) {
+					int count = countStringPushes(original, oldIndex);
+					// Fold a counted literal-only concatenation into one string push,
+					// e.g. PUSH_STRING "a", PUSH_STRING "b", CONCAT 2 -> PUSH_STRING
+					// "ab". Numeric literals are deliberately not folded here because
+					// their string representation depends on runtime formatting state.
+					optimizedQueue.add(countedStringConcat);
+					mapFoldedRange(indexMapping, oldIndex, count + 1, newIndex);
+					oldIndex += count + 1;
+					newIndex++;
+					modified = true;
+					continue;
+				}
 				if ((oldIndex + 2) < originalSize) {
 					Tuple nextTuple = original.get(oldIndex + 1);
 					Tuple opTuple = original.get(oldIndex + 2);
@@ -1987,7 +2016,7 @@ private boolean peepholeOptimizePass() {
 		return true;
 	}
 
-	private boolean[] addressTargets(java.util.List<Tuple> tuples, int tupleCount) {
+	private boolean[] addressTargets(List<Tuple> tuples, int tupleCount) {
 		boolean[] targets = new boolean[tupleCount];
 		for (Tuple tuple : tuples) {
 			Address address = tuple.getAddress();
@@ -2020,6 +2049,47 @@ private Object literalValue(Tuple tuple) {
 		}
 	}
 
+	private Tuple foldCountedStringConcat(List<Tuple> original, int oldIndex) {
+		Tuple firstTuple = original.get(oldIndex);
+		if (firstTuple.getOpcode() != Opcode.PUSH_STRING) {
+			return null;
+		}
+
+		int tupleCount = original.size();
+		StringBuilder folded = new StringBuilder();
+		int itemCount = 0;
+		int currentIndex = oldIndex;
+		while (currentIndex < tupleCount && original.get(currentIndex).getOpcode() == Opcode.PUSH_STRING) {
+			folded.append(((Tuple.PushStringTuple) original.get(currentIndex)).getValue());
+			itemCount++;
+			currentIndex++;
+		}
+
+		if (itemCount < 2 || currentIndex >= tupleCount) {
+			return null;
+		}
+
+		Tuple operation = original.get(currentIndex);
+		if (operation.getOpcode() != Opcode.CONCAT || !(operation instanceof Tuple.CountTuple)) {
+			return null;
+		}
+
+		long count = ((Tuple.CountTuple) operation).getCount();
+		if (count != itemCount) {
+			return null;
+		}
+
+		return createLiteralPush(folded.toString(), firstTuple.getLineNumber());
+	}
+
+	private int countStringPushes(List<Tuple> original, int oldIndex) {
+		int currentIndex = oldIndex;
+		while (currentIndex < original.size() && original.get(currentIndex).getOpcode() == Opcode.PUSH_STRING) {
+			currentIndex++;
+		}
+		return currentIndex - oldIndex;
+	}
+
 	private Object foldBinary(Object left, Object right, Tuple operation) {
 		Opcode opcode = operation.getOpcode();
 		if (opcode == null) {
diff --git a/src/main/java/io/jawk/intermediate/Opcode.java b/src/main/java/io/jawk/intermediate/Opcode.java
@@ -189,11 +189,13 @@ public enum Opcode {
 	 */
 	LENGTH,
 	/**
-	 * Pop and concatenate two strings from the top-of-stack; push the result onto
-	 * the stack.
+	 * Pops and concatenates N strings from the top-of-stack; push the result onto
+	 * the stack. The number of items is passed in as a tuple argument.
+	 * <p>
+	 * Argument: # of items (N)
 	 * <p>
-	 * Stack before: x y ...<br/>
-	 * Stack after: x-concatenated-with-y ...
+	 * Stack before: x1 x2 x3 .. xN ...<br/>
+	 * Stack after: x1-concatenated-through-xN ...
 	 */
 	CONCAT,
 	/**
diff --git a/src/test/java/io/jawk/AwkTupleOptimizationTest.java b/src/test/java/io/jawk/AwkTupleOptimizationTest.java
@@ -139,6 +139,34 @@ public void foldsLiteralStringConcatenation() throws Exception {
 		assertTrue("Expected folded literal push of foobar", hasLiteralPush(tuples, "foobar"));
 	}
 
+	@Test
+	public void foldsChainedLiteralStringConcatenation() throws Exception {
+		String script = "BEGIN { print \"foo\" \"bar\" \"baz\" \"qux\" }\n";
+		AwkTestSupport
+				.awkTest("folds chained literal string concatenation")
+				.script(script)
+				.expect("foobarbazqux\n")
+				.runAndAssert();
+
+		AwkProgram tuples = new Awk().compile(script);
+		List<Opcode> opcodes = collectOpcodes(tuples);
+		assertFalse("Chained literal concatenation should eliminate CONCAT tuple", opcodes.contains(Opcode.CONCAT));
+		assertTrue("Expected folded literal push of foobarbazqux", hasLiteralPush(tuples, "foobarbazqux"));
+	}
+
+	@Test
+	public void compilesChainedStringConcatenationAsSingleCountedConcat() throws Exception {
+		String script = "BEGIN { s1 = \"alpha\"; s2 = \"beta\"; print s1 \"-\" s2 \":\" }\n";
+		AwkTestSupport
+				.awkTest("counted chained string concatenation")
+				.script(script)
+				.expect("alpha-beta:\n")
+				.runAndAssert();
+
+		AwkProgram tuples = new Awk().compile(script);
+		assertEquals("Expected one counted CONCAT for the mixed chain", 1, countOpcodeWithCount(tuples, Opcode.CONCAT, 4));
+	}
+
 	@Test
 	public void foldsScalarAssignmentPopIntoNonPushingAssignment() throws Exception {
 		String script = "BEGIN { a = -2; b = 2; c = 4; print a + b + c }\n";
@@ -204,6 +232,10 @@ public void doesNotFoldNumericConcatenation() throws Exception {
 		AwkProgram tuples = new Awk().compile(script);
 		List<Opcode> opcodes = collectOpcodes(tuples);
 		assertTrue("Numeric literal concatenation should preserve CONCAT tuple", opcodes.contains(Opcode.CONCAT));
+		assertEquals(
+				"Numeric literal concatenation should remain binary",
+				1,
+				countOpcodeWithCount(tuples, Opcode.CONCAT, 2));
 		assertFalse("Optimizer should not fold numeric/string concatenation", hasLiteralPush(tuples, "1x"));
 	}
 
@@ -570,6 +602,20 @@ private static int countOpcode(AwkProgram tuples, Opcode opcode) {
 		return count;
 	}
 
+	private static int countOpcodeWithCount(AwkProgram tuples, Opcode opcode, long expectedCount) {
+		int count = 0;
+		PositionTracker tracker = rawTuples(tuples).top();
+		while (!tracker.isEOF()) {
+			if (tracker.opcode() == opcode
+					&& tracker.current() instanceof Tuple.CountTuple
+					&& ((Tuple.CountTuple) tracker.current()).getCount() == expectedCount) {
+				count++;
+			}
+			tracker.next();
+		}
+		return count;
+	}
+
 	private static String dumpTuples(AwkProgram tuples) throws Exception {
 		ByteArrayOutputStream out = new ByteArrayOutputStream();
 		try (PrintStream ps = new PrintStream(out, true, StandardCharsets.UTF_8.name())) {
