1//! How to read arbitrary but specific selection of arbitrary channels.
2//! This is not a zero-cost abstraction.
3
4use crate::image::recursive::*;
5use crate::block::samples::*;
6use crate::image::*;
7use crate::math::*;
8use crate::meta::header::*;
9use crate::error::*;
10use crate::block::UncompressedBlock;
11use crate::image::read::layers::{ChannelsReader, ReadChannels};
12use crate::block::chunk::TileCoordinates;
13
14use std::marker::PhantomData;
15use crate::io::Read;
16
17
18/// Can be attached one more channel reader.
19/// Call `required` or `optional` on this object to declare another channel to be read from the file.
20/// Call `collect_pixels` at last to define how the previously declared pixels should be stored.
21pub trait ReadSpecificChannel: Sized + CheckDuplicates {
22
23 /// A separate internal reader for the pixels. Will be of type `Recursive<_, SampleReader<_>>`,
24 /// depending on the pixels of the specific channel combination.
25 type RecursivePixelReader: RecursivePixelReader;
26
27 /// Create a separate internal reader for the pixels of the specific channel combination.
28 fn create_recursive_reader(&self, channels: &ChannelList) -> Result<Self::RecursivePixelReader>;
29
30 /// Plan to read an additional channel from the image, with the specified name.
31 /// If the channel cannot be found in the image when the image is read, the image will not be loaded.
32 /// The generic parameter can usually be inferred from the closure in `collect_pixels`.
33 fn required<Sample>(self, channel_name: impl Into<Text>) -> ReadRequiredChannel<Self, Sample> {
34 let channel_name = channel_name.into();
35 assert!(self.already_contains(&channel_name).not(), "a channel with the name `{}` is already defined", channel_name);
36 ReadRequiredChannel { channel_name, previous_channels: self, px: Default::default() }
37 }
38
39 /// Plan to read an additional channel from the image, with the specified name.
40 /// If the file does not contain this channel, the specified default sample will be returned instead.
41 /// You can check whether the channel has been loaded by
42 /// checking the presence of the optional channel description before instantiating your own image.
43 /// The generic parameter can usually be inferred from the closure in `collect_pixels`.
44 fn optional<Sample>(self, channel_name: impl Into<Text>, default_sample: Sample)
45 -> ReadOptionalChannel<Self, Sample>
46 {
47 let channel_name = channel_name.into();
48 assert!(self.already_contains(&channel_name).not(), "a channel with the name `{}` is already defined", channel_name);
49 ReadOptionalChannel { channel_name, previous_channels: self, default_sample }
50 }
51
52 /// Using two closures, define how to store the pixels.
53 /// The first closure creates an image, and the second closure inserts a single pixel.
54 /// The type of the pixel can be defined by the second closure;
55 /// it must be a tuple containing `f16`, `f32`, `u32` or `Sample` values.
56 /// See the examples for more information.
57 fn collect_pixels<Pixel, PixelStorage, CreatePixels, SetPixel>(
58 self, create_pixels: CreatePixels, set_pixel: SetPixel
59 ) -> CollectPixels<Self, Pixel, PixelStorage, CreatePixels, SetPixel>
60 where
61 <Self::RecursivePixelReader as RecursivePixelReader>::RecursivePixel: IntoTuple<Pixel>,
62 <Self::RecursivePixelReader as RecursivePixelReader>::RecursiveChannelDescriptions: IntoNonRecursive,
63 CreatePixels: Fn(
64 Vec2<usize>,
65 &<<Self::RecursivePixelReader as RecursivePixelReader>::RecursiveChannelDescriptions as IntoNonRecursive>::NonRecursive
66 ) -> PixelStorage,
67 SetPixel: Fn(&mut PixelStorage, Vec2<usize>, Pixel),
68 {
69 CollectPixels { read_channels: self, set_pixel, create_pixels, px: Default::default() }
70 }
71}
72
73/// A reader containing sub-readers for reading the pixel content of an image.
74pub trait RecursivePixelReader {
75
76 /// The channel descriptions from the image.
77 /// Will be converted to a tuple before being stored in `SpecificChannels<_, ChannelDescriptions>`.
78 type RecursiveChannelDescriptions;
79
80 /// Returns the channel descriptions based on the channels in the file.
81 fn get_descriptions(&self) -> Self::RecursiveChannelDescriptions;
82
83 /// The pixel type. Will be converted to a tuple at the end of the process.
84 type RecursivePixel: Copy + Default + 'static;
85
86 /// Read the line of pixels.
87 fn read_pixels<'s, FullPixel>(
88 &self, bytes: &'s[u8], pixels: &mut [FullPixel],
89 get_pixel: impl Fn(&mut FullPixel) -> &mut Self::RecursivePixel
90 );
91}
92
93// does not use the generic `Recursive` struct to reduce the number of angle brackets in the public api
94/// Used to read another specific channel from an image.
95/// Contains the previous `ReadChannels` objects.
96#[derive(Clone, Debug)]
97pub struct ReadOptionalChannel<ReadChannels, Sample> {
98 previous_channels: ReadChannels,
99 channel_name: Text,
100 default_sample: Sample,
101}
102
103// does not use the generic `Recursive` struct to reduce the number of angle brackets in the public api
104/// Used to read another specific channel from an image.
105/// Contains the previous `ReadChannels` objects.
106#[derive(Clone, Debug)]
107pub struct ReadRequiredChannel<ReadChannels, Sample> {
108 previous_channels: ReadChannels,
109 channel_name: Text,
110 px: PhantomData<Sample>,
111}
112
113/// Specifies how to collect all the specified channels into a number of individual pixels.
114#[derive(Copy, Clone, Debug)]
115pub struct CollectPixels<ReadChannels, Pixel, PixelStorage, CreatePixels, SetPixel> {
116 read_channels: ReadChannels,
117 create_pixels: CreatePixels,
118 set_pixel: SetPixel,
119 px: PhantomData<(Pixel, PixelStorage)>,
120}
121
122impl<Inner: CheckDuplicates, Sample> CheckDuplicates for ReadRequiredChannel<Inner, Sample> {
123 fn already_contains(&self, name: &Text) -> bool {
124 &self.channel_name == name || self.previous_channels.already_contains(name)
125 }
126}
127
128impl<Inner: CheckDuplicates, Sample> CheckDuplicates for ReadOptionalChannel<Inner, Sample> {
129 fn already_contains(&self, name: &Text) -> bool {
130 &self.channel_name == name || self.previous_channels.already_contains(name)
131 }
132}
133
134impl<'s, InnerChannels, Pixel, PixelStorage, CreatePixels, SetPixel: 's>
135ReadChannels<'s> for CollectPixels<InnerChannels, Pixel, PixelStorage, CreatePixels, SetPixel>
136 where
137 InnerChannels: ReadSpecificChannel,
138 <InnerChannels::RecursivePixelReader as RecursivePixelReader>::RecursivePixel: IntoTuple<Pixel>,
139 <InnerChannels::RecursivePixelReader as RecursivePixelReader>::RecursiveChannelDescriptions: IntoNonRecursive,
140 CreatePixels: Fn(Vec2<usize>, &<<InnerChannels::RecursivePixelReader as RecursivePixelReader>::RecursiveChannelDescriptions as IntoNonRecursive>::NonRecursive) -> PixelStorage,
141 SetPixel: Fn(&mut PixelStorage, Vec2<usize>, Pixel),
142{
143 type Reader = SpecificChannelsReader<
144 PixelStorage, &'s SetPixel,
145 InnerChannels::RecursivePixelReader,
146 Pixel,
147 >;
148
149 fn create_channels_reader(&'s self, header: &Header) -> Result<Self::Reader> {
150 if header.deep { return Err(Error::invalid(message:"`SpecificChannels` does not support deep data yet")) }
151
152 let pixel_reader: impl IntoTuple + IntoNonRecursive = self.read_channels.create_recursive_reader(&header.channels)?;
153 let channel_descriptions: < + IntoNonRecursive as RecursivePixelReader>::RecursiveChannelDescriptions as IntoNonRecursive>::NonRecursive = pixel_reader.get_descriptions().into_non_recursive();// TODO not call this twice
154
155 let create: &CreatePixels = &self.create_pixels;
156 let pixel_storage: PixelStorage = create(header.layer_size, &channel_descriptions);
157
158 Ok(SpecificChannelsReader {
159 set_pixel: &self.set_pixel,
160 pixel_storage,
161 pixel_reader,
162 px: Default::default()
163 })
164 }
165}
166
167/// The reader that holds the temporary data that is required to read some specified channels.
168#[derive(Copy, Clone, Debug)]
169pub struct SpecificChannelsReader<PixelStorage, SetPixel, PixelReader, Pixel> {
170 set_pixel: SetPixel,
171 pixel_storage: PixelStorage,
172 pixel_reader: PixelReader,
173 px: PhantomData<Pixel>
174}
175
176impl<PixelStorage, SetPixel, PxReader, Pixel>
177ChannelsReader for SpecificChannelsReader<PixelStorage, SetPixel, PxReader, Pixel>
178 where PxReader: RecursivePixelReader,
179 PxReader::RecursivePixel: IntoTuple<Pixel>,
180 PxReader::RecursiveChannelDescriptions: IntoNonRecursive,
181 SetPixel: Fn(&mut PixelStorage, Vec2<usize>, Pixel),
182{
183 type Channels = SpecificChannels<PixelStorage, <PxReader::RecursiveChannelDescriptions as IntoNonRecursive>::NonRecursive>;
184
185 fn filter_block(&self, tile: TileCoordinates) -> bool { tile.is_largest_resolution_level() } // TODO all levels
186
187 fn read_block(&mut self, header: &Header, block: UncompressedBlock) -> UnitResult {
188 let mut pixels = vec![PxReader::RecursivePixel::default(); block.index.pixel_size.width()]; // TODO allocate once in self
189
190 let byte_lines = block.data.chunks_exact(header.channels.bytes_per_pixel * block.index.pixel_size.width());
191 debug_assert_eq!(byte_lines.len(), block.index.pixel_size.height(), "invalid block lines split");
192
193 for (y_offset, line_bytes) in byte_lines.enumerate() { // TODO sampling
194 // this two-step copy method should be very cache friendly in theory, and also reduce sample_type lookup count
195 self.pixel_reader.read_pixels(line_bytes, &mut pixels, |px| px);
196
197 for (x_offset, pixel) in pixels.iter().enumerate() {
198 let set_pixel = &self.set_pixel;
199 set_pixel(&mut self.pixel_storage, block.index.pixel_position + Vec2(x_offset, y_offset), pixel.into_tuple());
200 }
201 }
202
203 Ok(())
204 }
205
206 fn into_channels(self) -> Self::Channels {
207 SpecificChannels { channels: self.pixel_reader.get_descriptions().into_non_recursive(), pixels: self.pixel_storage }
208 }
209}
210
211
212/// Read zero channels from an image. Call `with_named_channel` on this object
213/// to read as many channels as desired.
214pub type ReadZeroChannels = NoneMore;
215
216impl ReadSpecificChannel for NoneMore {
217 type RecursivePixelReader = NoneMore;
218 fn create_recursive_reader(&self, _: &ChannelList) -> Result<Self::RecursivePixelReader> { Ok(NoneMore) }
219}
220
221impl<DefaultSample, ReadChannels> ReadSpecificChannel for ReadOptionalChannel<ReadChannels, DefaultSample>
222 where ReadChannels: ReadSpecificChannel, DefaultSample: FromNativeSample + 'static,
223{
224 type RecursivePixelReader = Recursive<ReadChannels::RecursivePixelReader, OptionalSampleReader<DefaultSample>>;
225
226 fn create_recursive_reader(&self, channels: &ChannelList) -> Result<Self::RecursivePixelReader> {
227 debug_assert!(self.previous_channels.already_contains(&self.channel_name).not(), "duplicate channel name: {}", self.channel_name);
228
229 let inner_samples_reader: ::RecursivePixelReader = self.previous_channels.create_recursive_reader(channels)?;
230 let reader: Option> = channelsOption<(usize, &ChannelDescription)>.channels_with_byte_offset()
231 .find(|(_, channel: &&ChannelDescription)| channel.name == self.channel_name)
232 .map(|(channel_byte_offset: usize, channel: &ChannelDescription)| SampleReader {
233 channel_byte_offset, channel: channel.clone(),
234 px: Default::default()
235 });
236
237 Ok(Recursive::new(inner_samples_reader, value:OptionalSampleReader {
238 reader, default_sample: self.default_sample,
239 }))
240 }
241}
242
243impl<Sample, ReadChannels> ReadSpecificChannel for ReadRequiredChannel<ReadChannels, Sample>
244 where ReadChannels: ReadSpecificChannel, Sample: FromNativeSample + 'static
245{
246 type RecursivePixelReader = Recursive<ReadChannels::RecursivePixelReader, SampleReader<Sample>>;
247
248 fn create_recursive_reader(&self, channels: &ChannelList) -> Result<Self::RecursivePixelReader> {
249 let previous_samples_reader: ::RecursivePixelReader = self.previous_channels.create_recursive_reader(channels)?;
250 let (channel_byte_offset: usize, channel: &ChannelDescription) = channels.channels_with_byte_offset()
251 .find(|(_, channel)| channel.name == self.channel_name)
252 .ok_or_else(|| Error::invalid(message:format!(
253 "layer does not contain all of your specified channels (`{}` is missing)",
254 self.channel_name
255 )))?;
256
257 Ok(Recursive::new(inner:previous_samples_reader, value:SampleReader { channel_byte_offset, channel: channel.clone(), px: Default::default() }))
258 }
259}
260
261/// Reader for a single channel. Generic over the concrete sample type (f16, f32, u32).
262#[derive(Clone, Debug)]
263pub struct SampleReader<Sample> {
264
265 /// to be multiplied with line width!
266 channel_byte_offset: usize,
267
268 channel: ChannelDescription,
269 px: PhantomData<Sample>
270}
271
272/// Reader for a single channel. Generic over the concrete sample type (f16, f32, u32).
273/// Can also skip reading a channel if it could not be found in the image.
274#[derive(Clone, Debug)]
275pub struct OptionalSampleReader<DefaultSample> {
276 reader: Option<SampleReader<DefaultSample>>,
277 default_sample: DefaultSample,
278}
279
280impl<Sample: FromNativeSample> SampleReader<Sample> {
281 fn read_own_samples<'s, FullPixel>(
282 &self, bytes: &'s[u8], pixels: &mut [FullPixel],
283 get_sample: impl Fn(&mut FullPixel) -> &mut Sample
284 ){
285 let start_index = pixels.len() * self.channel_byte_offset;
286 let byte_count = pixels.len() * self.channel.sample_type.bytes_per_sample();
287 let mut own_bytes_reader = &mut &bytes[start_index .. start_index + byte_count]; // TODO check block size somewhere
288 let mut samples_out = pixels.iter_mut().map(|pixel| get_sample(pixel));
289
290 // match the type once for the whole line, not on every single sample
291 match self.channel.sample_type {
292 SampleType::F16 => read_and_convert_all_samples_batched(
293 &mut own_bytes_reader, &mut samples_out,
294 Sample::from_f16s
295 ),
296
297 SampleType::F32 => read_and_convert_all_samples_batched(
298 &mut own_bytes_reader, &mut samples_out,
299 Sample::from_f32s
300 ),
301
302 SampleType::U32 => read_and_convert_all_samples_batched(
303 &mut own_bytes_reader, &mut samples_out,
304 Sample::from_u32s
305 ),
306 }
307
308 debug_assert!(samples_out.next().is_none(), "not all samples have been converted");
309 debug_assert!(own_bytes_reader.is_empty(), "bytes left after reading all samples");
310 }
311}
312
313
314/// Does the same as `convert_batch(in_bytes.chunks().map(From::from_bytes))`, but vectorized.
315/// Reads the samples for one line, using the sample type specified in the file,
316/// and then converts those to the desired sample types.
317/// Uses batches to allow vectorization, converting multiple values with one instruction.
318fn read_and_convert_all_samples_batched<'t, From, To>(
319 mut in_bytes: impl Read,
320 out_samples: &mut impl ExactSizeIterator<Item=&'t mut To>,
321 convert_batch: fn(&[From], &mut [To])
322) where From: Data + Default + Copy, To: 't + Default + Copy
323{
324 // this is not a global! why is this warning triggered?
325 #[allow(non_upper_case_globals)]
326 const batch_size: usize = 16;
327
328 let total_sample_count = out_samples.len();
329 let batch_count = total_sample_count / batch_size;
330 let remaining_samples_count = total_sample_count % batch_size;
331
332 let len_error_msg = "sample count was miscalculated";
333 let byte_error_msg = "error when reading from in-memory slice";
334
335 // write samples from a given slice to the output iterator. should be inlined.
336 let output_n_samples = &mut move |samples: &[To]| {
337 for converted_sample in samples {
338 *out_samples.next().expect(len_error_msg) = *converted_sample;
339 }
340 };
341
342 // read samples from the byte source into a given slice. should be inlined.
343 // todo: use #[inline] when available
344 // error[E0658]: attributes on expressions are experimental,
345 // see issue #15701 <https://github.com/rust-lang/rust/issues/15701> for more information
346 let read_n_samples = &mut move |samples: &mut [From]| {
347 Data::read_slice(&mut in_bytes, samples).expect(byte_error_msg);
348 };
349
350 // temporary arrays with fixed size, operations should be vectorized within these arrays
351 let mut source_samples_batch: [From; batch_size] = Default::default();
352 let mut desired_samples_batch: [To; batch_size] = Default::default();
353
354 // first convert all whole batches, size statically known to be 16 element arrays
355 for _ in 0 .. batch_count {
356 read_n_samples(&mut source_samples_batch);
357 convert_batch(source_samples_batch.as_slice(), desired_samples_batch.as_mut_slice());
358 output_n_samples(&desired_samples_batch);
359 }
360
361 // then convert a partial remaining batch, size known only at runtime
362 if remaining_samples_count != 0 {
363 let source_samples_batch = &mut source_samples_batch[..remaining_samples_count];
364 let desired_samples_batch = &mut desired_samples_batch[..remaining_samples_count];
365
366 read_n_samples(source_samples_batch);
367 convert_batch(source_samples_batch, desired_samples_batch);
368 output_n_samples(desired_samples_batch);
369 }
370}
371
372#[cfg(test)]
373mod test {
374 use super::*;
375
376 #[test]
377 fn equals_naive_f32(){
378 for total_array_size in [3, 7, 30, 41, 120, 10_423] {
379 let input_f32s = (0..total_array_size).map(|_| rand::random::<f32>()).collect::<Vec<f32>>();
380 let in_f32s_bytes = input_f32s.iter().cloned().flat_map(f32::to_le_bytes).collect::<Vec<u8>>();
381
382 let mut out_f16_samples_batched = vec![
383 f16::from_f32(rand::random::<f32>());
384 total_array_size
385 ];
386
387 read_and_convert_all_samples_batched(
388 &mut in_f32s_bytes.as_slice(),
389 &mut out_f16_samples_batched.iter_mut(),
390 f16::from_f32s
391 );
392
393 let out_f16_samples_naive = input_f32s.iter()
394 .cloned().map(f16::from_f32);
395
396 assert!(out_f16_samples_naive.eq(out_f16_samples_batched));
397 }
398 }
399}
400
401
402impl RecursivePixelReader for NoneMore {
403 type RecursiveChannelDescriptions = NoneMore;
404 fn get_descriptions(&self) -> Self::RecursiveChannelDescriptions { NoneMore }
405
406 type RecursivePixel = NoneMore;
407
408 fn read_pixels<'s, FullPixel>(
409 &self, _: &'s[u8], _: &mut [FullPixel],
410 _: impl Fn(&mut FullPixel) -> &mut NoneMore
411 ){}
412}
413
414impl<Sample, InnerReader: RecursivePixelReader>
415 RecursivePixelReader
416 for Recursive<InnerReader, SampleReader<Sample>>
417 where Sample: FromNativeSample + 'static
418{
419 type RecursiveChannelDescriptions = Recursive<InnerReader::RecursiveChannelDescriptions, ChannelDescription>;
420 fn get_descriptions(&self) -> Self::RecursiveChannelDescriptions { Recursive::new(self.inner.get_descriptions(), self.value.channel.clone()) }
421
422 type RecursivePixel = Recursive<InnerReader::RecursivePixel, Sample>;
423
424 fn read_pixels<'s, FullPixel>(
425 &self, bytes: &'s[u8], pixels: &mut [FullPixel],
426 get_pixel: impl Fn(&mut FullPixel) -> &mut Self::RecursivePixel
427 ) {
428 self.value.read_own_samples(bytes, pixels, |px: &mut FullPixel| &mut get_pixel(px).value);
429 self.inner.read_pixels(bytes, pixels, |px: &mut FullPixel| &mut get_pixel(px).inner);
430 }
431}
432
433impl<Sample, InnerReader: RecursivePixelReader>
434RecursivePixelReader
435for Recursive<InnerReader, OptionalSampleReader<Sample>>
436 where Sample: FromNativeSample + 'static
437{
438 type RecursiveChannelDescriptions = Recursive<InnerReader::RecursiveChannelDescriptions, Option<ChannelDescription>>;
439 fn get_descriptions(&self) -> Self::RecursiveChannelDescriptions { Recursive::new(
440 self.inner.get_descriptions(), self.value.reader.as_ref().map(|reader| reader.channel.clone())
441 ) }
442
443 type RecursivePixel = Recursive<InnerReader::RecursivePixel, Sample>;
444
445 fn read_pixels<'s, FullPixel>(
446 &self, bytes: &'s[u8], pixels: &mut [FullPixel],
447 get_pixel: impl Fn(&mut FullPixel) -> &mut Self::RecursivePixel
448 ) {
449 if let Some(reader) = &self.value.reader {
450 reader.read_own_samples(bytes, pixels, |px| &mut get_pixel(px).value);
451 }
452 else {
453 // if this channel is optional and was not found in the file, fill the default sample
454 for pixel in pixels.iter_mut() {
455 get_pixel(pixel).value = self.value.default_sample;
456 }
457 }
458
459 self.inner.read_pixels(bytes, pixels, |px| &mut get_pixel(px).inner);
460 }
461}
462
463
464