removing ContractSpecifier from the AST tree

Author: Janther

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/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-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,7 @@ 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';
 
 function hasNodeIgnoreComment({ comments }: StrictAstNode): boolean {
   // Prettier sets SourceUnit's comments to undefined after assigning comments
@@ -82,6 +83,7 @@ function genericPrint(
       | SourceUnitMember
       | Statement
       | YulStatement
+      | ContractSpecifier
     >
   >,
   options: ParserOptions<AstNode>,