Cleaning up variant logic (#1209)

* exiting early when the variant is a string or an Identifier to avoid an extra check at the end

* adding a printVariant helper

* storing `variant` before switch statement

* helper function for variant creation

* Avoid unnecessary type casts

* updating all `createNonterminalVariant` based on Franco's feedback

---------

Co-authored-by: Franco Victorio <[email protected]>

Author: Janther

src/slang-nodes/ArgumentsDeclaration.ts

@@ -1,13 +1,27 @@
+import * as ast from '@nomicfoundation/slang/ast';
 import { NonterminalKind } from '@nomicfoundation/slang/cst';
 import { SlangNode } from './SlangNode.js';
 import { PositionalArgumentsDeclaration } from './PositionalArgumentsDeclaration.js';
 import { NamedArgumentsDeclaration } from './NamedArgumentsDeclaration.js';
 
-import type * as ast from '@nomicfoundation/slang/ast';
 import type { AstPath, Doc, ParserOptions } from 'prettier';
 import type { AstNode } from './types.d.ts';
 import type { PrintFunction } from '../types.d.ts';
 
+function createNonterminalVariant(
+  variant: ast.ArgumentsDeclaration['variant'],
+  options: ParserOptions<AstNode>
+): ArgumentsDeclaration['variant'] {
+  if (variant instanceof ast.PositionalArgumentsDeclaration) {
+    return new PositionalArgumentsDeclaration(variant, options);
+  }
+  if (variant instanceof ast.NamedArgumentsDeclaration) {
+    return new NamedArgumentsDeclaration(variant, options);
+  }
+  const exhaustiveCheck: never = variant;
+  return exhaustiveCheck;
+}
+
 export class ArgumentsDeclaration extends SlangNode {
   readonly kind = NonterminalKind.ArgumentsDeclaration;
 
@@ -16,22 +30,7 @@ export class ArgumentsDeclaration extends SlangNode {
   constructor(ast: ast.ArgumentsDeclaration, options: ParserOptions<AstNode>) {
     super(ast);
 
-    switch (ast.variant.cst.kind) {
-      case NonterminalKind.PositionalArgumentsDeclaration:
-        this.variant = new PositionalArgumentsDeclaration(
-          ast.variant as ast.PositionalArgumentsDeclaration,
-          options
-        );
-        break;
-      case NonterminalKind.NamedArgumentsDeclaration:
-        this.variant = new NamedArgumentsDeclaration(
-          ast.variant as ast.NamedArgumentsDeclaration,
-          options
-        );
-        break;
-      default:
-        throw new Error(`Unexpected variant: ${ast.variant.cst.kind}`);
-    }
+    this.variant = createNonterminalVariant(ast.variant, options);
 
     this.updateMetadata(this.variant);
   }

src/slang-nodes/ConstructorAttribute.ts

@@ -1,4 +1,5 @@
 import { NonterminalKind, TerminalNode } from '@nomicfoundation/slang/cst';
+import { printVariant } from '../slang-printers/print-variant.js';
 import { SlangNode } from './SlangNode.js';
 import { ModifierInvocation } from './ModifierInvocation.js';
 
@@ -15,17 +16,17 @@ export class ConstructorAttribute extends SlangNode {
   constructor(ast: ast.ConstructorAttribute, options: ParserOptions<AstNode>) {
     super(ast);
 
-    this.variant =
-      ast.variant instanceof TerminalNode
-        ? ast.variant.unparse()
-        : new ModifierInvocation(ast.variant, options);
+    const variant = ast.variant;
+    if (variant instanceof TerminalNode) {
+      this.variant = variant.unparse();
+      return;
+    }
+    this.variant = new ModifierInvocation(variant, options);
 
-    if (typeof this.variant !== 'string') this.updateMetadata(this.variant);
+    this.updateMetadata(this.variant);
   }
 
   print(path: AstPath<ConstructorAttribute>, print: PrintFunction): Doc {
-    return typeof this.variant === 'string'
-      ? this.variant
-      : path.call(print, 'variant');
+    return printVariant(this, path, print);
   }
 }

src/slang-nodes/ContractMember.ts

@@ -1,3 +1,4 @@
+import * as ast from '@nomicfoundation/slang/ast';
 import { NonterminalKind } from '@nomicfoundation/slang/cst';
 import { SlangNode } from './SlangNode.js';
 import { UsingDirective } from './UsingDirective.js';
@@ -14,11 +15,57 @@ import { StateVariableDefinition } from './StateVariableDefinition.js';
 import { ErrorDefinition } from './ErrorDefinition.js';
 import { UserDefinedValueTypeDefinition } from './UserDefinedValueTypeDefinition.js';
 
-import type * as ast from '@nomicfoundation/slang/ast';
 import type { AstPath, Doc, ParserOptions } from 'prettier';
 import type { AstNode } from './types.d.ts';
 import type { PrintFunction } from '../types.d.ts';
 
+function createNonterminalVariant(
+  variant: ast.ContractMember['variant'],
+  options: ParserOptions<AstNode>
+): ContractMember['variant'] {
+  if (variant instanceof ast.UsingDirective) {
+    return new UsingDirective(variant, options);
+  }
+  if (variant instanceof ast.FunctionDefinition) {
+    return new FunctionDefinition(variant, options);
+  }
+  if (variant instanceof ast.ConstructorDefinition) {
+    return new ConstructorDefinition(variant, options);
+  }
+  if (variant instanceof ast.ReceiveFunctionDefinition) {
+    return new ReceiveFunctionDefinition(variant, options);
+  }
+  if (variant instanceof ast.FallbackFunctionDefinition) {
+    return new FallbackFunctionDefinition(variant, options);
+  }
+  if (variant instanceof ast.UnnamedFunctionDefinition) {
+    return new UnnamedFunctionDefinition(variant, options);
+  }
+  if (variant instanceof ast.ModifierDefinition) {
+    return new ModifierDefinition(variant, options);
+  }
+  if (variant instanceof ast.StructDefinition) {
+    return new StructDefinition(variant, options);
+  }
+  if (variant instanceof ast.EnumDefinition) {
+    return new EnumDefinition(variant);
+  }
+  if (variant instanceof ast.EventDefinition) {
+    return new EventDefinition(variant, options);
+  }
+  if (variant instanceof ast.StateVariableDefinition) {
+    return new StateVariableDefinition(variant, options);
+  }
+  if (variant instanceof ast.ErrorDefinition) {
+    return new ErrorDefinition(variant, options);
+  }
+  if (variant instanceof ast.UserDefinedValueTypeDefinition) {
+    return new UserDefinedValueTypeDefinition(variant);
+  }
+  const exhaustiveCheck: never = variant;
+  return exhaustiveCheck;
+}
+
 export class ContractMember extends SlangNode {
   readonly kind = NonterminalKind.ContractMember;
 
@@ -40,84 +87,7 @@ export class ContractMember extends SlangNode {
   constructor(ast: ast.ContractMember, options: ParserOptions<AstNode>) {
     super(ast);
 
-    switch (ast.variant.cst.kind) {
-      case NonterminalKind.UsingDirective:
-        this.variant = new UsingDirective(
-          ast.variant as ast.UsingDirective,
-          options
-        );
-        break;
-      case NonterminalKind.FunctionDefinition:
-        this.variant = new FunctionDefinition(
-          ast.variant as ast.FunctionDefinition,
-          options
-        );
-        break;
-      case NonterminalKind.ConstructorDefinition:
-        this.variant = new ConstructorDefinition(
-          ast.variant as ast.ConstructorDefinition,
-          options
-        );
-        break;
-      case NonterminalKind.ReceiveFunctionDefinition:
-        this.variant = new ReceiveFunctionDefinition(
-          ast.variant as ast.ReceiveFunctionDefinition,
-          options
-        );
-        break;
-      case NonterminalKind.FallbackFunctionDefinition:
-        this.variant = new FallbackFunctionDefinition(
-          ast.variant as ast.FallbackFunctionDefinition,
-          options
-        );
-        break;
-      case NonterminalKind.UnnamedFunctionDefinition:
-        this.variant = new UnnamedFunctionDefinition(
-          ast.variant as ast.UnnamedFunctionDefinition,
-          options
-        );
-        break;
-      case NonterminalKind.ModifierDefinition:
-        this.variant = new ModifierDefinition(
-          ast.variant as ast.ModifierDefinition,
-          options
-        );
-        break;
-      case NonterminalKind.StructDefinition:
-        this.variant = new StructDefinition(
-          ast.variant as ast.StructDefinition,
-          options
-        );
-        break;
-      case NonterminalKind.EnumDefinition:
-        this.variant = new EnumDefinition(ast.variant as ast.EnumDefinition);
-        break;
-      case NonterminalKind.EventDefinition:
-        this.variant = new EventDefinition(
-          ast.variant as ast.EventDefinition,
-          options
-        );
-        break;
-      case NonterminalKind.StateVariableDefinition:
-        this.variant = new StateVariableDefinition(
-          ast.variant as ast.StateVariableDefinition,
-          options
-        );
-        break;
-      case NonterminalKind.ErrorDefinition:
-        this.variant = new ErrorDefinition(
-          ast.variant as ast.ErrorDefinition,
-          options
-        );
-        break;
-      case NonterminalKind.UserDefinedValueTypeDefinition:
-        this.variant = new UserDefinedValueTypeDefinition(
-          ast.variant as ast.UserDefinedValueTypeDefinition
-        );
-        break;
-      default:
-        throw new Error(`Unexpected variant: ${ast.variant.cst.kind}`);
-    }
+    this.variant = createNonterminalVariant(ast.variant, options);
 
     this.updateMetadata(this.variant);
   }

src/slang-nodes/ContractSpecifier.ts

@@ -1,13 +1,27 @@
+import * as ast from '@nomicfoundation/slang/ast';
 import { NonterminalKind } from '@nomicfoundation/slang/cst';
 import { SlangNode } from './SlangNode.js';
 import { InheritanceSpecifier } from './InheritanceSpecifier.js';
 import { StorageLayoutSpecifier } from './StorageLayoutSpecifier.js';
 
-import type * as ast from '@nomicfoundation/slang/ast';
 import type { AstPath, Doc, ParserOptions } from 'prettier';
 import type { AstNode } from './types.d.ts';
 import type { PrintFunction } from '../types.d.ts';
 
+function createNonterminalVariant(
+  variant: ast.ContractSpecifier['variant'],
+  options: ParserOptions<AstNode>
+): ContractSpecifier['variant'] {
+  if (variant instanceof ast.InheritanceSpecifier) {
+    return new InheritanceSpecifier(variant, options);
+  }
+  if (variant instanceof ast.StorageLayoutSpecifier) {
+    return new StorageLayoutSpecifier(variant, options);
+  }
+  const exhaustiveCheck: never = variant;
+  return exhaustiveCheck;
+}
+
 export class ContractSpecifier extends SlangNode {
   readonly kind = NonterminalKind.ContractSpecifier;
 
@@ -16,22 +30,8 @@ export class ContractSpecifier extends SlangNode {
   constructor(ast: ast.ContractSpecifier, options: ParserOptions<AstNode>) {
     super(ast);
 
-    switch (ast.variant.cst.kind) {
-      case NonterminalKind.InheritanceSpecifier:
-        this.variant = new InheritanceSpecifier(
-          ast.variant as ast.InheritanceSpecifier,
-          options
-        );
-        break;
-      case NonterminalKind.StorageLayoutSpecifier:
-        this.variant = new StorageLayoutSpecifier(
-          ast.variant as ast.StorageLayoutSpecifier,
-          options
-        );
-        break;
-      default:
-        throw new Error(`Unexpected variant: ${ast.variant.cst.kind}`);
-    }
+    this.variant = createNonterminalVariant(ast.variant, options);
+
     this.updateMetadata(this.variant);
   }
 

src/slang-nodes/ElementaryType.ts

@@ -1,4 +1,5 @@
 import { NonterminalKind, TerminalNode } from '@nomicfoundation/slang/cst';
+import { printVariant } from '../slang-printers/print-variant.js';
 import { SlangNode } from './SlangNode.js';
 import { AddressType } from './AddressType.js';
 
@@ -14,17 +15,17 @@ export class ElementaryType extends SlangNode {
   constructor(ast: ast.ElementaryType) {
     super(ast);
 
-    this.variant =
-      ast.variant instanceof TerminalNode
-        ? ast.variant.unparse()
-        : new AddressType(ast.variant);
+    const variant = ast.variant;
+    if (variant instanceof TerminalNode) {
+      this.variant = variant.unparse();
+      return;
+    }
+    this.variant = new AddressType(variant);
 
-    if (typeof this.variant !== 'string') this.updateMetadata(this.variant);
+    this.updateMetadata(this.variant);
   }
 
   print(path: AstPath<ElementaryType>, print: PrintFunction): Doc {
-    return typeof this.variant === 'string'
-      ? this.variant
-      : path.call(print, 'variant');
+    return printVariant(this, path, print);
   }
 }

src/slang-nodes/ExperimentalFeature.ts

@@ -1,8 +1,4 @@
-import {
-  NonterminalKind,
-  TerminalKind,
-  TerminalNode
-} from '@nomicfoundation/slang/cst';
+import { NonterminalKind, TerminalNode } from '@nomicfoundation/slang/cst';
 import { SlangNode } from './SlangNode.js';
 import { StringLiteral } from './StringLiteral.js';
 import { Identifier } from './Identifier.js';
@@ -20,13 +16,14 @@ export class ExperimentalFeature extends SlangNode {
   constructor(ast: ast.ExperimentalFeature, options: ParserOptions<AstNode>) {
     super(ast);
 
-    this.variant =
-      ast.variant instanceof TerminalNode
-        ? new Identifier(ast.variant)
-        : new StringLiteral(ast.variant, options);
+    const variant = ast.variant;
+    if (variant instanceof TerminalNode) {
+      this.variant = new Identifier(variant);
+      return;
+    }
+    this.variant = new StringLiteral(variant, options);
 
-    if (this.variant.kind !== TerminalKind.Identifier)
-      this.updateMetadata(this.variant);
+    this.updateMetadata(this.variant);
   }
 
   print(path: AstPath<ExperimentalFeature>, print: PrintFunction): Doc {