Add support for NatSpec comments

Author: vibern0

.eslintrc

@@ -15,6 +15,7 @@
     },
     "env": {
        "es6": true,
+       "mocha": true,
        "node": true
     }
 }

src/ASTBuilder.js

@@ -1,4 +1,5 @@
 const antlr4 = require('./antlr4/index')
+const parseComments = require('./natspec')
 
 function toText(ctx) {
   if (ctx !== null) {
@@ -116,13 +117,22 @@ const transformAST = {
 
   ContractDefinition(ctx) {
     const name = toText(ctx.identifier())
+    let natspec = null
+    let kind
+    if (ctx.natSpec()) {
+      natspec = parseComments(toText(ctx.getChild(0)))
+      kind = toText(ctx.getChild(1))
+    } else {
+      kind = toText(ctx.getChild(0))
+    }
     this._currentContract = name
 
     return {
+      natspec,
       name,
       baseContracts: this.visit(ctx.inheritanceSpecifier()),
       subNodes: this.visit(ctx.contractPart()),
-      kind: toText(ctx.getChild(0))
+      kind
     }
   },
 
@@ -215,7 +225,13 @@ const transformAST = {
       stateMutability = toText(ctx.modifierList().stateMutability(0))
     }
 
+    let natspec = null
+    if (ctx.natSpec()) {
+      natspec = parseComments(toText(ctx.getChild(0)))
+    }
+
     return {
+      natspec,
       name,
       parameters,
       returnParameters,
@@ -902,7 +918,12 @@ const transformAST = {
   },
 
   EventDefinition(ctx) {
+    let natspec = null
+    if (ctx.natSpec()) {
+      natspec = parseComments(toText(ctx.getChild(0)))
+    }
     return {
+      natspec,
       name: toText(ctx.identifier()),
       parameters: this.visit(ctx.eventParameterList()),
       isAnonymous: !!ctx.AnonymousKeyword()

src/natspec.js

@@ -0,0 +1,51 @@
+/**
+ * Parse a full text comment to an object that can be easly consumed
+ * @param comment full text comment
+ */
+module.exports = function(comment) {
+  // remove comments definers
+  // the comments can be multiline of single line
+  if (comment.substring(0, 3) === '/**') {
+    // if it is a multiline, remove line breaks, "/**"" (the comment begin)
+    // "*/"" (the comment end), and "* @" at the begining of each comment
+    comment = comment.replace(/(\r|\n|\/\*\*|\*\/|\*(?=[ ]?@))/g, '')
+  } else if (comment.substring(0, 3) === '///') {
+    // if it is a single line,
+    // remove just line breakers and "///" (the comment begin)
+    comment = comment.replace(/(\r|\n|\/\/\/)/g, '')
+  }
+
+  // split text by comment type
+  const splitComments = comment.split(
+    /@(title|author|notice|dev|param|return)/g
+  )
+  // let's start a map
+  const resultComments = {}
+
+  // iterate through all the split comments
+  // start in the second result (because the first one is usually nothing)
+  // then iterate each two elements, because since the split is done using
+  // the natspec type
+  for (let x = 1; x < splitComments.length; x += 2) {
+    const key = splitComments[x]
+    const value = splitComments[x + 1].trim()
+
+    // if the comment if type 'param' we need to extract the first word
+    // which is the variable name
+    if (key === 'param') {
+      const sep = value.indexOf(' ')
+      const paramName = value.substring(0, sep)
+      const paramValue = value.substring(sep + 1)
+
+      if (resultComments['params'] == null) {
+        resultComments['params'] = {}
+      }
+
+      resultComments['params'][paramName] = paramValue
+    } else {
+      // finally, if the value was found, update it
+      resultComments[key] = value
+    }
+  }
+  return resultComments
+}

test/ast.js

@@ -84,6 +84,7 @@ describe('AST', () => {
     var ast = parseContract("contract test {}")
     assert.deepEqual(ast, {
       "type": "ContractDefinition",
+      "natspec": null,
       "name": "test",
       "baseContracts": [],
       "subNodes": [],
@@ -94,6 +95,7 @@ describe('AST', () => {
     ast = parseContract("contract test is foo, bar {}")
     assert.deepEqual(ast, {
       "type": "ContractDefinition",
+      "natspec": null,
       "name": "test",
       "baseContracts": [
         {
@@ -121,6 +123,7 @@ describe('AST', () => {
     ast = parseContract("library test {}")
     assert.deepEqual(ast, {
       "type": "ContractDefinition",
+      "natspec": null,
       "name": "test",
       "baseContracts": [],
       "subNodes": [],
@@ -131,6 +134,7 @@ describe('AST', () => {
     ast = parseContract("interface test {}")
     assert.deepEqual(ast, {
       "type": "ContractDefinition",
+      "natspec": null,
       "name": "test",
       "baseContracts": [],
       "subNodes": [],
@@ -171,6 +175,7 @@ describe('AST', () => {
     var ast = parseNode("function foo(uint a) pure {}")
     assert.deepEqual(ast, {
       "type": "FunctionDefinition",
+      "natspec": null,
       "name": "foo",
       "parameters": [
         {
@@ -200,6 +205,7 @@ describe('AST', () => {
     ast = parseNode("function foo(uint a) pure returns (uint256) {}")
     assert.deepEqual(ast, {
       "type": "FunctionDefinition",
+      "natspec": null,
       "name": "foo",
       "parameters": [
         {
@@ -1091,6 +1097,7 @@ describe('AST', () => {
     assert.deepEqual(ast, {
       "type": "EventDefinition",
       "name": "Foo",
+      "natspec": null,
       "parameters": [
         {
           "type": "VariableDeclaration",
@@ -1490,4 +1497,121 @@ describe('AST', () => {
       "type": "AssemblyIf"
     })
   })
+
+  it("NatSpec multi line", function () {
+    const ast = parser.parse(
+`/**
+  * @dev This is the Sum contract.
+  * @title Sum Contract
+  * @author username
+  */
+contract Sum { }`
+    );
+    assert.deepEqual(ast.children[0], {
+      type: "ContractDefinition",
+      natspec: {
+        dev: "This is the Sum contract.",
+        title: "Sum Contract",
+        author: "username",
+      },
+      name: "Sum",
+      baseContracts: [],
+      subNodes: [],
+      kind: "contract",
+    })
+  })
+
+  it("NatSpec single line", function () {
+    const ast = parser.parse(
+`/// @dev This is the Sum contract.
+/// @title Sum Contract
+/// @author username
+contract Sum { }`
+    );
+    assert.deepEqual(ast.children[0], {
+      type: "ContractDefinition",
+      natspec: {
+        dev: "This is the Sum contract.",
+        title: "Sum Contract",
+        author: "username",
+      },
+      name: "Sum",
+      baseContracts: [],
+      subNodes: [],
+      kind: "contract",
+    })
+  })
+
+  it("NatSpec multi line event", function () {
+    const ast = parseNode(
+`/**
+  * @dev This method says hello
+  * @param user the user address
+  */
+  event sayHello(address user);`
+    );
+    assert.deepEqual(ast.natspec, {
+      dev: "This method says hello",
+      params: {
+        user: "the user address"
+      },
+    })
+  })
+
+  it("NatSpec multi line function", function () {
+    const ast = parseNode(
+`/**
+  * @dev This method transfer fund to other user
+  * @param from the address extract funds
+  * @param to the user address to give funds
+  * @param amount the amount to transfer
+  */
+ function transfer(address from, address to, uint256 amount) public {}`);
+
+    assert.deepEqual(ast.natspec, {
+      dev: "This method transfer fund to other user",
+      params: {
+        from: "the address extract funds",
+        to: "the user address to give funds",
+        amount: "the amount to transfer",
+      },
+    })
+  })
+  it("NatSpec multi line multiple functions in contract", function () {
+    const ast = parser.parse(
+`/**
+  * @dev The ERC20 contract
+  */
+ contract ERC20 {
+    /**
+     * @dev This method transfer fund to other user
+     * @param from the address extract funds
+     * @param to the user address to give funds
+     * @param amount the amount to transfer
+     */
+    function transfer(address from, address to, uint256 amount) public {}
+    /**
+     * @dev This method gets the approved amount
+     * @param user the user address to verify
+     * @return the approved amount
+     */
+    function approved(address user) public view returns(uint256) {}
+ }`);
+    const methods = ast.children[0].subNodes;
+    assert.deepEqual(methods[0].natspec, {
+      dev: 'This method transfer fund to other user',
+      params: {
+        from: 'the address extract funds',
+        to: 'the user address to give funds',
+        amount: 'the amount to transfer',
+      },
+    });
+    assert.deepEqual(methods[1].natspec, {
+      dev: 'This method gets the approved amount',
+      params: {
+        user: 'the user address to verify',
+      },
+      return: 'the approved amount',
+    });
+  })
 })