luaj4jvm: Add support for static invokedynamic call sites

Finally! Originally abandoned for its complexity, this was surprisingly
manageable task. Of course, most of the linker has been rewritten since then...

Other related changes:
* Codified the ad-hoc compiler passes (assertions only, for now)
* Added flag system for variables and converted mutability tracking to it
  * LuaLocalVars must NOT be mutated after IR_GEN pass, because the further
    passes can and will run more than once
* Disabled upvalue type tracking, it is unsafe with mutable upvalues
  * Upcoming VARIABLE_TRACKING pass will allow re-enabling it

Author: bensku

lua4jvm/src/main/java/fi/benjami/code4jvm/lua/LuaVm.java

@@ -11,6 +11,7 @@
 import org.antlr.v4.runtime.Recognizer;
 import org.antlr.v4.runtime.misc.ParseCancellationException;
 
+import fi.benjami.code4jvm.lua.compiler.CompilerPass;
 import fi.benjami.code4jvm.lua.compiler.IrCompiler;
 import fi.benjami.code4jvm.lua.compiler.LuaScope;
 import fi.benjami.code4jvm.lua.compiler.LuaSyntaxException;
@@ -87,7 +88,10 @@ public void syntaxError(Recognizer<?, ?> recognizer,
 		// Perform semantic analysis and compile to IR
 		var rootScope = LuaScope.chunkRoot();
 		var visitor = new IrCompiler(name, rootScope);
-		return new LuaModule(name, visitor.visitChunk(tree), (LuaLocalVar) rootScope.resolve("_ENV"));
+		CompilerPass.setCurrent(CompilerPass.IR_GEN);
+		var root = visitor.visitChunk(tree);
+		CompilerPass.setCurrent(null);
+		return new LuaModule(name, root, (LuaLocalVar) rootScope.resolve("_ENV"));
 	}
 
 	public LuaModule compile(String chunk) {
@@ -96,10 +100,9 @@ public LuaModule compile(String chunk) {
 
 	public LuaFunction load(LuaModule module, LuaTable env) {
 		// Instantiate the module
-		module.env().markMutable(); // Initial assignment by VM
 		var type = LuaType.function(
 				// TODO _ENV mutability tracking
-				List.of(new UpvalueTemplate(module.env(), module.env().mutable() ? LuaType.UNKNOWN : LuaType.TABLE, module.env().mutable())),
+				List.of(new UpvalueTemplate(module.env(), LuaType.TABLE)),
 				List.of(),
 				module.root(),
 				module.name(),

lua4jvm/src/main/java/fi/benjami/code4jvm/lua/compiler/CompilerPass.java

@@ -0,0 +1,43 @@
+package fi.benjami.code4jvm.lua.compiler;
+
+public enum CompilerPass {
+
+	/**
+	 * In this phase, lua4jvm IR is generated.
+	 */
+	IR_GEN,
+	
+	/**
+	 * In this phase, variables are traced to determine their mutability and
+	 * other properties. TODO not yet implemented
+	 */
+	VARIABLE_TRACING,
+	
+	/**
+	 * In analysis phase, types that can be statically inferred are inferred
+	 * to generate better code later.
+	 */
+	TYPE_ANALYSIS,
+	
+	/**
+	 * Code generation. In this pass, code4jvm IR is generated. Actual bytecode
+	 * generation is done later by code4jvm, which has its own set of internal
+	 * passes.
+	 */
+	CODEGEN
+	;
+	
+	private static final ThreadLocal<CompilerPass> current = new ThreadLocal<>();
+	
+	public static void setCurrent(CompilerPass pass) {
+		current.set(pass);
+	}
+	
+	public boolean active() {
+		return current.get() == this;
+	}
+	
+	public boolean inactive() {
+		return current.get() != this;
+	}
+}

lua4jvm/src/main/java/fi/benjami/code4jvm/lua/compiler/FunctionCompiler.java

@@ -18,6 +18,7 @@
 import fi.benjami.code4jvm.lua.ir.LuaType;
 import fi.benjami.code4jvm.lua.ir.UpvalueTemplate;
 import fi.benjami.code4jvm.lua.linker.LuaLinker;
+import fi.benjami.code4jvm.lua.runtime.LuaBox;
 import fi.benjami.code4jvm.lua.runtime.LuaFunction;
 import fi.benjami.code4jvm.statement.Return;
 import fi.benjami.code4jvm.typedef.ClassDef;
@@ -35,7 +36,9 @@ public record CacheKey(
 			List<LuaType> argTypes,
 			List<LuaType> upvalueTypes,
 			boolean truncateReturn
-	) {}
+	) {
+		
+	}
 
 	private static final MethodHandles.Lookup LOOKUP = MethodHandles.lookup();
 
@@ -60,8 +63,13 @@ public static MethodHandle callTarget(LuaType[] argTypes, LuaFunction function,
 
 		// Compile and load the function code, or use something that is already cached
 		var compiledFunc = function.type().specializations().computeIfAbsent(cacheKey, t -> {
+			CompilerPass.setCurrent(CompilerPass.TYPE_ANALYSIS);
 			var ctx = LuaContext.forFunction(function.owner(), function.type(), truncateReturn, argTypes);
-			var code = generateCode(ctx, function.type(), argTypes, upvalueTypes);
+			
+			CompilerPass.setCurrent(CompilerPass.CODEGEN);
+			var code = generateCode(ctx, function.type(), argTypes, upvalueTypes, function.upvalues());
+			CompilerPass.setCurrent(null);
+			
 			try {
 				// Load the class with single-use class loader
 				// Using hidden classes would be preferable, but JVM hides them from stack frames
@@ -126,7 +134,7 @@ public static MethodHandle callTarget(LuaType[] argTypes, LuaFunction function,
 	}
 
 	private static byte[] generateCode(LuaContext ctx, LuaType.Function type,
-			LuaType[] argTypes, LuaType[] upvalueTypes) {
+			LuaType[] argTypes, LuaType[] upvalueTypes, Object[] upvalues) {
 		// Create class that wraps the method acting as function body
 		var def = ClassDef.create(toClassName(type.moduleName()), Access.PUBLIC);
 		def.sourceFile(type.moduleName());
@@ -188,9 +196,11 @@ private static byte[] generateCode(LuaContext ctx, LuaType.Function type,
 		// Read upvalues from fields to local variables
 		for (var i = 0; i < upvalueTypes.length; i++) {
 			var template = type.upvalues().get(i);
+			// If upvalue is mutable, its actual type is unknown - and JVM type is LuaBox
+			var actualType = upvalues[i] instanceof LuaBox ? LuaBox.TYPE : upvalueTypes[i].backingType();
 			var value = method.add(template.variable().name(), method.self()
-					.getField(upvalueTypes[i].backingType(), template.variable().name()));
-			ctx.addUpvalue(template.variable(), value);
+					.getField(actualType, template.variable().name()));
+			ctx.addUpvalue(template.variable(), value, upvalues[i]);
 		}
 
 		// Emit Lua code as JVM bytecode

lua4jvm/src/main/java/fi/benjami/code4jvm/lua/compiler/LuaContext.java

@@ -1,6 +1,7 @@
 package fi.benjami.code4jvm.lua.compiler;
 
 import java.util.ArrayList;
+import java.util.EnumSet;
 import java.util.IdentityHashMap;
 import java.util.List;
 import java.util.Map;
@@ -23,6 +24,7 @@ public static LuaContext forFunction(LuaVm vm, LuaType.Function type, boolean tr
 		var ctx = new LuaContext(vm, truncateReturn);
 		for (var upvalue : type.upvalues()) {
 			ctx.recordType(upvalue.variable(), upvalue.type());
+			ctx.setFlag(upvalue.variable(), VariableFlag.ASSIGNED); // Compiler generates code to assign upvalues
 		}
 
 		// Add types of function arguments
@@ -34,11 +36,13 @@ public static LuaContext forFunction(LuaVm vm, LuaType.Function type, boolean tr
 		}
 		var acceptedArgs = type.acceptedArgs();
 		for (var i = 0; i < normalArgs; i++) {
+			var arg = acceptedArgs.get(i);
 			if (argTypes.length > i) {
-				ctx.recordType(acceptedArgs.get(i), argTypes[i]);
+				ctx.recordType(arg, argTypes[i]);
 			} else {
-				ctx.recordType(acceptedArgs.get(i), LuaType.NIL);
+				ctx.recordType(arg, LuaType.NIL);
 			}
+			ctx.setFlag(arg, VariableFlag.ASSIGNED); // JVM assigns arguments to these
 		}
 
 		// Compute types of local variables and the return type
@@ -60,7 +64,14 @@ public static LuaContext forFunction(LuaVm vm, LuaType.Function type, boolean tr
 	/**
 	 * Local variables that are, in fact, upvalues.
 	 */
-	private final Map<LuaLocalVar, Variable> upvalues;
+	private final Map<LuaLocalVar, Object> upvalues;
+	
+	/**
+	 * Flags for local variables. This is used to circumvent the issue the
+	 * fact that local variables cannot be mutated, since the same objects may
+	 * be used for multiple compilation runs.
+	 */
+	private final Map<LuaLocalVar, EnumSet<VariableFlag>> variableFlags;
 
 	/**
 	 * Data given to JVM when the function is loaded as a hidden class.
@@ -91,6 +102,7 @@ public LuaContext(LuaVm owner, boolean truncateReturn) {
 		this.typeTable = new IdentityHashMap<>();
 		this.variables = new IdentityHashMap<>();
 		this.upvalues = new IdentityHashMap<>();
+		this.variableFlags = new IdentityHashMap<>();
 		this.classData = new ArrayList<>();
 		this.cache = new IdentityHashMap<>();
 		this.truncateReturn = truncateReturn;
@@ -124,15 +136,26 @@ public void addFunctionArg(LuaLocalVar arg, Variable variable) {
 		variables.put(arg, variable);
 	}
 
-	public void addUpvalue(LuaLocalVar arg, Variable variable) {
-		variables.put(arg, variable);
-		upvalues.put(arg, variable);
+	/**
+	 * Adds an upvalue to local variables at this context.
+	 * @param localVar Lua local variable that represents the upvalue.
+	 * @param jvmVar JVM variable that represents it in currently compiled method.
+	 * @param value Known value of the upvalue at compilation time. This can be
+	 * used by compiler if it is known to be constant.
+	 */
+	public void addUpvalue(LuaLocalVar localVar, Variable jvmVar, Object value) {
+		variables.put(localVar, jvmVar);
+		upvalues.put(localVar, value);
 	}
 
 	public boolean isUpvalue(LuaLocalVar localVar) {
 		return upvalues.containsKey(localVar);
 	}
 
+	public Object getUpvalue(LuaLocalVar localVar) {
+		return upvalues.get(localVar);
+	}
+	
 	public LuaType variableType(LuaVariable variable) {
 		if (variable instanceof LuaLocalVar) {
 			assert typeTable.containsKey(variable) : variable;
@@ -156,13 +179,28 @@ public Variable resolveLocalVar(LuaLocalVar variable) {
 		var backingVar = variables.get(variable);
 		if (backingVar == null) {
 			var type = typeTable.get(variable);
-			var useBox = variable.upvalue() && variable.mutable();
+			var useBox = variable.upvalue() && hasFlag(variable, VariableFlag.MUTABLE);
 			backingVar = Variable.create(useBox ? LuaBox.TYPE : type.backingType(), variable.name());
 			variables.put(variable, backingVar);
 		}
 		return backingVar;
 	}
 
+	public boolean hasFlag(LuaLocalVar variable, VariableFlag flag) {
+		assert flag.lockedPass == null || flag.lockedPass.inactive();
+		var set = variableFlags.get(variable);
+		return set != null ? set.contains(flag) : false;
+	}
+	
+	public void setFlag(LuaLocalVar variable, VariableFlag flag) {
+		var set = variableFlags.get(variable);
+		if (set == null) {
+			set = EnumSet.of(flag);
+		} else {
+			set.add(flag);
+		}
+	}
+	
 	public void returnTypes(LuaType... types) {
 		if (returnTypes == null) {
 			returnTypes = types;

lua4jvm/src/main/java/fi/benjami/code4jvm/lua/compiler/VariableFlag.java

@@ -0,0 +1,37 @@
+package fi.benjami.code4jvm.lua.compiler;
+
+import fi.benjami.code4jvm.lua.ir.LuaLocalVar;
+
+/**
+ * Flags that can be assigned to {@link LuaLocalVar local variables} in
+ * {@link LuaContext compiler contexts}. Thus, these are unique per compilation
+ * run.
+ */
+public enum VariableFlag {
+	
+	/**
+	 * Variable is assigned to once or more.
+	 */
+	ASSIGNED(null),
+	
+	/**
+	 * Variable is mutable; that is, it is assigned to at least twice.
+	 */
+	MUTABLE(CompilerPass.TYPE_ANALYSIS)
+	;
+	
+	/**
+	 * The flag must not be checked during this pass, but can be set.
+	 * Null to disable this check. When assertions are enabled,
+	 * {@link LuaContext#hasFlag(LuaLocalVar, VariableFlag)} checks this.
+	 * 
+	 * <p>Locked passes are used to ensure that flags are not being mutated
+	 * and read at the same time. If unintentional, such actions could lead to
+	 * subtle but nasty bugs.
+	 */
+	public final CompilerPass lockedPass;
+	
+	VariableFlag(CompilerPass lockedPass) {
+		this.lockedPass = lockedPass;
+	}
+}

lua4jvm/src/main/java/fi/benjami/code4jvm/lua/debug/LinkerTrace.java

@@ -9,4 +9,6 @@ public class LinkerTrace {
 	public Object callable;
 
 	public Object currentPrototype;
+	
+	public int stableTargets;
 }

lua4jvm/src/main/java/fi/benjami/code4jvm/lua/ir/LuaLocalVar.java

@@ -1,6 +1,7 @@
 package fi.benjami.code4jvm.lua.ir;
 
 import fi.benjami.code4jvm.Type;
+import fi.benjami.code4jvm.lua.compiler.CompilerPass;
 
 public final class LuaLocalVar implements LuaVariable {
 
@@ -9,7 +10,6 @@ public final class LuaLocalVar implements LuaVariable {
 	public static final LuaLocalVar VARARGS = new LuaLocalVar("...");
 
 	private final String name;
-	private int mutationSites;
 	private boolean upvalue;
 
 	public LuaLocalVar(String name) {
@@ -20,28 +20,18 @@ public String name() {
 		return name;
 	}
 
-	@Override
-	public void markMutable() {
-		mutationSites++;
-	}
-	
-	/**
-	 * Whether or not this local variable is ever assigned to after its initial
-	 * assignment. This includes mutations by blocks that inherit it as upvalue
-	 * (to be precise, Lua upvalues are essentially external local variables).
-	 */
-	public boolean mutable() {
-		return mutationSites > 1;
-	}
+	// The below mutations are safe, because they're done in IR compilation phase that is run once only
 
 	public void markUpvalue() {
+		assert CompilerPass.IR_GEN.active();
 		upvalue = true;
 	}
 
 	/**
 	 * Whether or not this local variable is an upvalue for some block.
 	 */
 	public boolean upvalue() {
+		assert CompilerPass.IR_GEN.inactive();
 		return upvalue;
 	}
 

lua4jvm/src/main/java/fi/benjami/code4jvm/lua/ir/LuaType.java

@@ -140,28 +140,6 @@ public Type backingType() {
 		public String toString() {
 			return "LuaType.Function[upvalues=" + upvalues + ", acceptedArgs=" + acceptedArgs + ", body=" + body + "]";
 		}
-
-		@Override
-		public int hashCode() {
-			return Objects.hash(acceptedArgs, body, specializations, upvalues);
-		}
-
-		@Override
-		public boolean equals(Object obj) {
-			if (this == obj) {
-				return true;
-			}
-			if (obj == null) {
-				return false;
-			}
-			if (getClass() != obj.getClass()) {
-				return false;
-			}
-			Function other = (Function) obj;
-			return Objects.equals(acceptedArgs, other.acceptedArgs) && Objects.equals(body, other.body)
-					&& Objects.equals(specializations, other.specializations)
-					&& Objects.equals(upvalues, other.upvalues);
-		}
 
 	}
 

lua4jvm/src/main/java/fi/benjami/code4jvm/lua/ir/LuaVariable.java

@@ -2,5 +2,4 @@
 
 public sealed interface LuaVariable permits LuaLocalVar, TableField {
 
-	void markMutable();
 }

lua4jvm/src/main/java/fi/benjami/code4jvm/lua/ir/TableField.java

@@ -5,8 +5,4 @@ public record TableField(
 		IrNode field
 ) implements LuaVariable {
 
-	@Override
-	public void markMutable() {
-		// Do nothing, table fields are always mutable
-	}
 }