Placing comments in the middle of a member access chained expression

src/slang-comments/handlers/handle-member-access-expression-comments.ts

@@ -0,0 +1,37 @@
+import { NonterminalKind } from '@nomicfoundation/slang/cst';
+import { util } from 'prettier';
+
+import type { HandlerParams } from './types.d.ts';
+
+const { addLeadingComment, addTrailingComment } = util;
+
+export default function handleMemberAccessExpressionComments({
+  precedingNode,
+  enclosingNode,
+  followingNode,
+  comment
+}: HandlerParams): boolean {
+  if (enclosingNode?.kind !== NonterminalKind.MemberAccessExpression) {
+    return false;
+  }
+
+  if (followingNode === enclosingNode.operand) {
+    addLeadingComment(followingNode, comment);
+    return true;
+  }
+
+  if (
+    precedingNode === enclosingNode.operand &&
+    followingNode === enclosingNode.member
+  ) {
+    addTrailingComment(precedingNode, comment);
+    return true;
+  }
+
+  if (precedingNode === enclosingNode.member) {
+    addTrailingComment(precedingNode, comment);
+    return true;
+  }
+
+  return false;
+}

src/slang-comments/handlers/index.ts

@@ -5,6 +5,7 @@ import handleElseBranchComments from './handle-else-branch-comments.js';
 import handleIfStatementComments from './handle-if-statement-comments.js';
 import handleInterfaceDefinitionComments from './handle-interface-definition-comments.js';
 import handleLibraryDefinitionComments from './handle-library-definition-comments.js';
+import handleMemberAccessExpressionComments from './handle-member-access-expression-comments.js';
 import handleModifierInvocationComments from './handle-modifier-invocation-comments.js';
 import handleParametersDeclarationComments from './handle-parameters-declaration-comments.js';
 import handlePositionalArgumentsDeclarationComments from './handle-positional-arguments-declaration-comments.js';
@@ -22,6 +23,7 @@ export default [
   handleIfStatementComments,
   handleInterfaceDefinitionComments,
   handleLibraryDefinitionComments,
+  handleMemberAccessExpressionComments,
   handleModifierInvocationComments,
   handleParametersDeclarationComments,
   handlePositionalArgumentsDeclarationComments,

src/slang-nodes/FunctionCallExpression.ts

@@ -1,8 +1,6 @@
 import { NonterminalKind } from '@nomicfoundation/slang/cst';
-import { doc } from 'prettier';
-import { isLabel } from '../slang-utils/is-label.js';
-import { printGroupAndIndentIfBreakPair } from '../slang-printers/print-group-and-indent-if-break-pair.js';
 import { extractVariant } from '../slang-utils/extract-variant.js';
+import { printPossibleMemberAccessChainItem } from '../slang-printers/print-member-access-chain-item.js';
 import { SlangNode } from './SlangNode.js';
 import { Expression } from './Expression.js';
 import { ArgumentsDeclaration } from './ArgumentsDeclaration.js';
@@ -12,8 +10,6 @@ import type { AstPath, Doc, ParserOptions } from 'prettier';
 import type { CollectedMetadata, PrintFunction } from '../types.d.ts';
 import type { AstNode } from './types.d.ts';
 
-const { label } = doc.builders;
-
 export class FunctionCallExpression extends SlangNode {
   readonly kind = NonterminalKind.FunctionCallExpression;
 
@@ -39,20 +35,9 @@ export class FunctionCallExpression extends SlangNode {
   }
 
   print(path: AstPath<FunctionCallExpression>, print: PrintFunction): Doc {
-    const operand = path.call(print, 'operand');
-    const argumentsDoc = path.call(print, 'arguments');
-
-    // If we are at the end of a MemberAccessChain we should indent the
-    // arguments accordingly.
-    if (isLabel(operand) && operand.label === 'MemberAccessChain') {
-      // We wrap the expression in a label in case there is an IndexAccess or
-      // a FunctionCall following this IndexAccess.
-      return label(
-        'MemberAccessChain',
-        printGroupAndIndentIfBreakPair(operand.contents, argumentsDoc)
-      );
-    }
-
-    return [operand, argumentsDoc].flat();
+    return printPossibleMemberAccessChainItem(
+      path.call(print, 'operand'),
+      path.call(print, 'arguments')
+    );
   }
 }

src/slang-nodes/IndexAccessExpression.ts

@@ -1,9 +1,7 @@
 import { NonterminalKind } from '@nomicfoundation/slang/cst';
-import { doc } from 'prettier';
 import { printSeparatedItem } from '../slang-printers/print-separated-item.js';
-import { printGroupAndIndentIfBreakPair } from '../slang-printers/print-group-and-indent-if-break-pair.js';
-import { isLabel } from '../slang-utils/is-label.js';
 import { extractVariant } from '../slang-utils/extract-variant.js';
+import { printPossibleMemberAccessChainItem } from '../slang-printers/print-member-access-chain-item.js';
 import { SlangNode } from './SlangNode.js';
 import { Expression } from './Expression.js';
 import { IndexAccessEnd } from './IndexAccessEnd.js';
@@ -13,8 +11,6 @@ import type { AstPath, Doc, ParserOptions } from 'prettier';
 import type { CollectedMetadata, PrintFunction } from '../types.d.ts';
 import type { AstNode } from './types.d.ts';
 
-const { label } = doc.builders;
-
 export class IndexAccessExpression extends SlangNode {
   readonly kind = NonterminalKind.IndexAccessExpression;
 
@@ -47,24 +43,10 @@ export class IndexAccessExpression extends SlangNode {
   }
 
   print(path: AstPath<IndexAccessExpression>, print: PrintFunction): Doc {
-    const operand = path.call(print, 'operand');
-    const indexDoc = [
+    return printPossibleMemberAccessChainItem(path.call(print, 'operand'), [
       '[',
       printSeparatedItem([path.call(print, 'start'), path.call(print, 'end')]),
       ']'
-    ];
-
-    // If we are at the end of a MemberAccessChain we should indent the
-    // arguments accordingly.
-    if (isLabel(operand) && operand.label === 'MemberAccessChain') {
-      // We wrap the expression in a label in case there is an IndexAccess or
-      // a FunctionCall following this IndexAccess.
-      return label(
-        'MemberAccessChain',
-        printGroupAndIndentIfBreakPair(operand.contents, indexDoc)
-      );
-    }
-
-    return [operand, indexDoc].flat();
+    ]);
   }
 }

src/slang-nodes/MemberAccessExpression.ts

@@ -3,6 +3,7 @@ import { doc } from 'prettier';
 import { isLabel } from '../slang-utils/is-label.js';
 import { extractVariant } from '../slang-utils/extract-variant.js';
 import { isChainableExpression } from '../slang-utils/is-chainable-expression.js';
+import { memberAccessChainLabel } from '../slang-printers/print-member-access-chain-item.js';
 import { SlangNode } from './SlangNode.js';
 import { Expression } from './Expression.js';
 import { TerminalNode } from './TerminalNode.js';
@@ -14,6 +15,8 @@ import type { AstNode, ChainableExpression, StrictAstNode } from './types.d.ts';
 
 const { group, indent, label, softline } = doc.builders;
 
+const separatorLabel = Symbol('separator');
+
 function isEndOfChain(
   node: ChainableExpression,
   path: AstPath<StrictAstNode>
@@ -91,13 +94,13 @@ function isEndOfChain(
  * be printed.
  */
 function processChain(chain: Doc[]): Doc {
-  const firstSeparatorIndex = chain.findIndex(
-    (element) => isLabel(element) && element.label === 'separator'
+  const firstSeparatorIndex = chain.findIndex((element) =>
+    isLabel(element, separatorLabel)
   );
 
   // We wrap the expression in a label in case there is an IndexAccess or
   // a FunctionCall following this MemberAccess.
-  return label('MemberAccessChain', [
+  return label(memberAccessChainLabel, [
     // The doc[] before the first separator
     chain.slice(0, firstSeparatorIndex),
     // The doc[] containing the rest of the chain
@@ -128,9 +131,14 @@ export class MemberAccessExpression extends SlangNode {
   }
 
   print(path: AstPath<MemberAccessExpression>, print: PrintFunction): Doc {
+    let operandDoc = path.call(print, 'operand');
+    if (Array.isArray(operandDoc)) {
+      operandDoc = operandDoc.flat();
+    }
+
     const document = [
-      path.call(print, 'operand'),
-      label('separator', [softline, '.']),
+      operandDoc,
+      label(separatorLabel, [softline, '.']),
       path.call(print, 'member')
     ].flat();
 

src/slang-printers/print-member-access-chain-item.ts

@@ -0,0 +1,27 @@
+import { doc } from 'prettier';
+import { isLabel } from '../slang-utils/is-label.js';
+import { printGroupAndIndentIfBreakPair } from './print-group-and-indent-if-break-pair.js';
+
+import type { Doc } from 'prettier';
+
+const { label } = doc.builders;
+
+export const memberAccessChainLabel = Symbol('MemberAccessChain');
+
+export function printPossibleMemberAccessChainItem(
+  operand: Doc,
+  predicate: Doc
+): Doc {
+  // If we are at the end of a MemberAccessChain we should indent the
+  // arguments accordingly.
+  if (isLabel(operand, memberAccessChainLabel)) {
+    // We wrap the expression in a label in case there is an IndexAccess or
+    // a FunctionCall following Node.
+    return label(
+      memberAccessChainLabel,
+      printGroupAndIndentIfBreakPair(operand.contents, predicate)
+    );
+  }
+
+  return [operand, predicate].flat();
+}

src/slang-utils/is-label.ts

@@ -1,5 +1,11 @@
 import type { Doc, doc } from 'prettier';
 
-export function isLabel(document: Doc): document is doc.builders.Label {
-  return (document as doc.builders.DocCommand).type === 'label';
+export function isLabel(
+  document: Doc,
+  value: symbol
+): document is doc.builders.Label {
+  return (
+    (document as doc.builders.DocCommand).type === 'label' &&
+    (document as doc.builders.Label).label === value
+  );
 }

tests/format/MemberAccess/MemberAccess.sol

@@ -62,6 +62,19 @@ contract MemberAccess {
             abi.encodeWithSelector(f.selector),
             abi.encode(returned1)
         );
+// Comments in between the chain
+        game.
+// CONFIG
+// Do not touch
+config. // Window
+// Resolution 1000
+resolveWindow = 1000;
+config.defaultDelay = 0;
+begin
+    .// Comment 1
+    functionCall(/* Comment about parameter */ parameter)
+    .// Comment 2
+    end;
     }
 }
 

tests/format/MemberAccess/__snapshots__/format.test.js.snap

@@ -70,6 +70,19 @@ contract MemberAccess {
             abi.encodeWithSelector(f.selector),
             abi.encode(returned1)
         );
+// Comments in between the chain
+        game.
+// CONFIG
+// Do not touch
+config. // Window
+// Resolution 1000
+resolveWindow = 1000;
+config.defaultDelay = 0;
+begin
+    .// Comment 1
+    functionCall(/* Comment about parameter */ parameter)
+    .// Comment 2
+    end;
     }
 }
 
@@ -206,6 +219,17 @@ contract MemberAccess {
             abi.encodeWithSelector(f.selector),
             abi.encode(returned1)
         );
+        // Comments in between the chain
+        game
+        // CONFIG
+        // Do not touch
+            .config // Window
+            // Resolution 1000
+            .resolveWindow = 1000;
+        config.defaultDelay = 0;
+        begin // Comment 1
+            .functionCall(/* Comment about parameter */ parameter) // Comment 2
+            .end;
     }
 }