Refactor graph traversal methods to use new computeTraversalPartitions API

iidm/iidm-impl/src/main/java/com/powsybl/iidm/network/impl/BusBreakerTopologyModel.java

@@ -16,10 +16,7 @@
 import com.powsybl.iidm.network.util.Identifiables;
 import com.powsybl.iidm.network.util.Networks;
 import com.powsybl.iidm.network.util.ShortIdDictionary;
-import com.powsybl.math.graph.TraversalType;
-import com.powsybl.math.graph.TraverseResult;
-import com.powsybl.math.graph.UndirectedGraphImpl;
-import com.powsybl.math.graph.UndirectedGraphListener;
+import com.powsybl.math.graph.*;
 import org.anarres.graphviz.builder.GraphVizAttribute;
 import org.anarres.graphviz.builder.GraphVizEdge;
 import org.anarres.graphviz.builder.GraphVizGraph;
@@ -34,6 +31,7 @@
 import java.nio.file.Path;
 import java.security.SecureRandom;
 import java.util.*;
+import java.util.concurrent.atomic.AtomicInteger;
 import java.util.function.Function;
 import java.util.function.Predicate;
 import java.util.stream.Collectors;
@@ -230,30 +228,31 @@ private void updateCache() {
 
             // mapping between configured buses and merged buses
             Map<ConfiguredBus, MergedBus> mapping = new IdentityHashMap<>();
+            AtomicInteger busNum = new AtomicInteger(0);
+
+            graph.getConnectedComponents(
+                sw -> !sw.isOpen(),
+                new UndirectedGraph.ConnectedComponentCollector<Set<ConfiguredBus>, ConfiguredBus>() {
+                    @Override
+                    public Set<ConfiguredBus> createComponent() {
+                        return new LinkedHashSet<>(1);
+                    }
+
+                    @Override
+                    public void addVertex(Set<ConfiguredBus> component, int vertexIndex, ConfiguredBus bus) {
+                        component.add(bus);
+                    }
 
-            boolean[] encountered = new boolean[graph.getVertexCapacity()];
-            Arrays.fill(encountered, false);
-            int busNum = 0;
-            for (int v : graph.getVertices()) {
-                if (!encountered[v]) {
-                    final Set<ConfiguredBus> busSet = new LinkedHashSet<>(1);
-                    busSet.add(graph.getVertexObject(v));
-                    graph.traverse(v, TraversalType.DEPTH_FIRST, (v1, e, v2) -> {
-                        SwitchImpl aSwitch = graph.getEdgeObject(e);
-                        if (aSwitch.isOpen()) {
-                            return TraverseResult.TERMINATE_PATH;
-                        } else {
-                            busSet.add(graph.getVertexObject(v2));
-                            return TraverseResult.CONTINUE;
-                        }
-                    }, encountered);
-                    if (isBusValid(busSet)) {
-                        MergedBus mergedBus = createMergedBus(busNum++, busSet);
-                        mergedBuses.put(mergedBus.getId(), mergedBus);
-                        busSet.forEach(bus -> mapping.put(bus, mergedBus));
+                    @Override
+                    public boolean isComponentValid(Set<ConfiguredBus> component) {
+                        return isBusValid(component);
                     }
                 }
-            }
+            ).forEach(busSet -> {
+                MergedBus mergedBus = createMergedBus(busNum.getAndIncrement(), busSet);
+                mergedBuses.put(mergedBus.getId(), mergedBus);
+                busSet.forEach(bus -> mapping.put(bus, mergedBus));
+            });
 
             variants.get().cache = new BusCache(mergedBuses, mapping);
         }
@@ -293,7 +292,7 @@ MergedBus getMergedBus(ConfiguredBus bus) {
     }
 
     final CalculatedBusTopology calculatedBusTopology
-            = new CalculatedBusTopology();
+        = new CalculatedBusTopology();
 
     private static final class VariantImpl implements Variant {
 

iidm/iidm-impl/src/main/java/com/powsybl/iidm/network/impl/DcTopologyModel.java

@@ -13,10 +13,7 @@
 import com.powsybl.commons.ref.RefChain;
 import com.powsybl.iidm.network.*;
 import com.powsybl.iidm.network.util.Identifiables;
-import com.powsybl.math.graph.TraversalType;
-import com.powsybl.math.graph.TraverseResult;
-import com.powsybl.math.graph.UndirectedGraphImpl;
-import com.powsybl.math.graph.UndirectedGraphListener;
+import com.powsybl.math.graph.*;
 
 import java.util.*;
 import java.util.function.Function;
@@ -239,29 +236,31 @@ private DcBusCache getCache() {
             // mapping between DC nodes ID and DC buses
             Map<String, DcBusImpl> dcNodeIdToDcBus = new HashMap<>();
 
-            boolean[] encountered = new boolean[graph.getVertexCapacity()];
-            Arrays.fill(encountered, false);
-            for (int v : graph.getVertices()) {
-                if (!encountered[v]) {
-                    final Set<DcNodeImpl> dcNodeSet = new LinkedHashSet<>(1);
-                    dcNodeSet.add(graph.getVertexObject(v));
-                    graph.traverse(v, TraversalType.DEPTH_FIRST, (v1, e, v2) -> {
-                        DcSwitchImpl dcSwitch = graph.getEdgeObject(e);
-                        if (dcSwitch.isOpen()) {
-                            return TraverseResult.TERMINATE_PATH;
-                        } else {
-                            dcNodeSet.add(graph.getVertexObject(v2));
-                            return TraverseResult.CONTINUE;
-                        }
-                    }, encountered);
-
-                    if (isDcBusValid(dcNodeSet)) {
-                        DcBusImpl dcBus = createDcBus(dcNodeSet);
-                        dcBuses.put(dcBus.getId(), dcBus);
-                        dcNodeSet.forEach(dcNode -> dcNodeIdToDcBus.put(dcNode.getId(), dcBus));
+            graph.getConnectedComponents(
+                sw -> !sw.isOpen(),
+                new UndirectedGraph.ConnectedComponentCollector<Set<DcNodeImpl>, DcNodeImpl>() {
+
+                    @Override
+                    public Set<DcNodeImpl> createComponent() {
+                        return new LinkedHashSet<>(1);
+                    }
+
+                    @Override
+                    public void addVertex(Set<DcNodeImpl> component, int vertexIndex, DcNodeImpl dcNode) {
+                        component.add(dcNode);
+                    }
+
+                    @Override
+                    public boolean isComponentValid(Set<DcNodeImpl> component) {
+                        return isDcBusValid(component);
                     }
                 }
-            }
+
+            ).forEach(dcNodeSet -> {
+                DcBusImpl dcBus = createDcBus(dcNodeSet);
+                dcBuses.put(dcBus.getId(), dcBus);
+                dcNodeSet.forEach(dcNode -> dcNodeIdToDcBus.put(dcNode.getId(), dcBus));
+            });
 
             variants.get().cache = new DcBusCache(dcBuses, dcNodeIdToDcBus);
             return variants.get().cache;

iidm/iidm-impl/src/main/java/com/powsybl/iidm/network/impl/NodeBreakerTopologyModel.java

@@ -259,41 +259,6 @@ protected BusChecker getBusChecker() {
             return CALCULATED_BUS_CHECKER;
         }
 
-        private void traverse(int n, boolean[] encountered, Predicate<SwitchImpl> terminate, Map<String, CalculatedBus> id2bus, CalculatedBus[] node2bus) {
-            if (!encountered[n]) {
-                final TIntArrayList nodes = new TIntArrayList(1);
-                nodes.add(n);
-                Traverser traverser = (n1, e, n2) -> {
-                    SwitchImpl aSwitch = graph.getEdgeObject(e);
-                    if (aSwitch != null && terminate.test(aSwitch)) {
-                        return TraverseResult.TERMINATE_PATH;
-                    }
-
-                    if (!encountered[n2]) {
-                        // We need to check this as the traverser might be called twice with the same n2 but with different edges.
-                        // Note that the "encountered" array is used and maintained inside graph::traverse method, hence we should not update it.
-                        nodes.add(n2);
-                    }
-                    return TraverseResult.CONTINUE;
-                };
-                graph.traverse(n, TraversalType.DEPTH_FIRST, traverser, encountered);
-
-                // check that the component is a bus
-                String busId = Identifiables.getUniqueId(NAMING_STRATEGY.getId(voltageLevel, nodes), getNetwork().getIndex()::contains);
-                CopyOnWriteArrayList<NodeTerminal> terminals = new CopyOnWriteArrayList<>();
-                for (int i = 0; i < nodes.size(); i++) {
-                    int n2 = nodes.getQuick(i);
-                    NodeTerminal terminal2 = graph.getVertexObject(n2);
-                    if (terminal2 != null) {
-                        terminals.add(terminal2);
-                    }
-                }
-                if (getBusChecker().isValid(graph, nodes, terminals)) {
-                    addBus(nodes, id2bus, node2bus, busId, terminals);
-                }
-            }
-        }
-
         private void addBus(TIntArrayList nodes, Map<String, CalculatedBus> id2bus, CalculatedBus[] node2bus,
                             String busId, CopyOnWriteArrayList<NodeTerminal> terminals) {
             String busName = NAMING_STRATEGY.getName(voltageLevel, nodes);
@@ -312,10 +277,49 @@ protected void updateCache(final Predicate<SwitchImpl> terminate) {
             LOGGER.trace("Update bus topology of voltage level {}", voltageLevel.getId());
             Map<String, CalculatedBus> id2bus = new LinkedHashMap<>();
             CalculatedBus[] node2bus = new CalculatedBus[graph.getVertexCapacity()];
-            boolean[] encountered = new boolean[graph.getVertexCapacity()];
-            for (int v : graph.getVertices()) {
-                traverse(v, encountered, terminate, id2bus, node2bus);
-            }
+
+            graph.getConnectedComponents(
+                sw -> !terminate.test(sw),
+                new UndirectedGraph.ConnectedComponentCollector<TIntArrayList, NodeTerminal>() {
+                    @Override
+                    public TIntArrayList createComponent() {
+                        return new TIntArrayList(1);
+                    }
+
+                    @Override
+                    public void addVertex(TIntArrayList nodes, int nodeIndex, NodeTerminal terminal) {
+                        nodes.add(nodeIndex);
+                    }
+
+                    @Override
+                    public boolean isComponentValid(TIntArrayList nodes) {
+                        CopyOnWriteArrayList<NodeTerminal> terminals = new CopyOnWriteArrayList<>();
+                        for (int i = 0; i < nodes.size(); i++) {
+                            NodeTerminal terminal = graph.getVertexObject(nodes.getQuick(i));
+                            if (terminal != null) {
+                                terminals.add(terminal);
+                            }
+                        }
+                        return getBusChecker().isValid(graph, nodes, terminals);
+                    }
+                }
+            ).forEach(nodes -> {
+                String busId = Identifiables.getUniqueId(
+                    NAMING_STRATEGY.getId(voltageLevel, nodes),
+                    getNetwork().getIndex()::contains
+                );
+                CopyOnWriteArrayList<NodeTerminal> terminals = new CopyOnWriteArrayList<>();
+                for (int i = 0; i < nodes.size(); i++) {
+                    int nodeIndex = nodes.getQuick(i);
+                    NodeTerminal terminal = graph.getVertexObject(nodeIndex);
+                    if (terminal != null) {
+                        terminals.add(terminal);
+                    }
+                }
+
+                addBus(nodes, id2bus, node2bus, busId, terminals);
+            });
+
             busCache = new BusCache(node2bus, id2bus);
             LOGGER.trace("Found buses {}", id2bus.values());
         }

math/src/main/java/com/powsybl/math/graph/UndirectedGraph.java

@@ -331,6 +331,18 @@ public interface UndirectedGraph<V, E> {
      */
     boolean traverse(int v, TraversalType traversalType, Traverser traverser, boolean[] verticesEncountered);
 
+    <C> List<C> getConnectedComponents(Predicate<E> isTraversable, ConnectedComponentCollector<C, V> collector);
+
+    interface ConnectedComponentCollector<C, V> {
+        C createComponent();
+
+        void addVertex(C component, int vertexIndex, V vertexObject);
+
+        default boolean isComponentValid(C component) {
+            return true;
+        }
+    }
+
     /**
      * Traverse the entire graph, starting at the specified vertex v.
      * This method allocates a boolean array and calls {@link #traverse(int, TraversalType, Traverser, boolean[])}.

math/src/main/java/com/powsybl/math/graph/UndirectedGraphImpl.java

@@ -579,6 +579,43 @@ public boolean traverse(int v, TraversalType traversalType, Traverser traverser,
         return keepGoing;
     }
 
+    @Override
+    public <C> List<C> getConnectedComponents(Predicate<E> isTraversable, ConnectedComponentCollector<C, V> collector) {
+        Objects.requireNonNull(isTraversable);
+        Objects.requireNonNull(collector);
+
+        List<C> components = new ArrayList<>();
+        boolean[] encountered = new boolean[vertices.size()];
+
+        for (int v = 0; v < vertices.size(); v++) {
+            Vertex<V> vertex = vertices.get(v);
+            if (vertex != null && !encountered[v]) {
+                C component = collector.createComponent();
+
+                collector.addVertex(component, v, vertex.getObject());
+
+                traverse(v, TraversalType.BREADTH_FIRST, (v1, e, v2) -> {
+                    E edgeObject = edges.get(e).getObject();
+                    if (edgeObject == null || isTraversable.test(edgeObject)) {
+                        if (!encountered[v2]) {
+                            Vertex<V> destVertex = vertices.get(v2);
+                            if (destVertex != null) {
+                                collector.addVertex(component, v2, destVertex.getObject());
+                            }
+                        }
+                        return TraverseResult.CONTINUE;
+                    }
+                    return TraverseResult.TERMINATE_PATH;
+                }, encountered);
+
+                if (collector.isComponentValid(component)) {
+                    components.add(component);
+                }
+            }
+        }
+        return components;
+    }
+
     @Override
     public boolean traverse(int v, TraversalType traversalType, Traverser traverser) {
         boolean[] encountered = new boolean[vertices.size()];