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
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
#[macro_export]
macro_rules! expand_u32 {
    ($bitstream:ident; $($rest:tt)*) => {
        $bitstream.read_bits(2)
            .and_then(|selector| $crate::expand_u32!(@expand $bitstream, selector, 0; $($rest)*,))
    };
    (@expand $bitstream:ident, $selector:ident, $counter:expr;) => {
        unreachable!()
    };
    (@expand $bitstream:ident, $selector:ident, $counter:expr; $c:literal, $($rest:tt)*) => {
        if $selector == $counter {
            $crate::read_bits!($bitstream, $c)
        } else {
            $crate::expand_u32!(@expand $bitstream, $selector, $counter + 1; $($rest)*)
        }
    };
    (@expand $bitstream:ident, $selector:ident, $counter:expr; u($n:literal), $($rest:tt)*) => {
        if $selector == $counter {
            $crate::read_bits!($bitstream, u($n))
        } else {
            $crate::expand_u32!(@expand $bitstream, $selector, $counter + 1; $($rest)*)
        }
    };
    (@expand $bitstream:ident, $selector:ident, $counter:expr; $c:literal + u($n:literal), $($rest:tt)*) => {
        if $selector == $counter {
            $crate::read_bits!($bitstream, $c + u($n))
        } else {
            $crate::expand_u32!(@expand $bitstream, $selector, $counter + 1; $($rest)*)
        }
    };
}

#[macro_export]
macro_rules! read_bits {
    ($bistream:ident, $c:literal $(, $ctx:expr)?) => {
        $crate::Result::Ok($c)
    };
    ($bitstream:ident, u($n:literal) $(, $ctx:expr)?) => {
        $bitstream.read_bits($n)
    };
    ($bitstream:ident, u($n:literal); UnpackSigned $(, $ctx:expr)?) => {
        $bitstream.read_bits($n).map($crate::unpack_signed)
    };
    ($bitstream:ident, $c:literal + u($n:literal) $(, $ctx:expr)?) => {
        $bitstream.read_bits($n).map(|v| v.wrapping_add($c))
    };
    ($bitstream:ident, $c:literal + u($n:literal); UnpackSigned $(, $ctx:expr)?) => {
        $bitstream.read_bits($n).map(|v| $crate::unpack_signed(v.wrapping_add($c)))
    };
    ($bitstream:ident, U32($($args:tt)+) $(, $ctx:expr)?) => {
        $crate::expand_u32!($bitstream; $($args)+)
    };
    ($bitstream:ident, U32($($args:tt)+); UnpackSigned $(, $ctx:expr)?) => {
        $crate::expand_u32!($bitstream; $($args)+).map($crate::unpack_signed)
    };
    ($bitstream:ident, U64 $(, $ctx:expr)?) => {
        $bitstream.read_u64()
    };
    ($bitstream:ident, U64; UnpackSigned $(, $ctx:expr)?) => {
        read_bits!($bitstream, U64 $(, $ctx)?).map($crate::unpack_signed_u64)
    };
    ($bitstream:ident, F16 $(, $ctx:expr)?) => {
        $bitstream.read_f16_as_f32()
    };
    ($bitstream:ident, Bool $(, $ctx:expr)?) => {
        $bitstream.read_bool()
    };
    ($bitstream:ident, Enum($enumtype:ty) $(, $ctx:expr)?) => {
        $crate::read_bits!($bitstream, U32(0, 1, 2 + u(4), 18 + u(6)))
            .and_then(|v| {
                <$enumtype as TryFrom<u32>>::try_from(v).map_err(|_| $crate::Error::InvalidEnum {
                    name: stringify!($enumtype),
                    value: v,
                })
            })
    };
    ($bitstream:ident, ZeroPadToByte $(, $ctx:expr)?) => {
        $bitstream.zero_pad_to_byte()
    };
    ($bitstream:ident, Bundle($bundle:ty)) => {
        $bitstream.read_bundle::<$bundle>()
    };
    ($bitstream:ident, Bundle($bundle:ty), $ctx:expr) => {
        $bitstream.read_bundle_with_ctx::<$bundle, _>($ctx)
    };
    ($bitstream:ident, Vec[$($inner:tt)*]; $count:expr $(, $ctx:expr)?) => {
        {
            let count = $count as usize;
            (0..count)
                .into_iter()
                .map(|_| $crate::read_bits!($bitstream, $($inner)* $(, $ctx)?))
                .collect::<::std::result::Result<Vec<_>, _>>()
        }
    };
    ($bitstream:ident, Array[$($inner:tt)*]; $count:expr $(, $ctx:expr)?) => {
        (|| -> ::std::result::Result<[_; $count], _> {
            let mut ret = [Default::default(); $count];
            for point in &mut ret {
                *point = match $crate::read_bits!($bitstream, $($inner)* $(, $ctx)?) {
                    ::std::result::Result::Ok(v) => v,
                    ::std::result::Result::Err(err) => return ::std::result::Result::Err(err),
                };
            }
            Ok(ret)
        })()
    };
}

#[macro_export]
macro_rules! make_def {
    (@ty; $c:literal) => { u32 };
    (@ty; u($n:literal)) => { u32 };
    (@ty; u($n:literal); UnpackSigned) => { i32 };
    (@ty; $c:literal + u($n:literal)) => { u32 };
    (@ty; $c:literal + u($n:literal); UnpackSigned) => { i32 };
    (@ty; U32($($args:tt)*)) => { u32 };
    (@ty; U32($($args:tt)*); UnpackSigned) => { i32 };
    (@ty; U64) => { u64 };
    (@ty; U64; UnpackSigned) => { i64 };
    (@ty; F16) => { f32 };
    (@ty; Bool) => { bool };
    (@ty; Enum($enum:ty)) => { $enum };
    (@ty; Bundle($bundle:ty)) => { $bundle };
    (@ty; Vec[$($inner:tt)*]; $count:expr) => { Vec<$crate::make_def!(@ty; $($inner)*)> };
    (@ty; Array[$($inner:tt)*]; $count:expr) => { [$crate::make_def!(@ty; $($inner)*); $count] };
    ($(#[$attrs:meta])* $v:vis struct $bundle_name:ident {
        $($(#[$fieldattrs:meta])* $vfield:vis $field:ident: ty($($expr:tt)*) $(ctx($ctx_for_field:expr))? $(cond($cond:expr))? $(default($def_expr:expr))? ,)*
    }) => {
        $(#[$attrs])*
        $v struct $bundle_name {
            $($(#[$fieldattrs])* $vfield $field: $crate::make_def!(@ty; $($expr)*),)*
        }
    };
}

#[macro_export]
macro_rules! make_parse {
    (@parse $bitstream:ident; cond($cond:expr); default($def_expr:expr); ty($($spec:tt)*); ctx($ctx:expr)) => {
        if $cond {
            $crate::read_bits!($bitstream, $($spec)*, $ctx)?
        } else {
            $def_expr
        }
    };
    (@parse $bitstream:ident; cond($cond:expr); ty($($spec:tt)*); ctx($ctx:expr)) => {
        if $cond {
            $crate::read_bits!($bitstream, $($spec)*, $ctx)?
        } else {
            $crate::BundleDefault::default_with_context($ctx)
        }
    };
    (@parse $bitstream:ident; $(default($def_expr:expr);)? ty($($spec:tt)*); ctx($ctx:expr)) => {
        $crate::read_bits!($bitstream, $($spec)*, $ctx)?
    };
    (@default; ; $ctx:expr) => {
        $crate::BundleDefault::default_with_context($ctx)
    };
    (@default; $def_expr:expr $(; $ctx:expr)?) => {
        $def_expr
    };
    (@select_ctx; $ctx_id:ident; $ctx:expr) => {
        $ctx
    };
    (@select_ctx; $ctx_id:ident;) => {
        $ctx_id
    };
    (@select_error_ty;) => {
        $crate::Error
    };
    (@select_error_ty; $err:ty) => {
        $err
    };
    ($bundle_name:ident $(error($err:ty))? {
        $($(#[$fieldattrs:meta])* $v:vis $field:ident: ty($($expr:tt)*) $(ctx($ctx_for_field:expr))? $(cond($cond:expr))? $(default($def_expr:expr))? ,)*
    }) => {
        impl<Ctx: Copy> $crate::Bundle<Ctx> for $bundle_name {
            type Error = $crate::make_parse!(@select_error_ty; $($err)?);

            #[allow(unused_variables)]
            fn parse(bitstream: &mut $crate::Bitstream, ctx: Ctx) -> ::std::result::Result<Self, Self::Error> where Self: Sized {
                $(
                    let $field: $crate::make_def!(@ty; $($expr)*) = $crate::make_parse!(
                        @parse bitstream;
                        $(cond($cond);)?
                        $(default($def_expr);)?
                        ty($($expr)*);
                        ctx($crate::make_parse!(@select_ctx; ctx; $($ctx_for_field)?))
                    );
                )*
                Ok(Self { $($field,)* })
            }
        }

        impl<Ctx: Copy> $crate::BundleDefault<Ctx> for $bundle_name {
            #[allow(unused_variables)]
            fn default_with_context(_ctx: Ctx) -> Self where Self: Sized {
                $(
                    let $field: $crate::make_def!(@ty; $($expr)*) = $crate::make_parse!(
                        @default;
                        $($def_expr)?;
                        $crate::make_parse!(@select_ctx; _ctx; $($ctx_for_field)?)
                    );
                )*
                Self { $($field,)* }
            }
        }
    };
    ($bundle_name:ident ctx($ctx_id:ident : $ctx:ty) $(error($err:ty))? {
        $($(#[$fieldattrs:meta])* $v:vis $field:ident: ty($($expr:tt)*) $(ctx($ctx_for_field:expr))? $(cond($cond:expr))? $(default($def_expr:expr))? ,)*
    }) => {
        impl $crate::Bundle<$ctx> for $bundle_name {
            type Error = $crate::make_parse!(@select_error_ty; $($err)?);

            #[allow(unused_variables)]
            fn parse(bitstream: &mut $crate::Bitstream, $ctx_id: $ctx) -> ::std::result::Result<Self, Self::Error> where Self: Sized {
                $(
                    let $field: $crate::make_def!(@ty; $($expr)*) = $crate::make_parse!(
                        @parse bitstream;
                        $(cond($cond);)?
                        $(default($def_expr);)?
                        ty($($expr)*);
                        ctx($crate::make_parse!(@select_ctx; $ctx_id; $($ctx_for_field)?))
                    );
                )*
                Ok(Self { $($field,)* })
            }
        }

        impl $crate::BundleDefault<$ctx> for $bundle_name {
            #[allow(unused_variables)]
            fn default_with_context($ctx_id: $ctx) -> Self where Self: Sized {
                $(
                    let $field: $crate::make_def!(@ty; $($expr)*) = $crate::make_parse!(
                        @default;
                        $($def_expr)?;
                        $crate::make_parse!(@select_ctx; $ctx_id; $($ctx_for_field)?)
                    );
                )*
                Self { $($field,)* }
            }
        }
    };
}

#[macro_export]
macro_rules! define_bundle {
    (
        $(
            $(#[$attrs:meta])*
            $v:vis struct $bundle_name:ident
            $(aligned($aligned:literal))?
            $(ctx($ctx_id:ident : $ctx:ty))?
            $(error($err:ty))?
            {
                $($body:tt)*
            }
        )*
    ) => {
        $(
            $crate::make_def!($(#[$attrs])* $v struct $bundle_name { $($body)* });
            $crate::make_parse!($bundle_name $(aligned($aligned))? $(ctx($ctx_id: $ctx))? $(error($err))? { $($body)* });
        )*
    };
}

/// Perform `UnpackSigned` for `u32`, as specified in the JPEG XL specification.
#[inline]
pub fn unpack_signed(x: u32) -> i32 {
    let bit = x & 1;
    let base = x >> 1;
    let flip = 0u32.wrapping_sub(bit);
    (base ^ flip) as i32
}

/// Perform `UnpackSigned` for `u64`, as specified in the JPEG XL specification.
#[inline]
pub fn unpack_signed_u64(x: u64) -> i64 {
    let bit = x & 1;
    let base = x >> 1;
    let flip = 0u64.wrapping_sub(bit);
    (base ^ flip) as i64
}