From 9ea2ce666f8eb99fe97b43d57987943a740c7c2f Mon Sep 17 00:00:00 2001
From: Sophie Frasnedo <sophie.frasnedo@rte-france.com>
Date: Thu, 7 Dec 2023 14:39:44 +0100
Subject: [PATCH] Refactor traverse method to avoid code duplication

Signed-off-by: Sophie Frasnedo <sophie.frasnedo@rte-france.com>
---
 .../nad/build/iidm/VoltageLevelFilter.java    | 30 ++++++-------------
 1 file changed, 9 insertions(+), 21 deletions(-)

diff --git a/network-area-diagram/src/main/java/com/powsybl/nad/build/iidm/VoltageLevelFilter.java b/network-area-diagram/src/main/java/com/powsybl/nad/build/iidm/VoltageLevelFilter.java
index bbe10b515..39b66572c 100644
--- a/network-area-diagram/src/main/java/com/powsybl/nad/build/iidm/VoltageLevelFilter.java
+++ b/network-area-diagram/src/main/java/com/powsybl/nad/build/iidm/VoltageLevelFilter.java
@@ -91,41 +91,29 @@ public static Collection<VoltageLevel> getNextDepthVoltageLevels(Network network
         return voltageLevelSet;
     }
 
-    private static void traverseVoltageLevels(Set<VoltageLevel> voltageLevelsDepth, int depth, Set<VoltageLevel> visitedVoltageLevels) {
+    private static void traverseVoltageLevels(Set<VoltageLevel> voltageLevelsDepth, int depth, Set<VoltageLevel> visitedVoltageLevels, Predicate<VoltageLevel> predicate) {
         if (depth < 0) {
             return;
         }
         Set<VoltageLevel> nextDepthVoltageLevels = new HashSet<>();
         for (VoltageLevel vl : voltageLevelsDepth) {
-            if (!visitedVoltageLevels.contains(vl)) {
+            if (!visitedVoltageLevels.contains(vl) && predicate.test(vl)) {
                 visitedVoltageLevels.add(vl);
                 vl.visitEquipments(new VlVisitor(nextDepthVoltageLevels, visitedVoltageLevels));
             }
         }
-        traverseVoltageLevels(nextDepthVoltageLevels, depth - 1, visitedVoltageLevels);
+        traverseVoltageLevels(nextDepthVoltageLevels, depth - 1, visitedVoltageLevels, predicate);
+    }
+
+    private static void traverseVoltageLevels(Set<VoltageLevel> voltageLevelsDepth, int depth, Set<VoltageLevel> visitedVoltageLevels) {
+        traverseVoltageLevels(voltageLevelsDepth, depth, visitedVoltageLevels, NO_FILTER);
     }
 
     private static void traverseVoltageLevelsWithVoltageFilter(Set<VoltageLevel> voltageLevelsDepth, int depth, Set<VoltageLevel> visitedVoltageLevels,
                                                                double lowNominalVoltageBound, double highNominalVoltageBound) {
-        if (depth < 0) {
-            return;
-        }
-
-        Set<VoltageLevel> nextDepthVoltageLevels = new HashSet<>();
-        for (VoltageLevel vl : voltageLevelsDepth) {
-            if (!visitedVoltageLevels.contains(vl)
-                    && vl.getNominalV() >= lowNominalVoltageBound
-                    && vl.getNominalV() <= highNominalVoltageBound) {
-                traverseVoltageLevel(visitedVoltageLevels, nextDepthVoltageLevels, vl);
-            }
-        }
-        traverseVoltageLevelsWithVoltageFilter(nextDepthVoltageLevels, depth - 1, visitedVoltageLevels, lowNominalVoltageBound, highNominalVoltageBound);
-
-    }
 
-    private static void traverseVoltageLevel(Set<VoltageLevel> visitedVoltageLevels, Set<VoltageLevel> nextDepthVoltageLevels, VoltageLevel vl) {
-        visitedVoltageLevels.add(vl);
-        vl.visitEquipments(new VlVisitor(nextDepthVoltageLevels, visitedVoltageLevels));
+        Predicate<VoltageLevel> voltageFilterPredicate = voltageLevel -> voltageLevel.getNominalV() >= lowNominalVoltageBound && voltageLevel.getNominalV() <= highNominalVoltageBound;
+        traverseVoltageLevels(voltageLevelsDepth, depth, visitedVoltageLevels, voltageFilterPredicate);
     }
 
     public static VoltageLevelFilter createNominalVoltageFilter(Network network, List<String> voltageLevelIds,