fix: yul function definition error

jeasonstudio · Jan 25, 2024 · 18cdebd · 18cdebd
1 parent 556bd4a
commit 18cdebd
Show file tree

Hide file tree

Showing 3 changed files with 30 additions and 22 deletions.
diff --git a/src/ast/yul/yul-function-definition.ts b/src/ast/yul/yul-function-definition.ts
@@ -10,7 +10,7 @@ export class YulFunctionDefinition extends BaseNode {
   body: YulBlock | null = null;
   constructor(ctx: YulFunctionDefinitionContext, visitor: SolidityParserVisitor<any>) {
     super(ctx, visitor);
-    this.name = ctx._YulIdentifier!.text!;
+    this.name = ctx.YulIdentifier(0)!.getText();
     this.parameters = ctx._arguments.map((arg) => arg.text!);
     this.returnParameters = ctx._returnParameters.map((arg) => arg.text!);
     this.body = ctx._body?.accept(visitor) ?? null;

diff --git a/src/tests/__snapshots__/prettier.test.ts.snap b/src/tests/__snapshots__/prettier.test.ts.snap
@@ -346,7 +346,7 @@ contract AbiEncode {
           _salt
         )
       }
-      function pos(length) -> pos {
+      function allocate(length) -> pos {
         pos := mload(0x40)
       }
     }

diff --git a/src/tests/yul.test.ts b/src/tests/yul.test.ts
@@ -3,45 +3,51 @@ import { test, expect } from 'vitest';
 
 test('assemblyStatement', () => {
   expect(createParse((p) => p.assemblyStatement())(`assembly { let x := 123 }`)).toMatchObject({
-    yulStatements: [{ type: 'YulVariableDeclaration', identifiers: ['x'] }],
+    yulStatements: [{ expression: { type: 'YulVariableDeclaration', identifiers: ['x'] } }],
   });
   expect(
     createParse((p) => p.assemblyStatement())(`assembly { f.b, foo, bar := mul(x, y) }`),
   ).toMatchObject({
-    yulStatements: [{ type: 'YulAssignment', paths: ['f.b', 'foo', 'bar'] }],
+    yulStatements: [{ expression: { type: 'YulAssignment', paths: ['f.b', 'foo', 'bar'] } }],
   });
   expect(createParse((p) => p.assemblyStatement())(`assembly { mul(x, y) }`)).toMatchObject({
-    yulStatements: [{ type: 'YulFunctionCall', identifier: 'mul', expressions: ['x', 'y'] }],
+    yulStatements: [
+      { expression: { type: 'YulFunctionCall', identifier: 'mul', expressions: ['x', 'y'] } },
+    ],
   });
   expect(createParse((p) => p.assemblyStatement())(`assembly { if foo {} }`)).toMatchObject({
-    yulStatements: [{ type: 'YulIfStatement', condition: 'foo', body: {} }],
+    yulStatements: [{ expression: { type: 'YulIfStatement', condition: 'foo', body: {} } }],
   });
   expect(createParse((p) => p.assemblyStatement())(`assembly { for {} bar {} {} }`)).toMatchObject({
     yulStatements: [
       {
-        type: 'YulForStatement',
-        initializationBlock: {},
-        conditionExpression: 'bar',
-        loopBlock: {},
-        body: {},
+        expression: {
+          type: 'YulForStatement',
+          initializationBlock: {},
+          conditionExpression: 'bar',
+          loopBlock: {},
+          body: {},
+        },
       },
     ],
   });
   expect(
     createParse((p) => p.assemblyStatement())(`assembly { leave break continue }`),
   ).toMatchObject({
-    yulStatements: ['leave', 'break', 'continue'],
+    yulStatements: [{ expression: 'leave' }, { expression: 'break' }, { expression: 'continue' }],
   });
   expect(
     createParse((p) => p.assemblyStatement())(`assembly { switch x case 0 {} default {} }`),
   ).toMatchObject({
     yulStatements: [
       {
-        type: 'YulSwitchStatement',
-        expression: 'x',
-        switchCases: [{ case: '0' }],
-        default: true,
-        body: {},
+        expression: {
+          type: 'YulSwitchStatement',
+          expression: 'x',
+          switchCases: [{ case: '0' }],
+          default: true,
+          body: {},
+        },
       },
     ],
   });
@@ -50,11 +56,13 @@ test('assemblyStatement', () => {
   ).toMatchObject({
     yulStatements: [
       {
-        type: 'YulFunctionDefinition',
-        name: 'foo',
-        parameters: ['x', 'y'],
-        returnParameters: ['a', 'b'],
-        body: {},
+        expression: {
+          type: 'YulFunctionDefinition',
+          name: 'foo',
+          parameters: ['x', 'y'],
+          returnParameters: ['a', 'b'],
+          body: {},
+        },
       },
     ],
   });