CST: Handle named function declarations (#230)

JohnnyMorganz · web-flow · commit e7f5fb6a9a69 · 2025-04-29T20:54:44.000+02:00
Reapplication of #225, looks like it was merged into the assignment branch and I accidentally force pushed over it --- This PR extends the CST support to handle named function declarations (function name() or local function name()). As part of this PR, we also correct the comma positioning that separates arguments in the function definition, by using the comma positions stored on the CST node.
diff --git a/batteries/syntax/ast_types.luau b/batteries/syntax/ast_types.luau
@@ -266,6 +266,21 @@ export type AstStatCompoundAssign = {
 	value: AstExpr,
 }
 
+export type AstStatFunction = {
+	tag: "function",
+	["function"]: Token<"function">,
+	name: AstExpr,
+	body: AstFunctionBody,
+}
+
+export type AstStatLocalFunction = {
+	tag: "localfunction",
+	["local"]: Token<"local">,
+	["function"]: Token<"function">,
+	name: AstLocal,
+	body: AstFunctionBody,
+}
+
 export type AstStat =
 	| AstStatBlock
 	| AstStatIf
@@ -280,5 +295,7 @@ export type AstStat =
 	| AstStatForIn
 	| AstStatAssign
 	| AstStatCompoundAssign
+	| AstStatFunction
+	| AstStatLocalFunction
 
 return {}
diff --git a/batteries/syntax/visitor.luau b/batteries/syntax/visitor.luau
@@ -13,6 +13,8 @@ type Visitor = {
 	visitForIn: (T.AstStatForIn) -> boolean,
 	visitAssign: (T.AstStatAssign) -> boolean,
 	visitCompoundAssign: (T.AstStatCompoundAssign) -> boolean,
+	visitFunction: (T.AstStatFunction) -> boolean,
+	visitLocalFunction: (T.AstStatLocalFunction) -> boolean,
 
 	visitLocalReference: (T.AstExprLocal) -> boolean,
 	visitGlobal: (T.AstExprGlobal) -> boolean,
@@ -50,6 +52,8 @@ local defaultVisitor: Visitor = {
 	visitForIn = alwaysVisit :: any,
 	visitAssign = alwaysVisit :: any,
 	visitCompoundAssign = alwaysVisit :: any,
+	visitFunction = alwaysVisit :: any,
+	visitLocalFunction = alwaysVisit :: any,
 
 	visitLocalReference = alwaysVisit :: any,
 	visitGlobal = alwaysVisit :: any,
@@ -290,6 +294,23 @@ local function visitAnonymousFunction(node: T.AstExprAnonymousFunction, visitor:
 	end
 end
 
+local function visitFunction(node: T.AstStatFunction, visitor: Visitor)
+	if visitor.visitFunction(node) then
+		visitToken(node["function"], visitor)
+		visitExpression(node.name, visitor)
+		visitFunctionBody(node.body, visitor)
+	end
+end
+
+local function visitLocalFunction(node: T.AstStatLocalFunction, visitor: Visitor)
+	if visitor.visitLocalFunction(node) then
+		visitToken(node["local"], visitor)
+		visitToken(node["function"], visitor)
+		visitLocal(node.name, visitor)
+		visitFunctionBody(node.body, visitor)
+	end
+end
+
 local function visitTableItem(node: T.AstExprTableItem, visitor: Visitor)
 	if visitor.visitTableItem(node) then
 		if node.kind == "list" then
@@ -412,6 +433,10 @@ function visitStatement(statement: T.AstStat, visitor: Visitor)
 		visitAssign(statement, visitor)
 	elseif statement.tag == "compoundassign" then
 		visitCompoundAssign(statement, visitor)
+	elseif statement.tag == "function" then
+		visitFunction(statement, visitor)
+	elseif statement.tag == "localfunction" then
+		visitLocalFunction(statement, visitor)
 	else
 		exhaustiveMatch(statement.tag)
 	end
diff --git a/luau/src/luau.cpp b/luau/src/luau.cpp
@@ -307,7 +307,7 @@ struct AstSerialize : public Luau::AstVisitor
         lua_pushstring(L, name.value);
     }
 
-    void serialize(Luau::AstLocal* local)
+    void serialize(Luau::AstLocal* local, bool createToken = true)
     {
         lua_rawcheckstack(L, 2);
 
@@ -324,8 +324,11 @@ struct AstSerialize : public Luau::AstVisitor
             lua_pushvalue(L, -2);
             lua_settable(L, localTableIndex);
 
-            serializeToken(local->location.begin, local->name.value);
-            lua_setfield(L, -2, "name");
+            if (createToken)
+            {
+                serializeToken(local->location.begin, local->name.value);
+                lua_setfield(L, -2, "name");
+            }
 
             if (local->shadow)
                 serialize(local->shadow);
@@ -789,11 +792,13 @@ struct AstSerialize : public Luau::AstVisitor
 
     void serializeFunctionBody(Luau::AstExprFunction* node)
     {
+        const auto* cstNode = lookupCstNode<Luau::CstExprFunction>(node);
+
         lua_rawcheckstack(L, 3);
         lua_createtable(L, 0, 7);
 
         if (node->self)
-            serialize(node->self);
+            serialize(node->self, /* createToken= */ false);
         else
             lua_pushnil(L);
         lua_setfield(L, -2, "self");
@@ -804,8 +809,7 @@ struct AstSerialize : public Luau::AstVisitor
             lua_setfield(L, -2, "openParens");
         }
 
-        // TODO: separators
-        serializePunctuated(node->args, {}, ",");
+        serializePunctuated(node->args, cstNode ? cstNode->argsCommaPositions : Luau::AstArray<Luau::Position>{}, ",");
         lua_setfield(L, -2, "parameters");
 
         // TODO: generics, return types, etc.
@@ -1283,29 +1287,41 @@ struct AstSerialize : public Luau::AstVisitor
     void serializeStat(Luau::AstStatFunction* node)
     {
         lua_rawcheckstack(L, 2);
-        lua_createtable(L, 0, preambleSize + 2);
+        lua_createtable(L, 0, preambleSize + 3);
 
         serializeNodePreamble(node, "function");
 
+        serializeToken(node->location.begin, "function");
+        lua_setfield(L, -2, "function");
+
         node->name->visit(this);
         lua_setfield(L, -2, "name");
 
-        node->func->visit(this);
-        lua_setfield(L, -2, "function");
+        serializeFunctionBody(node->func);
+        lua_setfield(L, -2, "body");
     }
 
     void serializeStat(Luau::AstStatLocalFunction* node)
     {
         lua_rawcheckstack(L, 2);
-        lua_createtable(L, 0, preambleSize + 2);
+        lua_createtable(L, 0, preambleSize + 4);
 
         serializeNodePreamble(node, "localfunction");
 
+        serializeToken(node->location.begin, "local");
+        lua_setfield(L, -2, "local");
+
+        if (const auto cstNode = lookupCstNode<Luau::CstStatLocalFunction>(node))
+        {
+            serializeToken(cstNode->functionKeywordPosition, "function");
+            lua_setfield(L, -2, "function");
+        }
+
         serialize(node->name);
         lua_setfield(L, -2, "name");
 
-        node->func->visit(this);
-        lua_setfield(L, -2, "function");
+        serializeFunctionBody(node->func);
+        lua_setfield(L, -2, "body");
     }
 
     void serializeStat(Luau::AstStatTypeAlias* node)
diff --git a/tests/astSerializerTests/function-declaration-1.luau b/tests/astSerializerTests/function-declaration-1.luau
@@ -0,0 +1,5 @@
+function x()
+	call()
+end
+
+function x.y:z() end
diff --git a/tests/astSerializerTests/function-declaration-2.luau b/tests/astSerializerTests/function-declaration-2.luau
@@ -0,0 +1,3 @@
+-- stylua: ignore
+function x(x  ,  y,  z)
+end
diff --git a/tests/astSerializerTests/local-function-declaration-1.luau b/tests/astSerializerTests/local-function-declaration-1.luau
@@ -0,0 +1,3 @@
+local function x()
+	call(1)
+end
diff --git a/tests/testAstSerializer.spec.luau b/tests/testAstSerializer.spec.luau
@@ -132,7 +132,10 @@ local function test_roundtrippableAst()
 		"tests/astSerializerTests/assignment-1.luau",
 		"tests/astSerializerTests/break-continue-1.luau",
 		"tests/astSerializerTests/compound-assignment-1.luau",
+		"tests/astSerializerTests/function-declaration-1.luau",
+		"tests/astSerializerTests/function-declaration-2.luau",
 		"tests/astSerializerTests/generic-for-loop-1.luau",
+		"tests/astSerializerTests/local-function-declaration-1.luau",
 		"tests/astSerializerTests/numeric-for-loop-1.luau",
 		"tests/astSerializerTests/while-1.luau",
 		"tests/astSerializerTests/repeat-until-1.luau",

-Original file line number
+Diff line change
@@ @@ -0,0 +1,5 @@ @@
 +function x()
 +	call()
 +end
++
 +function x.y:z() end
Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,3 @@`
	`1`	`+-- stylua: ignore`
	`2`	`+function x(x , y, z)`
	`3`	`+end`
Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,3 @@`
	`1`	`+local function x()`
	`2`	`+ call(1)`
	`3`	`+end`