From c4d64058a908f42f1858bd967290b674a09538d7 Mon Sep 17 00:00:00 2001
From: Geoffroy Jamgotchian <geoffroy.jamgotchian@rte-france.com>
Date: Fri, 30 Jun 2023 09:43:11 +0200
Subject: [PATCH] Clean indexed terms when removing an equation (#808)

Signed-off-by: Geoffroy Jamgotchian <geoffroy.jamgotchian@rte-france.com>
---
 .../openloadflow/equations/EquationSystem.java  | 17 +++++++++++++++++
 .../equations/EquationSystemTest.java           |  8 ++++++++
 2 files changed, 25 insertions(+)

diff --git a/src/main/java/com/powsybl/openloadflow/equations/EquationSystem.java b/src/main/java/com/powsybl/openloadflow/equations/EquationSystem.java
index 6d1bde8ef3..efc9e023d9 100644
--- a/src/main/java/com/powsybl/openloadflow/equations/EquationSystem.java
+++ b/src/main/java/com/powsybl/openloadflow/equations/EquationSystem.java
@@ -146,12 +146,29 @@ public boolean hasEquation(int num, E type) {
         return equations.containsKey(p);
     }
 
+    private void deindexTerm(EquationTerm<V, E> term) {
+        if (term.getElementType() != null && term.getElementNum() != -1) {
+            List<EquationTerm<V, E>> termsForThisElement = equationTermsByElement.get(Pair.of(term.getElementType(), term.getElementNum()));
+            if (termsForThisElement != null) {
+                termsForThisElement.remove(term);
+            }
+        }
+        for (EquationTerm<V, E> child : term.getChildren()) {
+            deindexTerm(child);
+        }
+    }
+
     public Equation<V, E> removeEquation(int num, E type) {
         Pair<Integer, E> p = Pair.of(num, type);
         Equation<V, E> equation = equations.remove(p);
         if (equation != null) {
             Pair<ElementType, Integer> element = Pair.of(type.getElementType(), num);
             equationsByElement.get(element).remove(equation);
+            if (equationTermsByElement != null) {
+                for (EquationTerm<V, E> term : equation.getTerms()) {
+                    deindexTerm(term);
+                }
+            }
             notifyEquationChange(equation, EquationEventType.EQUATION_REMOVED);
         }
         return equation;
diff --git a/src/test/java/com/powsybl/openloadflow/equations/EquationSystemTest.java b/src/test/java/com/powsybl/openloadflow/equations/EquationSystemTest.java
index b37e2a924a..63f3bd7643 100644
--- a/src/test/java/com/powsybl/openloadflow/equations/EquationSystemTest.java
+++ b/src/test/java/com/powsybl/openloadflow/equations/EquationSystemTest.java
@@ -274,9 +274,17 @@ void removeEquationTest() {
                 .create();
         assertEquals(13, equationSystem.getEquations().size());
         assertEquals(3, equationSystem.getEquations(ElementType.BUS, 1).size());
+        assertEquals(4, equationSystem.getEquationTerms(ElementType.BRANCH, 0).size());
+        assertEquals(4, equationSystem.getEquationTerms(ElementType.BRANCH, 1).size());
+        assertEquals(4, equationSystem.getEquationTerms(ElementType.BRANCH, 2).size());
+        assertEquals(4, equationSystem.getEquationTerms(ElementType.BRANCH, 3).size());
 
         equationSystem.removeEquation(1, AcEquationType.BUS_TARGET_P);
         assertEquals(12, equationSystem.getEquations().size());
         assertEquals(2, equationSystem.getEquations(ElementType.BUS, 1).size());
+        assertEquals(3, equationSystem.getEquationTerms(ElementType.BRANCH, 0).size());
+        assertEquals(3, equationSystem.getEquationTerms(ElementType.BRANCH, 1).size());
+        assertEquals(3, equationSystem.getEquationTerms(ElementType.BRANCH, 2).size());
+        assertEquals(4, equationSystem.getEquationTerms(ElementType.BRANCH, 3).size());
     }
 }