From 05454380d92f0a3cbceae805cd467d2f567cd25a Mon Sep 17 00:00:00 2001
From: Alexey Andreev <konsoletyper@gmail.com>
Date: Mon, 9 Oct 2023 20:35:54 +0200
Subject: [PATCH] classlib: fix issue in TreeMap iterator remove method

---
 .../teavm/classlib/java/util/TTreeMap.java    |  8 +++++-
 .../teavm/classlib/java/util/TreeMapTest.java | 27 ++++++++++++++++---
 2 files changed, 31 insertions(+), 4 deletions(-)

diff --git a/classlib/src/main/java/org/teavm/classlib/java/util/TTreeMap.java b/classlib/src/main/java/org/teavm/classlib/java/util/TTreeMap.java
index 429e50d94..30775eff7 100644
--- a/classlib/src/main/java/org/teavm/classlib/java/util/TTreeMap.java
+++ b/classlib/src/main/java/org/teavm/classlib/java/util/TTreeMap.java
@@ -732,7 +732,13 @@ public class TTreeMap<K, V> extends TAbstractMap<K, V> implements TCloneable, TS
             if (last == null) {
                 throw new TNoSuchElementException();
             }
-            owner.root = owner.deleteNode(owner.root, last.getKey());
+            var newRoot = owner.deleteNode(owner.root, last.getKey());
+            if (owner.root != newRoot) {
+                owner.root = newRoot;
+                var newPath = owner.pathToNext(last.getKey(), reverse);
+                System.arraycopy(newPath, 0, path, 0, newPath.length);
+                depth = newPath.length;
+            }
             modCount = ++owner.modCount;
             last = null;
         }
diff --git a/tests/src/test/java/org/teavm/classlib/java/util/TreeMapTest.java b/tests/src/test/java/org/teavm/classlib/java/util/TreeMapTest.java
index 42a50aac1..6b13f2ae8 100644
--- a/tests/src/test/java/org/teavm/classlib/java/util/TreeMapTest.java
+++ b/tests/src/test/java/org/teavm/classlib/java/util/TreeMapTest.java
@@ -68,7 +68,7 @@ import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.teavm.junit.TeaVMTestRunner;
 
-@SuppressWarnings({ "UnnecessaryTemporaryOnConversionToString", "SuspiciousMethodCalls" })
+@SuppressWarnings("SuspiciousMethodCalls")
 @RunWith(TeaVMTestRunner.class)
 public class TreeMapTest {
 
@@ -148,9 +148,9 @@ public class TreeMapTest {
     public void test_ConstructorLjava_util_Map() {
         // Test for method java.util.TreeMap(java.util.Map)
         TreeMap<Object, Object> myTreeMap = new TreeMap<>(new HashMap<>(tm));
-        assertTrue("Map is incorrect size", myTreeMap.size() == objArray.length);
+        assertEquals("Map is incorrect size", objArray.length, myTreeMap.size());
         for (Object element : objArray) {
-            assertTrue("Map has incorrect mappings", myTreeMap.get(element.toString()).equals(element));
+            assertEquals("Map has incorrect mappings", myTreeMap.get(element.toString()), element);
         }
     }
 
@@ -707,4 +707,25 @@ public class TreeMapTest {
         assertEquals("{10=119}", map.subMap(10, 29).toString());
         assertEquals("{}", map.subMap(29, 100).toString());
     }
+
+    @Test
+    public void iteratorRemove() {
+        var keys = new String[] { "a", "b", "c", "d", "e", "f", "g" };
+        for (var i = 1; i < keys.length; ++i) {
+            for (var j = 0; j < i; ++j) {
+                var map = new TreeMap<String, Integer>();
+                for (var k = 0; k < i; ++k) {
+                    map.put(keys[k], k);
+                }
+                var iter = map.keySet().iterator();
+                for (var k = 0; k < i; ++k) {
+                    assertEquals(keys[k], iter.next());
+                    if (k == j) {
+                        iter.remove();
+                    }
+                }
+                assertFalse(iter.hasNext());
+            }
+        }
+    }
 }