From 58a216cb85a31579c13d078f115e84d04627fc6a Mon Sep 17 00:00:00 2001
From: Alex Rea <git@area.io>
Date: Thu, 16 Nov 2017 17:39:04 +0000
Subject: [PATCH] Only remove pure/view/constant from functions

Uses the preprocessor walking over the AST rather than a regex
---
 lib/app.js          | 10 +++-------
 lib/preprocessor.js |  9 +++++++++
 2 files changed, 12 insertions(+), 7 deletions(-)

diff --git a/lib/app.js b/lib/app.js
index ae9bbad..8020bbe 100644
--- a/lib/app.js
+++ b/lib/app.js
@@ -8,6 +8,7 @@ const istanbul = require('istanbul');
 const getInstrumentedVersion = require('./instrumentSolidity.js');
 const CoverageMap = require('./coverageMap.js');
 const defaultTruffleConfig = require('./truffleConfig.js');
+const preprocessor = require('./preprocessor');
 
 const isWin = /^win/.test(process.platform);
 
@@ -284,16 +285,11 @@ class App {
    */
   postProcessPure(env) {
     shell.ls(`${env}/**/*.sol`).forEach(file => {
-      const pureRe = /\spure\s/gi;
-      const viewRe = /\sview\s/gi;
-      const constantRe = /\sconstant\s/gi;
       const contractPath = this.platformNeutralPath(file);
       let contract = fs.readFileSync(contractPath).toString();
-      contract = contract.replace(pureRe, ' ');
-      contract = contract.replace(viewRe, ' ');
-      contract = contract.replace(constantRe, ' ');
+      contract = preprocessor.run(contract);
       fs.writeFileSync(contractPath, contract);
-    })
+    });
   }
 
   /**
diff --git a/lib/preprocessor.js b/lib/preprocessor.js
index dabaad6..c6dcf22 100644
--- a/lib/preprocessor.js
+++ b/lib/preprocessor.js
@@ -50,6 +50,15 @@ module.exports.run = function r(contract) {
           contract = blockWrap(contract, node.body);
           keepRunning = true;
           this.stopTraversal();
+        } else if (node.type === 'FunctionDeclaration' && node.modifiers) {
+          // We want to remove constant / pure / view from functions
+          for (let i = 0; i < node.modifiers.length; i++) {
+            if (['pure', 'constant', 'view'].indexOf(node.modifiers[i].name) > -1) {
+              contract = contract.slice(0, node.modifiers[i].start) + contract.slice(node.modifiers[i].end);
+              keepRunning = true;
+              this.stopTraversal();
+            }
+          }
         }
       },
     });