From e531e788df581af49fcb5d7d21892528f5e823cf Mon Sep 17 00:00:00 2001
From: bachinger <bachinger@google.com>
Date: Wed, 9 Mar 2022 16:58:46 +0000
Subject: [PATCH] Transform map of AdsLoader.State to a Bundle

This fixes an exception thrown when parceling becasue the type can not be found
(expects the AdsPlaybackState to be Serializable). Transforming the map and the
ad playback states to a Bundle fixes the problem.

#minor-release

PiperOrigin-RevId: 433491993
---
 .../media3/common/AdPlaybackState.java        | 30 +++++++++++++++++++
 .../ImaServerSideAdInsertionMediaSource.java  | 26 ++++++++++------
 2 files changed, 47 insertions(+), 9 deletions(-)

diff --git a/libraries/common/src/main/java/androidx/media3/common/AdPlaybackState.java b/libraries/common/src/main/java/androidx/media3/common/AdPlaybackState.java
index b713f32ef3..ba3caf1d56 100644
--- a/libraries/common/src/main/java/androidx/media3/common/AdPlaybackState.java
+++ b/libraries/common/src/main/java/androidx/media3/common/AdPlaybackState.java
@@ -827,6 +827,36 @@ public final class AdPlaybackState implements Bundleable {
         adsId, adGroups, adResumePositionUs, contentDurationUs, removedAdGroupCount);
   }
 
+  /**
+   * Returns a copy of the ad playback state with the given ads ID.
+   *
+   * @param adsId The new ads ID.
+   * @param adPlaybackState The ad playback state to copy.
+   * @return The new ad playback state.
+   */
+  public static AdPlaybackState fromAdPlaybackState(Object adsId, AdPlaybackState adPlaybackState) {
+    AdGroup[] adGroups =
+        new AdGroup[adPlaybackState.adGroupCount - adPlaybackState.removedAdGroupCount];
+    for (int i = 0; i < adGroups.length; i++) {
+      AdGroup adGroup = adPlaybackState.adGroups[i];
+      adGroups[i] =
+          new AdGroup(
+              adGroup.timeUs,
+              adGroup.count,
+              Arrays.copyOf(adGroup.states, adGroup.states.length),
+              Arrays.copyOf(adGroup.uris, adGroup.uris.length),
+              Arrays.copyOf(adGroup.durationsUs, adGroup.durationsUs.length),
+              adGroup.contentResumeOffsetUs,
+              adGroup.isServerSideInserted);
+    }
+    return new AdPlaybackState(
+        adsId,
+        adGroups,
+        adPlaybackState.adResumePositionUs,
+        adPlaybackState.contentDurationUs,
+        adPlaybackState.removedAdGroupCount);
+  }
+
   @Override
   public boolean equals(@Nullable Object o) {
     if (this == o) {
diff --git a/libraries/exoplayer_ima/src/main/java/androidx/media3/exoplayer/ima/ImaServerSideAdInsertionMediaSource.java b/libraries/exoplayer_ima/src/main/java/androidx/media3/exoplayer/ima/ImaServerSideAdInsertionMediaSource.java
index 5f7dabeff1..4bb90ea914 100644
--- a/libraries/exoplayer_ima/src/main/java/androidx/media3/exoplayer/ima/ImaServerSideAdInsertionMediaSource.java
+++ b/libraries/exoplayer_ima/src/main/java/androidx/media3/exoplayer/ima/ImaServerSideAdInsertionMediaSource.java
@@ -327,23 +327,31 @@ public final class ImaServerSideAdInsertionMediaSource extends CompositeMediaSou
       @Override
       public Bundle toBundle() {
         Bundle bundle = new Bundle();
-        bundle.putSerializable(keyForField(FIELD_AD_PLAYBACK_STATES), adPlaybackStates);
+        Bundle adPlaybackStatesBundle = new Bundle();
+        for (Map.Entry<String, AdPlaybackState> entry : adPlaybackStates.entrySet()) {
+          adPlaybackStatesBundle.putBundle(entry.getKey(), entry.getValue().toBundle());
+        }
+        bundle.putBundle(keyForField(FIELD_AD_PLAYBACK_STATES), adPlaybackStatesBundle);
         return bundle;
       }
 
       /** Object that can restore {@link AdsLoader.State} from a {@link Bundle}. */
       public static final Bundleable.Creator<State> CREATOR = State::fromBundle;
 
-      @SuppressWarnings("unchecked")
       private static State fromBundle(Bundle bundle) {
         @Nullable
-        Map<String, AdPlaybackState> adPlaybackStateMap =
-            (Map<String, AdPlaybackState>)
-                bundle.getSerializable(keyForField(FIELD_AD_PLAYBACK_STATES));
-        return new State(
-            adPlaybackStateMap != null
-                ? ImmutableMap.copyOf(adPlaybackStateMap)
-                : ImmutableMap.of());
+        ImmutableMap.Builder<String, AdPlaybackState> adPlaybackStateMap =
+            new ImmutableMap.Builder<>();
+        Bundle adPlaybackStateBundle =
+            checkNotNull(bundle.getBundle(keyForField(FIELD_AD_PLAYBACK_STATES)));
+        for (String key : adPlaybackStateBundle.keySet()) {
+          AdPlaybackState adPlaybackState =
+              AdPlaybackState.CREATOR.fromBundle(
+                  checkNotNull(adPlaybackStateBundle.getBundle(key)));
+          adPlaybackStateMap.put(
+              key, AdPlaybackState.fromAdPlaybackState(/* adsId= */ key, adPlaybackState));
+        }
+        return new State(adPlaybackStateMap.buildOrThrow());
       }
 
       private static String keyForField(@FieldNumber int field) {