fix: Expression statements not emitting local variables correctly

pom.xml

@@ -3,7 +3,7 @@
 	<groupId>me.coley</groupId>
 	<artifactId>recaf</artifactId>
 	<url>https://github.com/Col-E/Recaf/</url>
-	<version>2.16.0</version>
+	<version>2.16.1</version>
 	<name>Recaf</name>
 	<description>A modern java bytecode editor</description>
 	<!-- Variables -->

src/main/java/me/coley/recaf/Recaf.java

@@ -30,7 +30,7 @@
  * @author Matt
  */
 public class Recaf {
-	public static final String VERSION = "2.16.0";
+	public static final String VERSION = "2.16.1";
 	public static final String DOC_URL = "https://col-e.github.io/Recaf/documentation.html";
 	public static final int ASM_VERSION = Opcodes.ASM9;
 	private static Controller currentController;

src/main/java/me/coley/recaf/parse/bytecode/JavassistCompiler.java

@@ -4,6 +4,9 @@
 import javassist.bytecode.Bytecode;
 import javassist.bytecode.LocalVariableAttribute;
 import javassist.compiler.*;
+import javassist.compiler.ast.Declarator;
+import javassist.compiler.ast.Stmnt;
+import org.objectweb.asm.Type;
 import org.objectweb.asm.tree.LocalVariableNode;
 
 import java.lang.reflect.Field;
@@ -50,31 +53,38 @@ public static CtMethod compileMethod(CtClass declaring, String src) throws Canno
 	 * 		Declaring method that will contain the expression.
 	 * @param expression
 	 * 		Source of the expression.
-	 * @param variables Variable information to populate.
+	 * @param existingVars
+	 * 		Variable information to populate.
+	 * @param varCache
+	 * Variable name and index information.
 	 *
-	 * @return Compiled method.
+	 * @return Compiled expression.
 	 *
 	 * @throws CannotCompileException
 	 * 		When a compilation error occurred.
 	 */
-	public static Bytecode compileExpression(CtClass declaring, CtBehavior containerMethod, String expression,
-											 List<LocalVariableNode> variables) throws CannotCompileException {
+	public static CompilationResult compileExpression(CtClass declaring, CtBehavior containerMethod, String expression,
+													  List<LocalVariableNode> existingVars, VariableNameCache varCache)
+			throws CannotCompileException {
 		try {
-			InternalJavac compiler = new InternalJavac(declaring);
-			populateVariables(compiler, variables);
+			InternalJavac compiler = new InternalJavac(declaring, varCache);
+			populateVariables(compiler, existingVars);
 			populateVariables(compiler, containerMethod);
-			// TODO: Output variables so we can have one expression compile, then following code can access the vars
-			//  - compiler.stable.append("varName", new Declarator(type, internal, arrayDim, index, symbol));
 			compiler.compileStmnt(expression);
-			return compiler.getGeneratedBytecode();
+			patchGeneratedVars(compiler, varCache);
+			return new CompilationResult(compiler.getGeneratedBytecode(), compiler.getLastCompiledSymbols());
 		} catch (CompileError e) {
 			throw new CannotCompileException(e);
 		}
 	}
 
+	private static void patchGeneratedVars(InternalJavac compiler, VariableNameCache varCache) {
+
+	}
+
 	private static void populateVariables(InternalJavac compiler, List<LocalVariableNode> variables) {
 		JvstCodeGen gen = compiler.getGen();
-		SymbolTable symbolTable = compiler.getSTable();
+		SymbolTable symbolTable = compiler.getRootSTable();
 		for (LocalVariableNode variable : variables) {
 			try {
 				gen.recordVariable(variable.desc, variable.name, variable.index, symbolTable);
@@ -89,7 +99,7 @@ private static void populateVariables(InternalJavac compiler, CtBehavior contain
 				containerMethod.getMethodInfo().getCodeAttribute().getAttribute(LocalVariableAttribute.tag);
 		if (variables != null) {
 			JvstCodeGen gen = compiler.getGen();
-			SymbolTable symbolTable = compiler.getSTable();
+			SymbolTable symbolTable = compiler.getRootSTable();
 			for (int i = 0; i < variables.tableLength(); i++) {
 				int index = variables.index(i);
 				String signature = variables.signature(i);
@@ -103,6 +113,39 @@ private static void populateVariables(InternalJavac compiler, CtBehavior contain
 		}
 	}
 
+	/**
+	 * Compilation results.
+	 */
+	public static class CompilationResult {
+		private final Bytecode bytecode;
+		private final SymbolTable symbols;
+
+		/**
+		 * @param bytecode
+		 * 		Generated bytecode.
+		 * @param symbols
+		 * 		Generated symbols, if any.
+		 */
+		public CompilationResult(Bytecode bytecode, SymbolTable symbols) {
+			this.bytecode = bytecode;
+			this.symbols = symbols;
+		}
+
+		/**
+		 * @return Generated bytecode.
+		 */
+		public Bytecode getBytecode() {
+			return bytecode;
+		}
+
+		/**
+		 * @return Generated symbols, if any.
+		 */
+		public SymbolTable getSymbols() {
+			return symbols;
+		}
+	}
+
 	/**
 	 * An extension of Javassist's {@link Javac} that exposes some internal structures
 	 * needed to properly inject local variable information.
@@ -113,9 +156,70 @@ private static class InternalJavac extends Javac {
 		private static final Field fGen;
 		private static final Field fSTable;
 		private static final Field fBytecode;
+		private final VariableNameCache varCache;
+		private SymbolTable lastCompiledSymbols;
 
-		public InternalJavac(CtClass declaring) {
+		public InternalJavac(CtClass declaring, VariableNameCache varCache) {
 			super(declaring);
+			this.varCache = varCache;
+		}
+
+		@Override
+		public void compileStmnt(String src) throws CompileError {
+			Parser p = new Parser(new Lex(src));
+			lastCompiledSymbols = new SymbolTable(getRootSTable());
+			while (p.hasMore()) {
+				Stmnt s = p.parseStatement(lastCompiledSymbols);
+				// Patch the index so the following "accept" call doesn't generate with the wrong var index
+				if (varCache != null && s.getLeft() instanceof Declarator)
+					patchDeclarator(varCache, (Declarator) s.getLeft());
+				// Generate bytecode
+				if (s != null)
+					s.accept(getGen());
+			}
+		}
+
+		private void patchDeclarator(VariableNameCache varCache, Declarator dec) {
+			String name = dec.getLeft().toString();
+			// Update variable index if it exists already
+			try {
+				int index = varCache.getIndex(name);
+				dec.setLocalVar(index);
+				return;
+			} catch (Exception ignored) {
+				// ignored
+			}
+			// Otherwise define it
+			String desc = dec.getClassName();
+			if (desc == null) {
+				switch (dec.getType()) {
+					case TokenId.BOOLEAN:
+					case TokenId.BYTE:
+					case TokenId.SHORT:
+					case TokenId.INT:
+						desc = "I";
+						break;
+					case TokenId.CHAR:
+						desc = "C";
+						break;
+					case TokenId.FLOAT:
+						desc = "F";
+						break;
+					case TokenId.DOUBLE:
+						desc = "D";
+						break;
+					case TokenId.LONG:
+						desc = "J";
+						break;
+					default:
+						throw new IllegalArgumentException("Unknown primitive type for expression defined var");
+				}
+			}
+			Type type = Type.getType(desc);
+			int index = varCache.getAndIncrementNext(name, type);
+			dec.setLocalVar(index);
+			dec.setClassName(type.getClassName());
+			setMaxLocals(index);
 		}
 
 		public JvstCodeGen getGen() {
@@ -134,14 +238,18 @@ public Bytecode getGeneratedBytecode() {
 			}
 		}
 
-		public SymbolTable getSTable() {
+		public SymbolTable getRootSTable() {
 			try {
 				return (SymbolTable) fSTable.get(this);
 			} catch (IllegalAccessException ex) {
 				throw new IllegalStateException(ex);
 			}
 		}
 
+		public SymbolTable getLastCompiledSymbols() {
+			return lastCompiledSymbols;
+		}
+
 		static {
 			try {
 				fGen = Javac.class.getDeclaredField("gen");

src/main/java/me/coley/recaf/parse/bytecode/MethodAssembler.java

@@ -112,10 +112,11 @@ public MethodNode compile(ParseResult<RootAST> result) throws AssemblerException
 		if (config.verify)
 			lastVerifier = verify(node);
 		if (config.variables) {
- 			// Compute variable information
+			// Compute variable information
 			VariableGenerator variables = new VariableGenerator(variableNames, compilation, node);
 			variables.computeVariables(lastVerifier);
-			node.localVariables = variables.getVariables();
+			if (variables.getVariables() != null)
+				node.localVariables.addAll(variables.getVariables());
 		}
 		// Call complete to notify we are done
 		compilation.onCompletion();

src/main/java/me/coley/recaf/parse/bytecode/MethodCompilation.java

@@ -3,26 +3,22 @@
 import javassist.ClassPool;
 import javassist.CtBehavior;
 import javassist.CtClass;
-import javassist.bytecode.Bytecode;
 import me.coley.recaf.control.Controller;
 import me.coley.recaf.metadata.Comments;
 import me.coley.recaf.parse.bytecode.ast.AST;
 import me.coley.recaf.parse.bytecode.ast.LabelAST;
 import me.coley.recaf.parse.bytecode.ast.MethodDefinitionAST;
 import me.coley.recaf.parse.bytecode.ast.RootAST;
 import me.coley.recaf.parse.bytecode.exception.AssemblerException;
-import org.objectweb.asm.Opcodes;
+import me.coley.recaf.util.TypeUtil;
 import org.objectweb.asm.tree.AbstractInsnNode;
 import org.objectweb.asm.tree.LabelNode;
 import org.objectweb.asm.tree.LocalVariableNode;
 import org.objectweb.asm.tree.MethodNode;
 
 import java.io.ByteArrayInputStream;
 import java.io.InputStream;
-import java.util.Collections;
-import java.util.HashMap;
-import java.util.List;
-import java.util.Map;
+import java.util.*;
 
 /**
  * AST compilation context.
@@ -37,7 +33,7 @@ public final class MethodCompilation {
     private final Controller controller;
 	private final List<LocalVariableNode> priorVars;
 	private final Comments comments = new Comments();
-    private final Map<String, LabelNode> nameToLabel = new HashMap<>();
+    private final Map<String, LabelNode> nameToLabel = new LinkedHashMap<>();
     private final Map<LabelNode, LabelAST> labelToAst = new HashMap<>();
     private final Map<AbstractInsnNode, AST> insnToAst = new HashMap<>();
     private final Map<Integer, AbstractInsnNode> lineToInsn = new HashMap<>();
@@ -57,16 +53,16 @@ public final class MethodCompilation {
 	 * @param priorVars
 	 * 		Prior variable information.
 	 */
-    public MethodCompilation(ParseResult<RootAST> ast, MethodDefinitionAST methodDefinition,
+	public MethodCompilation(ParseResult<RootAST> ast, MethodDefinitionAST methodDefinition,
 							 MethodNode node, String declaringType, Controller controller,
 							 List<LocalVariableNode> priorVars) {
-        this.ast = ast;
-        this.methodDefinition = methodDefinition;
-        this.node = node;
-        this.declaringType = declaringType;
-        this.controller = controller;
-        this.priorVars = priorVars;
-    }
+		this.ast = ast;
+		this.methodDefinition = methodDefinition;
+		this.node = node;
+		this.declaringType = declaringType;
+		this.controller = controller;
+		this.priorVars = priorVars;
+	}
 
     /**
      * @return  root AST parse result.
@@ -161,18 +157,31 @@ public void addComment(String comment) {
 	 * 		When the expression cannot be compiled.
 	 */
 	public void addExpression(String expression, AST ast) throws AssemblerException {
-		// TODO: Since instructions and AST have a one to one mapping we may want to make a map of lines to expressions
-		//       if something goes wrong in the code generated by the expression.
+		// TODO: Since instructions and AST have a one to one mapping we may want to make a map of lines
+		//       to expressions if something goes wrong in the code generated by the expression.
 		try {
 			InputStream stream = new ByteArrayInputStream(controller.getWorkspace().getRawClass(declaringType));
 			CtClass declaring = ClassPool.getDefault().makeClass(stream);
 			CtBehavior containerMethod = declaring.getMethod(node.name, node.desc);
 			// Compile with Javassist
-			Bytecode bytecode = JavassistCompiler.compileExpression(declaring, containerMethod, expression, priorVars);
+			JavassistCompiler.CompilationResult result =
+					JavassistCompiler.compileExpression(declaring, containerMethod, expression,
+							priorVars, variableNames);
 			// Translate to ASM
 			JavassistASMTranslator translator = new JavassistASMTranslator();
-			translator.visit(declaring, bytecode.toCodeAttribute());
+			translator.visit(declaring, result.getBytecode().toCodeAttribute());
 			node.instructions.add(translator.getInstructions());
+			// Add variables
+			Object[] entries = nameToLabel.keySet().toArray();
+			result.getSymbols().forEach((name, dec) -> {
+				String desc = dec.getClassName();
+				if (TypeUtil.isPrimitiveClassName(desc))
+					desc = TypeUtil.classToPrimitive(desc);
+				int index = dec.getLocalVar();
+				LabelNode start = nameToLabel.get(String.valueOf(entries[0]));
+				LabelNode end = nameToLabel.get(String.valueOf(entries[nameToLabel.size() - 1]));
+				node.localVariables.add(new LocalVariableNode(name, desc, null, start, end, index));
+			});
 		} catch (Exception ex) {
 			String msg = ex.getMessage();
 			if (msg == null)

src/main/java/me/coley/recaf/parse/bytecode/VariableNameCache.java

@@ -108,6 +108,51 @@ void visit(RootAST root) throws AssemblerException {
 		}
 	}
 
+	/**
+	 * Adds a variable to the cache.
+	 *
+	 * @param name
+	 * 		Variable name.
+	 * @param type
+	 * 		Variable type.
+	 *
+	 * @return Assigned index of the new variable.
+	 */
+	public int getAndIncrementNext(String name, Type type) {
+		int size = type.getSize();
+		int ret = getNextFreeVar(next, size);
+		// Update used indices
+		usedRawIndices.add(ret);
+		if (size > 1)
+			usedRawIndices.add(ret + 1);
+		nameToIndex.put(name, ret);
+		// Update next and max values
+		next = getNextFreeVar(next, 1);
+		maxIndex = next;
+		return ret;
+	}
+
+	/**
+	 * Finds the next free variable.
+	 *
+	 * @param start
+	 * 		Starting position.
+	 * @param size
+	 * 		Size of variable to check for free space.
+	 *
+	 * @return Next free index.
+	 */
+	public int getNextFreeVar(int start, int size) {
+		int temp = start;
+		if (size == 1)
+			while (isIndexUsed(temp))
+				temp++;
+		else
+			while (isIndexUsed(temp) && isIndexUsed(temp + 1))
+				temp++;
+		return temp;
+	}
+
 	/**
 	 * Finds the increment needed to fit the next variable slot. Will skip already used values.
 	 *