tiff-reader 0.5.0

//! Strip-based data access for TIFF images.

use std::sync::Arc;

#[cfg(feature = "rayon")]
use rayon::prelude::*;

use crate::block_decode;
use crate::cache::{BlockCache, BlockKey, BlockKind};
use crate::error::{Error, Result};
use crate::header::ByteOrder;
use crate::ifd::{Ifd, RasterLayout};
use crate::source::TiffSource;
use crate::{read_gdal_block_payload, GdalStructuralMetadata, Window};

const TAG_JPEG_TABLES: u16 = 347;

pub(crate) fn read_window(
    source: &dyn TiffSource,
    ifd: &Ifd,
    byte_order: ByteOrder,
    cache: &BlockCache,
    window: Window,
    gdal_structural_metadata: Option<&GdalStructuralMetadata>,
) -> Result<Vec<u8>> {
    let layout = ifd.raster_layout()?;
    if window.is_empty() {
        return Ok(Vec::new());
    }

    let output_len = window.output_len(&layout)?;
    let mut output = vec![0u8; output_len];
    let window_row_end = window.row_end();
    let output_row_bytes = window.cols * layout.pixel_stride_bytes();

    let relevant_specs = collect_strip_specs_for_window(ifd, &layout, window, None)?;

    #[cfg(not(feature = "rayon"))]
    let decoded_blocks: Result<Vec<_>> = relevant_specs
        .iter()
        .map(|&spec| {
            read_strip_block(
                source,
                ifd,
                byte_order,
                cache,
                spec,
                &layout,
                gdal_structural_metadata,
            )
            .map(|block| (spec, block))
        })
        .collect();

    #[cfg(feature = "rayon")]
    let decoded_blocks: Result<Vec<_>> = relevant_specs
        .par_iter()
        .map(|&spec| {
            read_strip_block(
                source,
                ifd,
                byte_order,
                cache,
                spec,
                &layout,
                gdal_structural_metadata,
            )
            .map(|block| (spec, block))
        })
        .collect();

    for (spec, block) in decoded_blocks? {
        let block = &*block;
        let block_row_end = spec.row_start + spec.rows_in_strip;
        let copy_row_start = spec.row_start.max(window.row_off);
        let copy_row_end = block_row_end.min(window_row_end);

        if layout.planar_configuration == 1 {
            let src_row_bytes = layout.row_bytes();
            let copy_bytes_per_row = window.cols * layout.pixel_stride_bytes();
            for row in copy_row_start..copy_row_end {
                let src_row_index = row - spec.row_start;
                let dest_row_index = row - window.row_off;
                let src_offset =
                    src_row_index * src_row_bytes + window.col_off * layout.pixel_stride_bytes();
                let dest_offset = dest_row_index * output_row_bytes;
                output[dest_offset..dest_offset + copy_bytes_per_row]
                    .copy_from_slice(&block[src_offset..src_offset + copy_bytes_per_row]);
            }
        } else {
            let src_row_bytes = layout.sample_plane_row_bytes();
            for row in copy_row_start..copy_row_end {
                let src_row_index = row - spec.row_start;
                let dest_row_index = row - window.row_off;
                let src_row =
                    &block[src_row_index * src_row_bytes..(src_row_index + 1) * src_row_bytes];
                let dest_row = &mut output
                    [dest_row_index * output_row_bytes..(dest_row_index + 1) * output_row_bytes];
                for col in window.col_off..window.col_end() {
                    let src = &src_row
                        [col * layout.bytes_per_sample..(col + 1) * layout.bytes_per_sample];
                    let dest_col_index = col - window.col_off;
                    let pixel_base = dest_col_index * layout.pixel_stride_bytes()
                        + spec.plane * layout.bytes_per_sample;
                    dest_row[pixel_base..pixel_base + layout.bytes_per_sample].copy_from_slice(src);
                }
            }
        }
    }

    Ok(output)
}

pub(crate) fn read_window_band(
    source: &dyn TiffSource,
    ifd: &Ifd,
    byte_order: ByteOrder,
    cache: &BlockCache,
    window: Window,
    band_index: usize,
    gdal_structural_metadata: Option<&GdalStructuralMetadata>,
) -> Result<Vec<u8>> {
    let layout = ifd.raster_layout()?;
    if band_index >= layout.samples_per_pixel {
        return Err(Error::BandIndexOutOfBounds {
            index: band_index,
            band_count: layout.samples_per_pixel,
        });
    }
    if window.is_empty() {
        return Ok(Vec::new());
    }

    let output_len = window.band_output_len(&layout)?;
    let mut output = vec![0u8; output_len];
    let window_row_end = window.row_end();
    let output_row_bytes = window.cols * layout.bytes_per_sample;

    let relevant_specs = collect_strip_specs_for_window(ifd, &layout, window, Some(band_index))?;

    #[cfg(not(feature = "rayon"))]
    let decoded_blocks: Result<Vec<_>> = relevant_specs
        .iter()
        .map(|&spec| {
            read_strip_block(
                source,
                ifd,
                byte_order,
                cache,
                spec,
                &layout,
                gdal_structural_metadata,
            )
            .map(|block| (spec, block))
        })
        .collect();

    #[cfg(feature = "rayon")]
    let decoded_blocks: Result<Vec<_>> = relevant_specs
        .par_iter()
        .map(|&spec| {
            read_strip_block(
                source,
                ifd,
                byte_order,
                cache,
                spec,
                &layout,
                gdal_structural_metadata,
            )
            .map(|block| (spec, block))
        })
        .collect();

    for (spec, block) in decoded_blocks? {
        let block = &*block;
        let block_row_end = spec.row_start + spec.rows_in_strip;
        let copy_row_start = spec.row_start.max(window.row_off);
        let copy_row_end = block_row_end.min(window_row_end);

        if layout.planar_configuration == 1 {
            let src_row_bytes = layout.row_bytes();
            let band_offset = band_index * layout.bytes_per_sample;
            for row in copy_row_start..copy_row_end {
                let src_row_index = row - spec.row_start;
                let dest_row_index = row - window.row_off;
                let src_row =
                    &block[src_row_index * src_row_bytes..(src_row_index + 1) * src_row_bytes];
                let dest_row = &mut output
                    [dest_row_index * output_row_bytes..(dest_row_index + 1) * output_row_bytes];
                for col in window.col_off..window.col_end() {
                    let src_base = col * layout.pixel_stride_bytes() + band_offset;
                    let dest_col_index = col - window.col_off;
                    let dest_base = dest_col_index * layout.bytes_per_sample;
                    dest_row[dest_base..dest_base + layout.bytes_per_sample]
                        .copy_from_slice(&src_row[src_base..src_base + layout.bytes_per_sample]);
                }
            }
        } else {
            let src_row_bytes = layout.sample_plane_row_bytes();
            let copy_bytes_per_row = window.cols * layout.bytes_per_sample;
            for row in copy_row_start..copy_row_end {
                let src_row_index = row - spec.row_start;
                let dest_row_index = row - window.row_off;
                let src_offset =
                    src_row_index * src_row_bytes + window.col_off * layout.bytes_per_sample;
                let dest_offset = dest_row_index * output_row_bytes;
                output[dest_offset..dest_offset + copy_bytes_per_row]
                    .copy_from_slice(&block[src_offset..src_offset + copy_bytes_per_row]);
            }
        }
    }

    Ok(output)
}

fn collect_strip_specs_for_window(
    ifd: &Ifd,
    layout: &RasterLayout,
    window: Window,
    band_index: Option<usize>,
) -> Result<Vec<StripBlockSpec>> {
    let offsets = ifd
        .strip_offsets()
        .ok_or(Error::TagNotFound(crate::ifd::TAG_STRIP_OFFSETS))?;
    let counts = ifd
        .strip_byte_counts()
        .ok_or(Error::TagNotFound(crate::ifd::TAG_STRIP_BYTE_COUNTS))?;
    if offsets.len() != counts.len() {
        return Err(Error::InvalidImageLayout(format!(
            "StripOffsets has {} entries but StripByteCounts has {}",
            offsets.len(),
            counts.len()
        )));
    }

    let rows_per_strip = ifd.rows_per_strip().unwrap_or(ifd.height());
    if rows_per_strip == 0 {
        return Err(Error::InvalidImageLayout(
            "RowsPerStrip must be greater than zero".into(),
        ));
    }
    let rows_per_strip = rows_per_strip as usize;
    let strips_per_plane = layout.height.div_ceil(rows_per_strip);
    let expected = match layout.planar_configuration {
        1 => strips_per_plane,
        2 => strips_per_plane * layout.samples_per_pixel,
        planar => return Err(Error::UnsupportedPlanarConfiguration(planar)),
    };
    if offsets.len() != expected {
        return Err(Error::InvalidImageLayout(format!(
            "expected {expected} strips, found {}",
            offsets.len()
        )));
    }

    let first_strip = window.row_off / rows_per_strip;
    let last_strip = window
        .row_end()
        .div_ceil(rows_per_strip)
        .min(strips_per_plane);
    let plane_range = if layout.planar_configuration == 1 {
        0..1
    } else if let Some(band_index) = band_index {
        band_index..band_index + 1
    } else {
        0..layout.samples_per_pixel
    };
    let spec_count = (last_strip - first_strip).saturating_mul(plane_range.end - plane_range.start);
    let mut specs = Vec::with_capacity(spec_count);

    for plane in plane_range {
        for plane_strip_index in first_strip..last_strip {
            let strip_index = if layout.planar_configuration == 1 {
                plane_strip_index
            } else {
                plane * strips_per_plane + plane_strip_index
            };
            let row_start = plane_strip_index * rows_per_strip;
            let rows_in_strip = rows_per_strip.min(layout.height.saturating_sub(row_start));
            specs.push(StripBlockSpec {
                index: strip_index,
                plane,
                row_start,
                offset: offsets[strip_index],
                byte_count: counts[strip_index],
                rows_in_strip,
            });
        }
    }

    Ok(specs)
}

#[derive(Clone, Copy)]
struct StripBlockSpec {
    index: usize,
    plane: usize,
    row_start: usize,
    offset: u64,
    byte_count: u64,
    rows_in_strip: usize,
}

fn read_strip_block(
    source: &dyn TiffSource,
    ifd: &Ifd,
    byte_order: ByteOrder,
    cache: &BlockCache,
    spec: StripBlockSpec,
    layout: &RasterLayout,
    gdal_structural_metadata: Option<&GdalStructuralMetadata>,
) -> Result<Arc<Vec<u8>>> {
    let cache_key = BlockKey {
        ifd_index: ifd.index,
        kind: BlockKind::Strip,
        block_index: spec.index,
    };
    if let Some(cached) = cache.get(&cache_key) {
        return Ok(cached);
    }

    let compressed = if gdal_structural_metadata.is_some() {
        Vec::new()
    } else if let Some(bytes) = source.as_slice() {
        let start = usize::try_from(spec.offset).map_err(|_| Error::OffsetOutOfBounds {
            offset: spec.offset,
            length: spec.byte_count,
            data_len: bytes.len() as u64,
        })?;
        let len = usize::try_from(spec.byte_count).map_err(|_| Error::OffsetOutOfBounds {
            offset: spec.offset,
            length: spec.byte_count,
            data_len: bytes.len() as u64,
        })?;
        let end = start.checked_add(len).ok_or(Error::OffsetOutOfBounds {
            offset: spec.offset,
            length: spec.byte_count,
            data_len: bytes.len() as u64,
        })?;
        if end > bytes.len() {
            return Err(Error::OffsetOutOfBounds {
                offset: spec.offset,
                length: spec.byte_count,
                data_len: bytes.len() as u64,
            });
        }
        bytes[start..end].to_vec()
    } else {
        let len = usize::try_from(spec.byte_count).map_err(|_| Error::OffsetOutOfBounds {
            offset: spec.offset,
            length: spec.byte_count,
            data_len: source.len(),
        })?;
        source.read_exact_at(spec.offset, len)?
    };

    let compressed = match gdal_structural_metadata {
        Some(metadata) => {
            read_gdal_block_payload(source, metadata, byte_order, spec.offset, spec.byte_count)?
        }
        None => compressed,
    };

    let jpeg_tables = ifd
        .tag(TAG_JPEG_TABLES)
        .and_then(|tag| tag.value.as_bytes());
    let decoded = block_decode::decode_compressed_block(block_decode::BlockDecodeRequest {
        ifd,
        layout: *layout,
        byte_order,
        compressed: &compressed,
        index: spec.index,
        jpeg_tables,
        block_width: layout.width,
        block_height: spec.rows_in_strip,
    })?;
    Ok(cache.insert(cache_key, decoded))
}