package joshng.util.concurrent;

import com.google.common.base.Objects;
import com.google.common.collect.ForwardingMap;
import org.pcollections.HashTreePMap;
import org.pcollections.PMap;

import java.util.Map;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.atomic.AtomicReferenceFieldUpdater;
import java.util.function.Function;

/**
 * User: josh
 * Date: 7/31/14
 * Time: 8:08 AM
 */
public class AtomicMap<K, V> extends ForwardingMap<K,V> implements ConcurrentMap<K, V> {
  private static final AtomicReferenceFieldUpdater<AtomicMap, PMap> storageUpdater
          = AtomicReferenceFieldUpdater.newUpdater(AtomicMap.class, PMap.class, "storage");
  private volatile PMap<K, V> storage = HashTreePMap.empty();

  public static <K, V> AtomicMap<K, V> newAtomicMap() {
    return new AtomicMap<>();
  }

  @Override protected Map<K, V> delegate() {
    return storage;
  }

  public V put(K key, V value) {
    return (V) storageUpdater.getAndUpdate(this, s -> s.plus(key, value)).get(key);
  }

  @Override public void putAll(Map<? extends K, ? extends V> values) {
    if (!values.isEmpty()) {
      storageUpdater.getAndUpdate(this, s -> s.plusAll(values));
    }
  }

  @Override public V putIfAbsent(K key, V value) {
    PMap<K, V> currentStorage = storage;
    V currentValue;
    while ((currentValue = currentStorage.get(key)) == null && !storageUpdater.compareAndSet(this, currentStorage, currentStorage.plus(key, value))) {
      currentStorage = storage;
    }
    return currentValue;
  }

  @Override public V computeIfAbsent(K key, Function<? super K, ? extends V> mappingFunction) {
    PMap<K, V> currentStorage = storage;
    V value;
    while ((value = currentStorage.get(key)) == null
            && !storageUpdater.compareAndSet(this, currentStorage, currentStorage.plus(key, value = mappingFunction.apply(key)))) {
      currentStorage = storage;
    }
    return value;

  }

  @Override public V remove(Object key) {
    K k = (K) key;
    V removed;
    PMap<K, V> currentStorage;
    do {
      currentStorage = storage;
      removed = currentStorage.get(k);
    } while (removed != null && !storageUpdater.compareAndSet(this, currentStorage, currentStorage.minus(k)));
    return removed;
  }

  @Override public boolean remove(Object key, Object value) {
    PMap<K, V> currentStorage;
    boolean matched;
    do {
      currentStorage = storage;
    } while ((matched = Objects.equal(currentStorage.get(key), value))
            && !storageUpdater.compareAndSet(this, currentStorage, currentStorage.minus(key)));
    return matched;
  }

  @Override public boolean replace(K key, V oldValue, V newValue) {
    PMap<K, V> currentStorage;
    boolean matched;
    do {
      currentStorage = storage;
    } while ((matched = Objects.equal(currentStorage.get(key), oldValue))
            && !storageUpdater.compareAndSet(this, currentStorage, currentStorage.plus(key, newValue)));
    return matched;
  }

  @Override public V replace(K key, V value) {
    PMap<K, V> currentStorage;
    V currentValue;
    do {
      currentStorage = storage;
    } while ((currentValue = currentStorage.get(key)) != null && !storageUpdater.compareAndSet(this, currentStorage, currentStorage.plus(key, value)));
    return currentValue;
  }
}

@@ -892,16 +892,49 @@
 e;%0A  %7D%0A%0A
+  @SuppressWarnings(%22unchecked%22)%0A
   public
@@ -1039,32 +1039,65 @@
 .get(key);%0A  %7D%0A%0A
+  @SuppressWarnings(%22unchecked%22)%0A
   @Override publ
@@ -1083,32 +1083,34 @@
 ed%22)%0A  @Override
+%0A 
  public void put
@@ -1982,24 +1982,57 @@
 alue;%0A%0A  %7D%0A%0A
+  @SuppressWarnings(%22unchecked%22)%0A
   @Override 
@@ -2022,32 +2022,34 @@
 ed%22)%0A  @Override
+%0A 
  public V remove
@@ -2336,32 +2336,65 @@
 n removed;%0A  %7D%0A%0A
+  @SuppressWarnings(%22unchecked%22)%0A
   @Override publ
@@ -2380,32 +2380,34 @@
 ed%22)%0A  @Override
+%0A 
  public boolean 
@@ -2433,32 +2433,51 @@
 Object value) %7B%0A
+    K k = (K) key;%0A
     PMap%3CK, V%3E c
@@ -2611,18 +2611,16 @@
 ge.get(k
-ey
 ), value

clean up unchecked AtomicMap warnings