[NavigationBar] Set badges to be null instead of removing them from the sparse array so that removing badges before restoring badge states will not override the current state

PiperOrigin-RevId: 546027617

Author: imhappi

lib/java/com/google/android/material/badge/BadgeUtils.java

@@ -266,10 +266,7 @@ public static ParcelableSparseArray createParcelableBadgeStates(
     for (int i = 0; i < badgeDrawables.size(); i++) {
       int key = badgeDrawables.keyAt(i);
       BadgeDrawable badgeDrawable = badgeDrawables.valueAt(i);
-      if (badgeDrawable == null) {
-        throw new IllegalArgumentException("badgeDrawable cannot be null");
-      }
-      badgeStates.put(key, badgeDrawable.getSavedState());
+      badgeStates.put(key, badgeDrawable != null ? badgeDrawable.getSavedState() : null);
     }
     return badgeStates;
   }
@@ -291,10 +288,10 @@ public static SparseArray<BadgeDrawable> createBadgeDrawablesFromSavedStates(
     for (int i = 0; i < badgeStates.size(); i++) {
       int key = badgeStates.keyAt(i);
       BadgeState.State savedState = (BadgeState.State) badgeStates.valueAt(i);
-      if (savedState == null) {
-        throw new IllegalArgumentException("BadgeDrawable's savedState cannot be null");
+      BadgeDrawable badgeDrawable = null;
+      if (savedState != null) {
+        badgeDrawable = BadgeDrawable.createFromSavedState(context, savedState);
       }
-      BadgeDrawable badgeDrawable = BadgeDrawable.createFromSavedState(context, savedState);
       badgeDrawables.put(key, badgeDrawable);
     }
     return badgeDrawables;

lib/java/com/google/android/material/navigation/NavigationBarMenuView.java

@@ -869,7 +869,10 @@ void restoreBadgeDrawables(SparseArray<BadgeDrawable> badgeDrawables) {
     }
     if (buttons != null) {
       for (NavigationBarItemView itemView : buttons) {
-        itemView.setBadge(this.badgeDrawables.get(itemView.getId()));
+        BadgeDrawable badge = this.badgeDrawables.get(itemView.getId());
+        if (badge != null) {
+          itemView.setBadge(badge);
+        }
       }
     }
   }
@@ -903,14 +906,11 @@ BadgeDrawable getOrCreateBadge(int menuItemId) {
 
   void removeBadge(int menuItemId) {
     validateMenuItemId(menuItemId);
-    BadgeDrawable badgeDrawable = badgeDrawables.get(menuItemId);
     NavigationBarItemView itemView = findItemView(menuItemId);
     if (itemView != null) {
       itemView.removeBadge();
     }
-    if (badgeDrawable != null) {
-      badgeDrawables.remove(menuItemId);
-    }
+    badgeDrawables.put(menuItemId, null);
   }
 
   private void setBadgeIfNeeded(@NonNull NavigationBarItemView child) {