removing ElementaryType from the AST tree

Author: Janther

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/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/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/slangPrinter.ts

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