libav-ruby 0.1.0

data/lib/libav/audio_encoder.rb

@@ -0,0 +1,326 @@
+# frozen_string_literal: true
+
+module LibAV
+  class AudioEncoder
+    CODECS = {
+      aac: %w[aac],
+      opus: %w[libopus opus],
+      pcm_f32le: %w[pcm_f32le]
+    }.freeze
+    CODEC_IDS = {
+      pcm_f32le: 65_557,
+      aac: 86_018,
+      opus: 86_076
+    }.freeze
+    SAMPLE_FORMATS = {
+      aac: "fltp",
+      opus: "flt",
+      pcm_f32le: "flt"
+    }.freeze
+    DEFAULT_FRAME_SIZE = 1_024
+
+    attr_reader :stream, :time_base, :sample_rate, :channels
+
+    def initialize(output, codec: :aac, sample_rate: 48_000, channels: 2,
+                   bitrate: 192_000, **)
+      @output = output
+      @codec_name = codec.to_sym
+      @sample_rate = positive_integer(sample_rate, :sample_rate)
+      @channels = positive_integer(channels, :channels)
+      raise ArgumentError, "channels must be between 1 and 64" if @channels > 64
+
+      @time_base = Rational(1, @sample_rate)
+      @codec = find_encoder
+      @stream = output.new_stream(@codec)
+      @context = Native.codec.avcodec_alloc_context3(@codec)
+      raise EncodingError, "avcodec_alloc_context3 returned NULL" if @context.null?
+
+      configure_context(bitrate)
+      open_context
+      allocate_conversion
+      @sample_index = 0
+    rescue StandardError
+      close
+      raise
+    end
+
+    def push(pcm)
+      raise ClosedError, "audio encoder is closed" if @closed
+
+      bytes = normalize_pcm(pcm)
+      convert_to_fifo(bytes, bytes.bytesize / input_sample_bytes)
+      encode_complete_frames
+      self
+    end
+
+    def flush
+      return if @flushed || @closed
+
+      flush_resampler
+      encode_complete_frames
+      encode_fifo_tail
+      send_frame(nil)
+      @flushed = true
+    end
+
+    def close
+      return if @closed
+
+      Native.util.av_audio_fifo_free(@fifo) if @fifo && !@fifo.null?
+      free_swr
+      Native.util.av_channel_layout_uninit(@input_layout) if @input_layout
+      free_pointer(@frame, :av_frame_free, Native.util)
+      free_pointer(@packet, :av_packet_free, Native.codec)
+      free_pointer(@context, :avcodec_free_context, Native.codec)
+      @fifo = @frame = @packet = @context = @swr = nil
+      @closed = true
+    end
+
+    private
+
+    def configure_context(bitrate)
+      sample_format = Native.util.av_get_sample_fmt(SAMPLE_FORMATS.fetch(@codec_name))
+      Native::Layout.write_int(@context, :codec_context, :sample_rate, @sample_rate)
+      Native::Layout.write_int(@context, :codec_context, :sample_fmt, sample_format)
+      Native::Layout.write_int64(@context, :codec_context, :bit_rate, Integer(bitrate)) if bitrate
+      Native::Layout.write_rational(@context, :codec_context, :time_base, time_base)
+      Native.util.av_channel_layout_default(
+        Native::Layout.field_pointer(@context, :codec_context, :ch_layout), @channels
+      )
+      Native::Layout.write_rational(stream, :stream, :time_base, time_base)
+
+      return unless @output.global_header?
+
+      flags = Native::Layout.read_int(@context, :codec_context, :flags)
+      Native::Layout.write_int(
+        @context, :codec_context, :flags, flags | Native::AV_CODEC_FLAG_GLOBAL_HEADER
+      )
+    end
+
+    def open_context
+      Native.check(Native.codec.avcodec_open2(@context, @codec, nil), "avcodec_open2(audio)")
+      parameters = Native::Layout.read_pointer(stream, :stream, :codecpar)
+      Native.check(
+        Native.codec.avcodec_parameters_from_context(parameters, @context),
+        "avcodec_parameters_from_context(audio)"
+      )
+      reported_size = Native::Layout.read_int(@context, :codec_context, :frame_size)
+      @variable_frames = reported_size.zero?
+      @frame_capacity = @variable_frames ? DEFAULT_FRAME_SIZE : reported_size
+    end
+
+    def allocate_conversion
+      @frame = Native.util.av_frame_alloc
+      @packet = Native.codec.av_packet_alloc
+      raise EncodingError, "failed to allocate an audio AVFrame or AVPacket" if
+        @frame.null? || @packet.null?
+
+      @sample_format = Native::Layout.read_int(@context, :codec_context, :sample_fmt)
+      Native::Layout.write_int(@frame, :frame, :format, @sample_format)
+      Native::Layout.write_int(@frame, :frame, :sample_rate, @sample_rate)
+      Native::Layout.write_int(@frame, :frame, :nb_samples, @frame_capacity)
+      Native.util.av_channel_layout_default(
+        Native::Layout.field_pointer(@frame, :frame, :ch_layout), @channels
+      )
+      Native.check(Native.util.av_frame_get_buffer(@frame, 0), "av_frame_get_buffer(audio)")
+
+      @input_layout = FFI::MemoryPointer.new(:char, Native::AV_CHANNEL_LAYOUT_SIZE, true)
+      Native.util.av_channel_layout_default(@input_layout, @channels)
+      swr_holder = FFI::MemoryPointer.new(:pointer)
+      result = Native.resample.swr_alloc_set_opts2(
+        swr_holder,
+        Native::Layout.field_pointer(@context, :codec_context, :ch_layout),
+        @sample_format, @sample_rate,
+        @input_layout, Native::AV_SAMPLE_FMT_FLT, @sample_rate,
+        0, nil
+      )
+      Native.check(result, "swr_alloc_set_opts2")
+      @swr = swr_holder.read_pointer
+      raise EncodingError, "swr_alloc_set_opts2 returned NULL" if @swr.null?
+
+      Native.check(Native.resample.swr_init(@swr), "swr_init")
+      @fifo = Native.util.av_audio_fifo_alloc(
+        @sample_format, @channels, @frame_capacity
+      )
+      raise EncodingError, "av_audio_fifo_alloc returned NULL" if @fifo.null?
+    end
+
+    def encode_complete_frames
+      while Native.util.av_audio_fifo_size(@fifo) >= @frame_capacity
+        encode_fifo_samples(@frame_capacity)
+      end
+    end
+
+    def encode_fifo_tail
+      samples = Native.util.av_audio_fifo_size(@fifo)
+      return if samples.zero?
+
+      encode_fifo_samples(samples, pad_to: @variable_frames ? samples : @frame_capacity)
+    end
+
+    def encode_fifo_samples(samples, pad_to: samples)
+      Native.check(Native.util.av_frame_make_writable(@frame), "av_frame_make_writable(audio)")
+      frame_data = Native::Layout.read_pointer(@frame, :frame, :extended_data)
+      read = Native.util.av_audio_fifo_read(@fifo, frame_data, samples)
+      Native.check(read, "av_audio_fifo_read")
+      raise EncodingError, "AVAudioFifo returned #{read} samples, expected #{samples}" unless
+        read == samples
+
+      if pad_to > samples
+        Native.check(
+          Native.util.av_samples_set_silence(
+            frame_data, samples, pad_to - samples, @channels, @sample_format
+          ),
+          "av_samples_set_silence"
+        )
+      end
+
+      Native::Layout.write_int(@frame, :frame, :nb_samples, pad_to)
+      Native::Layout.write_int64(@frame, :frame, :pts, @sample_index)
+      Native::Layout.write_int64(@frame, :frame, :duration, pad_to)
+      @sample_index += pad_to
+      send_frame(@frame)
+    end
+
+    def convert_to_fifo(bytes, samples)
+      source = FFI::MemoryPointer.from_string(bytes)
+      input_data = FFI::MemoryPointer.new(:pointer)
+      input_data.write_pointer(source)
+      capacity = Native.resample.swr_get_out_samples(@swr, samples)
+      Native.check(capacity, "swr_get_out_samples")
+      return if capacity.zero?
+
+      with_sample_buffer(capacity) do |output_data|
+        converted = Native.resample.swr_convert(
+          @swr, output_data, capacity, input_data, samples
+        )
+        Native.check(converted, "swr_convert")
+        write_fifo(output_data, converted)
+      end
+    end
+
+    def flush_resampler
+      loop do
+        capacity = Native.resample.swr_get_out_samples(@swr, 0)
+        Native.check(capacity, "swr_get_out_samples(flush)")
+        break if capacity.zero?
+
+        converted = with_sample_buffer(capacity) do |output_data|
+          result = Native.resample.swr_convert(@swr, output_data, capacity, nil, 0)
+          Native.check(result, "swr_convert(flush)")
+          write_fifo(output_data, result)
+          result
+        end
+        break if converted.zero?
+      end
+    end
+
+    def write_fifo(output_data, samples)
+      return if samples.zero?
+
+      required = Native.util.av_audio_fifo_size(@fifo) + samples
+      Native.check(Native.util.av_audio_fifo_realloc(@fifo, required), "av_audio_fifo_realloc")
+      written = Native.util.av_audio_fifo_write(@fifo, output_data, samples)
+      Native.check(written, "av_audio_fifo_write")
+      raise EncodingError, "AVAudioFifo wrote #{written} samples, expected #{samples}" unless
+        written == samples
+    end
+
+    def with_sample_buffer(samples)
+      holder = FFI::MemoryPointer.new(:pointer)
+      linesize = FFI::MemoryPointer.new(:int)
+      Native.check(
+        Native.util.av_samples_alloc_array_and_samples(
+          holder, linesize, @channels, samples, @sample_format, 0
+        ),
+        "av_samples_alloc_array_and_samples"
+      )
+      data = holder.read_pointer
+      yield data
+    ensure
+      Native.util.av_freep(data) if data && !data.null?
+      Native.util.av_freep(holder) if holder && !holder.read_pointer.null?
+    end
+
+    def send_frame(frame)
+      result = Native.codec.avcodec_send_frame(@context, frame)
+      drain_packets if result == Native::AVERROR_EAGAIN
+      result = Native.codec.avcodec_send_frame(@context, frame) if result == Native::AVERROR_EAGAIN
+      Native.check(result, "avcodec_send_frame(audio)")
+      drain_packets
+    end
+
+    def drain_packets
+      loop do
+        result = Native.codec.avcodec_receive_packet(@context, @packet)
+        break if Native.control_flow?(result)
+
+        Native.check(result, "avcodec_receive_packet(audio)")
+        begin
+          duration = Native::Layout.read_int64(@packet, :packet, :duration)
+          Native::Layout.write_int64(@packet, :packet, :duration, @frame_capacity) unless
+            duration.positive?
+          @output.write_packet(@packet, time_base, stream)
+        ensure
+          Native.codec.av_packet_unref(@packet)
+        end
+      end
+    end
+
+    def normalize_pcm(pcm)
+      bytes = if pcm.is_a?(Array)
+                pcm.pack("e*")
+              else
+                pcm.to_str.b
+              end
+      alignment = @channels * FFI.type_size(:float)
+      raise ArgumentError, "PCM data must contain complete #{@channels}-channel f32 samples" unless
+        (bytes.bytesize % alignment).zero?
+
+      bytes.dup
+    rescue NoMethodError, TypeError
+      raise ArgumentError, "PCM must be packed f32 data or an Array of samples"
+    end
+
+    def input_sample_bytes
+      @channels * FFI.type_size(:float)
+    end
+
+    def find_encoder
+      names = CODECS.fetch(@codec_name) do
+        raise ArgumentError, "unsupported audio codec: #{@codec_name.inspect}"
+      end
+      names.each do |name|
+        encoder = Native.codec.avcodec_find_encoder_by_name(name)
+        return encoder unless encoder.null?
+      end
+      encoder = Native.codec.avcodec_find_encoder(CODEC_IDS.fetch(@codec_name))
+      return encoder unless encoder.null?
+
+      raise EncodingError, "FFmpeg does not provide a #{names.join(" or ")} encoder"
+    end
+
+    def positive_integer(value, name)
+      value = Integer(value)
+      return value if value.positive?
+
+      raise ArgumentError, "#{name} must be positive"
+    end
+
+    def free_swr
+      return unless @swr && !@swr.null?
+
+      holder = FFI::MemoryPointer.new(:pointer)
+      holder.write_pointer(@swr)
+      Native.resample.swr_free(holder)
+    end
+
+    def free_pointer(pointer, function, library)
+      return unless pointer && !pointer.null?
+
+      holder = FFI::MemoryPointer.new(:pointer)
+      holder.write_pointer(pointer)
+      library.public_send(function, holder)
+    end
+  end
+end

data/lib/libav/audio_frame.rb

@@ -0,0 +1,25 @@
+# frozen_string_literal: true
+
+module LibAV
+  class AudioFrame
+    attr_reader :sample_rate, :channels, :samples, :data, :pts
+
+    def initialize(sample_rate:, channels:, samples:, data:, pts: nil)
+      @sample_rate = Integer(sample_rate)
+      @channels = Integer(channels)
+      @samples = Integer(samples)
+      @data = data.to_str
+      @data = @data.dup unless @data.encoding == Encoding::BINARY
+      @data.force_encoding(Encoding::BINARY)
+      @pts = pts
+    end
+
+    def to_str
+      data
+    end
+
+    def duration
+      Rational(samples, sample_rate)
+    end
+  end
+end

data/lib/libav/audio_reader.rb

@@ -0,0 +1,336 @@
+# frozen_string_literal: true
+
+module LibAV
+  class AudioReader
+    attr_reader :sample_rate, :channels, :duration
+
+    def self.open(path)
+      reader = new(path)
+      return reader unless block_given?
+
+      begin
+        yield reader
+      ensure
+        reader.close
+      end
+    end
+
+    def initialize(path)
+      @path = File.path(path)
+      @format_holder = FFI::MemoryPointer.new(:pointer)
+      Native.check(
+        Native.format.avformat_open_input(@format_holder, @path, nil, nil),
+        "avformat_open_input"
+      )
+      @format_context = @format_holder.read_pointer
+      Native.check(
+        Native.format.avformat_find_stream_info(@format_context, nil),
+        "avformat_find_stream_info"
+      )
+      open_audio_stream
+      allocate_decode_state
+      load_metadata
+      reset_decode_state
+    rescue StandardError
+      close
+      raise
+    end
+
+    def read_frame
+      raise ClosedError, "audio reader is closed" if closed?
+
+      loop do
+        return unless receive_next_frame
+
+        frame = convert_frame
+        next if before_seek_target?(frame)
+
+        frame = trim_to_seek_target(frame)
+        next if frame.samples.zero?
+
+        @seek_target = nil
+        return frame
+      end
+    end
+
+    def read_samples
+      read_frame&.data
+    end
+
+    def each_frame
+      return enum_for(__method__) unless block_given?
+
+      while (frame = read_frame)
+        yield frame, frame.pts
+      end
+      self
+    end
+
+    def each_chunk
+      return enum_for(__method__) unless block_given?
+
+      each_frame { |frame, pts| yield frame.data, pts }
+      self
+    end
+
+    def seek(seconds)
+      raise ClosedError, "audio reader is closed" if closed?
+
+      seconds = Float(seconds)
+      raise ArgumentError, "seek time must be nonnegative" if seconds.negative? || !seconds.finite?
+
+      timestamp = (seconds / @time_base).round
+      Native.check(
+        Native.format.av_seek_frame(
+          @format_context, @stream_index, timestamp, Native::AVSEEK_FLAG_BACKWARD
+        ),
+        "av_seek_frame"
+      )
+      Native.codec.avcodec_flush_buffers(@codec_context)
+      Native.util.av_frame_unref(@frame)
+      Native.codec.av_packet_unref(@packet)
+      free_swr
+      @conversion_signature = nil
+      reset_decode_state
+      @seek_target = seconds
+      self
+    end
+
+    def close
+      return if closed?
+
+      free_swr
+      Native.util.av_channel_layout_uninit(@output_layout) if @output_layout
+      free_pointer(@frame, :av_frame_free, Native.util)
+      free_pointer(@packet, :av_packet_free, Native.codec)
+      free_pointer(@codec_context, :avcodec_free_context, Native.codec)
+      if @format_context && !@format_context.null?
+        @format_holder.write_pointer(@format_context)
+        Native.format.avformat_close_input(@format_holder)
+      end
+      @format_context = @codec_context = @frame = @packet = @swr = nil
+      @closed = true
+      nil
+    end
+
+    def closed?
+      @closed == true
+    end
+
+    private
+
+    def open_audio_stream
+      decoder_holder = FFI::MemoryPointer.new(:pointer)
+      @stream_index = Native.format.av_find_best_stream(
+        @format_context, Native::AVMEDIA_TYPE_AUDIO, -1, -1, decoder_holder, 0
+      )
+      Native.check(@stream_index, "av_find_best_stream(audio)")
+      decoder = decoder_holder.read_pointer
+      raise Error, "FFmpeg did not return an audio decoder" if decoder.null?
+
+      @stream = stream_at(@stream_index)
+      @codec_context = Native.codec.avcodec_alloc_context3(decoder)
+      raise Error, "avcodec_alloc_context3 returned NULL" if @codec_context.null?
+
+      parameters = Native::Layout.read_pointer(@stream, :stream, :codecpar)
+      Native.check(
+        Native.codec.avcodec_parameters_to_context(@codec_context, parameters),
+        "avcodec_parameters_to_context(audio)"
+      )
+      Native.check(
+        Native.codec.avcodec_open2(@codec_context, decoder, nil),
+        "avcodec_open2(audio)"
+      )
+    end
+
+    def allocate_decode_state
+      @frame = Native.util.av_frame_alloc
+      @packet = Native.codec.av_packet_alloc
+      raise Error, "failed to allocate an audio AVFrame or AVPacket" if
+        @frame.null? || @packet.null?
+
+      @destination_data = FFI::MemoryPointer.new(:pointer)
+    end
+
+    def load_metadata
+      @sample_rate = Native::Layout.read_int(@codec_context, :codec_context, :sample_rate)
+      context_layout = Native::Layout.field_pointer(@codec_context, :codec_context, :ch_layout)
+      @channels = context_layout.get_int32(4)
+      raise Error, "invalid audio metadata: #{@sample_rate} Hz, #{@channels} channels" unless
+        @sample_rate.positive? && @channels.positive?
+
+      @time_base = Native::Layout.read_rational(@stream, :stream, :time_base)
+      stream_duration = Native::Layout.read_int64(@stream, :stream, :duration)
+      @duration = if valid_timestamp?(stream_duration)
+                    stream_duration * @time_base
+                  else
+                    format_duration
+                  end
+      @output_layout = FFI::MemoryPointer.new(:char, Native::AV_CHANNEL_LAYOUT_SIZE, true)
+      Native.util.av_channel_layout_default(@output_layout, @channels)
+    end
+
+    def receive_next_frame
+      loop do
+        result = Native.codec.avcodec_receive_frame(@codec_context, @frame)
+        return true if result.zero?
+        return false if result == Native::AVERROR_EOF
+
+        Native.check(result, "avcodec_receive_frame(audio)") unless
+          result == Native::AVERROR_EAGAIN
+        return false unless feed_decoder
+      end
+    end
+
+    def feed_decoder
+      if @input_eof
+        return true if send_flush_packet
+
+        return false
+      end
+
+      loop do
+        result = Native.format.av_read_frame(@format_context, @packet)
+        if result == Native::AVERROR_EOF
+          @input_eof = true
+          return send_flush_packet
+        end
+        Native.check(result, "av_read_frame")
+
+        packet_stream = Native::Layout.read_int(@packet, :packet, :stream_index)
+        if packet_stream == @stream_index
+          begin
+            Native.check(
+              Native.codec.avcodec_send_packet(@codec_context, @packet),
+              "avcodec_send_packet(audio)"
+            )
+          ensure
+            Native.codec.av_packet_unref(@packet)
+          end
+          return true
+        end
+        Native.codec.av_packet_unref(@packet)
+      end
+    end
+
+    def send_flush_packet
+      return false if @flush_sent
+
+      result = Native.codec.avcodec_send_packet(@codec_context, nil)
+      Native.check(result, "avcodec_send_packet(audio flush)") unless result == Native::AVERROR_EOF
+      @flush_sent = true
+      true
+    end
+
+    def convert_frame
+      samples = Native::Layout.read_int(@frame, :frame, :nb_samples)
+      frame_format = Native::Layout.read_int(@frame, :frame, :format)
+      input_rate = Native::Layout.read_int(@frame, :frame, :sample_rate)
+      input_rate = @sample_rate unless input_rate.positive?
+      input_layout = Native::Layout.field_pointer(@frame, :frame, :ch_layout)
+      input_layout = Native::Layout.field_pointer(
+        @codec_context, :codec_context, :ch_layout
+      ) unless input_layout.get_int32(4).positive?
+      prepare_conversion(frame_format, input_rate, input_layout)
+
+      capacity = samples + Native.resample.swr_get_delay(@swr, input_rate)
+      capacity = samples if capacity < samples
+      destination = FFI::MemoryPointer.new(:float, capacity * @channels)
+      @destination_data.put_pointer(0, destination)
+      converted = Native.resample.swr_convert(
+        @swr, @destination_data, capacity,
+        Native::Layout.read_pointer(@frame, :frame, :extended_data), samples
+      )
+      Native.check(converted, "swr_convert(audio decode)")
+
+      timestamp = Native::Layout.read_int64(@frame, :frame, :best_effort_timestamp)
+      timestamp = Native::Layout.read_int64(@frame, :frame, :pts) unless valid_timestamp?(timestamp)
+      pts = valid_timestamp?(timestamp) ? timestamp * @time_base : nil
+      byte_size = converted * @channels * FFI.type_size(:float)
+      AudioFrame.new(
+        sample_rate: @sample_rate, channels: @channels, samples: converted,
+        data: destination.read_string_length(byte_size), pts:
+      )
+    end
+
+    def prepare_conversion(frame_format, input_rate, input_layout)
+      signature = [frame_format, input_rate, input_layout.get_int32(4)]
+      return if signature == @conversion_signature
+
+      free_swr
+      holder = FFI::MemoryPointer.new(:pointer)
+      Native.check(
+        Native.resample.swr_alloc_set_opts2(
+          holder,
+          @output_layout, Native::AV_SAMPLE_FMT_FLT, @sample_rate,
+          input_layout, frame_format, input_rate,
+          0, nil
+        ),
+        "swr_alloc_set_opts2(audio decode)"
+      )
+      @swr = holder.read_pointer
+      raise Error, "swr_alloc_set_opts2 returned NULL" if @swr.null?
+
+      Native.check(Native.resample.swr_init(@swr), "swr_init(audio decode)")
+      @conversion_signature = signature
+    end
+
+    def before_seek_target?(frame)
+      return false unless @seek_target && frame.pts
+
+      frame.pts + frame.duration <= @seek_target
+    end
+
+    def trim_to_seek_target(frame)
+      return frame unless @seek_target && frame.pts && frame.pts < @seek_target
+
+      skipped_samples = ((@seek_target - frame.pts) * sample_rate).ceil
+      skipped_samples = [skipped_samples, frame.samples].min
+      remaining_samples = frame.samples - skipped_samples
+      byte_offset = skipped_samples * channels * FFI.type_size(:float)
+      AudioFrame.new(
+        sample_rate:, channels:, samples: remaining_samples,
+        data: frame.data.byteslice(byte_offset..),
+        pts: frame.pts + Rational(skipped_samples, sample_rate)
+      )
+    end
+
+    def reset_decode_state
+      @input_eof = false
+      @flush_sent = false
+    end
+
+    def stream_at(index)
+      streams = Native::Layout.read_pointer(@format_context, :format_context, :streams)
+      streams.get_pointer(index * FFI.type_size(:pointer))
+    end
+
+    def format_duration
+      value = Native::Layout.read_int64(@format_context, :format_context, :duration)
+      return unless valid_timestamp?(value)
+
+      Rational(value, Native::AV_TIME_BASE)
+    end
+
+    def valid_timestamp?(value)
+      value && value != Native::AV_NOPTS_VALUE && value >= 0
+    end
+
+    def free_swr
+      return unless @swr && !@swr.null?
+
+      holder = FFI::MemoryPointer.new(:pointer)
+      holder.write_pointer(@swr)
+      Native.resample.swr_free(holder)
+      @swr = nil
+    end
+
+    def free_pointer(pointer, function, library)
+      return unless pointer && !pointer.null?
+
+      holder = FFI::MemoryPointer.new(:pointer)
+      holder.write_pointer(pointer)
+      library.public_send(function, holder)
+    end
+  end
+end