CST: Handle for loops (#223)

This PR extends the CST support to include all the necessary tokens in
numeric (`for i = 1, 10, 2 do`) and for-in (`for ... in ... do`) loops.

As mentioned previously, a future change will standardise the AST
property names to switch away from using keywords

Author: JohnnyMorganz

batteries/syntax/ast_types.luau

@@ -226,6 +226,31 @@ export type AstStatLocal = {
 	values: Punctuated<AstExpr>,
 }
 
+export type AstStatFor = {
+	tag: "for",
+	["for"]: Token<"for">,
+	variable: Token<string>,
+	equals: Token<"=">,
+	from: AstExpr,
+	toComma: Token<",">,
+	to: AstExpr,
+	stepComma: Token<",">?,
+	step: AstExpr?,
+	["do"]: Token<"do">,
+	body: AstStatBlock,
+	["end"]: Token<"end">,
+}
+
+export type AstStatForIn = {
+	tag: "forin",
+	["for"]: Token<"for">,
+	variables: Punctuated<Token<string>>,
+	["in"]: Token<"in">,
+	values: Punctuated<Token<string>>,
+	["do"]: Token<"do">,
+	body: AstStatBlock,
+	["end"]: Token<"end">,
+}
 export type AstStat =
 	| AstStatBlock
 	| AstStatIf
@@ -236,5 +261,7 @@ export type AstStat =
 	| AstStatReturn
 	| AstStatExpr
 	| AstStatLocal
+	| AstStatFor
+	| AstStatForIn
 
 return {}

batteries/syntax/visitor.luau

@@ -9,6 +9,8 @@ type Visitor = {
 	visitRepeat: (T.AstStatRepeat) -> boolean,
 	visitReturn: (T.AstStatReturn) -> boolean,
 	visitLocalDeclaration: (T.AstStatLocal) -> boolean,
+	visitFor: (T.AstStatFor) -> boolean,
+	visitForIn: (T.AstStatForIn) -> boolean,
 
 	visitLocalReference: (T.AstExprLocal) -> boolean,
 	visitGlobal: (T.AstExprGlobal) -> boolean,
@@ -42,6 +44,8 @@ local defaultVisitor: Visitor = {
 	visitRepeat = alwaysVisit :: any,
 	visitReturn = alwaysVisit :: any,
 	visitLocalDeclaration = alwaysVisit :: any,
+	visitFor = alwaysVisit :: any,
+	visitForIn = alwaysVisit :: any,
 
 	visitLocalReference = alwaysVisit :: any,
 	visitGlobal = alwaysVisit :: any,
@@ -149,6 +153,38 @@ local function visitLocalStatement(node: T.AstStatLocal, visitor: Visitor)
 	end
 end
 
+local function visitFor(node: T.AstStatFor, visitor: Visitor)
+	if visitor.visitFor(node) then
+		visitToken(node["for"], visitor)
+		visitLocal(node.variable, visitor)
+		visitToken(node.equals, visitor)
+		visitExpression(node.from, visitor)
+		visitToken(node.toComma, visitor)
+		visitExpression(node.to, visitor)
+		if node.stepComma then
+			visitToken(node.stepComma, visitor)
+		end
+		if node.step then
+			visitExpression(node.step, visitor)
+		end
+		visitToken(node["do"], visitor)
+		visitBlock(node.body, visitor)
+		visitToken(node["end"], visitor)
+	end
+end
+
+local function visitForIn(node: T.AstStatForIn, visitor: Visitor)
+	if visitor.visitForIn(node) then
+		visitToken(node["for"], visitor)
+		visitPunctuated(node.variables, visitor, visitLocal)
+		visitToken(node["in"], visitor)
+		visitPunctuated(node.values, visitor, visitExpression)
+		visitToken(node["do"], visitor)
+		visitBlock(node.body, visitor)
+		visitToken(node["end"], visitor)
+	end
+end
+
 local function visitString(node: T.AstExprConstantString, visitor: Visitor)
 	if visitor.visitString(node) then
 		visitor.visitToken(node)
@@ -348,6 +384,10 @@ function visitStatement(statement: T.AstStat, visitor: Visitor)
 		visitToken(statement, visitor)
 	elseif statement.tag == "repeat" then
 		visitRepeat(statement, visitor)
+	elseif statement.tag == "for" then
+		visitFor(statement, visitor)
+	elseif statement.tag == "forin" then
+		visitForIn(statement, visitor)
 	else
 		exhaustiveMatch(statement.tag)
 	end

luau/src/luau.cpp

@@ -1146,59 +1146,96 @@ struct AstSerialize : public Luau::AstVisitor
     void serializeStat(Luau::AstStatFor* node)
     {
         lua_rawcheckstack(L, 2);
-        lua_createtable(L, 0, preambleSize + 6);
+        lua_createtable(L, 0, preambleSize + 11);
+
+        const auto cstNode = lookupCstNode<Luau::CstStatFor>(node);
 
         serializeNodePreamble(node, "for");
 
+        serializeToken(node->location.begin, "for");
+        lua_setfield(L, -2, "for");
+
         serialize(node->var);
         lua_setfield(L, -2, "variable");
 
+        if (cstNode)
+        {
+            serializeToken(cstNode->equalsPosition, "=");
+            lua_setfield(L, -2, "equals");
+        }
+
         node->from->visit(this);
         lua_setfield(L, -2, "from");
 
+        if (cstNode)
+        {
+            serializeToken(cstNode->endCommaPosition, ",");
+            lua_setfield(L, -2, "toComma");
+        }
+
         node->to->visit(this);
         lua_setfield(L, -2, "to");
 
-        node->step->visit(this);
-        lua_setfield(L, -2, "step");
+        if (cstNode && cstNode->stepCommaPosition)
+        {
+            serializeToken(*cstNode->stepCommaPosition, ",");
+            lua_setfield(L, -2, "stepComma");
+        }
 
-        node->body->visit(this);
-        lua_setfield(L, -2, "body");
+        if (node->step)
+            node->step->visit(this);
+        else
+            lua_pushnil(L);
+        lua_setfield(L, -2, "step");
 
         if (node->hasDo)
-            serialize(node->doLocation);
+            serializeToken(node->doLocation.begin, "do");
         else
             lua_pushnil(L);
-        lua_setfield(L, -2, "doLocation");
+        lua_setfield(L, -2, "do");
+
+        node->body->visit(this);
+        lua_setfield(L, -2, "body");
+
+        serializeToken(node->body->location.end, "end");
+        lua_setfield(L, -2, "end");
     }
 
     void serializeStat(Luau::AstStatForIn* node)
     {
         lua_rawcheckstack(L, 2);
-        lua_createtable(L, 0, preambleSize + 5);
+        lua_createtable(L, 0, preambleSize + 7);
 
-        serializeNodePreamble(node, "forin");
+        const auto cstNode = lookupCstNode<Luau::CstStatForIn>(node);
 
-        serializeLocals(node->vars);
-        lua_setfield(L, -2, "variables");
+        serializeNodePreamble(node, "forin");
 
-        serializeExprs(node->values);
-        lua_setfield(L, -2, "values");
+        serializeToken(node->location.begin, "for");
+        lua_setfield(L, -2, "for");
 
-        node->body->visit(this);
-        lua_setfield(L, -2, "body");
+        serializePunctuated(node->vars, cstNode ? cstNode->varsCommaPositions : Luau::AstArray<Luau::Position>{}, ",");
+        lua_setfield(L, -2, "variables");
 
         if (node->hasIn)
-            serialize(node->inLocation);
+            serializeToken(node->inLocation.begin, "in");
         else
             lua_pushnil(L);
-        lua_setfield(L, -2, "inLocation");
+        lua_setfield(L, -2, "in");
+
+        serializePunctuated(node->values, cstNode ? cstNode->valuesCommaPositions : Luau::AstArray<Luau::Position>{}, ",");
+        lua_setfield(L, -2, "values");
 
         if (node->hasDo)
-            serialize(node->doLocation);
+            serializeToken(node->doLocation.begin, "do");
         else
             lua_pushnil(L);
-        lua_setfield(L, -2, "doLocation");
+        lua_setfield(L, -2, "do");
+
+        node->body->visit(this);
+        lua_setfield(L, -2, "body");
+
+        serializeToken(node->body->location.end, "end");
+        lua_setfield(L, -2, "end");
     }
 
     void serializeStat(Luau::AstStatAssign* node)

tests/astSerializerTests/generic-for-loop-1.luau

@@ -0,0 +1,11 @@
+for index, value in pairs(list) do
+	call(index, value)
+end
+
+for index, value in next, list do
+	call(index, value)
+end
+
+for index, value in list do
+	call(index, value)
+end

tests/astSerializerTests/numeric-for-loop-1.luau

@@ -0,0 +1,7 @@
+for index = 1, 10 do
+	call(index)
+end
+for _ = start, final do
+end
+for _ = 1, 10, 2 do
+end

tests/testAstSerializer.spec.luau

@@ -130,6 +130,8 @@ local function test_roundtrippableAst()
 		"examples/time_example.luau",
 		"examples/writeFile.luau",
 		"tests/astSerializerTests/break-continue-1.luau",
+		"tests/astSerializerTests/generic-for-loop-1.luau",
+		"tests/astSerializerTests/numeric-for-loop-1.luau",
 		"tests/astSerializerTests/while-1.luau",
 		"tests/astSerializerTests/repeat-until-1.luau",
 	}