Skip to content

Commit 805e0e9

Browse files
benskubensku
committed
lua4jvm: Fix Lua upvalues
Previously, upvalues were captured at function creation time. From language design point of view, it is a sensible choice. It also goes directly against the Lua spec. Oops... The fix introduces mutability tracking for ALL local variables and boxes mutable upvalues so that they work according to the spec. Upvalues that do not change are unchanged, which is nice for performance reasons. Next up: Try to get constant invokedynamic call sites for upvalue functions that never change after their creation.
1 parent 3e5933e commit 805e0e9

16 files changed

+200
-42
lines changed

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

+3-1
Original file line numberDiff line numberDiff line change
@@ -96,8 +96,10 @@ public LuaModule compile(String chunk) {
9696

9797
public LuaFunction load(LuaModule module, LuaTable env) {
9898
// Instantiate the module
99+
module.env().markMutable(); // Initial assignment by VM
99100
var type = LuaType.function(
100-
List.of(new UpvalueTemplate(module.env(), LuaType.TABLE)),
101+
// TODO _ENV mutability tracking
102+
List.of(new UpvalueTemplate(module.env(), module.env().mutable() ? LuaType.UNKNOWN : LuaType.TABLE, module.env().mutable())),
101103
List.of(),
102104
module.root(),
103105
module.name(),

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

+1-1
Original file line numberDiff line numberDiff line change
@@ -190,7 +190,7 @@ private static byte[] generateCode(LuaContext ctx, LuaType.Function type,
190190
var template = type.upvalues().get(i);
191191
var value = method.add(template.variable().name(), method.self()
192192
.getField(upvalueTypes[i].backingType(), template.variable().name()));
193-
ctx.addFunctionArg(template.variable(), value);
193+
ctx.addUpvalue(template.variable(), value);
194194
}
195195

196196
// Emit Lua code as JVM bytecode

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

+22-2
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,6 @@
11
package fi.benjami.code4jvm.lua.compiler;
22

33
import java.util.ArrayList;
4-
import java.util.Arrays;
54
import java.util.IdentityHashMap;
65
import java.util.List;
76
import java.util.Map;
@@ -15,6 +14,7 @@
1514
import fi.benjami.code4jvm.lua.ir.LuaVariable;
1615
import fi.benjami.code4jvm.lua.ir.TableField;
1716
import fi.benjami.code4jvm.lua.ir.expr.LuaConstant;
17+
import fi.benjami.code4jvm.lua.runtime.LuaBox;
1818

1919
public class LuaContext {
2020

@@ -57,6 +57,11 @@ public static LuaContext forFunction(LuaVm vm, LuaType.Function type, boolean tr
5757
*/
5858
private final Map<LuaLocalVar, Variable> variables;
5959

60+
/**
61+
* Local variables that are, in fact, upvalues.
62+
*/
63+
private final Map<LuaLocalVar, Variable> upvalues;
64+
6065
/**
6166
* Data given to JVM when the function is loaded as a hidden class.
6267
* This is used for creating constants of arbitrary kind, which can then
@@ -85,6 +90,7 @@ public LuaContext(LuaVm owner, boolean truncateReturn) {
8590
assert owner != null;
8691
this.typeTable = new IdentityHashMap<>();
8792
this.variables = new IdentityHashMap<>();
93+
this.upvalues = new IdentityHashMap<>();
8894
this.classData = new ArrayList<>();
8995
this.cache = new IdentityHashMap<>();
9096
this.truncateReturn = truncateReturn;
@@ -118,6 +124,15 @@ public void addFunctionArg(LuaLocalVar arg, Variable variable) {
118124
variables.put(arg, variable);
119125
}
120126

127+
public void addUpvalue(LuaLocalVar arg, Variable variable) {
128+
variables.put(arg, variable);
129+
upvalues.put(arg, variable);
130+
}
131+
132+
public boolean isUpvalue(LuaLocalVar localVar) {
133+
return upvalues.containsKey(localVar);
134+
}
135+
121136
public LuaType variableType(LuaVariable variable) {
122137
if (variable instanceof LuaLocalVar) {
123138
assert typeTable.containsKey(variable) : variable;
@@ -133,11 +148,16 @@ public LuaType variableType(LuaVariable variable) {
133148
}
134149
}
135150

151+
public boolean hasBeenAssigned(LuaLocalVar variable) {
152+
return variables.containsKey(variable);
153+
}
154+
136155
public Variable resolveLocalVar(LuaLocalVar variable) {
137156
var backingVar = variables.get(variable);
138157
if (backingVar == null) {
139158
var type = typeTable.get(variable);
140-
backingVar = Variable.create(type.backingType(), variable.name());
159+
var useBox = variable.upvalue() && variable.mutable();
160+
backingVar = Variable.create(useBox ? LuaBox.TYPE : type.backingType(), variable.name());
141161
variables.put(variable, backingVar);
142162
}
143163
return backingVar;

‎lua4jvm/src/main/java/fi/benjami/code4jvm/lua/compiler/LuaScope.java

+7-13
Original file line numberDiff line numberDiff line change
@@ -11,23 +11,21 @@
1111
import fi.benjami.code4jvm.lua.ir.TableField;
1212
import fi.benjami.code4jvm.lua.ir.expr.LuaConstant;
1313
import fi.benjami.code4jvm.lua.ir.expr.VariableExpr;
14-
import fi.benjami.code4jvm.lua.ir.expr.FunctionDeclExpr.Upvalue;
15-
import fi.benjami.code4jvm.lua.ir.stmt.LoopStmt;
1614

1715
public class LuaScope {
1816

1917
public static LuaScope chunkRoot() {
2018
var scope = new LuaScope(null, true);
2119
var env = scope.declare("_ENV");
22-
scope.upvalues.put("_ENV", new Upvalue(env, null));
20+
scope.upvalues.put("_ENV", env);
2321
return scope;
2422
}
2523

2624
private final LuaScope parent;
2725
private final boolean functionRoot;
2826

2927
private final Map<String, LuaLocalVar> locals;
30-
private final Map<String, Upvalue> upvalues;
28+
private final Map<String, LuaLocalVar> upvalues;
3129

3230
/**
3331
* Reference to current loop or null, used for break'ing out of loop.
@@ -66,15 +64,11 @@ public LuaVariable resolve(String name) {
6664
if (result != null) {
6765
// Local variable or upvalue
6866
if (result.isUpvalue()) {
69-
// Upvalue: record it and create a local variable
70-
var inside = new LuaLocalVar(name);
71-
locals.put(name, inside);
72-
var outside = result.variable();
73-
upvalues.put(name, new Upvalue(inside, outside));
74-
return inside;
75-
} else {
76-
return result.variable(); // Local variable
67+
locals.put(name, result.variable());
68+
upvalues.put(name, result.variable());
69+
result.variable().markUpvalue();
7770
}
71+
return result.variable(); // Local variable
7872
} else {
7973
// Neither local variable or upvalue; take a look at _ENV table
8074
var env = resolve("_ENV");
@@ -99,7 +93,7 @@ private ResolveResult resolveLocal(String name) {
9993
}
10094
}
10195

102-
public List<Upvalue> upvalues() {
96+
public List<LuaLocalVar> upvalues() {
10397
return new ArrayList<>(upvalues.values());
10498
}
10599

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

+38-3
Original file line numberDiff line numberDiff line change
@@ -2,12 +2,47 @@
22

33
import fi.benjami.code4jvm.Type;
44

5-
public record LuaLocalVar(
6-
String name
7-
) implements LuaVariable {
5+
public final class LuaLocalVar implements LuaVariable {
86

97
public static final Type TYPE = Type.of(LuaLocalVar.class);
108

119
public static final LuaLocalVar VARARGS = new LuaLocalVar("...");
10+
11+
private final String name;
12+
private int mutationSites;
13+
private boolean upvalue;
14+
15+
public LuaLocalVar(String name) {
16+
this.name = name;
17+
}
18+
19+
public String name() {
20+
return name;
21+
}
22+
23+
@Override
24+
public void markMutable() {
25+
mutationSites++;
26+
}
27+
28+
/**
29+
* Whether or not this local variable is ever assigned to after its initial
30+
* assignment. This includes mutations by blocks that inherit it as upvalue
31+
* (to be precise, Lua upvalues are essentially external local variables).
32+
*/
33+
public boolean mutable() {
34+
return mutationSites > 1;
35+
}
36+
37+
public void markUpvalue() {
38+
upvalue = true;
39+
}
40+
41+
/**
42+
* Whether or not this local variable is an upvalue for some block.
43+
*/
44+
public boolean upvalue() {
45+
return upvalue;
46+
}
1247

1348
}

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

+1
Original file line numberDiff line numberDiff line change
@@ -2,4 +2,5 @@
22

33
public sealed interface LuaVariable permits LuaLocalVar, TableField {
44

5+
void markMutable();
56
}

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

+7-1
Original file line numberDiff line numberDiff line change
@@ -3,4 +3,10 @@
33
public record TableField(
44
IrNode table,
55
IrNode field
6-
) implements LuaVariable {}
6+
) implements LuaVariable {
7+
8+
@Override
9+
public void markMutable() {
10+
// Do nothing, table fields are always mutable
11+
}
12+
}

‎lua4jvm/src/main/java/fi/benjami/code4jvm/lua/ir/UpvalueTemplate.java

+7-1
Original file line numberDiff line numberDiff line change
@@ -15,5 +15,11 @@ public record UpvalueTemplate(
1515
* {@link LuaFunction#upvalueTypes final types} that are known after
1616
* the function has been instantiated, this may be unknown.
1717
*/
18-
LuaType type
18+
LuaType type,
19+
20+
/**
21+
* Whether or not the upvalue variable is assigned to after its initial
22+
* assignment.
23+
*/
24+
boolean mutable
1925
) {}

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

+12
Original file line numberDiff line numberDiff line change
@@ -7,8 +7,10 @@
77
import fi.benjami.code4jvm.Value;
88
import fi.benjami.code4jvm.block.Block;
99
import fi.benjami.code4jvm.call.CallTarget;
10+
import fi.benjami.code4jvm.call.FixedCallTarget;
1011
import fi.benjami.code4jvm.lua.compiler.LuaContext;
1112
import fi.benjami.code4jvm.lua.ir.IrNode;
13+
import fi.benjami.code4jvm.lua.ir.LuaLocalVar;
1214
import fi.benjami.code4jvm.lua.ir.LuaType;
1315
import fi.benjami.code4jvm.lua.linker.CallSiteOptions;
1416
import fi.benjami.code4jvm.lua.linker.LuaLinker;
@@ -33,6 +35,16 @@ public Value emit(LuaContext ctx, Block block, String intrinsicId) {
3335
var returnType = cache.returnType();
3436

3537
// TODO constant bootstrap is broken due to upvalues
38+
// FixedCallTarget bootstrap;
39+
// if (function instanceof VariableExpr variable // function is a variable read
40+
// && variable.source() instanceof LuaLocalVar localVar // from local variable
41+
// && localVar.upvalue() && !localVar.mutable() // that will be stable between calls to this function
42+
// && TODO we also need to check that 1) linker has used LuaType.TARGET_HAS_CHANGED (to prove upvalue's block hasn't been re-executed)
43+
// && TODO 2) cache key includes identity of the upvalue, not just its type! (this is quite tricky)
44+
// ) {
45+
//
46+
// }
47+
//
3648
var bootstrap = LuaLinker.BOOTSTRAP_DYNAMIC;
3749
var lastMultiVal = !args.isEmpty() && MultiVals.canReturnMultiVal(args.get(args.size() - 1));
3850
var options = new CallSiteOptions(ctx.owner(), argTypes, ctx.allowSpread(), lastMultiVal, intrinsicId);

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

+3-8
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,7 @@ public record FunctionDeclExpr(
2626
*/
2727
String name,
2828

29-
List<Upvalue> upvalues,
29+
List<LuaLocalVar> upvalues,
3030

3131
/**
3232
* Arguments inside the new function.
@@ -38,11 +38,6 @@ public record FunctionDeclExpr(
3838
*/
3939
LuaBlock body
4040
) implements IrNode {
41-
42-
public record Upvalue(
43-
LuaLocalVar inside,
44-
LuaLocalVar outside
45-
) {}
4641

4742
@Override
4843
public Value emit(LuaContext ctx, Block block) {
@@ -52,7 +47,7 @@ public Value emit(LuaContext ctx, Block block) {
5247
// Copy local variables to upvalues array
5348
var upvalueValues = block.add(Type.OBJECT.array(1).newInstance(Constant.of(upvalues.size())));
5449
for (var i = 0; i < upvalues.size(); i++) {
55-
var value = ctx.resolveLocalVar(upvalues.get(i).outside());
50+
var value = ctx.resolveLocalVar(upvalues.get(i));
5651
block.add(ArrayAccess.set(upvalueValues, Constant.of(i), value.cast(Type.OBJECT)));
5752
}
5853

@@ -64,7 +59,7 @@ public Value emit(LuaContext ctx, Block block) {
6459
public LuaType.Function outputType(LuaContext ctx) {
6560
// Upvalue template has the variable INSIDE declared function, with type of OUTSIDE variable
6661
var upvalueTemplates = upvalues.stream()
67-
.map(upvalue -> new UpvalueTemplate(upvalue.inside(), ctx.variableType(upvalue.outside())))
62+
.map(upvalue -> new UpvalueTemplate(upvalue, upvalue.mutable() ? LuaType.UNKNOWN : ctx.variableType(upvalue), upvalue.mutable()))
6863
.toList();
6964
return ctx.cached(this, LuaType.function(upvalueTemplates, arguments, body, moduleName, name));
7065
}

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

+9-1
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@
1212
import fi.benjami.code4jvm.lua.ir.TableField;
1313
import fi.benjami.code4jvm.lua.linker.CallSiteOptions;
1414
import fi.benjami.code4jvm.lua.linker.LuaLinker;
15+
import fi.benjami.code4jvm.lua.runtime.LuaBox;
1516
import fi.benjami.code4jvm.lua.runtime.TableAccess;
1617

1718
/**
@@ -25,7 +26,14 @@ public record VariableExpr(
2526
@Override
2627
public Value emit(LuaContext ctx, Block block) {
2728
if (source instanceof LuaLocalVar localVar) {
28-
return ctx.resolveLocalVar(localVar);
29+
if (localVar.upvalue() && localVar.mutable()) {
30+
// Mutable upvalues have to use LuaBoxes
31+
// TODO cast should not be needed - potential code4jvm bug
32+
return block.add(ctx.resolveLocalVar(localVar).cast(LuaBox.TYPE).getField(outputType(ctx).backingType(), "value"));
33+
} else {
34+
// Normal JVM local variable
35+
return ctx.resolveLocalVar(localVar);
36+
}
2937
} else if (source instanceof TableField tableField) {
3038
var table = tableField.table().emit(ctx, block);
3139
var field = tableField.field().emit(ctx, block);

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

+23-5
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616
import fi.benjami.code4jvm.lua.ir.TableField;
1717
import fi.benjami.code4jvm.lua.ir.expr.FunctionCallExpr;
1818
import fi.benjami.code4jvm.lua.ir.expr.VariableExpr;
19+
import fi.benjami.code4jvm.lua.runtime.LuaBox;
1920
import fi.benjami.code4jvm.lua.runtime.LuaTable;
2021
import fi.benjami.code4jvm.lua.runtime.MultiVals;
2122

@@ -97,8 +98,19 @@ public Value emit(LuaContext ctx, Block block) {
9798
private Statement setVariable(LuaContext ctx, LuaVariable variable, Value value) {
9899
return block -> {
99100
if (variable instanceof LuaLocalVar localVar) {
100-
var jvmVar = ctx.resolveLocalVar(localVar);
101-
block.add(jvmVar.set(value.cast(jvmVar.type())));
101+
if (localVar.upvalue() && localVar.mutable()) {
102+
// Mutable upvalues need to be put to LuaBoxes
103+
if (!ctx.hasBeenAssigned(localVar)) {
104+
// First assignment? Initialize box!
105+
var box = block.add(LuaBox.TYPE.newInstance());
106+
block.add(ctx.resolveLocalVar(localVar).set(box));
107+
}
108+
block.add(ctx.resolveLocalVar(localVar).putField("value", value.cast(Type.OBJECT)));
109+
} else {
110+
// Normal local variable assignment
111+
var jvmVar = ctx.resolveLocalVar(localVar);
112+
block.add(jvmVar.set(value.cast(jvmVar.type())));
113+
}
102114
} else if (variable instanceof TableField tableField) {
103115
// Just call the setter
104116
// TODO invokedynamic to TableAccess.CONSTANT_SET once it has some optimizations
@@ -115,7 +127,9 @@ private Statement setVariable(LuaContext ctx, LuaVariable variable, Value value)
115127
public LuaType outputType(LuaContext ctx) {
116128
var normalSources = spread ? sources.size() - 1 : sources.size();
117129
for (var i = 0; i < Math.min(normalSources, targets.size()); i++) {
118-
ctx.recordType(targets.get(i), sources.get(i).outputType(ctx));
130+
var target = targets.get(i);
131+
ctx.recordType(target, sources.get(i).outputType(ctx));
132+
target.markMutable();
119133
}
120134

121135
if (spread) {
@@ -128,12 +142,16 @@ public LuaType outputType(LuaContext ctx) {
128142
// Tuple -> types for individual variables
129143
// UNKNOWN -> current behavior
130144
// anything else -> first multiValType, rest NIL
131-
ctx.recordType(targets.get(i), LuaType.UNKNOWN);
145+
var target = targets.get(i);
146+
ctx.recordType(target, LuaType.UNKNOWN);
147+
target.markMutable();
132148
}
133149
} else {
134150
// If there are leftover targets, set them to nil
135151
for (var i = normalSources; i < targets.size(); i++) {
136-
ctx.recordType(targets.get(i), LuaType.NIL);
152+
var target = targets.get(i);
153+
ctx.recordType(target, LuaType.NIL);
154+
target.markMutable();
137155
}
138156
}
139157
return LuaType.NIL;

0 commit comments

Comments
 (0)