From 3f8450f060260870778b1d4de88af695247ef807 Mon Sep 17 00:00:00 2001
From: Richard Ramos <info@richardramos.me>
Date: Fri, 4 May 2018 16:17:12 -0400
Subject: [PATCH] Extra functionality for deciding which elements to include in
 a graph

---
 lib/cmd.js        |  8 +++++++-
 lib/cmds/graph.js | 36 ++++++++++++++++++++++--------------
 lib/index.js      |  6 +++---
 3 files changed, 32 insertions(+), 18 deletions(-)

diff --git a/lib/cmd.js b/lib/cmd.js
index ebc82b24..669086fe 100644
--- a/lib/cmd.js
+++ b/lib/cmd.js
@@ -191,11 +191,17 @@ class Cmd {
   graph() {
     program
       .command('graph [environment]')
+      .option('--skip-undeployed', 'Graph will not include undeployed contracts')
+      .option('--skip-functions', 'Graph will not include functions')
+      .option('--skip-events', 'Graph will not include events')
       .description('generates documentation based on the smart contracts configured')
       .action(function (env, options) {
         embark.graph({
           env: env || 'development',
-          logFile: options.logfile
+          logFile: options.logfile,
+          skipUndeployed: options.skipUndeployed,
+          skipFunctions: options.skipFunctions,
+          skipEvents: options.skipEvents
         });
       });
   }
diff --git a/lib/cmds/graph.js b/lib/cmds/graph.js
index 7621fe65..8ae2eda5 100644
--- a/lib/cmds/graph.js
+++ b/lib/cmds/graph.js
@@ -6,15 +6,16 @@ class GraphGenerator {
         this.engine = engine;
     }
 
-    generate() {
+    generate(options) {
         let id = 0;
         let contractString = "";
         let relationshipString = "";
         let idMapping = {};
         let contractInheritance = {};
 
-        
         for (let contract in this.engine.contractsManager.contracts) {
+            if(options.skipUndeployed && !this.engine.contractsManager.contracts[contract].deploy) continue;
+
             id++;
 
             idMapping[contract] = id;
@@ -30,23 +31,27 @@ class GraphGenerator {
                 contractLabel += ": " + this.engine.contractsManager.contracts[contract].instanceOf;
                 tooltip += " instance of " + this.engine.contractsManager.contracts[contract].instanceOf; 
             } else {
-                contractLabel += "|";
+                if(!(options.skipFunctions === true && options.skipEvents === true)) contractLabel += "|";
 
                 for(let i = 0; i < this.engine.contractsManager.contracts[contract].abiDefinition.length; i++){
-                    switch(this.engine.contractsManager.contracts[contract].abiDefinition[i].type){
+                    let abiDef = this.engine.contractsManager.contracts[contract].abiDefinition[i];
+                    if(abiDef.type == 'event' && options.skipEvents) continue;
+                    if(['constructor', 'fallback'].indexOf(abiDef.type) > -1 && options.skipFunctions) continue;
+
+                    switch(abiDef.type){
                         case 'fallback': 
                             contractLabel += "«fallback»()\\l";
                             break;
-                            case 'constructor':
+                        case 'constructor':
                             contractLabel += "«constructor»(";
-                            this.engine.contractsManager.contracts[contract].abiDefinition[i].inputs.forEach(function(elem, index){
+                            abiDef.inputs.forEach(function(elem, index){
                                 contractLabel += (index == 0 ? "" : ", ") + elem.type;
                             });
                             contractLabel += ")\\l";
                             break;
-                            case 'event':
-                            contractLabel += "«event»" + this.engine.contractsManager.contracts[contract].abiDefinition[i].name  + "(";
-                            this.engine.contractsManager.contracts[contract].abiDefinition[i].inputs.forEach(function(elem, index){
+                        case 'event':
+                            contractLabel += "«event»" + abiDef.name  + "(";
+                            abiDef.inputs.forEach(function(elem, index){
                                 contractLabel += (index == 0 ? "" : ", ") + elem.type;
                             });
                             contractLabel += ")\\l";
@@ -56,7 +61,7 @@ class GraphGenerator {
                 }
 
                 let fHashes = this.engine.contractsManager.contracts[contract].functionHashes;
-                if(fHashes != {} && fHashes != undefined){
+                if(fHashes != {} && fHashes != undefined && !options.skipFunctions){
                     for(let method in this.engine.contractsManager.contracts[contract].functionHashes){
                         contractLabel += method + '\\l';
                     }
@@ -73,20 +78,23 @@ class GraphGenerator {
 
         }
 
-        for (let c in this.engine.contractsManager.contractDependencies){
+        for (let c in this.engine.contractsManager.contractDependencies){            
             let contractDependencies = Array.from(new Set(this.engine.contractsManager.contractDependencies[c]));
-            contractDependencies.forEach(function(d){
+            contractDependencies.forEach((d) => {
                 if(idMapping[c] !== undefined && idMapping[d] !== undefined){
-                    relationshipString += `${idMapping[d]}->${idMapping[c]}[constraint=true, arrowtail=diamond, tooltip="${c} uses ${d}"]\n`;
+                    if(options.skipUndeployed && this.engine.contractsManager.contracts[c].deploy && this.engine.contractsManager.contracts[d].deploy){
+                        relationshipString += `${idMapping[d]}->${idMapping[c]}[constraint=true, arrowtail=diamond, tooltip="${c} uses ${d}"]\n`;
+                    }
                 }
             });
         }
 
         for (let c in contractInheritance){
+            if(options.skipUndeployed && !this.engine.contractsManager.contracts[contractInheritance[c]].deploy) continue;
+
             relationshipString += `${idMapping[contractInheritance[c]]}->${idMapping[c]}[tooltip="${c} instance of ${contractInheritance[c]}"]\n`;
         }
 
-
         let dot = `
         digraph Contracts {
             node[shape=record,style=filled]
diff --git a/lib/index.js b/lib/index.js
index 1f73d771..81481e7d 100644
--- a/lib/index.js
+++ b/lib/index.js
@@ -254,7 +254,7 @@ class Embark {
           callback(err);
         });
       }
-    ], function (err, _result) {
+    ],  (err, _result) => {
       if (err) {
         engine.logger.error(err.message);
         engine.logger.info(err.stack);
@@ -262,9 +262,9 @@ class Embark {
 
         const GraphGenerator = require('./cmds/graph.js');
         let graphGen = new GraphGenerator(engine);
-        graphGen.generate();
+        graphGen.generate(options);
 
-        engine.logger.info("Done".underline);
+        engine.logger.info("Done. ./diagram.svg generated".underline);
         process.exit();
       }
     });