Add support for custom storage layout

Author: fvictorio

src/ASTBuilder.ts

@@ -54,6 +54,11 @@ export class ASTBuilder
 
     this._currentContract = name
 
+    const customLayoutStorageList = ctx.customStorageLayout_list()
+    if (customLayoutStorageList.length > 1) {
+      throw new Error('Only one custom storage layout is allowed per contract')
+    }
+
     const node: AST.ContractDefinition = {
       type: 'ContractDefinition',
       name,
@@ -64,6 +69,12 @@ export class ASTBuilder
       kind,
     }
 
+    if (customLayoutStorageList.length === 1) {
+      node.storageLayout = this.visitExpression(
+        customLayoutStorageList[0].expression()
+      )
+    }
+
     return this._addMeta(node, ctx)
   }
 

src/ast-types.ts

@@ -38,6 +38,7 @@ export interface ContractDefinition extends BaseASTNode {
   baseContracts: InheritanceSpecifier[]
   kind: string
   subNodes: BaseASTNode[]
+  storageLayout?: Expression
 }
 
 export interface InheritanceSpecifier extends BaseASTNode {

test/ast.ts

@@ -440,6 +440,78 @@ describe('AST', () => {
       subNodes: [],
       kind: 'interface',
     })
+
+    // custom layout storage, no inheritance
+    ast = parseContract('contract test layout at 1 {}')
+    assert.deepEqual(ast, {
+      type: 'ContractDefinition',
+      name: 'test',
+      baseContracts: [],
+      subNodes: [],
+      kind: 'contract',
+      storageLayout: {
+        number: '1',
+        subdenomination: null,
+        type: 'NumberLiteral',
+      },
+    })
+
+    // custom layout storage after inheritance specifier
+    ast = parseContract('contract test is foo layout at 1+2 {}')
+    assert.deepEqual(ast, {
+      type: 'ContractDefinition',
+      name: 'test',
+      baseContracts: [
+        {
+          type: 'InheritanceSpecifier',
+          baseName: {
+            type: 'UserDefinedTypeName',
+            namePath: 'foo',
+          },
+          arguments: [],
+        },
+      ],
+      subNodes: [],
+      kind: 'contract',
+      storageLayout: {
+        type: 'BinaryOperation',
+        operator: '+',
+        left: {
+          type: 'NumberLiteral',
+          number: '1',
+          subdenomination: null,
+        },
+        right: {
+          type: 'NumberLiteral',
+          number: '2',
+          subdenomination: null,
+        },
+      },
+    })
+
+    // custom layout storage before inheritance specifier
+    ast = parseContract('contract test layout at 10 is foo {}')
+    assert.deepEqual(ast, {
+      type: 'ContractDefinition',
+      name: 'test',
+      baseContracts: [
+        {
+          type: 'InheritanceSpecifier',
+          baseName: {
+            type: 'UserDefinedTypeName',
+            namePath: 'foo',
+          },
+          arguments: [],
+        },
+      ],
+      subNodes: [],
+      kind: 'contract',
+      storageLayout: {
+        type: 'NumberLiteral',
+        number: '10',
+        subdenomination: null,
+      },
+    })
   })
 
   it('FunctionDefinition constructor case', () => {