Add support for compund statements

Author: levno-710

src/prometheus/ast.lua

@@ -26,6 +26,15 @@ local AstKind = {
 	PassSelfFunctionCallStatement = "PassSelfFunctionCallStatement";
 	AssignmentStatement = "AssignmentStatement";
 
+	-- LuaU Compound Statements
+	CompoundAddStatement = "CompoundAddStatement";
+	CompoundSubStatement = "CompoundSubStatement";
+	CompoundMulStatement = "CompoundMulStatement";
+	CompoundDivStatement = "CompoundDivStatement";
+	CompoundModStatement = "CompoundModStatement";
+	CompoundPowStatement = "CompoundPowStatement";
+	CompoundConcatStatement = "CompoundConcatStatement";
+
 	-- Assignment Index
 	AssignmentIndexing = "AssignmentIndexing";
 	AssignmentVariable = "AssignmentVariable";  
@@ -217,6 +226,62 @@ function Ast.AssignmentStatement(lhs, rhs)
 	}
 end
 
+function Ast.CompoundAddStatement(lhs, rhs)
+	return {
+		kind = AstKind.CompoundAddStatement,
+		lhs = lhs,
+		rhs = rhs,
+	}
+end
+
+function Ast.CompoundSubStatement(lhs, rhs)
+	return {
+		kind = AstKind.CompoundSubStatement,
+		lhs = lhs,
+		rhs = rhs,
+	}
+end
+
+function Ast.CompoundMulStatement(lhs, rhs)
+	return {
+		kind = AstKind.CompoundMulStatement,
+		lhs = lhs,
+		rhs = rhs,
+	}
+end
+
+function Ast.CompoundDivStatement(lhs, rhs)
+	return {
+		kind = AstKind.CompoundDivStatement,
+		lhs = lhs,
+		rhs = rhs,
+	}
+end
+
+function Ast.CompoundPowStatement(lhs, rhs)
+	return {
+		kind = AstKind.CompoundPowStatement,
+		lhs = lhs,
+		rhs = rhs,
+	}
+end
+
+function Ast.CompoundModStatement(lhs, rhs)
+	return {
+		kind = AstKind.CompoundModStatement,
+		lhs = lhs,
+		rhs = rhs,
+	}
+end
+
+function Ast.CompoundConcatStatement(lhs, rhs)
+	return {
+		kind = AstKind.CompoundConcatStatement,
+		lhs = lhs,
+		rhs = rhs,
+	}
+end
+
 function Ast.FunctionCallStatement(base, args)
 	return {
 		kind = AstKind.FunctionCallStatement,

src/prometheus/enums.lua

@@ -69,7 +69,7 @@ Enums.Conventions = {
 		Symbols = {
 			"+",  "-",  "*",  "/",  "%",  "^",  "#",
 			"==", "~=", "<=", ">=", "<",  ">",  "=",
-			"+=", "-=", "/=", "%=", "^=", "..=",
+			"+=", "-=", "/=", "%=", "^=", "..=", "*=",
 			"(",  ")",  "{",  "}",  "[",  "]",
 			";",  ":",  ",",  ".",  "..", "...",
 			"::", "->", "?",  "|",  "&", 

src/prometheus/parser.lua

@@ -404,6 +404,45 @@ function Parser:statement(scope, currentLoop)
 			if(expr.kind == AstKind.VariableExpression) then
 				expr.kind = AstKind.AssignmentVariable
 			end
+
+			if(self.luaVersion == LuaVersion.LuaU) then
+				-- LuaU Compound Assignment
+				if(consume(self, TokenKind.Symbol, "+=")) then
+					local rhs = self:expression(scope);
+					return Ast.CompoundAddStatement(expr, rhs);
+				end
+
+				if(consume(self, TokenKind.Symbol, "-=")) then
+					local rhs = self:expression(scope);
+					return Ast.CompoundSubStatement(expr, rhs);
+				end
+
+				if(consume(self, TokenKind.Symbol, "*=")) then
+					local rhs = self:expression(scope);
+					return Ast.CompoundMulStatement(expr, rhs);
+				end
+
+				if(consume(self, TokenKind.Symbol, "/=")) then
+					local rhs = self:expression(scope);
+					return Ast.CompoundDivStatement(expr, rhs);
+				end
+
+				if(consume(self, TokenKind.Symbol, "%=")) then
+					local rhs = self:expression(scope);
+					return Ast.CompoundModStatement(expr, rhs);
+				end
+
+				if(consume(self, TokenKind.Symbol, "^=")) then
+					local rhs = self:expression(scope);
+					return Ast.CompoundPowStatement(expr, rhs);
+				end
+
+				if(consume(self, TokenKind.Symbol, "..=")) then
+					local rhs = self:expression(scope);
+					return Ast.CompoundConcatStatement(expr, rhs);
+				end
+			end
+
 			local lhs = {
 				expr
 			}

src/prometheus/unparser.lua

@@ -15,6 +15,7 @@ local config = require("config");
 local Ast    = require("prometheus.ast");
 local Enums  = require("prometheus.enums");
 local util = require("prometheus.util");
+local logger = require("logger");
 
 local lookupify = util.lookupify;
 local LuaVersion = Enums.LuaVersion;
@@ -104,7 +105,7 @@ end
 
 function Unparser:unparse(ast)
 	if(ast.kind ~= AstKind.TopNode) then
-		error("Unparser:unparse expects a TopNode as first argument")
+		logger:error("Unparser:unparse expects a TopNode as first argument")
 	end
 
 	return self:unparseBlock(ast.body);
@@ -381,9 +382,26 @@ function Unparser:unparseStatement(statement, tabbing)
 				code = code .. "," .. self:optionalWhitespace() .. exprcode;
 			end
 		end
-	-- While Statement
+	elseif self.luaVersion == LuaVersion.LuaU then
+		if statement.kind == AstKind.CompoundAddStatement then
+			code = code .. self:unparseExpression(statement.lhs, tabbing) .. self:optionalWhitespace() .. "+=" .. self:optionalWhitespace() .. self:unparseExpression(statement.rhs, tabbing);
+		elseif statement.kind == AstKind.CompoundSubStatement then
+			code = code .. self:unparseExpression(statement.lhs, tabbing) .. self:optionalWhitespace() .. "-=" .. self:optionalWhitespace() .. self:unparseExpression(statement.rhs, tabbing);
+		elseif statement.kind == AstKind.CompoundMulStatement then
+			code = code .. self:unparseExpression(statement.lhs, tabbing) .. self:optionalWhitespace() .. "*=" .. self:optionalWhitespace() .. self:unparseExpression(statement.rhs, tabbing);
+		elseif statement.kind == AstKind.CompoundDivStatement then
+			code = code .. self:unparseExpression(statement.lhs, tabbing) .. self:optionalWhitespace() .. "/=" .. self:optionalWhitespace() .. self:unparseExpression(statement.rhs, tabbing);
+		elseif statement.kind == AstKind.CompoundModStatement then
+			code = code .. self:unparseExpression(statement.lhs, tabbing) .. self:optionalWhitespace() .. "%=" .. self:optionalWhitespace() .. self:unparseExpression(statement.rhs, tabbing);
+		elseif statement.kind == AstKind.CompoundPowStatement then
+			code = code .. self:unparseExpression(statement.lhs, tabbing) .. self:optionalWhitespace() .. "^=" .. self:optionalWhitespace() .. self:unparseExpression(statement.rhs, tabbing);
+		elseif statement.kind == AstKind.CompoundConcatStatement then
+			code = code .. self:unparseExpression(statement.lhs, tabbing) .. self:optionalWhitespace() .. "..=" .. self:optionalWhitespace() .. self:unparseExpression(statement.rhs, tabbing);
+		else
+			logger:error(string.format("\"%s\" is not a valid unparseable statement in %s!", statement.kind, self.luaVersion));
+		end
 	else
-		error(string.format("\"%s\" is not a valid unparseable statement", statement.kind));
+		logger:error(string.format("\"%s\" is not a valid unparseable statement in %s!", statement.kind, self.luaVersion));
 	end
 
 	return self:tabs(tabbing, false) .. code;
@@ -839,7 +857,7 @@ function Unparser:unparseExpression(expression, tabbing)
 		return code .. self:optionalWhitespace((p and "," or "") .. self:newline() .. self:tabs(tabbing)) .. "}";
 	end
 
-	error(string.format("\"%s\" is not a valid unparseable expression", expression.kind));
+	logger:error(string.format("\"%s\" is not a valid unparseable expression", expression.kind));
 end
 
 return Unparser