private DAG groupLoops()

in compiler/optimizer/src/main/java/org/apache/nemo/compiler/optimizer/pass/compiletime/reshaping/LoopExtractionPass.java [86:144]
51 lines of code
12 McCabe index (conditional complexity)

  private DAG<IRVertex, IREdge> groupLoops(final DAG<IRVertex, IREdge> dag, final Integer depth) {
    if (depth <= 0) {
      return dag;
    } else {
      final DAGBuilder<IRVertex, IREdge> builder = new DAGBuilder<>();
      for (IRVertex irVertex : dag.getTopologicalSort()) {
        if (irVertex instanceof SourceVertex) { // Source vertex: no incoming edges
          if (dag.isCompositeVertex(irVertex) && dag.getLoopStackDepthOf(irVertex).equals(depth)) {
            // when src is inside a loop
            final LoopVertex assignedLoopVertex = dag.getAssignedLoopVertexOf(irVertex);
            builder.addVertex(assignedLoopVertex, dag);
            connectElementToLoop(dag, builder, irVertex, assignedLoopVertex);
          } else {
            builder.addVertex(irVertex, dag);
          }

        } else if (irVertex instanceof OperatorVertex) { // Operator vertex
          final OperatorVertex operatorVertex = (OperatorVertex) irVertex;
          // If this is Composite && depth is appropriate. == If this belongs to a loop.
          if (dag.isCompositeVertex(operatorVertex) && dag.getLoopStackDepthOf(operatorVertex).equals(depth)) {
            final LoopVertex assignedLoopVertex = dag.getAssignedLoopVertexOf(operatorVertex);
            builder.addVertex(assignedLoopVertex, dag);
            connectElementToLoop(dag, builder, operatorVertex, assignedLoopVertex); // something -> loop
          } else { // Otherwise: it is not composite || depth inappropriate. == If this is just an operator.
            builder.addVertex(operatorVertex, dag);
            dag.getIncomingEdgesOf(operatorVertex).forEach(irEdge -> {
              if (dag.isCompositeVertex(irEdge.getSrc())) {
                // connecting with a loop: loop -> operator.
                final LoopVertex srcLoopVertex = dag.getAssignedLoopVertexOf(irEdge.getSrc());
                srcLoopVertex.addDagOutgoingEdge(irEdge);
                final IREdge edgeFromLoop = new IREdge(
                  irEdge.getPropertyValue(CommunicationPatternProperty.class).get(), srcLoopVertex, operatorVertex);
                irEdge.copyExecutionPropertiesTo(edgeFromLoop);
                builder.connectVertices(edgeFromLoop);
                srcLoopVertex.mapEdgeWithLoop(edgeFromLoop, irEdge);
              } else { // connecting outside the composite loop: operator -> operator.
                builder.connectVertices(irEdge);
              }
            });
          }

        } else if (irVertex instanceof LoopVertex) { // Loop vertices of higher depth (nested loops).
          final LoopVertex loopVertex = (LoopVertex) irVertex;
          if (dag.isCompositeVertex(loopVertex)) { // the loopVertex belongs to another loop.
            final LoopVertex assignedLoopVertex = dag.getAssignedLoopVertexOf(loopVertex);
            connectElementToLoop(dag, builder, loopVertex, assignedLoopVertex); // something -> loop
          } else { // it cannot be just at the operator level, as it had more depth.
            throw new UnsupportedOperationException("This loop (" + loopVertex + ") shouldn't be of this depth");
          }

        } else {
          throw new UnsupportedOperationException("Unknown vertex type: " + irVertex);
        }
      }

      // Recursive calls for lower depths.
      return groupLoops(loopRolling(builder.build()), depth - 1);
    }
  }