[Formatter] Fix comment insertion inside a function declaration withi…

…n a global assignment statement

server/src/parser/formatter.ts

@@ -13,6 +13,8 @@ import {
   GotoStatement,
   IfClause,
   IfStatement,
+  isAssignmentStatement,
+  isFunctionDeclaration,
   isIfStatement,
   isStatementWithBody,
   LabelStatement,
@@ -131,6 +133,13 @@ export default class Formatter {
           this.insertComment(comment, (currStmt as any).body);
           return;
         }
+      } else if (compareResult === BoundsCompareResult.CONTAINS && isAssignmentStatement(currStmt)) {
+        for (const expression of currStmt.init) {
+          if (isFunctionDeclaration(expression) && boundsCompare(comment.loc!, expression.loc!) === BoundsCompareResult.CONTAINS) {
+            this.insertComment(comment, expression.body);
+            return;
+          }
+        }
       } else if (compareResult === BoundsCompareResult.BEFORE) {
         // This comment is before the current statement -- insert it at i
         body.splice(i, 0, comment);

server/src/parser/statements.ts

@@ -33,6 +33,14 @@ export function isIfStatement(stmt: any): stmt is IfStatement {
   return stmt?.type === 'IfStatement';
 }
 
+export function isAssignmentStatement(stmt: any): stmt is AssignmentStatement {
+  return stmt?.type === 'AssignmentStatement';
+}
+
+export function isFunctionDeclaration(stmt: any): stmt is FunctionDeclaration {
+  return stmt?.type === 'FunctionDeclaration';
+}
+
 export type IfClause = ASTNode & {
   type: 'IfClause',
   condition: Expression,

server/src/parser/test/formatter.test.ts

@@ -242,7 +242,25 @@ end
       eq(format(input), input);
     });
 
-    it('keeps comments around and inside an assigned function', () => {
+    it('keeps comments around and inside an assigned function (case of global assignment)', () => {
+      const input = `
+-- comment before an assigned function
+a = function(b)
+  -- print(b)
+  do_something(b)
+  --[[
+    do_another_thing(b)
+    do_another_thing(b + 1)
+  ]]
+  do_something_totally_different(b)
+  -- print(b - 1)
+end
+-- comment after an assigned function
+        `.trim();
+      eq(format(input), input);
+    });
+
+    it('keeps comments around and inside an assigned function (case of local assignment)', () => {
       const input = `
 -- comment before an assigned function
 local a = function(b)