removing ContractMember, SourceUnitMember, Statement, and YulStatement from the AST tree (#1284)

Author: Janther

src/slang-comments/handlers/handle-else-branch-comments.ts

@@ -17,15 +17,12 @@ export default function handleElseBranchComments({
 
   if (
     followingNode === enclosingNode.body &&
-    followingNode.variant.kind === NonterminalKind.IfStatement
+    followingNode.kind === NonterminalKind.IfStatement
   ) {
-    if (followingNode.variant.body.variant.kind === NonterminalKind.Block) {
-      addCollectionFirstComment(
-        followingNode.variant.body.variant.statements,
-        comment
-      );
+    if (followingNode.body.kind === NonterminalKind.Block) {
+      addCollectionFirstComment(followingNode.body.statements, comment);
     } else {
-      addLeadingComment(followingNode.variant.body.variant, comment);
+      addLeadingComment(followingNode.body, comment);
     }
     return true;
   }

src/slang-comments/handlers/handle-if-statement-comments.ts

@@ -39,15 +39,15 @@ export default function handleIfStatementComments({
     precedingNode === enclosingNode.body &&
     followingNode === enclosingNode.elseBranch
   ) {
-    addTrailingComment(precedingNode.variant, comment);
+    addTrailingComment(precedingNode, comment);
     return true;
   }
 
   if (followingNode.kind === NonterminalKind.IfStatement) {
-    if (followingNode.body.variant.kind === NonterminalKind.Block) {
-      addCollectionFirstComment(followingNode.body.variant.statements, comment);
+    if (followingNode.body.kind === NonterminalKind.Block) {
+      addCollectionFirstComment(followingNode.body.statements, comment);
     } else {
-      addLeadingComment(followingNode.body.variant, comment);
+      addLeadingComment(followingNode.body, comment);
     }
     return true;
   }
@@ -56,8 +56,8 @@ export default function handleIfStatementComments({
   // before the consequent without brackets on, such as
   // if (a) /* comment */ true
   if (enclosingNode.body === followingNode) {
-    if (followingNode.variant.kind === NonterminalKind.Block) {
-      addCollectionFirstComment(followingNode.variant.statements, comment);
+    if (followingNode.kind === NonterminalKind.Block) {
+      addCollectionFirstComment(followingNode.statements, comment);
     } else {
       addLeadingComment(followingNode, comment);
     }

src/slang-comments/handlers/handle-while-statement-comments.ts

@@ -36,8 +36,8 @@ export default function handleWhileStatementComments({
   }
 
   if (enclosingNode.body === followingNode) {
-    if (followingNode.variant.kind === NonterminalKind.Block) {
-      addCollectionFirstComment(followingNode.variant.statements, comment);
+    if (followingNode.kind === NonterminalKind.Block) {
+      addCollectionFirstComment(followingNode.statements, comment);
     } else {
       addLeadingComment(followingNode, comment);
     }

src/slang-nodes/ContractDefinition.ts

@@ -41,12 +41,12 @@ export class ContractDefinition extends SlangNode {
     // Older versions of Solidity defined a constructor as a function having
     // the same name as the contract.
     if (!satisfies(options.compiler, '>=0.5.0')) {
-      for (const { variant } of this.members.items) {
+      for (const member of this.members.items) {
         if (
-          variant.kind === NonterminalKind.FunctionDefinition &&
-          variant.name.variant.value !== this.name.value
+          member.kind === NonterminalKind.FunctionDefinition &&
+          member.name.variant.value !== this.name.value
         ) {
-          variant.cleanModifierInvocationArguments();
+          member.cleanModifierInvocationArguments();
         }
       }
     }

src/slang-nodes/ContractMember.ts

@@ -1,6 +1,6 @@
 import * as ast from '@nomicfoundation/slang/ast';
 import { NonterminalKind } from '@nomicfoundation/slang/cst';
-import { PolymorphicNode } from './PolymorphicNode.js';
+import { SlangNode } from './SlangNode.js';
 import { UsingDirective } from './UsingDirective.js';
 import { FunctionDefinition } from './FunctionDefinition.js';
 import { ConstructorDefinition } from './ConstructorDefinition.js';
@@ -65,7 +65,7 @@ function createNonterminalVariant(
   return exhaustiveCheck;
 }
 
-export class ContractMember extends PolymorphicNode {
+export class ContractMember extends SlangNode {
   readonly kind = NonterminalKind.ContractMember;
 
   variant:

src/slang-nodes/ContractMembers.ts

@@ -2,6 +2,7 @@ import { NonterminalKind } from '@nomicfoundation/slang/cst';
 import { doc } from 'prettier';
 import { printSeparatedItem } from '../slang-printers/print-separated-item.js';
 import { printPreservingEmptyLines } from '../slang-printers/print-preserving-empty-lines.js';
+import { extractVariant } from '../slang-utils/extract-variant.js';
 import { SlangNode } from './SlangNode.js';
 import { ContractMember } from './ContractMember.js';
 
@@ -15,12 +16,14 @@ const { hardline } = doc.builders;
 export class ContractMembers extends SlangNode {
   readonly kind = NonterminalKind.ContractMembers;
 
-  items: ContractMember[];
+  items: ContractMember['variant'][];
 
   constructor(ast: ast.ContractMembers, options: ParserOptions<AstNode>) {
     super(ast, true);
 
-    this.items = ast.items.map((item) => new ContractMember(item, options));
+    this.items = ast.items.map((item) =>
+      extractVariant(new ContractMember(item, options))
+    );
   }
 
   print(

src/slang-nodes/DoWhileStatement.ts

@@ -2,6 +2,7 @@ import { NonterminalKind } from '@nomicfoundation/slang/cst';
 import { doc } from 'prettier';
 import { printSeparatedItem } from '../slang-printers/print-separated-item.js';
 import { printVariant } from '../slang-printers/print-variant.js';
+import { extractVariant } from '../slang-utils/extract-variant.js';
 import { SlangNode } from './SlangNode.js';
 import { Statement } from './Statement.js';
 import { Expression } from './Expression.js';
@@ -16,24 +17,24 @@ const { line } = doc.builders;
 export class DoWhileStatement extends SlangNode {
   readonly kind = NonterminalKind.DoWhileStatement;
 
-  body: Statement;
+  body: Statement['variant'];
 
   condition: Expression;
 
   constructor(ast: ast.DoWhileStatement, options: ParserOptions<AstNode>) {
     super(ast);
 
-    this.body = new Statement(ast.body, options);
+    this.body = extractVariant(new Statement(ast.body, options));
     this.condition = new Expression(ast.condition, options);
 
     this.updateMetadata(this.body, this.condition);
   }
 
   print(path: AstPath<DoWhileStatement>, print: PrintFunction): Doc {
-    const body = path.call(printVariant(print), 'body');
+    const body = path.call(print, 'body');
     return [
       'do',
-      this.body.variant.kind === NonterminalKind.Block
+      this.body.kind === NonterminalKind.Block
         ? [' ', body, ' ']
         : printSeparatedItem(body, { firstSeparator: line }),
       'while (',

src/slang-nodes/ElseBranch.ts

@@ -1,7 +1,7 @@
 import { NonterminalKind } from '@nomicfoundation/slang/cst';
 import { createKindCheckFunction } from '../slang-utils/create-kind-check-function.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 { Statement } from './Statement.js';
 
@@ -18,12 +18,12 @@ const isIfStatementOrBlock = createKindCheckFunction([
 export class ElseBranch extends SlangNode {
   readonly kind = NonterminalKind.ElseBranch;
 
-  body: Statement;
+  body: Statement['variant'];
 
   constructor(ast: ast.ElseBranch, options: ParserOptions<AstNode>) {
     super(ast);
 
-    this.body = new Statement(ast.body, options);
+    this.body = extractVariant(new Statement(ast.body, options));
 
     this.updateMetadata(this.body);
   }
@@ -32,8 +32,8 @@ export class ElseBranch extends SlangNode {
     return [
       'else',
       printIndentedGroupOrSpacedDocument(
-        path.call(printVariant(print), 'body'),
-        !isIfStatementOrBlock(this.body.variant)
+        path.call(print, 'body'),
+        !isIfStatementOrBlock(this.body)
       )
     ];
   }

src/slang-nodes/ForStatement.ts

@@ -3,6 +3,7 @@ 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';
 import { ForStatementCondition } from './ForStatementCondition.js';
@@ -25,7 +26,7 @@ export class ForStatement extends SlangNode {
 
   iterator?: Expression;
 
-  body: Statement;
+  body: Statement['variant'];
 
   constructor(ast: ast.ForStatement, options: ParserOptions<AstNode>) {
     super(ast);
@@ -38,7 +39,7 @@ export class ForStatement extends SlangNode {
     if (ast.iterator) {
       this.iterator = new Expression(ast.iterator, options);
     }
-    this.body = new Statement(ast.body, options);
+    this.body = extractVariant(new Statement(ast.body, options));
 
     this.updateMetadata(
       this.initialization,
@@ -63,8 +64,8 @@ export class ForStatement extends SlangNode {
       }),
       ')',
       printIndentedGroupOrSpacedDocument(
-        path.call(printVariant(print), 'body'),
-        this.body.variant.kind !== NonterminalKind.Block
+        path.call(print, 'body'),
+        this.body.kind !== NonterminalKind.Block
       )
     ];
   }

src/slang-nodes/IfStatement.ts

@@ -4,6 +4,7 @@ import { printSeparatedItem } from '../slang-printers/print-separated-item.js';
 import { printIndentedGroupOrSpacedDocument } from '../slang-printers/print-indented-group-or-spaced-document.js';
 import { isBlockComment } from '../slang-utils/is-comment.js';
 import { printVariant } from '../slang-printers/print-variant.js';
+import { extractVariant } from '../slang-utils/extract-variant.js';
 import { SlangNode } from './SlangNode.js';
 import { Expression } from './Expression.js';
 import { Statement } from './Statement.js';
@@ -21,15 +22,15 @@ export class IfStatement extends SlangNode {
 
   condition: Expression;
 
-  body: Statement;
+  body: Statement['variant'];
 
   elseBranch?: ElseBranch;
 
   constructor(ast: ast.IfStatement, options: ParserOptions<AstNode>) {
     super(ast);
 
     this.condition = new Expression(ast.condition, options);
-    this.body = new Statement(ast.body, options);
+    this.body = extractVariant(new Statement(ast.body, options));
     if (ast.elseBranch) {
       this.elseBranch = new ElseBranch(ast.elseBranch, options);
     }
@@ -38,13 +39,13 @@ export class IfStatement extends SlangNode {
   }
 
   print(path: AstPath<IfStatement>, print: PrintFunction): Doc {
-    const { kind: bodyKind, comments: bodyComments } = this.body.variant;
+    const { kind: bodyKind, comments: bodyComments } = this.body;
     return [
       'if (',
       printSeparatedItem(path.call(printVariant(print), 'condition')),
       ')',
       printIndentedGroupOrSpacedDocument(
-        path.call(printVariant(print), 'body'),
+        path.call(print, 'body'),
         bodyKind !== NonterminalKind.Block,
         // `if` within `if`
         { shouldBreak: bodyKind === NonterminalKind.IfStatement }