CST: Handle repeat, break and continue (#210)

This PR adds support for serializing all the the repeat stmt node, as
well as `break` and `continue` last stmt

For break and continue, it is fairly simple - we serialize the tokens
directly.

For repeat, we serialize the `repeat` and `until` tokens using the CST
data, alongside the condition and body.

Author: JohnnyMorganz

batteries/syntax/ast_types.luau

@@ -195,6 +195,18 @@ export type AstStatWhile = {
 	["end"]: Token<"end">,
 }
 
+export type AstStatRepeat = {
+	tag: "repeat",
+	["repeat"]: Token<"repeat">,
+	body: AstStatBlock,
+	["until"]: Token<"until">,
+	condition: AstExpr,
+}
+
+export type AstStatBreak = Token<"break"> & { tag: "break" }
+
+export type AstStatContinue = Token<"continue"> & { tag: "continue" }
+
 export type AstStatReturn = {
 	tag: "return",
 	["return"]: Token<"return">,
@@ -214,6 +226,15 @@ export type AstStatLocal = {
 	values: Punctuated<AstExpr>,
 }
 
-export type AstStat = AstStatBlock | AstStatIf | AstStatWhile | AstStatReturn | AstStatExpr | AstStatLocal
+export type AstStat =
+	| AstStatBlock
+	| AstStatIf
+	| AstStatWhile
+	| AstStatRepeat
+	| AstStatBreak
+	| AstStatContinue
+	| AstStatReturn
+	| AstStatExpr
+	| AstStatLocal
 
 return {}

batteries/syntax/visitor.luau

@@ -6,6 +6,7 @@ type Visitor = {
 	visitBlock: (T.AstStatBlock) -> boolean,
 	visitIf: (T.AstStatIf) -> boolean,
 	visitWhile: (T.AstStatWhile) -> boolean,
+	visitRepeat: (T.AstStatRepeat) -> boolean,
 	visitReturn: (T.AstStatReturn) -> boolean,
 	visitLocalDeclaration: (T.AstStatLocal) -> boolean,
 
@@ -38,6 +39,7 @@ local defaultVisitor: Visitor = {
 	visitBlock = alwaysVisit :: any,
 	visitIf = alwaysVisit :: any,
 	visitWhile = alwaysVisit :: any,
+	visitRepeat = alwaysVisit :: any,
 	visitReturn = alwaysVisit :: any,
 	visitLocalDeclaration = alwaysVisit :: any,
 
@@ -120,6 +122,15 @@ local function visitWhile(node: T.AstStatWhile, visitor: Visitor)
 	end
 end
 
+local function visitRepeat(node: T.AstStatRepeat, visitor: Visitor)
+	if visitor.visitRepeat(node) then
+		visitToken(node["repeat"], visitor)
+		visitBlock(node.body, visitor)
+		visitToken(node["until"], visitor)
+		visitExpression(node.condition, visitor)
+	end
+end
+
 local function visitReturn(node: T.AstStatReturn, visitor: Visitor)
 	if visitor.visitReturn(node) then
 		visitToken(node["return"], visitor)
@@ -331,6 +342,12 @@ function visitStatement(statement: T.AstStat, visitor: Visitor)
 		visitReturn(statement, visitor)
 	elseif statement.tag == "while" then
 		visitWhile(statement, visitor)
+	elseif statement.tag == "break" then
+		visitToken(statement, visitor)
+	elseif statement.tag == "continue" then
+		visitToken(statement, visitor)
+	elseif statement.tag == "repeat" then
+		visitRepeat(statement, visitor)
 	else
 		exhaustiveMatch(statement.tag)
 	end

luau/src/luau.cpp

@@ -1063,26 +1063,33 @@ struct AstSerialize : public Luau::AstVisitor
 
         serializeNodePreamble(node, "repeat");
 
-        node->condition->visit(this);
-        lua_setfield(L, -2, "condition");
+        serializeToken(node->location.begin, "repeat");
+        lua_setfield(L, -2, "repeat");
 
         node->body->visit(this);
         lua_setfield(L, -2, "body");
+
+        if (const auto cstNode = lookupCstNode<Luau::CstStatRepeat>(node))
+        {
+            serializeToken(cstNode->untilPosition, "until");
+            lua_setfield(L, -2, "until");
+        }
+
+        node->condition->visit(this);
+        lua_setfield(L, -2, "condition");
     }
 
     void serializeStat(Luau::AstStatBreak* node)
     {
         lua_rawcheckstack(L, 2);
-        lua_createtable(L, 0, preambleSize);
-
+        serializeToken(node->location.begin, "break", preambleSize);
         serializeNodePreamble(node, "break");
     }
 
     void serializeStat(Luau::AstStatContinue* node)
     {
         lua_rawcheckstack(L, 2);
-        lua_createtable(L, 0, preambleSize);
-
+        serializeToken(node->location.begin, "continue", preambleSize);
         serializeNodePreamble(node, "continue");
     }
 

tests/astSerializerTests/break-continue-1.luau

@@ -0,0 +1,7 @@
+while true do
+	break
+end
+
+while true do
+	continue
+end

tests/astSerializerTests/repeat-until-1.luau

@@ -0,0 +1,3 @@
+repeat
+	call()
+until condition

tests/testAstSerializer.spec.luau

@@ -129,7 +129,9 @@ local function test_roundtrippableAst()
 		"examples/parsing.luau",
 		"examples/time_example.luau",
 		"examples/writeFile.luau",
+		"tests/astSerializerTests/break-continue-1.luau",
 		"tests/astSerializerTests/while-1.luau",
+		"tests/astSerializerTests/repeat-until-1.luau",
 	}
 
 	for _, path in files do