1 // Copyright 2024 Google LLC
2 //
3 // Licensed under the Apache License, Version 2.0 (the "License");
4 // you may not use this file except in compliance with the License.
5 // You may obtain a copy of the License at
6 //
7 //     http://www.apache.org/licenses/LICENSE-2.0
8 //
9 // Unless required by applicable law or agreed to in writing, software
10 // distributed under the License is distributed on an "AS IS" BASIS,
11 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 // See the License for the specific language governing permissions and
13 // limitations under the License.
14 
15 use crate::decoder::*;
16 use crate::*;
17 
18 pub const MAX_AV1_LAYER_COUNT: usize = 4;
19 
20 #[derive(Debug, Default)]
21 pub struct DecodeSample {
22     pub item_id: u32, // 1-based. 0 if it comes from a track.
23     pub offset: u64,
24     pub size: usize,
25     pub spatial_id: u8,
26     pub sync: bool,
27 }
28 
29 impl DecodeSample {
partial_data<'a>( &'a self, io: &'a mut Box<impl decoder::IO + ?Sized>, buffer: &'a Option<Vec<u8>>, size: usize, ) -> AvifResult<&'a [u8]>30     pub fn partial_data<'a>(
31         &'a self,
32         io: &'a mut Box<impl decoder::IO + ?Sized>,
33         buffer: &'a Option<Vec<u8>>,
34         size: usize,
35     ) -> AvifResult<&'a [u8]> {
36         match buffer {
37             Some(x) => {
38                 let start_offset = usize_from_u64(self.offset)?;
39                 let end_offset = checked_add!(start_offset, size)?;
40                 let range = start_offset..end_offset;
41                 check_slice_range(x.len(), &range)?;
42                 Ok(&x[range])
43             }
44             None => {
45                 let data = io.read(self.offset, size)?;
46                 if data.len() != size {
47                     Err(AvifError::TruncatedData)
48                 } else {
49                     Ok(data)
50                 }
51             }
52         }
53     }
54 
data<'a>( &'a self, io: &'a mut Box<impl decoder::IO + ?Sized>, buffer: &'a Option<Vec<u8>>, ) -> AvifResult<&'a [u8]>55     pub fn data<'a>(
56         &'a self,
57         io: &'a mut Box<impl decoder::IO + ?Sized>,
58         buffer: &'a Option<Vec<u8>>,
59     ) -> AvifResult<&'a [u8]> {
60         self.partial_data(io, buffer, self.size)
61     }
62 }
63 
64 #[derive(Debug, Default)]
65 pub struct DecodeInput {
66     pub samples: Vec<DecodeSample>,
67     pub all_layers: bool,
68     pub category: Category,
69 }
70 
71 #[derive(Clone, Copy, Debug, Default)]
72 pub struct Grid {
73     pub rows: u32,
74     pub columns: u32,
75     pub width: u32,
76     pub height: u32,
77 }
78 
79 #[derive(Debug, Default)]
80 pub struct TileInfo {
81     pub tile_count: u32,
82     pub decoded_tile_count: u32,
83     pub grid: Grid,
84 }
85 
86 impl TileInfo {
is_grid(&self) -> bool87     pub fn is_grid(&self) -> bool {
88         self.grid.rows > 0 && self.grid.columns > 0
89     }
90 
grid_tile_count(&self) -> AvifResult<u32>91     pub fn grid_tile_count(&self) -> AvifResult<u32> {
92         if self.is_grid() {
93             checked_mul!(self.grid.rows, self.grid.columns)
94         } else {
95             Ok(1)
96         }
97     }
98 
decoded_row_count(&self, image_height: u32, tile_height: u32) -> u3299     pub fn decoded_row_count(&self, image_height: u32, tile_height: u32) -> u32 {
100         if self.decoded_tile_count == 0 {
101             return 0;
102         }
103         if self.decoded_tile_count == self.tile_count || !self.is_grid() {
104             return image_height;
105         }
106         std::cmp::min(
107             (self.decoded_tile_count / self.grid.columns) * tile_height,
108             image_height,
109         )
110     }
111 
is_fully_decoded(&self) -> bool112     pub fn is_fully_decoded(&self) -> bool {
113         self.tile_count == self.decoded_tile_count
114     }
115 }
116 
117 #[derive(Default)]
118 pub struct Tile {
119     pub width: u32,
120     pub height: u32,
121     pub operating_point: u8,
122     pub image: Image,
123     pub input: DecodeInput,
124     pub codec_index: usize,
125     pub codec_config: CodecConfiguration,
126 }
127 
128 impl Tile {
create_from_item( item: &mut Item, allow_progressive: bool, image_count_limit: u32, size_hint: u64, ) -> AvifResult<Tile>129     pub fn create_from_item(
130         item: &mut Item,
131         allow_progressive: bool,
132         image_count_limit: u32,
133         size_hint: u64,
134     ) -> AvifResult<Tile> {
135         if size_hint != 0 && item.size as u64 > size_hint {
136             return Err(AvifError::BmffParseFailed("exceeded size_hint".into()));
137         }
138         let mut tile = Tile {
139             width: item.width,
140             height: item.height,
141             operating_point: item.operating_point(),
142             image: Image::default(),
143             codec_config: item
144                 .codec_config()
145                 .ok_or(AvifError::BmffParseFailed("missing av1C property".into()))?
146                 .clone(),
147             ..Tile::default()
148         };
149         let mut layer_sizes: [usize; MAX_AV1_LAYER_COUNT] = [0; MAX_AV1_LAYER_COUNT];
150         let mut layer_count: usize = 0;
151         let a1lx = item.a1lx();
152         let has_a1lx = a1lx.is_some();
153         if let Some(a1lx) = a1lx {
154             let mut remaining_size: usize = item.size;
155             for i in 0usize..3 {
156                 layer_count += 1;
157                 if a1lx[i] > 0 {
158                     // >= instead of > because there must be room for the last layer
159                     if a1lx[i] >= remaining_size {
160                         return Err(AvifError::BmffParseFailed(format!(
161                             "a1lx layer index [{i}] does not fit in item size"
162                         )));
163                     }
164                     layer_sizes[i] = a1lx[i];
165                     remaining_size -= a1lx[i];
166                 } else {
167                     layer_sizes[i] = remaining_size;
168                     remaining_size = 0;
169                     break;
170                 }
171             }
172             if remaining_size > 0 {
173                 assert!(layer_count == 3);
174                 layer_count += 1;
175                 layer_sizes[3] = remaining_size;
176             }
177         }
178         let lsel;
179         let has_lsel;
180         match item.lsel() {
181             Some(x) => {
182                 lsel = *x;
183                 has_lsel = true;
184             }
185             None => {
186                 lsel = 0;
187                 has_lsel = false;
188             }
189         }
190         // Progressive images offer layers via the a1lxProp, but don't specify a layer selection with
191         // lsel.
192         item.progressive = has_a1lx && (!has_lsel || lsel == 0xFFFF);
193         let base_item_offset = if item.extents.len() == 1 { item.extents[0].offset } else { 0 };
194         if has_lsel && lsel != 0xFFFF {
195             // Layer selection. This requires that the underlying AV1 codec decodes all layers, and
196             // then only returns the requested layer as a single frame. To the user of libavif,
197             // this appears to be a single frame.
198             tile.input.all_layers = true;
199             let mut sample_size: usize = 0;
200             let layer_id = usize_from_u16(lsel)?;
201             if layer_count > 0 {
202                 // Optimization: If we're selecting a layer that doesn't require the entire image's
203                 // payload (hinted via the a1lx box).
204                 if layer_id >= layer_count {
205                     return Err(AvifError::InvalidImageGrid(
206                         "lsel layer index not found in a1lx.".into(),
207                     ));
208                 }
209                 let layer_id_plus_1 = layer_id + 1;
210                 for layer_size in layer_sizes.iter().take(layer_id_plus_1) {
211                     checked_incr!(sample_size, *layer_size);
212                 }
213             } else {
214                 // This layer payload subsection is not known. Use the whole payload.
215                 sample_size = item.size;
216             }
217             let sample = DecodeSample {
218                 item_id: item.id,
219                 offset: base_item_offset,
220                 size: sample_size,
221                 spatial_id: lsel as u8,
222                 sync: true,
223             };
224             tile.input.samples.push(sample);
225         } else if item.progressive && allow_progressive {
226             // Progressive image. Decode all layers and expose them all to the
227             // user.
228             if image_count_limit != 0 && layer_count as u32 > image_count_limit {
229                 return Err(AvifError::BmffParseFailed(
230                     "exceeded image_count_limit (progressive)".into(),
231                 ));
232             }
233             tile.input.all_layers = true;
234             let mut offset = 0;
235             for (i, layer_size) in layer_sizes.iter().take(layer_count).enumerate() {
236                 let sample = DecodeSample {
237                     item_id: item.id,
238                     offset: checked_add!(base_item_offset, offset)?,
239                     size: *layer_size,
240                     spatial_id: 0xff,
241                     sync: i == 0, // Assume all layers depend on the first layer.
242                 };
243                 tile.input.samples.push(sample);
244                 offset = checked_add!(offset, *layer_size as u64)?;
245             }
246         } else {
247             // Typical case: Use the entire item's payload for a single frame output
248             let sample = DecodeSample {
249                 item_id: item.id,
250                 offset: base_item_offset,
251                 size: item.size,
252                 // Legal spatial_id values are [0,1,2,3], so this serves as a sentinel value for
253                 // "do not filter by spatial_id"
254                 spatial_id: 0xff,
255                 sync: true,
256             };
257             tile.input.samples.push(sample);
258         }
259         Ok(tile)
260     }
261 
create_from_track( track: &Track, mut image_count_limit: u32, size_hint: u64, category: Category, ) -> AvifResult<Tile>262     pub fn create_from_track(
263         track: &Track,
264         mut image_count_limit: u32,
265         size_hint: u64,
266         category: Category,
267     ) -> AvifResult<Tile> {
268         let mut tile = Tile {
269             width: track.width,
270             height: track.height,
271             operating_point: 0, // No way to set operating point via tracks
272             input: DecodeInput {
273                 category,
274                 ..DecodeInput::default()
275             },
276             ..Tile::default()
277         };
278         let sample_table = &track.sample_table.unwrap_ref();
279 
280         if image_count_limit != 0 {
281             for (chunk_index, _chunk_offset) in sample_table.chunk_offsets.iter().enumerate() {
282                 // Figure out how many samples are in this chunk.
283                 let sample_count = sample_table.get_sample_count_of_chunk(chunk_index as u32);
284                 if sample_count == 0 {
285                     return Err(AvifError::BmffParseFailed(
286                         "chunk with 0 samples found".into(),
287                     ));
288                 }
289                 if sample_count > image_count_limit {
290                     return Err(AvifError::BmffParseFailed(
291                         "exceeded image_count_limit".into(),
292                     ));
293                 }
294                 image_count_limit -= sample_count;
295             }
296         }
297 
298         let mut sample_size_index: usize = 0;
299         for (chunk_index, chunk_offset) in sample_table.chunk_offsets.iter().enumerate() {
300             // Figure out how many samples are in this chunk.
301             let sample_count = sample_table.get_sample_count_of_chunk(chunk_index as u32);
302             if sample_count == 0 {
303                 return Err(AvifError::BmffParseFailed(
304                     "chunk with 0 samples found".into(),
305                 ));
306             }
307 
308             let mut sample_offset = *chunk_offset;
309             for _ in 0..sample_count {
310                 let sample_size = sample_table.sample_size(sample_size_index)?;
311                 let sample_size_hint = checked_add!(sample_offset, sample_size as u64)?;
312                 if size_hint != 0 && sample_size_hint > size_hint {
313                     return Err(AvifError::BmffParseFailed("exceeded size_hint".into()));
314                 }
315                 let sample = DecodeSample {
316                     item_id: 0,
317                     offset: sample_offset,
318                     size: sample_size,
319                     // Legal spatial_id values are [0,1,2,3], so this serves as a sentinel value for "do
320                     // not filter by spatial_id"
321                     spatial_id: 0xff,
322                     // Assume first sample is always sync (in case stss box was missing).
323                     sync: tile.input.samples.is_empty(),
324                 };
325                 tile.input.samples.push(sample);
326                 checked_incr!(sample_offset, sample_size as u64);
327                 checked_incr!(sample_size_index, 1);
328             }
329         }
330         for sync_sample_number in &sample_table.sync_samples {
331             let index = usize_from_u32(*sync_sample_number)?;
332             // sample_table.sync_samples is 1-based.
333             if index == 0 || index > tile.input.samples.len() {
334                 return Err(AvifError::BmffParseFailed(format!(
335                     "invalid sync sample number {}",
336                     index
337                 )));
338             }
339             tile.input.samples[index - 1].sync = true;
340         }
341         Ok(tile)
342     }
343 
max_sample_size(&self) -> usize344     pub fn max_sample_size(&self) -> usize {
345         match self.input.samples.iter().max_by_key(|sample| sample.size) {
346             Some(sample) => sample.size,
347             None => 0,
348         }
349     }
350 }
351