lua4jvm: Implement VARIABLE_TRACING pass and re-enable upvalue typing

This should make local functions efficient enough for numerical code!

While testing said numerical code, it turned out that lua4jvm's support for
using ints for numerical code was, especially mixed with doubles, quite
buggy. The bugs are now fixed.

Author: bensku

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

@@ -101,8 +101,8 @@ public LuaModule compile(String chunk) {
 	public LuaFunction load(LuaModule module, LuaTable env) {
 		// Instantiate the module
 		var type = LuaType.function(
-				// TODO _ENV mutability tracking
-				List.of(new UpvalueTemplate(module.env(), LuaType.TABLE)),
+				// TODO _ENV mutability tracking - we need LuaContext, which is a bit tricky here...
+				List.of(new UpvalueTemplate(module.env(), LuaType.UNKNOWN, true)),
 				List.of(),
 				module.root(),
 				module.name(),

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

@@ -8,13 +8,17 @@ public enum CompilerPass {
 	IR_GEN,
 
 	/**
-	 * In this phase, variables are traced to determine their mutability and
-	 * other properties. TODO not yet implemented
+	 * Return tracking, to determine if lua4jvm has to insert empty returns.
+	 */
+	RETURN_TRACKING,
+	
+	/**
+	 * Variable flagging, based on e.g. their mutability.
 	 */
 	VARIABLE_TRACING,
 
 	/**
-	 * In analysis phase, types that can be statically inferred are inferred
+	 * In analysis pass, types that can be statically inferred are inferred
 	 * to generate better code later.
 	 */
 	TYPE_ANALYSIS,

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

@@ -63,7 +63,6 @@ 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);
 
 			CompilerPass.setCurrent(CompilerPass.CODEGEN);

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

@@ -448,7 +448,8 @@ public IrNode visitStringConcat(StringConcatContext ctx) {
 	@Override
 	public IrNode visitNumberLiteral(NumberLiteralContext ctx) {
 		var value = Double.valueOf(ctx.Numeral().getText());
-		return new LuaConstant(value.intValue() == value ? value.intValue() : value);
+		// Use Math.rint() to handle very large doubles safely
+		return Math.rint(value) == value ? new LuaConstant(value.intValue()) : new LuaConstant(value);
 	}
 
 	@Override

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

@@ -45,8 +45,13 @@ public static LuaContext forFunction(LuaVm vm, LuaType.Function type, boolean tr
 			ctx.setFlag(arg, VariableFlag.ASSIGNED); // JVM assigns arguments to these
 		}
 
+		// Do variable flagging BEFORE type analysis, we need that mutability information
+		type.body().flagVariables(ctx);
+		
 		// Compute types of local variables and the return type
+		CompilerPass.setCurrent(CompilerPass.TYPE_ANALYSIS);
 		type.body().outputType(ctx);
+		CompilerPass.setCurrent(null);
 		return ctx;
 	}
 

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

@@ -17,7 +17,7 @@ public enum VariableFlag {
 	/**
 	 * Variable is mutable; that is, it is assigned to at least twice.
 	 */
-	MUTABLE(CompilerPass.TYPE_ANALYSIS)
+	MUTABLE(CompilerPass.VARIABLE_TRACING)
 	;
 
 	/**

lua4jvm/src/main/java/fi/benjami/code4jvm/lua/ffi/JavaFunction.java

@@ -82,7 +82,9 @@ public Target matchToArgs(LuaType[] argTypes, String intrinsicId) {
 			if (target.intrinsicId != null && !target.intrinsicId.equals(intrinsicId)) {
 				continue; // Intrinsic not allowed by caller
 			}
-			if (checkArgs(target, argTypes) == MatchResult.SUCCESS) {
+			var result = checkArgs(target, argTypes);
+			if (result == MatchResult.SUCCESS || result == MatchResult.INT_DOUBLE_CAST_NEEDED) {
+				// Linker calls MethodHandle#cast(...), which casts ints to doubles if needed
 				return target;
 			}
 		}
@@ -99,7 +101,8 @@ private enum MatchResult {
 		SUCCESS,
 		TOO_FEW_ARGS,
 		ARG_TYPE_MISMATCH,
-		VARARGS_TYPE_MISMATCH
+		VARARGS_TYPE_MISMATCH,
+		INT_DOUBLE_CAST_NEEDED
 	}
 
 	private MatchResult checkArgs(Target target, LuaType[] argTypes) {
@@ -113,12 +116,18 @@ private MatchResult checkArgs(Target target, LuaType[] argTypes) {
 		}
 
 		// Check types of arguments
+		var intDoubleCast = false;
 		for (var i = 0; i < requiredArgs; i++) {
 			var arg = target.arguments.get(i);
 			if (!arg.type.isAssignableFrom(argTypes[i])) {
 				// Allow nil instead of expected type if nullability is allowed
-				if (!arg.nullable ||!argTypes[i].equals(LuaType.NIL)) {
-					return MatchResult.ARG_TYPE_MISMATCH;
+				if (!arg.nullable || !argTypes[i].equals(LuaType.NIL)) {
+					if (argTypes[i].equals(LuaType.INTEGER) && arg.type.equals(LuaType.FLOAT)) {
+						// We'll need to cast ints to doubles using MethodHandle magic
+						intDoubleCast = true;
+					} else {
+						return MatchResult.ARG_TYPE_MISMATCH;
+					}
 				}
 			}
 		}
@@ -133,6 +142,10 @@ private MatchResult checkArgs(Target target, LuaType[] argTypes) {
 			}
 		}
 
+		if (intDoubleCast) {
+			return MatchResult.INT_DOUBLE_CAST_NEEDED;
+		}
+		
 		return MatchResult.SUCCESS;
 	}
 }

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

@@ -31,6 +31,11 @@ public boolean hasReturn() {
 		return node.hasReturn();
 	}
 
+	@Override
+	public void flagVariables(LuaContext ctx) {
+		node.flagVariables(ctx);
+	}
+	
 	@Override
 	public IrNode concreteNode() {
 		return node.concreteNode();

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

@@ -14,6 +14,10 @@ default boolean hasReturn() {
 		return false;
 	}
 
+	default void flagVariables(LuaContext ctx) {
+		// No-op
+	}
+	
 	default IrNode concreteNode() {
 		return this;
 	}

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

@@ -39,4 +39,11 @@ public boolean hasReturn() {
 		return false;
 	}
 
+	@Override
+	public void flagVariables(LuaContext ctx) {
+		for (var node : nodes) {
+			node.flagVariables(ctx);
+		}
+	}
+	
 }