// Copyright 2012 The Chromium Authors
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.


#ifndef MEDIA_BASE_AUDIO_BUS_H_
#define MEDIA_BASE_AUDIO_BUS_H_

#include <stdint.h>

#include <memory>
#include <vector>

#include "base/compiler_specific.h"
#include "base/containers/span.h"
#include "base/functional/callback.h"
#include "base/memory/aligned_memory.h"
#include "base/memory/raw_ptr_exclusion.h"
#include "base/memory/raw_span.h"
#include "base/numerics/checked_math.h"
#include "base/numerics/safe_conversions.h"
#include "media/base/media_export.h"

namespace media {
class AudioParameters;

// Represents a sequence of audio frames containing frames() audio samples for
// each of channels() channels. The data is stored as a set of contiguous
// float arrays with one array per channel. The memory for the arrays is either
// allocated and owned by the AudioBus or it is provided to one of the factory
// methods. AudioBus guarantees that it allocates memory such that float array
// for each channel is aligned by AudioBus::kChannelAlignment bytes and it
// requires the same for memory passed to its Wrap...() factory methods.
class MEDIA_EXPORT AudioBus {
 public:
  using BitstreamData = base::span<uint8_t>;
  using Channel = base::span<float>;
  using ConstChannel = base::span<const float>;
  using ChannelVector = std::vector<Channel>;

  // Guaranteed alignment of each channel's data; use 16-byte alignment for easy
  // SSE optimizations.
  static constexpr size_t kChannelAlignment = 16;

  // Creates a new AudioBus and allocates |channels| of length |frames|.  Uses
  // channels() and frames_per_buffer() from AudioParameters if given.
  static std::unique_ptr<AudioBus> Create(int channels, int frames);
  static std::unique_ptr<AudioBus> Create(const AudioParameters& params);

  // Creates a new AudioBus with the given number of channels, but zero length.
  // Clients are expected to subsequently call SetChannelData() and set_frames()
  // to wrap externally allocated memory.
  static std::unique_ptr<AudioBus> CreateWrapper(int channels);

  // Creates a new AudioBus by wrapping an existing block of memory.  Block must
  // be at least CalculateMemorySize() bytes in size.  |data| must outlive the
  // returned AudioBus.  |data| must be aligned by kChannelAlignment.
  static std::unique_ptr<AudioBus> WrapMemory(int channels,
                                              int frames,
                                              base::span<float> data);
  static std::unique_ptr<AudioBus> WrapMemory(const AudioParameters& params,
                                              base::span<uint8_t> data);
  static std::unique_ptr<AudioBus> WrapMemory(const AudioParameters& params,
                                              base::span<float> data);

  // Based on the given number of channels and frames, calculates the minimum
  // required size in bytes of a contiguous block of memory to be passed to
  // AudioBus for storage of the audio data.
  // Uses channels() and frames_per_buffer() from AudioParameters if given.
  static size_t CalculateMemorySize(int channels, int frames);
  static size_t CalculateMemorySize(const AudioParameters& params);

  // Checks if buffer is properly aligned to be used in `SetChannelData()`
  static bool IsAligned(void* ptr);
  static bool IsAligned(base::span<float> span);

  ~AudioBus();

  AudioBus(const AudioBus&) = delete;
  AudioBus& operator=(const AudioBus&) = delete;

  // Methods that are expected to be called after AudioBus::CreateWrapper() in
  // order to wrap externally allocated memory.
  // To avoid cases where channel sizes and number of frames don't match,
  // `set_frames()` must be called before setting channel data.
  // Note: It is illegal to call these methods when using a factory method other
  // than CreateWrapper().
  void set_frames(int frames);
  void SetChannelData(int channel, Channel data);
  void SetAllChannels(const ChannelVector& channel_data);

  // Method optionally called after AudioBus::CreateWrapper().
  // Runs |deleter| when on |this|' destruction, freeing external data
  // referenced by SetChannelData().
  // Note: It is illegal to call this method when using a factory method other
  // than CreateWrapper().
  void SetWrappedDataDeleter(base::OnceClosure deleter);

  // Methods for compressed bitstream formats. The data size may not be equal to
  // the capacity of the AudioBus. Also, the frame count may not be equal to the
  // capacity of the AudioBus. Thus, we need extra methods to access the real
  // data size and frame count for bitstream formats.
  bool is_bitstream_format() const { return is_bitstream_format_; }
  void set_is_bitstream_format(bool is_bitstream_format) {
    if (is_bitstream_format) {
      // Don't allow bitstreams if we don't have a continuous chunk of memory.
      // This happens for busses created by CreateWrapper() and WrapVector().
      CHECK(!reserved_memory_.empty());
    }
    is_bitstream_format_ = is_bitstream_format;
  }
  void SetBitstreamSize(size_t data_size);
  int GetBitstreamFrames() const;
  void SetBitstreamFrames(size_t frames);

  // Returns the currently used bitstream data.
  BitstreamData bitstream_data() const { return bitstream_data_; }

  // Overwrites every sample stored in this AudioBus instance with values
  // from a given interleaved `source` with expected layout
  // [ch0, ch1, ..., chN, ch0, ch1, ...]. The sample values are converted to
  // float values by means of the method provided by the SourceSampleTypeTraits.
  // If `zero_remaining_frames` is true, frames not overwritten by the contents
  // of `source` will be zeroed. If it is false, `source` must have the exact
  // size to hold `frames() * channels()` elements.
  template <class SourceSampleTypeTraits>
  void FromInterleaved(
      base::span<const typename SourceSampleTypeTraits::ValueType> source,
      bool zero_remaining_frames = false);

  // Similar to FromInterleaved...(), but overwrites the frames starting at a
  // given offset `write_offset`, without zero'ing other frames.
  template <class SourceSampleTypeTraits>
  void FromInterleavedPartial(
      base::span<const typename SourceSampleTypeTraits::ValueType> source,
      size_t write_offset);

  // Fills `dest` with the sample values in this AudioBus instance. Converts the
  // samples to the format specified by `TargetSampleTypeTraits` and places them
  // in interleaved format.
  // Note: `dest` must have the exact size to hold `frames() * channels()`
  // elements.
  template <class TargetSampleTypeTraits>
  void ToInterleaved(
      base::span<typename TargetSampleTypeTraits::ValueType> dest) const;

  // Similar to ToInterleaved...(), but reads the frames starting at a given
  // `read_offset`.
  // Note: `dest` must have a multiple of `channels()` elements, but it does not
  // need to be big enough to hold all remaining frames past `read_offset`.
  template <class TargetSampleTypeTraits>
  void ToInterleavedPartial(
      size_t read_offset,
      base::span<typename TargetSampleTypeTraits::ValueType> dest) const;

  // Helpers delegating to their respective "byte-less" function, included for
  // convenience. Handles up casting the byte spans safely into spans of the
  // appropriate sample type.
  // Prefer using "byte-less" functions directly.
  template <class SourceSampleTypeTraits>
  void FromInterleavedBytes(base::span<const uint8_t> source,
                            bool zero_remaining_frames = false);
  template <class SourceSampleTypeTraits>
  void FromInterleavedBytesPartial(base::span<const uint8_t> source,
                                   size_t write_offset);
  template <class TargetSampleTypeTraits>
  void ToInterleavedBytes(base::span<uint8_t> dest) const;
  template <class TargetSampleTypeTraits>
  void ToInterleavedBytesPartial(size_t read_offset,
                                 base::span<uint8_t> dest) const;

  // Helper method for copying channel data from one AudioBus to another.  Both
  // AudioBus object must have the same frames() and channels().
  void CopyTo(AudioBus* dest) const;

  // Similar to above, but clips values to [-1, 1] during the copy process.
  void CopyAndClipTo(AudioBus* dest) const;

  // Helper method to copy frames from one AudioBus to another. Both AudioBus
  // objects must have the same number of channels(). |source_start_frame| is
  // the starting offset. |dest_start_frame| is the starting offset in |dest|.
  // |frame_count| is the number of frames to copy.
  void CopyPartialFramesTo(int source_start_frame,
                           int frame_count,
                           int dest_start_frame,
                           AudioBus* dest) const;

  // Returns a raw pointer to the requested channel.  Pointer is guaranteed to
  // have a 16-byte alignment.  Warning: Do not rely on having sane (i.e. not
  // inf, nan, or between [-1.0, 1.0]) values in the channel data.
  Channel channel(int channel) {
    CHECK(!is_bitstream_format_);
    return channel_data_[channel];
  }
  ConstChannel channel(int channel) const {
    CHECK(!is_bitstream_format_);
    return channel_data_[channel];
  }

  // Convenience function to allow range-based for-loops.
  const ChannelVector& AllChannels() const;

  // Returns a copy of `channels_`, with `subspan()` applied to each channel.
  // Note: The returned channels might not be aligned, depending on `offset`.
  ChannelVector AllChannelsSubspan(size_t offset, size_t count) const;

  // Returns the number of channels.
  int channels() const { return static_cast<int>(channel_data_.size()); }
  // Returns the number of frames.
  // Note: for bitstream formats, use GetBitstreamFrames() to get the actual
  // number of encoded frames. However, `frames()` remains useful in determining
  // the amount of `reserved_memory_` this bus has.
  int frames() const { return base::checked_cast<int>(frames_); }

  // Helper method for zeroing out all channels of audio data.
  void Zero();
  void ZeroFrames(int frames);
  void ZeroFramesPartial(int start_frame, int frames);

  // Checks if all frames are zero.
  bool AreFramesZero() const;

  // Scale internal channel values by |volume| >= 0.  If an invalid value
  // is provided, no adjustment is done.
  void Scale(float volume);

  // Swaps channels identified by |a| and |b|.  The caller needs to make sure
  // the channels are valid.
  void SwapChannels(int a, int b);

 private:
  AudioBus(int channels, int frames);
  AudioBus(int channels, int frames, base::span<float> data);
  explicit AudioBus(int channels);

  void ZeroBitstream();

  // Helper method for building |channel_data_| from a block of memory.  |data|
  // must be at least CalculateMemorySize(...) bytes in size.
  void BuildChannelData(int channels, base::span<float> data);

  template <class SourceSampleTypeTraits>
  static void CopyConvertFromInterleavedSourceToAudioBus(
      base::span<const typename SourceSampleTypeTraits::ValueType> source,
      size_t write_offset,
      AudioBus* dest);

  template <class TargetSampleTypeTraits>
  static void CopyConvertFromAudioBusToInterleavedTarget(
      const AudioBus* source,
      size_t read_offset,
      base::span<typename TargetSampleTypeTraits::ValueType> dest);

  template <typename T>
  static size_t get_frame_count(base::span<T> data, size_t channels) {
    CHECK_EQ(data.size() % channels, 0u);
    return data.size() / channels;
  }

  // Contiguous block of channel memory.
  base::AlignedHeapArray<float> data_;

  // Chunk of binary data for bitstream formats.
  // This might point towards external memory, or `data_`.
  base::raw_span<uint8_t> reserved_memory_;

  // View over `reserved_memory_`, which represents the chunk of memory which
  // is actively reserved to hold bitstream data. The size of this memory can
  // be adjusted using SetBitstreamDataSize().
  base::raw_span<uint8_t> bitstream_data_;

  // Whether the data is compressed bitstream or not.
  bool is_bitstream_format_ = false;
  // The PCM frame count for a compressed bitstream.
  size_t bitstream_frames_ = 0;

  // One float pointer per channel pointing to a contiguous block of memory for
  // that channel. If the memory is owned by this instance, this will
  // point to the memory in |data_|. Otherwise, it may point to memory provided
  // by the client.
  // TODO(crbug.com/385028986): Convert to `base::raw_span`
  RAW_PTR_EXCLUSION ChannelVector channel_data_;

  size_t frames_ = 0u;

  // Protect SetChannelData(), set_frames() and SetWrappedDataDeleter() for use
  // by CreateWrapper().
  const bool is_wrapper_ = false;

  // Run on destruction. Frees memory to the data set via SetChannelData().
  // Only used with CreateWrapper().
  base::OnceClosure wrapped_data_deleter_cb_;
};

// Delegates to FromInterleavedPartial()
template <class SourceSampleTypeTraits>
void AudioBus::FromInterleaved(
    base::span<const typename SourceSampleTypeTraits::ValueType> source,
    bool zero_remaining_frames) {
  const size_t source_frame_count = get_frame_count(source, channels());
  CHECK_LE(source_frame_count, frames_);

  FromInterleavedPartial<SourceSampleTypeTraits>(source, 0u);

  const size_t remaining_frames = frames_ - source_frame_count;
  if (!remaining_frames) {
    return;
  }

  // If not using `zero_remaining_frames`, `source` should have the exact size.
  CHECK(zero_remaining_frames);
  ZeroFramesPartial(base::checked_cast<int>(source_frame_count),
                    base::checked_cast<int>(remaining_frames));
}

template <class SourceSampleTypeTraits>
void AudioBus::FromInterleavedPartial(
    base::span<const typename SourceSampleTypeTraits::ValueType> source,
    size_t write_offset) {
  const size_t frame_count = get_frame_count(source, channels());
  const size_t total_offset =
      base::CheckAdd(frame_count, write_offset).ValueOrDie();
  CHECK_LE(total_offset, frames_);

  CopyConvertFromInterleavedSourceToAudioBus<SourceSampleTypeTraits>(
      source, write_offset, this);
}

// Delegates to ToInterleavedPartial()
template <class TargetSampleTypeTraits>
void AudioBus::ToInterleaved(
    base::span<typename TargetSampleTypeTraits::ValueType> dest) const {
  const size_t frames_count = get_frame_count(dest, channels());
  CHECK_EQ(frames_count, frames_);
  ToInterleavedPartial<TargetSampleTypeTraits>(0u, dest);
}

template <class TargetSampleTypeTraits>
void AudioBus::ToInterleavedPartial(
    size_t read_offset,
    base::span<typename TargetSampleTypeTraits::ValueType> dest) const {
  const size_t frame_count = get_frame_count(dest, channels());
  const size_t total_offset =
      base::CheckAdd(frame_count, read_offset).ValueOrDie();
  CHECK_LE(total_offset, frames_);
  CopyConvertFromAudioBusToInterleavedTarget<TargetSampleTypeTraits>(
      this, read_offset, dest);
}

template <class SourceSampleTypeTraits>
void AudioBus::CopyConvertFromInterleavedSourceToAudioBus(
    base::span<const typename SourceSampleTypeTraits::ValueType> source,
    size_t write_offset,
    AudioBus* dest) {
  const size_t channels = dest->channels();
  const size_t frame_count = get_frame_count(source, channels);
  const size_t total_offset =
      base::CheckAdd(frame_count, write_offset).ValueOrDie();
  CHECK_LE(total_offset, static_cast<size_t>(dest->frames()));

  for (size_t ch = 0; ch < channels; ++ch) {
    auto channel_data =
        dest->channel(static_cast<int>(ch)).subspan(write_offset, frame_count);
    for (size_t dest_idx = 0, src_idx = ch; dest_idx < frame_count;
         ++dest_idx, src_idx += channels) {
      auto source_sample = source[src_idx];
      channel_data[dest_idx] = SourceSampleTypeTraits::ToFloat(source_sample);
    }
  }
}

template <class TargetSampleTypeTraits>
void AudioBus::CopyConvertFromAudioBusToInterleavedTarget(
    const AudioBus* source,
    size_t read_offset,
    base::span<typename TargetSampleTypeTraits::ValueType> dest) {
  const size_t channels = source->channels();
  const size_t frame_count = get_frame_count(dest, channels);
  const size_t total_offset =
      base::CheckAdd(frame_count, read_offset).ValueOrDie();
  CHECK_LE(total_offset, static_cast<size_t>(source->frames()));

  for (size_t ch = 0; ch < channels; ++ch) {
    auto channel_data =
        source->channel(static_cast<int>(ch)).subspan(read_offset, frame_count);
    for (size_t src_idx = 0, dest_idx = ch; src_idx < frame_count;
         ++src_idx, dest_idx += channels) {
      float source_sample = channel_data[src_idx];
      dest[dest_idx] = TargetSampleTypeTraits::FromFloat(source_sample);
    }
  }
}

template <class SourceSampleTypeTraits>
void AudioBus::FromInterleavedBytes(base::span<const uint8_t> source,
                                    bool zero_remaining_frames) {
  FromInterleaved<SourceSampleTypeTraits>(
      base::subtle::reinterpret_span<
          const typename SourceSampleTypeTraits::ValueType>(source),
      zero_remaining_frames);
}
template <class SourceSampleTypeTraits>
void AudioBus::FromInterleavedBytesPartial(base::span<const uint8_t> source,
                                           size_t write_offset) {
  FromInterleavedPartial<SourceSampleTypeTraits>(
      base::subtle::reinterpret_span<
          const typename SourceSampleTypeTraits::ValueType>(source),
      write_offset);
}
template <class TargetSampleTypeTraits>
void AudioBus::ToInterleavedBytes(base::span<uint8_t> dest) const {
  ToInterleaved<TargetSampleTypeTraits>(
      base::subtle::reinterpret_span<
          typename TargetSampleTypeTraits::ValueType>(dest));
}
template <class TargetSampleTypeTraits>
void AudioBus::ToInterleavedBytesPartial(size_t read_offset,
                                         base::span<uint8_t> dest) const {
  ToInterleavedPartial<TargetSampleTypeTraits>(
      read_offset, base::subtle::reinterpret_span<
                       typename TargetSampleTypeTraits::ValueType>(dest));
}

}  // namespace media

#endif  // MEDIA_BASE_AUDIO_BUS_H_