diff --git a/java/src/com/android/inputmethod/keyboard/KeyboardSwitcher.java b/java/src/com/android/inputmethod/keyboard/KeyboardSwitcher.java
index c02b37b82..d38f25496 100644
--- a/java/src/com/android/inputmethod/keyboard/KeyboardSwitcher.java
+++ b/java/src/com/android/inputmethod/keyboard/KeyboardSwitcher.java
@@ -81,8 +81,6 @@ public class KeyboardSwitcher implements KeyboardState.SwitchActions,
     private final HashMap<KeyboardId, SoftReference<LatinKeyboard>> mKeyboardCache =
             new HashMap<KeyboardId, SoftReference<LatinKeyboard>>();
 
-    private KeyboardLayoutState mSavedKeyboardState = new KeyboardLayoutState();
-
     /** mIsAutoCorrectionActive indicates that auto corrected word will be input instead of
      * what user actually typed. */
     private boolean mIsAutoCorrectionActive;
@@ -92,56 +90,6 @@ public class KeyboardSwitcher implements KeyboardState.SwitchActions,
 
     private static final KeyboardSwitcher sInstance = new KeyboardSwitcher();
 
-    // TODO: Move this to KeyboardState.
-    private class KeyboardLayoutState {
-        private boolean mIsValid;
-        private boolean mIsAlphabetMode;
-        private boolean mIsShiftLocked;
-        private boolean mIsShifted;
-
-        public void save() {
-            mIsAlphabetMode = isAlphabetMode();
-            if (mIsAlphabetMode) {
-                mIsShiftLocked = mState.isShiftLocked();
-                mIsShifted = !mIsShiftLocked && mState.isShiftedOrShiftLocked();
-            } else {
-                mIsShiftLocked = false;
-                mIsShifted = isSymbolShifted();
-            }
-            mIsValid = true;
-            if (DEBUG_STATE) {
-                Log.d(TAG, "save: alphabet=" + mIsAlphabetMode + " shiftLocked=" + mIsShiftLocked
-                        + " shift=" + mIsShifted);
-            }
-        }
-
-        public void restore() {
-            if (DEBUG_STATE) {
-                Log.d(TAG, "restore: valid=" + mIsValid + " alphabet=" + mIsAlphabetMode
-                        + " shiftLocked=" + mIsShiftLocked + " shift=" + mIsShifted);
-            }
-            if (!mIsValid || mIsAlphabetMode) {
-                setAlphabetKeyboard();
-            } else {
-                if (mIsShifted) {
-                    setSymbolsShiftedKeyboard();
-                } else {
-                    setSymbolsKeyboard();
-                }
-            }
-
-            if (!mIsValid) return;
-            mIsValid = false;
-
-            if (mIsAlphabetMode) {
-                setShiftLocked(mIsShiftLocked);
-                if (!mIsShiftLocked) {
-                    setShifted(mIsShifted ? MANUAL_SHIFT : UNSHIFT);
-                }
-            }
-        }
-    }
-
     public static KeyboardSwitcher getInstance() {
         return sInstance;
     }
@@ -194,7 +142,7 @@ public class KeyboardSwitcher implements KeyboardState.SwitchActions,
             mSymbolsShiftedKeyboardId = getKeyboardId(editorInfo, true, true, settingsValues);
             mState.onLoadKeyboard(mResources.getString(R.string.layout_switch_back_symbols));
             mPrevMainKeyboardWasShiftLocked = false;
-            mSavedKeyboardState.restore();
+            mState.onRestoreKeyboardState();
         } catch (RuntimeException e) {
             Log.w(TAG, "loading keyboard failed: " + mMainKeyboardId, e);
             LatinImeLogger.logOnException(mMainKeyboardId.toString(), e);
@@ -203,7 +151,7 @@ public class KeyboardSwitcher implements KeyboardState.SwitchActions,
 
     public void saveKeyboardState() {
         if (mCurrentId != null) {
-            mSavedKeyboardState.save();
+            mState.onSaveKeyboardState(isAlphabetMode(), isSymbolShifted());
         }
     }
 
diff --git a/java/src/com/android/inputmethod/keyboard/internal/KeyboardState.java b/java/src/com/android/inputmethod/keyboard/internal/KeyboardState.java
index a632a27b1..ee5ef9164 100644
--- a/java/src/com/android/inputmethod/keyboard/internal/KeyboardState.java
+++ b/java/src/com/android/inputmethod/keyboard/internal/KeyboardState.java
@@ -17,11 +17,15 @@
 package com.android.inputmethod.keyboard.internal;
 
 import android.text.TextUtils;
+import android.util.Log;
 
 import com.android.inputmethod.keyboard.Keyboard;
 
 // TODO: Add unit tests
 public class KeyboardState {
+    private static final String TAG = KeyboardState.class.getSimpleName();
+    private static final boolean DEBUG_STATE = false;
+
     public interface SwitchActions {
         public void setAlphabetKeyboard();
         public static final int UNSHIFT = 0;
@@ -52,6 +56,15 @@ public class KeyboardState {
 
     private final SwitchActions mSwitchActions;
 
+    private final SavedKeyboardState mSavedKeyboardState = new SavedKeyboardState();
+
+    private static class SavedKeyboardState {
+        public boolean mIsValid;
+        public boolean mIsAlphabetMode;
+        public boolean mIsShiftLocked;
+        public boolean mIsShifted;
+    }
+
     public KeyboardState(SwitchActions switchActions) {
         mSwitchActions = switchActions;
     }
@@ -64,6 +77,52 @@ public class KeyboardState {
         mSymbolKeyState.onRelease();
     }
 
+    // TODO: Get rid of isAlphabetMode and isSymbolShifted arguments.
+    public void onSaveKeyboardState(boolean isAlphabetMode, boolean isSymbolShifted) {
+        final SavedKeyboardState state = mSavedKeyboardState;
+        state.mIsAlphabetMode = isAlphabetMode;
+        if (isAlphabetMode) {
+            state.mIsShiftLocked = isShiftLocked();
+            state.mIsShifted = !state.mIsShiftLocked && isShiftedOrShiftLocked();
+        } else {
+            state.mIsShiftLocked = false;
+            state.mIsShifted = isSymbolShifted;
+        }
+        state.mIsValid = true;
+        if (DEBUG_STATE) {
+            Log.d(TAG, "save: alphabet=" + state.mIsAlphabetMode
+                    + " shiftLocked=" + state.mIsShiftLocked + " shift=" + state.mIsShifted);
+        }
+    }
+
+    public void onRestoreKeyboardState() {
+        final SavedKeyboardState state = mSavedKeyboardState;
+        if (DEBUG_STATE) {
+            Log.d(TAG, "restore: valid=" + state.mIsValid + " alphabet=" + state.mIsAlphabetMode
+                    + " shiftLocked=" + state.mIsShiftLocked + " shift=" + state.mIsShifted);
+        }
+        if (!state.mIsValid || state.mIsAlphabetMode) {
+            mSwitchActions.setAlphabetKeyboard();
+        } else {
+            if (state.mIsShifted) {
+                mSwitchActions.setSymbolsShiftedKeyboard();
+            } else {
+                mSwitchActions.setSymbolsKeyboard();
+            }
+        }
+
+        if (!state.mIsValid) return;
+        state.mIsValid = false;
+
+        if (state.mIsAlphabetMode) {
+            mSwitchActions.setShiftLocked(state.mIsShiftLocked);
+            if (!state.mIsShiftLocked) {
+                mSwitchActions.setShifted(
+                        state.mIsShifted ? SwitchActions.MANUAL_SHIFT : SwitchActions.UNSHIFT);
+            }
+        }
+    }
+
     // TODO: Get rid of this method
     public void onSetKeyboard(boolean isAlphabetMode) {
         mSwitchState = isAlphabetMode ? SWITCH_STATE_ALPHA : SWITCH_STATE_SYMBOL_BEGIN;