removing ForStatementInitialization and ForStatementCondition from the AST tree

Author: Janther

src/slang-nodes/ForStatement.ts

@@ -2,7 +2,6 @@ import { NonterminalKind } from '@nomicfoundation/slang/cst';
 import { doc } from 'prettier';
 import { printSeparatedList } from '../slang-printers/print-separated-list.js';
 import { printIndentedGroupOrSpacedDocument } from '../slang-printers/print-indented-group-or-spaced-document.js';
-import { printVariant } from '../slang-printers/print-variant.js';
 import { extractVariant } from '../slang-utils/extract-variant.js';
 import { SlangNode } from './SlangNode.js';
 import { ForStatementInitialization } from './ForStatementInitialization.js';
@@ -20,9 +19,9 @@ const { line } = doc.builders;
 export class ForStatement extends SlangNode {
   readonly kind = NonterminalKind.ForStatement;
 
-  initialization: ForStatementInitialization;
+  initialization: ForStatementInitialization['variant'];
 
-  condition: ForStatementCondition;
+  condition: ForStatementCondition['variant'];
 
   iterator?: Expression['variant'];
 
@@ -31,11 +30,12 @@ export class ForStatement extends SlangNode {
   constructor(ast: ast.ForStatement, options: ParserOptions<AstNode>) {
     super(ast);
 
-    this.initialization = new ForStatementInitialization(
-      ast.initialization,
-      options
+    this.initialization = extractVariant(
+      new ForStatementInitialization(ast.initialization, options)
+    );
+    this.condition = extractVariant(
+      new ForStatementCondition(ast.condition, options)
     );
-    this.condition = new ForStatementCondition(ast.condition, options);
     if (ast.iterator) {
       this.iterator = extractVariant(new Expression(ast.iterator, options));
     }
@@ -50,8 +50,8 @@ export class ForStatement extends SlangNode {
   }
 
   print(path: AstPath<ForStatement>, print: PrintFunction): Doc {
-    const initialization = path.call(printVariant(print), 'initialization');
-    const condition = path.call(printVariant(print), 'condition');
+    const initialization = path.call(print, 'initialization');
+    const condition = path.call(print, 'condition');
     const iterator = path.call(print, 'iterator');
 
     return [

src/slang-nodes/ForStatementCondition.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 { ExpressionStatement } from './ExpressionStatement.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 ForStatementCondition extends PolymorphicNode {
+export class ForStatementCondition extends SlangNode {
   readonly kind = NonterminalKind.ForStatementCondition;
 
   variant: ExpressionStatement | TerminalNode;

src/slang-nodes/ForStatementInitialization.ts

@@ -3,7 +3,7 @@ import {
   NonterminalKind,
   TerminalNode as SlangTerminalNode
 } from '@nomicfoundation/slang/cst';
-import { PolymorphicNode } from './PolymorphicNode.js';
+import { SlangNode } from './SlangNode.js';
 import { ExpressionStatement } from './ExpressionStatement.js';
 import { VariableDeclarationStatement } from './VariableDeclarationStatement.js';
 import { TupleDeconstructionStatement } from './TupleDeconstructionStatement.js';
@@ -32,7 +32,7 @@ function createNonterminalVariant(
   return exhaustiveCheck;
 }
 
-export class ForStatementInitialization extends PolymorphicNode {
+export class ForStatementInitialization extends SlangNode {
   readonly kind = NonterminalKind.ForStatementInitialization;
 
   variant:

src/slang-printers/print-binary-operation.ts

@@ -38,7 +38,7 @@ export const shouldNotIndent = (
   isStatementWithoutIndentedOperation(node) ||
   (node.kind === NonterminalKind.ExpressionStatement &&
     (path.getNode(index + 1) as StrictAstNode).kind ===
-      NonterminalKind.ForStatementCondition);
+      NonterminalKind.ForStatement);
 
 export const binaryIndentRulesBuilder =
   (shouldIndent: (node: BinaryOperation) => boolean) =>

src/slangPrinter.ts

@@ -21,6 +21,8 @@ import type { ContractSpecifier } from './slang-nodes/ContractSpecifier.js';
 import type { ElementaryType } from './slang-nodes/ElementaryType.js';
 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';
 
 function hasNodeIgnoreComment({ comments }: StrictAstNode): boolean {
   // Prettier sets SourceUnit's comments to undefined after assigning comments
@@ -90,6 +92,8 @@ function genericPrint(
       | ElementaryType
       | ExperimentalFeature
       | Expression
+      | ForStatementInitialization
+      | ForStatementCondition
     >
   >,
   options: ParserOptions<AstNode>,