Ensuring that restoreInstanceState is being called promptly for synchronously bound page

Change-Id: I0e71c29f553ad360ec42a6a0b2529d16cbddd437
diff --git a/src/com/android/launcher2/CellLayout.java b/src/com/android/launcher2/CellLayout.java
index c028ff1..97d9bc1 100644
--- a/src/com/android/launcher2/CellLayout.java
+++ b/src/com/android/launcher2/CellLayout.java
@@ -17,9 +17,9 @@
 package com.android.launcher2;
 
 import android.animation.Animator;
+import android.animation.AnimatorListenerAdapter;
 import android.animation.AnimatorSet;
 import android.animation.ObjectAnimator;
-import android.animation.AnimatorListenerAdapter;
 import android.animation.TimeInterpolator;
 import android.animation.ValueAnimator;
 import android.animation.ValueAnimator.AnimatorUpdateListener;
@@ -37,8 +37,10 @@
 import android.graphics.drawable.ColorDrawable;
 import android.graphics.drawable.Drawable;
 import android.graphics.drawable.NinePatchDrawable;
+import android.os.Parcelable;
 import android.util.AttributeSet;
 import android.util.Log;
+import android.util.SparseArray;
 import android.view.MotionEvent;
 import android.view.View;
 import android.view.ViewDebug;
@@ -532,6 +534,10 @@
         return false;
     }
 
+    public void restoreInstanceState(SparseArray<Parcelable> states) {
+        dispatchRestoreInstanceState(states);
+    }
+
     @Override
     public void cancelLongPress() {
         super.cancelLongPress();
diff --git a/src/com/android/launcher2/Launcher.java b/src/com/android/launcher2/Launcher.java
index 70c641b..ae2ca3b 100644
--- a/src/com/android/launcher2/Launcher.java
+++ b/src/com/android/launcher2/Launcher.java
@@ -278,6 +278,8 @@
     private static Drawable.ConstantState[] sVoiceSearchIcon = new Drawable.ConstantState[2];
     private static Drawable.ConstantState[] sAppMarketIcon = new Drawable.ConstantState[2];
 
+    private final ArrayList<Integer> mSynchronouslyBoundPages = new ArrayList<Integer>();
+
     static final ArrayList<String> sDumpLogs = new ArrayList<String>();
 
     // We only want to get the SharedPreferences once since it does an FS stat each time we get
@@ -1413,9 +1415,11 @@
     }
 
     @Override
-    protected void onRestoreInstanceState(Bundle savedInstanceState) {
-        // Do not call super here
-        mSavedInstanceState = savedInstanceState;
+    public void onRestoreInstanceState(Bundle state) {
+        super.onRestoreInstanceState(state);
+        for (int page: mSynchronouslyBoundPages) {
+            mWorkspace.restoreInstanceStateForChild(page);
+        }
     }
 
     @Override
@@ -3314,6 +3318,10 @@
         }
     }
 
+    public void onPageBoundSynchronously(int page) {
+        mSynchronouslyBoundPages.add(page);
+    }
+
     /**
      * Callback saying that there aren't any more items to bind.
      *
@@ -3329,10 +3337,7 @@
             mSavedState = null;
         }
 
-        if (mSavedInstanceState != null) {
-            super.onRestoreInstanceState(mSavedInstanceState);
-            mSavedInstanceState = null;
-        }
+        mWorkspace.restoreInstanceStateForRemainingPages();
 
         // If we received the result of any pending adds while the loader was running (e.g. the
         // widget configuration forced an orientation change), process them now.
diff --git a/src/com/android/launcher2/LauncherModel.java b/src/com/android/launcher2/LauncherModel.java
index ab29fc6..40a9e21 100644
--- a/src/com/android/launcher2/LauncherModel.java
+++ b/src/com/android/launcher2/LauncherModel.java
@@ -156,6 +156,7 @@
         public boolean isAllAppsVisible();
         public boolean isAllAppsButtonRank(int rank);
         public void bindSearchablesChanged();
+        public void onPageBoundSynchronously(int page);
     }
 
     LauncherModel(LauncherApplication app, IconCache iconCache) {
@@ -1614,6 +1615,17 @@
             // Load items on the current page
             bindWorkspaceItems(oldCallbacks, currentWorkspaceItems, currentAppWidgets,
                     currentFolders, null);
+            if (isLoadingSynchronously) {
+                r = new Runnable() {
+                    public void run() {
+                        Callbacks callbacks = tryGetCallbacks(oldCallbacks);
+                        if (callbacks != null) {
+                            callbacks.onPageBoundSynchronously(currentScreen);
+                        }
+                    }
+                };
+                runOnMainThread(r);
+            }
 
             // Load all the remaining pages (if we are loading synchronously, we want to defer this
             // work until after the first render)
diff --git a/src/com/android/launcher2/Workspace.java b/src/com/android/launcher2/Workspace.java
index 331e86f..44b9f68 100644
--- a/src/com/android/launcher2/Workspace.java
+++ b/src/com/android/launcher2/Workspace.java
@@ -47,6 +47,7 @@
 import android.util.AttributeSet;
 import android.util.DisplayMetrics;
 import android.util.Log;
+import android.util.SparseArray;
 import android.view.Display;
 import android.view.MotionEvent;
 import android.view.View;
@@ -232,6 +233,9 @@
     private int mLastReorderX = -1;
     private int mLastReorderY = -1;
 
+    private SparseArray<Parcelable> mSavedStates;
+    private final ArrayList<Integer> mRestoredPages = new ArrayList<Integer>();
+
     // These variables are used for storing the initial and final values during workspace animations
     private int mSavedScrollX;
     private float mSavedRotationY;
@@ -3411,6 +3415,32 @@
     }
 
     @Override
+    protected void dispatchRestoreInstanceState(SparseArray<Parcelable> container) {
+        // We don't dispatch restoreInstanceState to our children using this code path.
+        // Some pages will be restored immediately as their items are bound immediately, and 
+        // others we will need to wait until after their items are bound.
+        mSavedStates = container;
+    }
+
+    public void restoreInstanceStateForChild(int child) {
+        if (mSavedStates != null) {
+            mRestoredPages.add(child);
+            CellLayout cl = (CellLayout) getChildAt(child);
+            cl.restoreInstanceState(mSavedStates);
+        }
+    }
+
+    public void restoreInstanceStateForRemainingPages() {
+        int count = getChildCount();
+        for (int i = 0; i < count; i++) {
+            if (!mRestoredPages.contains(i)) {
+                restoreInstanceStateForChild(i);
+            }
+        }
+        mRestoredPages.clear();
+    }
+
+    @Override
     public void scrollLeft() {
         if (!isSmall() && !mIsSwitchingState) {
             super.scrollLeft();