lua4jvm: Add table iteration support (WIP)

pairs() is broken, possibly due to linker issues... Other things from
this commit:

* Intrinsic Java methods; used for optimizing pairs() loops
* InternalLib - non-public functions that Lua code can indirectly access

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

@@ -3,6 +3,7 @@
 import java.lang.invoke.MethodHandle;
 import java.util.List;
 
+import fi.benjami.code4jvm.lua.linker.CallSiteOptions;
 import fi.benjami.code4jvm.lua.ir.LuaType;
 
 /**
@@ -62,16 +63,25 @@ public record Target(
 			/**
 			 * Java method that should be called for this target.
 			 */
-			MethodHandle method
+			MethodHandle method,
+
+			/**
+			 * Intrinsic id. If non-null, this target is ignored unless the
+			 * {@link CallSiteOptions#intrinsicId call site} has same id set.
+			 */
+			String intrinsicId
 	) {}
 
 	public record Arg(String name, LuaType type, boolean nullable) {}
 
 	// TODO support functions for generating errors
 
-	public Target matchToArgs(LuaType[] argTypes) {
+	public Target matchToArgs(LuaType[] argTypes, String intrinsicId) {
 		// Try all targets in order
 		for (var target : targets) {
+			if (target.intrinsicId != null && !target.intrinsicId.equals(intrinsicId)) {
+				continue; // Intrinsic not allowed by caller
+			}
 			if (checkArgs(target, argTypes) == MatchResult.SUCCESS) {
 				return target;
 			}

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

@@ -99,7 +99,9 @@ private JavaFunction.Target toFunctionTarget(Method method) {
 			throw new IllegalStateException("lookup provided for LuaBinder has insufficient access", e);
 		}
 
-		return new JavaFunction.Target(injectedArgs, args, method.isVarArgs(), LuaType.of(returnType), multipleReturns, handle);
+		var intrinsic = method.getAnnotation(LuaIntrinsic.class);
+		var intrinsicId = intrinsic != null ? intrinsic.value() : null;
+		return new JavaFunction.Target(injectedArgs, args, method.isVarArgs(), LuaType.of(returnType), multipleReturns, handle, intrinsicId);
 	}
 
 	private InjectedArg toInjectedArg(Class<?> type, String source) {

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

@@ -0,0 +1,13 @@
+package fi.benjami.code4jvm.lua.ffi;
+
+import java.lang.annotation.ElementType;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+import java.lang.annotation.Target;
+
+@Retention(RetentionPolicy.RUNTIME)
+@Target(ElementType.METHOD)
+public @interface LuaIntrinsic {
+
+	String value();
+}

lua4jvm/src/main/java/fi/benjami/code4jvm/lua/ir/expr/FunctionCallExpr.java

@@ -23,6 +23,10 @@ private record CachedCall(LuaType[] argTypes, LuaType returnType) {}
 
 	@Override
 	public Value emit(LuaContext ctx, Block block) {
+		return emit(ctx, block, null);
+	}
+
+	public Value emit(LuaContext ctx, Block block, String intrinsicId) {
 		// Get expensive-to-compute types from cache
 		var cache = (CachedCall) ctx.getCache(this);
 		var argTypes = cache.argTypes();
@@ -31,7 +35,7 @@ public Value emit(LuaContext ctx, Block block) {
 		// TODO constant bootstrap is broken due to upvalues
 		var bootstrap = LuaLinker.BOOTSTRAP_DYNAMIC;
 		var lastMultiVal = !args.isEmpty() && MultiVals.canReturnMultiVal(args.get(args.size() - 1));
-		var options = new CallSiteOptions(ctx.owner(), argTypes, ctx.allowSpread(), lastMultiVal);
+		var options = new CallSiteOptions(ctx.owner(), argTypes, ctx.allowSpread(), lastMultiVal, intrinsicId);
 		bootstrap = bootstrap.withCapturedArgs(ctx.addClassData(options));
 
 		// Evaluate arguments to values (function is first argument)

lua4jvm/src/main/java/fi/benjami/code4jvm/lua/ir/stmt/IteratorForStmt.java

@@ -15,6 +15,7 @@
 import fi.benjami.code4jvm.lua.ir.LuaBlock;
 import fi.benjami.code4jvm.lua.ir.LuaLocalVar;
 import fi.benjami.code4jvm.lua.ir.LuaType;
+import fi.benjami.code4jvm.lua.ir.expr.FunctionCallExpr;
 import fi.benjami.code4jvm.lua.linker.CallSiteOptions;
 import fi.benjami.code4jvm.lua.linker.LuaLinker;
 import fi.benjami.code4jvm.lua.stdlib.LuaException;
@@ -29,7 +30,7 @@ public record IteratorForStmt(
 		List<LuaLocalVar> loopVars,
 		List<IrNode> iterable
 ) implements IrNode {
-
+	
 	@Override
 	public Value emit(LuaContext ctx, Block block) {
 		// TODO code4jvm's LoopBlock is dangerously useless; return here if/when it is fixed
@@ -47,9 +48,15 @@ public Value emit(LuaContext ctx, Block block) {
 			// Before loop body, call the iterable to (hopefully) produce an array of:
 			// iterator function, state, initial value for control variable, (TODO closing value)
 			// TODO Java iterable interop?
-			ctx.setAllowSpread(true);
-			var iterator = iterable.get(0).emit(ctx, init);
-			ctx.setAllowSpread(false);
+			var first = iterable.get(0);
+			Value iterator;
+			if (first instanceof FunctionCallExpr call) {
+				ctx.setAllowSpread(true);
+				iterator = call.emit(ctx, block, "iteratorFor");
+				ctx.setAllowSpread(false);
+			} else {
+				iterator = first.emit(ctx, block);
+			}
 
 			// We might've gotten a multival of next, state, control or only some of those
 			// Set state, control to null as they are technically optional
@@ -68,6 +75,7 @@ public Value emit(LuaContext ctx, Block block) {
 				inner.add(next, ArrayAccess.get(array, Constant.of(0))); // This must exist, but TODO improve error messages
 				inner.add(Jump.to(init, Jump.Target.END, Condition.equal(length, Constant.of(1))));
 				inner.add(ArrayAccess.get(array, Constant.of(1)));
+				inner.add(Jump.to(init, Jump.Target.END, Condition.equal(length, Constant.of(2))));
 				inner.add(control, ArrayAccess.get(array, Constant.of(2)));
 			});
 			innerInit.fallback(inner -> {
@@ -96,9 +104,9 @@ public Value emit(LuaContext ctx, Block block) {
 		}
 		block.add(init);
 
-		// In loop body, call next(state, control)
-		// (types are unknown because we can't yet track them for multivals)
+		// In loop body, call next(state, control)		
 		var bootstrap = LuaLinker.BOOTSTRAP_DYNAMIC;
+		// Types are unknown because we can't yet track them for multivals
 		var options = new CallSiteOptions(ctx.owner(), new LuaType[] {LuaType.UNKNOWN, LuaType.UNKNOWN}, true, false);
 		bootstrap = bootstrap.withCapturedArgs(ctx.addClassData(options));
 		var target = CallTarget.dynamic(bootstrap, Type.OBJECT, "_", Type.OBJECT, Type.OBJECT);

lua4jvm/src/main/java/fi/benjami/code4jvm/lua/linker/CallSiteOptions.java

@@ -29,9 +29,20 @@ public record CallSiteOptions(
 		 * as the its last argument. This forces the linker to inspect the
 		 * arguments and (attempt to) spread it over arguments.
 		 */
-		boolean spreadArguments
+		boolean spreadArguments,
+
+		/**
+		 * Intrinsic id of this call site. If non-null, when the call target is
+		 * a Java function, its targets with same intrinsic id are selected in
+		 * addition to 
+		 */
+		String intrinsicId
 ) {
 
+	public CallSiteOptions(LuaVm owner, LuaType[] types, boolean spreadResults, boolean spreadArguments) {
+		this(owner, types, spreadResults, spreadArguments, null);
+	}
+
 	/**
 	 * Creates call site options for a non-function call.
 	 * @param types Types at call site. Use UNKNOWNs if not known or needed.

lua4jvm/src/main/java/fi/benjami/code4jvm/lua/linker/LuaLinker.java

@@ -35,7 +35,7 @@ public class LuaLinker {
 	private static final MethodHandles.Lookup LOOKUP = MethodHandles.lookup();
 	public static final Type TYPE = Type.of(LuaLinker.class);
 
-	static final MethodHandle TARGET_HAS_CHANGED, PROTOTYPE_HAS_CHANGED, TYPE_HAS_CHANGED,
+	public static final MethodHandle TARGET_HAS_CHANGED, PROTOTYPE_HAS_CHANGED, TYPE_HAS_CHANGED,
 			ARRAY_FIRST, SHAPE_ARRAYS, UPDATE_SITE, NEW_WRAPPER_EX, WRAPPER_EX_CAUSE;
 
 	private static final class WrapperException extends RuntimeException {
@@ -126,6 +126,7 @@ public static LuaCallTarget linkCall(LuaCallSite meta, Object callable, Object..
 			}
 		} else if (callable instanceof JavaFunction function) {
 			// Java method exposed to Lua via FFI
+			var intrinsicId = meta.options.intrinsicId();
 			var specializedTypes = compiledTypes;
 			JavaFunction.Target funcTarget;
 			if (meta.hasUnknownTypes) {
@@ -134,24 +135,24 @@ public static LuaCallTarget linkCall(LuaCallSite meta, Object callable, Object..
 					// Use them! Runtime types are never LESS applicable than compile-time types
 					// so we don't need to check anything else
 					specializedTypes = Arrays.stream(args).map(LuaType::of).toArray(LuaType[]::new);
-					funcTarget = function.matchToArgs(specializedTypes);
+					funcTarget = function.matchToArgs(specializedTypes, intrinsicId);
 					meta.usesRuntimeTypes = true;
 				} else {
 					// Too many type changes; we'd prefer to use compile-time types
-					funcTarget = function.matchToArgs(compiledTypes);
+					funcTarget = function.matchToArgs(compiledTypes, intrinsicId);
 					if (funcTarget == null) {
 						// But it is entirely possible that we can't!
 						// Performance be damned, a call that has correct types at runtime must not fail
 						specializedTypes = Arrays.stream(args).map(LuaType::of).toArray(LuaType[]::new);
-						funcTarget = function.matchToArgs(specializedTypes);
+						funcTarget = function.matchToArgs(specializedTypes, intrinsicId);
 						meta.usesRuntimeTypes = true;
 					} else {
 						meta.usesRuntimeTypes = false;
 					}
 				}
 			} else {
 				// All argument types are known compile-time
-				funcTarget = function.matchToArgs(compiledTypes);
+				funcTarget = function.matchToArgs(compiledTypes, intrinsicId);
 				meta.usesRuntimeTypes = false;
 			}
 

lua4jvm/src/main/java/fi/benjami/code4jvm/lua/runtime/LuaTable.java

@@ -304,4 +304,114 @@ public int arraySize() {
 	public Object shape() {
 		return shape;
 	}
+
+	/**
+	 * Gets the next entry in this table.
+	 * 
+	 * <p>This method supports stateless iteration, much like Lua's next()
+	 * function. When writing Java, this is probably not what you want!
+	 * Stateful iterators produced by {@link #iterator()} are more efficient,
+	 * potentially significantly so.
+	 * @param prevKey Key of the previous entry, or null to start from scratch.
+	 * @return A pair of key and value.
+	 */
+	public Object[] next(Object prevKey) {
+		if (prevKey == null) {
+			if (arraySize != 0) {				
+				// First call, array has at least one member
+				return new Object[] {1d, getArray(1)};
+			} else {
+				// First call, no array members -> return "first" table member
+				for (var i = arrayCapacity; i < table.length; i++) {
+					if (table[i] != null) {
+						return new Object[] {keys[i - arrayCapacity], table[i]};
+					}
+				}
+			}
+		}
+
+		if (prevKey instanceof Double index && index < arraySize + 1) {
+			// Iterate the array in order as long as we have elements
+			var newIndex = index + 1;
+			return new Object[] {newIndex, getArray((int) newIndex)};
+		}
+
+		// Out of array members
+		for (var i = getSlot(prevKey) + 1; i < table.length; i++) {
+			if (table[i] != null) {
+				return new Object[] {keys[i - arrayCapacity], table[i]};
+			}
+		}
+		return null; // Table end
+	}
+
+	/**
+	 * Creates a stateful table iterator of this table.
+	 * 
+	 * <p><b>Note:</b> Table iterators are not compatible with
+	 * {@link java.util.Iterator}!
+	 * @return A table iterator.
+	 */
+	public Iterator iterator() {
+		return new Iterator();
+	}
+
+	/**
+	 * A stateful table iterator.
+	 * 
+	 * <p>Typical usage:
+	 * <pre>
+	 * var it = table.iterator();
+	 * while (it.next()) {
+	 *     var key = it.key();
+	 *     var value = it.value();
+	 * }
+	 * </pre>
+	 *
+	 */
+	public class Iterator {
+
+		private boolean array;
+		private int index;
+
+		private Iterator() {
+			this.array = false;
+			this.index = 0;
+		}
+
+		public boolean next() {
+			return array ? nextArray() : nextTable();
+		}
+
+		private boolean nextArray() {
+			index++;
+			if (index >= arraySize) {
+				// Reached array end
+				array = false;
+				index = arrayCapacity - 1; // Jump over possible unused array space
+				return nextTable(); // Table might or might not have entries
+			}
+			return true;
+		}
+
+		private boolean nextTable() {
+			// Iterate over empty space until we find next entry
+			for (var i = index + 1; i < table.length; i++) {
+				if (table[i] != null) {
+					index = i;
+					return true;
+				}
+			}
+			index = table.length;
+			return false;
+		}
+
+		public Object key() {
+			return array ? (double) index : keys[index];
+		}
+
+		public Object value() {
+			return table[index];
+		}
+	}
 }

lua4jvm/src/main/java/fi/benjami/code4jvm/lua/stdlib/BasicLib.java

@@ -12,10 +12,14 @@
 import fi.benjami.code4jvm.lua.ffi.LuaLibrary;
 import fi.benjami.code4jvm.lua.ffi.Nullable;
 import fi.benjami.code4jvm.lua.ir.LuaType;
+import fi.benjami.code4jvm.lua.linker.CallSiteOptions;
+import fi.benjami.code4jvm.lua.linker.LuaCallSite;
+import fi.benjami.code4jvm.lua.linker.LuaLinker;
 import fi.benjami.code4jvm.lua.ffi.Inject;
 import fi.benjami.code4jvm.lua.ffi.JavaFunction;
 import fi.benjami.code4jvm.lua.ffi.LuaBinder;
 import fi.benjami.code4jvm.lua.ffi.LuaExport;
+import fi.benjami.code4jvm.lua.ffi.LuaIntrinsic;
 import fi.benjami.code4jvm.lua.runtime.LuaFunction;
 import fi.benjami.code4jvm.lua.runtime.LuaTable;
 
@@ -189,4 +193,39 @@ public static String type(Object value) {
 		return LuaType.of(value).name();
 	}
 
+	private static final JavaFunction INTRINSIC_ITERATOR = InternalLib.FUNCTIONS.get("intrinsicIterator"),
+			TABLE_ITERATOR = InternalLib.FUNCTIONS.get("tableIterator");
+
+	@LuaExport("pairs")
+	@LuaIntrinsic("iteratorFor")
+	private static Object[] pairsStateful(@Inject LuaVm vm, Object iterable) throws Throwable {
+		if (iterable instanceof LuaTable table
+				&& (table.metatable() == null || table.metatable().get("__pairs") == null)) {
+			// Normal Lua table; let's cheat a bit and use a stateful table iterator (intrinsic path)
+			return new Object[] {INTRINSIC_ITERATOR, table.iterator()};
+		}
+
+		// Aside of the above fast path, delegate to normal pairs
+		return pairs(vm, iterable);
+	}
+
+	@LuaExport("pairs")
+	private static Object[] pairs(@Inject LuaVm vm, Object iterable) throws Throwable {
+		if (iterable instanceof LuaTable table) {
+			if (table.metatable() != null) {
+				var metamethod = table.metatable().get("__pairs");
+				if (metamethod != null) {
+					// Call __pairs and use whatever it returns as an iterator
+					var target = LuaLinker.linkCall(new LuaCallSite(null, CallSiteOptions.nonFunction(vm, LuaType.TABLE)),
+							metamethod, table);
+					return (Object[]) target.target().invoke(metamethod, table);
+				}
+			}
+			// No __pairs, just iterate over the table normally (non-intrinsic path)
+			return new Object[] {TABLE_ITERATOR, table};
+		} else {
+			throw new LuaException("value not iterable");
+		}
+	}
+
 }