dag.ts 4.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214
  1. export class Edge {
  2. inputNode?: Node;
  3. outputNode?: Node;
  4. _linkTo(node: Node, direction: number) {
  5. if (direction <= 0) {
  6. node.inputEdges.push(this);
  7. }
  8. if (direction >= 0) {
  9. node.outputEdges.push(this);
  10. }
  11. node.edges.push(this);
  12. }
  13. link(inputNode: Node, outputNode: Node) {
  14. if (!inputNode) {
  15. throw Error('inputNode is required');
  16. }
  17. if (!outputNode) {
  18. throw Error('outputNode is required');
  19. }
  20. this.unlink();
  21. this.inputNode = inputNode;
  22. this.outputNode = outputNode;
  23. this._linkTo(inputNode, 1);
  24. this._linkTo(outputNode, -1);
  25. return this;
  26. }
  27. unlink() {
  28. let pos;
  29. const inode = this.inputNode;
  30. const onode = this.outputNode;
  31. if (!(inode && onode)) {
  32. return;
  33. }
  34. pos = inode.edges.indexOf(this);
  35. if (pos > -1) {
  36. inode.edges.splice(pos, 1);
  37. }
  38. pos = onode.edges.indexOf(this);
  39. if (pos > -1) {
  40. onode.edges.splice(pos, 1);
  41. }
  42. pos = inode.outputEdges.indexOf(this);
  43. if (pos > -1) {
  44. inode.outputEdges.splice(pos, 1);
  45. }
  46. pos = onode.inputEdges.indexOf(this);
  47. if (pos > -1) {
  48. onode.inputEdges.splice(pos, 1);
  49. }
  50. }
  51. }
  52. export class Node {
  53. name: string;
  54. edges: Edge[];
  55. inputEdges: Edge[];
  56. outputEdges: Edge[];
  57. constructor(name: string) {
  58. this.name = name;
  59. this.edges = [];
  60. this.inputEdges = [];
  61. this.outputEdges = [];
  62. }
  63. getEdgeFrom(from: string | Node): Edge | null | undefined {
  64. if (!from) {
  65. return null;
  66. }
  67. if (typeof from === 'object') {
  68. return this.inputEdges.find((e) => e.inputNode?.name === from.name);
  69. }
  70. return this.inputEdges.find((e) => e.inputNode?.name === from);
  71. }
  72. getEdgeTo(to: string | Node): Edge | null | undefined {
  73. if (!to) {
  74. return null;
  75. }
  76. if (typeof to === 'object') {
  77. return this.outputEdges.find((e) => e.outputNode?.name === to.name);
  78. }
  79. return this.outputEdges.find((e) => e.outputNode?.name === to);
  80. }
  81. getOptimizedInputEdges(): Edge[] {
  82. const toBeRemoved: any[] = [];
  83. this.inputEdges.forEach((e) => {
  84. const inputEdgesNodes = e.inputNode?.inputEdges.map((e) => e.inputNode);
  85. inputEdgesNodes?.forEach((n) => {
  86. const edgeToRemove = n?.getEdgeTo(this.name);
  87. if (edgeToRemove) {
  88. toBeRemoved.push(edgeToRemove);
  89. }
  90. });
  91. });
  92. return this.inputEdges.filter((e) => toBeRemoved.indexOf(e) === -1);
  93. }
  94. }
  95. export class Graph {
  96. nodes: any = {};
  97. constructor() {}
  98. createNode(name: string): Node {
  99. const n = new Node(name);
  100. this.nodes[name] = n;
  101. return n;
  102. }
  103. createNodes(names: string[]): Node[] {
  104. const nodes: Node[] = [];
  105. names.forEach((name) => {
  106. nodes.push(this.createNode(name));
  107. });
  108. return nodes;
  109. }
  110. link(input: string | string[] | Node | Node[], output: string | string[] | Node | Node[]): Edge[] {
  111. let inputArr = [];
  112. let outputArr = [];
  113. const inputNodes: Node[] = [];
  114. const outputNodes: Node[] = [];
  115. if (input instanceof Array) {
  116. inputArr = input;
  117. } else {
  118. inputArr = [input];
  119. }
  120. if (output instanceof Array) {
  121. outputArr = output;
  122. } else {
  123. outputArr = [output];
  124. }
  125. for (let n = 0; n < inputArr.length; n++) {
  126. const i = inputArr[n];
  127. if (typeof i === 'string') {
  128. const n = this.getNode(i);
  129. if (!n) {
  130. throw Error(`cannot link input node named ${i} since it doesn't exist in graph`);
  131. }
  132. inputNodes.push(n);
  133. } else {
  134. inputNodes.push(i);
  135. }
  136. }
  137. for (let n = 0; n < outputArr.length; n++) {
  138. const i = outputArr[n];
  139. if (typeof i === 'string') {
  140. const n = this.getNode(i);
  141. if (!n) {
  142. throw Error(`cannot link output node named ${i} since it doesn't exist in graph`);
  143. }
  144. outputNodes.push(n);
  145. } else {
  146. outputNodes.push(i);
  147. }
  148. }
  149. const edges: Edge[] = [];
  150. inputNodes.forEach((input) => {
  151. outputNodes.forEach((output) => {
  152. edges.push(this.createEdge().link(input, output));
  153. });
  154. });
  155. return edges;
  156. }
  157. createEdge(): Edge {
  158. return new Edge();
  159. }
  160. getNode(name: string): Node {
  161. return this.nodes[name];
  162. }
  163. }
  164. export const printGraph = (g: Graph) => {
  165. Object.keys(g.nodes).forEach((name) => {
  166. const n = g.nodes[name];
  167. let outputEdges = n.outputEdges.map((e: Edge) => e.outputNode?.name).join(', ');
  168. if (!outputEdges) {
  169. outputEdges = '<none>';
  170. }
  171. let inputEdges = n.inputEdges.map((e: Edge) => e.inputNode?.name).join(', ');
  172. if (!inputEdges) {
  173. inputEdges = '<none>';
  174. }
  175. console.log(`${n.name}:\n - links to: ${outputEdges}\n - links from: ${inputEdges}`);
  176. });
  177. };