Serenity Operating System
1/*
2 * Copyright (c) 2023, kleines Filmröllchen <filmroellchen@serenityos.org>
3 *
4 * SPDX-License-Identifier: BSD-2-Clause
5 */
6
7#include "QOALoader.h"
8#include "Loader.h"
9#include "LoaderError.h"
10#include "QOATypes.h"
11#include <AK/Array.h>
12#include <AK/Assertions.h>
13#include <AK/Endian.h>
14#include <AK/FixedArray.h>
15#include <AK/MemoryStream.h>
16#include <AK/Stream.h>
17#include <AK/Types.h>
18#include <LibCore/File.h>
19
20namespace Audio {
21
22QOALoaderPlugin::QOALoaderPlugin(NonnullOwnPtr<AK::SeekableStream> stream)
23 : LoaderPlugin(move(stream))
24{
25}
26
27Result<NonnullOwnPtr<QOALoaderPlugin>, LoaderError> QOALoaderPlugin::create(StringView path)
28{
29 auto stream = LOADER_TRY(Core::BufferedFile::create(LOADER_TRY(Core::File::open(path, Core::File::OpenMode::Read))));
30 auto loader = make<QOALoaderPlugin>(move(stream));
31
32 LOADER_TRY(loader->initialize());
33
34 return loader;
35}
36
37Result<NonnullOwnPtr<QOALoaderPlugin>, LoaderError> QOALoaderPlugin::create(Bytes buffer)
38{
39 auto loader = make<QOALoaderPlugin>(make<FixedMemoryStream>(buffer));
40
41 LOADER_TRY(loader->initialize());
42
43 return loader;
44}
45
46MaybeLoaderError QOALoaderPlugin::initialize()
47{
48 TRY(parse_header());
49 TRY(reset());
50 return {};
51}
52
53MaybeLoaderError QOALoaderPlugin::parse_header()
54{
55 u32 header_magic = LOADER_TRY(m_stream->read_value<BigEndian<u32>>());
56 if (header_magic != QOA::magic)
57 return LoaderError { LoaderError::Category::Format, 0, "QOA header: Magic number must be 'qoaf'" };
58
59 m_total_samples = LOADER_TRY(m_stream->read_value<BigEndian<u32>>());
60
61 return {};
62}
63
64MaybeLoaderError QOALoaderPlugin::load_one_frame(Span<Sample>& target, IsFirstFrame is_first_frame)
65{
66 QOA::FrameHeader header = LOADER_TRY(m_stream->read_value<QOA::FrameHeader>());
67
68 if (header.num_channels > 8)
69 dbgln("QOALoader: Warning: QOA frame at {} has more than 8 channels ({}), this is not supported by the reference implementation.", LOADER_TRY(m_stream->tell()) - sizeof(QOA::FrameHeader), header.num_channels);
70 if (header.num_channels == 0)
71 return LoaderError { LoaderError::Category::Format, LOADER_TRY(m_stream->tell()), "QOA frame: Number of channels must be greater than 0" };
72 if (header.sample_count > QOA::max_frame_samples)
73 return LoaderError { LoaderError::Category::Format, LOADER_TRY(m_stream->tell()), "QOA frame: Too many samples in frame" };
74
75 // We weren't given a large enough buffer; signal that we didn't write anything and return.
76 if (header.sample_count > target.size()) {
77 target = target.trim(0);
78 LOADER_TRY(m_stream->seek(-sizeof(QOA::frame_header_size), AK::SeekMode::FromCurrentPosition));
79 return {};
80 }
81
82 target = target.trim(header.sample_count);
83
84 auto lms_states = LOADER_TRY(FixedArray<QOA::LMSState>::create(header.num_channels));
85 for (size_t channel = 0; channel < header.num_channels; ++channel) {
86 auto history_packed = LOADER_TRY(m_stream->read_value<BigEndian<u64>>());
87 auto weights_packed = LOADER_TRY(m_stream->read_value<BigEndian<u64>>());
88 lms_states[channel] = { history_packed, weights_packed };
89 }
90
91 // We pre-allocate very large arrays here, but that's the last allocation of the QOA loader!
92 // Everything else is just shuffling data around.
93 // (We will also be using all of the arrays in every frame but the last one.)
94 auto channels = LOADER_TRY((FixedArray<Array<i16, QOA::max_frame_samples>>::create(header.num_channels)));
95
96 // There's usually (and at maximum) 256 slices per channel, but less at the very end.
97 // If the final slice would be partial, we still need to decode it; integer division would tell us that this final slice doesn't exist.
98 auto const slice_count = static_cast<size_t>(ceil(static_cast<double>(header.sample_count) / static_cast<double>(QOA::slice_samples)));
99 VERIFY(slice_count <= QOA::max_slices_per_frame);
100
101 // Observe the loop nesting: Slices are channel-interleaved.
102 for (size_t slice = 0; slice < slice_count; ++slice) {
103 for (size_t channel = 0; channel < header.num_channels; ++channel) {
104 auto slice_samples = channels[channel].span().slice(slice * QOA::slice_samples, QOA::slice_samples);
105 TRY(read_one_slice(lms_states[channel], slice_samples));
106 }
107 }
108
109 if (is_first_frame == IsFirstFrame::Yes) {
110 m_num_channels = header.num_channels;
111 m_sample_rate = header.sample_rate;
112 } else {
113 if (m_sample_rate != header.sample_rate)
114 return LoaderError { LoaderError::Category::Unimplemented, LOADER_TRY(m_stream->tell()), "QOA: Differing sample rate in non-initial frame" };
115 if (m_num_channels != header.num_channels)
116 m_has_uniform_channel_count = false;
117 }
118
119 switch (header.num_channels) {
120 case 1:
121 for (size_t sample = 0; sample < header.sample_count; ++sample)
122 target[sample] = Sample { static_cast<float>(channels[0][sample]) / static_cast<float>(NumericLimits<i16>::max()) };
123 break;
124 // FIXME: Combine surround channels sensibly, FlacLoader has the same simplification at the moment.
125 case 2:
126 case 3:
127 case 4:
128 case 5:
129 case 6:
130 case 7:
131 case 8:
132 default:
133 for (size_t sample = 0; sample < header.sample_count; ++sample) {
134 target[sample] = {
135 static_cast<float>(channels[0][sample]) / static_cast<float>(NumericLimits<i16>::max()),
136 static_cast<float>(channels[1][sample]) / static_cast<float>(NumericLimits<i16>::max()),
137 };
138 }
139 break;
140 }
141
142 return {};
143}
144
145ErrorOr<Vector<FixedArray<Sample>>, LoaderError> QOALoaderPlugin::load_chunks(size_t samples_to_read_from_input)
146{
147 ssize_t const remaining_samples = static_cast<ssize_t>(m_total_samples - m_loaded_samples);
148 if (remaining_samples <= 0)
149 return Vector<FixedArray<Sample>> {};
150 size_t const samples_to_read = min(samples_to_read_from_input, remaining_samples);
151 auto is_first_frame = m_loaded_samples == 0 ? IsFirstFrame::Yes : IsFirstFrame::No;
152
153 Vector<FixedArray<Sample>> frames;
154 size_t current_loaded_samples = 0;
155
156 while (current_loaded_samples < samples_to_read) {
157 auto samples = LOADER_TRY(FixedArray<Sample>::create(QOA::max_frame_samples));
158 auto slice_to_load_into = samples.span();
159 TRY(this->load_one_frame(slice_to_load_into, is_first_frame));
160 is_first_frame = IsFirstFrame::No;
161 VERIFY(slice_to_load_into.size() <= QOA::max_frame_samples);
162 current_loaded_samples += slice_to_load_into.size();
163 if (slice_to_load_into.size() != samples.size()) {
164 auto smaller_samples = LOADER_TRY(FixedArray<Sample>::create(slice_to_load_into));
165 samples.swap(smaller_samples);
166 }
167 LOADER_TRY(frames.try_append(move(samples)));
168
169 if (slice_to_load_into.size() != samples.size())
170 break;
171 }
172 m_loaded_samples += current_loaded_samples;
173
174 return frames;
175}
176
177MaybeLoaderError QOALoaderPlugin::reset()
178{
179 LOADER_TRY(m_stream->seek(QOA::header_size, AK::SeekMode::SetPosition));
180 m_loaded_samples = 0;
181 // Read the first frame, then seek back to the beginning. This is necessary since the first frame contains the sample rate and channel count.
182 auto frame_samples = LOADER_TRY(FixedArray<Sample>::create(QOA::max_frame_samples));
183 auto span = frame_samples.span();
184 LOADER_TRY(load_one_frame(span, IsFirstFrame::Yes));
185
186 LOADER_TRY(m_stream->seek(QOA::header_size, AK::SeekMode::SetPosition));
187 m_loaded_samples = 0;
188 return {};
189}
190
191MaybeLoaderError QOALoaderPlugin::seek(int sample_index)
192{
193 if (sample_index == 0 && m_loaded_samples == 0)
194 return {};
195 // A QOA file consists of 8 bytes header followed by a number of usually fixed-size frames.
196 // This fixed bitrate allows us to seek in constant time.
197 if (!m_has_uniform_channel_count)
198 return LoaderError { LoaderError::Category::Unimplemented, LOADER_TRY(m_stream->tell()), "QOA with non-uniform channel count is currently not seekable"sv };
199 /// FIXME: Change the Loader API to use size_t.
200 VERIFY(sample_index >= 0);
201 // We seek to the frame "before"; i.e. the frame that contains that sample.
202 auto const frame_of_sample = static_cast<size_t>(AK::floor<double>(static_cast<double>(sample_index) / static_cast<double>(QOA::max_frame_samples)));
203 auto const frame_size = QOA::frame_header_size + m_num_channels * (QOA::lms_state_size + sizeof(QOA::PackedSlice) * QOA::max_slices_per_frame);
204 auto const byte_index = QOA::header_size + frame_of_sample * frame_size;
205 LOADER_TRY(m_stream->seek(byte_index, AK::SeekMode::SetPosition));
206 m_loaded_samples = frame_of_sample * QOA::max_frame_samples;
207 return {};
208}
209
210MaybeLoaderError QOALoaderPlugin::read_one_slice(QOA::LMSState& lms_state, Span<i16>& samples)
211{
212 VERIFY(samples.size() == QOA::slice_samples);
213
214 auto packed_slice = LOADER_TRY(m_stream->read_value<BigEndian<u64>>());
215 auto unpacked_slice = unpack_slice(packed_slice);
216
217 for (size_t i = 0; i < QOA::slice_samples; ++i) {
218 auto const residual = unpacked_slice.residuals[i];
219 auto const predicted = lms_state.predict();
220 auto const dequantized = QOA::dequantization_table[unpacked_slice.scale_factor_index][residual];
221 auto const reconstructed = clamp(predicted + dequantized, QOA::sample_minimum, QOA::sample_maximum);
222 samples[i] = static_cast<i16>(reconstructed);
223 lms_state.update(reconstructed, dequantized);
224 }
225
226 return {};
227}
228
229QOA::UnpackedSlice QOALoaderPlugin::unpack_slice(QOA::PackedSlice packed_slice)
230{
231 size_t const scale_factor_index = (packed_slice >> 60) & 0b1111;
232 Array<u8, 20> residuals = {};
233 auto shifted_slice = packed_slice << 4;
234
235 for (size_t i = 0; i < QOA::slice_samples; ++i) {
236 residuals[i] = static_cast<u8>((shifted_slice >> 61) & 0b111);
237 shifted_slice <<= 3;
238 }
239
240 return {
241 .scale_factor_index = scale_factor_index,
242 .residuals = residuals,
243 };
244}
245
246i16 QOALoaderPlugin::qoa_divide(i16 value, i16 scale_factor)
247{
248 auto const reciprocal = QOA::reciprocal_table[scale_factor];
249 auto const n = (value * reciprocal + (1 << 15)) >> 16;
250 // Rounding away from zero gives better quantization for small values.
251 auto const n_rounded = n + (static_cast<int>(value > 0) - static_cast<int>(value < 0)) - (static_cast<int>(n > 0) - static_cast<int>(n < 0));
252 return static_cast<i16>(n_rounded);
253}
254
255}