diff --git a/java/src/com/android/inputmethod/keyboard/KeyboardLayoutSet.java b/java/src/com/android/inputmethod/keyboard/KeyboardLayoutSet.java
index 5d9b8a712813dcd2fd0f4877268fe031fe27e265..d71dc5956cc7d62534cf615b1e02abd4e4b806dc 100644
--- a/java/src/com/android/inputmethod/keyboard/KeyboardLayoutSet.java
+++ b/java/src/com/android/inputmethod/keyboard/KeyboardLayoutSet.java
@@ -34,7 +34,7 @@ import com.android.inputmethod.compat.EditorInfoCompatUtils;
 import com.android.inputmethod.compat.InputMethodSubtypeCompatUtils;
 import com.android.inputmethod.keyboard.internal.KeyboardBuilder;
 import com.android.inputmethod.keyboard.internal.KeyboardParams;
-import com.android.inputmethod.keyboard.internal.KeysCache;
+import com.android.inputmethod.keyboard.internal.UniqueKeysCache;
 import com.android.inputmethod.latin.InputAttributes;
 import com.android.inputmethod.latin.R;
 import com.android.inputmethod.latin.RichInputMethodSubtype;
@@ -86,7 +86,7 @@ public final class KeyboardLayoutSet {
     private static final HashMap<KeyboardId, SoftReference<Keyboard>> sKeyboardCache =
             new HashMap<>();
     @Nonnull
-    private static final KeysCache sKeysCache = new KeysCache();
+    private static final UniqueKeysCache sUniqueKeysCache = UniqueKeysCache.newInstance();
     private final static HashMap<InputMethodSubtype, Integer> sScriptIdsForSubtypes =
             new HashMap<>();
 
@@ -144,7 +144,7 @@ public final class KeyboardLayoutSet {
 
     private static void clearKeyboardCache() {
         sKeyboardCache.clear();
-        sKeysCache.clear();
+        sUniqueKeysCache.clear();
     }
 
     public static int getScriptId(final Resources resources,
@@ -219,10 +219,8 @@ public final class KeyboardLayoutSet {
         }
 
         final KeyboardBuilder<KeyboardParams> builder =
-                new KeyboardBuilder<>(mContext, new KeyboardParams());
-        if (id.isAlphabetKeyboard()) {
-            builder.setAutoGenerate(sKeysCache);
-        }
+                new KeyboardBuilder<>(mContext, new KeyboardParams(sUniqueKeysCache));
+        sUniqueKeysCache.setEnabled(id.isAlphabetKeyboard());
         builder.setAllowRedundantMoreKes(elementParams.mAllowRedundantMoreKeys);
         final int keyboardXmlId = elementParams.mKeyboardXmlId;
         builder.load(keyboardXmlId, id);
diff --git a/java/src/com/android/inputmethod/keyboard/internal/KeyboardBuilder.java b/java/src/com/android/inputmethod/keyboard/internal/KeyboardBuilder.java
index 5743ef9675c5ce2ba7f91bb4133f7e82fbdd11d8..2b07e1d46671e154d0ce2c2f3173eab26cb5e227 100644
--- a/java/src/com/android/inputmethod/keyboard/internal/KeyboardBuilder.java
+++ b/java/src/com/android/inputmethod/keyboard/internal/KeyboardBuilder.java
@@ -161,10 +161,6 @@ public class KeyboardBuilder<KP extends KeyboardParams> {
         params.GRID_HEIGHT = res.getInteger(R.integer.config_keyboard_grid_height);
     }
 
-    public void setAutoGenerate(final KeysCache keysCache) {
-        mParams.mKeysCache = keysCache;
-    }
-
     public void setAllowRedundantMoreKes(final boolean enabled) {
         mParams.mAllowRedundantMoreKeys = enabled;
     }
diff --git a/java/src/com/android/inputmethod/keyboard/internal/KeyboardParams.java b/java/src/com/android/inputmethod/keyboard/internal/KeyboardParams.java
index 4326876351265f96ed54845d999af03d5f0532d1..738d6a40050c593d4626b19bc1a31784880a8045 100644
--- a/java/src/com/android/inputmethod/keyboard/internal/KeyboardParams.java
+++ b/java/src/com/android/inputmethod/keyboard/internal/KeyboardParams.java
@@ -77,9 +77,8 @@ public class KeyboardParams {
     @Nonnull
     public final KeyStylesSet mKeyStyles = new KeyStylesSet(mTextsSet);
 
-    // TODO: Make this @Nonnull
-    @Nullable
-    public KeysCache mKeysCache;
+    @Nonnull
+    private final UniqueKeysCache mUniqueKeysCache;
     public boolean mAllowRedundantMoreKeys;
 
     public int mMostCommonKeyHeight = 0;
@@ -103,6 +102,14 @@ public class KeyboardParams {
         }
     };
 
+    public KeyboardParams() {
+        this(UniqueKeysCache.NO_CACHE);
+    }
+
+    public KeyboardParams(@Nonnull final UniqueKeysCache keysCache) {
+        mUniqueKeysCache = keysCache;
+    }
+
     protected void clearKeys() {
         mSortedKeys.clear();
         mShiftKeys.clear();
@@ -110,9 +117,7 @@ public class KeyboardParams {
     }
 
     public void onAddKey(@Nonnull final Key newKey) {
-        // To avoid possible null pointer access.
-        final KeysCache keysCache = mKeysCache;
-        final Key key = (keysCache != null) ? keysCache.get(newKey) : newKey;
+        final Key key = mUniqueKeysCache.getUniqueKey(newKey);
         final boolean isSpacer = key.isSpacer();
         if (isSpacer && key.getWidth() == 0) {
             // Ignore zero width {@link Spacer}.
@@ -140,16 +145,11 @@ public class KeyboardParams {
         for (final Key key : mSortedKeys) {
             lettersOnBaseLayout.addLetter(key);
         }
-        // To avoid possible null pointer access.
-        final KeysCache keysCache = mKeysCache;
         final ArrayList<Key> allKeys = new ArrayList<>(mSortedKeys);
         mSortedKeys.clear();
         for (final Key key : allKeys) {
             final Key filteredKey = Key.removeRedundantMoreKeys(key, lettersOnBaseLayout);
-            if (keysCache != null) {
-                keysCache.replace(key, filteredKey);
-            }
-            mSortedKeys.add(filteredKey);
+            mSortedKeys.add(mUniqueKeysCache.getUniqueKey(filteredKey));
         }
     }
 
diff --git a/java/src/com/android/inputmethod/keyboard/internal/KeysCache.java b/java/src/com/android/inputmethod/keyboard/internal/KeysCache.java
deleted file mode 100644
index 6ad450c291cf80497483e8c5a0a46b7a0bb6064b..0000000000000000000000000000000000000000
--- a/java/src/com/android/inputmethod/keyboard/internal/KeysCache.java
+++ /dev/null
@@ -1,50 +0,0 @@
-/*
- * Copyright (C) 2012 The Android Open Source Project
- *
- * 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 com.android.inputmethod.keyboard.internal;
-
-import com.android.inputmethod.keyboard.Key;
-
-import java.util.HashMap;
-
-// TODO: Rename more appropriate name.
-public final class KeysCache {
-    private final HashMap<Key, Key> mMap = new HashMap<>();
-
-    public void clear() {
-        mMap.clear();
-    }
-
-    // TODO: Rename more descriptive name.
-    public Key get(final Key key) {
-        final Key existingKey = mMap.get(key);
-        if (existingKey != null) {
-            // Reuse the existing element that equals to "key" without adding "key" to the map.
-            return existingKey;
-        }
-        mMap.put(key, key);
-        return key;
-    }
-
-    // TODO: Rename more descriptive name.
-    public Key replace(final Key oldKey, final Key newKey) {
-        if (oldKey.equals(newKey)) {
-            return oldKey;
-        }
-        mMap.remove(oldKey);
-        return get(newKey);
-    }
-}
diff --git a/java/src/com/android/inputmethod/keyboard/internal/UniqueKeysCache.java b/java/src/com/android/inputmethod/keyboard/internal/UniqueKeysCache.java
new file mode 100644
index 0000000000000000000000000000000000000000..5b329dce4b65dd94eb7aaa1e7a618c680916615d
--- /dev/null
+++ b/java/src/com/android/inputmethod/keyboard/internal/UniqueKeysCache.java
@@ -0,0 +1,81 @@
+/*
+ * Copyright (C) 2014 The Android Open Source Project
+ *
+ * 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 com.android.inputmethod.keyboard.internal;
+
+import com.android.inputmethod.keyboard.Key;
+
+import java.util.HashMap;
+
+import javax.annotation.Nonnull;
+
+public abstract class UniqueKeysCache {
+    public abstract void setEnabled(boolean enabled);
+    public abstract void clear();
+    public abstract @Nonnull Key getUniqueKey(@Nonnull Key key);
+
+    @Nonnull
+    public static final UniqueKeysCache NO_CACHE = new UniqueKeysCache() {
+        @Override
+        public void setEnabled(boolean enabled) {}
+
+        @Override
+        public void clear() {}
+
+        @Override
+        public Key getUniqueKey(Key key) { return key; }
+    };
+
+    @Nonnull
+    public static UniqueKeysCache newInstance() {
+        return new UniqueKeysCacheImpl();
+    }
+
+    private static final class UniqueKeysCacheImpl extends UniqueKeysCache {
+        private final HashMap<Key, Key> mCache;
+
+        private boolean mEnabled;
+
+        UniqueKeysCacheImpl() {
+            mCache = new HashMap<>();
+        }
+
+        @Override
+        public void setEnabled(final boolean enabled) {
+            mEnabled = enabled;
+        }
+
+        @Override
+        public void clear() {
+            mCache.clear();
+        }
+
+        @Override
+        public Key getUniqueKey(final Key key) {
+            if (!mEnabled) {
+                return key;
+            }
+            final Key existingKey = mCache.get(key);
+            if (existingKey != null) {
+                // Reuse the existing object that equals to "key" without adding "key" to
+                // the cache.
+                return existingKey;
+            }
+            mCache.put(key, key);
+            return key;
+        }
+    }
+}
diff --git a/tests/src/com/android/inputmethod/keyboard/layout/tests/LayoutTestsBase.java b/tests/src/com/android/inputmethod/keyboard/layout/tests/LayoutTestsBase.java
index a8c4ac8faaa0a35eba96bf8c61400e8b6334a86a..27519ee93d5ab7becb95a13d2f9f227dcb1ebf87 100644
--- a/tests/src/com/android/inputmethod/keyboard/layout/tests/LayoutTestsBase.java
+++ b/tests/src/com/android/inputmethod/keyboard/layout/tests/LayoutTestsBase.java
@@ -120,31 +120,13 @@ abstract class LayoutTestsBase extends KeyboardLayoutSetTestsBase {
 
     // TODO: Add phone, phone symbols, number, number password layout tests.
 
-    public final void testAlphabet() {
+    public final void testLayouts() {
         doKeyboardTests(KeyboardId.ELEMENT_ALPHABET);
-    }
-
-    public final void testAlphabetAutomaticShifted() {
         doKeyboardTests(KeyboardId.ELEMENT_ALPHABET_AUTOMATIC_SHIFTED);
-    }
-
-    public final void testAlphabetManualShifted() {
         doKeyboardTests(KeyboardId.ELEMENT_ALPHABET_MANUAL_SHIFTED);
-    }
-
-    public final void testAlphabetShiftLocked() {
         doKeyboardTests(KeyboardId.ELEMENT_ALPHABET_SHIFT_LOCKED);
-    }
-
-    public final void testAlphabetShiftLockShifted() {
         doKeyboardTests(KeyboardId.ELEMENT_ALPHABET_SHIFT_LOCK_SHIFTED);
-    }
-
-    public final void testSymbols() {
         doKeyboardTests(KeyboardId.ELEMENT_SYMBOLS);
-    }
-
-    public final void testSymbolsShifted() {
         doKeyboardTests(KeyboardId.ELEMENT_SYMBOLS_SHIFTED);
     }