From 62d2db1d60a6161f82f14e4752ed44056e64c185 Mon Sep 17 00:00:00 2001
From: Mario Bertschler <bmario@google.com>
Date: Tue, 16 Jan 2018 14:10:28 -0800
Subject: [PATCH] Fix for badge updates in PredictionRowView

Bug: 72054359
Change-Id: I6d435c380d182ba99f9627c049f7f36ad331ee1d
---
 .../allapps/AllAppsContainerView.java         | 36 ++++++++++++-------
 1 file changed, 24 insertions(+), 12 deletions(-)

diff --git a/src/com/android/launcher3/allapps/AllAppsContainerView.java b/src/com/android/launcher3/allapps/AllAppsContainerView.java
index a40f8b3809..47dc566dfc 100644
--- a/src/com/android/launcher3/allapps/AllAppsContainerView.java
+++ b/src/com/android/launcher3/allapps/AllAppsContainerView.java
@@ -364,18 +364,30 @@ public class AllAppsContainerView extends RelativeLayout implements DragSource,
     public void updateIconBadges(Set<PackageUserKey> updatedBadges) {
         final PackageUserKey packageUserKey = new PackageUserKey(null, null);
         for (int j = 0; j < mAH.length; j++) {
-            if (mAH[j].recyclerView != null) {
-                final int n = mAH[j].recyclerView.getChildCount();
-                for (int i = 0; i < n; i++) {
-                    View child = mAH[j].recyclerView.getChildAt(i);
-                    if (!(child instanceof BubbleTextView) || !(child.getTag() instanceof ItemInfo)) {
-                        continue;
-                    }
-                    ItemInfo info = (ItemInfo) child.getTag();
-                    if (packageUserKey.updateFromItemInfo(info) && updatedBadges.contains(packageUserKey)) {
-                        ((BubbleTextView) child).applyBadgeState(info, true /* animate */);
-                    }
-                }
+            updateIconBadges(updatedBadges, packageUserKey, mAH[j].recyclerView);
+        }
+        if (mHeader != null) {
+            updateIconBadges(updatedBadges, packageUserKey, mHeader.getPredictionRow());
+        }
+    }
+
+    private void updateIconBadges(Set<PackageUserKey> updatedBadges, PackageUserKey packageUserKey,
+            ViewGroup parent) {
+        if (parent == null) {
+            return;
+        }
+        final int n = parent.getChildCount();
+        for (int i = 0; i < n; i++) {
+            View child = parent.getChildAt(i);
+            if (child instanceof PredictionRowView) {
+                updateIconBadges(updatedBadges, packageUserKey, (PredictionRowView) child);
+            }
+            if (!(child instanceof BubbleTextView) || !(child.getTag() instanceof ItemInfo)) {
+                continue;
+            }
+            ItemInfo info = (ItemInfo) child.getTag();
+            if (packageUserKey.updateFromItemInfo(info) && updatedBadges.contains(packageUserKey)) {
+                ((BubbleTextView) child).applyBadgeState(info, true /* animate */);
             }
         }
     }