Ability to pass a FactoryInstance to WaveformAudioBufferSink which will create custom WaveformBar instances

libraries/exoplayer/src/main/java/androidx/media3/exoplayer/audio/WaveformAudioBufferSink.java

@@ -89,6 +89,13 @@ public interface Listener {
     void onNewWaveformBar(int channelIndex, WaveformBar waveformBar);
   }
 
+  public static class WaveformBarFactory {
+    /** Called when a new waveform bar has to be created */
+    WaveformBar createWaveformBar() {
+      return new WaveformBar();
+    }
+  }
+
   private final int barsPerSecond;
   private final Listener listener;
   private final SparseArray<WaveformBar> outputChannels;
@@ -98,6 +105,13 @@ public interface Listener {
   private @MonotonicNonNull ChannelMixingMatrix channelMixingMatrix;
   private int samplesPerBar;
 
+  private WaveformBarFactory waveformBarFactory;
+
+  /**
+   * Default factory instance that creates standard WaveformBar objects.
+   */
+  private static final WaveformBarFactory DEFAULT_WAVEFORM_BAR_FACTORY = new WaveformBarFactory();
+
   /**
    * Creates a new instance.
    *
@@ -108,13 +122,28 @@ public interface Listener {
    * @param listener The listener to be notified when a new waveform bar has been generated.
    */
   public WaveformAudioBufferSink(int barsPerSecond, int outputChannelCount, Listener listener) {
+    this(barsPerSecond, outputChannelCount, listener, DEFAULT_WAVEFORM_BAR_FACTORY);
+  }
+
+  /**
+   * Creates a new instance.
+   *
+   * @param barsPerSecond The number of bars that should be generated per each second of audio.
+   * @param outputChannelCount The number of channels that the output waveform should contain. If
+   *     this is different than the number of input channels, the audio will be mixed using the
+   *     {@linkplain ChannelMixingMatrix#createForConstantGain default mixing matrix}.
+   * @param listener The listener to be notified when a new waveform bar has been generated.
+   * @param waveformBarFactory The factory to create new waveform bars.
+   */
+  public WaveformAudioBufferSink(int barsPerSecond, int outputChannelCount, Listener listener, WaveformBarFactory waveformBarFactory) {
     this.barsPerSecond = barsPerSecond;
     this.listener = listener;
+    this.waveformBarFactory = waveformBarFactory;
     mixingBuffer =
         ByteBuffer.allocate(Util.getPcmFrameSize(C.ENCODING_PCM_FLOAT, outputChannelCount));
     outputChannels = new SparseArray<>(outputChannelCount);
     for (int i = 0; i < outputChannelCount; i++) {
-      outputChannels.append(i, new WaveformBar());
+      outputChannels.append(i, waveformBarFactory.createWaveformBar());
     }
   }
 
@@ -149,7 +178,7 @@ public void handleBuffer(ByteBuffer buffer) {
         bar.addSample(mixingBuffer.getFloat());
         if (bar.getSampleCount() >= samplesPerBar) {
           listener.onNewWaveformBar(i, bar);
-          outputChannels.set(i, new WaveformBar());
+          outputChannels.set(i, waveformBarFactory.createWaveformBar());
         }
       }
     }

libraries/exoplayer/src/test/java/androidx/media3/exoplayer/audio/WaveformAudioBufferSinkTest.java

@@ -146,6 +146,60 @@ public void handleBuffer_monoToMonoFloat_callbackHasExpectedValue() throws Excep
         .of(0.4249);
   }
 
+  @Test
+  public void handleBuffer_customWaveformBarFactory_usesCustomFactory() throws Exception {
+    // Custom WaveformBar that tracks when it's created
+    class TrackedWaveformBar extends WaveformAudioBufferSink.WaveformBar {
+    }
+
+    // Custom factory that creates tracked bars
+    class TrackedWaveformBarFactory extends WaveformAudioBufferSink.WaveformBarFactory {
+
+      private int creationCount = 0;
+
+      @Override
+      WaveformAudioBufferSink.WaveformBar createWaveformBar() {
+        ++creationCount;
+        return new TrackedWaveformBar();
+      }
+
+      int getCreationCount() {
+        return creationCount;
+      }
+    }
+
+    TrackedWaveformBarFactory customFactory = new TrackedWaveformBarFactory();
+    List<WaveformAudioBufferSink.WaveformBar> receivedBars = new ArrayList<>();
+    CountDownLatch countDownLatch = new CountDownLatch(2); // Expecting 2 bars to be created
+
+    WaveformAudioBufferSink waveformAudioBufferSink =
+        new WaveformAudioBufferSink(
+            /* barsPerSecond= */ 2,
+            /* outputChannelCount= */ 1,
+            (channelIndex, bar) -> {
+              receivedBars.add(bar);
+              countDownLatch.countDown();
+            },
+            customFactory);
+
+    // Create a buffer that will trigger bar generation twice
+    ByteBuffer byteBuffer = ByteBuffer.allocate(8);
+    byteBuffer.putShort(0, (short) 1000);
+    byteBuffer.putShort(2, (short) 2000);
+    byteBuffer.putShort(4, (short) 3000);
+    byteBuffer.putShort(6, (short) 4000);
+
+    int sampleRateHz = 4; // 4 samples total, 2 bars per second = 2 samples per bar
+    waveformAudioBufferSink.flush(sampleRateHz, /* inputChannelCount= */ 1, C.ENCODING_PCM_16BIT);
+    waveformAudioBufferSink.handleBuffer(byteBuffer);
+
+    assertThat(countDownLatch.await(TIMEOUT_MS, MILLISECONDS)).isTrue();
+
+    // Verify that the custom factory was used
+    assertThat(customFactory.getCreationCount()).isEqualTo(3);
+    assertThat(receivedBars).hasSize(2);
+  }
+
   private ImmutableList<WaveformAudioBufferSink.WaveformBar> calculateChannelWaveformBars(
       ByteBuffer byteBuffer,
       int inputChannelCount,