removing FunctionBody from the AST tree

Author: Janther

src/slang-nodes/FallbackFunctionDefinition.ts

@@ -1,5 +1,6 @@
 import { NonterminalKind } from '@nomicfoundation/slang/cst';
 import { printFunctionWithBody } from '../slang-printers/print-function.js';
+import { extractVariant } from '../slang-utils/extract-variant.js';
 import { SlangNode } from './SlangNode.js';
 import { ParametersDeclaration } from './ParametersDeclaration.js';
 import { FallbackFunctionAttributes } from './FallbackFunctionAttributes.js';
@@ -20,7 +21,7 @@ export class FallbackFunctionDefinition extends SlangNode {
 
   returns?: ReturnsDeclaration;
 
-  body: FunctionBody;
+  body: FunctionBody['variant'];
 
   constructor(
     ast: ast.FallbackFunctionDefinition,
@@ -33,7 +34,7 @@ export class FallbackFunctionDefinition extends SlangNode {
     if (ast.returns) {
       this.returns = new ReturnsDeclaration(ast.returns, options);
     }
-    this.body = new FunctionBody(ast.body, options);
+    this.body = extractVariant(new FunctionBody(ast.body, options));
 
     this.updateMetadata(
       this.parameters,

src/slang-nodes/FunctionBody.ts

@@ -2,15 +2,15 @@ import {
   NonterminalKind,
   TerminalNode as SlangTerminalNode
 } from '@nomicfoundation/slang/cst';
-import { PolymorphicNode } from './PolymorphicNode.js';
+import { SlangNode } from './SlangNode.js';
 import { Block } from './Block.js';
 import { TerminalNode } from './TerminalNode.js';
 
 import type * as ast from '@nomicfoundation/slang/ast';
 import type { ParserOptions } from 'prettier';
 import type { AstNode } from './types.d.ts';
 
-export class FunctionBody extends PolymorphicNode {
+export class FunctionBody extends SlangNode {
   readonly kind = NonterminalKind.FunctionBody;
 
   variant: Block | TerminalNode;

src/slang-nodes/FunctionDefinition.ts

@@ -2,6 +2,7 @@ import { satisfies } from 'semver';
 import { NonterminalKind } from '@nomicfoundation/slang/cst';
 import { printFunctionWithBody } from '../slang-printers/print-function.js';
 import { printVariant } from '../slang-printers/print-variant.js';
+import { extractVariant } from '../slang-utils/extract-variant.js';
 import { SlangNode } from './SlangNode.js';
 import { FunctionName } from './FunctionName.js';
 import { ParametersDeclaration } from './ParametersDeclaration.js';
@@ -25,7 +26,7 @@ export class FunctionDefinition extends SlangNode {
 
   returns?: ReturnsDeclaration;
 
-  body: FunctionBody;
+  body: FunctionBody['variant'];
 
   constructor(ast: ast.FunctionDefinition, options: ParserOptions<AstNode>) {
     super(ast);
@@ -36,7 +37,7 @@ export class FunctionDefinition extends SlangNode {
     if (ast.returns) {
       this.returns = new ReturnsDeclaration(ast.returns, options);
     }
-    this.body = new FunctionBody(ast.body, options);
+    this.body = extractVariant(new FunctionBody(ast.body, options));
 
     this.updateMetadata(
       this.name,

src/slang-nodes/ModifierDefinition.ts

@@ -1,5 +1,6 @@
 import { NonterminalKind } from '@nomicfoundation/slang/cst';
 import { printFunctionWithBody } from '../slang-printers/print-function.js';
+import { extractVariant } from '../slang-utils/extract-variant.js';
 import { SlangNode } from './SlangNode.js';
 import { TerminalNode } from './TerminalNode.js';
 import { ParametersDeclaration } from './ParametersDeclaration.js';
@@ -21,7 +22,7 @@ export class ModifierDefinition extends SlangNode {
 
   attributes: ModifierAttributes;
 
-  body: FunctionBody;
+  body: FunctionBody['variant'];
 
   constructor(ast: ast.ModifierDefinition, options: ParserOptions<AstNode>) {
     super(ast);
@@ -31,7 +32,7 @@ export class ModifierDefinition extends SlangNode {
       this.parameters = new ParametersDeclaration(ast.parameters, options);
     }
     this.attributes = new ModifierAttributes(ast.attributes);
-    this.body = new FunctionBody(ast.body, options);
+    this.body = extractVariant(new FunctionBody(ast.body, options));
 
     this.updateMetadata(this.parameters, this.attributes, this.body);
 

src/slang-nodes/ReceiveFunctionDefinition.ts

@@ -1,5 +1,6 @@
 import { NonterminalKind } from '@nomicfoundation/slang/cst';
 import { printFunctionWithBody } from '../slang-printers/print-function.js';
+import { extractVariant } from '../slang-utils/extract-variant.js';
 import { SlangNode } from './SlangNode.js';
 import { ParametersDeclaration } from './ParametersDeclaration.js';
 import { ReceiveFunctionAttributes } from './ReceiveFunctionAttributes.js';
@@ -17,7 +18,7 @@ export class ReceiveFunctionDefinition extends SlangNode {
 
   attributes: ReceiveFunctionAttributes;
 
-  body: FunctionBody;
+  body: FunctionBody['variant'];
 
   constructor(
     ast: ast.ReceiveFunctionDefinition,
@@ -27,7 +28,7 @@ export class ReceiveFunctionDefinition extends SlangNode {
 
     this.parameters = new ParametersDeclaration(ast.parameters, options);
     this.attributes = new ReceiveFunctionAttributes(ast.attributes, options);
-    this.body = new FunctionBody(ast.body, options);
+    this.body = extractVariant(new FunctionBody(ast.body, options));
 
     this.updateMetadata(this.parameters, this.attributes, this.body);
 

src/slang-nodes/UnnamedFunctionDefinition.ts

@@ -1,5 +1,6 @@
 import { NonterminalKind } from '@nomicfoundation/slang/cst';
 import { printFunctionWithBody } from '../slang-printers/print-function.js';
+import { extractVariant } from '../slang-utils/extract-variant.js';
 import { SlangNode } from './SlangNode.js';
 import { ParametersDeclaration } from './ParametersDeclaration.js';
 import { UnnamedFunctionAttributes } from './UnnamedFunctionAttributes.js';
@@ -17,7 +18,7 @@ export class UnnamedFunctionDefinition extends SlangNode {
 
   attributes: UnnamedFunctionAttributes;
 
-  body: FunctionBody;
+  body: FunctionBody['variant'];
 
   constructor(
     ast: ast.UnnamedFunctionDefinition,
@@ -27,7 +28,7 @@ export class UnnamedFunctionDefinition extends SlangNode {
 
     this.parameters = new ParametersDeclaration(ast.parameters, options);
     this.attributes = new UnnamedFunctionAttributes(ast.attributes, options);
-    this.body = new FunctionBody(ast.body, options);
+    this.body = extractVariant(new FunctionBody(ast.body, options));
 
     this.updateMetadata(this.parameters, this.attributes, this.body);
 

src/slang-printers/print-function.ts

@@ -1,13 +1,9 @@
 import { NonterminalKind } from '@nomicfoundation/slang/cst';
 import { doc } from 'prettier';
 import { joinExisting } from '../slang-utils/join-existing.js';
-import { printVariant } from './print-variant.js';
 
 import type { AstPath, Doc } from 'prettier';
-import type {
-  FunctionLike,
-  StrictPolymorphicNode
-} from '../slang-nodes/types.d.ts';
+import type { FunctionLike } from '../slang-nodes/types.d.ts';
 import type { PrintFunction } from '../types.d.ts';
 import type { FunctionBody } from '../slang-nodes/FunctionBody.js';
 
@@ -19,7 +15,9 @@ export function printFunction(
   path: AstPath<FunctionLike>,
   print: PrintFunction
 ): Doc {
-  const body = (node as Extract<FunctionLike, { body: FunctionBody }>).body;
+  const body = (
+    node as Extract<FunctionLike, { body: FunctionBody['variant'] }>
+  ).body;
 
   return group([
     functionName,
@@ -30,9 +28,7 @@ export function printFunction(
           path.call(print, 'attributes'),
           path.call(print, 'returns')
         ]),
-        body && (!body.variant || body.variant.kind === NonterminalKind.Block)
-          ? dedent(line)
-          : ''
+        body && body.kind === NonterminalKind.Block ? dedent(line) : ''
       ])
     )
   ]);
@@ -46,12 +42,6 @@ export function printFunctionWithBody(
 ): Doc {
   return [
     printFunction(functionName, node, path, print),
-    node.kind !== NonterminalKind.ConstructorDefinition
-      ? (
-          path as AstPath<
-            Extract<FunctionLike, { body: StrictPolymorphicNode }>
-          >
-        ).call(printVariant(print), 'body')
-      : path.call(print, 'body')
+    path.call(print, 'body')
   ];
 }

src/slangPrinter.ts

@@ -23,6 +23,7 @@ import type { ExperimentalFeature } from './slang-nodes/ExperimentalFeature.js';
 import type { Expression } from './slang-nodes/Expression.js';
 import type { ForStatementInitialization } from './slang-nodes/ForStatementInitialization.js';
 import type { ForStatementCondition } from './slang-nodes/ForStatementCondition.js';
+import type { FunctionBody } from './slang-nodes/FunctionBody.js';
 
 function hasNodeIgnoreComment({ comments }: StrictAstNode): boolean {
   // Prettier sets SourceUnit's comments to undefined after assigning comments
@@ -94,6 +95,7 @@ function genericPrint(
       | Expression
       | ForStatementInitialization
       | ForStatementCondition
+      | FunctionBody
     >
   >,
   options: ParserOptions<AstNode>,