1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145

use std::io::Read;

use jxl_bitstream::{Bitstream, Bundle, unpack_signed};
use jxl_image::ImageHeader;

use crate::Result;

#[derive(Debug)]
pub struct Patches {
    pub patches: Vec<PatchRef>,
}

#[derive(Debug)]
pub struct PatchRef {
    pub ref_idx: u32,
    pub x0: u32,
    pub y0: u32,
    pub width: u32,
    pub height: u32,
    pub patch_targets: Vec<PatchTarget>,
}

#[derive(Debug)]
pub struct PatchTarget {
    pub x: i32,
    pub y: i32,
    pub blending: Vec<BlendingModeInformation>,
}

#[derive(Debug)]
pub struct BlendingModeInformation {
    pub mode: PatchBlendMode,
    pub alpha_channel: u32,
    pub clamp: bool,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[repr(u8)]
pub enum PatchBlendMode {
    None = 0,
    Replace,
    Add,
    Mul,
    BlendAbove,
    BlendBelow,
    MulAddAbove,
    MulAddBelow,
}

impl TryFrom<u32> for PatchBlendMode {
    type Error = jxl_bitstream::Error;

    fn try_from(value: u32) -> std::result::Result<Self, Self::Error> {
        use PatchBlendMode::*;

        Ok(match value {
            0 => PatchBlendMode::None,
            1 => Replace,
            2 => Add,
            3 => Mul,
            4 => BlendAbove,
            5 => BlendBelow,
            6 => MulAddAbove,
            7 => MulAddBelow,
            _ => return Err(jxl_bitstream::Error::InvalidEnum { name: "PatchBlendMode", value }),
        })
    }
}

impl PatchBlendMode {
    #[inline]
    pub fn use_alpha(self) -> bool {
        matches!(self, Self::BlendAbove | Self::BlendBelow | Self::MulAddAbove | Self::MulAddBelow)
    }
}

impl Bundle<&ImageHeader> for Patches {
    type Error = crate::Error;

    fn parse<R: Read>(bitstream: &mut Bitstream<R>, image_header: &ImageHeader) -> Result<Self> {
        let num_extra = image_header.metadata.ec_info.len();
        let alpha_channel_indices = image_header.metadata.ec_info.iter()
            .enumerate()
            .filter_map(|(idx, info)| info.is_alpha().then_some(idx as u32))
            .collect::<Vec<_>>();

        let mut decoder = jxl_coding::Decoder::parse(bitstream, 10)?;
        decoder.begin(bitstream)?;

        let num_patches = decoder.read_varint(bitstream, 0)?;
        let patches = std::iter::repeat_with(|| -> Result<_> {
            let ref_idx = decoder.read_varint(bitstream, 1)?;
            let x0 = decoder.read_varint(bitstream, 3)?;
            let y0 = decoder.read_varint(bitstream, 3)?;
            let width = decoder.read_varint(bitstream, 2)? + 1;
            let height = decoder.read_varint(bitstream, 2)? + 1;
            let count = decoder.read_varint(bitstream, 7)? + 1;

            let mut prev_xy = None;
            let patch_targets = std::iter::repeat_with(|| -> Result<_> {
                let (x, y) = if let Some((px, py)) = prev_xy {
                    let dx = decoder.read_varint(bitstream, 6)?;
                    let dy = decoder.read_varint(bitstream, 6)?;
                    let dx = unpack_signed(dx);
                    let dy = unpack_signed(dy);
                    (dx + px, dy + py)
                } else {
                    (decoder.read_varint(bitstream, 4)? as i32, decoder.read_varint(bitstream, 4)? as i32)
                };
                prev_xy = Some((x, y));

                let blending = std::iter::repeat_with(|| -> Result<_> {
                    let raw_mode = decoder.read_varint(bitstream, 5)?;
                    let mode = PatchBlendMode::try_from(raw_mode)?;
                    let alpha_channel = if raw_mode >= 4 && alpha_channel_indices.len() >= 2 {
                        decoder.read_varint(bitstream, 8)?
                    } else {
                        alpha_channel_indices.first().copied().unwrap_or_default()
                    };
                    let clamp = if raw_mode >= 3 {
                        decoder.read_varint(bitstream, 9)? != 0
                    } else {
                        false
                    };

                    Ok(BlendingModeInformation { mode, alpha_channel, clamp })
                }).take(num_extra + 1).collect::<Result<Vec<_>>>()?;

                Ok(PatchTarget { x, y, blending })
            }).take(count as usize).collect::<Result<Vec<_>>>()?;

            Ok(PatchRef {
                ref_idx,
                x0,
                y0,
                width,
                height,
                patch_targets,
            })
        }).take(num_patches as usize).collect::<Result<Vec<_>>>()?;

        decoder.finalize()?;
        Ok(Self { patches })
    }
}