Handle if stmt and nil

Enough to parse the writeFile test case

Author: JohnnyMorganz

examples/codemods/identity.luau

@@ -14,3 +14,4 @@ codemod.apply("examples/json.luau", identityTransform)
 codemod.apply("examples/main.luau", identityTransform)
 codemod.apply("examples/parsing.luau", identityTransform)
 codemod.apply("examples/time_example.luau", identityTransform)
+codemod.apply("examples/writeFile.luau", identityTransform)

luau/src/luau.cpp

@@ -472,9 +472,7 @@ struct AstSerialize : public Luau::AstVisitor
 
     void serialize(Luau::AstExprConstantNil* node)
     {
-        lua_rawcheckstack(L, 2);
-        lua_createtable(L, 0, preambleSize);
-
+        serializeToken(node->location.begin, "nil", preambleSize);
         serializeNodePreamble(node, "nil");
     }
 
@@ -825,29 +823,41 @@ struct AstSerialize : public Luau::AstVisitor
 
         serializeNodePreamble(node, "conditional");
 
+        serializeToken(node->location.begin, "if");
+        lua_setfield(L, -2, "if");
+
         node->condition->visit(this);
         lua_setfield(L, -2, "condition");
 
+        serializeToken(node->thenLocation->begin, "then");
+        lua_setfield(L, -2, "then");
+
         node->thenbody->visit(this);
         lua_setfield(L, -2, "consequent");
 
         if (node->elsebody)
+        {
+            LUAU_ASSERT(node->elseLocation);
+            serializeToken(node->elseLocation->begin, "else");
+            lua_setfield(L, -2, "else");
+
             node->elsebody->visit(this);
-        else
-            lua_pushnil(L);
-        lua_setfield(L, -2, "antecedent");
+            lua_setfield(L, -2, "antecedent");
 
-        if (node->thenLocation)
-            serialize(*node->thenLocation);
+            serializeToken(node->elsebody->location.end, "end");
+            lua_setfield(L, -2, "end");
+        }
         else
+        {
             lua_pushnil(L);
-        lua_setfield(L, -2, "thenLocation");
+            lua_setfield(L, -2, "else");
 
-        if (node->elseLocation)
-            serialize(*node->elseLocation);
-        else
             lua_pushnil(L);
-        lua_setfield(L, -2, "elseLocation");
+            lua_setfield(L, -2, "antecedent");
+
+            serializeToken(node->thenbody->location.end, "end");
+            lua_setfield(L, -2, "end");
+        }
     }
 
     void serializeStat(Luau::AstStatWhile* node)

std/codemod/ast_types.luau

@@ -43,6 +43,8 @@ export type AstLocal = {
 	name: Token<string>,
 }
 
+export type AstExprConstantNil = Token<"nil"> & { tag: "nil" }
+
 export type AstExprConstantBool = Token<"true" | "false"> & { tag: "boolean", value: boolean }
 
 export type AstExprConstantNumber = Token<string> & { tag: "number", value: number }
@@ -118,6 +120,7 @@ export type AstExprBinary = {
 }
 
 export type AstExpr =
+	| AstExprConstantNil
 	| AstExprConstantBool
 	| AstExprConstantNumber
 	| AstExprConstantString
@@ -134,6 +137,17 @@ export type AstStatBlock = {
 	statements: { AstStat },
 }
 
+export type AstStatIf = {
+	tag: "conditional",
+	["if"]: Token<"if">,
+	condition: AstExpr,
+	["then"]: Token<"then">,
+	consequent: AstStatBlock,
+	["else"]: Token<"else">, -- TODO: this could be elseif!
+	antecedent: AstStatBlock,
+	["end"]: Token<"end">,
+}
+
 export type AstStatReturn = {
 	tag: "return",
 	["return"]: Token<"return">,
@@ -153,6 +167,6 @@ export type AstStatLocal = {
 	values: Punctuated<AstExpr>,
 }
 
-export type AstStat = AstStatBlock | AstStatReturn | AstStatExpr | AstStatLocal
+export type AstStat = AstStatBlock | AstStatIf | AstStatReturn | AstStatExpr | AstStatLocal
 
 return {}

std/codemod/visitor.luau

@@ -4,6 +4,7 @@ local T = require("./ast_types")
 
 type Visitor = {
 	visitBlock: (T.AstStatBlock) -> boolean,
+	visitIf: (T.AstStatIf) -> boolean,
 	visitReturn: (T.AstStatReturn) -> boolean,
 	visitLocalDeclaration: (T.AstStatLocal) -> boolean,
 
@@ -16,6 +17,7 @@ type Visitor = {
 	visitIndexName: (T.AstExprIndexName) -> boolean,
 
 	visitToken: (T.Token) -> boolean,
+	visitNil: (T.AstExprConstantNil) -> boolean,
 	visitString: (T.AstExprConstantString) -> boolean,
 	visitBoolean: (T.AstExprConstantBool) -> boolean,
 	visitNumber: (T.AstExprConstantNumber) -> boolean,
@@ -28,6 +30,7 @@ end
 
 local defaultVisitor: Visitor = {
 	visitBlock = alwaysVisit :: any,
+	visitIf = alwaysVisit :: any,
 	visitReturn = alwaysVisit :: any,
 	visitLocalDeclaration = alwaysVisit :: any,
 
@@ -40,6 +43,7 @@ local defaultVisitor: Visitor = {
 	visitIndexName = alwaysVisit :: any,
 
 	visitToken = alwaysVisit :: any,
+	visitNil = alwaysVisit :: any,
 	visitString = alwaysVisit :: any,
 	visitBoolean = alwaysVisit :: any,
 	visitNumber = alwaysVisit :: any,
@@ -77,6 +81,23 @@ local function visitBlock(block: T.AstStatBlock, visitor: Visitor)
 	end
 end
 
+local function visitIf(node: T.AstStatIf, visitor: Visitor)
+	if visitor.visitIf(node) then
+		visitToken(node["if"], visitor)
+		visitExpression(node.condition, visitor)
+		visitToken(node["then"], visitor)
+		visitBlock(node.consequent, visitor)
+		-- TODO: special handling for elseif?
+		if node["else"] then
+			visitToken(node["else"], visitor)
+		end
+		if node.antecedent then
+			visitBlock(node.antecedent, visitor)
+		end
+		visitToken(node["end"], visitor)
+	end
+end
+
 local function visitReturn(node: T.AstStatReturn, visitor: Visitor)
 	if visitor.visitReturn(node) then
 		visitToken(node["return"], visitor)
@@ -101,6 +122,12 @@ local function visitString(node: T.AstExprConstantString, visitor: Visitor)
 	end
 end
 
+local function visitNil(node: T.AstExprConstantNil, visitor: Visitor)
+	if visitor.visitNil(node) then
+		visitToken(node, visitor)
+	end
+end
+
 local function visitBoolean(node: T.AstExprConstantBool, visitor: Visitor)
 	if visitor.visitBoolean(node) then
 		visitToken(node, visitor)
@@ -189,7 +216,9 @@ local function visitIndexName(node: T.AstExprIndexName, visitor: Visitor)
 end
 
 function visitExpression(expression: T.AstExpr, visitor: Visitor)
-	if expression.tag == "boolean" then
+	if expression.tag == "nil" then
+		visitNil(expression, visitor)
+	elseif expression.tag == "boolean" then
 		visitBoolean(expression, visitor)
 	elseif expression.tag == "number" then
 		visitNumber(expression, visitor)
@@ -215,6 +244,8 @@ end
 function visitStatement(statement: T.AstStat, visitor: Visitor)
 	if statement.tag == "block" then
 		visitBlock(statement, visitor)
+	elseif statement.tag == "conditional" then
+		visitIf(statement, visitor)
 	elseif statement.tag == "expression" then
 		visitExpression(statement.expression, visitor)
 	elseif statement.tag == "local" then