package com.tinkerpop.blueprints.util.wrappers.batch.cache;

import com.tinkerpop.blueprints.Graph;
import com.tinkerpop.blueprints.Vertex;

import java.util.HashMap;
import java.util.Map;

/**
 * (c) Matthias Broecheler (me@matthiasb.com)
 */

public class StringIDVertexCache implements VertexCache {

    private static final int INITIAL_CAPACITY = 1000;

    private final Map<String, Object> map;
    private final StringCompression compression;

    public StringIDVertexCache(final StringCompression compression) {
        if (compression == null) throw new IllegalArgumentException("Compression expected.");
        this.compression = compression;
        map = new HashMap<String, Object>(INITIAL_CAPACITY);
    }

    public StringIDVertexCache() {
        this(StringCompression.NO_COMPRESSION);
    }

    @Override
    public Object getEntry(Object externalId) {
        String id = compression.compress(externalId.toString());
        return map.get(id);
    }

    @Override
    public void set(Vertex vertex, Object externalId) {
        setId(vertex,externalId);
    }

    @Override
    public void setId(Object vertexId, Object externalId) {
        String id = compression.compress(externalId.toString());
        map.put(id, vertexId);
    }

    @Override
    public boolean contains(Object externalId) {
        return map.containsKey(compression.compress(externalId.toString()));
    }

    @Override
    public void newTransaction() {
        for (Map.Entry<String, Object> entry : map.entrySet()) {
            if (entry.getValue() instanceof Vertex) {
                Vertex v = (Vertex) entry.getValue();
                entry.setValue(v.getId());
            }
        }
    }
}
@@ -177,19 +177,67 @@
 va.util.
-Map
+HashSet;%0Aimport java.util.Map;%0Aimport java.util.Set
 ;%0A%0A/**%0A 
@@ -403,50 +403,156 @@
 00;%0A
-%0A    private final Map%3CString, Object%3E map
+    private static final int INITIAL_TX_CAPACITY = 100;%0A%0A    private final Map%3CString, Object%3E map;%0A    private final Set%3CString%3E mapKeysInCurrentTx
 ;%0A  
@@ -864,16 +864,87 @@
 ACITY);%0A
+        mapKeysInCurrentTx = new HashSet%3CString%3E(INITIAL_TX_CAPACITY);%0A
     %7D%0A%0A 
@@ -1471,24 +1471,60 @@
  vertexId);%0A
+        mapKeysInCurrentTx.add(id);%0A
     %7D%0A%0A    @
@@ -1731,92 +1731,125 @@
 or (
-Map.Entry%3CString, Object%3E entry : map.entrySet()) %7B%0A            if (entry.getValue()
+String id : mapKeysInCurrentTx) %7B%0A            Object o = map.get(id);%0A            assert null != o;%0A            if (o
  ins
@@ -1905,25 +1905,9 @@
 tex)
- entry.getValue()
+o
 ;%0A  
@@ -1924,23 +1924,20 @@
     
-entry.setValue(
+map.put(id, 
 v.ge
@@ -1972,11 +1972,48 @@
   %7D%0A
+        mapKeysInCurrentTx.clear();%0A
     %7D%0A%7D
+%0A

Optimize StringIDVertexCache