removing ContractSpecifier, ElementaryType, and ExperimentalFeature from the AST tree

src/slang-comments/handlers/handle-contract-definition-comments.ts

@@ -48,7 +48,7 @@ export default function handleContractDefinitionComments({
         return true;
       }
       const lastContractSpecifier =
-        precedingNode.items[precedingNode.items.length - 1].variant;
+        precedingNode.items[precedingNode.items.length - 1];
       // If the last ContractSpecifier's an InheritanceSpecifier, the comment
       // is appended to the last InheritanceType.
       if (lastContractSpecifier.kind === NonterminalKind.InheritanceSpecifier) {

src/slang-comments/handlers/handle-contract-specifiers-comments.ts

@@ -15,16 +15,13 @@ export default function handleContractSpecifiersComments({
     return false;
   }
 
-  if (
-    precedingNode &&
-    precedingNode.kind === NonterminalKind.ContractSpecifier
-  ) {
-    if (precedingNode.variant.kind === NonterminalKind.InheritanceSpecifier) {
-      addCollectionLastComment(precedingNode.variant.types, comment);
+  if (precedingNode) {
+    if (precedingNode.kind === NonterminalKind.InheritanceSpecifier) {
+      addCollectionLastComment(precedingNode.types, comment);
       return true;
     }
-    if (precedingNode.variant.kind === NonterminalKind.StorageLayoutSpecifier) {
-      addTrailingComment(precedingNode.variant.expression, comment);
+    if (precedingNode.kind === NonterminalKind.StorageLayoutSpecifier) {
+      addTrailingComment(precedingNode.expression, comment);
       return true;
     }
   }

src/slang-nodes/ContractSpecifier.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 { InheritanceSpecifier } from './InheritanceSpecifier.js';
 import { StorageLayoutSpecifier } from './StorageLayoutSpecifier.js';
 
@@ -21,7 +21,7 @@ function createNonterminalVariant(
   return exhaustiveCheck;
 }
 
-export class ContractSpecifier extends PolymorphicNode {
+export class ContractSpecifier extends SlangNode {
   readonly kind = NonterminalKind.ContractSpecifier;
 
   variant: InheritanceSpecifier | StorageLayoutSpecifier;

src/slang-nodes/ContractSpecifiers.ts

@@ -2,7 +2,7 @@ import { NonterminalKind } from '@nomicfoundation/slang/cst';
 import { doc } from 'prettier';
 import { sortContractSpecifiers } from '../slang-utils/sort-contract-specifiers.js';
 import { printSeparatedList } from '../slang-printers/print-separated-list.js';
-import { printVariant } from '../slang-printers/print-variant.js';
+import { extractVariant } from '../slang-utils/extract-variant.js';
 import { SlangNode } from './SlangNode.js';
 import { ContractSpecifier } from './ContractSpecifier.js';
 
@@ -16,18 +16,20 @@ const { group, ifBreak, line, softline } = doc.builders;
 export class ContractSpecifiers extends SlangNode {
   readonly kind = NonterminalKind.ContractSpecifiers;
 
-  items: ContractSpecifier[];
+  items: ContractSpecifier['variant'][];
 
   constructor(ast: ast.ContractSpecifiers, options: ParserOptions<AstNode>) {
     super(ast, true);
 
-    this.items = ast.items.map((item) => new ContractSpecifier(item, options));
+    this.items = ast.items.map((item) =>
+      extractVariant(new ContractSpecifier(item, options))
+    );
 
     this.items.sort(sortContractSpecifiers);
   }
 
   print(path: AstPath<ContractSpecifiers>, print: PrintFunction): Doc {
-    const [specifier1, specifier2] = path.map(printVariant(print), 'items');
+    const [specifier1, specifier2] = path.map(print, 'items');
 
     if (specifier1 === undefined) return '';
 

src/slang-nodes/ElementaryType.ts

@@ -2,13 +2,13 @@ import {
   NonterminalKind,
   TerminalNode as SlangTerminalNode
 } from '@nomicfoundation/slang/cst';
-import { PolymorphicNode } from './PolymorphicNode.js';
+import { SlangNode } from './SlangNode.js';
 import { AddressType } from './AddressType.js';
 import { TerminalNode } from './TerminalNode.js';
 
 import type * as ast from '@nomicfoundation/slang/ast';
 
-export class ElementaryType extends PolymorphicNode {
+export class ElementaryType extends SlangNode {
   readonly kind = NonterminalKind.ElementaryType;
 
   variant: AddressType | TerminalNode;

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

src/slang-nodes/ExperimentalPragma.ts

@@ -1,5 +1,5 @@
 import { NonterminalKind } from '@nomicfoundation/slang/cst';
-import { printVariant } from '../slang-printers/print-variant.js';
+import { extractVariant } from '../slang-utils/extract-variant.js';
 import { SlangNode } from './SlangNode.js';
 import { ExperimentalFeature } from './ExperimentalFeature.js';
 
@@ -11,17 +11,19 @@ import type { AstNode } from './types.d.ts';
 export class ExperimentalPragma extends SlangNode {
   readonly kind = NonterminalKind.ExperimentalPragma;
 
-  feature: ExperimentalFeature;
+  feature: ExperimentalFeature['variant'];
 
   constructor(ast: ast.ExperimentalPragma, options: ParserOptions<AstNode>) {
     super(ast);
 
-    this.feature = new ExperimentalFeature(ast.feature, options);
+    this.feature = extractVariant(
+      new ExperimentalFeature(ast.feature, options)
+    );
 
     this.updateMetadata(this.feature);
   }
 
   print(path: AstPath<ExperimentalPragma>, print: PrintFunction): Doc {
-    return ['experimental ', path.call(printVariant(print), 'feature')];
+    return ['experimental ', path.call(print, 'feature')];
   }
 }

src/slang-nodes/Expression.ts

@@ -3,6 +3,7 @@ import {
   NonterminalKind,
   TerminalNode as SlangTerminalNode
 } from '@nomicfoundation/slang/cst';
+import { extractVariant } from '../slang-utils/extract-variant.js';
 import { PolymorphicNode } from './PolymorphicNode.js';
 import { AssignmentExpression } from './AssignmentExpression.js';
 import { ConditionalExpression } from './ConditionalExpression.js';
@@ -39,7 +40,7 @@ import type { AstNode } from './types.d.ts';
 function createNonterminalVariant(
   variant: Exclude<ast.Expression['variant'], SlangTerminalNode>,
   options: ParserOptions<AstNode>
-): Exclude<Expression['variant'], TerminalNode> {
+): Expression['variant'] {
   if (variant instanceof ast.AssignmentExpression) {
     return new AssignmentExpression(variant, options);
   }
@@ -119,7 +120,7 @@ function createNonterminalVariant(
     return new StringExpression(variant, options);
   }
   if (variant instanceof ast.ElementaryType) {
-    return new ElementaryType(variant);
+    return extractVariant(new ElementaryType(variant));
   }
   const exhaustiveCheck: never = variant;
   return exhaustiveCheck;
@@ -155,7 +156,7 @@ export class Expression extends PolymorphicNode {
     | HexNumberExpression
     | DecimalNumberExpression
     | StringExpression
-    | ElementaryType
+    | ElementaryType['variant']
     | TerminalNode;
 
   constructor(ast: ast.Expression, options: ParserOptions<AstNode>) {

src/slang-nodes/MappingKeyType.ts

@@ -1,5 +1,6 @@
 import * as ast from '@nomicfoundation/slang/ast';
 import { NonterminalKind } from '@nomicfoundation/slang/cst';
+import { extractVariant } from '../slang-utils/extract-variant.js';
 import { PolymorphicNode } from './PolymorphicNode.js';
 import { ElementaryType } from './ElementaryType.js';
 import { IdentifierPath } from './IdentifierPath.js';
@@ -8,7 +9,7 @@ function createNonterminalVariant(
   variant: ast.MappingKeyType['variant']
 ): MappingKeyType['variant'] {
   if (variant instanceof ast.ElementaryType) {
-    return new ElementaryType(variant);
+    return extractVariant(new ElementaryType(variant));
   }
   if (variant instanceof ast.IdentifierPath) {
     return new IdentifierPath(variant);
@@ -20,7 +21,7 @@ function createNonterminalVariant(
 export class MappingKeyType extends PolymorphicNode {
   readonly kind = NonterminalKind.MappingKeyType;
 
-  variant: ElementaryType | IdentifierPath;
+  variant: ElementaryType['variant'] | IdentifierPath;
 
   constructor(ast: ast.MappingKeyType) {
     super(ast);

src/slang-nodes/SlangNode.ts

@@ -1,7 +1,7 @@
 import {
+  TerminalNode as SlangTerminalNode,
   TerminalKind,
-  TerminalKindExtensions,
-  TerminalNode
+  TerminalKindExtensions
 } from '@nomicfoundation/slang/cst';
 import { MultiLineComment } from '../slang-nodes/MultiLineComment.js';
 import { MultiLineNatSpecComment } from '../slang-nodes/MultiLineNatSpecComment.js';
@@ -10,6 +10,7 @@ import { SingleLineNatSpecComment } from '../slang-nodes/SingleLineNatSpecCommen
 
 import type { Comment, StrictAstNode } from '../slang-nodes/types.d.ts';
 import type { AstLocation, SlangAstNode } from '../types.d.ts';
+import type { TerminalNode } from './TerminalNode.js';
 
 const offsets = new Map<number, number>();
 const comments: Comment[] = [];
@@ -42,10 +43,10 @@ export class SlangNode {
   loc: AstLocation;
 
   constructor(
-    ast: SlangAstNode | TerminalNode,
+    ast: SlangAstNode | SlangTerminalNode,
     enclosePeripheralComments = false
   ) {
-    if (ast instanceof TerminalNode) {
+    if (ast instanceof SlangTerminalNode) {
       const offset = offsets.get(ast.id) || 0;
       this.loc = {
         start: offset,
@@ -121,7 +122,9 @@ export class SlangNode {
     };
   }
 
-  updateMetadata(...childNodes: (StrictAstNode | undefined)[]): void {
+  updateMetadata(
+    ...childNodes: (StrictAstNode | TerminalNode | undefined)[]
+  ): void {
     const { loc } = this;
     // calculate correct loc object
     if (loc.leadingOffset === 0) {

src/slang-nodes/StorageLayoutSpecifier.ts

@@ -35,9 +35,7 @@ export class StorageLayoutSpecifier extends SlangNode {
         firstSeparator: line,
         // If this is the second ContractSpecifier we have to delegate printing
         // the line to the ContractSpecifiers node.
-        lastSeparator: path.callParent((parentPath) => parentPath.isFirst)
-          ? line
-          : ''
+        lastSeparator: path.isFirst ? line : ''
       })
     ];
   }

src/slang-nodes/TypeName.ts

@@ -1,5 +1,6 @@
 import * as ast from '@nomicfoundation/slang/ast';
 import { NonterminalKind } from '@nomicfoundation/slang/cst';
+import { extractVariant } from '../slang-utils/extract-variant.js';
 import { PolymorphicNode } from './PolymorphicNode.js';
 import { ArrayTypeName } from './ArrayTypeName.js';
 import { FunctionType } from './FunctionType.js';
@@ -24,7 +25,7 @@ function createNonterminalVariant(
     return new MappingType(variant, options);
   }
   if (variant instanceof ast.ElementaryType) {
-    return new ElementaryType(variant);
+    return extractVariant(new ElementaryType(variant));
   }
   if (variant instanceof ast.IdentifierPath) {
     return new IdentifierPath(variant);
@@ -40,7 +41,7 @@ export class TypeName extends PolymorphicNode {
     | ArrayTypeName
     | FunctionType
     | MappingType
-    | ElementaryType
+    | ElementaryType['variant']
     | IdentifierPath;
 
   constructor(ast: ast.TypeName, options: ParserOptions<AstNode>) {

src/slang-nodes/UserDefinedValueTypeDefinition.ts

@@ -1,5 +1,5 @@
 import { NonterminalKind } from '@nomicfoundation/slang/cst';
-import { printVariant } from '../slang-printers/print-variant.js';
+import { extractVariant } from '../slang-utils/extract-variant.js';
 import { SlangNode } from './SlangNode.js';
 import { TerminalNode } from './TerminalNode.js';
 import { ElementaryType } from './ElementaryType.js';
@@ -13,13 +13,13 @@ export class UserDefinedValueTypeDefinition extends SlangNode {
 
   name: TerminalNode;
 
-  valueType: ElementaryType;
+  valueType: ElementaryType['variant'];
 
   constructor(ast: ast.UserDefinedValueTypeDefinition) {
     super(ast);
 
     this.name = new TerminalNode(ast.name);
-    this.valueType = new ElementaryType(ast.valueType);
+    this.valueType = extractVariant(new ElementaryType(ast.valueType));
 
     this.updateMetadata(this.valueType);
   }
@@ -32,7 +32,7 @@ export class UserDefinedValueTypeDefinition extends SlangNode {
       'type ',
       path.call(print, 'name'),
       ' is ',
-      path.call(printVariant(print), 'valueType'),
+      path.call(print, 'valueType'),
       ';'
     ];
   }

src/slang-utils/sort-contract-specifiers.ts

@@ -3,8 +3,8 @@ import { NonterminalKind } from '@nomicfoundation/slang/cst';
 import type { ContractSpecifier } from '../slang-nodes/ContractSpecifier.ts';
 
 export function sortContractSpecifiers(
-  { variant: { kind: aKind } }: ContractSpecifier,
-  { variant: { kind: bKind } }: ContractSpecifier
+  { kind: aKind }: ContractSpecifier['variant'],
+  { kind: bKind }: ContractSpecifier['variant']
 ): number {
   // OverrideSpecifiers before ModifierInvocation
   if (

src/slangPrinter.ts

@@ -17,6 +17,9 @@ import type { ContractMember } from './slang-nodes/ContractMember.js';
 import type { SourceUnitMember } from './slang-nodes/SourceUnitMember.js';
 import type { Statement } from './slang-nodes/Statement.js';
 import type { YulStatement } from './slang-nodes/YulStatement.js';
+import type { ContractSpecifier } from './slang-nodes/ContractSpecifier.js';
+import type { ElementaryType } from './slang-nodes/ElementaryType.js';
+import type { ExperimentalFeature } from './slang-nodes/ExperimentalFeature.js';
 
 function hasNodeIgnoreComment({ comments }: StrictAstNode): boolean {
   // Prettier sets SourceUnit's comments to undefined after assigning comments
@@ -82,6 +85,9 @@ function genericPrint(
       | SourceUnitMember
       | Statement
       | YulStatement
+      | ContractSpecifier
+      | ElementaryType
+      | ExperimentalFeature
     >
   >,
   options: ParserOptions<AstNode>,