Add support for user-defined operators

Author: fvictorio

src/ASTBuilder.ts

@@ -533,32 +533,48 @@ export class ASTBuilder
       typeName = this.visitTypeName(ctxTypeName)
     }
 
-    const isGlobal = ctx.GlobalKeyword() !== undefined;
+    const isGlobal = ctx.GlobalKeyword() !== undefined
+
+    const usingForObjectCtx = ctx.usingForObject()
+
+    const userDefinedTypeNameCtx = usingForObjectCtx.userDefinedTypeName()
 
-    // the object of the `usingForDeclaration` can be a single identifier
-    // (the library name) or a group of functions:
-    //   using Lib for uint;
-    //   using { f } for uint;
     let node: AST.UsingForDeclaration
-    const usingForObject = ctx.usingForObject()
-    const firstChild = this._toText(usingForObject.getChild(0))
-    if (firstChild === '{') {
+    if (userDefinedTypeNameCtx !== undefined) {
+      // using Lib for ...
       node = {
         type: 'UsingForDeclaration',
         isGlobal,
         typeName,
-        libraryName: null,
-        functions: usingForObject
-          .userDefinedTypeName()
-          .map((x) => this._toText(x)),
+        libraryName: this._toText(userDefinedTypeNameCtx),
+        functions: [],
+        operators: [],
       }
     } else {
+      // using { } for ...
+      const usingForObjectDirectives = usingForObjectCtx.usingForObjectDirective()
+      const functions: string[] = []
+      const operators: Array<string | null> = []
+
+      for (const usingForObjectDirective of usingForObjectDirectives) {
+        functions.push(
+          this._toText(usingForObjectDirective.userDefinedTypeName())
+        )
+        const operator = usingForObjectDirective.userDefinableOperators()
+        if (operator !== undefined) {
+          operators.push(this._toText(operator))
+        } else {
+          operators.push(null)
+        }
+      }
+
       node = {
         type: 'UsingForDeclaration',
         isGlobal,
         typeName,
-        libraryName: this._toText(usingForObject.userDefinedTypeName(0)),
-        functions: [],
+        libraryName: null,
+        functions,
+        operators,
       }
     }
 
@@ -1726,7 +1742,7 @@ export class ASTBuilder
         value: this._toText(ctx.BooleanLiteral()!) === 'true',
       }
 
-      return this._addMeta(node, ctx);
+      return this._addMeta(node, ctx)
     }
 
     if (ctx.DecimalNumber()) {

src/ast-types.ts

@@ -152,6 +152,9 @@ export interface UsingForDeclaration extends BaseASTNode {
   type: 'UsingForDeclaration'
   typeName: TypeName | null
   functions: string[]
+  // for each item in `functions`, the item with the same index in `operators`
+  // will be the defined operator, or null if it's just an attached function
+  operators: Array<string | null>
   libraryName: string | null
   isGlobal: boolean;
 }

test/ast.ts

@@ -52,6 +52,7 @@ describe('AST', () => {
         },
         libraryName: 'Lib',
         functions: [],
+        operators: [],
       })
 
       ast = parseNode('using Lib for *;')
@@ -61,6 +62,7 @@ describe('AST', () => {
         typeName: null,
         libraryName: 'Lib',
         functions: [],
+        operators: [],
       })
 
       ast = parseNode('using Lib for S;')
@@ -73,6 +75,7 @@ describe('AST', () => {
         },
         libraryName: 'Lib',
         functions: [],
+        operators: [],
       })
 
       ast = parseNode('using L.Lib for S;')
@@ -85,6 +88,7 @@ describe('AST', () => {
         },
         libraryName: 'L.Lib',
         functions: [],
+        operators: [],
       })
     })
 
@@ -100,6 +104,7 @@ describe('AST', () => {
         },
         libraryName: 'Lib',
         functions: [],
+        operators: [],
       })
 
       ast = parseNode('using Lib for * global;')
@@ -109,6 +114,7 @@ describe('AST', () => {
         typeName: null,
         libraryName: 'Lib',
         functions: [],
+        operators: [],
       })
 
       ast = parseNode('using Lib for S global;')
@@ -121,6 +127,7 @@ describe('AST', () => {
         },
         libraryName: 'Lib',
         functions: [],
+        operators: [],
       })
 
       ast = parseNode('using L.Lib for S global;')
@@ -133,6 +140,7 @@ describe('AST', () => {
         },
         libraryName: 'L.Lib',
         functions: [],
+        operators: [],
       })
     })
 
@@ -148,6 +156,7 @@ describe('AST', () => {
         },
         libraryName: null,
         functions: ['f'],
+        operators: [null],
       })
 
       ast = parseNode('using { f, g } for uint;')
@@ -161,6 +170,7 @@ describe('AST', () => {
         },
         libraryName: null,
         functions: ['f', 'g'],
+        operators: [null, null],
       })
 
       ast = parseNode('using { f } for *;')
@@ -170,6 +180,7 @@ describe('AST', () => {
         typeName: null,
         libraryName: null,
         functions: ['f'],
+        operators: [null],
       })
 
       ast = parseNode('using { f } for S;')
@@ -182,6 +193,7 @@ describe('AST', () => {
         },
         libraryName: null,
         functions: ['f'],
+        operators: [null],
       })
     })
 
@@ -197,6 +209,7 @@ describe('AST', () => {
         },
         libraryName: null,
         functions: ['f'],
+        operators: [null],
       })
 
       ast = parseNode('using { f, g } for uint global;')
@@ -210,6 +223,7 @@ describe('AST', () => {
         },
         libraryName: null,
         functions: ['f', 'g'],
+        operators: [null, null],
       })
 
       ast = parseNode('using { f } for * global;')
@@ -219,6 +233,7 @@ describe('AST', () => {
         typeName: null,
         libraryName: null,
         functions: ['f'],
+        operators: [null],
       })
 
       ast = parseNode('using { f } for S global;')
@@ -231,6 +246,67 @@ describe('AST', () => {
         },
         libraryName: null,
         functions: ['f'],
+        operators: [null],
+      })
+    })
+
+    describe('user-defined operators', function () {
+      it('defining a single operator', function () {
+        const ast = parseNode('using { add as + } for Fixed18 global;')
+        assert.deepEqual(ast, {
+          type: 'UsingForDeclaration',
+          isGlobal: true,
+          typeName: {
+            type: 'UserDefinedTypeName',
+            namePath: 'Fixed18',
+          },
+          libraryName: null,
+          functions: ['add'],
+          operators: ['+'],
+        })
+      })
+
+      it('defining two operators', function () {
+        const ast = parseNode('using { add as +, sub as - } for Fixed18 global;')
+        assert.deepEqual(ast, {
+          type: 'UsingForDeclaration',
+          isGlobal: true,
+          typeName: {
+            type: 'UserDefinedTypeName',
+            namePath: 'Fixed18',
+          },
+          libraryName: null,
+          functions: ['add', 'sub'],
+          operators: ['+', '-'],
+        })
+      })
+
+      it('mixing attached functions and user-defined operators', function () {
+        let ast = parseNode('using { add, sub as - } for Fixed18 global;')
+        assert.deepEqual(ast, {
+          type: 'UsingForDeclaration',
+          isGlobal: true,
+          typeName: {
+            type: 'UserDefinedTypeName',
+            namePath: 'Fixed18',
+          },
+          libraryName: null,
+          functions: ['add', 'sub'],
+          operators: [null, '-'],
+        })
+
+        ast = parseNode('using { add as +, sub } for Fixed18 global;')
+        assert.deepEqual(ast, {
+          type: 'UsingForDeclaration',
+          isGlobal: true,
+          typeName: {
+            type: 'UserDefinedTypeName',
+            namePath: 'Fixed18',
+          },
+          libraryName: null,
+          functions: ['add', 'sub'],
+          operators: ['+', null],
+        })
       })
     })
   })