Implemented forEachEntry() and forEachEntryWhile()

src/main/java/net/openhft/chronicle/hash/ChronicleHash.java

@@ -18,6 +18,8 @@
 
 package net.openhft.chronicle.hash;
 
+import net.openhft.chronicle.hash.function.Consumer;
+import net.openhft.chronicle.hash.function.Predicate;
 import net.openhft.chronicle.map.ChronicleMap;
 import net.openhft.chronicle.set.ChronicleSet;
 
@@ -28,7 +30,7 @@ import java.io.File;
  * This interface defines common {@link ChronicleMap} and {@link ChronicleSet}, related to off-heap
  * memory management and file-mapping. Not usable by itself.
  */
-public interface ChronicleHash<K> extends Closeable {
+public interface ChronicleHash<K, C extends KeyContext<K>> extends Closeable {
     /**
      * Returns the file this hash container mapped to, i. e. when it is created by
      * {@link ChronicleHashBuilder#create()} call, or {@code null} if it is purely in-memory,
@@ -42,13 +44,43 @@ public interface ChronicleHash<K> extends Closeable {
 
     long longSize();
 
-    KeyContext context(K key);
+    C context(K key);
 
     /**
      * @return the class of {@code <K>}
      */
     Class<K> keyClass();
 
+    /**
+     * Checks the given predicate on each entry in this {@code ChronicleHash} until all entries
+     * have been processed or the predicate returns {@code false} for some entry, or throws
+     * an {@code Exception}. Exceptions thrown by the predicate are relayed to the caller.
+     *
+     * <p>The order in which the entries will be processed is unspecified. It might differ from
+     * the order of iteration via {@code Iterator} returned by any method of this
+     * {@code ChronicleHash} or it's collection view.
+     *
+     * <p>If the {@code ChronicleHash} is empty, this method returns {@code true} immediately.
+     *
+     * @param predicate the predicate to be checked for each entry
+     * @return {@code true} if the predicate returned {@code true} for all entries of
+     * the {@code ChronicleHash}, {@code false} if it returned {@code false} for the entry
+     */
+    boolean forEachEntryWhile(Predicate<? super C> predicate);
+
+    /**
+     * Performs the given action for each entry in this {@code ChronicleHash} until all entries have
+     * been processed or the action throws an {@code Exception}. Exceptions thrown by the action are
+     * relayed to the caller.
+     *
+     * <p>The order in which the entries will be processed is unspecified. It might differ from
+     * the order of iteration via {@code Iterator} returned by any method of this
+     * {@code ChronicleHash} or it's collection view.
+     *
+     * @param action the action to be performed for each entry
+     */
+    void forEachEntry(Consumer<? super C> action);
+
     /**
      * Releases the off-heap memory, used by this hash container and resources, used by replication,
      * if any. However, if hash container (hence off-heap memory, used by it) is mapped to the file

src/main/java/net/openhft/chronicle/hash/ChronicleHashBuilder.java

@@ -65,7 +65,7 @@ import java.util.concurrent.TimeUnit;
  * @param <B> the concrete builder type, i. e. {@link ChronicleMapBuilder}
  *            or {@link ChronicleSetBuilder}
  */
-public interface ChronicleHashBuilder<K, H extends ChronicleHash<K>,
+public interface ChronicleHashBuilder<K, H extends ChronicleHash<K, ?>,
         B extends ChronicleHashBuilder<K, H, B>> extends Cloneable {
 
     /**

src/main/java/net/openhft/chronicle/hash/KeyContext.java

@@ -31,8 +31,21 @@ import java.util.ConcurrentModificationException;
  *     <li>All methods throw {@code IllegalStateException} immediately, if called after
  *     {@link #close()} called on this context and returned.</li>
  * </ul>
+ *
+ * @param <K> key type
  */
-public interface KeyContext extends AutoCloseable, InterProcessReadWriteUpdateLock {
+public interface KeyContext<K> extends AutoCloseable, InterProcessReadWriteUpdateLock {
+
+    /**
+     * Returns the key object of this context. The instance might be reused, so you shouldn't
+     * save the returned object and use it after the context is closed or the callback (to which
+     * this context is provided) return. Might acquire {@link #readLock} before reading the key,
+     * if the context is not locked yet.
+     *
+     * @return the key object of this context
+     */
+    @NotNull
+    K key();
 
     /**
      * Returns the entry bytes, if the key is present in the hash, otherwise throws
@@ -78,8 +91,7 @@ public interface KeyContext extends AutoCloseable, InterProcessReadWriteUpdateLo
      * is already held.
      *
      * @return {@code true} if the key was actually removed (hash was updated),
-     * {@code false} if the key is already absent in the hash, or the remove operation was rejected
-     * by the hash for some reason
+     * {@code false} if the key is already absent in the hash
      * @throws IllegalStateException if the context is locked for read, but not for update, and
      * update or write lock is actually required for {@code remove()} in this {@code KeyContext}
      * implementation

src/main/java/net/openhft/chronicle/hash/function/Consumer.java

@@ -0,0 +1,38 @@
+/*
+ * Copyright 2014 Higher Frequency Trading http://www.higherfrequencytrading.com
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package net.openhft.chronicle.hash.function;
+
+/**
+ * Represents an operation that accepts a single input argument and returns no
+ * result. Unlike most other functional interfaces, {@code Consumer} is expected
+ * to operate via side-effects.
+ *
+ * <p>This is a functional interface whose functional method is {@link #accept(Object)}.
+ *
+ * <p>This is a copy of Java 8's {@code java.util.function.Consumer} interface.
+ *
+ * @param <T> the type of the input to the operation
+ */
+public interface Consumer<T> {
+
+    /**
+     * Performs this operation on the given argument.
+     *
+     * @param t the input argument
+     */
+    void accept(T t);
+}

src/main/java/net/openhft/chronicle/hash/function/Function.java

@@ -1,12 +1,13 @@
-package net.openhft.chronicle.map;
+package net.openhft.chronicle.hash.function;
 
 import java.io.Serializable;
 
 /**
  * Represents a function that accepts one argument and produces a result.
  *
- * <p>This is a <a href="package-summary.html">functional interface</a> whose functional method is {@link
- * #apply(Object)}.
+ * <p>This is a functional interface whose functional method is {@link #apply(Object)}.
+ *
+ * <p>This is a copy of Java 8's {@code java.util.function.Function} interface.
  *
  * @param <T> the type of the input to the function
  * @param <R> the type of the result of the function

src/main/java/net/openhft/chronicle/hash/function/Predicate.java

@@ -0,0 +1,39 @@
+/*
+ * Copyright 2014 Higher Frequency Trading http://www.higherfrequencytrading.com
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package net.openhft.chronicle.hash.function;
+
+
+/**
+ * Represents a predicate (boolean-valued function) of one argument.
+ *
+ * <p>This is a functional interface whose functional method is {@link #test(Object)}.
+ *
+ * <p>This is a copy of Java 8's {@code java.util.function.Predicate} interface.
+ *
+ * @param <T> the type of the input to the predicate
+ */
+public interface Predicate<T> {
+
+    /**
+     * Evaluates this predicate on the given argument.
+     *
+     * @param t the input argument
+     * @return {@code true} if the input argument matches the predicate,
+     * otherwise {@code false}
+     */
+    boolean test(T t);
+}

src/main/java/net/openhft/chronicle/hash/function/package-info.java

@@ -0,0 +1,21 @@
+/*
+ * Copyright 2014 Higher Frequency Trading http://www.higherfrequencytrading.com
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+/**
+ * Copies of some interfaces from Java 8 used by Chronicle Map project
+ * for Java 7+.
+ */
+package net.openhft.chronicle.hash.function;

src/main/java/net/openhft/chronicle/hash/replication/LateUpdateException.java

@@ -16,11 +16,13 @@
 
 package net.openhft.chronicle.hash.replication;
 
+import net.openhft.chronicle.hash.ChronicleHash;
+
 /**
- * Thrown if update to replicated {@link net.openhft.chronicle.hash.ChronicleHash} is rejected
- * due to local time is lesser than the time of the latest update to the entry. This most likely
- * indicates bugs or largely inconsistent timing across replicating nodes. If this is tolerated,
- * you should catch {@code LateUpdateException} and ignore it.
+ * Thrown if update to replicated {@link ChronicleHash} is rejected due to local time is lesser
+ * than the time of the latest update to the entry. This most likely indicates bugs or largely
+ * inconsistent timing across replicating nodes. If this is tolerated, you should catch
+ * {@code LateUpdateException} and ignore it.
  */
 public class LateUpdateException extends RuntimeException {
     private static final long serialVersionUID = 0L;

src/main/java/net/openhft/chronicle/map/AbstractChronicleMap.java

@@ -17,6 +17,10 @@
 package net.openhft.chronicle.map;
 
 import net.openhft.chronicle.hash.KeyContext;
+import net.openhft.chronicle.hash.function.Consumer;
+import net.openhft.chronicle.hash.function.Function;
+import net.openhft.chronicle.hash.function.Predicate;
+import net.openhft.chronicle.map.ReplicatedChronicleMap.ReplicatedContext;
 import org.jetbrains.annotations.NotNull;
 
 import java.io.File;
@@ -26,6 +30,7 @@ import java.util.*;
 
 import static java.util.Collections.emptyList;
 import static net.openhft.chronicle.map.Objects.requireNonNull;
+import static net.openhft.chronicle.map.VanillaContext.SearchState.DELETED;
 
 abstract class AbstractChronicleMap<K, V> extends AbstractMap<K, V>
         implements ChronicleMap<K, V>, Serializable {
@@ -39,14 +44,14 @@ abstract class AbstractChronicleMap<K, V> extends AbstractMap<K, V>
     @Override
     @SuppressWarnings("unchecked")
     public final boolean containsKey(Object key) {
-        try (MapKeyContext<V> c = context((K) key)) {
+        try (MapKeyContext<K, V> c = context((K) key)) {
             return c.containsKey();
         }
     }
 
     @Override
     public final V put(K key, V value) {
-        try (MapKeyContext<V> c = context(key)) {
+        try (MapKeyContext<K, V> c = context(key)) {
             // We cannot read the previous value under read lock, because then we will need
             // to release the read lock -> acquire write lock, the value might be updated in
             // between, that will break ConcurrentMap.put() atomicity guarantee. So, we acquire
@@ -58,14 +63,14 @@ abstract class AbstractChronicleMap<K, V> extends AbstractMap<K, V>
         }
     }
 
-    V prevValueOnPut(MapKeyContext<V> context) {
+    V prevValueOnPut(MapKeyContext<K, V> context) {
         return context.getUsing(null);
     }
 
     @Override
     public final V putIfAbsent(K key, V value) {
         checkValue(value);
-        try (MapKeyContext<V> c = context(key)) {
+        try (MapKeyContext<K, V> c = context(key)) {
             // putIfAbsent() shouldn't actually put most of the time,
             // so check if the key is present under read lock first:
             if (c.readLock().tryLock()) {
@@ -99,7 +104,7 @@ abstract class AbstractChronicleMap<K, V> extends AbstractMap<K, V>
 
     @Override
     public final V getUsing(K key, V usingValue) {
-        try (MapKeyContext<V> c = context(key)) {
+        try (MapKeyContext<K, V> c = context(key)) {
             return c.getUsing(usingValue);
         }
     }
@@ -142,7 +147,7 @@ abstract class AbstractChronicleMap<K, V> extends AbstractMap<K, V>
     @Override
     public <R> R getMapped(K key, @NotNull Function<? super V, R> function) {
         requireNonNull(function);
-        try (MapKeyContext<V> c = context(key)) {
+        try (MapKeyContext<K, V> c = context(key)) {
             return c.containsKey() ? function.apply(c.get()) : null;
         }
     }
@@ -150,7 +155,7 @@ abstract class AbstractChronicleMap<K, V> extends AbstractMap<K, V>
     @Override
     public V putMapped(K key, @NotNull UnaryOperator<V> unaryOperator) {
         requireNonNull(unaryOperator);
-        try (MapKeyContext<V> c = context(key)) {
+        try (MapKeyContext<K, V> c = context(key)) {
             // putMapped() should find a value, update & put most of the time,
             // so don't try to check key presence under read lock first,
             // as in putIfAbsent()/acquireUsing(), start with update lock:
@@ -178,7 +183,7 @@ abstract class AbstractChronicleMap<K, V> extends AbstractMap<K, V>
     @Override
     @SuppressWarnings("unchecked")
     public final V remove(Object key) {
-        try (MapKeyContext<V> c = context((K) key)) {
+        try (MapKeyContext<K, V> c = context((K) key)) {
             // We cannot read the previous value under read lock, because then we will need
             // to release the read lock -> acquire write lock, the value might be updated in
             // between, that will break ConcurrentMap.remove() atomicity guarantee. So, we acquire
@@ -190,7 +195,7 @@ abstract class AbstractChronicleMap<K, V> extends AbstractMap<K, V>
         }
     }
 
-    V prevValueOnRemove(MapKeyContext<V> context) {
+    V prevValueOnRemove(MapKeyContext<K, V> context) {
         return context.getUsing(null);
     }
 
@@ -200,7 +205,7 @@ abstract class AbstractChronicleMap<K, V> extends AbstractMap<K, V>
         if (value == null)
             return false; // CHM compatibility; General ChronicleMap policy is to throw NPE
         checkValue((V) value);
-        try (MapKeyContext<V> c = context((K) key)) {
+        try (MapKeyContext<K, V> c = context((K) key)) {
             // remove(key, value) should find the entry & remove most of the time,
             // so don't try to check key presence and value equivalence under read lock first,
             // as in putIfAbsent()/acquireUsing(), start with update lock:
@@ -212,7 +217,7 @@ abstract class AbstractChronicleMap<K, V> extends AbstractMap<K, V>
     @Override
     public final V replace(K key, V value) {
         checkValue(value);
-        try (MapKeyContext<V> c = context(key)) {
+        try (MapKeyContext<K, V> c = context(key)) {
             // replace(key, value) should find the key & put the value most of the time,
             // so don't try to check key presence under read lock first,
             // as in putIfAbsent()/acquireUsing(), start with update lock:
@@ -232,7 +237,7 @@ abstract class AbstractChronicleMap<K, V> extends AbstractMap<K, V>
     public final boolean replace(K key, V oldValue, V newValue) {
         checkValue(oldValue);
         checkValue(newValue);
-        try (MapKeyContext<V> c = context(key)) {
+        try (MapKeyContext<K, V> c = context(key)) {
             // replace(key, old, new) should find the entry & put new value most of the time,
             // so don't try to check key presence and value equivalence under read lock first,
             // as in putIfAbsent()/acquireUsing(), start with update lock:
@@ -261,7 +266,6 @@ abstract class AbstractChronicleMap<K, V> extends AbstractMap<K, V>
     class EntryIterator implements Iterator<Entry<K, V>>, HashLookup.EntryConsumer {
         /**
          * Very inefficient (esp. if segments are large), but CORRECT implementation
-         * TODO map.forEachEntry(Consumer<MapKeyContext<V>>)
          */
         private final Thread ownerThread = Thread.currentThread();
         private int segmentIndex = actualSegments() - 1;
@@ -347,7 +351,7 @@ abstract class AbstractChronicleMap<K, V> extends AbstractMap<K, V>
 
         @Override
         public V setValue(V value) {
-            try (MapKeyContext<V> c = AbstractChronicleMap.this.context(getKey())) {
+            try (MapKeyContext<K, V> c = AbstractChronicleMap.this.context(getKey())) {
                 c.put(value);
             }
             return super.setValue(value);
@@ -401,4 +405,81 @@ abstract class AbstractChronicleMap<K, V> extends AbstractMap<K, V>
             AbstractChronicleMap.this.clear();
         }
     }
+
+    @Override
+    public void forEachEntry(final Consumer<? super MapKeyContext<K, V>> action) {
+        forEachEntryWhile(new Predicate<MapKeyContext<K, V>>() {
+            @Override
+            public boolean test(MapKeyContext<K, V> c) {
+                action.accept(c);
+                return true;
+            }
+        });
+    }
+
+    @Override
+    public boolean forEachEntryWhile(final Predicate<? super MapKeyContext<K, V>> predicate) {
+        boolean interrupt = false;
+        iteration:
+        try (VanillaContext<K, ?, ?, V, ?, ?> c = mapContext()) {
+            ForEachWhilePredicate<K, V> hashLookupPredicate =
+                    new ForEachWhilePredicate<>(c, predicate);
+            for (int segmentIndex = actualSegments() - 1; segmentIndex >= 0; segmentIndex--) {
+                c.segmentIndex = segmentIndex;
+                c.forEachEntry = true;
+                try {
+                    c.updateLock().lock();
+                    if (c.size() == 0)
+                        continue;
+                    c.initSegment();
+                    c.hashLookup.forEachRemoving(hashLookupPredicate);
+                    if (hashLookupPredicate.shouldBreak) {
+                        interrupt = true;
+                        break iteration;
+                    }
+                } finally {
+                    c.forEachEntry = false;
+                    c.closeSegmentIndex();
+                }
+            }
+        }
+        return !interrupt;
+    }
+
+
+    private static class ForEachWhilePredicate<K, V> implements HashLookup.EntryPredicate {
+        private final VanillaContext<K, ?, ?, V, ?, ?> c;
+        private final Predicate<? super MapKeyContext<K, V>> predicate;
+        boolean shouldBreak = false;
+
+        public ForEachWhilePredicate(VanillaContext<K, ?, ?, V, ?, ?> c,
+                                     Predicate<? super MapKeyContext<K, V>> predicate) {
+            this.c = c;
+            this.predicate = predicate;
+            shouldBreak = false;
+        }
+
+        @Override
+        public boolean remove(long hash, long pos) {
+            c.pos = pos;
+            c.initKeyFromPos();
+            try {
+                if (!c.containsKey()) // for replicated map
+                    return false;
+                shouldBreak = predicate.test(c);
+                // release all exclusive locks: possibly if context.remove() is performed
+                // in the callback, or acquired manually
+                while (c.writeLockCount > 0) {
+                    c.writeLock().unlock();
+                }
+                if (!c.isUpdateLocked()) {
+                    throw new IllegalStateException("Shouldn't release update lock " +
+                            "inside forEachEntry() callback");
+                }
+                return c.state == DELETED && !(c instanceof ReplicatedContext);
+            } finally {
+                c.closeKeySearch();
+            }
+        }
+    }
 }

src/main/java/net/openhft/chronicle/map/BytesChronicleMap.java

@@ -87,7 +87,8 @@ public class BytesChronicleMap extends AbstractChronicleMap<Bytes, Bytes> {
 
     @NotNull
     @Override
-    public MapKeyContext<Bytes> acquireContext(@NotNull Bytes key, @NotNull Bytes usingValue) {
+    public MapKeyContext<Bytes, Bytes> acquireContext(
+            @NotNull Bytes key, @NotNull Bytes usingValue) {
         throw new UnsupportedOperationException();
     }
 

src/main/java/net/openhft/chronicle/map/ChronicleMap.java

@@ -19,6 +19,9 @@
 package net.openhft.chronicle.map;
 
 import net.openhft.chronicle.hash.ChronicleHash;
+import net.openhft.chronicle.hash.KeyContext;
+import net.openhft.chronicle.hash.function.Consumer;
+import net.openhft.chronicle.hash.function.Function;
 import net.openhft.chronicle.hash.serialization.BytesReader;
 import net.openhft.lang.io.Bytes;
 import net.openhft.lang.io.serialization.BytesMarshaller;
@@ -53,10 +56,8 @@ import java.util.concurrent.ConcurrentMap;
  * @param <K> the map key type
  * @param <V> the map value type
  */
-public interface ChronicleMap<K, V> extends ConcurrentMap<K, V>, ChronicleHash<K> {
-
-    @Override
-    MapKeyContext<V> context(K key);
+public interface ChronicleMap<K, V>
+        extends ConcurrentMap<K, V>, ChronicleHash<K, MapKeyContext<K, V>> {
 
     /**
      * Returns the value to which the specified key is mapped, or {@code null} if this map contains
@@ -140,7 +141,7 @@ public interface ChronicleMap<K, V> extends ConcurrentMap<K, V>, ChronicleHash<K
     V acquireUsing(@NotNull K key, V usingValue);
 
     @NotNull
-    MapKeyContext<V> acquireContext(@NotNull K key, @NotNull V usingValue);
+    MapKeyContext<K, V> acquireContext(@NotNull K key, @NotNull V usingValue);
 
     /**
      * Apply a mapping to the value returned by a key and return a result. A read lock is assumed.

src/main/java/net/openhft/chronicle/map/HashLookup.java

@@ -218,8 +218,11 @@ class HashLookup {
     }
 
     private long remove0() {
-        long posToRemove = searchPos;
-        long entryToRemove = readEntry(searchPos);
+        return remove00(searchPos);
+    }
+
+    private long remove00(long posToRemove) {
+        long entryToRemove = readEntry(posToRemove);
         long posToShift = posToRemove;
         while (true) {
             posToShift = step(posToShift);
@@ -285,4 +288,26 @@ class HashLookup {
                 action.accept(key(entry), value(entry));
         }
     }
+
+    static interface EntryPredicate {
+        boolean remove(long hash, long pos);
+    }
+
+    public void forEachRemoving(EntryPredicate predicate) {
+        long pos = 0L;
+        while (!empty(readEntry(pos))) {
+            pos = step(pos);
+        }
+        long startPos = pos;
+        do {
+            pos = step(pos);
+            long entry = readEntry(pos);
+            if (!empty(entry)) {
+                if (predicate.remove(key(entry), value(entry))) {
+                    remove00(pos);
+                    pos = stepBack(pos);
+                }
+            }
+        } while (pos != startPos);
+    }
 }

src/main/java/net/openhft/chronicle/map/MapKeyContext.java

@@ -18,7 +18,7 @@ package net.openhft.chronicle.map;
 
 import net.openhft.chronicle.hash.KeyContext;
 
-public interface MapKeyContext<V> extends KeyContext {
+public interface MapKeyContext<K, V> extends KeyContext<K> {
 
     long valueOffset();
 

src/main/java/net/openhft/chronicle/map/ReplicationHubFindByName.java

@@ -70,7 +70,7 @@ class ReplicationHubFindByName<K> implements FindByName {
 
         int withChannelId = nextFreeChannel.incrementAndGet();
 
-        try (MapKeyContext<MapInstanceBuilder> c = map.context(config.name)) {
+        try (MapKeyContext<CharSequence, MapInstanceBuilder> c = map.context(config.name)) {
             c.updateLock().lock();
             MapInstanceBuilder value =
                     config.replicatedViaChannel(replicationHub.createChannel(withChannelId));

src/main/java/net/openhft/chronicle/map/StatelessChronicleMap.java

@@ -20,6 +20,9 @@ package net.openhft.chronicle.map;
 
 import com.sun.jdi.connect.spi.ClosedConnectionException;
 import net.openhft.chronicle.hash.RemoteCallTimeoutException;
+import net.openhft.chronicle.hash.function.Consumer;
+import net.openhft.chronicle.hash.function.Function;
+import net.openhft.chronicle.hash.function.Predicate;
 import net.openhft.chronicle.hash.serialization.BytesReader;
 import net.openhft.lang.io.ByteBufferBytes;
 import net.openhft.lang.io.Bytes;
@@ -223,6 +226,18 @@ class StatelessChronicleMap<K, V> implements ChronicleMap<K, V>, Closeable, Clon
         return kClass;
     }
 
+    @Override
+    public boolean forEachEntryWhile(Predicate<? super MapKeyContext<K, V>> predicate) {
+        // TODO implement!
+        throw new UnsupportedOperationException();
+    }
+
+    @Override
+    public void forEachEntry(Consumer<? super MapKeyContext<K, V>> action) {
+        // TODO implement!
+        throw new UnsupportedOperationException();
+    }
+
     @Override
     public Class<V> valueClass() {
         return vClass;
@@ -535,7 +550,7 @@ class StatelessChronicleMap<K, V> implements ChronicleMap<K, V>, Closeable, Clon
     }
 
     @Override
-    public MapKeyContext<V> context(K key) {
+    public MapKeyContext<K, V> context(K key) {
         throw new UnsupportedOperationException("Contexts are not supported by stateless clients");
     }
 
@@ -578,7 +593,7 @@ class StatelessChronicleMap<K, V> implements ChronicleMap<K, V>, Closeable, Clon
 
     @NotNull
     @Override
-    public MapKeyContext<V> acquireContext(@NotNull K key, @NotNull V usingValue) {
+    public MapKeyContext<K, V> acquireContext(@NotNull K key, @NotNull V usingValue) {
         throw new UnsupportedOperationException("Contexts are not supported by stateless clients");
     }
 
@@ -1894,7 +1909,6 @@ class StatelessChronicleMap<K, V> implements ChronicleMap<K, V>, Closeable, Clon
             return null;
 
         return (R) readObject(transactionId, startTime);
-
     }
 }
 

src/main/java/net/openhft/chronicle/map/TcpReplicator.java

@@ -18,6 +18,7 @@
 
 package net.openhft.chronicle.map;
 
+import net.openhft.chronicle.hash.function.Function;
 import net.openhft.chronicle.hash.replication.RemoteNodeValidator;
 import net.openhft.chronicle.hash.replication.TcpTransportAndNetworkConfig;
 import net.openhft.chronicle.hash.replication.ThrottlingConfig;

src/main/java/net/openhft/chronicle/map/VanillaChronicleMap.java

@@ -493,7 +493,7 @@ class VanillaChronicleMap<K, KI, MKI extends MetaBytesInterop<K, ? super KI>,
 
     @NotNull
     @Override
-    public final MapKeyContext<V> acquireContext(K key, V usingValue) {
+    public final MapKeyContext<K, V> acquireContext(K key, V usingValue) {
         VanillaContext<K, KI, MKI, V, VI, MVI> c = acquireContext(key);
         // TODO optimize to update lock in certain cases
         c.writeLock().lock();

src/main/java/net/openhft/chronicle/map/VanillaContext.java

@@ -44,7 +44,7 @@ import static net.openhft.lang.io.NativeBytes.UNSAFE;
 
 class VanillaContext<K, KI, MKI extends MetaBytesInterop<K, ? super KI>,
         V, VI, MVI extends MetaBytesInterop<V, ? super VI>>
-        implements MapKeyContext<V> {
+        implements MapKeyContext<K, V> {
 
     static interface ContextFactory<T extends VanillaContext> {
         T createContext(VanillaContext root, int indexInContextCache);
@@ -433,6 +433,18 @@ class VanillaContext<K, KI, MKI extends MetaBytesInterop<K, ? super KI>,
         return keySize;
     }
 
+    @NotNull
+    @Override
+    public K key() {
+        if (key != null)
+            return key;
+        initKeySearch();
+        initMapAndContextLocals();
+        mapAndContextLocalsReference.mapAndContextLocals.reusableKey =
+                key0(mapAndContextLocalsReference.mapAndContextLocals.reusableKey);
+        assert key != null;
+        return key;
+    }
 
     /////////////////////////////////////////////////
     // Key hash
@@ -1654,7 +1666,7 @@ class VanillaContext<K, KI, MKI extends MetaBytesInterop<K, ? super KI>,
         if (containsKey()) {
             initEntrySizeInChunks();
             upgradeToWriteLock();
-            hashLookup.remove();
+            hashLookupRemove();
             free(pos, entrySizeInChunks);
             size(size() - 1L);
             state = DELETED;
@@ -1664,6 +1676,13 @@ class VanillaContext<K, KI, MKI extends MetaBytesInterop<K, ? super KI>,
         }
     }
 
+    boolean forEachEntry;
+
+    void hashLookupRemove() {
+        if (!forEachEntry)
+            hashLookup.remove();
+    }
+
     void initRemoveDependencies() {
         initSegment();
         initLocks();
@@ -1901,9 +1920,13 @@ class VanillaContext<K, KI, MKI extends MetaBytesInterop<K, ? super KI>,
     }
 
     K immutableKey() {
+        return key0(null);
+    }
+
+    K key0(K usingKey) {
         initKeyReader();
         entry.position(keyOffset);
-        key = keyReader.read(entry, keySize, null);
+        key = keyReader.read(entry, keySize, usingKey);
         return key;
     }
 

src/main/java/net/openhft/chronicle/set/ChronicleSet.java

@@ -19,8 +19,9 @@
 package net.openhft.chronicle.set;
 
 import net.openhft.chronicle.hash.ChronicleHash;
+import net.openhft.chronicle.hash.KeyContext;
 
 import java.util.Set;
 
-public interface ChronicleSet<K> extends Set<K>, ChronicleHash<K> {
+public interface ChronicleSet<K> extends Set<K>, ChronicleHash<K, KeyContext<K>> {
 }

src/main/java/net/openhft/chronicle/set/SetFromMap.java

@@ -19,6 +19,8 @@
 package net.openhft.chronicle.set;
 
 import net.openhft.chronicle.hash.KeyContext;
+import net.openhft.chronicle.hash.function.Consumer;
+import net.openhft.chronicle.hash.function.Predicate;
 import net.openhft.chronicle.map.ChronicleMap;
 
 import java.io.File;
@@ -130,6 +132,16 @@ class SetFromMap<E> extends AbstractSet<E>
         return m.keyClass();
     }
 
+    @Override
+    public boolean forEachEntryWhile(Predicate<? super KeyContext<E>> predicate) {
+        return m.forEachEntryWhile(predicate);
+    }
+
+    @Override
+    public void forEachEntry(Consumer<? super KeyContext<E>> action) {
+        m.forEachEntry(action);
+    }
+
     @Override
     public File file() {
         return m.file();

src/main/java/net/openhft/xstream/converters/VanillaChronicleMapConverter.java

@@ -19,6 +19,10 @@ package net.openhft.xstream.converters;
 
 import com.thoughtworks.xstream.converters.MarshallingContext;
 import com.thoughtworks.xstream.io.HierarchicalStreamWriter;
+import net.openhft.chronicle.hash.function.Consumer;
+import net.openhft.chronicle.hash.function.Predicate;
+import net.openhft.chronicle.map.ChronicleMap;
+import net.openhft.chronicle.map.MapKeyContext;
 import net.openhft.lang.model.constraints.NotNull;
 
 import java.util.Map;
@@ -33,27 +37,26 @@ public class VanillaChronicleMapConverter<K, V> extends AbstractChronicleMapConv
     }
 
     @Override
-    public void marshal(Object o, HierarchicalStreamWriter writer, MarshallingContext
+    public void marshal(Object o, final HierarchicalStreamWriter writer, final MarshallingContext
             marshallingContext) {
+        ((ChronicleMap<K, V>) o).forEachEntry(new Consumer<MapKeyContext<K,V>>() {
+            @Override
+            public void accept(MapKeyContext<K, V> c) {
+                writer.startNode("entry");
+                {
+                    final Object key = c.key();
+                    writer.startNode(key.getClass().getName());
+                    marshallingContext.convertAnother(key);
+                    writer.endNode();
 
-        for (Map.Entry e : (Iterable<Map.Entry>) ((Map) o).entrySet()) {
-
-            writer.startNode("entry");
-            {
-                final Object key = e.getKey();
-                writer.startNode(key.getClass().getName());
-                marshallingContext.convertAnother(key);
-                writer.endNode();
-
-                Object value = e.getValue();
-                writer.startNode(value.getClass().getName());
-                marshallingContext.convertAnother(value);
+                    Object value = c.get();
+                    writer.startNode(value.getClass().getName());
+                    marshallingContext.convertAnother(value);
+                    writer.endNode();
+                }
                 writer.endNode();
             }
-            writer.endNode();
-        }
-
+        });
     }
-
 }