LibAudio: Automatically write a FLAC seektable

This contains a seekpoint every 2 seconds, allowing our own players to
work better.

Userland/Libraries/LibAudio/FlacTypes.h

@@ -37,6 +37,8 @@ using FlacFrameHeaderCRC = Crypto::Checksum::CRC8<flac_polynomial>;
 static constexpr u16 ibm_polynomial = 0xA001;
 using IBMCRC = Crypto::Checksum::CRC16<ibm_polynomial>;
 
+static constexpr size_t flac_seekpoint_size = (64 + 64 + 16) / 8;
+
 // 11.8 BLOCK_TYPE (7 bits)
 enum class FlacMetadataBlockType : u8 {
     STREAMINFO = 0,     // Important data about the audio format

Userland/Libraries/LibAudio/FlacWriter.cpp

@@ -61,6 +61,8 @@ ErrorOr<void> FlacWriter::finalize()
         TRY(bit_stream.align_to_byte_boundary());
     }
 
+    TRY(flush_seektable());
+
     // TODO: Write the audio data MD5 to the header.
 
     m_stream->close();
@@ -124,6 +126,53 @@ ErrorOr<void> FlacWriter::set_metadata(Metadata const& metadata)
     return add_metadata_block(move(vorbis_block), 0);
 }
 
+size_t FlacWriter::max_number_of_seekpoints() const
+{
+    if (m_last_padding.has_value())
+        return m_last_padding->size / flac_seekpoint_size;
+
+    if (!m_cached_metadata_blocks.is_empty() && m_cached_metadata_blocks.last().type == FlacMetadataBlockType::PADDING)
+        return m_cached_metadata_blocks.last().length / flac_seekpoint_size;
+
+    return 0;
+}
+
+void FlacWriter::sample_count_hint(size_t sample_count)
+{
+    constexpr StringView oom_warning = "FLAC Warning: Couldn't use sample hint to reserve {} bytes padding; ignoring hint."sv;
+
+    auto const samples_per_seekpoint = m_sample_rate * seekpoint_period_seconds;
+    auto seekpoint_count = round_to<size_t>(static_cast<double>(sample_count) / samples_per_seekpoint);
+    // Round seekpoint count down to an even number, so that the seektable byte size is divisible by 4.
+    // One seekpoint is 18 bytes, which isn't divisible by 4.
+    seekpoint_count &= ~1;
+    auto const seektable_size = seekpoint_count * flac_seekpoint_size;
+
+    // Only modify the trailing padding block; other padding blocks are intentionally untouched.
+    if (!m_cached_metadata_blocks.is_empty() && m_cached_metadata_blocks.last().type == FlacMetadataBlockType::PADDING) {
+        auto padding_block = m_cached_metadata_blocks.last();
+        auto result = padding_block.data.try_resize(seektable_size);
+        padding_block.length = padding_block.data.size();
+        // Fuzzers and inputs with wrong large sample counts often hit this.
+        if (result.is_error())
+            dbgln(oom_warning, seektable_size);
+    } else {
+        auto empty_buffer = ByteBuffer::create_zeroed(seektable_size);
+        if (empty_buffer.is_error()) {
+            dbgln(oom_warning, seektable_size);
+            return;
+        }
+        FlacRawMetadataBlock padding {
+            .is_last_block = true,
+            .type = FlacMetadataBlockType::PADDING,
+            .length = static_cast<u32>(empty_buffer.value().size()),
+            .data = empty_buffer.release_value(),
+        };
+        // If we can't add padding, we're out of luck.
+        (void)add_metadata_block(move(padding));
+    }
+}
+
 ErrorOr<void> FlacWriter::write_header()
 {
     ByteBuffer data;
@@ -158,6 +207,19 @@ ErrorOr<void> FlacWriter::write_header()
     };
     TRY(add_metadata_block(move(streaminfo_block), 0));
 
+    // Add default padding if necessary.
+    if (m_cached_metadata_blocks.last().type != FlacMetadataBlockType::PADDING) {
+        auto padding_data = ByteBuffer::create_zeroed(default_padding);
+        if (!padding_data.is_error()) {
+            TRY(add_metadata_block({
+                .is_last_block = true,
+                .type = FlacMetadataBlockType::PADDING,
+                .length = default_padding,
+                .data = padding_data.release_value(),
+            }));
+        }
+    }
+
     TRY(m_stream->write_until_depleted(flac_magic.bytes()));
     m_streaminfo_start_index = TRY(m_stream->tell());
 
@@ -166,11 +228,18 @@ ErrorOr<void> FlacWriter::write_header()
         // Correct is_last_block flag here to avoid index shenanigans in add_metadata_block.
         auto const is_last_block = i == m_cached_metadata_blocks.size() - 1;
         block.is_last_block = is_last_block;
+        if (is_last_block) {
+            m_last_padding = LastPadding {
+                .start = TRY(m_stream->tell()),
+                .size = block.length,
+            };
+        }
 
         TRY(write_metadata_block(block));
     }
 
     m_cached_metadata_blocks.clear();
+    m_frames_start_index = TRY(m_stream->tell());
     return {};
 }
 
@@ -187,8 +256,50 @@ ErrorOr<void> FlacWriter::add_metadata_block(FlacRawMetadataBlock block, Optiona
     return {};
 }
 
-ErrorOr<void> FlacWriter::write_metadata_block(FlacRawMetadataBlock const& block)
+ErrorOr<void> FlacWriter::write_metadata_block(FlacRawMetadataBlock& block)
 {
+    if (m_state == WriteState::FormatFinalized) {
+        if (!m_last_padding.has_value())
+            return Error::from_string_view("No (more) padding available to write block into"sv);
+
+        auto const last_padding = m_last_padding.release_value();
+        if (block.length > last_padding.size)
+            return Error::from_string_view("Late metadata block doesn't fit in available padding"sv);
+
+        auto const current_position = TRY(m_stream->tell());
+        ScopeGuard guard = [&] { (void)m_stream->seek(current_position, SeekMode::SetPosition); };
+        TRY(m_stream->seek(last_padding.start, SeekMode::SetPosition));
+
+        // No more padding after this: the new block is the last.
+        auto new_size = last_padding.size - block.length;
+        if (new_size == 0)
+            block.is_last_block = true;
+
+        TRY(m_stream->write_value(block));
+
+        // If the size is zero, we don't need to write a new padding block.
+        // If the size is between 1 and 3, we have empty space that cannot be marked with an empty padding block, so we must abort.
+        // Other code should make sure that this never happens; e.g. our seektable only has sizes divisible by 4 anyways.
+        // If the size is 4, we have no padding, but the padding block header can be written without any subsequent payload.
+        if (new_size >= 4) {
+            FlacRawMetadataBlock new_padding_block {
+                .is_last_block = true,
+                .type = FlacMetadataBlockType::PADDING,
+                .length = static_cast<u32>(new_size),
+                .data = TRY(ByteBuffer::create_zeroed(new_size)),
+            };
+            m_last_padding = LastPadding {
+                .start = TRY(m_stream->tell()),
+                .size = new_size,
+            };
+            TRY(m_stream->write_value(new_padding_block));
+        } else if (new_size != 0) {
+            return Error::from_string_view("Remaining padding is not divisible by 4, there will be some stray zero bytes!"sv);
+        }
+
+        return {};
+    }
+
     return m_stream->write_value(block);
 }
 
@@ -204,6 +315,52 @@ ErrorOr<void> FlacRawMetadataBlock::write_to_stream(Stream& stream) const
     return {};
 }
 
+ErrorOr<void> FlacWriter::flush_seektable()
+{
+    if (m_cached_seektable.size() == 0)
+        return {};
+
+    auto max_seekpoints = max_number_of_seekpoints();
+    if (max_seekpoints < m_cached_seektable.size()) {
+        dbgln("FLAC Warning: There are {} seekpoints, but we only have space for {}. Some seekpoints will be dropped.", m_cached_seektable.size(), max_seekpoints);
+        // Drop seekpoints in regular intervals to space out the loss of seek precision.
+        auto const points_to_drop = m_cached_seektable.size() - max_seekpoints;
+        auto const drop_interval = static_cast<double>(m_cached_seektable.size()) / static_cast<double>(points_to_drop);
+        double ratio = 0.;
+        for (size_t i = 0; i < m_cached_seektable.size(); ++i) {
+            // Avoid dropping the first seekpoint.
+            if (ratio > drop_interval) {
+                m_cached_seektable.seek_points().remove(i);
+                --i;
+                ratio -= drop_interval;
+            }
+            ++ratio;
+        }
+        // Account for integer division imprecisions.
+        if (max_seekpoints < m_cached_seektable.size())
+            m_cached_seektable.seek_points().shrink(max_seekpoints);
+    }
+
+    auto seektable_data = TRY(ByteBuffer::create_zeroed(m_cached_seektable.size() * flac_seekpoint_size));
+    FixedMemoryStream seektable_stream { seektable_data.bytes() };
+
+    for (auto const& seekpoint : m_cached_seektable.seek_points()) {
+        // https://www.ietf.org/archive/id/draft-ietf-cellar-flac-08.html#name-seekpoint
+        TRY(seektable_stream.write_value<BigEndian<u64>>(seekpoint.sample_index));
+        TRY(seektable_stream.write_value<BigEndian<u64>>(seekpoint.byte_offset));
+        // This is probably wrong for the last frame, but it doesn't seem to matter.
+        TRY(seektable_stream.write_value<BigEndian<u16>>(block_size));
+    }
+
+    FlacRawMetadataBlock seektable {
+        .is_last_block = false,
+        .type = FlacMetadataBlockType::SEEKTABLE,
+        .length = static_cast<u32>(seektable_data.size()),
+        .data = move(seektable_data),
+    };
+    return write_metadata_block(seektable);
+}
+
 // If the given sample count is uncommon, this function will return one of the uncommon marker block sizes.
 // The caller has to handle and add these later manually.
 static BlockSizeCategory to_common_block_size(u16 sample_count)
@@ -394,10 +551,24 @@ ErrorOr<void> FlacWriter::write_frame()
         }
     }
 
-    return write_frame_for(subframe_samples, channel_type);
+    auto const sample_index = m_sample_count;
+    auto const frame_start_byte = TRY(write_frame_for(subframe_samples, channel_type));
+
+    // Insert a seekpoint if necessary.
+    auto const seekpoint_period_samples = m_sample_rate * seekpoint_period_seconds;
+    auto const last_seekpoint = m_cached_seektable.seek_point_before(sample_index);
+    if (!last_seekpoint.has_value() || static_cast<double>(sample_index - last_seekpoint->sample_index) >= seekpoint_period_samples) {
+        dbgln_if(FLAC_ENCODER_DEBUG, "Inserting seekpoint at sample index {} frame start {}", sample_index, frame_start_byte);
+        TRY(m_cached_seektable.insert_seek_point({
+            .sample_index = sample_index,
+            .byte_offset = frame_start_byte - m_frames_start_index,
+        }));
+    }
+
+    return {};
 }
 
-ErrorOr<void> FlacWriter::write_frame_for(ReadonlySpan<Vector<i64, block_size>> subblock, FlacFrameChannelType channel_type)
+ErrorOr<size_t> FlacWriter::write_frame_for(ReadonlySpan<Vector<i64, block_size>> subblock, FlacFrameChannelType channel_type)
 {
     auto sample_count = subblock.first().size();
 
@@ -406,7 +577,6 @@ ErrorOr<void> FlacWriter::write_frame_for(ReadonlySpan<Vector<i64, block_size>>
         .sample_count = static_cast<u16>(sample_count),
         .sample_or_frame_index = static_cast<u32>(m_current_frame),
         .blocking_strategy = BlockingStrategy::Fixed,
-        // FIXME: We should brute-force channel coupling for stereo.
         .channels = channel_type,
         .bit_depth = static_cast<u8>(m_bits_per_sample),
         // Calculated for us during header write.
@@ -444,7 +614,7 @@ ErrorOr<void> FlacWriter::write_frame_for(ReadonlySpan<Vector<i64, block_size>>
     m_current_frame++;
     m_sample_count += sample_count;
 
-    return {};
+    return frame_start_offset;
 }
 
 ErrorOr<void> FlacWriter::write_subframe(ReadonlySpan<i64> subframe, BigEndianOutputBitStream& bit_stream, u8 bits_per_sample)

Userland/Libraries/LibAudio/FlacWriter.h

@@ -55,6 +55,10 @@ class FlacWriter : public Encoder {
     // After how many useless (i.e. worse than current optimal) Rice parameters to abort parameter search.
     // Note that due to the zig-zag search, we start with searching the parameters that are most likely to be good.
     static constexpr size_t useless_parameter_threshold = 2;
+    // How often a seek point is inserted.
+    static constexpr double seekpoint_period_seconds = 2.0;
+    // Default padding reserved for seek points; enough for almost 4 minutes of audio.
+    static constexpr size_t default_padding = 2048;
 
     enum class WriteState {
         // Header has not been written at all, audio data cannot be written.
@@ -83,6 +87,12 @@ public:
     ErrorOr<void> set_sample_rate(u32 sample_rate);
     ErrorOr<void> set_bits_per_sample(u16 bits_per_sample);
 
+    // The FLAC encoder by default tries to reserve some space for seek points,
+    // but that may not be enough if more than approximately four minutes of audio are stored.
+    // The sample count hint can be used to instruct the FLAC encoder on how much space to reserve for seek points,
+    // which will both reduce the padding for small files and allow the FLAC encoder to write seek points at the end of large files.
+    virtual void sample_count_hint(size_t sample_count) override;
+
     virtual ErrorOr<void> set_metadata(Metadata const& metadata) override;
 
     ErrorOr<void> finalize_header_format();
@@ -92,7 +102,8 @@ private:
     ErrorOr<void> write_header();
 
     ErrorOr<void> write_frame();
-    ErrorOr<void> write_frame_for(ReadonlySpan<Vector<i64, block_size>> subblock, FlacFrameChannelType channel_type);
+    // Returns the frame start byte offset, to be used for creating a seektable.
+    ErrorOr<size_t> write_frame_for(ReadonlySpan<Vector<i64, block_size>> subblock, FlacFrameChannelType channel_type);
     ErrorOr<void> write_subframe(ReadonlySpan<i64> subframe, BigEndianOutputBitStream& bit_stream, u8 bits_per_sample);
     ErrorOr<void> write_lpc_subframe(FlacLPCEncodedSubframe lpc_subframe, BigEndianOutputBitStream& bit_stream, u8 bits_per_sample);
     ErrorOr<void> write_verbatim_subframe(ReadonlySpan<i64> subframe, BigEndianOutputBitStream& bit_stream, u8 bits_per_sample);
@@ -104,7 +115,12 @@ private:
     ErrorOr<Optional<FlacLPCEncodedSubframe>> encode_fixed_lpc(FlacFixedLPC order, ReadonlySpan<i64> subframe, size_t current_min_cost, u8 bits_per_sample);
 
     ErrorOr<void> add_metadata_block(FlacRawMetadataBlock block, Optional<size_t> insertion_index = {});
-    ErrorOr<void> write_metadata_block(FlacRawMetadataBlock const& block);
+    // Depending on whether the header is finished or not, we either write to the current position for an unfinished header,
+    // or we write to the start of the last padding and adjust that padding block.
+    ErrorOr<void> write_metadata_block(FlacRawMetadataBlock& block);
+    // Determine how many seekpoints we can write depending on the size of our final padding.
+    size_t max_number_of_seekpoints() const;
+    ErrorOr<void> flush_seektable();
 
     NonnullOwnPtr<SeekableStream> m_stream;
     WriteState m_state { WriteState::HeaderUnwritten };
@@ -122,9 +138,21 @@ private:
     size_t m_sample_count { 0 };
     // Remember where the STREAMINFO block was written in the stream.
     size_t m_streaminfo_start_index;
+    // Start of the first frame, used for calculating seektable byte offsets.
+    size_t m_frames_start_index;
+
+    struct LastPadding {
+        size_t start;
+        size_t size;
+    };
+    // Remember last PADDING block data, since we overwrite part of it with "late" metadata blocks.
+    Optional<LastPadding> m_last_padding;
 
     // Raw metadata blocks that will be written out before header finalization.
     Vector<FlacRawMetadataBlock> m_cached_metadata_blocks;
+
+    // The full seektable, may be fully or partially written.
+    SeekTable m_cached_seektable {};
 };
 
 }

Userland/Libraries/LibAudio/GenericTypes.cpp

@@ -21,6 +21,11 @@ ReadonlySpan<SeekPoint> SeekTable::seek_points() const
     return m_seek_points.span();
 }
 
+Vector<SeekPoint>& SeekTable::seek_points()
+{
+    return m_seek_points;
+}
+
 Optional<SeekPoint const&> SeekTable::seek_point_before(u64 sample_index) const
 {
     if (m_seek_points.is_empty())

Userland/Libraries/LibAudio/GenericTypes.h

@@ -66,6 +66,7 @@ public:
 
     size_t size() const;
     ReadonlySpan<SeekPoint> seek_points() const;
+    Vector<SeekPoint>& seek_points();
 
     ErrorOr<void> insert_seek_point(SeekPoint);