Using PrintableNode as a more correct type for StrictAstNode

Janther · Janther · commit 6396280112bd · 2026-04-09T17:29:25.000-04:00
diff --git a/src/slang-comments/handlers/handle-storage-layout-specifier-comments.ts b/src/slang-comments/handlers/handle-storage-layout-specifier-comments.ts
@@ -1,10 +1,48 @@
 import { NonterminalKind } from '@nomicfoundation/slang/cst';
 import { util } from 'prettier';
+import { createKindCheckFunction } from '../../slang-utils/create-kind-check-function.js';
+import { TerminalNode } from '../../slang-nodes/TerminalNode.js';
 
+import type { Expression } from '../../slang-nodes/Expression.js';
+import type { PrintableNode } from '../../slang-nodes/types.d.ts';
 import type { HandlerParams } from './types.d.ts';
 
 const { addLeadingComment } = util;
 
+const isExpression = createKindCheckFunction([
+  NonterminalKind.AssignmentExpression,
+  NonterminalKind.ConditionalExpression,
+  NonterminalKind.OrExpression,
+  NonterminalKind.AndExpression,
+  NonterminalKind.EqualityExpression,
+  NonterminalKind.InequalityExpression,
+  NonterminalKind.BitwiseOrExpression,
+  NonterminalKind.BitwiseXorExpression,
+  NonterminalKind.BitwiseAndExpression,
+  NonterminalKind.ShiftExpression,
+  NonterminalKind.AdditiveExpression,
+  NonterminalKind.MultiplicativeExpression,
+  NonterminalKind.ExponentiationExpression,
+  NonterminalKind.PostfixExpression,
+  NonterminalKind.PrefixExpression,
+  NonterminalKind.FunctionCallExpression,
+  NonterminalKind.CallOptionsExpression,
+  NonterminalKind.MemberAccessExpression,
+  NonterminalKind.IndexAccessExpression,
+  NonterminalKind.NewExpression,
+  NonterminalKind.TupleExpression,
+  NonterminalKind.TypeExpression,
+  NonterminalKind.ArrayExpression,
+  NonterminalKind.HexNumberExpression,
+  NonterminalKind.DecimalNumberExpression,
+  NonterminalKind.StringLiteral,
+  NonterminalKind.StringLiterals,
+  NonterminalKind.HexStringLiteral,
+  NonterminalKind.HexStringLiterals,
+  NonterminalKind.UnicodeStringLiterals,
+  NonterminalKind.AddressType
+]) as (node: PrintableNode) => node is Expression['variant'];
+
 export default function handleStorageLayoutSpecifierComments({
   enclosingNode,
   followingNode,
@@ -14,7 +52,10 @@ export default function handleStorageLayoutSpecifierComments({
     return false;
   }
 
-  if (followingNode?.kind === NonterminalKind.Expression) {
+  if (
+    followingNode &&
+    (isExpression(followingNode) || followingNode instanceof TerminalNode)
+  ) {
     addLeadingComment(followingNode, comment);
     return true;
   }
diff --git a/src/slang-comments/handlers/types.d.ts b/src/slang-comments/handlers/types.d.ts
@@ -1,9 +1,9 @@
-import type { Comment, StrictAstNode } from '../../slang-nodes/types.d.ts';
+import type { Comment, PrintableNode } from '../../slang-nodes/types.d.ts';
 
 interface HandlerParams {
   text: string;
-  precedingNode?: StrictAstNode;
-  enclosingNode?: StrictAstNode;
-  followingNode?: StrictAstNode;
+  precedingNode?: PrintableNode;
+  enclosingNode?: PrintableNode;
+  followingNode?: PrintableNode;
   comment: Comment;
 }
diff --git a/src/slang-nodes/CommentNode.ts b/src/slang-nodes/CommentNode.ts
@@ -1,6 +1,6 @@
 import type { TerminalNode } from '@nomicfoundation/slang/cst';
 import type { Location } from '../types.d.ts';
-import type { StrictAstNode } from './types.d.ts';
+import type { PrintableNode } from './types.d.ts';
 
 export class CommentNode {
   loc: Location;
@@ -13,11 +13,11 @@ export class CommentNode {
 
   placement?: 'endOfLine' | 'ownLine' | 'remaining';
 
-  precedingNode?: StrictAstNode;
+  precedingNode?: PrintableNode;
 
-  enclosingNode?: StrictAstNode;
+  enclosingNode?: PrintableNode;
 
-  followingNode?: StrictAstNode;
+  followingNode?: PrintableNode;
 
   constructor(ast: TerminalNode, offset: number) {
     this.loc = {
diff --git a/src/slang-nodes/ConditionalExpression.ts b/src/slang-nodes/ConditionalExpression.ts
@@ -8,13 +8,13 @@ import { Expression } from './Expression.js';
 import type * as ast from '@nomicfoundation/slang/ast';
 import type { AstPath, Doc, ParserOptions } from 'prettier';
 import type { CollectedMetadata, PrintFunction } from '../types.d.ts';
-import type { AstNode, StrictAstNode } from './types.d.ts';
+import type { AstNode, PrintableNode } from './types.d.ts';
 
 const { group, hardline, ifBreak, indent, line } = doc.builders;
 
 function experimentalTernaries(
   node: ConditionalExpression,
-  path: AstPath<StrictAstNode>,
+  path: AstPath<PrintableNode>,
   print: PrintFunction,
   { useTabs, tabWidth }: ParserOptions<AstNode>
 ): Doc {
@@ -77,7 +77,7 @@ function experimentalTernaries(
 }
 
 function traditionalTernaries(
-  path: AstPath<StrictAstNode>,
+  path: AstPath<PrintableNode>,
   print: PrintFunction
 ): Doc {
   return group([
diff --git a/src/slang-nodes/MemberAccessExpression.ts b/src/slang-nodes/MemberAccessExpression.ts
@@ -11,15 +11,15 @@ import { TerminalNode } from './TerminalNode.js';
 import type * as ast from '@nomicfoundation/slang/ast';
 import type { AstPath, Doc, ParserOptions } from 'prettier';
 import type { CollectedMetadata, PrintFunction } from '../types.d.ts';
-import type { AstNode, ChainableExpression, StrictAstNode } from './types.d.ts';
+import type { AstNode, ChainableExpression, PrintableNode } from './types.d.ts';
 
 const { group, indent, label, softline } = doc.builders;
 
 const separatorLabel = Symbol('separator');
 
 function isEndOfChain(
   node: ChainableExpression,
-  path: AstPath<StrictAstNode>
+  path: AstPath<PrintableNode>
 ): boolean {
   for (let i = 1, current = node, parent; ; i++, current = parent) {
     parent = path.getNode(i)!;
diff --git a/src/slang-nodes/SlangNode.ts b/src/slang-nodes/SlangNode.ts
@@ -14,7 +14,7 @@ import type {
   CollectedMetadata,
   SlangAstNode
 } from '../types.d.ts';
-import type { Comment, StrictAstNode } from './types.d.ts';
+import type { Comment, PrintableNode } from './types.d.ts';
 import type { TerminalNode } from './TerminalNode.ts';
 
 function reversedIterator<T>(children: T[]): Iterable<T> {
@@ -121,7 +121,7 @@ export abstract class SlangNode {
   }
 
   updateMetadata(
-    ...childNodes: (StrictAstNode | TerminalNode | undefined)[]
+    ...childNodes: (PrintableNode | TerminalNode | undefined)[]
   ): void {
     const { loc } = this;
     // calculate correct loc object
diff --git a/src/slang-nodes/types.d.ts b/src/slang-nodes/types.d.ts
@@ -474,6 +474,8 @@ export type StrictPolymorphicNode = Extract<
   { variant: StrictAstNode | TerminalNode }
 >;
 
+export type PrintableNode = Exclude<StrictAstNode, StrictPolymorphicNode>;
+
 export type Collection = Extract<StrictAstNode, { items: unknown[] }>;
 
 export type StrictCollection = Extract<
diff --git a/src/slang-printers/create-binary-operation-printer.ts b/src/slang-printers/create-binary-operation-printer.ts
@@ -6,15 +6,15 @@ import type { AstPath, Doc, ParserOptions } from 'prettier';
 import type {
   AstNode,
   BinaryOperation,
-  StrictAstNode
+  PrintableNode
 } from '../slang-nodes/types.d.ts';
 import type { PrintFunction } from '../types.d.ts';
 
 const { group, line } = doc.builders;
 
 function rightOperandPrint(
   { operator, leftOperand }: BinaryOperation,
-  path: AstPath<StrictAstNode>,
+  path: AstPath<PrintableNode>,
   print: PrintFunction,
   options: ParserOptions<AstNode>
 ): Doc {
@@ -37,15 +37,15 @@ function rightOperandPrint(
 
 export const createBinaryOperationPrinter =
   (
-    groupRulesBuilder: (path: AstPath<StrictAstNode>) => (document: Doc) => Doc,
+    groupRulesBuilder: (path: AstPath<PrintableNode>) => (document: Doc) => Doc,
     indentRulesBuilder: (
       node: BinaryOperation,
-      path: AstPath<StrictAstNode>
+      path: AstPath<PrintableNode>
     ) => (document: Doc) => Doc
   ) =>
   (
     node: BinaryOperation,
-    path: AstPath<StrictAstNode>,
+    path: AstPath<PrintableNode>,
     print: PrintFunction,
     options: ParserOptions<AstNode>
   ): Doc => {
diff --git a/src/slang-printers/print-binary-operation.ts b/src/slang-printers/print-binary-operation.ts
@@ -8,15 +8,15 @@ import type { AstPath, Doc, ParserOptions } from 'prettier';
 import type {
   AstNode,
   BinaryOperation,
-  StrictAstNode
+  PrintableNode
 } from '../slang-nodes/types.d.ts';
 import type { PrintFunction } from '../types.d.ts';
 
 const { group, indent } = doc.builders;
 
 export const binaryGroupRulesBuilder =
   (shouldGroup: (node: BinaryOperation) => boolean) =>
-  (path: AstPath<StrictAstNode>) =>
+  (path: AstPath<PrintableNode>) =>
   (document: Doc): Doc => {
     const parent = path.parent!;
     if (!isBinaryOperation(parent)) return group(document);
@@ -31,8 +31,8 @@ const isStatementWithoutIndentedOperation = createKindCheckFunction([
 ]);
 
 export const shouldNotIndent = (
-  node: StrictAstNode,
-  path: AstPath<StrictAstNode>,
+  node: PrintableNode,
+  path: AstPath<PrintableNode>,
   index: number
 ): boolean =>
   isStatementWithoutIndentedOperation(node) ||
@@ -41,7 +41,7 @@ export const shouldNotIndent = (
 
 export const binaryIndentRulesBuilder =
   (shouldIndent: (node: BinaryOperation) => boolean) =>
-  (node: BinaryOperation, path: AstPath<StrictAstNode>) =>
+  (node: BinaryOperation, path: AstPath<PrintableNode>) =>
   (document: Doc): Doc => {
     for (let i = 1, current = node, parent; ; i++, current = parent) {
       parent = path.getNode(i)!;
@@ -57,7 +57,7 @@ export const printBinaryOperation = (
   shouldGroupAndIndent: (node: BinaryOperation) => boolean
 ): ((
   node: BinaryOperation,
-  path: AstPath<StrictAstNode>,
+  path: AstPath<PrintableNode>,
   print: PrintFunction,
   options: ParserOptions<AstNode>
 ) => Doc) =>
diff --git a/src/slang-printers/print-comments.ts b/src/slang-printers/print-comments.ts
@@ -7,7 +7,7 @@ import type { AstPath, Doc, ParserOptions } from 'prettier';
 import type {
   AstNode,
   Comment,
-  StrictAstNode
+  PrintableNode
 } from '../slang-nodes/types.d.ts';
 
 const { hardline, line } = doc.builders;
@@ -17,8 +17,8 @@ function isPrintable(comment: Comment): boolean {
 }
 
 export function printComments(
-  node: StrictAstNode,
-  path: AstPath<StrictAstNode>,
+  node: PrintableNode,
+  path: AstPath<PrintableNode>,
   options: ParserOptions<AstNode>
 ): Doc[] {
   const lastPrintableIndex = (node.comments ?? []).findLastIndex(isPrintable);
diff --git a/src/slang-printers/print-logical-operation.ts b/src/slang-printers/print-logical-operation.ts
@@ -8,14 +8,14 @@ import {
 } from './print-binary-operation.js';
 
 import type { AstPath, Doc } from 'prettier';
-import type { BinaryOperation, StrictAstNode } from '../slang-nodes/types.d.ts';
+import type { BinaryOperation, PrintableNode } from '../slang-nodes/types.d.ts';
 
 const { indent } = doc.builders;
 
 const logicalGroupRulesBuilder = binaryGroupRulesBuilder(() => false);
 
 const logicalIndentRulesBuilder =
-  (node: BinaryOperation, path: AstPath<StrictAstNode>) =>
+  (node: BinaryOperation, path: AstPath<PrintableNode>) =>
   (document: Doc): Doc => {
     for (let i = 1, current = node, parent; ; i++, current = parent) {
       parent = path.getNode(i)!;
diff --git a/src/slang-utils/create-kind-check-function.ts b/src/slang-utils/create-kind-check-function.ts
@@ -3,11 +3,11 @@ import type {
   NonterminalKind,
   TerminalKind
 } from '@nomicfoundation/slang/cst';
-import type { StrictAstNode } from '../slang-nodes/types.d.ts';
+import type { PrintableNode } from '../slang-nodes/types.d.ts';
 
 export function createKindCheckFunction(
   kindsArray: (keyof typeof TerminalKind | keyof typeof NonterminalKind)[]
-): (node: StrictAstNode | Node) => boolean {
+): (node: PrintableNode | Node) => boolean {
   const kinds = new Set(kindsArray);
-  return (node: StrictAstNode | Node): boolean => kinds.has(node.kind);
+  return (node: PrintableNode | Node): boolean => kinds.has(node.kind);
 }
diff --git a/src/slang-utils/is-binary-operation.ts b/src/slang-utils/is-binary-operation.ts
@@ -1,7 +1,7 @@
 import { NonterminalKind } from '@nomicfoundation/slang/cst';
 import { createKindCheckFunction } from './create-kind-check-function.js';
 
-import type { BinaryOperation, StrictAstNode } from '../slang-nodes/types.d.ts';
+import type { BinaryOperation, PrintableNode } from '../slang-nodes/types.d.ts';
 
 export const isBinaryOperation = createKindCheckFunction([
   NonterminalKind.AdditiveExpression,
@@ -16,4 +16,4 @@ export const isBinaryOperation = createKindCheckFunction([
   NonterminalKind.AndExpression,
   NonterminalKind.OrExpression,
   NonterminalKind.ShiftExpression
-]) as (node: StrictAstNode) => node is BinaryOperation;
+]) as (node: PrintableNode) => node is BinaryOperation;
diff --git a/src/slang-utils/is-chainable-expression.ts b/src/slang-utils/is-chainable-expression.ts
@@ -3,11 +3,11 @@ import { createKindCheckFunction } from './create-kind-check-function.js';
 
 import type {
   ChainableExpression,
-  StrictAstNode
+  PrintableNode
 } from '../slang-nodes/types.d.ts';
 
 export const isChainableExpression = createKindCheckFunction([
   NonterminalKind.FunctionCallExpression,
   NonterminalKind.IndexAccessExpression,
   NonterminalKind.MemberAccessExpression
-]) as (node: StrictAstNode) => node is ChainableExpression;
+]) as (node: PrintableNode) => node is ChainableExpression;
diff --git a/src/slangPrinter.ts b/src/slangPrinter.ts
@@ -2,14 +2,10 @@ import { isBlockComment } from './slang-utils/is-comment.js';
 import { locEnd, locStart } from './slang-utils/loc.js';
 
 import type { AstPath, Doc, ParserOptions } from 'prettier';
-import type {
-  AstNode,
-  StrictAstNode,
-  StrictPolymorphicNode
-} from './slang-nodes/types.d.ts';
+import type { AstNode, PrintableNode } from './slang-nodes/types.d.ts';
 import type { PrintFunction } from './types.d.ts';
 
-function hasNodeIgnoreComment({ comments }: StrictAstNode): boolean {
+function hasNodeIgnoreComment({ comments }: PrintableNode): boolean {
   // Prettier sets SourceUnit's comments to undefined after assigning comments
   // to each node.
   return Boolean(
@@ -22,19 +18,20 @@ function hasNodeIgnoreComment({ comments }: StrictAstNode): boolean {
   );
 }
 
-function ignoreComments(path: AstPath<StrictAstNode>): void {
+function ignoreComments(path: AstPath<PrintableNode>): void {
   const node = path.node;
   // We ignore anything that is not an object
   if (node === null || typeof node !== 'object') return;
 
-  let key: keyof StrictAstNode;
+  let key: keyof PrintableNode;
   for (key in node) {
     switch (key) {
       // We ignore `kind` and `loc` since these are added by the parser.
       // `updateMetadata` is an internal function.
       case 'kind':
       case 'loc':
       case 'updateMetadata':
+      case 'print':
         break;
       // The key `comments` will contain every comment for this node.
       case 'comments':
@@ -56,10 +53,8 @@ function ignoreComments(path: AstPath<StrictAstNode>): void {
   }
 }
 
-// Nodes take care of undefined and string properties so we can restrict path
-// to AstPath<StrictAstNode>
 function genericPrint(
-  path: AstPath<Exclude<StrictAstNode, StrictPolymorphicNode>>,
+  path: AstPath<PrintableNode>,
   options: ParserOptions<AstNode>,
   print: PrintFunction
 ): Doc {
diff --git a/src/types.d.ts b/src/types.d.ts
diff --git a/tests/format/ContractDefinitions/ContractDefinitions.sol b/tests/format/ContractDefinitions/ContractDefinitions.sol
diff --git a/tests/format/ContractDefinitions/__snapshots__/format.test.js.snap b/tests/format/ContractDefinitions/__snapshots__/format.test.js.snap
