From d5314d2a8539bdac75e626e74c9877bc6b5b217d Mon Sep 17 00:00:00 2001
From: Nick O'Leary <nick.oleary@gmail.com>
Date: Tue, 16 Feb 2021 20:46:41 +0000
Subject: [PATCH] Add select-up/downstream-nodes action to editor

---
 .../editor-client/src/js/keymap.json          |  2 +
 .../@node-red/editor-client/src/js/nodes.js   | 93 ++++++++++++++++---
 .../editor-client/src/js/ui/view-tools.js     | 37 ++++++++
 3 files changed, 117 insertions(+), 15 deletions(-)

diff --git a/packages/node_modules/@node-red/editor-client/src/js/keymap.json b/packages/node_modules/@node-red/editor-client/src/js/keymap.json
index 5d8930ebf..aa357a832 100644
--- a/packages/node_modules/@node-red/editor-client/src/js/keymap.json
+++ b/packages/node_modules/@node-red/editor-client/src/js/keymap.json
@@ -44,6 +44,8 @@
         "ctrl-z": "core:undo",
         "ctrl-y": "core:redo",
         "ctrl-a": "core:select-all-nodes",
+        "alt-s u": "core:select-upstream-nodes",
+        "alt-s d": "core:select-downstream-nodes",
         "shift-?": "core:show-help",
         "w": "core:scroll-view-up",
         "d": "core:scroll-view-right",
diff --git a/packages/node_modules/@node-red/editor-client/src/js/nodes.js b/packages/node_modules/@node-red/editor-client/src/js/nodes.js
index 236f7fa57..d39fe48ff 100644
--- a/packages/node_modules/@node-red/editor-client/src/js/nodes.js
+++ b/packages/node_modules/@node-red/editor-client/src/js/nodes.js
@@ -18,9 +18,11 @@ RED.nodes = (function() {
     var node_defs = {};
     var nodes = {};
     var nodeTabMap = {};
+    var linkTabMap = {};
 
     var configNodes = {};
     var links = [];
+    var nodeLinks = {};
     var defaultWorkspace;
     var workspaces = {};
     var workspacesOrder =[];
@@ -33,7 +35,7 @@ RED.nodes = (function() {
     var initialLoad;
 
     var dirty = false;
-
+RED.ll = function() { return nodeLinks}
     function setDirty(d) {
         dirty = d;
         RED.events.emit("workspace:dirty",{dirty:dirty});
@@ -286,6 +288,9 @@ RED.nodes = (function() {
                 n.i = nextId+1;
             }
             nodes[n.id] = n;
+            if (!nodeLinks[n.id]) {
+                nodeLinks[n.id] = {in:[],out:[]};
+            }
             if (nodeTabMap[n.z]) {
                 nodeTabMap[n.z][n.id] = n;
             } else {
@@ -296,6 +301,22 @@ RED.nodes = (function() {
     }
     function addLink(l) {
         links.push(l);
+        if (l.source) {
+            // Possible the node hasn't been added yet
+            if (!nodeLinks[l.source.id]) {
+                nodeLinks[l.source.id] = {in:[],out:[]};
+            }
+            nodeLinks[l.source.id].out.push(l);
+        }
+        if (l.target) {
+            if (!nodeLinks[l.target.id]) {
+                nodeLinks[l.target.id] = {in:[],out:[]};
+            }
+            nodeLinks[l.target.id].in.push(l);
+        }
+        if (l.source.z === l.target.z && linkTabMap[l.source.z]) {
+            linkTabMap[l.source.z].push(l);
+        }
         RED.events.emit("links:add",l);
     }
 
@@ -320,6 +341,7 @@ RED.nodes = (function() {
         } else if (id in nodes) {
             node = nodes[id];
             delete nodes[id]
+            delete nodeLinks[id];
             if (nodeTabMap[node.z]) {
                 delete nodeTabMap[node.z][node.id];
             }
@@ -407,6 +429,24 @@ RED.nodes = (function() {
         var index = links.indexOf(l);
         if (index != -1) {
             links.splice(index,1);
+            if (l.source) {
+                var sIndex = nodeLinks[l.source.id].out.indexOf(l)
+                if (sIndex !== -1) {
+                    nodeLinks[l.source.id].out.splice(sIndex,1)
+                }
+            }
+            if (l.target) {
+                var tIndex = nodeLinks[l.target.id].in.indexOf(l)
+                if (tIndex !== -1) {
+                    nodeLinks[l.target.id].in.splice(tIndex,1)
+                }
+            }
+            if (l.source.z === l.target.z && linkTabMap[l.source.z]) {
+                index = linkTabMap[l.source.z].indexOf(l);
+                if (index !== -1) {
+                    linkTabMap[l.source.z].splice(index,1)
+                }
+            }
         }
         RED.events.emit("links:remove",l);
     }
@@ -414,6 +454,7 @@ RED.nodes = (function() {
     function addWorkspace(ws,targetIndex) {
         workspaces[ws.id] = ws;
         nodeTabMap[ws.id] = {};
+        linkTabMap[ws.id] = [];
 
         ws._def = RED.nodes.getType('tab');
         if (targetIndex === undefined) {
@@ -437,6 +478,7 @@ RED.nodes = (function() {
         if (ws) {
             delete workspaces[id];
             delete nodeTabMap[id];
+            delete linkTabMap[id];
             workspacesOrder.splice(workspacesOrder.indexOf(id),1);
             var i;
             var node;
@@ -502,6 +544,7 @@ RED.nodes = (function() {
         }
         subflows[sf.id] = sf;
         nodeTabMap[sf.id] = {};
+        linkTabMap[sf.id] = [];
 
         RED.nodes.registerType("subflow:"+sf.id, {
             defaults:{
@@ -581,7 +624,7 @@ RED.nodes = (function() {
         var stack = [node];
         while(stack.length !== 0) {
             var n = stack.shift();
-            var childLinks = links.filter(function(d) { return (d.source === n) || (d.target === n);});
+            var childLinks = [].concat(nodeLinks[n.id].in).concat(nodeLinks[n.id].out);
             for (var i=0;i<childLinks.length;i++) {
                 var child = (childLinks[i].source === n)?childLinks[i].target:childLinks[i].source;
                 var id = child.id;
@@ -1884,9 +1927,37 @@ RED.nodes = (function() {
     }
     function filterLinks(filter) {
         var result = [];
-
-        for (var n=0;n<links.length;n++) {
-            var link = links[n];
+        var candidateLinks = [];
+        var hasCandidates = false;
+        var filterSZ = filter.source && filter.source.z;
+        var filterTZ = filter.target && filter.target.z;
+        var filterZ;
+        if (filterSZ || filterTZ) {
+            if (filterSZ === filterTZ) {
+                filterZ = filterSZ;
+            } else {
+                filterZ = (filterSZ === undefined)?filterTZ:filterSZ
+            }
+        }
+        if (filterZ) {
+            candidateLinks = linkTabMap[filterZ] || [];
+            hasCandidates = true;
+        } else if (filter.source && filter.source.hasOwnProperty("id")) {
+            if (nodeLinks[filter.source.id]) {
+                hasCandidates = true;
+                candidateLinks = candidateLinks.concat(nodeLinks[filter.source.id].out)
+            }
+        } else if (filter.target && filter.target.hasOwnProperty("id")) {
+            if (nodeLinks[filter.target.id]) {
+                hasCandidates = true;
+                candidateLinks = candidateLinks.concat(nodeLinks[filter.target.id].in)
+            }
+        }
+        if (!hasCandidates) {
+            candidateLinks = links;
+        }
+        for (var n=0;n<candidateLinks.length;n++) {
+            var link = candidateLinks[n];
             if (filter.source) {
                 if (filter.source.hasOwnProperty("id") && link.source.id !== filter.source.id) {
                     continue;
@@ -1943,6 +2014,8 @@ RED.nodes = (function() {
         nodes = {};
         links = [];
         nodeTabMap = {};
+        linkTabMap = {};
+        nodeLinks = {};
         configNodes = {};
         workspacesOrder = [];
         groups = {};
@@ -1968,16 +2041,6 @@ RED.nodes = (function() {
         RED.sidebar.info.refresh();
 
         RED.events.emit("workspace:clear");
-
-        // var node_defs = {};
-        // var nodes = {};
-        // var configNodes = {};
-        // var links = [];
-        // var defaultWorkspace;
-        // var workspaces = {};
-        // var workspacesOrder =[];
-        // var subflows = {};
-        // var loadedFlowVersion = null;
     }
 
     function addGroup(group) {
diff --git a/packages/node_modules/@node-red/editor-client/src/js/ui/view-tools.js b/packages/node_modules/@node-red/editor-client/src/js/ui/view-tools.js
index e25b570b9..27078538d 100644
--- a/packages/node_modules/@node-red/editor-client/src/js/ui/view-tools.js
+++ b/packages/node_modules/@node-red/editor-client/src/js/ui/view-tools.js
@@ -17,6 +17,40 @@
 RED.view.tools = (function() {
 
 
+    function selectDownstream() {
+        selectStream('source','target')
+    }
+    function selectUpstream() {
+        selectStream('target','source')
+    }
+    function selectStream(from,to) {
+        var selection = RED.view.selection();
+        var visited = new Set();
+        if (selection.nodes && selection.nodes.length > 0) {
+            var nodes = [];
+            selection.nodes.forEach(function(n) {
+                visited.add(n);
+                var filter = {};
+                filter[from] = {id:n.id}
+                var initialLinks = RED.nodes.filterLinks(filter);
+                nodes = nodes.concat(initialLinks.map(function(n){return n[to]}));
+            })
+            while(nodes.length > 0) {
+                var n = nodes.shift();
+                visited.add(n);
+                var filter = {};
+                filter[from] = {id:n.id}
+                var links = RED.nodes.filterLinks(filter);
+                links.forEach(function(l) {
+                    if (!visited.has(l[to])) {
+                        nodes.push(l[to]);
+                    }
+                })
+            }
+            RED.view.select({nodes:Array.from(visited)});
+        }
+    }
+
     function alignToGrid() {
         var selection = RED.view.selection();
         if (selection.nodes) {
@@ -206,6 +240,9 @@ RED.view.tools = (function() {
             RED.actions.add("core:step-selection-right", function() { moveSelection(RED.view.gridSize(),0);});
             RED.actions.add("core:step-selection-down", function() { moveSelection(0,RED.view.gridSize());});
             RED.actions.add("core:step-selection-left", function() { moveSelection(-RED.view.gridSize(),0);});
+
+            RED.actions.add("core:select-downstream-nodes", selectDownstream);
+            RED.actions.add("core:select-upstream-nodes", selectUpstream);
         },
         /**
          * Aligns all selected nodes to the current grid