diff --git a/extensions/ima/src/main/java/com/google/android/exoplayer2/ext/ima/ImaAdsLoader.java b/extensions/ima/src/main/java/com/google/android/exoplayer2/ext/ima/ImaAdsLoader.java
index 4f20535ff1..a8573f31e3 100644
--- a/extensions/ima/src/main/java/com/google/android/exoplayer2/ext/ima/ImaAdsLoader.java
+++ b/extensions/ima/src/main/java/com/google/android/exoplayer2/ext/ima/ImaAdsLoader.java
@@ -797,7 +797,7 @@ public final class ImaAdsLoader extends Player.DefaultEventListener implements A
     long[] adGroupTimesUs = getAdGroupTimesUs(adsManager.getAdCuePoints());
     adPlaybackState = new AdPlaybackState(adGroupTimesUs);
     int adGroupIndexForPosition =
-        getAdGroupIndexForPosition(adGroupTimesUs, C.msToUs(pendingContentPositionMs));
+        adPlaybackState.getAdGroupIndexForPositionUs(C.msToUs(pendingContentPositionMs));
     if (adGroupIndexForPosition == 0) {
       podIndexOffset = 0;
     } else if (adGroupIndexForPosition == C.INDEX_UNSET) {
@@ -952,23 +952,6 @@ public final class ImaAdsLoader extends Player.DefaultEventListener implements A
     }
   }
 
-  /**
-   * Returns the index of the ad group that should be played before playing the content at {@code
-   * playbackPositionUs} when starting playback for the first time. This is the latest ad group at
-   * or before the specified playback position. If the first ad is after the playback position,
-   * returns {@link C#INDEX_UNSET}.
-   */
-  private int getAdGroupIndexForPosition(long[] adGroupTimesUs, long playbackPositionUs) {
-    for (int i = 0; i < adGroupTimesUs.length; i++) {
-      long adGroupTimeUs = adGroupTimesUs[i];
-      // A postroll ad is after any position in the content.
-      if (adGroupTimeUs == C.TIME_END_OF_SOURCE || playbackPositionUs < adGroupTimeUs) {
-        return i == 0 ? C.INDEX_UNSET : (i - 1);
-      }
-    }
-    return adGroupTimesUs.length == 0 ? C.INDEX_UNSET : (adGroupTimesUs.length - 1);
-  }
-
   /**
    * Returns the next ad index in the specified ad group to load, or {@link C#INDEX_UNSET} if all
    * ads in the ad group have loaded.
diff --git a/library/core/src/main/java/com/google/android/exoplayer2/Timeline.java b/library/core/src/main/java/com/google/android/exoplayer2/Timeline.java
index 4b58a1075c..50a3e66880 100644
--- a/library/core/src/main/java/com/google/android/exoplayer2/Timeline.java
+++ b/library/core/src/main/java/com/google/android/exoplayer2/Timeline.java
@@ -413,53 +413,30 @@ public abstract class Timeline {
      * @return Whether the ad group at index {@code adGroupIndex} has been played.
      */
     public boolean hasPlayedAdGroup(int adGroupIndex) {
-      AdPlaybackState.AdGroup adGroup = adPlaybackState.adGroups[adGroupIndex];
-      return adGroup.getFirstAdIndexToPlay() == adGroup.count;
+      return !adPlaybackState.adGroups[adGroupIndex].hasUnplayedAds();
     }
 
     /**
      * Returns the index of the ad group at or before {@code positionUs}, if that ad group is
-     * unplayed. Returns {@link C#INDEX_UNSET} if the ad group before {@code positionUs} has been
-     * played, or if there is no such ad group.
+     * unplayed. Returns {@link C#INDEX_UNSET} if the ad group at or before {@code positionUs} has
+     * no ads remaining to be played, or if there is no such ad group.
      *
      * @param positionUs The position at or before which to find an ad group, in microseconds.
      * @return The index of the ad group, or {@link C#INDEX_UNSET}.
      */
     public int getAdGroupIndexForPositionUs(long positionUs) {
-      long[] adGroupTimesUs = adPlaybackState.adGroupTimesUs;
-      if (adGroupTimesUs == null) {
-        return C.INDEX_UNSET;
-      }
-      // Use a linear search as the array elements may not be increasing due to TIME_END_OF_SOURCE.
-      // In practice we expect there to be few ad groups so the search shouldn't be expensive.
-      int index = adGroupTimesUs.length - 1;
-      while (index >= 0 && (adGroupTimesUs[index] == C.TIME_END_OF_SOURCE
-          || adGroupTimesUs[index] > positionUs)) {
-        index--;
-      }
-      return index >= 0 && !hasPlayedAdGroup(index) ? index : C.INDEX_UNSET;
+      return adPlaybackState.getAdGroupIndexForPositionUs(positionUs);
     }
 
     /**
-     * Returns the index of the next unplayed ad group after {@code positionUs}. Returns
-     * {@link C#INDEX_UNSET} if there is no such ad group.
+     * Returns the index of the next ad group after {@code positionUs} that has ads remaining to be
+     * played. Returns {@link C#INDEX_UNSET} if there is no such ad group.
      *
      * @param positionUs The position after which to find an ad group, in microseconds.
      * @return The index of the ad group, or {@link C#INDEX_UNSET}.
      */
     public int getAdGroupIndexAfterPositionUs(long positionUs) {
-      long[] adGroupTimesUs = adPlaybackState.adGroupTimesUs;
-      if (adGroupTimesUs == null) {
-        return C.INDEX_UNSET;
-      }
-      // Use a linear search as the array elements may not be increasing due to TIME_END_OF_SOURCE.
-      // In practice we expect there to be few ad groups so the search shouldn't be expensive.
-      int index = 0;
-      while (index < adGroupTimesUs.length && adGroupTimesUs[index] != C.TIME_END_OF_SOURCE
-          && (positionUs >= adGroupTimesUs[index] || hasPlayedAdGroup(index))) {
-        index++;
-      }
-      return index < adGroupTimesUs.length ? index : C.INDEX_UNSET;
+      return adPlaybackState.getAdGroupIndexAfterPositionUs(positionUs);
     }
 
     /**
diff --git a/library/core/src/main/java/com/google/android/exoplayer2/source/ads/AdPlaybackState.java b/library/core/src/main/java/com/google/android/exoplayer2/source/ads/AdPlaybackState.java
index 2f6f81d6b9..8654e94bdb 100644
--- a/library/core/src/main/java/com/google/android/exoplayer2/source/ads/AdPlaybackState.java
+++ b/library/core/src/main/java/com/google/android/exoplayer2/source/ads/AdPlaybackState.java
@@ -91,6 +91,11 @@ public final class AdPlaybackState {
       return nextAdIndexToPlay;
     }
 
+    /** Returns whether the ad group has at least one ad that still needs to be played. */
+    public boolean hasUnplayedAds() {
+      return count == C.LENGTH_UNSET || getFirstAdIndexToPlay() < count;
+    }
+
     /**
      * Returns a new instance with the ad count set to {@code count}. This method may only be called
      * if this instance's ad count has not yet been specified.
@@ -270,6 +275,44 @@ public final class AdPlaybackState {
     this.contentDurationUs = contentDurationUs;
   }
 
+  /**
+   * Returns the index of the ad group at or before {@code positionUs}, if that ad group is
+   * unplayed. Returns {@link C#INDEX_UNSET} if the ad group at or before {@code positionUs} has no
+   * ads remaining to be played, or if there is no such ad group.
+   *
+   * @param positionUs The position at or before which to find an ad group, in microseconds.
+   * @return The index of the ad group, or {@link C#INDEX_UNSET}.
+   */
+  public int getAdGroupIndexForPositionUs(long positionUs) {
+    // Use a linear search as the array elements may not be increasing due to TIME_END_OF_SOURCE.
+    // In practice we expect there to be few ad groups so the search shouldn't be expensive.
+    int index = adGroupTimesUs.length - 1;
+    while (index >= 0
+        && (adGroupTimesUs[index] == C.TIME_END_OF_SOURCE || adGroupTimesUs[index] > positionUs)) {
+      index--;
+    }
+    return index >= 0 && adGroups[index].hasUnplayedAds() ? index : C.INDEX_UNSET;
+  }
+
+  /**
+   * Returns the index of the next ad group after {@code positionUs} that has ads remaining to be
+   * played. Returns {@link C#INDEX_UNSET} if there is no such ad group.
+   *
+   * @param positionUs The position after which to find an ad group, in microseconds.
+   * @return The index of the ad group, or {@link C#INDEX_UNSET}.
+   */
+  public int getAdGroupIndexAfterPositionUs(long positionUs) {
+    // Use a linear search as the array elements may not be increasing due to TIME_END_OF_SOURCE.
+    // In practice we expect there to be few ad groups so the search shouldn't be expensive.
+    int index = 0;
+    while (index < adGroupTimesUs.length
+        && adGroupTimesUs[index] != C.TIME_END_OF_SOURCE
+        && (positionUs >= adGroupTimesUs[index] || !adGroups[index].hasUnplayedAds())) {
+      index++;
+    }
+    return index < adGroupTimesUs.length ? index : C.INDEX_UNSET;
+  }
+
   /**
    * Returns an instance with the number of ads in {@code adGroupIndex} resolved to {@code adCount}.
    * The ad count must be greater than zero.