Skip to content

Commit 57656d3

Browse files
levno-710levno-710
committed
Add Vararg support and fix bug
1 parent 55537b5 commit 57656d3

File tree

1 file changed

+43
-7
lines changed

1 file changed

+43
-7
lines changed

‎src/prometheus/compiler/compiler.lua‎

Lines changed: 43 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -120,7 +120,8 @@ function Compiler:compile(ast)
120120
local _, newproxyVar = newGlobalScope:resolve("newproxy");
121121
local _, setmetatableVar = newGlobalScope:resolve("setmetatable");
122122
local _, getmetatableVar = newGlobalScope:resolve("getmetatable");
123-
123+
local _, selectVar = newGlobalScope:resolve("select");
124+
124125
psc:addReferenceToHigherScope(newGlobalScope, getfenvVar, 2);
125126
psc:addReferenceToHigherScope(newGlobalScope, tableVar);
126127
psc:addReferenceToHigherScope(newGlobalScope, unpackVar);
@@ -136,6 +137,7 @@ function Compiler:compile(ast)
136137
self.newproxyVar = self.scope:addVariable();
137138
self.setmetatableVar = self.scope:addVariable();
138139
self.getmetatableVar = self.scope:addVariable();
140+
self.selectVar = self.scope:addVariable();
139141

140142
self.containerFuncScope = Scope:new(self.scope);
141143
self.whileScope = Scope:new(self.containerFuncScope);
@@ -274,6 +276,7 @@ function Compiler:compile(ast)
274276
Ast.VariableExpression(self.scope, self.newproxyVar),
275277
Ast.VariableExpression(self.scope, self.setmetatableVar),
276278
Ast.VariableExpression(self.scope, self.getmetatableVar),
279+
Ast.VariableExpression(self.scope, self.selectVar),
277280
unpack(util.shuffle({
278281
Ast.VariableExpression(self.scope, self.containerFuncVar),
279282
Ast.VariableExpression(self.scope, self.createClosureVar),
@@ -302,6 +305,7 @@ function Compiler:compile(ast)
302305
Ast.VariableExpression(newGlobalScope, newproxyVar);
303306
Ast.VariableExpression(newGlobalScope, setmetatableVar);
304307
Ast.VariableExpression(newGlobalScope, getmetatableVar);
308+
Ast.VariableExpression(newGlobalScope, selectVar);
305309
})};
306310
}, psc), newGlobalScope);
307311
end
@@ -924,6 +928,9 @@ function Compiler:compileFunction(node, funcDepth)
924928
funcDepth = funcDepth + 1;
925929
local oldActiveBlock = self.activeBlock;
926930

931+
local upperVarargReg = self.varargReg;
932+
self.varargReg = nil;
933+
927934
local upvalueExpressions = {};
928935
local upvalueIds = {};
929936
local usedRegs = {};
@@ -972,6 +979,19 @@ function Compiler:compileFunction(node, funcDepth)
972979
scope:addReferenceToHigherScope(self.containerFuncScope, self.argsVar);
973980
self:addStatement(self:setRegister(scope, argReg, Ast.IndexExpression(Ast.VariableExpression(self.containerFuncScope, self.argsVar), Ast.NumberExpression(i))), {argReg}, {}, false);
974981
end
982+
else
983+
self.varargReg = self:allocRegister(true);
984+
scope:addReferenceToHigherScope(self.containerFuncScope, self.argsVar);
985+
scope:addReferenceToHigherScope(self.scope, self.selectVar);
986+
scope:addReferenceToHigherScope(self.scope, self.unpackVar);
987+
self:addStatement(self:setRegister(scope, self.varargReg, Ast.TableConstructorExpression({
988+
Ast.TableEntry(Ast.FunctionCallExpression(Ast.VariableExpression(self.scope, self.selectVar), {
989+
Ast.NumberExpression(i);
990+
Ast.FunctionCallExpression(Ast.VariableExpression(self.scope, self.unpackVar), {
991+
Ast.VariableExpression(self.containerFuncScope, self.argsVar),
992+
});
993+
})),
994+
})), {self.varargReg}, {}, false);
975995
end
976996
end
977997

@@ -982,6 +1002,10 @@ function Compiler:compileFunction(node, funcDepth)
9821002
self.activeBlock.advanceToNextBlock = false;
9831003
end
9841004

1005+
if(self.varargReg) then
1006+
self:freeRegister(self.varargReg, true);
1007+
end
1008+
self.varargReg = upperVarargReg;
9851009
self.getUpvalueId = oldGetUpvalueId;
9861010

9871011
self:popRegisterUsageInfo();
@@ -1029,7 +1053,7 @@ function Compiler:compileStatement(statement, funcDepth)
10291053
local regs = {};
10301054

10311055
for i, expr in ipairs(statement.args) do
1032-
if i == #statement.args and (expr.kind == AstKind.FunctionCallExpression or expr.kind == AstKind.PassSelfFunctionCallExpression) then
1056+
if i == #statement.args and (expr.kind == AstKind.FunctionCallExpression or expr.kind == AstKind.PassSelfFunctionCallExpression or expr.kind == AstKind.VarargExpression) then
10331057
local reg = self:compileExpression(expr, funcDepth, self.RETURN_ALL)[1];
10341058
table.insert(entries, Ast.TableEntry(Ast.FunctionCallExpression(
10351059
self:unpack(scope),
@@ -1107,7 +1131,7 @@ function Compiler:compileStatement(statement, funcDepth)
11071131
local args = {};
11081132

11091133
for i, expr in ipairs(statement.args) do
1110-
if i == #statement.args and (expr.kind == AstKind.FunctionCallExpression or expr.kind == AstKind.PassSelfFunctionCallExpression) then
1134+
if i == #statement.args and (expr.kind == AstKind.FunctionCallExpression or expr.kind == AstKind.PassSelfFunctionCallExpression or expr.kind == AstKind.VarargExpression) then
11111135
local reg = self:compileExpression(expr, funcDepth, self.RETURN_ALL)[1];
11121136
table.insert(args, Ast.FunctionCallExpression(
11131137
self:unpack(scope),
@@ -1138,7 +1162,7 @@ function Compiler:compileStatement(statement, funcDepth)
11381162
local regs = { baseReg };
11391163

11401164
for i, expr in ipairs(statement.args) do
1141-
if i == #statement.args and (expr.kind == AstKind.FunctionCallExpression or expr.kind == AstKind.PassSelfFunctionCallExpression) then
1165+
if i == #statement.args and (expr.kind == AstKind.FunctionCallExpression or expr.kind == AstKind.PassSelfFunctionCallExpression or expr.kind == AstKind.VarargExpression) then
11421166
local reg = self:compileExpression(expr, funcDepth, self.RETURN_ALL)[1];
11431167
table.insert(args, Ast.FunctionCallExpression(
11441168
self:unpack(scope),
@@ -1803,7 +1827,7 @@ function Compiler:compileExpression(expression, funcDepth, numReturns)
18031827
local regs = { baseReg };
18041828

18051829
for i, expr in ipairs(expression.args) do
1806-
if i == #expression.args and (expr.kind == AstKind.FunctionCallExpression or expr.kind == AstKind.PassSelfFunctionCallExpression) then
1830+
if i == #expression.args and (expr.kind == AstKind.FunctionCallExpression or expr.kind == AstKind.PassSelfFunctionCallExpression or expr.kind == AstKind.VarargExpression) then
18071831
local reg = self:compileExpression(expr, funcDepth, self.RETURN_ALL)[1];
18081832
table.insert(args, Ast.FunctionCallExpression(
18091833
self:unpack(scope),
@@ -1816,7 +1840,7 @@ function Compiler:compileExpression(expression, funcDepth, numReturns)
18161840
end
18171841
end
18181842

1819-
if(numReturns > 1 or returnAll) then
1843+
if(returnAll or numReturns > 1) then
18201844
local tmpReg = self:allocRegister(false);
18211845

18221846
self:addStatement(self:setRegister(scope, tmpReg, Ast.StringExpression(expression.passSelfFunctionName)), {tmpReg}, {}, false);
@@ -2064,7 +2088,7 @@ function Compiler:compileExpression(expression, funcDepth, numReturns)
20642088
local entryRegs = {};
20652089
for i, entry in ipairs(expression.entries) do
20662090
if(entry.kind == AstKind.TableEntry) then
2067-
if i == #expression.entries and entry.kind == AstKind.FunctionCallExpression then
2091+
if i == #expression.entries and entry.kind == AstKind.FunctionCallExpression or expression.kind == AstKind.PassSelfFunctionCallExpression or expression.kind == AstKind.VarargExpression then
20682092
local reg = self:compileExpression(entry.value, funcDepth, self.RETURN_ALL)[1];
20692093
table.insert(entries, Ast.TableEntry(Ast.FunctionCallExpression(
20702094
self:unpack(scope),
@@ -2107,6 +2131,18 @@ function Compiler:compileExpression(expression, funcDepth, numReturns)
21072131
return regs;
21082132
end
21092133

2134+
if(expression.kind == AstKind.VarargExpression) then
2135+
if numReturns == self.RETURN_ALL then
2136+
return {self.varargReg};
2137+
end
2138+
local regs = {};
2139+
for i=1, numReturns do
2140+
regs[i] = self:allocRegister(false);
2141+
self:addStatement(self:setRegister(scope, regs[i], Ast.IndexExpression(self:register(scope, self.varargReg), Ast.NumberExpression(i))), {regs[i]}, {self.varargReg}, false);
2142+
end
2143+
return regs;
2144+
end
2145+
21102146
logger:error(string.format("%s is not an compileable expression!", expression.kind));
21112147
end
21122148

0 commit comments

Comments
 (0)