write_loop_wav.hpp

#pragma once
#include <dr_wav.h>
 
#include <algorithm>
#include <cstdint>
#include <filesystem>
#include <fstream>
#include <string_view>
#include <vector>
 
namespace ossia
{
struct wav_loop_info
{
  double tempo{};
  int beats{};
  int meter_numerator{4};
  int meter_denominator{4};
  bool loop{true};
};

inline bool write_loop_wav(
    std::string_view path_utf8, int channels, int rate, int64_t frames,
    const float* const* data, const wav_loop_info& info)
{
  if(channels <= 0 || rate <= 0 || frames < 0 || path_utf8.empty())
    return false;
 
  const std::filesystem::path path{std::u8string_view{
      reinterpret_cast<const char8_t*>(path_utf8.data()), path_utf8.size()}};
  std::ofstream file{path, std::ios::binary | std::ios::trunc};
  if(!file)
    return false;
 
  auto fail = [&] {
    file.close();
    std::error_code ec;
    std::filesystem::remove(path, ec);
    return false;
  };
 
  // RIFF sizes are 32 bits; the margin covers the header and metadata chunks.
  constexpr uint64_t riff_max_data = 0xFFFFFFFFull - 65536;
  const uint64_t data_bytes = uint64_t(frames) * uint64_t(channels) * sizeof(float);
 
  drwav_data_format format{};
  format.container
      = data_bytes > riff_max_data ? drwav_container_rf64 : drwav_container_riff;
  format.format = DR_WAVE_FORMAT_IEEE_FLOAT;
  format.channels = channels;
  format.sampleRate = rate;
  format.bitsPerSample = 32;
 
  drwav_metadata metadata[2]{};
  drwav_uint32 metadata_count = 0;
 
  // The sampler chunk counts in 32 bits even in RF64: a loop it cannot end
  // where the file ends is left out rather than written wrong.
  drwav_smpl_loop sampler_loop{};
  if(info.loop && frames > 0 && uint64_t(frames - 1) <= 0xFFFFFFFFull)
  {
    sampler_loop.type = drwav_smpl_loop_type_forward;
    sampler_loop.firstSampleOffset = 0;
    sampler_loop.lastSampleOffset = drwav_uint32(frames - 1);
 
    auto& smpl = metadata[metadata_count++];
    smpl.type = drwav_metadata_type_smpl;
    smpl.data.smpl.samplePeriodNanoseconds = drwav_uint32(1e9 / rate);
    smpl.data.smpl.midiUnityNote = 60;
    smpl.data.smpl.sampleLoopCount = 1;
    smpl.data.smpl.pLoops = &sampler_loop;
  }
 
  if(info.tempo > 0.)
  {
    auto& acid = metadata[metadata_count++];
    acid.type = drwav_metadata_type_acid;
    acid.data.acid.flags = info.loop ? drwav_acid_flag_stretch : drwav_acid_flag_one_shot;
    acid.data.acid.midiUnityNote = 60;
    acid.data.acid.reserved1 = 0x80;
    acid.data.acid.numBeats = drwav_uint32(info.beats > 0 ? info.beats : 0);
    acid.data.acid.meterNumerator = drwav_uint16(info.meter_numerator);
    acid.data.acid.meterDenominator = drwav_uint16(info.meter_denominator);
    acid.data.acid.tempo = float(info.tempo);
  }
 
  auto on_write = [](void* user, const void* src, size_t bytes) -> size_t {
    auto& f = *static_cast<std::ofstream*>(user);
    f.write(static_cast<const char*>(src), std::streamsize(bytes));
    return f ? bytes : 0;
  };
  auto on_seek = [](void* user, int offset, drwav_seek_origin origin) -> drwav_bool32 {
    auto& f = *static_cast<std::ofstream*>(user);
    switch(origin)
    {
      case DRWAV_SEEK_SET:
        f.seekp(offset, std::ios::beg);
        break;
      case DRWAV_SEEK_CUR:
        f.seekp(offset, std::ios::cur);
        break;
      default:
        f.seekp(offset, std::ios::end);
        break;
    }
    return bool(f);
  };
 
  drwav wav{};
  if(!drwav_init_write_with_metadata(
         &wav, &format, on_write, on_seek, &file, nullptr,
         metadata_count > 0 ? metadata : nullptr, metadata_count))
    return fail();
 
  constexpr int64_t chunk = 4096;
  std::vector<float> interleaved(std::size_t(chunk * channels));
  int64_t written = 0;
  while(written < frames)
  {
    const int64_t n = std::min(chunk, frames - written);
    for(int64_t i = 0; i < n; i++)
      for(int c = 0; c < channels; c++)
        interleaved[std::size_t(i * channels + c)] = data[c][written + i];
    if(drwav_write_pcm_frames(&wav, drwav_uint64(n), interleaved.data()) != drwav_uint64(n))
      break;
    written += n;
  }
 
  const bool finalized = drwav_uninit(&wav) == DRWAV_SUCCESS;
  file.flush();
  if(written != frames || !finalized || !file)
    return fail();
  file.close();
  if(!file)
    return fail();
  return true;
}
}