From ec2a52900e2c7aff14e9845ebcbb4772937ce838 Mon Sep 17 00:00:00 2001
From: Sergey Kapralov <skapral@mera.ru>
Date: Thu, 25 Aug 2016 11:13:08 +0300
Subject: [PATCH] TMap.compute fixed and tested

---
 .../org/teavm/classlib/java/util/TMap.java    |  9 ++-
 .../classlib/java/util/HashtableTest.java     | 57 +++++++++++++++++++
 2 files changed, 61 insertions(+), 5 deletions(-)

diff --git a/classlib/src/main/java/org/teavm/classlib/java/util/TMap.java b/classlib/src/main/java/org/teavm/classlib/java/util/TMap.java
index 2d078a6fd..30387d785 100644
--- a/classlib/src/main/java/org/teavm/classlib/java/util/TMap.java
+++ b/classlib/src/main/java/org/teavm/classlib/java/util/TMap.java
@@ -105,15 +105,14 @@ public interface TMap<K, V> {
         V newValue = remappingFunction.apply(key, oldValue);
         if (oldValue != null) {
             if (newValue != null) {
-                return put(key, newValue);
+                put(key, newValue);
             } else {
-                return remove(key);
+                remove(key);
             }
         } else if (newValue != null) {
-            return put(key, newValue);
-        } else {
-            return null;
+            put(key, newValue);
         }
+        return newValue;
     }
 
     default V merge(K key, V value, BiFunction<? super V, ? super V, ? extends V> remappingFunction) {
diff --git a/tests/src/test/java/org/teavm/classlib/java/util/HashtableTest.java b/tests/src/test/java/org/teavm/classlib/java/util/HashtableTest.java
index 8478b05bd..deb94921d 100644
--- a/tests/src/test/java/org/teavm/classlib/java/util/HashtableTest.java
+++ b/tests/src/test/java/org/teavm/classlib/java/util/HashtableTest.java
@@ -693,6 +693,63 @@ public class HashtableTest {
         }
     }
 
+    @Test
+    public void test_computeUpdatesValueIfPresent() {
+        Hashtable<String, String> ht10 = new Hashtable<>();
+        for(int i = 0; i < 10; i++) {
+            ht10.put("Key" + i, "Val" + i);
+        }
+        
+        String newVal = ht10.compute("Key5", (k,v) -> "changed");
+        assertEquals("changed", newVal);
+        assertEquals(10, ht10.size());
+        
+        for(int i = 0; i < 10; i++) {
+            if(i == 5) {
+                assertEquals("Value was incorrectly changed", "changed", ht10.get("Key" + i));
+            } else {
+                assertEquals("Value was unexpectedly changed", "Val" + i, ht10.get("Key" + i));
+            }
+        }
+    }
+
+    @Test
+    public void test_computePutsNewEntryIfKeyIsAbsent() {
+        Hashtable<String, String> ht10 = new Hashtable<>();
+        for(int i = 0; i < 10; i++) {
+            ht10.put("Key" + i, "Val" + i);
+        }
+        
+        String newVal = ht10.compute("absent key", (k,v) -> "added");
+        assertEquals("added", newVal);
+        assertEquals(11, ht10.size());
+        
+        for(int i = 0; i < 10; i++) {
+            assertEquals("Value was unexpectedly changed", "Val" + i, ht10.get("Key" + i));
+        }
+        assertEquals("New value expected","added", ht10.get("absent key"));
+    }
+
+    @Test
+    public void test_computeRemovesEntryOnNullMapping() {
+        Hashtable<String, String> ht10 = new Hashtable<>();
+        for(int i = 0; i < 10; i++) {
+            ht10.put("Key" + i, "Val" + i);
+        }
+        
+        String newVal = ht10.compute("Key5", (k,v) -> null);
+        assertEquals(null, newVal);
+        assertEquals(9, ht10.size());
+        
+        for(int i = 0; i < 10; i++) {
+            if(i == 5) {
+                assertEquals("Value was unexpectedly present in map", null, ht10.get("Key" + i));
+            } else {
+                assertEquals("Value was unexpectedly changed", "Val" + i, ht10.get("Key" + i));
+            }
+        }
+    }
+    
     @SuppressWarnings("unchecked")
     protected <K, V> Hashtable<K, V> hashtableClone(Hashtable<K, V> s) {
         return (Hashtable<K, V>) s.clone();