commit eb11db3022c46017b5a35206cba7a3d7f78a705e parent ac3269deac064c201be57dcc40a8571ce0efd997 Author: Kagami Sascha Rosylight <saschanaz@outlook.com> Date: Thu, 9 Oct 2025 09:06:56 +0000 Bug 1987386 - Vendor jxl-rs r=sylvestre,supply-chain-reviewers Differential Revision: https://phabricator.services.mozilla.com/D264092 Diffstat:
259 files changed, 41186 insertions(+), 0 deletions(-)
diff --git a/Cargo.lock b/Cargo.lock @@ -145,6 +145,12 @@ dependencies = [ ] [[package]] +name = "array-init" +version = "2.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3d62b7694a562cdf5a74227903507c56ab2cc8bdd1f781ed5cb4cf9c9f810bfc" + +[[package]] name = "arraydeque" version = "0.5.1" source = "registry+https://github.com/rust-lang/crates.io-index" @@ -2611,6 +2617,7 @@ dependencies = [ "gkrust_utils", "http_sfv", "idna_glue", + "image_jxl", "ipcclientcerts", "ipdl_utils", "jog", @@ -3391,6 +3398,13 @@ dependencies = [ ] [[package]] +name = "image_jxl" +version = "0.1.0" +dependencies = [ + "jxl", +] + +[[package]] name = "indexmap" version = "2.11.4" source = "registry+https://github.com/rust-lang/crates.io-index" @@ -3617,6 +3631,33 @@ dependencies = [ ] [[package]] +name = "jxl" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "710d04dc49b46298fa2258f3fd082bcafb318e93e45bfbffb2b2419d615964e7" +dependencies = [ + "array-init", + "byteorder", + "half 2.5.0", + "jxl_macros", + "num-derive", + "num-traits", + "thiserror 2.0.12", +] + +[[package]] +name = "jxl_macros" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5169e84285ee08cee04f115f2658cafba62e7ff54e8eb91a3842129ca12b003f" +dependencies = [ + "proc-macro-error2", + "proc-macro2", + "quote", + "syn", +] + +[[package]] name = "keccak" version = "0.1.4" source = "registry+https://github.com/rust-lang/crates.io-index" @@ -5539,6 +5580,28 @@ dependencies = [ ] [[package]] +name = "proc-macro-error-attr2" +version = "2.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "96de42df36bb9bba5542fe9f1a054b8cc87e172759a1868aa05c1f3acc89dfc5" +dependencies = [ + "proc-macro2", + "quote", +] + +[[package]] +name = "proc-macro-error2" +version = "2.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "11ec05c52be0a07b08061f7dd003e7d7092e0472bc731b4af7bb1ef876109802" +dependencies = [ + "proc-macro-error-attr2", + "proc-macro2", + "quote", + "syn", +] + +[[package]] name = "proc-macro2" version = "1.0.101" source = "registry+https://github.com/rust-lang/crates.io-index" diff --git a/image/rust/jxl/Cargo.toml b/image/rust/jxl/Cargo.toml @@ -0,0 +1,12 @@ +[package] +name = "image_jxl" +version = "0.1.0" +authors = [ + "Martin Bruse <zondolfin@gmail.com>", + "Kagami Sascha Rosylight <saschanaz@outlook.com>", +] +edition = "2021" +license = "MPL-2.0" + +[dependencies] +jxl = "0.1.1" diff --git a/image/rust/jxl/src/lib.rs b/image/rust/jxl/src/lib.rs @@ -0,0 +1,3 @@ +// This Source Code Form is subject to the terms of the Mozilla Public +// License, v. 2.0. If a copy of the MPL was not distributed with this +// file, You can obtain one at http://mozilla.org/MPL/2.0/. diff --git a/python/mozbuild/mozbuild/vendor/vendor_rust.py b/python/mozbuild/mozbuild/vendor/vendor_rust.py @@ -339,6 +339,8 @@ Please commit or stash these changes before vendoring, or re-run with `--ignore- "qlog", ], "BSD-3-Clause": [ + "jxl", + "jxl_macros", "subtle", "uritemplate-next", ], diff --git a/supply-chain/audits.toml b/supply-chain/audits.toml @@ -3569,6 +3569,12 @@ criteria = "safe-to-deploy" delta = "0.1.32 -> 0.1.33" notes = "No unsafe added, only non-trivial change is switching to the getrandom crate on Windows." +[[audits.jxl_macros]] +who = "Kagami Sascha Rosylight <saschanaz@outlook.com>" +criteria = "safe-to-deploy" +version = "0.1.1" +notes = "No unsafe block, no ambient capabilities like file system access (there's one but it's behind a test feature)" + [[audits.keccak]] who = "Simon Friedberger <simon@mozilla.com>" criteria = "safe-to-deploy" @@ -4803,6 +4809,18 @@ who = "Simon Friedberger <simon@mozilla.com>" criteria = "safe-to-deploy" version = "0.9.1" +[[audits.proc-macro-error-attr2]] +who = "Kagami Sascha Rosylight <saschanaz@outlook.com>" +criteria = "safe-to-deploy" +version = "2.0.0" +notes = "No unsafe block." + +[[audits.proc-macro-error2]] +who = "Kagami Sascha Rosylight <saschanaz@outlook.com>" +criteria = "safe-to-deploy" +version = "2.0.1" +notes = "No unsafe block with a lovely `#![forbid(unsafe_code)]`." + [[audits.proc-macro-hack]] who = "Mike Hommey <mh+mozilla@glandium.org>" criteria = "safe-to-deploy" diff --git a/supply-chain/imports.lock b/supply-chain/imports.lock @@ -1529,6 +1529,12 @@ who = "Johan Andersson <opensource@embark-studios.com>" criteria = "safe-to-deploy" version = "1.0.58" +[[audits.embark-studios.audits.array-init]] +who = "Johan Andersson <opensource@embark-studios.com>" +criteria = "safe-to-deploy" +version = "2.1.0" +notes = "Some unsafe usage but with safety comments and tests, and appear sound but didn't do detailed evaluation. No ambient capabilities" + [[audits.embark-studios.audits.cargo_metadata]] who = "Johan Andersson <opensource@embark-studios.com>" criteria = "safe-to-deploy" @@ -2145,6 +2151,13 @@ delta = "0.3.0 -> 0.4.0" notes = "No unsafe" aggregated-from = "https://chromium.googlesource.com/chromium/src/+/main/third_party/rust/chromium_crates_io/supply-chain/audits.toml?format=TEXT" +[[audits.google.audits.jxl]] +who = "Luca Versari <veluca@google.com>" +criteria = "safe-to-deploy" +version = "0.1.1" +notes = "Based on ub-risk-1 by joshlf@google.com and the lack of filesystem usage outside tests." +aggregated-from = "https://raw.githubusercontent.com/google/rust-crate-audits/main/manual-sources/additional-audits.toml" + [[audits.google.audits.litemap]] who = "Manish Goregaokar <manishearth@google.com>" criteria = "safe-to-deploy" diff --git a/third_party/rust/array-init/.cargo-checksum.json b/third_party/rust/array-init/.cargo-checksum.json @@ -0,0 +1 @@ +{"files":{"CHANGELOG.md":"78bdea34b6b568b97a3402d3a2aa601df4034fdedb33c1743f2027dae44ae20e","Cargo.toml":"dee4dc28ef0303e618c8a9321061edbd3f47b5c0bc4dd26f69d7824a6068ca57","LICENSE-APACHE":"c8d9a0d15dd76ca3bf277b6bf6da56799e266eac60bdc321a97ebc6d76d5153c","LICENSE-MIT":"e27fb2953c088c71285a4f2f54a0ac53323460ee7c2b1b838d563bd2687a38af","README.md":"f4146d1604179218bb51cd60d130fc30e0ae36ce04b3814f0faeb652835f33d9","src/lib.rs":"743ce846d76b8559ee1417ef218a2c4490eff5b02907872dd44074556a8b458d"},"package":"3d62b7694a562cdf5a74227903507c56ab2cc8bdd1f781ed5cb4cf9c9f810bfc"} +\ No newline at end of file diff --git a/third_party/rust/array-init/CHANGELOG.md b/third_party/rust/array-init/CHANGELOG.md @@ -0,0 +1,30 @@ +# Changelog + +All notable changes to this project will be documented in this file. + +The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), +and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). + + +## 2.1.0 +### Added +- Introduced an MSRV: Rust 1.51 +- Added `map_array_init` function ([#38](https://github.com/Manishearth/array-init/pull/38)) + +## v2.0.1 (2022-06-24) +### Added +- Added `from_iter_reversed` function ([#30](https://github.com/Manishearth/array-init/issues/30)) + +## v2.0.0 (2021-03-29) +### Breaking +- Removed `IsArray` trait (not necessary anymore with const generics) + +## v1.1.0 (yanked) +### Breaking +- Removed `const-generics` feature flag. The MSRV is now rust 1.51 + +## v1.0.0 (2020-10-14) +### Added + - Added a `try_array_init` function which initializes an array with a callable that may fail. + - Added a `const-generics` feature which uses rust (unstable) `const-generics` feature to implement the initializer functions for all array sizes. + - Added documentation diff --git a/third_party/rust/array-init/Cargo.toml b/third_party/rust/array-init/Cargo.toml @@ -0,0 +1,37 @@ +# THIS FILE IS AUTOMATICALLY GENERATED BY CARGO +# +# When uploading crates to the registry Cargo will automatically +# "normalize" Cargo.toml files for maximal compatibility +# with all versions of Cargo and also rewrite `path` dependencies +# to registry (e.g., crates.io) dependencies. +# +# If you are reading this file be aware that the original Cargo.toml +# will likely look very different (and much more reasonable). +# See Cargo.toml.orig for the original contents. + +[package] +edition = "2018" +name = "array-init" +version = "2.1.0" +authors = [ + "Manish Goregaokar <manishsmail@gmail.com>", + "Michal 'vorner' Vaner <vorner@vorner.cz>", +] +exclude = [".travis.yml"] +description = "Safe wrapper for initializing fixed-size arrays" +documentation = "https://docs.rs/crate/array-init" +readme = "README.md" +keywords = [ + "abstraction", + "initialization", + "no_std", +] +categories = [ + "data-structures", + "no-std", +] +license = "MIT OR Apache-2.0" +repository = "https://github.com/Manishearth/array-init/" + +[package.metadata] +msrv = "1.51" diff --git a/third_party/rust/array-init/LICENSE-APACHE b/third_party/rust/array-init/LICENSE-APACHE @@ -0,0 +1,201 @@ + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + +TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + +1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + +2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + +3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + +4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + +5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + +6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + +7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + +8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + +9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + +END OF TERMS AND CONDITIONS + +APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + +Copyright 2017-2020 The array-init developers + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. diff --git a/third_party/rust/array-init/LICENSE-MIT b/third_party/rust/array-init/LICENSE-MIT @@ -0,0 +1,25 @@ +Copyright (c) 2017-2020 The array-init developers + +Permission is hereby granted, free of charge, to any +person obtaining a copy of this software and associated +documentation files (the "Software"), to deal in the +Software without restriction, including without +limitation the rights to use, copy, modify, merge, +publish, distribute, sublicense, and/or sell copies of +the Software, and to permit persons to whom the Software +is furnished to do so, subject to the following +conditions: + +The above copyright notice and this permission notice +shall be included in all copies or substantial portions +of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF +ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED +TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A +PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT +SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY +CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION +OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR +IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER +DEALINGS IN THE SOFTWARE. diff --git a/third_party/rust/array-init/README.md b/third_party/rust/array-init/README.md @@ -0,0 +1,62 @@ +# array-init + +[](https://crates.io/crates/array-init) +[](https://docs.rs/array-init) + +The `array-init` crate allows you to initialize arrays +with an initializer closure that will be called +once for each element until the array is filled. + +This way you do not need to default-fill an array +before running initializers. Rust currently only +lets you either specify all initializers at once, +individually (`[a(), b(), c(), ...]`), or specify +one initializer for a `Copy` type (`[a(); N]`), +which will be called once with the result copied over. + +Care is taken not to leak memory shall the initialization +fail. + +## Examples: + +```rust +// Initialize an array of length 50 containing +// successive squares + +let arr: [usize; 50] = array_init::array_init(|i| i * i); + +// Initialize an array from an iterator +// producing an array of [1,2,3,4] repeated + +let four = [1,2,3,4]; +let mut iter = four.iter().copied().cycle(); +let arr: [u32; 50] = array_init::from_iter(iter).unwrap(); + +// Closures can also mutate state. We guarantee that they will be called +// in order from lower to higher indices. + +let mut last = 1u64; +let mut secondlast = 0; +let fibonacci: [u64; 50] = array_init::array_init(|_| { + let this = last + secondlast; + secondlast = last; + last = this; + this +}); +``` + +## Minimum Supported Rust Version (MSRV) + +`array-init` will only increase the MSRV on a new major +or minor release, but not for patch releases. +Any changes of the MSRV will be announced in the changelog. +When increasing the MSRV, the new Rust version must have been +released at least six months ago. The current MSRV is 1.51.0. +MSRV changes can be expected to happen conservatively. + +## Licensing + +Licensed under either of Apache License, Version 2.0 or MIT license at your option. + +Unless you explicitly state otherwise, any contribution intentionally submitted for inclusion in this crate by you, as defined in the Apache-2.0 license, shall be dual licensed as above, without any additional terms or conditions. + diff --git a/third_party/rust/array-init/src/lib.rs b/third_party/rust/array-init/src/lib.rs @@ -0,0 +1,490 @@ +#![no_std] + +//! The `array-init` crate allows you to initialize arrays +//! with an initializer closure that will be called +//! once for each element until the array is filled. +//! +//! This way you do not need to default-fill an array +//! before running initializers. Rust currently only +//! lets you either specify all initializers at once, +//! individually (`[a(), b(), c(), ...]`), or specify +//! one initializer for a `Copy` type (`[a(); N]`), +//! which will be called once with the result copied over. +//! +//! Care is taken not to leak memory shall the initialization +//! fail. +//! +//! # Examples: +//! ```rust +//! # #![allow(unused)] +//! # extern crate array_init; +//! # +//! // Initialize an array of length 50 containing +//! // successive squares +//! +//! let arr: [u32; 50] = array_init::array_init(|i: usize| (i * i) as u32); +//! +//! // Initialize an array from an iterator +//! // producing an array of [1,2,3,4] repeated +//! +//! let four = [1,2,3,4]; +//! let mut iter = four.iter().copied().cycle(); +//! let arr: [u32; 50] = array_init::from_iter(iter).unwrap(); +//! +//! // Closures can also mutate state. We guarantee that they will be called +//! // in order from lower to higher indices. +//! +//! let mut last = 1u64; +//! let mut secondlast = 0; +//! let fibonacci: [u64; 50] = array_init::array_init(|_| { +//! let this = last + secondlast; +//! secondlast = last; +//! last = this; +//! this +//! }); +//! ``` + +use ::core::{ + mem::{self, MaybeUninit}, + ptr, slice, +}; + +#[inline] +/// Initialize an array given an initializer expression. +/// +/// The initializer is given the index of the element. It is allowed +/// to mutate external state; we will always initialize the elements in order. +/// +/// # Examples +/// +/// ```rust +/// # #![allow(unused)] +/// # extern crate array_init; +/// # +/// // Initialize an array of length 50 containing +/// // successive squares +/// let arr: [usize; 50] = array_init::array_init(|i| i * i); +/// +/// assert!(arr.iter().enumerate().all(|(i, &x)| x == i * i)); +/// ``` +pub fn array_init<F, T, const N: usize>(mut initializer: F) -> [T; N] +where + F: FnMut(usize) -> T, +{ + enum Unreachable {} + + try_array_init( + // monomorphise into an infallible version + move |i| -> Result<T, Unreachable> { Ok(initializer(i)) }, + ) + .unwrap_or_else( + // zero-cost unwrap + |unreachable| match unreachable { /* ! */ }, + ) +} + +#[inline] +/// Initialize an array given an iterator +/// +/// We will iterate until the array is full or the iterator is exhausted. Returns +/// `None` if the iterator is exhausted before we can fill the array. +/// +/// - Once the array is full, extra elements from the iterator (if any) +/// won't be consumed. +/// +/// # Examples +/// +/// ```rust +/// # #![allow(unused)] +/// # extern crate array_init; +/// # +/// // Initialize an array from an iterator +/// // producing an array of [1,2,3,4] repeated +/// +/// let four = [1,2,3,4]; +/// let mut iter = four.iter().copied().cycle(); +/// let arr: [u32; 10] = array_init::from_iter(iter).unwrap(); +/// assert_eq!(arr, [1, 2, 3, 4, 1, 2, 3, 4, 1, 2]); +/// ``` +pub fn from_iter<Iterable, T, const N: usize>(iterable: Iterable) -> Option<[T; N]> +where + Iterable: IntoIterator<Item = T>, +{ + try_array_init_impl::<_, _, T, N, 1>({ + let mut iterator = iterable.into_iter(); + move |_| iterator.next().ok_or(()) + }) + .ok() +} + +#[inline] +/// Initialize an array in reverse given an iterator +/// +/// We will iterate until the array is full or the iterator is exhausted. Returns +/// `None` if the iterator is exhausted before we can fill the array. +/// +/// - Once the array is full, extra elements from the iterator (if any) +/// won't be consumed. +/// +/// # Examples +/// +/// ```rust +/// # #![allow(unused)] +/// # extern crate array_init; +/// # +/// // Initialize an array from an iterator +/// // producing an array of [4,3,2,1] repeated, finishing with 1. +/// +/// let four = [1,2,3,4]; +/// let mut iter = four.iter().copied().cycle(); +/// let arr: [u32; 10] = array_init::from_iter_reversed(iter).unwrap(); +/// assert_eq!(arr, [2, 1, 4, 3, 2, 1, 4, 3, 2, 1]); +/// ``` +pub fn from_iter_reversed<Iterable, T, const N: usize>(iterable: Iterable) -> Option<[T; N]> +where + Iterable: IntoIterator<Item = T>, +{ + try_array_init_impl::<_, _, T, N, -1>({ + let mut iterator = iterable.into_iter(); + move |_| iterator.next().ok_or(()) + }) + .ok() +} + +#[inline] +/// Initialize an array given an initializer expression that may fail. +/// +/// The initializer is given the index (between 0 and `N - 1` included) of the element, and returns a `Result<T, Err>,`. It is allowed +/// to mutate external state; we will always initialize from lower to higher indices. +/// +/// # Examples +/// +/// ```rust +/// # #![allow(unused)] +/// # extern crate array_init; +/// # +/// #[derive(PartialEq,Eq,Debug)] +/// struct DivideByZero; +/// +/// fn inv(i : usize) -> Result<f64,DivideByZero> { +/// if i == 0 { +/// Err(DivideByZero) +/// } else { +/// Ok(1./(i as f64)) +/// } +/// } +/// +/// // If the initializer does not fail, we get an initialized array +/// let arr: [f64; 3] = array_init::try_array_init(|i| inv(3-i)).unwrap(); +/// assert_eq!(arr,[1./3., 1./2., 1./1.]); +/// +/// // The initializer fails +/// let res : Result<[f64;4], DivideByZero> = array_init::try_array_init(|i| inv(3-i)); +/// assert_eq!(res,Err(DivideByZero)); +/// ``` +pub fn try_array_init<Err, F, T, const N: usize>(initializer: F) -> Result<[T; N], Err> +where + F: FnMut(usize) -> Result<T, Err>, +{ + try_array_init_impl::<Err, F, T, N, 1>(initializer) +} + +#[inline] +/// Initialize an array given a source array and a mapping expression. The size of the source array +/// is the same as the size of the returned array. +/// +/// The mapper is given an element from the source array and maps it to an element in the +/// destination. +/// +/// # Examples +/// +/// ```rust +/// # #![allow(unused)] +/// # extern crate array_init; +/// # +/// // Initialize an array of length 50 containing successive squares +/// let arr: [usize; 50] = array_init::array_init(|i| i * i); +/// +/// // Map each usize element to a u64 element. +/// let u64_arr: [u64; 50] = array_init::map_array_init(&arr, |element| *element as u64); +/// +/// assert!(u64_arr.iter().enumerate().all(|(i, &x)| x == (i * i) as u64)); +/// ``` +pub fn map_array_init<M, T, U, const N: usize>(source: &[U; N], mut mapper: M) -> [T; N] +where + M: FnMut(&U) -> T, +{ + // # Safety + // - The array size is known at compile time so we are certain that both the source and + // desitination have the same size. If the two arrays are of the same size we know that a + // valid index for one would be a valid index for the other. + array_init(|index| unsafe { mapper(source.get_unchecked(index)) }) +} + +#[inline] +fn try_array_init_impl<Err, F, T, const N: usize, const D: i8>( + mut initializer: F, +) -> Result<[T; N], Err> +where + F: FnMut(usize) -> Result<T, Err>, +{ + // The implementation differentiates two cases: + // A) `T` does not need to be dropped. Even if the initializer panics + // or returns `Err` we will not leak memory. + // B) `T` needs to be dropped. We must keep track of which elements have + // been initialized so far, and drop them if we encounter a panic or `Err` midway. + if !mem::needs_drop::<T>() { + let mut array: MaybeUninit<[T; N]> = MaybeUninit::uninit(); + // pointer to array = *mut [T; N] <-> *mut T = pointer to first element + let mut ptr_i = array.as_mut_ptr() as *mut T; + + // # Safety + // + // - for D > 0, we are within the array since we start from the + // beginning of the array, and we have `0 <= i < N`. + // - for D < 0, we start at the end of the array and go back one + // place before writing, going back N times in total, finishing + // at the start of the array. + unsafe { + if D < 0 { + ptr_i = ptr_i.add(N); + } + for i in 0..N { + let value_i = initializer(i)?; + // We overwrite *ptr_i previously undefined value without reading or dropping it. + if D < 0 { + ptr_i = ptr_i.sub(1); + } + ptr_i.write(value_i); + if D > 0 { + ptr_i = ptr_i.add(1); + } + } + Ok(array.assume_init()) + } + } else { + // else: `mem::needs_drop::<T>()` + + /// # Safety + /// + /// - `base_ptr[.. initialized_count]` must be a slice of init elements... + /// + /// - ... that must be sound to `ptr::drop_in_place` if/when + /// `UnsafeDropSliceGuard` is dropped: "symbolic ownership" + struct UnsafeDropSliceGuard<Item> { + base_ptr: *mut Item, + initialized_count: usize, + } + + impl<Item> Drop for UnsafeDropSliceGuard<Item> { + fn drop(self: &'_ mut Self) { + unsafe { + // # Safety + // + // - the contract of the struct guarantees that this is sound + ptr::drop_in_place(slice::from_raw_parts_mut( + self.base_ptr, + self.initialized_count, + )); + } + } + } + + // If the `initializer(i)` call panics, `panic_guard` is dropped, + // dropping `array[.. initialized_count]` => no memory leak! + // + // # Safety + // + // 1. - For D > 0, by construction, array[.. initiliazed_count] only + // contains init elements, thus there is no risk of dropping + // uninit data; + // - For D < 0, by construction, array[N - initialized_count..] only + // contains init elements. + // + // 2. - for D > 0, we are within the array since we start from the + // beginning of the array, and we have `0 <= i < N`. + // - for D < 0, we start at the end of the array and go back one + // place before writing, going back N times in total, finishing + // at the start of the array. + // + unsafe { + let mut array: MaybeUninit<[T; N]> = MaybeUninit::uninit(); + // pointer to array = *mut [T; N] <-> *mut T = pointer to first element + let mut ptr_i = array.as_mut_ptr() as *mut T; + if D < 0 { + ptr_i = ptr_i.add(N); + } + let mut panic_guard = UnsafeDropSliceGuard { + base_ptr: ptr_i, + initialized_count: 0, + }; + + for i in 0..N { + // Invariant: `i` elements have already been initialized + panic_guard.initialized_count = i; + // If this panics or fails, `panic_guard` is dropped, thus + // dropping the elements in `base_ptr[.. i]` for D > 0 or + // `base_ptr[N - i..]` for D < 0. + let value_i = initializer(i)?; + // this cannot panic + // the previously uninit value is overwritten without being read or dropped + if D < 0 { + ptr_i = ptr_i.sub(1); + panic_guard.base_ptr = ptr_i; + } + ptr_i.write(value_i); + if D > 0 { + ptr_i = ptr_i.add(1); + } + } + // From now on, the code can no longer `panic!`, let's take the + // symbolic ownership back + mem::forget(panic_guard); + + Ok(array.assume_init()) + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn seq() { + let seq: [usize; 5] = array_init(|i| i); + assert_eq!(&[0, 1, 2, 3, 4], &seq); + } + + #[test] + fn array_from_iter() { + let array = [0, 1, 2, 3, 4]; + let seq: [usize; 5] = from_iter(array.iter().copied()).unwrap(); + assert_eq!(array, seq,); + } + + #[test] + fn array_init_no_drop() { + DropChecker::with(|drop_checker| { + let result: Result<[_; 5], ()> = try_array_init(|i| { + if i < 3 { + Ok(drop_checker.new_element()) + } else { + Err(()) + } + }); + assert!(result.is_err()); + }); + } + + #[test] + fn from_iter_no_drop() { + DropChecker::with(|drop_checker| { + let iterator = (0..3).map(|_| drop_checker.new_element()); + let result: Option<[_; 5]> = from_iter(iterator); + assert!(result.is_none()); + }); + } + + #[test] + fn from_iter_reversed_no_drop() { + DropChecker::with(|drop_checker| { + let iterator = (0..3).map(|_| drop_checker.new_element()); + let result: Option<[_; 5]> = from_iter_reversed(iterator); + assert!(result.is_none()); + }); + } + + #[test] + fn test_513_seq() { + let seq: [usize; 513] = array_init(|i| i); + assert_eq!( + [ + 0, 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, 283, 284, 285, 286, 287, 288, 289, 290, 291, 292, 293, 294, 295, 296, + 297, 298, 299, 300, 301, 302, 303, 304, 305, 306, 307, 308, 309, 310, 311, 312, + 313, 314, 315, 316, 317, 318, 319, 320, 321, 322, 323, 324, 325, 326, 327, 328, + 329, 330, 331, 332, 333, 334, 335, 336, 337, 338, 339, 340, 341, 342, 343, 344, + 345, 346, 347, 348, 349, 350, 351, 352, 353, 354, 355, 356, 357, 358, 359, 360, + 361, 362, 363, 364, 365, 366, 367, 368, 369, 370, 371, 372, 373, 374, 375, 376, + 377, 378, 379, 380, 381, 382, 383, 384, 385, 386, 387, 388, 389, 390, 391, 392, + 393, 394, 395, 396, 397, 398, 399, 400, 401, 402, 403, 404, 405, 406, 407, 408, + 409, 410, 411, 412, 413, 414, 415, 416, 417, 418, 419, 420, 421, 422, 423, 424, + 425, 426, 427, 428, 429, 430, 431, 432, 433, 434, 435, 436, 437, 438, 439, 440, + 441, 442, 443, 444, 445, 446, 447, 448, 449, 450, 451, 452, 453, 454, 455, 456, + 457, 458, 459, 460, 461, 462, 463, 464, 465, 466, 467, 468, 469, 470, 471, 472, + 473, 474, 475, 476, 477, 478, 479, 480, 481, 482, 483, 484, 485, 486, 487, 488, + 489, 490, 491, 492, 493, 494, 495, 496, 497, 498, 499, 500, 501, 502, 503, 504, + 505, 506, 507, 508, 509, 510, 511, 512 + ], + seq + ); + } + + use self::drop_checker::DropChecker; + mod drop_checker { + use ::core::cell::Cell; + + pub(super) struct DropChecker { + slots: [Cell<bool>; 512], + next_uninit_slot: Cell<usize>, + } + + pub(super) struct Element<'drop_checker> { + slot: &'drop_checker Cell<bool>, + } + + impl Drop for Element<'_> { + fn drop(self: &'_ mut Self) { + assert!(self.slot.replace(false), "Double free!"); + } + } + + impl DropChecker { + pub(super) fn with(f: impl FnOnce(&Self)) { + let drop_checker = Self::new(); + f(&drop_checker); + drop_checker.assert_no_leaks(); + } + + pub(super) fn new_element(self: &'_ Self) -> Element<'_> { + let i = self.next_uninit_slot.get(); + self.next_uninit_slot.set(i + 1); + self.slots[i].set(true); + Element { + slot: &self.slots[i], + } + } + + fn new() -> Self { + Self { + slots: crate::array_init(|_| Cell::new(false)), + next_uninit_slot: Cell::new(0), + } + } + + fn assert_no_leaks(self: Self) { + let leak_count: usize = self.slots[..self.next_uninit_slot.get()] + .iter() + .map(|slot| usize::from(slot.get() as u8)) + .sum(); + assert_eq!(leak_count, 0, "{} elements leaked", leak_count); + } + } + } +} diff --git a/third_party/rust/jxl/.cargo-checksum.json b/third_party/rust/jxl/.cargo-checksum.json @@ -0,0 +1 @@ +{"files":{"Cargo.lock":"4122cabcc0c4babbf22ef69ce1a4a61da479a15a66a48199fd12ea9351e3a3b8","Cargo.toml":"4f6d33443dd9248d0734bdcdb25eb110f404942fa7f424ebca06f6a6cab37e8e","README.md":"3bb0a5cac8cebc1b90a75a94f981ba895cc7afff28ce57b4d88697ba9ee0a92a","resources/test/8x8_noise.jxl":"613e1142124cb8ce4b1673fa405762dce18c3d8223a86dc450edec15b196fdc0","resources/test/all.html":"03ed0f4d6e1f9782eeb5b01864d95678f5ec6dd0ce153189da41d81f2503ca90","resources/test/basic.jxl":"9af7cfa95117b19d5d08190ee2a1986bbddf60cdd0739bc4f4866754b3c80bcf","resources/test/candle.jxl":"c8ca9c3156e0d3a8cd2ff3f01dd8be42f65e83bfd7eef261864af4faf93aab8a","resources/test/candle_license.txt":"e851b9d168cd8b23848aa89247189d7e9bcad9cca27ca9fbe56666822b2f0f16","resources/test/conformance_test_images/alpha_nonpremultiplied.jxl":"15acbe3edbfd5a75c7609726ae60526ffc812642b5dd6be8475f0b990ce9b1db","resources/test/conformance_test_images/alpha_premultiplied.jxl":"5028e630dc358bf4031cbf06e8a80f2bfbfa21fd0d088af74d83b473d3dd540a","resources/test/conformance_test_images/alpha_triangles.jxl":"19ac7752a23ad2b22814064cb6b62a581b48be18ed73b5ccc2340888c114d2c9","resources/test/conformance_test_images/animation_icos4d.jxl":"6df934e61b07779626988daa0c4a9f05f89e6142a030f915e8b8ea52c37b3fe2","resources/test/conformance_test_images/animation_icos4d_5.jxl":"6df934e61b07779626988daa0c4a9f05f89e6142a030f915e8b8ea52c37b3fe2","resources/test/conformance_test_images/animation_newtons_cradle.jxl":"35a13d365f5eb88c3249afe073b7d8b0174d74e38f7edaa8805ca8ab46a59ab4","resources/test/conformance_test_images/animation_spline.jxl":"87793cac33d05eaa380011e3b0754ff6f228967431a126fd0a0ace2106940c79","resources/test/conformance_test_images/animation_spline_5.jxl":"87793cac33d05eaa380011e3b0754ff6f228967431a126fd0a0ace2106940c79","resources/test/conformance_test_images/bench_oriented_brg.jxl":"e223fed907c6238622b2b6ec1c80609050c9d2db4d759cf3ca6f0db304cbb82a","resources/test/conformance_test_images/bench_oriented_brg_5.jxl":"e223fed907c6238622b2b6ec1c80609050c9d2db4d759cf3ca6f0db304cbb82a","resources/test/conformance_test_images/bicycles.jxl":"8ef03ccccb45bdb5f3d37c9a4c392f745e3e4ad7f4ffc2cbc1e1ba68487a75a6","resources/test/conformance_test_images/bike.jxl":"639efe13b89868688234ebb9ec36a0ed848ba01bde43681d3d37438d2e0c0370","resources/test/conformance_test_images/bike_5.jxl":"639efe13b89868688234ebb9ec36a0ed848ba01bde43681d3d37438d2e0c0370","resources/test/conformance_test_images/blendmodes.jxl":"b22d487440dd591fbadd4e75398506070497283f38525433961b1e809c75c860","resources/test/conformance_test_images/blendmodes_5.jxl":"b22d487440dd591fbadd4e75398506070497283f38525433961b1e809c75c860","resources/test/conformance_test_images/cafe.jxl":"cb4603bd089483fc1d63a039a3b758ec510f0090b74ce245b88b1e8feb54f0eb","resources/test/conformance_test_images/cafe_5.jxl":"cb4603bd089483fc1d63a039a3b758ec510f0090b74ce245b88b1e8feb54f0eb","resources/test/conformance_test_images/cmyk_layers.jxl":"d732c8836bf1abeadf310d2e07387a32813ed4690d32650c1c25e541b80eed4a","resources/test/conformance_test_images/delta_palette.jxl":"00e24cc453cdf84897d62b0aafc7a9f7024205bbce1922e99d7ad0003759ae7c","resources/test/conformance_test_images/grayscale.jxl":"78fbbba852e99946d187dcf0bcbd7fb0e7c22be2f0852523aaae6ed91e7e3c39","resources/test/conformance_test_images/grayscale_5.jxl":"78fbbba852e99946d187dcf0bcbd7fb0e7c22be2f0852523aaae6ed91e7e3c39","resources/test/conformance_test_images/grayscale_jpeg.jxl":"346cc92b55864b0efb9f85cdddd02b586473ff18b03718e5a02fa03cb3e86844","resources/test/conformance_test_images/grayscale_jpeg_5.jxl":"346cc92b55864b0efb9f85cdddd02b586473ff18b03718e5a02fa03cb3e86844","resources/test/conformance_test_images/grayscale_public_university.jxl":"b1c6fe79167858f86a23e0f32c1c5ee7602dae6f82991c1dd53bcb89a7d78606","resources/test/conformance_test_images/lossless_pfm.jxl":"61ae52b5851ab2e156aec1d22502e8e1ec0cdf6bc0d6a3956ac6d7d6d2969d5e","resources/test/conformance_test_images/lz77_flower.jxl":"41046f0de4dd36177aefa9b53fe6c3cd35a8e2c215a8d41ca89a6a9bc7f27541","resources/test/conformance_test_images/noise.jxl":"20cddadc3fdbae331a912dba0d14e6326008094d7a16d006961ccdeac2bc9950","resources/test/conformance_test_images/noise_5.jxl":"20cddadc3fdbae331a912dba0d14e6326008094d7a16d006961ccdeac2bc9950","resources/test/conformance_test_images/opsin_inverse.jxl":"5c25ba126e6fb3536a12ca191b92f1c3be86ee74683688ed56c94804db5d2e16","resources/test/conformance_test_images/opsin_inverse_5.jxl":"5c25ba126e6fb3536a12ca191b92f1c3be86ee74683688ed56c94804db5d2e16","resources/test/conformance_test_images/patches.jxl":"3d8d3847ba40844b1d829366065261f34cc4b4acc3d9006342c64127c1052a3e","resources/test/conformance_test_images/patches_5.jxl":"3d8d3847ba40844b1d829366065261f34cc4b4acc3d9006342c64127c1052a3e","resources/test/conformance_test_images/progressive.jxl":"bb4722ab780dce609a7b007fe7fe235962b351691f60ffb2cde7b4808c8601e1","resources/test/conformance_test_images/progressive_5.jxl":"bb4722ab780dce609a7b007fe7fe235962b351691f60ffb2cde7b4808c8601e1","resources/test/conformance_test_images/spot.jxl":"69d48f7ace25683db9eb908f112cef112bd1583870abab0fbf166018889893fa","resources/test/conformance_test_images/sunset_logo.jxl":"6617480923e1fdef555e165a1e7df9ca648068dd0bdbc41a22c0e4213392d834","resources/test/conformance_test_images/upsampling.jxl":"d3d97aa5f91401ad6df137b647164e402c7a99d66f8a6f72734562bf844b69d2","resources/test/conformance_test_images/upsampling_5.jxl":"d3d97aa5f91401ad6df137b647164e402c7a99d66f8a6f72734562bf844b69d2","resources/test/cropped_traffic_light.jxl":"28429b199c41fb54eda40986cd80dc5ddcf93c6253b9ce8eec5017e3c9d5a419","resources/test/extra_channels.jxl":"2514fa63fb8daa2c673d94f8ece21e02c462a930d47ebea479a976c157ec5bb0","resources/test/grayscale_patches_modular.jxl":"fef4f3870d7f3c77104d81ad4b03e5970cae165650bd9ffa63fd1d8ed0a5c3f0","resources/test/grayscale_patches_var_dct.jxl":"64f64ace4b696b782b01210402a265bb030ef8f1e3bf9f41b236141b47dc90e8","resources/test/green_queen_modular_e3.jxl":"2397130dc451dcadce749a15af0bf6169f737d6588bb3fde6c2e333770b664f5","resources/test/green_queen_vardct_e3.jxl":"a8d66bad1caf61558568af95091429e0d20d4b9d72e8b64a1de172ae40d2d573","resources/test/has_permutation.jxl":"28fac5b7e4e5c040e025a6bf453815d50e6d3aad1c9cd12998cc94470893f9ef","resources/test/has_permutation_with_container.jxl":"933b84748b6707d910771de2ef9dcc34d737d4eb2521ba2ce8a810e58de2ce60","resources/test/multiple_lf_420.jxl":"50a232faefaccccc5a0617a0a5d59a2c233652c66fe818cae12a2acb61b02881","resources/test/orientation1_identity.jxl":"c15df202da4941c3a6157a94a6f80d071b0cea807160eaaa01b6220e85b8c590","resources/test/orientation2_flip_horizontal.jxl":"f0be126c96dedcded1ec9cacc81372be98da7578c85f2499896d8c30027bbb7e","resources/test/orientation3_rotate_180.jxl":"afb4079648e84c57d9283c08599e15af803c41e8b30c3df0424c0a099b943573","resources/test/orientation4_flip_vertical.jxl":"df4ab89d8e32a7ab88d4a6371ab308d8a5ff2eaa9cb038f961342af3333fe492","resources/test/orientation5_transpose.jxl":"dbec1f57131b1a84ae315337163d53f51a601317a5339beaad7c88b8fb825c85","resources/test/orientation6_rotate_90_cw.jxl":"6709e3ec050ef8dc0c6644e1184f33cf126a4c4f11e362e302124b59c040ca99","resources/test/orientation7_anti_transpose.jxl":"cccd67a2a771f6cbcb0a3213181676123443c18a9da13827cbdd0e09fe061b98","resources/test/orientation8_rotate_90_ccw.jxl":"68f3efc9c9c3ed6206e0d8b6044c3959d0b54b2378408542941e345ac667d38c","resources/test/pq_gradient.jxl":"b37c9cd69b62458a3d5fb4efc2efe9f6da8e434db4772d9e03cbbf0d50945193","resources/test/small_grayscale_patches_modular.jxl":"4f199e787b044bc0529728c759f6e9f4dbcc18b111fdf289260c1d96bb903742","resources/test/small_grayscale_patches_modular_with_icc.jxl":"d02c1e46169a0c3b0f69a89153f3ce5c321b85551154bfb9704c1ef8ba6320ba","resources/test/spline_on_first_frame.jxl":"70f753c0de4ccc28859b3aa5c6810030784c1b6989831c81b9f72d4ca9776193","resources/test/splines.jxl":"63d8c776c5f99f161ca0b832c3c5465d6dafff379ad95adc5cfed2d30251d8a0","resources/test/splines.pfm":"5eb25b45f187487a40ed9f00d8fec25fb2ac1e46c78802cf7ed5c610e6aa33b0","resources/test/squeeze_edge.jxl":"371e5da5bd20b4bad9df1945e788496183e6adda8a5c628dc69e968cd24e5811","resources/test/with_icc.jxl":"d102c703bfc4208748723656371381b7a7e33690e43806ec114f5c0b10b11b4b","resources/test/zoltan_tasi_unsplash.jxl":"b73170a9a1b783b822deae84067b1ea93dfef8f455435e0216ceadb000b22b13","src/api.rs":"88598cfa55ee205992d0ccbcab3ece7c32605a816f34ca1b1d341cc1d077825b","src/api/color.rs":"f01c5e809d958878c9cceba2de3cd202dd47460a35a9c385db5e93d8583227a7","src/api/data_types.rs":"b9bfa6382df3d474e73e6ef499b9fdf377ba20814a8ec813865ede7376caba70","src/api/decoder.rs":"fe8d5c3988eccb130fc06991fad4afb247c32824e04e1c2367688baa51c46083","src/api/inner.rs":"7370fc2a8a7d7fa1cff4a8a8684417eef472cec0b73d8932d483d351fde5fe80","src/api/inner/box_parser.rs":"4da3b1fc9b6dcc17933cd71d34808e80fcdfe39d8e6a7d9dca25873b0297c1fa","src/api/inner/codestream_parser.rs":"18e9ee597a987827dff4272fa64bd3cfa5bfc9a2f3ce79cf1749db8bc99fd2ef","src/api/inner/codestream_parser/non_section.rs":"88f7626ca1e1f8b35ea526370afa6c7119c2b9864f74e00743502543e0c0cc99","src/api/inner/codestream_parser/sections.rs":"75e4680a944d89ceedc7d077caa8b2e6bb5d451ccaf063ec8b697d82d3199d2b","src/api/inner/process.rs":"74f9546d5aa6c61df86e565fad73a2a73f054c14ac5b667c96042dfcc88ced1d","src/api/input.rs":"0ddf4639be5c6bb1ddfab04e3a61f0f2f2aa95b481abbfc9166596daad2467e5","src/api/options.rs":"7cb1ae523f9e81a2294c18c9882809f9b2089a70ea411848824934254b924947","src/api/output.rs":"1a4ecadc500956e653aaffc1ada9afa4683e1882048fcd601ddcff8e5a7fc934","src/api/signature.rs":"48b9e603ea082b08df7984abb9cd94ca332f690912e3dd1c3c362ff6480d0dd4","src/api/test.rs":"bae439d9995883c41bff0a02faa755514f4502c1b2bb335a6e3918f0a453c981","src/bit_reader.rs":"1386808b7bd8a7580127db184b3a707c8a6fd5fa22b2f46d5ca6082981f57976","src/color/mod.rs":"731c12b91c995da547d970a9f4f8ab5ffa477b4e82fbb23dbdd3921ca88cf985","src/color/tf.rs":"13695face0d158e9452863bd6315bdfb0994406a407f3b73ba8053d2696217ca","src/container/box_header.rs":"45f738a29794a179383b8ad547b0bd93a4d977ad253f6741a5a73758d0c79032","src/container/mod.rs":"e67172b6ae3cc1d22fd0721aa160efb047490e3dc8020d24fe5fae5b7f4cc7e6","src/container/parse.rs":"b7724b8f85e9626a849da9d8394f5c254e410c207d7e73341ca61376b2dc426e","src/entropy_coding/ans.rs":"289420e3eedeecde9ef421b0f4e1b7e286ac3fbbe342b1f3ebe811d34de692b8","src/entropy_coding/context_map.rs":"12ea755842a6e552a7252b15b8f208b72222a16546880c1a99a15b2bae9feaa7","src/entropy_coding/decode.rs":"47e9e651fb57ef8e47632bf0dbd424805de672eb6a31d2477618dbd305aa4b1d","src/entropy_coding/huffman.rs":"c89bdbb0904ce7a5ac4949be9ec42db9ffffe76ffb3ac7d586763d1e48314226","src/entropy_coding/hybrid_uint.rs":"3303c6919fc549a9cf67b5351dba68542a3f2f9e14579d9bf42e73d5bb637b3a","src/entropy_coding/mod.rs":"a3b56042f4c0c2837ff0e14eafe08bba0b5f472923b0ef6623f3024a01babdf7","src/error.rs":"437c8c9feab45f0878ff42b43670e0f9389ab44656ab2bf3e162f7020ec149f5","src/features/blending.rs":"ed5a50c5ed03902dfeafb869efeef11f247082e8f615e96a79a7347fda9c213d","src/features/epf.rs":"9944ab6be576756dc59b868a3ef409c5287bc97c8891695493cef01669aa702f","src/features/mod.rs":"cc50b40b44542465baf73e15450564a0b585849d1a4a7e99e7bb717387040a74","src/features/noise.rs":"655b9ad502d6584ea02f10c9e504c841b6a250166afe852a6c787cea7d7ca0e4","src/features/patches.rs":"c20fee43a4f2eecc8532a4eddee12e17b910b4cfc49d4e74a01d65612c6cc2d0","src/features/spline.rs":"493fae2581a86621b6875d6f5385d7114e451cf1cac96b61ffd09ff1aef2cdff","src/frame.rs":"5848226142bb3e71353784b5aeaea4e6cd552831f156fe96327d712a844e6929","src/frame/adaptive_lf_smoothing.rs":"412f1f916545d4a5da95183cbe3621070153fdefdfce7fcb1ec3ba3df2005915","src/frame/block_context_map.rs":"27912adabf61cd009710e8fc4c5ac4ee972a4f2fc3498e0f6e1012ef113cb7dd","src/frame/coeff_order.rs":"9c2dfde17f6db9dd4b2d51c9b5cb7dd7276e6e08dad9e903d6ee9f09fbd67a67","src/frame/color_correlation_map.rs":"126bcdbb438381adc9942fb2555dad7c3a205c010d55c781ae645c45bcbd1561","src/frame/group.rs":"613955a6602e941bc58ce5b3dea3312d272ee06a82f41a22650c3ce75ca5d6c7","src/frame/modular.rs":"bf72f0e0949f647955dbc8606895c25341391280f4e1b6733166863cb3e569e2","src/frame/modular/borrowed_buffers.rs":"dd2ee7baa005ce77e3329fe92b8cb191adb8cddd738fc2232930f68b1e74f828","src/frame/modular/decode.rs":"2a7981fd111ab26d8829eb36c85d6deed797905253ae92e7ccba8fe367374ea3","src/frame/modular/predict.rs":"66ca1da729bd5029205632221dc833292cdda8a600ca7af73ddad6cdc913a6c7","src/frame/modular/transforms.rs":"e46ace89a4db1c01f8ef71126591a9bbed2df715c1b22649331a9a9837374a8e","src/frame/modular/transforms/apply.rs":"a6df010473b21a4735c99b8c58df134ab3b0a714ccdbeb15bb739728d0d9bcf2","src/frame/modular/transforms/palette.rs":"42e66faa19c89f2115facaf664d08a5cf956a2db8c37276143d77c223247edd4","src/frame/modular/transforms/rct.rs":"3b68808702e154d88b81f797bc0f6c1f38a24f2bd91b8705184bc5bfa73ece62","src/frame/modular/transforms/squeeze.rs":"008e13d6d7167602345bea851a6ed9f152755482b35ba4c45b927e360ff2ff0f","src/frame/modular/tree.rs":"135f9398d30d65596cccfeca44f1bd92351217d96860afe0d3545abe7c02f41d","src/frame/quant_weights.rs":"859e98e4c209462624e26279cf9662fed607806fe73cb40b5b8bca296fcd28b4","src/frame/quantizer.rs":"c27eb7481c2b7a08e252da5db793594227b2249004490abce80cc8370e3d9b68","src/frame/transform_map.rs":"512df6f5c83a1348223f817bd4db9da684da1feaf1514de9973079de388214b6","src/headers/bit_depth.rs":"38f0d7fc0ec3470d4455045578e6b8ba1afd3afe84fa83d5a205a0e202510897","src/headers/color_encoding.rs":"085b1f73adbe3760fe79a6de28b3fba9444461c9efc3f82051d7e161bdabba02","src/headers/encodings.rs":"f3550ae26e4981ec08b8c7520f41c8beebaf5a0dba69db6665041d147d968e70","src/headers/extra_channels.rs":"00e6d64039be2406e67a005b37726b96f02b8f9b0329fadc2380bc45ac936e70","src/headers/frame_header.rs":"1ee08a8bbac2c3e088ea4462d8ab35d6597aa0080f7e898b6a0a23ed7b6e76e1","src/headers/image_metadata.rs":"82f4a34cf098736d845379ebbef3f5531774c2d49d6cdce1f22ffd403e4a3ebe","src/headers/mod.rs":"be4423908d8b22e6b056da39f99be57722442b609602affc95b0fa615332027c","src/headers/modular.rs":"1b542cbd2929fc1223a4eb5b4e0fb200d9b31c42e931a4a0dc4c5f7b4cc3a96e","src/headers/permutation.rs":"8d43a7589d1e46baebaa2f2b71d751b0901f02ffb3087b147da7e30246c9a96d","src/headers/size.rs":"84886382fb1e2ecf50e1561c58d60a56982619410f1eeff2faec166f6c22d1dd","src/headers/transform_data.rs":"2accdd3d3d4d2c7f4a6a2f1a0daa5fb9c68726d3e5f673b23d7978210e56dbae","src/icc.rs":"838dba0e98801651d51db2df96af836a610b71b2e0d20db22c993cac5e3f119a","src/icc/header.rs":"0fbddb05f171e3db4c019b40dc70423776e68a39a9e8b3e9c99448fc808ac8be","src/icc/stream.rs":"31215cf53affed0e367d1bb629852f56d55071a9ffc6b4ab10ef695bd247fe54","src/icc/tag.rs":"dbdb930e0c243acd97c0ffbe3daffbd3b1cb12a7c19ebb59baf6db6d248afb82","src/image.rs":"e956f553e50782569ed6462c3aaa0b99b655d04f6d3b51a53304fb2625c582b6","src/lib.rs":"cf06728a64a3e2ebfbecea21c0e8b3db7649c119d8e09c79deb3126d3a1261de","src/render/internal.rs":"c04b4492964ff2d2db7fade997aafa1d254ff6bc29e5a90648b46557534cc28d","src/render/mod.rs":"71eb2c2886a64663978a6184e8172c5e3606b3b443b57170d0881812b429eb2c","src/render/simple_pipeline/mod.rs":"25bbee0c625a96233a8286c20540a19e5dda105692434e842dab00ee46652921","src/render/simple_pipeline/save.rs":"bea6310cb75886b25df34ba2090c1d1a18adac786312ec4618ccd37069ee753b","src/render/stages/blending.rs":"5930abbf0eb69d1e3faf6b8f3d22469cb1f3ea301a28cf96aaa95bae8dcd9837","src/render/stages/chroma_upsample.rs":"9b3b693f393b6b0da233104eb227ccb5da8b6382e4218b63da3e99558302350d","src/render/stages/convert.rs":"e818856f44b798d8487c9dababe813f9aef070922d3a697ec6efc326a43762a4","src/render/stages/epf.rs":"285957d7256ad4347fa1053fa54f9982cac4e7efbbd7d0269c7355d14894ab28","src/render/stages/extend.rs":"8bc77daab980daead54785e121275ab33296c87f1859b2e61cd1905048d7c02f","src/render/stages/from_linear.rs":"25c3bf1cda6d6cf4680816eb0e973499b637f6be43b37f92d03403752b9f8962","src/render/stages/gaborish.rs":"36fb7a2c13e490b9f0a178b4588c13e68e6937009d55a0bfeacec9e63c48b14b","src/render/stages/mod.rs":"d143920ed5bdae327b6d3dccd13830ff33d097d00d6cd4bb5bf4412b42451983","src/render/stages/nearest_neighbor.rs":"11f900584cd8b86a5da079641b6546066fc56c9681a74d9e3d2657ea94681fc5","src/render/stages/noise.rs":"12264d901c95e98dd376f9c495d2543810f5d5e6fc1fa42baa36e5d7cb5a49f1","src/render/stages/patches.rs":"bad7a2ddb218d378169e84dedbc11f74af1281f45773c1d832fd35206a2dfe97","src/render/stages/splines.rs":"c3cbd58b17713aaf515adeb24d3265442795027a6f6c5386344185d118a4cc17","src/render/stages/spot.rs":"422fd7f064ba549005bf980c0cf3a6dbae0314c9978bcf2ed249f82941c391cc","src/render/stages/to_linear.rs":"e47b69b3e7dc74f25b3cff96849fe75f8264a06eb768e31cf494f334bc3bb855","src/render/stages/upsample.rs":"be36c375110f09254032dbaab9443473abf599755737ff052973c3e6fd2dd422","src/render/stages/xyb.rs":"1dd37362471eb961eea0064c2971f5bf93f49386e15b093fd48281864334e6ca","src/render/stages/ycbcr.rs":"6be2eed3ffba5f1d486d0f81f09f506184392024156ab2889e4486ef3f6b8892","src/render/test.rs":"9e63a1b905df6a797017a2f12d617d1f1e53be0fde445409f64c8398d25b979d","src/simd/mod.rs":"b0095e72ad673cf84824496a7a2eb5a88b2996e264adc87b11921e14c1b69e8f","src/simd/scalar.rs":"89c476b51497601767dbe29462b485b29d92ef23b004710d542fdaa083054a36","src/simd/x86_64/avx.rs":"0d629a5c5f6aef6ab1ab9a1d7793cd561e0dff4ce17355011a8fba7adce94ff2","src/simd/x86_64/avx512.rs":"e1403a936b5e22edf18ba6ac8375071e6db15130e8b8374050827a67aa2de0fc","src/simd/x86_64/mod.rs":"4e300384298af7878c245c42fbbd9c8f78680831ce22595ad10ea9a479e412dc","src/util.rs":"675cb07448a0dd075963db922fc21b95c4f09389e168fc36d06fc78d7266559e","src/util/bits.rs":"7b00f8488bb7b14766f236d56b79bd0413c647dd7918aa60f06f11b3bfb64cac","src/util/concat_slice.rs":"d6766e827e8ac9558ddcb1d9bb8a079d76c83e09cc61b121e0f39c610f2cea43","src/util/fast_math.rs":"eed12ce710f79777a26143aed4a2f9027356d81423f8fcf2edc9bb1d79e3ddb3","src/util/linalg.rs":"9289c3a6d4d1b82913d54103553282cdcee8ed5882ca3b7c59e9d14e46f3b588","src/util/log2.rs":"816977b20df5567d8127841b740c05ebfc8d2c92fc416dce9a72388c7e2e4c45","src/util/ndarray.rs":"32528f16885f72e7817348da7c106361728c7d598beb84381f9bcb42a460ee94","src/util/rational_poly.rs":"27c6cef01d9cc3115abbd6dc12ac862f6e82342a97ea0d19943896d2cf59650c","src/util/shift_right_ceil.rs":"d59dc1d5c577c6b47fa9639c74c3a5b973553274e2f6aef64c696db63442a2bd","src/util/test.rs":"96866d931205ead7f57d77f77ff27c448851f2335f21d7747e3770bde605c1f5","src/util/tracing_wrappers.rs":"d0d6ca6ca12c09e7e38dee589523730f997d9b5c950433d28d3ab43698fbb67b","src/util/vec_helpers.rs":"b3fef6f24f3cbb7522f8d0cf0a8085802ab0b6c5f32a7683d5697a823fd62392","src/util/xorshift128plus.rs":"24f2cc7b3600e478ce3919258b4dc90199e917106ebebe0ac9d67f11ac9599c3","src/var_dct.rs":"f9c36280190b4bfaee7fdd2819f0fa8ebdff0e88959b6e82eaf4cd03f3e0e23b","src/var_dct/dct.rs":"60ad89d06d8b2f8f20612138bd6b45aee9db7f710f536592db95bc52b452bcf4","src/var_dct/dct_scales.rs":"9ba712638b7678b1acb849f482d12dd038a673486740a4779264c81469b59821","src/var_dct/dct_slow.rs":"485e5183a7c7391822ca5e7647048f9f62d34b24c6daa777e7a1d0d5b5cf21ec","src/var_dct/transform.rs":"175b31355b1146a548245dba27719b6ce73c7c5d2cc550e11edbc5d9dcffa277"},"package":"710d04dc49b46298fa2258f3fd082bcafb318e93e45bfbffb2b2419d615964e7"} +\ No newline at end of file diff --git a/third_party/rust/jxl/Cargo.lock b/third_party/rust/jxl/Cargo.lock @@ -0,0 +1,682 @@ +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. +version = 4 + +[[package]] +name = "aho-corasick" +version = "1.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8e60d3430d3a69478ad0993f19238d2df97c507009a52b3c10addcd7f6bcb916" +dependencies = [ + "memchr", +] + +[[package]] +name = "anstream" +version = "0.6.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "64e15c1ab1f89faffbf04a634d5e1962e9074f2741eef6d97f3c4e322426d526" +dependencies = [ + "anstyle", + "anstyle-parse", + "anstyle-query", + "anstyle-wincon", + "colorchoice", + "is_terminal_polyfill", + "utf8parse", +] + +[[package]] +name = "anstyle" +version = "1.0.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1bec1de6f59aedf83baf9ff929c98f2ad654b97c9510f4e70cf6f661d49fd5b1" + +[[package]] +name = "anstyle-parse" +version = "0.2.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "eb47de1e80c2b463c735db5b217a0ddc39d612e7ac9e2e96a5aed1f57616c1cb" +dependencies = [ + "utf8parse", +] + +[[package]] +name = "anstyle-query" +version = "1.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6d36fc52c7f6c869915e99412912f22093507da8d9e942ceaf66fe4b7c14422a" +dependencies = [ + "windows-sys", +] + +[[package]] +name = "anstyle-wincon" +version = "3.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5bf74e1b6e971609db8ca7a9ce79fd5768ab6ae46441c572e46cf596f59e57f8" +dependencies = [ + "anstyle", + "windows-sys", +] + +[[package]] +name = "arbitrary" +version = "1.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7d5a26814d8dcb93b0e5a0ff3c6d80a8843bafb21b39e8e18a6f05471870e110" + +[[package]] +name = "arbtest" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2a3be567977128c0f71ad1462d9624ccda712193d124e944252f0c5789a06d46" +dependencies = [ + "arbitrary", +] + +[[package]] +name = "array-init" +version = "2.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3d62b7694a562cdf5a74227903507c56ab2cc8bdd1f781ed5cb4cf9c9f810bfc" + +[[package]] +name = "autocfg" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c4b4d0bd25bd0b74681c0ad21497610ce1b7c91b1022cd21c80c6fbdd9476b0" + +[[package]] +name = "byteorder" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fd0f2584146f6f2ef48085050886acf353beff7305ebd1ae69500e27c67f64b" + +[[package]] +name = "cfg-if" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "baf1de4339761588bc0619e3cbc0120ee582ebb74b53b4efbf79117bd2da40fd" + +[[package]] +name = "colorchoice" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d3fd119d74b830634cea2a0f58bbd0d54540518a14397557951e79340abc28c0" + +[[package]] +name = "crunchy" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7a81dae078cea95a014a339291cec439d2f232ebe854a9d672b796c6afafa9b7" + +[[package]] +name = "env_filter" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4f2c92ceda6ceec50f43169f9ee8424fe2db276791afde7b2cd8bc084cb376ab" +dependencies = [ + "log", +] + +[[package]] +name = "env_logger" +version = "0.11.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e13fa619b91fb2381732789fc5de83b45675e882f66623b7d8cb4f643017018d" +dependencies = [ + "anstream", + "anstyle", + "env_filter", + "log", +] + +[[package]] +name = "getrandom" +version = "0.2.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c4567c8db10ae91089c99af84c68c38da3ec2f087c3f82960bcdbf3656b6f4d7" +dependencies = [ + "cfg-if", + "libc", + "wasi", +] + +[[package]] +name = "half" +version = "2.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6dd08c532ae367adf81c312a4580bc67f1d0fe8bc9c460520283f4c0ff277888" +dependencies = [ + "cfg-if", + "crunchy", +] + +[[package]] +name = "is_terminal_polyfill" +version = "1.70.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7943c866cc5cd64cbc25b2e01621d07fa8eb2a1a23160ee81ce38704e97b8ecf" + +[[package]] +name = "jxl" +version = "0.1.1" +dependencies = [ + "arbtest", + "array-init", + "byteorder", + "half", + "jxl_macros", + "num-derive", + "num-traits", + "paste", + "rand", + "rand_xorshift", + "test-log", + "thiserror", + "tracing", +] + +[[package]] +name = "jxl_macros" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5169e84285ee08cee04f115f2658cafba62e7ff54e8eb91a3842129ca12b003f" +dependencies = [ + "proc-macro-error2", + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "lazy_static" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe" + +[[package]] +name = "libc" +version = "0.2.159" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "561d97a539a36e26a9a5fad1ea11a3039a67714694aaa379433e580854bc3dc5" + +[[package]] +name = "log" +version = "0.4.22" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a7a70ba024b9dc04c27ea2f0c0548feb474ec5c54bba33a7f72f873a39d07b24" + +[[package]] +name = "matchers" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8263075bb86c5a1b1427b5ae862e8889656f126e9f77c484496e8b47cf5c5558" +dependencies = [ + "regex-automata 0.1.10", +] + +[[package]] +name = "memchr" +version = "2.7.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "78ca9ab1a0babb1e7d5695e3530886289c18cf2f87ec19a575a0abdce112e3a3" + +[[package]] +name = "nu-ansi-term" +version = "0.46.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "77a8165726e8236064dbb45459242600304b42a5ea24ee2948e18e023bf7ba84" +dependencies = [ + "overload", + "winapi", +] + +[[package]] +name = "num-derive" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed3955f1a9c7c0c15e092f9c887db08b1fc683305fdf6eb6684f22555355e202" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "num-traits" +version = "0.2.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "071dfc062690e90b734c0b2273ce72ad0ffa95f0c74596bc250dcfd960262841" +dependencies = [ + "autocfg", +] + +[[package]] +name = "once_cell" +version = "1.19.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3fdb12b2476b595f9358c5161aa467c2438859caa136dec86c26fdd2efe17b92" + +[[package]] +name = "overload" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b15813163c1d831bf4a13c3610c05c0d03b39feb07f7e09fa234dac9b15aaf39" + +[[package]] +name = "paste" +version = "1.0.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "57c0d7b74b563b49d38dae00a0c37d4d6de9b432382b2892f0574ddcae73fd0a" + +[[package]] +name = "pin-project-lite" +version = "0.2.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bda66fc9667c18cb2758a2ac84d1167245054bcf85d5d1aaa6923f45801bdd02" + +[[package]] +name = "ppv-lite86" +version = "0.2.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "77957b295656769bb8ad2b6a6b09d897d94f05c41b069aede1fcdaa675eaea04" +dependencies = [ + "zerocopy", +] + +[[package]] +name = "proc-macro-error-attr2" +version = "2.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "96de42df36bb9bba5542fe9f1a054b8cc87e172759a1868aa05c1f3acc89dfc5" +dependencies = [ + "proc-macro2", + "quote", +] + +[[package]] +name = "proc-macro-error2" +version = "2.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "11ec05c52be0a07b08061f7dd003e7d7092e0472bc731b4af7bb1ef876109802" +dependencies = [ + "proc-macro-error-attr2", + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "proc-macro2" +version = "1.0.92" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "37d3544b3f2748c54e147655edb5025752e2303145b5aefb3c3ea2c78b973bb0" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "quote" +version = "1.0.37" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b5b9d34b8991d19d98081b46eacdd8eb58c6f2b201139f7c5f643cc155a633af" +dependencies = [ + "proc-macro2", +] + +[[package]] +name = "rand" +version = "0.8.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "34af8d1a0e25924bc5b7c43c079c942339d8f0a8b57c39049bef581b46327404" +dependencies = [ + "libc", + "rand_chacha", + "rand_core", +] + +[[package]] +name = "rand_chacha" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6c10a63a0fa32252be49d21e7709d4d4baf8d231c2dbce1eaa8141b9b127d88" +dependencies = [ + "ppv-lite86", + "rand_core", +] + +[[package]] +name = "rand_core" +version = "0.6.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ec0be4795e2f6a28069bec0b5ff3e2ac9bafc99e6a9a7dc3547996c5c816922c" +dependencies = [ + "getrandom", +] + +[[package]] +name = "rand_xorshift" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d25bf25ec5ae4a3f1b92f929810509a2f53d7dca2f50b794ff57e3face536c8f" +dependencies = [ + "rand_core", +] + +[[package]] +name = "regex" +version = "1.10.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4219d74c6b67a3654a9fbebc4b419e22126d13d2f3c4a07ee0cb61ff79a79619" +dependencies = [ + "aho-corasick", + "memchr", + "regex-automata 0.4.7", + "regex-syntax 0.8.4", +] + +[[package]] +name = "regex-automata" +version = "0.1.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6c230d73fb8d8c1b9c0b3135c5142a8acee3a0558fb8db5cf1cb65f8d7862132" +dependencies = [ + "regex-syntax 0.6.29", +] + +[[package]] +name = "regex-automata" +version = "0.4.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "38caf58cc5ef2fed281f89292ef23f6365465ed9a41b7a7754eb4e26496c92df" +dependencies = [ + "aho-corasick", + "memchr", + "regex-syntax 0.8.4", +] + +[[package]] +name = "regex-syntax" +version = "0.6.29" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f162c6dd7b008981e4d40210aca20b4bd0f9b60ca9271061b07f78537722f2e1" + +[[package]] +name = "regex-syntax" +version = "0.8.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7a66a03ae7c801facd77a29370b4faec201768915ac14a721ba36f20bc9c209b" + +[[package]] +name = "sharded-slab" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f40ca3c46823713e0d4209592e8d6e826aa57e928f09752619fc696c499637f6" +dependencies = [ + "lazy_static", +] + +[[package]] +name = "syn" +version = "2.0.90" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "919d3b74a5dd0ccd15aeb8f93e7006bd9e14c295087c9896a110f490752bcf31" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "test-log" +version = "0.2.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3dffced63c2b5c7be278154d76b479f9f9920ed34e7574201407f0b14e2bbb93" +dependencies = [ + "env_logger", + "test-log-macros", + "tracing-subscriber", +] + +[[package]] +name = "test-log-macros" +version = "0.2.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5999e24eaa32083191ba4e425deb75cdf25efefabe5aaccb7446dd0d4122a3f5" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "thiserror" +version = "2.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c006c85c7651b3cf2ada4584faa36773bd07bac24acfb39f3c431b36d7e667aa" +dependencies = [ + "thiserror-impl", +] + +[[package]] +name = "thiserror-impl" +version = "2.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f077553d607adc1caf65430528a576c757a71ed73944b66ebb58ef2bbd243568" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "thread_local" +version = "1.1.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8b9ef9bad013ada3808854ceac7b46812a6465ba368859a37e2100283d2d719c" +dependencies = [ + "cfg-if", + "once_cell", +] + +[[package]] +name = "tracing" +version = "0.1.40" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c3523ab5a71916ccf420eebdf5521fcef02141234bbc0b8a49f2fdc4544364ef" +dependencies = [ + "pin-project-lite", + "tracing-attributes", + "tracing-core", +] + +[[package]] +name = "tracing-attributes" +version = "0.1.27" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "34704c8d6ebcbc939824180af020566b01a7c01f80641264eba0999f6c2b6be7" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "tracing-core" +version = "0.1.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c06d3da6113f116aaee68e4d601191614c9053067f9ab7f6edbcb161237daa54" +dependencies = [ + "once_cell", + "valuable", +] + +[[package]] +name = "tracing-log" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ee855f1f400bd0e5c02d150ae5de3840039a3f54b025156404e34c23c03f47c3" +dependencies = [ + "log", + "once_cell", + "tracing-core", +] + +[[package]] +name = "tracing-subscriber" +version = "0.3.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ad0f048c97dbd9faa9b7df56362b8ebcaa52adb06b498c050d2f4e32f90a7a8b" +dependencies = [ + "matchers", + "nu-ansi-term", + "once_cell", + "regex", + "sharded-slab", + "thread_local", + "tracing", + "tracing-core", + "tracing-log", +] + +[[package]] +name = "unicode-ident" +version = "1.0.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e91b56cd4cadaeb79bbf1a5645f6b4f8dc5bde8834ad5894a8db35fda9efa1fe" + +[[package]] +name = "utf8parse" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "06abde3611657adf66d383f00b093d7faecc7fa57071cce2578660c9f1010821" + +[[package]] +name = "valuable" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "830b7e5d4d90034032940e4ace0d9a9a057e7a45cd94e6c007832e39edb82f6d" + +[[package]] +name = "wasi" +version = "0.11.0+wasi-snapshot-preview1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9c8d87e72b64a3b4db28d11ce29237c246188f4f51057d65a7eab63b7987e423" + +[[package]] +name = "winapi" +version = "0.3.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5c839a674fcd7a98952e593242ea400abe93992746761e38641405d28b00f419" +dependencies = [ + "winapi-i686-pc-windows-gnu", + "winapi-x86_64-pc-windows-gnu", +] + +[[package]] +name = "winapi-i686-pc-windows-gnu" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ac3b87c63620426dd9b991e5ce0329eff545bccbbb34f3be09ff6fb6ab51b7b6" + +[[package]] +name = "winapi-x86_64-pc-windows-gnu" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f" + +[[package]] +name = "windows-sys" +version = "0.52.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "282be5f36a8ce781fad8c8ae18fa3f9beff57ec1b52cb3de0789201425d9a33d" +dependencies = [ + "windows-targets", +] + +[[package]] +name = "windows-targets" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9b724f72796e036ab90c1021d4780d4d3d648aca59e491e6b98e725b84e99973" +dependencies = [ + "windows_aarch64_gnullvm", + "windows_aarch64_msvc", + "windows_i686_gnu", + "windows_i686_gnullvm", + "windows_i686_msvc", + "windows_x86_64_gnu", + "windows_x86_64_gnullvm", + "windows_x86_64_msvc", +] + +[[package]] +name = "windows_aarch64_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32a4622180e7a0ec044bb555404c800bc9fd9ec262ec147edd5989ccd0c02cd3" + +[[package]] +name = "windows_aarch64_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09ec2a7bb152e2252b53fa7803150007879548bc709c039df7627cabbd05d469" + +[[package]] +name = "windows_i686_gnu" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8e9b5ad5ab802e97eb8e295ac6720e509ee4c243f69d781394014ebfe8bbfa0b" + +[[package]] +name = "windows_i686_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0eee52d38c090b3caa76c563b86c3a4bd71ef1a819287c19d586d7334ae8ed66" + +[[package]] +name = "windows_i686_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "240948bc05c5e7c6dabba28bf89d89ffce3e303022809e73deaefe4f6ec56c66" + +[[package]] +name = "windows_x86_64_gnu" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "147a5c80aabfbf0c7d901cb5895d1de30ef2907eb21fbbab29ca94c5b08b1a78" + +[[package]] +name = "windows_x86_64_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "24d5b23dc417412679681396f2b49f3de8c1473deb516bd34410872eff51ed0d" + +[[package]] +name = "windows_x86_64_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" + +[[package]] +name = "zerocopy" +version = "0.7.35" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1b9b4fd18abc82b8136838da5d50bae7bdea537c574d8dc1a34ed098d6c166f0" +dependencies = [ + "byteorder", + "zerocopy-derive", +] + +[[package]] +name = "zerocopy-derive" +version = "0.7.35" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fa4f8080344d4671fb4e831a13ad1e68092748387dfc4f55e356242fae12ce3e" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] diff --git a/third_party/rust/jxl/Cargo.toml b/third_party/rust/jxl/Cargo.toml @@ -0,0 +1,88 @@ +# THIS FILE IS AUTOMATICALLY GENERATED BY CARGO +# +# When uploading crates to the registry Cargo will automatically +# "normalize" Cargo.toml files for maximal compatibility +# with all versions of Cargo and also rewrite `path` dependencies +# to registry (e.g., crates.io) dependencies. +# +# If you are reading this file be aware that the original Cargo.toml +# will likely look very different (and much more reasonable). +# See Cargo.toml.orig for the original contents. + +[package] +edition = "2024" +name = "jxl" +version = "0.1.1" +authors = ["Luca Versari <veluca93@gmail.com>"] +build = false +autolib = false +autobins = false +autoexamples = false +autotests = false +autobenches = false +description = "High performance Rust implementation of a JPEG XL decoder" +readme = "README.md" +keywords = [ + "jpeg-xl", + "decoder", +] +categories = ["multimedia::images"] +license = "BSD-3-Clause" +repository = "https://github.com/libjxl/jxl-rs" +resolver = "2" + +[lib] +name = "jxl" +path = "src/lib.rs" + +[dependencies.array-init] +version = "2.0.0" + +[dependencies.byteorder] +version = "1.4.3" + +[dependencies.half] +version = "2.4.1" + +[dependencies.jxl_macros] +version = "=0.1.1" + +[dependencies.num-derive] +version = "0.4" + +[dependencies.num-traits] +version = "0.2.14" + +[dependencies.thiserror] +version = "2.0" + +[dependencies.tracing] +version = "0.1.40" +optional = true + +[dev-dependencies.arbtest] +version = "0.3.2" + +[dev-dependencies.jxl_macros] +version = "=0.1.1" +features = ["test"] + +[dev-dependencies.paste] +version = "1.0.15" + +[dev-dependencies.rand] +version = "0.8.5" + +[dev-dependencies.rand_xorshift] +version = "0.3.0" + +[dev-dependencies.test-log] +version = "0.2.16" +features = ["trace"] + +[lints.clippy] +missing_safety_doc = "deny" +undocumented_unsafe_blocks = "deny" + +[lints.rust] +unsafe_op_in_unsafe_fn = "deny" diff --git a/third_party/rust/jxl/README.md b/third_party/rust/jxl/README.md @@ -0,0 +1,7 @@ +# JPEG XL in Rust + +This is a work-in-progress, and currently incomplete reimplementation of a JPEG XL decoder in Rust, which aims to be conforming, safe and fast. + +Refer to [the main libjxl repository](https://github.com/libjxl/libjxl) for +more information, including contributing instructions. + diff --git a/third_party/rust/jxl/resources/test/8x8_noise.jxl b/third_party/rust/jxl/resources/test/8x8_noise.jxl Binary files differ. diff --git a/third_party/rust/jxl/resources/test/all.html b/third_party/rust/jxl/resources/test/all.html @@ -0,0 +1,77 @@ +<!DOCTYPE html> +<html> + +<head> + <title>JXL Test Images</title> +</head> + +<body> + <img src="conformance_test_images/alpha_nonpremultiplied.jxl" /> + <img src="conformance_test_images/alpha_premultiplied.jxl" /> + <img src="conformance_test_images/alpha_triangles.jxl" /> + <img src="conformance_test_images/animation_icos4d_5.jxl" /> + <img src="conformance_test_images/animation_icos4d.jxl" /> + <img src="conformance_test_images/animation_newtons_cradle.jxl" /> + <img src="conformance_test_images/animation_spline_5.jxl" /> + <img src="conformance_test_images/animation_spline.jxl" /> + <img src="conformance_test_images/bench_oriented_brg_5.jxl" /> + <img src="conformance_test_images/bench_oriented_brg.jxl" /> + <img src="conformance_test_images/bicycles.jxl" /> + <img src="conformance_test_images/bike_5.jxl" /> + <img src="conformance_test_images/bike.jxl" /> + <img src="conformance_test_images/blendmodes_5.jxl" /> + <img src="conformance_test_images/blendmodes.jxl" /> + <img src="conformance_test_images/cafe_5.jxl" /> + <img src="conformance_test_images/cafe.jxl" /> + <img src="conformance_test_images/cmyk_layers.jxl" /> + <img src="conformance_test_images/delta_palette.jxl" /> + <img src="conformance_test_images/grayscale_5.jxl" /> + <img src="conformance_test_images/grayscale_jpeg_5.jxl" /> + <img src="conformance_test_images/grayscale_jpeg.jxl" /> + <img src="conformance_test_images/grayscale.jxl" /> + <img src="conformance_test_images/grayscale_public_university.jxl" /> + <img src="conformance_test_images/lossless_pfm.jxl" /> + <img src="conformance_test_images/lz77_flower.jxl" /> + <img src="conformance_test_images/noise_5.jxl" /> + <img src="conformance_test_images/noise.jxl" /> + <img src="conformance_test_images/opsin_inverse_5.jxl" /> + <img src="conformance_test_images/opsin_inverse.jxl" /> + <img src="conformance_test_images/patches_5.jxl" /> + <img src="conformance_test_images/patches.jxl" /> + <img src="conformance_test_images/progressive_5.jxl" /> + <img src="conformance_test_images/progressive.jxl" /> + <img src="conformance_test_images/spot.jxl" /> + <img src="conformance_test_images/sunset_logo.jxl" /> + <img src="conformance_test_images/upsampling_5.jxl" /> + <img src="conformance_test_images/upsampling.jxl" /> + <img src="8x8_noise.jxl" /> + <img src="basic.jxl" /> + <img src="candle.jxl" /> + <img src="cropped_traffic_light.jxl" /> + <img src="extra_channels.jxl" /> + <img src="grayscale_patches_modular.jxl" /> + <img src="grayscale_patches_var_dct.jxl" /> + <img src="green_queen_modular_e3.jxl" /> + <img src="green_queen_vardct_e3.jxl" /> + <img src="has_permutation.jxl" /> + <img src="has_permutation_with_container.jxl" /> + <img src="multiple_lf_420.jxl" /> + <img src="orientation1_identity.jxl" /> + <img src="orientation2_flip_horizontal.jxl" /> + <img src="orientation3_rotate_180.jxl" /> + <img src="orientation4_flip_vertical.jxl" /> + <img src="orientation5_transpose.jxl" /> + <img src="orientation6_rotate_90_cw.jxl" /> + <img src="orientation7_anti_transpose.jxl" /> + <img src="orientation8_rotate_90_ccw.jxl" /> + <img src="pq_gradient.jxl" /> + <img src="small_grayscale_patches_modular.jxl" /> + <img src="small_grayscale_patches_modular_with_icc.jxl" /> + <img src="spline_on_first_frame.jxl" /> + <img src="splines.jxl" /> + <img src="squeeze_edge.jxl" /> + <img src="with_icc.jxl" /> + <img src="zoltan_tasi_unsplash.jxl" /> +</body> + +</html> diff --git a/third_party/rust/jxl/resources/test/basic.jxl b/third_party/rust/jxl/resources/test/basic.jxl Binary files differ. diff --git a/third_party/rust/jxl/resources/test/candle.jxl b/third_party/rust/jxl/resources/test/candle.jxl Binary files differ. diff --git a/third_party/rust/jxl/resources/test/candle_license.txt b/third_party/rust/jxl/resources/test/candle_license.txt @@ -0,0 +1,35 @@ +Original image at https://github.com/AcademySoftwareFoundation/openexr-images/blob/main/ScanLines/CandleGlass.exr + +Copyright (c) 2004, Industrial Light & Magic, a division of Lucasfilm +Entertainment Company Ltd. Portions contributed and copyright held by +others as indicated. All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are +met: + + * Redistributions of source code must retain the above + copyright notice, this list of conditions and the following + disclaimer. + + * Redistributions in binary form must reproduce the above + copyright notice, this list of conditions and the following + disclaimer in the documentation and/or other materials provided with + the distribution. + + * Neither the name of Industrial Light & Magic nor the names of + any other contributors to this software may be used to endorse or + promote products derived from this software without specific prior + written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS +IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, +THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR +PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR +CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, +EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, +PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR +PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF +LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING +NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS +SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. diff --git a/third_party/rust/jxl/resources/test/conformance_test_images/alpha_nonpremultiplied.jxl b/third_party/rust/jxl/resources/test/conformance_test_images/alpha_nonpremultiplied.jxl Binary files differ. diff --git a/third_party/rust/jxl/resources/test/conformance_test_images/alpha_premultiplied.jxl b/third_party/rust/jxl/resources/test/conformance_test_images/alpha_premultiplied.jxl Binary files differ. diff --git a/third_party/rust/jxl/resources/test/conformance_test_images/alpha_triangles.jxl b/third_party/rust/jxl/resources/test/conformance_test_images/alpha_triangles.jxl Binary files differ. diff --git a/third_party/rust/jxl/resources/test/conformance_test_images/animation_icos4d.jxl b/third_party/rust/jxl/resources/test/conformance_test_images/animation_icos4d.jxl Binary files differ. diff --git a/third_party/rust/jxl/resources/test/conformance_test_images/animation_icos4d_5.jxl b/third_party/rust/jxl/resources/test/conformance_test_images/animation_icos4d_5.jxl Binary files differ. diff --git a/third_party/rust/jxl/resources/test/conformance_test_images/animation_newtons_cradle.jxl b/third_party/rust/jxl/resources/test/conformance_test_images/animation_newtons_cradle.jxl Binary files differ. diff --git a/third_party/rust/jxl/resources/test/conformance_test_images/animation_spline.jxl b/third_party/rust/jxl/resources/test/conformance_test_images/animation_spline.jxl Binary files differ. diff --git a/third_party/rust/jxl/resources/test/conformance_test_images/animation_spline_5.jxl b/third_party/rust/jxl/resources/test/conformance_test_images/animation_spline_5.jxl Binary files differ. diff --git a/third_party/rust/jxl/resources/test/conformance_test_images/bench_oriented_brg.jxl b/third_party/rust/jxl/resources/test/conformance_test_images/bench_oriented_brg.jxl Binary files differ. diff --git a/third_party/rust/jxl/resources/test/conformance_test_images/bench_oriented_brg_5.jxl b/third_party/rust/jxl/resources/test/conformance_test_images/bench_oriented_brg_5.jxl Binary files differ. diff --git a/third_party/rust/jxl/resources/test/conformance_test_images/bicycles.jxl b/third_party/rust/jxl/resources/test/conformance_test_images/bicycles.jxl Binary files differ. diff --git a/third_party/rust/jxl/resources/test/conformance_test_images/bike.jxl b/third_party/rust/jxl/resources/test/conformance_test_images/bike.jxl Binary files differ. diff --git a/third_party/rust/jxl/resources/test/conformance_test_images/bike_5.jxl b/third_party/rust/jxl/resources/test/conformance_test_images/bike_5.jxl Binary files differ. diff --git a/third_party/rust/jxl/resources/test/conformance_test_images/blendmodes.jxl b/third_party/rust/jxl/resources/test/conformance_test_images/blendmodes.jxl Binary files differ. diff --git a/third_party/rust/jxl/resources/test/conformance_test_images/blendmodes_5.jxl b/third_party/rust/jxl/resources/test/conformance_test_images/blendmodes_5.jxl Binary files differ. diff --git a/third_party/rust/jxl/resources/test/conformance_test_images/cafe.jxl b/third_party/rust/jxl/resources/test/conformance_test_images/cafe.jxl Binary files differ. diff --git a/third_party/rust/jxl/resources/test/conformance_test_images/cafe_5.jxl b/third_party/rust/jxl/resources/test/conformance_test_images/cafe_5.jxl Binary files differ. diff --git a/third_party/rust/jxl/resources/test/conformance_test_images/cmyk_layers.jxl b/third_party/rust/jxl/resources/test/conformance_test_images/cmyk_layers.jxl Binary files differ. diff --git a/third_party/rust/jxl/resources/test/conformance_test_images/delta_palette.jxl b/third_party/rust/jxl/resources/test/conformance_test_images/delta_palette.jxl Binary files differ. diff --git a/third_party/rust/jxl/resources/test/conformance_test_images/grayscale.jxl b/third_party/rust/jxl/resources/test/conformance_test_images/grayscale.jxl Binary files differ. diff --git a/third_party/rust/jxl/resources/test/conformance_test_images/grayscale_5.jxl b/third_party/rust/jxl/resources/test/conformance_test_images/grayscale_5.jxl Binary files differ. diff --git a/third_party/rust/jxl/resources/test/conformance_test_images/grayscale_jpeg.jxl b/third_party/rust/jxl/resources/test/conformance_test_images/grayscale_jpeg.jxl Binary files differ. diff --git a/third_party/rust/jxl/resources/test/conformance_test_images/grayscale_jpeg_5.jxl b/third_party/rust/jxl/resources/test/conformance_test_images/grayscale_jpeg_5.jxl Binary files differ. diff --git a/third_party/rust/jxl/resources/test/conformance_test_images/grayscale_public_university.jxl b/third_party/rust/jxl/resources/test/conformance_test_images/grayscale_public_university.jxl Binary files differ. diff --git a/third_party/rust/jxl/resources/test/conformance_test_images/lossless_pfm.jxl b/third_party/rust/jxl/resources/test/conformance_test_images/lossless_pfm.jxl Binary files differ. diff --git a/third_party/rust/jxl/resources/test/conformance_test_images/lz77_flower.jxl b/third_party/rust/jxl/resources/test/conformance_test_images/lz77_flower.jxl Binary files differ. diff --git a/third_party/rust/jxl/resources/test/conformance_test_images/noise.jxl b/third_party/rust/jxl/resources/test/conformance_test_images/noise.jxl Binary files differ. diff --git a/third_party/rust/jxl/resources/test/conformance_test_images/noise_5.jxl b/third_party/rust/jxl/resources/test/conformance_test_images/noise_5.jxl Binary files differ. diff --git a/third_party/rust/jxl/resources/test/conformance_test_images/opsin_inverse.jxl b/third_party/rust/jxl/resources/test/conformance_test_images/opsin_inverse.jxl Binary files differ. diff --git a/third_party/rust/jxl/resources/test/conformance_test_images/opsin_inverse_5.jxl b/third_party/rust/jxl/resources/test/conformance_test_images/opsin_inverse_5.jxl Binary files differ. diff --git a/third_party/rust/jxl/resources/test/conformance_test_images/patches.jxl b/third_party/rust/jxl/resources/test/conformance_test_images/patches.jxl Binary files differ. diff --git a/third_party/rust/jxl/resources/test/conformance_test_images/patches_5.jxl b/third_party/rust/jxl/resources/test/conformance_test_images/patches_5.jxl Binary files differ. diff --git a/third_party/rust/jxl/resources/test/conformance_test_images/progressive.jxl b/third_party/rust/jxl/resources/test/conformance_test_images/progressive.jxl Binary files differ. diff --git a/third_party/rust/jxl/resources/test/conformance_test_images/progressive_5.jxl b/third_party/rust/jxl/resources/test/conformance_test_images/progressive_5.jxl Binary files differ. diff --git a/third_party/rust/jxl/resources/test/conformance_test_images/spot.jxl b/third_party/rust/jxl/resources/test/conformance_test_images/spot.jxl Binary files differ. diff --git a/third_party/rust/jxl/resources/test/conformance_test_images/sunset_logo.jxl b/third_party/rust/jxl/resources/test/conformance_test_images/sunset_logo.jxl Binary files differ. diff --git a/third_party/rust/jxl/resources/test/conformance_test_images/upsampling.jxl b/third_party/rust/jxl/resources/test/conformance_test_images/upsampling.jxl Binary files differ. diff --git a/third_party/rust/jxl/resources/test/conformance_test_images/upsampling_5.jxl b/third_party/rust/jxl/resources/test/conformance_test_images/upsampling_5.jxl Binary files differ. diff --git a/third_party/rust/jxl/resources/test/cropped_traffic_light.jxl b/third_party/rust/jxl/resources/test/cropped_traffic_light.jxl Binary files differ. diff --git a/third_party/rust/jxl/resources/test/extra_channels.jxl b/third_party/rust/jxl/resources/test/extra_channels.jxl Binary files differ. diff --git a/third_party/rust/jxl/resources/test/grayscale_patches_modular.jxl b/third_party/rust/jxl/resources/test/grayscale_patches_modular.jxl Binary files differ. diff --git a/third_party/rust/jxl/resources/test/grayscale_patches_var_dct.jxl b/third_party/rust/jxl/resources/test/grayscale_patches_var_dct.jxl Binary files differ. diff --git a/third_party/rust/jxl/resources/test/green_queen_modular_e3.jxl b/third_party/rust/jxl/resources/test/green_queen_modular_e3.jxl Binary files differ. diff --git a/third_party/rust/jxl/resources/test/green_queen_vardct_e3.jxl b/third_party/rust/jxl/resources/test/green_queen_vardct_e3.jxl Binary files differ. diff --git a/third_party/rust/jxl/resources/test/has_permutation.jxl b/third_party/rust/jxl/resources/test/has_permutation.jxl Binary files differ. diff --git a/third_party/rust/jxl/resources/test/has_permutation_with_container.jxl b/third_party/rust/jxl/resources/test/has_permutation_with_container.jxl Binary files differ. diff --git a/third_party/rust/jxl/resources/test/multiple_lf_420.jxl b/third_party/rust/jxl/resources/test/multiple_lf_420.jxl Binary files differ. diff --git a/third_party/rust/jxl/resources/test/orientation1_identity.jxl b/third_party/rust/jxl/resources/test/orientation1_identity.jxl Binary files differ. diff --git a/third_party/rust/jxl/resources/test/orientation2_flip_horizontal.jxl b/third_party/rust/jxl/resources/test/orientation2_flip_horizontal.jxl Binary files differ. diff --git a/third_party/rust/jxl/resources/test/orientation3_rotate_180.jxl b/third_party/rust/jxl/resources/test/orientation3_rotate_180.jxl Binary files differ. diff --git a/third_party/rust/jxl/resources/test/orientation4_flip_vertical.jxl b/third_party/rust/jxl/resources/test/orientation4_flip_vertical.jxl Binary files differ. diff --git a/third_party/rust/jxl/resources/test/orientation5_transpose.jxl b/third_party/rust/jxl/resources/test/orientation5_transpose.jxl Binary files differ. diff --git a/third_party/rust/jxl/resources/test/orientation6_rotate_90_cw.jxl b/third_party/rust/jxl/resources/test/orientation6_rotate_90_cw.jxl Binary files differ. diff --git a/third_party/rust/jxl/resources/test/orientation7_anti_transpose.jxl b/third_party/rust/jxl/resources/test/orientation7_anti_transpose.jxl Binary files differ. diff --git a/third_party/rust/jxl/resources/test/orientation8_rotate_90_ccw.jxl b/third_party/rust/jxl/resources/test/orientation8_rotate_90_ccw.jxl Binary files differ. diff --git a/third_party/rust/jxl/resources/test/pq_gradient.jxl b/third_party/rust/jxl/resources/test/pq_gradient.jxl Binary files differ. diff --git a/third_party/rust/jxl/resources/test/small_grayscale_patches_modular.jxl b/third_party/rust/jxl/resources/test/small_grayscale_patches_modular.jxl Binary files differ. diff --git a/third_party/rust/jxl/resources/test/small_grayscale_patches_modular_with_icc.jxl b/third_party/rust/jxl/resources/test/small_grayscale_patches_modular_with_icc.jxl Binary files differ. diff --git a/third_party/rust/jxl/resources/test/spline_on_first_frame.jxl b/third_party/rust/jxl/resources/test/spline_on_first_frame.jxl Binary files differ. diff --git a/third_party/rust/jxl/resources/test/splines.jxl b/third_party/rust/jxl/resources/test/splines.jxl Binary files differ. diff --git a/third_party/rust/jxl/resources/test/splines.pfm b/third_party/rust/jxl/resources/test/splines.pfm Binary files differ. diff --git a/third_party/rust/jxl/resources/test/squeeze_edge.jxl b/third_party/rust/jxl/resources/test/squeeze_edge.jxl Binary files differ. diff --git a/third_party/rust/jxl/resources/test/with_icc.jxl b/third_party/rust/jxl/resources/test/with_icc.jxl Binary files differ. diff --git a/third_party/rust/jxl/resources/test/zoltan_tasi_unsplash.jxl b/third_party/rust/jxl/resources/test/zoltan_tasi_unsplash.jxl Binary files differ. diff --git a/third_party/rust/jxl/src/api.rs b/third_party/rust/jxl/src/api.rs @@ -0,0 +1,65 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +// #![warn(missing_docs)] + +mod color; +mod data_types; +mod decoder; +mod inner; +mod input; +mod options; +mod output; +mod signature; +#[cfg(test)] +pub(crate) mod test; + +pub use color::*; +pub use data_types::*; +pub use decoder::*; +pub use inner::*; +pub use input::*; +pub use options::*; +pub use output::*; +pub use signature::*; + +use crate::headers::image_metadata::Orientation; + +/// This type represents the return value of a function that reads input from a bitstream. The +/// variant `Complete` indicates that the operation was completed successfully, and its return +/// value is available. The variant `NeedsMoreInput` indicates that more input is needed, and the +/// function should be called again. This variant comes with a `size_hint`, representing an +/// estimate of the number of additional bytes needed, and a `fallback`, representing additional +/// information that might be needed to call the function again (i.e. because it takes a decoder +/// object by value). +#[derive(Debug, PartialEq)] +pub enum ProcessingResult<T, U> { + Complete { result: T }, + NeedsMoreInput { size_hint: usize, fallback: U }, +} + +impl<T> ProcessingResult<T, ()> { + fn new( + result: Result<T, crate::error::Error>, + ) -> Result<ProcessingResult<T, ()>, crate::error::Error> { + match result { + Ok(v) => Ok(ProcessingResult::Complete { result: v }), + Err(crate::error::Error::OutOfBounds(v)) => Ok(ProcessingResult::NeedsMoreInput { + size_hint: v, + fallback: (), + }), + Err(e) => Err(e), + } + } +} + +#[derive(Clone)] +pub struct JxlBasicInfo { + pub size: (usize, usize), + pub bit_depth: JxlBitDepth, + pub orientation: Orientation, + pub extra_channels: Vec<JxlExtraChannel>, + pub animation: Option<JxlAnimation>, +} diff --git a/third_party/rust/jxl/src/api/color.rs b/third_party/rust/jxl/src/api/color.rs @@ -0,0 +1,1561 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +use std::{borrow::Cow, fmt}; + +use crate::{ + error::{Error, Result}, + headers::color_encoding::{ + ColorEncoding, ColorSpace, Primaries, RenderingIntent, TransferFunction, WhitePoint, + }, + util::{Matrix3x3, Vector3, inv_3x3_matrix, mul_3x3_matrix, mul_3x3_vector}, +}; + +// Bradford matrices for chromatic adaptation +const K_BRADFORD: Matrix3x3<f64> = [ + [0.8951, 0.2664, -0.1614], + [-0.7502, 1.7135, 0.0367], + [0.0389, -0.0685, 1.0296], +]; + +const K_BRADFORD_INV: Matrix3x3<f64> = [ + [0.9869929, -0.1470543, 0.1599627], + [0.4323053, 0.5183603, 0.0492912], + [-0.0085287, 0.0400428, 0.9684867], +]; + +pub fn compute_md5(data: &[u8]) -> [u8; 16] { + let mut sum = [0u8; 16]; + let mut data64 = data.to_vec(); + data64.push(128); + + // Add bytes such that ((size + 8) & 63) == 0 + let extra = (64 - ((data64.len() + 8) & 63)) & 63; + data64.resize(data64.len() + extra, 0); + + // Append length in bits as 64-bit little-endian + let bit_len = (data.len() as u64) << 3; + for i in (0..64).step_by(8) { + data64.push((bit_len >> i) as u8); + } + + const SINEPARTS: [u32; 64] = [ + 0xd76aa478, 0xe8c7b756, 0x242070db, 0xc1bdceee, 0xf57c0faf, 0x4787c62a, 0xa8304613, + 0xfd469501, 0x698098d8, 0x8b44f7af, 0xffff5bb1, 0x895cd7be, 0x6b901122, 0xfd987193, + 0xa679438e, 0x49b40821, 0xf61e2562, 0xc040b340, 0x265e5a51, 0xe9b6c7aa, 0xd62f105d, + 0x02441453, 0xd8a1e681, 0xe7d3fbc8, 0x21e1cde6, 0xc33707d6, 0xf4d50d87, 0x455a14ed, + 0xa9e3e905, 0xfcefa3f8, 0x676f02d9, 0x8d2a4c8a, 0xfffa3942, 0x8771f681, 0x6d9d6122, + 0xfde5380c, 0xa4beea44, 0x4bdecfa9, 0xf6bb4b60, 0xbebfbc70, 0x289b7ec6, 0xeaa127fa, + 0xd4ef3085, 0x04881d05, 0xd9d4d039, 0xe6db99e5, 0x1fa27cf8, 0xc4ac5665, 0xf4292244, + 0x432aff97, 0xab9423a7, 0xfc93a039, 0x655b59c3, 0x8f0ccc92, 0xffeff47d, 0x85845dd1, + 0x6fa87e4f, 0xfe2ce6e0, 0xa3014314, 0x4e0811a1, 0xf7537e82, 0xbd3af235, 0x2ad7d2bb, + 0xeb86d391, + ]; + + const SHIFT: [u32; 64] = [ + 7, 12, 17, 22, 7, 12, 17, 22, 7, 12, 17, 22, 7, 12, 17, 22, 5, 9, 14, 20, 5, 9, 14, 20, 5, + 9, 14, 20, 5, 9, 14, 20, 4, 11, 16, 23, 4, 11, 16, 23, 4, 11, 16, 23, 4, 11, 16, 23, 6, 10, + 15, 21, 6, 10, 15, 21, 6, 10, 15, 21, 6, 10, 15, 21, + ]; + + let mut a0: u32 = 0x67452301; + let mut b0: u32 = 0xefcdab89; + let mut c0: u32 = 0x98badcfe; + let mut d0: u32 = 0x10325476; + + for i in (0..data64.len()).step_by(64) { + let mut a = a0; + let mut b = b0; + let mut c = c0; + let mut d = d0; + + for j in 0..64 { + let (f, g) = if j < 16 { + ((b & c) | ((!b) & d), j) + } else if j < 32 { + ((d & b) | ((!d) & c), (5 * j + 1) & 0xf) + } else if j < 48 { + (b ^ c ^ d, (3 * j + 5) & 0xf) + } else { + (c ^ (b | (!d)), (7 * j) & 0xf) + }; + + let dg0 = data64[i + g * 4] as u32; + let dg1 = data64[i + g * 4 + 1] as u32; + let dg2 = data64[i + g * 4 + 2] as u32; + let dg3 = data64[i + g * 4 + 3] as u32; + let u = dg0 | (dg1 << 8) | (dg2 << 16) | (dg3 << 24); + + let f = f.wrapping_add(a).wrapping_add(SINEPARTS[j]).wrapping_add(u); + a = d; + d = c; + c = b; + b = b.wrapping_add((f << SHIFT[j]) | (f >> (32 - SHIFT[j]))); + } + + a0 = a0.wrapping_add(a); + b0 = b0.wrapping_add(b); + c0 = c0.wrapping_add(c); + d0 = d0.wrapping_add(d); + } + + sum[0] = a0 as u8; + sum[1] = (a0 >> 8) as u8; + sum[2] = (a0 >> 16) as u8; + sum[3] = (a0 >> 24) as u8; + sum[4] = b0 as u8; + sum[5] = (b0 >> 8) as u8; + sum[6] = (b0 >> 16) as u8; + sum[7] = (b0 >> 24) as u8; + sum[8] = c0 as u8; + sum[9] = (c0 >> 8) as u8; + sum[10] = (c0 >> 16) as u8; + sum[11] = (c0 >> 24) as u8; + sum[12] = d0 as u8; + sum[13] = (d0 >> 8) as u8; + sum[14] = (d0 >> 16) as u8; + sum[15] = (d0 >> 24) as u8; + sum +} + +#[allow(clippy::too_many_arguments)] +pub(crate) fn primaries_to_xyz( + rx: f32, + ry: f32, + gx: f32, + gy: f32, + bx: f32, + by: f32, + wx: f32, + wy: f32, +) -> Result<Matrix3x3<f64>, Error> { + // Validate white point coordinates + if !((0.0..=1.0).contains(&wx) && (wy > 0.0 && wy <= 1.0)) { + return Err(Error::IccInvalidWhitePoint( + wx, + wy, + "White point coordinates out of range ([0,1] for x, (0,1] for y)".to_string(), + )); + } + // Comment from libjxl: + // TODO(lode): also require rx, ry, gx, gy, bx, to be in range 0-1? ICC + // profiles in theory forbid negative XYZ values, but in practice the ACES P0 + // color space uses a negative y for the blue primary. + + // Construct the primaries matrix P. Its columns are the XYZ coordinates + // of the R, G, B primaries (derived from their chromaticities x, y, z=1-x-y). + // P = [[xr, xg, xb], + // [yr, yg, yb], + // [zr, zg, zb]] + let rz = 1.0 - rx as f64 - ry as f64; + let gz = 1.0 - gx as f64 - gy as f64; + let bz = 1.0 - bx as f64 - by as f64; + let p_matrix = [ + [rx as f64, gx as f64, bx as f64], + [ry as f64, gy as f64, by as f64], + [rz, gz, bz], + ]; + + let p_inv_matrix = inv_3x3_matrix(&p_matrix)?; + + // Convert reference white point (wx, wy) to XYZ form with Y=1 + // This is WhitePoint_XYZ_wp = [wx/wy, 1, (1-wx-wy)/wy] + let x_over_y_wp = wx as f64 / wy as f64; + let z_over_y_wp = (1.0 - wx as f64 - wy as f64) / wy as f64; + + if !x_over_y_wp.is_finite() || !z_over_y_wp.is_finite() { + return Err(Error::IccInvalidWhitePoint( + wx, + wy, + "Calculated X/Y or Z/Y for white point is not finite.".to_string(), + )); + } + let white_point_xyz_vec: Vector3<f64> = [x_over_y_wp, 1.0, z_over_y_wp]; + + // Calculate scaling factors S = [Sr, Sg, Sb] such that P * S = WhitePoint_XYZ_wp + // So, S = P_inv * WhitePoint_XYZ_wp + let s_vec = mul_3x3_vector(&p_inv_matrix, &white_point_xyz_vec); + + // Construct diagonal matrix S_diag from s_vec + let s_diag_matrix = [ + [s_vec[0], 0.0, 0.0], + [0.0, s_vec[1], 0.0], + [0.0, 0.0, s_vec[2]], + ]; + // The final RGB-to-XYZ matrix is P * S_diag + let result_matrix = mul_3x3_matrix(&p_matrix, &s_diag_matrix); + + Ok(result_matrix) +} + +pub(crate) fn adapt_to_xyz_d50(wx: f32, wy: f32) -> Result<Matrix3x3<f64>, Error> { + if !((0.0..=1.0).contains(&wx) && (wy > 0.0 && wy <= 1.0)) { + return Err(Error::IccInvalidWhitePoint( + wx, + wy, + "White point coordinates out of range ([0,1] for x, (0,1] for y)".to_string(), + )); + } + + // Convert white point (wx, wy) to XYZ with Y=1 + let x_over_y = wx as f64 / wy as f64; + let z_over_y = (1.0 - wx as f64 - wy as f64) / wy as f64; + + // Check for finiteness, as 1.0 / tiny float can overflow. + if !x_over_y.is_finite() || !z_over_y.is_finite() { + return Err(Error::IccInvalidWhitePoint( + wx, + wy, + "Calculated X/Y or Z/Y for white point is not finite.".to_string(), + )); + } + let w: Vector3<f64> = [x_over_y, 1.0, z_over_y]; + + // D50 white point in XYZ (Y=1 form) + // These are X_D50/Y_D50, 1.0, Z_D50/Y_D50 + let w50: Vector3<f64> = [0.96422, 1.0, 0.82521]; + + // Transform to LMS color space + let lms_source = mul_3x3_vector(&K_BRADFORD, &w); + let lms_d50 = mul_3x3_vector(&K_BRADFORD, &w50); + + // Check for invalid LMS values which would lead to division by zero + if lms_source.contains(&0.0) { + return Err(Error::IccInvalidWhitePoint( + wx, + wy, + "LMS components for source white point are zero, leading to division by zero." + .to_string(), + )); + } + + // Create diagonal scaling matrix in LMS space + let mut a_diag_matrix: Matrix3x3<f64> = [[0.0; 3]; 3]; + for i in 0..3 { + a_diag_matrix[i][i] = lms_d50[i] / lms_source[i]; + if !a_diag_matrix[i][i].is_finite() { + return Err(Error::IccInvalidWhitePoint( + wx, + wy, + format!("Diagonal adaptation matrix component {i} is not finite."), + )); + } + } + + // Combine transformations + let b_matrix = mul_3x3_matrix(&a_diag_matrix, &K_BRADFORD); + let final_adaptation_matrix = mul_3x3_matrix(&K_BRADFORD_INV, &b_matrix); + + Ok(final_adaptation_matrix) +} + +#[allow(clippy::too_many_arguments)] +pub(crate) fn primaries_to_xyz_d50( + rx: f32, + ry: f32, + gx: f32, + gy: f32, + bx: f32, + by: f32, + wx: f32, + wy: f32, +) -> Result<Matrix3x3<f64>, Error> { + // Get the matrix to convert RGB to XYZ, adapted to its native white point (wx, wy). + let rgb_to_xyz_native_wp_matrix = primaries_to_xyz(rx, ry, gx, gy, bx, by, wx, wy)?; + + // Get the chromatic adaptation matrix from the native white point (wx, wy) to D50. + let adaptation_to_d50_matrix = adapt_to_xyz_d50(wx, wy)?; + // This matrix converts XYZ values relative to white point (wx, wy) + // to XYZ values relative to D50. + + // Combine the matrices: M_RGBtoD50XYZ = M_AdaptToD50 * M_RGBtoNativeXYZ + // Applying M_RGBtoNativeXYZ first gives XYZ relative to native white point. + // Then applying M_AdaptToD50 converts these XYZ values to be relative to D50. + let result_matrix = mul_3x3_matrix(&adaptation_to_d50_matrix, &rgb_to_xyz_native_wp_matrix); + + Ok(result_matrix) +} + +#[allow(clippy::too_many_arguments)] +fn create_icc_rgb_matrix( + rx: f32, + ry: f32, + gx: f32, + gy: f32, + bx: f32, + by: f32, + wx: f32, + wy: f32, +) -> Result<Matrix3x3<f32>, Error> { + // TODO: think about if we need/want to change precision to f64 for some calculations here + let result_f64 = primaries_to_xyz_d50(rx, ry, gx, gy, bx, by, wx, wy)?; + Ok(std::array::from_fn(|r_idx| { + std::array::from_fn(|c_idx| result_f64[r_idx][c_idx] as f32) + })) +} + +#[derive(Clone, Debug, PartialEq)] +pub enum JxlWhitePoint { + D65, + E, + DCI, + Chromaticity { wx: f32, wy: f32 }, +} + +impl fmt::Display for JxlWhitePoint { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + JxlWhitePoint::D65 => f.write_str("D65"), + JxlWhitePoint::E => f.write_str("EER"), + JxlWhitePoint::DCI => f.write_str("DCI"), + JxlWhitePoint::Chromaticity { wx, wy } => write!(f, "{wx:.7};{wy:.7}"), + } + } +} + +impl JxlWhitePoint { + pub fn to_xy_coords(&self) -> (f32, f32) { + match self { + JxlWhitePoint::Chromaticity { wx, wy } => (*wx, *wy), + JxlWhitePoint::D65 => (0.3127, 0.3290), + JxlWhitePoint::DCI => (0.314, 0.351), + JxlWhitePoint::E => (1.0 / 3.0, 1.0 / 3.0), + } + } +} + +#[derive(Clone, Debug, PartialEq)] +pub enum JxlPrimaries { + SRGB, + BT2100, + P3, + Chromaticities { + rx: f32, + ry: f32, + gx: f32, + gy: f32, + bx: f32, + by: f32, + }, +} + +impl fmt::Display for JxlPrimaries { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + JxlPrimaries::SRGB => f.write_str("SRG"), + JxlPrimaries::BT2100 => f.write_str("202"), + JxlPrimaries::P3 => f.write_str("DCI"), + JxlPrimaries::Chromaticities { + rx, + ry, + gx, + gy, + bx, + by, + } => write!(f, "{rx:.7},{ry:.7};{gx:.7},{gy:.7};{bx:.7},{by:.7}"), + } + } +} + +impl JxlPrimaries { + pub fn to_xy_coords(&self) -> [(f32, f32); 3] { + match self { + JxlPrimaries::Chromaticities { + rx, + ry, + gx, + gy, + bx, + by, + } => [(*rx, *ry), (*gx, *gy), (*bx, *by)], + JxlPrimaries::SRGB => [ + // libjxl has these weird numbers for some reason. + (0.639_998_7, 0.330_010_15), + //(0.640, 0.330), // R + (0.300_003_8, 0.600_003_36), + //(0.300, 0.600), // G + (0.150_002_05, 0.059_997_204), + //(0.150, 0.060), // B + ], + JxlPrimaries::BT2100 => [ + (0.708, 0.292), // R + (0.170, 0.797), // G + (0.131, 0.046), // B + ], + JxlPrimaries::P3 => [ + (0.680, 0.320), // R + (0.265, 0.690), // G + (0.150, 0.060), // B + ], + } + } +} + +#[derive(Clone, Debug, PartialEq)] +pub enum JxlTransferFunction { + BT709, + Linear, + SRGB, + PQ, + DCI, + HLG, + Gamma(f32), +} + +impl fmt::Display for JxlTransferFunction { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + JxlTransferFunction::BT709 => f.write_str("709"), + JxlTransferFunction::Linear => f.write_str("Lin"), + JxlTransferFunction::SRGB => f.write_str("SRG"), + JxlTransferFunction::PQ => f.write_str("PeQ"), + JxlTransferFunction::DCI => f.write_str("DCI"), + JxlTransferFunction::HLG => f.write_str("HLG"), + JxlTransferFunction::Gamma(g) => write!(f, "g{g:.7}"), + } + } +} + +#[derive(Clone, Debug, PartialEq)] +pub enum JxlColorEncoding { + RgbColorSpace { + white_point: JxlWhitePoint, + primaries: JxlPrimaries, + transfer_function: JxlTransferFunction, + rendering_intent: RenderingIntent, + }, + GrayscaleColorSpace { + white_point: JxlWhitePoint, + transfer_function: JxlTransferFunction, + rendering_intent: RenderingIntent, + }, + XYB { + rendering_intent: RenderingIntent, + }, +} + +impl JxlColorEncoding { + pub fn from_internal(internal: &ColorEncoding) -> Result<Self> { + let rendering_intent = internal.rendering_intent; + if internal.color_space == ColorSpace::XYB { + if rendering_intent != RenderingIntent::Perceptual { + return Err(Error::InvalidRenderingIntent); + } + return Ok(Self::XYB { rendering_intent }); + } + + let white_point = match internal.white_point { + WhitePoint::D65 => JxlWhitePoint::D65, + WhitePoint::E => JxlWhitePoint::E, + WhitePoint::DCI => JxlWhitePoint::DCI, + WhitePoint::Custom => { + let (wx, wy) = internal.white.as_f32_coords(); + JxlWhitePoint::Chromaticity { wx, wy } + } + }; + let transfer_function = if internal.tf.have_gamma { + JxlTransferFunction::Gamma(internal.tf.gamma()) + } else { + match internal.tf.transfer_function { + TransferFunction::BT709 => JxlTransferFunction::BT709, + TransferFunction::Linear => JxlTransferFunction::Linear, + TransferFunction::SRGB => JxlTransferFunction::SRGB, + TransferFunction::PQ => JxlTransferFunction::PQ, + TransferFunction::DCI => JxlTransferFunction::DCI, + TransferFunction::HLG => JxlTransferFunction::HLG, + TransferFunction::Unknown => { + return Err(Error::InvalidColorEncoding); + } + } + }; + + if internal.color_space == ColorSpace::Gray { + return Ok(Self::GrayscaleColorSpace { + white_point, + transfer_function, + rendering_intent, + }); + } + + let primaries = match internal.primaries { + Primaries::SRGB => JxlPrimaries::SRGB, + Primaries::BT2100 => JxlPrimaries::BT2100, + Primaries::P3 => JxlPrimaries::P3, + Primaries::Custom => { + let (rx, ry) = internal.custom_primaries[0].as_f32_coords(); + let (gx, gy) = internal.custom_primaries[1].as_f32_coords(); + let (bx, by) = internal.custom_primaries[2].as_f32_coords(); + JxlPrimaries::Chromaticities { + rx, + ry, + gx, + gy, + bx, + by, + } + } + }; + + match internal.color_space { + ColorSpace::Gray | ColorSpace::XYB => unreachable!(), + ColorSpace::RGB => Ok(Self::RgbColorSpace { + white_point, + primaries, + transfer_function, + rendering_intent, + }), + ColorSpace::Unknown => Err(Error::InvalidColorSpace), + } + } + + fn create_icc_cicp_tag_data(&self, tags_data: &mut Vec<u8>) -> Result<Option<TagInfo>, Error> { + let JxlColorEncoding::RgbColorSpace { + white_point, + primaries, + transfer_function, + .. + } = self + else { + return Ok(None); + }; + + // Determine the CICP value for primaries. + let primaries_val: u8 = match (white_point, primaries) { + (JxlWhitePoint::D65, JxlPrimaries::SRGB) => 1, + (JxlWhitePoint::D65, JxlPrimaries::BT2100) => 9, + (JxlWhitePoint::D65, JxlPrimaries::P3) => 12, + (JxlWhitePoint::DCI, JxlPrimaries::P3) => 11, + _ => return Ok(None), + }; + + let tf_val = match transfer_function { + JxlTransferFunction::BT709 => 1, + JxlTransferFunction::Linear => 8, + JxlTransferFunction::SRGB => 13, + JxlTransferFunction::PQ => 16, + JxlTransferFunction::DCI => 17, + JxlTransferFunction::HLG => 18, + // Custom gamma cannot be represented. + JxlTransferFunction::Gamma(_) => return Ok(None), + }; + + let signature = b"cicp"; + let start_offset = tags_data.len() as u32; + tags_data.extend_from_slice(signature); + let data_len = tags_data.len(); + tags_data.resize(tags_data.len() + 4, 0); + write_u32_be(tags_data, data_len, 0)?; + tags_data.push(primaries_val); + tags_data.push(tf_val); + // Matrix Coefficients (RGB is non-constant luminance) + tags_data.push(0); + // Video Full Range Flag + tags_data.push(1); + + Ok(Some(TagInfo { + signature: *signature, + offset_in_tags_blob: start_offset, + size_unpadded: 12, + })) + } + + fn can_tone_map_for_icc(&self) -> bool { + let JxlColorEncoding::RgbColorSpace { + white_point, + primaries, + transfer_function, + .. + } = self + else { + return false; + }; + // This function determines if an ICC profile can be used for tone mapping. + // The logic is ported from the libjxl `CanToneMap` function. + // The core idea is that if the color space can be represented by a CICP tag + // in the ICC profile, then there's more freedom to use other parts of the + // profile (like the A2B0 LUT) for tone mapping. Otherwise, the profile must + // unambiguously describe the color space. + + // The conditions for being able to tone map are: + // 1. The color space must be RGB. + // 2. The transfer function must be either PQ (Perceptual Quantizer) or HLG (Hybrid Log-Gamma). + // 3. The combination of primaries and white point must be one that is commonly + // describable by a standard CICP value. This includes: + // a) P3 primaries with either a D65 or DCI white point. + // b) Any non-custom primaries, as long as the white point is D65. + + if let JxlPrimaries::Chromaticities { .. } = primaries { + return false; + } + + matches!( + transfer_function, + JxlTransferFunction::PQ | JxlTransferFunction::HLG + ) && (*white_point == JxlWhitePoint::D65 + || (*white_point == JxlWhitePoint::DCI && *primaries == JxlPrimaries::P3)) + } + + pub fn get_color_encoding_description(&self) -> String { + // Handle special known color spaces first + if let Some(common_name) = match self { + JxlColorEncoding::RgbColorSpace { + white_point: JxlWhitePoint::D65, + primaries: JxlPrimaries::SRGB, + transfer_function: JxlTransferFunction::SRGB, + rendering_intent: RenderingIntent::Perceptual, + } => Some("sRGB"), + JxlColorEncoding::RgbColorSpace { + white_point: JxlWhitePoint::D65, + primaries: JxlPrimaries::P3, + transfer_function: JxlTransferFunction::SRGB, + rendering_intent: RenderingIntent::Perceptual, + } => Some("DisplayP3"), + JxlColorEncoding::RgbColorSpace { + white_point: JxlWhitePoint::D65, + primaries: JxlPrimaries::BT2100, + transfer_function: JxlTransferFunction::PQ, + rendering_intent: RenderingIntent::Relative, + } => Some("Rec2100PQ"), + JxlColorEncoding::RgbColorSpace { + white_point: JxlWhitePoint::D65, + primaries: JxlPrimaries::BT2100, + transfer_function: JxlTransferFunction::HLG, + rendering_intent: RenderingIntent::Relative, + } => Some("Rec2100HLG"), + _ => None, + } { + return common_name.to_string(); + } + + // Build the string part by part for other case + let mut d = String::with_capacity(64); + + match self { + JxlColorEncoding::RgbColorSpace { + white_point, + primaries, + transfer_function, + rendering_intent, + } => { + d.push_str("RGB_"); + d.push_str(&white_point.to_string()); + d.push('_'); + d.push_str(&primaries.to_string()); + d.push('_'); + d.push_str(&rendering_intent.to_string()); + d.push('_'); + d.push_str(&transfer_function.to_string()); + } + JxlColorEncoding::GrayscaleColorSpace { + white_point, + transfer_function, + rendering_intent, + } => { + d.push_str("Gra_"); + d.push_str(&white_point.to_string()); + d.push('_'); + d.push_str(&rendering_intent.to_string()); + d.push('_'); + d.push_str(&transfer_function.to_string()); + } + JxlColorEncoding::XYB { rendering_intent } => { + d.push_str("XYB_"); + d.push_str(&rendering_intent.to_string()); + } + } + + d + } + + fn create_icc_header(&self) -> Result<Vec<u8>, Error> { + let mut header_data = vec![0u8; 128]; + + // Profile size - To be filled in at the end of profile creation. + write_u32_be(&mut header_data, 0, 0)?; + const CMM_TAG: &str = "jxl "; + // CMM Type + write_icc_tag(&mut header_data, 4, CMM_TAG)?; + + // Profile version - ICC v4.4 (0x04400000) + // Conformance tests have v4.3, libjxl produces v4.4 + write_u32_be(&mut header_data, 8, 0x04400000u32)?; + + let profile_class_str = match self { + JxlColorEncoding::XYB { .. } => "scnr", + _ => "mntr", + }; + write_icc_tag(&mut header_data, 12, profile_class_str)?; + + // Data color space + let data_color_space_str = match self { + JxlColorEncoding::GrayscaleColorSpace { .. } => "GRAY", + _ => "RGB ", + }; + write_icc_tag(&mut header_data, 16, data_color_space_str)?; + + // PCS - Profile Connection Space + // Corresponds to: if (kEnable3DToneMapping && CanToneMap(c)) + // Assuming kEnable3DToneMapping is true for this port for now. + const K_ENABLE_3D_ICC_TONEMAPPING: bool = true; + if K_ENABLE_3D_ICC_TONEMAPPING && self.can_tone_map_for_icc() { + write_icc_tag(&mut header_data, 20, "Lab ")?; + } else { + write_icc_tag(&mut header_data, 20, "XYZ ")?; + } + + // Date and Time - Placeholder values from libjxl + write_u16_be(&mut header_data, 24, 2019)?; // Year + write_u16_be(&mut header_data, 26, 12)?; // Month + write_u16_be(&mut header_data, 28, 1)?; // Day + write_u16_be(&mut header_data, 30, 0)?; // Hours + write_u16_be(&mut header_data, 32, 0)?; // Minutes + write_u16_be(&mut header_data, 34, 0)?; // Seconds + + write_icc_tag(&mut header_data, 36, "acsp")?; + write_icc_tag(&mut header_data, 40, "APPL")?; + + // Profile flags + write_u32_be(&mut header_data, 44, 0)?; + // Device manufacturer + write_u32_be(&mut header_data, 48, 0)?; + // Device model + write_u32_be(&mut header_data, 52, 0)?; + // Device attributes + write_u32_be(&mut header_data, 56, 0)?; + write_u32_be(&mut header_data, 60, 0)?; + + // Rendering Intent + let rendering_intent = match self { + JxlColorEncoding::RgbColorSpace { + rendering_intent, .. + } + | JxlColorEncoding::GrayscaleColorSpace { + rendering_intent, .. + } + | JxlColorEncoding::XYB { rendering_intent } => rendering_intent, + }; + write_u32_be(&mut header_data, 64, *rendering_intent as u32)?; + + // Whitepoint is fixed to D50 for ICC. + write_u32_be(&mut header_data, 68, 0x0000F6D6)?; + write_u32_be(&mut header_data, 72, 0x00010000)?; + write_u32_be(&mut header_data, 76, 0x0000D32D)?; + + // Profile Creator + write_icc_tag(&mut header_data, 80, CMM_TAG)?; + + // Profile ID (MD5 checksum) (offset 84) - 16 bytes. + // This is calculated at the end of profile creation and written here. + + // Reserved (offset 100-127) - already zeroed here. + + Ok(header_data) + } + + pub fn maybe_create_profile(&self) -> Result<Option<Vec<u8>>, Error> { + if let JxlColorEncoding::XYB { rendering_intent } = self + && *rendering_intent != RenderingIntent::Perceptual + { + return Err(Error::InvalidRenderingIntent); + } + let header = self.create_icc_header()?; + let mut tags_data: Vec<u8> = Vec::new(); + let mut collected_tags: Vec<TagInfo> = Vec::new(); + + // Create 'desc' (ProfileDescription) tag + let description_string = self.get_color_encoding_description(); + + let desc_tag_start_offset = tags_data.len() as u32; // 0 at this point ... + create_icc_mluc_tag(&mut tags_data, &description_string)?; + let desc_tag_unpadded_size = (tags_data.len() as u32) - desc_tag_start_offset; + pad_to_4_byte_boundary(&mut tags_data); + collected_tags.push(TagInfo { + signature: *b"desc", + offset_in_tags_blob: desc_tag_start_offset, + size_unpadded: desc_tag_unpadded_size, + }); + + // Create 'cprt' (Copyright) tag + let copyright_string = "CC0"; + let cprt_tag_start_offset = tags_data.len() as u32; + create_icc_mluc_tag(&mut tags_data, copyright_string)?; + let cprt_tag_unpadded_size = (tags_data.len() as u32) - cprt_tag_start_offset; + pad_to_4_byte_boundary(&mut tags_data); + collected_tags.push(TagInfo { + signature: *b"cprt", + offset_in_tags_blob: cprt_tag_start_offset, + size_unpadded: cprt_tag_unpadded_size, + }); + + match self { + JxlColorEncoding::GrayscaleColorSpace { white_point, .. } => { + let (wx, wy) = white_point.to_xy_coords(); + collected_tags.push(create_icc_xyz_tag( + &mut tags_data, + &cie_xyz_from_white_cie_xy(wx, wy)?, + )?); + } + _ => { + // Ok, in this case we will add the chad tag below + const D50: [f32; 3] = [0.964203f32, 1.0, 0.824905]; + collected_tags.push(create_icc_xyz_tag(&mut tags_data, &D50)?); + } + } + pad_to_4_byte_boundary(&mut tags_data); + if !matches!(self, JxlColorEncoding::GrayscaleColorSpace { .. }) { + let (wx, wy) = match self { + JxlColorEncoding::GrayscaleColorSpace { .. } => unreachable!(), + JxlColorEncoding::RgbColorSpace { white_point, .. } => white_point.to_xy_coords(), + JxlColorEncoding::XYB { .. } => JxlWhitePoint::D65.to_xy_coords(), + }; + let chad_matrix_f64 = adapt_to_xyz_d50(wx, wy)?; + let chad_matrix = std::array::from_fn(|r_idx| { + std::array::from_fn(|c_idx| chad_matrix_f64[r_idx][c_idx] as f32) + }); + collected_tags.push(create_icc_chad_tag(&mut tags_data, &chad_matrix)?); + pad_to_4_byte_boundary(&mut tags_data); + } + + if let JxlColorEncoding::RgbColorSpace { + white_point, + primaries, + .. + } = self + { + if let Some(tag_info) = self.create_icc_cicp_tag_data(&mut tags_data)? { + collected_tags.push(tag_info); + // Padding here not necessary, since we add 12 bytes to already 4-byte aligned + // buffer + // pad_to_4_byte_boundary(&mut tags_data); + } + + // Get colorant and white point coordinates to build the conversion matrix. + let primaries_coords = primaries.to_xy_coords(); + let (rx, ry) = primaries_coords[0]; + let (gx, gy) = primaries_coords[1]; + let (bx, by) = primaries_coords[2]; + let (wx, wy) = white_point.to_xy_coords(); + + // Calculate the RGB to XYZD50 matrix. + let m = create_icc_rgb_matrix(rx, ry, gx, gy, bx, by, wx, wy)?; + + // Extract the columns, which are the XYZ values for the R, G, and B primaries. + let r_xyz = [m[0][0], m[1][0], m[2][0]]; + let g_xyz = [m[0][1], m[1][1], m[2][1]]; + let b_xyz = [m[0][2], m[1][2], m[2][2]]; + + // Helper to create the raw data for any 'XYZ ' type tag. + let create_xyz_type_tag_data = + |tags: &mut Vec<u8>, xyz: &[f32; 3]| -> Result<u32, Error> { + let start_offset = tags.len(); + // The tag *type* is 'XYZ ' for all three + tags.extend_from_slice(b"XYZ "); + tags.extend_from_slice(&0u32.to_be_bytes()); + for &val in xyz { + append_s15_fixed_16(tags, val)?; + } + Ok((tags.len() - start_offset) as u32) + }; + + // Create the 'rXYZ' tag. + let r_xyz_tag_start_offset = tags_data.len() as u32; + let r_xyz_tag_unpadded_size = create_xyz_type_tag_data(&mut tags_data, &r_xyz)?; + pad_to_4_byte_boundary(&mut tags_data); + collected_tags.push(TagInfo { + signature: *b"rXYZ", // Making the *signature* is unique. + offset_in_tags_blob: r_xyz_tag_start_offset, + size_unpadded: r_xyz_tag_unpadded_size, + }); + + // Create the 'gXYZ' tag. + let g_xyz_tag_start_offset = tags_data.len() as u32; + let g_xyz_tag_unpadded_size = create_xyz_type_tag_data(&mut tags_data, &g_xyz)?; + pad_to_4_byte_boundary(&mut tags_data); + collected_tags.push(TagInfo { + signature: *b"gXYZ", + offset_in_tags_blob: g_xyz_tag_start_offset, + size_unpadded: g_xyz_tag_unpadded_size, + }); + + // Create the 'bXYZ' tag. + let b_xyz_tag_start_offset = tags_data.len() as u32; + let b_xyz_tag_unpadded_size = create_xyz_type_tag_data(&mut tags_data, &b_xyz)?; + pad_to_4_byte_boundary(&mut tags_data); + collected_tags.push(TagInfo { + signature: *b"bXYZ", + offset_in_tags_blob: b_xyz_tag_start_offset, + size_unpadded: b_xyz_tag_unpadded_size, + }); + } + if self.can_tone_map_for_icc() { + todo!("implement A2B0 and B2A0 tags when being able to tone map") + } else { + match self { + JxlColorEncoding::XYB { .. } => todo!("implement A2B0 and B2A0 tags"), + JxlColorEncoding::RgbColorSpace { + transfer_function, .. + } + | JxlColorEncoding::GrayscaleColorSpace { + transfer_function, .. + } => { + let trc_tag_start_offset = tags_data.len() as u32; + let trc_tag_unpadded_size = match transfer_function { + JxlTransferFunction::Gamma(g) => { + // Type 0 parametric curve: Y = X^gamma + let gamma = 1.0 / g; + create_icc_curv_para_tag(&mut tags_data, &[gamma], 0)? + } + JxlTransferFunction::SRGB => { + // Type 3 parametric curve for sRGB standard. + const PARAMS: [f32; 5] = + [2.4, 1.0 / 1.055, 0.055 / 1.055, 1.0 / 12.92, 0.04045]; + create_icc_curv_para_tag(&mut tags_data, &PARAMS, 3)? + } + JxlTransferFunction::BT709 => { + // Type 3 parametric curve for BT.709 standard. + const PARAMS: [f32; 5] = + [1.0 / 0.45, 1.0 / 1.099, 0.099 / 1.099, 1.0 / 4.5, 0.081]; + create_icc_curv_para_tag(&mut tags_data, &PARAMS, 3)? + } + JxlTransferFunction::Linear => { + // Type 3 can also represent a linear response (gamma=1.0). + const PARAMS: [f32; 5] = [1.0, 1.0, 0.0, 1.0, 0.0]; + create_icc_curv_para_tag(&mut tags_data, &PARAMS, 3)? + } + JxlTransferFunction::DCI => { + // Type 3 can also represent a pure power curve (gamma=2.6). + const PARAMS: [f32; 5] = [2.6, 1.0, 0.0, 1.0, 0.0]; + create_icc_curv_para_tag(&mut tags_data, &PARAMS, 3)? + } + JxlTransferFunction::HLG | JxlTransferFunction::PQ => { + let params = create_table_curve(64, transfer_function, false)?; + create_icc_curv_para_tag(&mut tags_data, params.as_slice(), 3)? + } + }; + pad_to_4_byte_boundary(&mut tags_data); + + match self { + JxlColorEncoding::GrayscaleColorSpace { .. } => { + // Grayscale profiles use a single 'kTRC' tag. + collected_tags.push(TagInfo { + signature: *b"kTRC", + offset_in_tags_blob: trc_tag_start_offset, + size_unpadded: trc_tag_unpadded_size, + }); + } + _ => { + // For RGB, rTRC, gTRC, and bTRC all point to the same curve data, + // an optimization to keep the profile size small. + collected_tags.push(TagInfo { + signature: *b"rTRC", + offset_in_tags_blob: trc_tag_start_offset, + size_unpadded: trc_tag_unpadded_size, + }); + collected_tags.push(TagInfo { + signature: *b"gTRC", + offset_in_tags_blob: trc_tag_start_offset, // Same offset + size_unpadded: trc_tag_unpadded_size, // Same size + }); + collected_tags.push(TagInfo { + signature: *b"bTRC", + offset_in_tags_blob: trc_tag_start_offset, // Same offset + size_unpadded: trc_tag_unpadded_size, // Same size + }); + } + } + } + } + } + + // Construct the Tag Table bytes + let mut tag_table_bytes: Vec<u8> = Vec::new(); + // First, the number of tags (u32) + tag_table_bytes.extend_from_slice(&(collected_tags.len() as u32).to_be_bytes()); + + let header_size = header.len() as u32; + // Each entry in the tag table on disk is 12 bytes: signature (4), offset (4), size (4) + let tag_table_on_disk_size = 4 + (collected_tags.len() as u32 * 12); + + for tag_info in &collected_tags { + tag_table_bytes.extend_from_slice(&tag_info.signature); + // The offset in the tag table is absolute from the start of the ICC profile file + let final_profile_offset_for_tag = + header_size + tag_table_on_disk_size + tag_info.offset_in_tags_blob; + tag_table_bytes.extend_from_slice(&final_profile_offset_for_tag.to_be_bytes()); + // In https://www.color.org/specification/ICC.1-2022-05.pdf, section 7.3.5 reads: + // + // "The value of the tag data element size shall be the number of actual data + // bytes and shall not include any padding at the end of the tag data element." + // + // The reference from conformance tests and libjxl use the padded size here instead. + + tag_table_bytes.extend_from_slice(&tag_info.size_unpadded.to_be_bytes()); + // In order to get byte_exact the same output as libjxl, remove the line above + // and uncomment the lines below + // let padded_size = tag_info.size_unpadded.next_multiple_of(4); + // tag_table_bytes.extend_from_slice(&padded_size.to_be_bytes()); + } + + // Assemble the final ICC profile parts: header + tag_table + tags_data + let mut final_icc_profile_data: Vec<u8> = + Vec::with_capacity(header.len() + tag_table_bytes.len() + tags_data.len()); + final_icc_profile_data.extend_from_slice(&header); + final_icc_profile_data.extend_from_slice(&tag_table_bytes); + final_icc_profile_data.extend_from_slice(&tags_data); + + // Update the profile size in the header (at offset 0) + let total_profile_size = final_icc_profile_data.len() as u32; + write_u32_be(&mut final_icc_profile_data, 0, total_profile_size)?; + + // Assemble the final ICC profile parts: header + tag_table + tags_data + let mut final_icc_profile_data: Vec<u8> = + Vec::with_capacity(header.len() + tag_table_bytes.len() + tags_data.len()); + final_icc_profile_data.extend_from_slice(&header); + final_icc_profile_data.extend_from_slice(&tag_table_bytes); + final_icc_profile_data.extend_from_slice(&tags_data); + + // Update the profile size in the header (at offset 0) + let total_profile_size = final_icc_profile_data.len() as u32; + write_u32_be(&mut final_icc_profile_data, 0, total_profile_size)?; + + // The MD5 checksum (Profile ID) must be computed on the profile with + // specific header fields zeroed out, as per the ICC specification. + let mut profile_for_checksum = final_icc_profile_data.clone(); + + if profile_for_checksum.len() >= 84 { + // Zero out Profile Flags at offset 44. + profile_for_checksum[44..48].fill(0); + // Zero out Rendering Intent at offset 64. + profile_for_checksum[64..68].fill(0); + // The Profile ID field at offset 84 is already zero at this stage. + } + + // Compute the MD5 hash on the modified profile data. + let checksum = compute_md5(&profile_for_checksum); + + // Write the 16-byte checksum into the "Profile ID" field of the *original* + // profile data buffer, starting at offset 84. + if final_icc_profile_data.len() >= 100 { + final_icc_profile_data[84..100].copy_from_slice(&checksum); + } + + Ok(Some(final_icc_profile_data)) + } + + pub fn srgb(grayscale: bool) -> Self { + if grayscale { + JxlColorEncoding::GrayscaleColorSpace { + white_point: JxlWhitePoint::D65, + transfer_function: JxlTransferFunction::SRGB, + rendering_intent: RenderingIntent::Relative, + } + } else { + JxlColorEncoding::RgbColorSpace { + white_point: JxlWhitePoint::D65, + primaries: JxlPrimaries::SRGB, + transfer_function: JxlTransferFunction::SRGB, + rendering_intent: RenderingIntent::Relative, + } + } + } +} + +#[derive(Clone)] +pub enum JxlColorProfile { + Icc(Vec<u8>), + Simple(JxlColorEncoding), +} + +impl JxlColorProfile { + pub fn as_icc(&self) -> Cow<'_, Vec<u8>> { + match self { + Self::Icc(x) => Cow::Borrowed(x), + Self::Simple(encoding) => Cow::Owned(encoding.maybe_create_profile().unwrap().unwrap()), + } + } +} + +// TODO: do we want/need to return errors from here? +pub trait JxlCmsTransformer { + /// Runs a single transform. The buffers each contain `num_pixels` x `num_channels` interleaved + /// floating point (0..1) samples, where `num_channels` is the number of color channels of + /// their respective color profiles. For CMYK data, 0 represents the maximum amount of ink + /// while 1 represents no ink. + fn do_transform(&mut self, input: &[f32], output: &mut [f32]); + + /// Runs a single transform in-place. The buffer contains `num_pixels` x `num_channels` + /// interleaved floating point (0..1) samples, where `num_channels` is the number of color + /// channels of the input and output color profiles. For CMYK data, 0 represents the maximum + /// amount of ink while 1 represents no ink. + fn do_transform_inplace(&mut self, inout: &mut [f32]); +} + +pub trait JxlCms { + /// Parses an ICC profile, returning a ColorEncoding and whether the ICC profile represents a + /// CMYK profile. + fn parse_icc(&mut self, icc: &[u8]) -> Result<(ColorEncoding, bool)>; + + /// Initializes `n` transforms (different transforms might be used in parallel) to + /// convert from color space `input` to colorspace `output`, assuming an intensity of 1.0 for + /// non-absolute luminance colorspaces of `intensity_target`. + /// It is an error to not return `n` transforms. + fn initialize_transforms( + &mut self, + n: usize, + max_pixels_per_transform: usize, + input: JxlColorProfile, + output: JxlColorProfile, + intensity_target: f32, + ) -> Result<Vec<Box<dyn JxlCmsTransformer>>>; +} + +/// Writes a u32 value in big-endian format to the slice at the given position. +fn write_u32_be(slice: &mut [u8], pos: usize, value: u32) -> Result<(), Error> { + if pos.checked_add(4).is_none_or(|end| end > slice.len()) { + return Err(Error::IccWriteOutOfBounds); + } + slice[pos..pos + 4].copy_from_slice(&value.to_be_bytes()); + Ok(()) +} + +/// Writes a u16 value in big-endian format to the slice at the given position. +fn write_u16_be(slice: &mut [u8], pos: usize, value: u16) -> Result<(), Error> { + if pos.checked_add(2).is_none_or(|end| end > slice.len()) { + return Err(Error::IccWriteOutOfBounds); + } + slice[pos..pos + 2].copy_from_slice(&value.to_be_bytes()); + Ok(()) +} + +/// Writes a 4-character ASCII tag string to the slice at the given position. +fn write_icc_tag(slice: &mut [u8], pos: usize, tag_str: &str) -> Result<(), Error> { + if tag_str.len() != 4 || !tag_str.is_ascii() { + return Err(Error::IccInvalidTagString(tag_str.to_string())); + } + if pos.checked_add(4).is_none_or(|end| end > slice.len()) { + return Err(Error::IccWriteOutOfBounds); + } + slice[pos..pos + 4].copy_from_slice(tag_str.as_bytes()); + Ok(()) +} + +/// Creates an ICC 'mluc' tag with a single "enUS" record. +/// +/// The input `text` must be ASCII, as it will be encoded as UTF-16BE by prepending +/// a null byte to each ASCII character. +fn create_icc_mluc_tag(tags: &mut Vec<u8>, text: &str) -> Result<(), Error> { + // libjxl comments that "The input text must be ASCII". + // We enforce this. + if !text.is_ascii() { + return Err(Error::IccMlucTextNotAscii(text.to_string())); + } + // Tag signature 'mluc' (4 bytes) + tags.extend_from_slice(b"mluc"); + // Reserved, must be 0 (4 bytes) + tags.extend_from_slice(&0u32.to_be_bytes()); + // Number of records (u32, 4 bytes) - Hardcoded to 1. + tags.extend_from_slice(&1u32.to_be_bytes()); + // Record size (u32, 4 bytes) - Each record descriptor is 12 bytes. + // (Language Code [2] + Country Code [2] + String Length [4] + String Offset [4]) + tags.extend_from_slice(&12u32.to_be_bytes()); + // Language Code (2 bytes) - "en" for English + tags.extend_from_slice(b"en"); + // Country Code (2 bytes) - "US" for United States + tags.extend_from_slice(b"US"); + // Length of the string (u32, 4 bytes) + // For ASCII text encoded as UTF-16BE, each char becomes 2 bytes. + let string_actual_byte_length = text.len() * 2; + tags.extend_from_slice(&(string_actual_byte_length as u32).to_be_bytes()); + // Offset of the string (u32, 4 bytes) + // The string data for this record starts at offset 28. + tags.extend_from_slice(&28u32.to_be_bytes()); + // The actual string data, encoded as UTF-16BE. + // For ASCII char 'X', UTF-16BE is 0x00 0x58. + for ascii_char_code in text.as_bytes() { + tags.push(0u8); + tags.push(*ascii_char_code); + } + + Ok(()) +} + +struct TagInfo { + signature: [u8; 4], + // Offset of this tag's data relative to the START of the `tags_data` block + offset_in_tags_blob: u32, + // Unpadded size of this tag's actual data content. + size_unpadded: u32, +} + +fn pad_to_4_byte_boundary(data: &mut Vec<u8>) { + data.resize(data.len().next_multiple_of(4), 0u8); +} + +/// Converts an f32 to s15Fixed16 format and appends it as big-endian bytes. +/// s15Fixed16 is a signed 32-bit number with 1 sign bit, 15 integer bits, +/// and 16 fractional bits. +fn append_s15_fixed_16(tags_data: &mut Vec<u8>, value: f32) -> Result<(), Error> { + // In libjxl, the following specific range check is used: (-32767.995f <= value) && (value <= 32767.995f) + // This is slightly tighter than the theoretical max positive s15.16 value. + // We replicate this for consistency. + if !(value.is_finite() && (-32767.995..=32767.995).contains(&value)) { + return Err(Error::IccValueOutOfRangeS15Fixed16(value)); + } + + // Multiply by 2^16 and round to nearest integer + let scaled_value = (value * 65536.0).round(); + // Cast to i32 for correct two's complement representation + let int_value = scaled_value as i32; + tags_data.extend_from_slice(&int_value.to_be_bytes()); + Ok(()) +} + +/// Creates the data for an ICC 'XYZ ' tag and appends it to `tags_data`. +/// The 'XYZ ' tag contains three s15Fixed16Number values. +fn create_icc_xyz_tag(tags_data: &mut Vec<u8>, xyz_color: &[f32; 3]) -> Result<TagInfo, Error> { + // Tag signature 'XYZ ' (4 bytes, note the trailing space) + let start_offset = tags_data.len() as u32; + let signature = b"XYZ "; + tags_data.extend_from_slice(signature); + + // Reserved, must be 0 (4 bytes) + tags_data.extend_from_slice(&0u32.to_be_bytes()); + + // XYZ data (3 * s15Fixed16Number = 3 * 4 bytes) + for &val in xyz_color { + append_s15_fixed_16(tags_data, val)?; + } + + Ok(TagInfo { + signature: *b"wtpt", + offset_in_tags_blob: start_offset, + size_unpadded: (tags_data.len() as u32) - start_offset, + }) +} + +fn create_icc_chad_tag( + tags_data: &mut Vec<u8>, + chad_matrix: &Matrix3x3<f32>, +) -> Result<TagInfo, Error> { + // The tag type signature "sf32" (4 bytes). + let signature = b"sf32"; + let start_offset = tags_data.len() as u32; + tags_data.extend_from_slice(signature); + + // A reserved field (4 bytes), which must be set to 0. + tags_data.extend_from_slice(&0u32.to_be_bytes()); + + // The 9 matrix elements as s15Fixed16Number values. + // m[0][0], m[0][1], m[0][2], m[1][0], ..., m[2][2] + for row_array in chad_matrix.iter() { + for &value in row_array.iter() { + append_s15_fixed_16(tags_data, value)?; + } + } + Ok(TagInfo { + signature: *b"chad", + offset_in_tags_blob: start_offset, + size_unpadded: (tags_data.len() as u32) - start_offset, + }) +} + +/// Converts CIE xy white point coordinates to CIE XYZ values (Y is normalized to 1.0). +fn cie_xyz_from_white_cie_xy(wx: f32, wy: f32) -> Result<[f32; 3], Error> { + // Check for wy being too close to zero to prevent division by zero or extreme values. + if wy.abs() < 1e-12 { + return Err(Error::IccInvalidWhitePointY(wy)); + } + let factor = 1.0 / wy; + let x_val = wx * factor; + let y_val = 1.0f32; + let z_val = (1.0 - wx - wy) * factor; + Ok([x_val, y_val, z_val]) +} + +/// Creates the data for an ICC `para` (parametricCurveType) tag. +/// It writes `12 + 4 * params.len()` bytes. +fn create_icc_curv_para_tag( + tags_data: &mut Vec<u8>, + params: &[f32], + curve_type: u16, +) -> Result<u32, Error> { + let start_offset = tags_data.len(); + // Tag type 'para' (4 bytes) + tags_data.extend_from_slice(b"para"); + // Reserved, must be 0 (4 bytes) + tags_data.extend_from_slice(&0u32.to_be_bytes()); + // Function type (u16, 2 bytes) + tags_data.extend_from_slice(&curve_type.to_be_bytes()); + // Reserved, must be 0 (u16, 2 bytes) + tags_data.extend_from_slice(&0u16.to_be_bytes()); + // Parameters (s15Fixed16Number each) + for ¶m in params { + append_s15_fixed_16(tags_data, param)?; + } + Ok((tags_data.len() - start_offset) as u32) +} + +fn display_from_encoded_pq(display_intensity_target: f32, mut e: f64) -> f64 { + const M1: f64 = 2610.0 / 16384.0; + const M2: f64 = (2523.0 / 4096.0) * 128.0; + const C1: f64 = 3424.0 / 4096.0; + const C2: f64 = (2413.0 / 4096.0) * 32.0; + const C3: f64 = (2392.0 / 4096.0) * 32.0; + // Handle the zero case directly. + if e == 0.0 { + return 0.0; + } + + // Handle negative inputs by using their absolute + // value for the calculation and reapplying the sign at the end. + let original_sign = e.signum(); + e = e.abs(); + + // Core PQ EOTF formula from ST 2084. + let xp = e.powf(1.0 / M2); + let num = (xp - C1).max(0.0); + let den = C2 - C3 * xp; + + // In release builds, a zero denominator would lead to `inf` or `NaN`, + // which is handled by the assertion below. For valid inputs (e in [0,1]), + // the denominator is always positive. + debug_assert!(den != 0.0, "PQ transfer function denominator is zero."); + + let d = (num / den).powf(1.0 / M1); + + // The result `d` should always be non-negative for non-negative inputs. + debug_assert!( + d >= 0.0, + "PQ intermediate value `d` should not be negative." + ); + + // The libjxl implementation includes a scaling factor. Note that `d` represents + // a value normalized to a 10,000 nit peak. + let scaled_d = d * (10000.0 / display_intensity_target as f64); + + // Re-apply the original sign. + scaled_d.copysign(original_sign) +} + +/// TF_HLG_Base class for BT.2100 HLG. +/// +/// This struct provides methods to convert between non-linear encoded HLG signals +/// and linear display-referred light, following the definitions in BT.2100-2. +/// +/// - **"display"**: linear light, normalized to [0, 1]. +/// - **"encoded"**: a non-linear HLG signal, nominally in [0, 1]. +/// - **"scene"**: scene-referred linear light, normalized to [0, 1]. +/// +/// The functions are designed to be unbounded to handle inputs outside the +/// nominal [0, 1] range, which can occur during color space conversions. Negative +/// inputs are handled by mirroring the function (`f(-x) = -f(x)`). +#[allow(non_camel_case_types)] +struct TF_HLG; + +impl TF_HLG { + // Constants for the HLG formula, as defined in BT.2100. + const A: f64 = 0.17883277; + const RA: f64 = 1.0 / Self::A; + const B: f64 = 1.0 - 4.0 * Self::A; + const C: f64 = 0.5599107295; + const INV_12: f64 = 1.0 / 12.0; + + /// Converts a non-linear encoded signal to a linear display value (EOTF). + /// + /// This corresponds to `DisplayFromEncoded(e) = OOTF(InvOETF(e))`. + /// Since the OOTF is simplified to an identity function, this is equivalent + /// to calling `inv_oetf(e)`. + #[inline] + fn display_from_encoded(e: f64) -> f64 { + Self::inv_oetf(e) + } + + /// Converts a linear display value to a non-linear encoded signal (inverse EOTF). + /// + /// This corresponds to `EncodedFromDisplay(d) = OETF(InvOOTF(d))`. + /// Since the InvOOTF is an identity function, this is equivalent to `oetf(d)`. + #[inline] + #[allow(dead_code)] + fn encoded_from_display(d: f64) -> f64 { + Self::oetf(d) + } + + /// The private HLG OETF, converting scene-referred light to a non-linear signal. + fn oetf(mut s: f64) -> f64 { + if s == 0.0 { + return 0.0; + } + let original_sign = s.signum(); + s = s.abs(); + + let e = if s <= Self::INV_12 { + (3.0 * s).sqrt() + } else { + Self::A * (12.0 * s - Self::B).ln() + Self::C + }; + + // The result should be positive for positive inputs. + debug_assert!(e > 0.0); + + e.copysign(original_sign) + } + + /// The private HLG inverse OETF, converting a non-linear signal back to scene-referred light. + fn inv_oetf(mut e: f64) -> f64 { + if e == 0.0 { + return 0.0; + } + let original_sign = e.signum(); + e = e.abs(); + + let s = if e <= 0.5 { + // The `* (1.0 / 3.0)` is slightly more efficient than `/ 3.0`. + e * e * (1.0 / 3.0) + } else { + (((e - Self::C) * Self::RA).exp() + Self::B) * Self::INV_12 + }; + + // The result should be non-negative for non-negative inputs. + debug_assert!(s >= 0.0); + + s.copysign(original_sign) + } +} + +/// Creates a lookup table for an ICC `curv` tag from a transfer function. +/// +/// This function generates a vector of 16-bit integers representing the response +/// of the HLG or PQ electro-optical transfer functions (EOTF). +/// +/// ### Arguments +/// * `n` - The number of entries in the lookup table. Must not exceed 4096. +/// * `tf` - The transfer function to model, either `TransferFunction::HLG` or `TransferFunction::PQ`. +/// * `tone_map` - A boolean to enable tone mapping for PQ curves. Currently a stub. +/// +/// ### Returns +/// A `Result` containing the `Vec<f32>` lookup table or an `Error`. +fn create_table_curve( + n: usize, + tf: &JxlTransferFunction, + tone_map: bool, +) -> Result<Vec<f32>, Error> { + // ICC Specification (v4.4, section 10.6) for `curveType` with `curv` + // processing elements states the table can have at most 4096 entries. + if n > 4096 { + return Err(Error::IccTableSizeExceeded(n)); + } + + if !matches!(tf, JxlTransferFunction::PQ | JxlTransferFunction::HLG) { + return Err(Error::IccUnsupportedTransferFunction); + } + + // The peak luminance for PQ decoding, as specified in the original C++ code. + const PQ_INTENSITY_TARGET: f64 = 10000.0; + // The target peak luminance for SDR, used if tone mapping is applied. + const DEFAULT_INTENSITY_TARGET: f64 = 255.0; // Placeholder value + + let mut table = Vec::with_capacity(n); + for i in 0..n { + // `x` represents the normalized input signal, from 0.0 to 1.0. + let x = i as f64 / (n - 1) as f64; + + // Apply the specified EOTF to get the linear light value `y`. + // The output `y` is normalized to the range [0.0, 1.0]. + let y = match tf { + JxlTransferFunction::HLG => TF_HLG::display_from_encoded(x), + JxlTransferFunction::PQ => { + // For PQ, the output of the EOTF is absolute luminance, so we + // normalize it back to [0, 1] relative to the peak luminance. + display_from_encoded_pq(PQ_INTENSITY_TARGET as f32, x) / PQ_INTENSITY_TARGET + } + _ => unreachable!(), // Already checked above. + }; + + // Apply tone mapping if requested. + if tone_map + && *tf == JxlTransferFunction::PQ + && PQ_INTENSITY_TARGET > DEFAULT_INTENSITY_TARGET + { + // TODO(firsching): add tone mapping here. (make y mutable for this) + // let linear_luminance = y * PQ_INTENSITY_TARGET; + // let tone_mapped_luminance = rec2408_tone_map(linear_luminance)?; + // y = tone_mapped_luminance / DEFAULT_INTENSITY_TARGET; + } + + // Clamp the final value to the valid range [0.0, 1.0]. This is + // particularly important for HLG, which can exceed 1.0. + let y_clamped = y.clamp(0.0, 1.0); + + // table.push((y_clamped * 65535.0).round() as u16); + table.push(y_clamped as f32); + } + + Ok(table) +} + +#[cfg(test)] +mod test { + use super::*; + + #[test] + fn test_md5() { + // Test vectors + let test_cases = vec![ + ("", "d41d8cd98f00b204e9800998ecf8427e"), + ( + "The quick brown fox jumps over the lazy dog", + "9e107d9d372bb6826bd81d3542a419d6", + ), + ("abc", "900150983cd24fb0d6963f7d28e17f72"), + ("message digest", "f96b697d7cb7938d525a2f31aaf161d0"), + ( + "abcdefghijklmnopqrstuvwxyz", + "c3fcd3d76192e4007dfb496cca67e13b", + ), + ( + "12345678901234567890123456789012345678901234567890123456789012345678901234567890", + "57edf4a22be3c955ac49da2e2107b67a", + ), + ]; + + for (input, expected) in test_cases { + let hash = compute_md5(input.as_bytes()); + let hex: String = hash.iter().map(|e| format!("{:02x}", e)).collect(); + assert_eq!(hex, expected, "Failed for input: '{}'", input); + } + } + + #[test] + fn test_description() { + assert_eq!( + JxlColorEncoding::srgb(false).get_color_encoding_description(), + "RGB_D65_SRG_Rel_SRG" + ); + assert_eq!( + JxlColorEncoding::srgb(true).get_color_encoding_description(), + "Gra_D65_Rel_SRG" + ); + assert_eq!( + JxlColorEncoding::RgbColorSpace { + white_point: JxlWhitePoint::D65, + primaries: JxlPrimaries::BT2100, + transfer_function: JxlTransferFunction::Gamma(1.7), + rendering_intent: RenderingIntent::Relative + } + .get_color_encoding_description(), + "RGB_D65_202_Rel_g1.7000000" + ); + assert_eq!( + JxlColorEncoding::RgbColorSpace { + white_point: JxlWhitePoint::D65, + primaries: JxlPrimaries::P3, + transfer_function: JxlTransferFunction::SRGB, + rendering_intent: RenderingIntent::Perceptual + } + .get_color_encoding_description(), + "DisplayP3" + ); + } +} diff --git a/third_party/rust/jxl/src/api/data_types.rs b/third_party/rust/jxl/src/api/data_types.rs @@ -0,0 +1,129 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +use crate::headers::extra_channels::ExtraChannel; + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum JxlColorType { + Grayscale, + GrayscaleAlpha, + Rgb, + Rgba, + Bgr, + Bgra, +} + +impl JxlColorType { + pub fn has_alpha(&self) -> bool { + match self { + Self::Grayscale => false, + Self::GrayscaleAlpha => true, + Self::Rgb | Self::Bgr => false, + Self::Rgba | Self::Bgra => true, + } + } + pub fn samples_per_pixel(&self) -> usize { + match self { + Self::Grayscale => 1, + Self::GrayscaleAlpha => 2, + Self::Rgb | Self::Bgr => 3, + Self::Rgba | Self::Bgra => 4, + } + } +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum Endianness { + LittleEndian, + BigEndian, +} + +impl Endianness { + pub fn native() -> Self { + #[cfg(target_endian = "little")] + { + Endianness::LittleEndian + } + #[cfg(target_endian = "big")] + { + Endianness::BigEndian + } + } +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum JxlDataFormat { + U8 { + bit_depth: u8, + }, + U16 { + endianness: Endianness, + bit_depth: u8, + }, + F16 { + endianness: Endianness, + }, + F32 { + endianness: Endianness, + }, +} + +impl JxlDataFormat { + pub fn bytes_per_sample(&self) -> usize { + match self { + Self::U8 { .. } => 1, + Self::U16 { .. } | Self::F16 { .. } => 2, + Self::F32 { .. } => 4, + } + } +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct JxlPixelFormat { + pub color_type: JxlColorType, + // None -> ignore + pub color_data_format: Option<JxlDataFormat>, + pub extra_channel_format: Vec<Option<JxlDataFormat>>, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum JxlBitDepth { + Int { + bits_per_sample: u32, + }, + Float { + bits_per_sample: u32, + exponent_bits_per_sample: u32, + }, +} + +impl JxlBitDepth { + pub fn bits_per_sample(&self) -> u32 { + match self { + JxlBitDepth::Int { bits_per_sample: b } => *b, + JxlBitDepth::Float { + bits_per_sample: b, .. + } => *b, + } + } +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct JxlExtraChannel { + pub ec_type: ExtraChannel, + pub alpha_associated: bool, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct JxlAnimation { + pub num_loops: u32, + pub have_timecodes: bool, +} + +#[derive(Clone, Debug)] +pub struct JxlFrameHeader { + pub name: String, + pub duration: Option<f64>, +} diff --git a/third_party/rust/jxl/src/api/decoder.rs b/third_party/rust/jxl/src/api/decoder.rs @@ -0,0 +1,323 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +use super::{ + JxlBasicInfo, JxlBitstreamInput, JxlCms, JxlColorProfile, JxlDecoderInner, JxlDecoderOptions, + JxlOutputBuffer, JxlPixelFormat, ProcessingResult, +}; +#[cfg(test)] +use crate::frame::Frame; +use crate::{api::JxlFrameHeader, error::Result}; +use states::*; +use std::marker::PhantomData; + +pub mod states { + pub trait JxlState {} + pub struct Initialized; + pub struct WithImageInfo; + pub struct WithFrameInfo; + impl JxlState for Initialized {} + impl JxlState for WithImageInfo {} + impl JxlState for WithFrameInfo {} +} + +// Q: do we plan to add support for box decoding? +// If we do, one way is to take a callback &[u8; 4] -> Box<dyn Write>. + +/// High level API using the typestate pattern to forbid invalid usage. +pub struct JxlDecoder<State: JxlState> { + inner: JxlDecoderInner, + _state: PhantomData<State>, +} + +#[cfg(test)] +pub type FrameCallback = dyn FnMut(&Frame, usize) -> Result<()>; + +impl<S: JxlState> JxlDecoder<S> { + fn wrap_inner(inner: JxlDecoderInner) -> Self { + Self { + inner, + _state: PhantomData, + } + } + + /// Returns a decoder that processes all frames by calling `callback(frame, frame_index)`. + #[cfg(test)] + pub fn with_frame_callback(mut self, callback: Box<FrameCallback>) -> Self { + self.inner = self.inner.with_frame_callback(callback); + self + } + + #[cfg(test)] + pub fn decoded_frames(&self) -> usize { + self.inner.decoded_frames() + } + + /// Rewinds a decoder to the start of the file, allowing past frames to be displayed again. + pub fn rewind(mut self) -> JxlDecoder<Initialized> { + self.inner.rewind(); + JxlDecoder::wrap_inner(self.inner) + } + + fn map_inner_processing_result<SuccessState: JxlState>( + self, + inner_result: ProcessingResult<(), ()>, + ) -> ProcessingResult<JxlDecoder<SuccessState>, Self> { + match inner_result { + ProcessingResult::Complete { .. } => ProcessingResult::Complete { + result: JxlDecoder::wrap_inner(self.inner), + }, + ProcessingResult::NeedsMoreInput { size_hint, .. } => { + ProcessingResult::NeedsMoreInput { + size_hint, + fallback: self, + } + } + } + } +} + +impl JxlDecoder<Initialized> { + pub fn new(options: JxlDecoderOptions) -> Self { + Self::wrap_inner(JxlDecoderInner::new(options, None)) + } + + pub fn new_with_cms(options: JxlDecoderOptions, cms: impl JxlCms + 'static) -> Self { + Self::wrap_inner(JxlDecoderInner::new(options, Some(Box::new(cms)))) + } + + pub fn process( + mut self, + input: &mut impl JxlBitstreamInput, + ) -> Result<ProcessingResult<JxlDecoder<WithImageInfo>, Self>> { + let inner_result = self.inner.process(input, None)?; + Ok(self.map_inner_processing_result(inner_result)) + } +} + +impl JxlDecoder<WithImageInfo> { + // TODO(veluca): once frame skipping is implemented properly, expose that in the API. + + /// Obtains the image's basic information. + pub fn basic_info(&self) -> &JxlBasicInfo { + self.inner.basic_info().unwrap() + } + + /// Retrieves the file's color profile. + pub fn embedded_color_profile(&self) -> &JxlColorProfile { + self.inner.embedded_color_profile().unwrap() + } + + /// Retrieves the current output color profile. + pub fn output_color_profile(&self) -> &JxlColorProfile { + self.inner.output_color_profile().unwrap() + } + + /// Specifies the preferred color profile to be used for outputting data. + /// Same semantics as JxlDecoderSetOutputColorProfile. + pub fn set_output_color_profile(&mut self, profile: &JxlColorProfile) -> Result<()> { + self.inner.set_output_color_profile(profile) + } + + pub fn current_pixel_format(&self) -> &JxlPixelFormat { + self.inner.current_pixel_format().unwrap() + } + + pub fn set_pixel_format(&mut self, pixel_format: JxlPixelFormat) { + self.inner.set_pixel_format(pixel_format); + } + + pub fn process( + mut self, + input: &mut impl JxlBitstreamInput, + ) -> Result<ProcessingResult<JxlDecoder<WithFrameInfo>, Self>> { + let inner_result = self.inner.process(input, None)?; + Ok(self.map_inner_processing_result(inner_result)) + } + + pub fn has_more_frames(&self) -> bool { + self.inner.has_more_frames() + } +} + +impl JxlDecoder<WithFrameInfo> { + /// Skip the current frame. + pub fn skip_frame( + mut self, + input: &mut impl JxlBitstreamInput, + ) -> Result<ProcessingResult<JxlDecoder<WithImageInfo>, Self>> { + let inner_result = self.inner.process(input, None)?; + Ok(self.map_inner_processing_result(inner_result)) + } + + // TODO: don't use the raw bitstream type; include name and extra channel blend info. + pub fn frame_header(&self) -> JxlFrameHeader { + self.inner.frame_header().unwrap() + } + + /// Number of passes we have full data for. + pub fn num_completed_passes(&self) -> usize { + self.inner.num_completed_passes().unwrap() + } + + /// Draws all the pixels we have data for. + pub fn flush_pixels(&mut self, buffers: &mut [JxlOutputBuffer<'_>]) -> Result<()> { + self.inner.flush_pixels(buffers) + } + + /// Guarantees to populate exactly the appropriate part of the buffers. + /// Wants one buffer for each non-ignored pixel type, i.e. color channels and each extra channel. + pub fn process<In: JxlBitstreamInput>( + mut self, + input: &mut In, + buffers: &mut [JxlOutputBuffer<'_>], + ) -> Result<ProcessingResult<JxlDecoder<WithImageInfo>, Self>> { + let inner_result = self.inner.process(input, Some(buffers))?; + Ok(self.map_inner_processing_result(inner_result)) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::api::JxlDecoderOptions; + use crate::api::test::create_output_buffers; + use jxl_macros::for_each_test_file; + use std::path::Path; + + #[test] + fn decode_small_chunks() { + arbtest::arbtest(|u| { + decode_test_data( + std::fs::read("resources/test/green_queen_vardct_e3.jxl") + .expect("Failed to read test file"), + u.arbitrary::<u8>().unwrap() as usize + 1, + ) + .unwrap(); + Ok(()) + }); + } + + fn decode_test_data(data: Vec<u8>, chunk_size: usize) -> Result<(), crate::error::Error> { + // Create decoder with default options + let options = JxlDecoderOptions::default(); + let mut initialized_decoder = JxlDecoder::<states::Initialized>::new(options); + + let mut input = data.as_slice(); + let mut chunk_input = &input[0..0]; + + // Process until we have image info + let mut decoder_with_image_info = loop { + chunk_input = &input[..(chunk_input.len().saturating_add(chunk_size)).min(input.len())]; + let available_before = chunk_input.len(); + let process_result = initialized_decoder.process(&mut chunk_input); + input = &input[(available_before - chunk_input.len())..]; + match process_result.unwrap() { + ProcessingResult::Complete { result } => break result, + ProcessingResult::NeedsMoreInput { fallback, .. } => { + if input.is_empty() { + panic!("Unexpected end of input while reading image info"); + } + initialized_decoder = fallback; + } + } + }; + + // Get basic info + let basic_info = decoder_with_image_info.basic_info().clone(); + assert!(basic_info.bit_depth.bits_per_sample() > 0); + + // Get image dimensions (after upsampling, which is the actual output size) + let (width, height) = basic_info.size; + assert!(width > 0); + assert!(height > 0); + + // Get pixel format info + let pixel_format = decoder_with_image_info.current_pixel_format().clone(); + let num_channels = pixel_format.color_type.samples_per_pixel(); + assert!(num_channels > 0); + + let mut frame_count = 0; + + loop { + // Process until we have frame info + let mut decoder_with_frame_info = loop { + chunk_input = + &input[..(chunk_input.len().saturating_add(chunk_size)).min(input.len())]; + let available_before = chunk_input.len(); + let process_result = decoder_with_image_info.process(&mut chunk_input); + input = &input[(available_before - chunk_input.len())..]; + match process_result.unwrap() { + ProcessingResult::Complete { result } => break result, + ProcessingResult::NeedsMoreInput { fallback, .. } => { + if input.is_empty() { + panic!("Unexpected end of input while reading frame info"); + } + decoder_with_image_info = fallback; + } + } + }; + decoder_with_frame_info.frame_header(); + + create_output_buffers!(basic_info, pixel_format, output_buffers, output_slices); + + decoder_with_image_info = loop { + chunk_input = + &input[..(chunk_input.len().saturating_add(chunk_size)).min(input.len())]; + let available_before = chunk_input.len(); + let process_result = + decoder_with_frame_info.process(&mut chunk_input, &mut output_slices); + input = &input[(available_before - chunk_input.len())..]; + match process_result.unwrap() { + ProcessingResult::Complete { result } => break result, + ProcessingResult::NeedsMoreInput { fallback, .. } => { + if input.is_empty() { + panic!("Unexpected end of input while decoding frame"); + } + decoder_with_frame_info = fallback; + } + } + }; + + // Verify we decoded something + if pixel_format.color_type == Rgb { + // For RGB, first buffer contains interleaved RGB data + assert!(!output_buffers.is_empty()); + assert_eq!(output_buffers[0].len(), width * height * 12); // 3 channels * 4 bytes + // Additional buffers for extra channels + for buffer in &output_buffers[1..] { + assert_eq!(buffer.len(), width * height * 4); + } + } else { + // For other formats, one buffer per channel + assert_eq!(output_buffers.len(), num_channels); + for buffer in &output_buffers { + assert_eq!(buffer.len(), width * height * 4); + } + } + + frame_count += 1; + + // Check if there are more frames + if !decoder_with_image_info.has_more_frames() { + break; + } + } + + // Ensure we decoded at least one frame + assert!(frame_count > 0, "No frames were decoded"); + + Ok(()) + } + + fn decode_test_file(path: &Path) -> Result<(), crate::error::Error> { + decode_test_data( + std::fs::read(path).expect("Failed to read test file"), + usize::MAX, + ) + } + + for_each_test_file!(decode_test_file); +} diff --git a/third_party/rust/jxl/src/api/inner.rs b/third_party/rust/jxl/src/api/inner.rs @@ -0,0 +1,110 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +#[cfg(test)] +use crate::api::FrameCallback; +use crate::{ + api::JxlFrameHeader, + error::{Error, Result}, +}; + +use super::{JxlBasicInfo, JxlCms, JxlColorProfile, JxlDecoderOptions, JxlPixelFormat}; +use box_parser::BoxParser; +use codestream_parser::CodestreamParser; + +mod box_parser; +mod codestream_parser; +mod process; + +/// Low-level, less-type-safe API. +pub struct JxlDecoderInner { + options: JxlDecoderOptions, + cms: Option<Box<dyn JxlCms>>, + box_parser: BoxParser, + codestream_parser: CodestreamParser, +} + +impl JxlDecoderInner { + /// Creates a new decoder with the given options and, optionally, CMS. + pub fn new(options: JxlDecoderOptions, cms: Option<Box<dyn JxlCms>>) -> Self { + JxlDecoderInner { + options, + cms, + box_parser: BoxParser::new(), + codestream_parser: CodestreamParser::new(), + } + } + + #[cfg(test)] + pub fn with_frame_callback(mut self, callback: Box<FrameCallback>) -> Self { + self.codestream_parser.frame_callback = Some(callback); + self + } + + #[cfg(test)] + pub fn decoded_frames(&self) -> usize { + self.codestream_parser.decoded_frames + } + + /// Obtains the image's basic information, if available. + pub fn basic_info(&self) -> Option<&JxlBasicInfo> { + self.codestream_parser.basic_info.as_ref() + } + + /// Retrieves the file's color profile, if available. + pub fn embedded_color_profile(&self) -> Option<&JxlColorProfile> { + self.codestream_parser.embedded_color_profile.as_ref() + } + + /// Retrieves the current output color profile, if available. + pub fn output_color_profile(&self) -> Option<&JxlColorProfile> { + self.codestream_parser.output_color_profile.as_ref() + } + + /// Specifies the preferred color profile to be used for outputting data. + /// Same semantics as JxlDecoderSetOutputColorProfile. + pub fn set_output_color_profile(&mut self, profile: &JxlColorProfile) -> Result<()> { + if let (JxlColorProfile::Icc(_), None) = (profile, &self.cms) { + return Err(Error::ICCOutputNoCMS); + } + unimplemented!() + } + + pub fn current_pixel_format(&self) -> Option<&JxlPixelFormat> { + self.codestream_parser.pixel_format.as_ref() + } + + pub fn set_pixel_format(&mut self, pixel_format: JxlPixelFormat) { + drop(pixel_format); + unimplemented!() + } + + pub fn frame_header(&self) -> Option<JxlFrameHeader> { + let frame_header = self.codestream_parser.frame.as_ref()?.header(); + Some(JxlFrameHeader { + name: frame_header.name.clone(), + duration: self + .codestream_parser + .animation + .as_ref() + .map(|anim| frame_header.duration(anim)), + }) + } + /// Number of passes we have full data for. + pub fn num_completed_passes(&self) -> Option<usize> { + None // TODO. + } + + /// Rewinds a decoder to the start of the file, allowing past frames to be displayed again. + pub fn rewind(&mut self) { + // TODO(veluca): keep track of frame offsets for skipping. + self.box_parser = BoxParser::new(); + self.codestream_parser = CodestreamParser::new(); + } + + pub fn has_more_frames(&self) -> bool { + self.codestream_parser.has_more_frames + } +} diff --git a/third_party/rust/jxl/src/api/inner/box_parser.rs b/third_party/rust/jxl/src/api/inner/box_parser.rs @@ -0,0 +1,175 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +use crate::error::{Error, Result}; + +use crate::api::{ + JxlBitstreamInput, JxlSignatureType, check_signature_internal, inner::process::SmallBuffer, +}; + +#[derive(Clone)] +enum ParseState { + SignatureNeeded, + BoxNeeded, + CodestreamBox(u64), + SkippableBox(u64), +} + +enum CodestreamBoxType { + None, + Jxlc, + Jxlp(u32), + LastJxlp, +} + +pub(super) struct BoxParser { + pub(super) box_buffer: SmallBuffer<128>, + state: ParseState, + box_type: CodestreamBoxType, +} + +impl BoxParser { + pub(super) fn new() -> Self { + BoxParser { + box_buffer: SmallBuffer::new(), + state: ParseState::SignatureNeeded, + box_type: CodestreamBoxType::None, + } + } + + // Reads input until the next byte of codestream is available. + // This function might over-read bytes. Thus, the contents of self.box_buffer should always be + // read after this function call. + // Returns the number of codestream bytes that will be available to be read after this call, + // including any bytes in self.box_buffer. + // Might return `u64::MAX`, indicating that the rest of the file is codestream. + pub(super) fn get_more_codestream( + &mut self, + input: &mut impl JxlBitstreamInput, + ) -> Result<u64> { + // TODO(veluca): consider moving most of this function into a function that is not generic. + loop { + match self.state.clone() { + ParseState::SignatureNeeded => { + self.box_buffer.refill(|b| input.read(b), None)?; + match check_signature_internal(&self.box_buffer)? { + None => return Err(Error::InvalidSignature), + Some(JxlSignatureType::Codestream) => { + self.state = ParseState::CodestreamBox(u64::MAX); + return Ok(u64::MAX); + } + Some(JxlSignatureType::Container) => { + self.box_buffer + .consume(JxlSignatureType::Container.signature().len()); + self.state = ParseState::BoxNeeded; + } + } + } + ParseState::CodestreamBox(b) => { + return Ok(b); + } + ParseState::SkippableBox(mut s) => { + let num = s.min(usize::MAX as u64) as usize; + let skipped = if !self.box_buffer.is_empty() { + self.box_buffer.consume(num) + } else { + input.skip(num)? + }; + if skipped == 0 { + return Err(Error::OutOfBounds(num)); + } + s -= skipped as u64; + if s == 0 { + self.state = ParseState::BoxNeeded; + } else { + self.state = ParseState::SkippableBox(s); + } + } + ParseState::BoxNeeded => { + self.box_buffer.refill(|b| input.read(b), None)?; + let min_len = match &self.box_buffer[..] { + [0, 0, 0, 1, ..] => 16, + _ => 8, + }; + if self.box_buffer.len() <= min_len { + return Err(Error::OutOfBounds(min_len - self.box_buffer.len())); + } + let ty: [_; 4] = self.box_buffer[4..8].try_into().unwrap(); + let extra_len = if &ty == b"jxlp" { 4 } else { 0 }; + if self.box_buffer.len() <= min_len + extra_len { + return Err(Error::OutOfBounds( + min_len + extra_len - self.box_buffer.len(), + )); + } + let box_len = match &self.box_buffer[..] { + [0, 0, 0, 1, ..] => { + u64::from_be_bytes(self.box_buffer[8..16].try_into().unwrap()) + } + _ => u32::from_be_bytes(self.box_buffer[0..4].try_into().unwrap()) as u64, + }; + // Per JXL spec: jxlc box with length 0 has special meaning "extends to EOF" + let content_len = if box_len == 0 && (&ty == b"jxlp" || &ty == b"jxlc") { + u64::MAX + } else { + if box_len <= (min_len + extra_len) as u64 { + return Err(Error::InvalidBox); + } + box_len - min_len as u64 - extra_len as u64 + }; + match &ty { + b"jxlc" => { + if matches!( + self.box_type, + CodestreamBoxType::Jxlp(..) | CodestreamBoxType::LastJxlp + ) { + return Err(Error::InvalidBox); + } + self.box_type = CodestreamBoxType::Jxlc; + self.state = ParseState::CodestreamBox(content_len); + } + b"jxlp" => { + let index = u32::from_be_bytes( + self.box_buffer[min_len..min_len + 4].try_into().unwrap(), + ); + let wanted_idx = match self.box_type { + CodestreamBoxType::Jxlc | CodestreamBoxType::LastJxlp => { + return Err(Error::InvalidBox); + } + CodestreamBoxType::None => 0, + CodestreamBoxType::Jxlp(i) => i + 1, + }; + let last = index & 0x80000000 != 0; + let idx = index & 0x7fffffff; + if idx != wanted_idx { + return Err(Error::InvalidBox); + } + self.box_type = if last { + CodestreamBoxType::LastJxlp + } else { + CodestreamBoxType::Jxlp(idx) + }; + self.state = ParseState::CodestreamBox(content_len); + } + _ => { + self.state = ParseState::SkippableBox(content_len); + } + } + self.box_buffer.consume(min_len + extra_len); + } + } + } + } + + pub(super) fn consume_codestream(&mut self, amount: u64) { + if let ParseState::CodestreamBox(cb) = &mut self.state { + *cb = cb.checked_sub(amount).unwrap(); + if *cb == 0 { + self.state = ParseState::BoxNeeded; + } + } else if amount != 0 { + unreachable!() + } + } +} diff --git a/third_party/rust/jxl/src/api/inner/codestream_parser.rs b/third_party/rust/jxl/src/api/inner/codestream_parser.rs @@ -0,0 +1,261 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +use std::{collections::VecDeque, io::IoSliceMut}; + +use sections::SectionState; + +#[cfg(test)] +use crate::api::FrameCallback; +use crate::{ + api::{ + JxlBasicInfo, JxlBitstreamInput, JxlColorProfile, JxlDecoderOptions, JxlOutputBuffer, + JxlPixelFormat, + inner::{box_parser::BoxParser, process::SmallBuffer}, + }, + error::{Error, Result}, + frame::{DecoderState, Frame, Section}, + headers::{Animation, FileHeader, frame_header::FrameHeader}, + icc::IncrementalIccReader, +}; + +mod non_section; +mod sections; + +struct SectionBuffer { + len: usize, + data: Vec<u8>, + section: Section, +} + +// This number should be big enough to guarantee that we can always make progress by reading +// fragments of size at most *half* of it, if not reading a section. +const NON_SECTION_CHUNK_SIZE: usize = 4096; + +pub(super) struct CodestreamParser { + // TODO(veluca): this would probably be cleaner with some kind of state enum. + pub(super) file_header: Option<FileHeader>, + icc_parser: Option<IncrementalIccReader>, + // These fields are populated once image information is available. + decoder_state: Option<DecoderState>, + pub(super) basic_info: Option<JxlBasicInfo>, + pub(super) animation: Option<Animation>, + pub(super) embedded_color_profile: Option<JxlColorProfile>, + pub(super) output_color_profile: Option<JxlColorProfile>, + pub(super) pixel_format: Option<JxlPixelFormat>, + + // These fields are populated when starting to decode a frame, and cleared once + // the frame is done. + frame_header: Option<FrameHeader>, + pub(super) frame: Option<Frame>, + + // Buffers. + non_section_buf: SmallBuffer<NON_SECTION_CHUNK_SIZE>, + non_section_bit_offset: u8, + sections: VecDeque<SectionBuffer>, + ready_section_data: usize, + skip_sections: bool, + // True when we need to process frames without copying them to output buffers, e.g. reference frames + process_without_output: bool, + + section_state: SectionState, + available_sections: Vec<SectionBuffer>, + + pub(super) has_more_frames: bool, + + #[cfg(test)] + pub frame_callback: Option<Box<FrameCallback>>, + #[cfg(test)] + pub decoded_frames: usize, +} + +impl CodestreamParser { + pub(super) fn new() -> Self { + Self { + file_header: None, + icc_parser: None, + decoder_state: None, + basic_info: None, + animation: None, + embedded_color_profile: None, + output_color_profile: None, + pixel_format: None, + frame_header: None, + frame: None, + non_section_buf: SmallBuffer::new(), + non_section_bit_offset: 0, + sections: VecDeque::new(), + ready_section_data: 0, + skip_sections: false, + process_without_output: false, + section_state: SectionState::new(0, 0), + available_sections: vec![], + has_more_frames: true, + #[cfg(test)] + frame_callback: None, + #[cfg(test)] + decoded_frames: 0, + } + } + + fn has_visible_frame(&self) -> bool { + if let Some(frame) = &self.frame { + frame.header().is_visible() + } else { + false + } + } + + pub(super) fn process<In: JxlBitstreamInput>( + &mut self, + box_parser: &mut BoxParser, + input: &mut In, + decode_options: &JxlDecoderOptions, + mut output_buffers: Option<&mut [JxlOutputBuffer]>, + ) -> Result<()> { + // If we have sections to read, read into sections; otherwise, read into the local buffer. + loop { + if !self.sections.is_empty() { + let regular_frame = self.has_visible_frame(); + if !self.process_without_output && output_buffers.is_none() { + self.skip_sections = true; + } + + if !self.skip_sections { + // This is just an estimate as there could be box bytes in the middle. + let mut readable_section_data = (self.non_section_buf.len() + + input.available_bytes()? + + self.ready_section_data) + .max(1); + // Ensure enough section buffers are available for reading available data. + for buf in self.sections.iter_mut() { + if buf.data.is_empty() { + buf.data.resize(buf.len, 0); + } + readable_section_data = + readable_section_data.saturating_sub(buf.data.len()); + if readable_section_data == 0 { + break; + } + } + // Read sections up to the end of the current box. + let mut available_codestream = match box_parser.get_more_codestream(input) { + Err(Error::OutOfBounds(_)) => 0, + Ok(c) => c as usize, + Err(e) => return Err(e), + }; + let mut section_buffers = vec![]; + let mut ready = self.ready_section_data; + for buf in self.sections.iter_mut() { + if buf.data.is_empty() { + break; + } + let len = buf.data.len(); + if len > ready { + let readable = (available_codestream + ready).min(len); + section_buffers.push(IoSliceMut::new(&mut buf.data[ready..readable])); + available_codestream = + available_codestream.saturating_sub(readable - ready); + if available_codestream == 0 { + break; + } + } + ready = ready.saturating_sub(len); + } + let mut buffers = &mut section_buffers[..]; + loop { + let num = if !box_parser.box_buffer.is_empty() { + box_parser.box_buffer.take(buffers) + } else { + input.read(buffers)? + }; + self.ready_section_data += num; + box_parser.consume_codestream(num as u64); + IoSliceMut::advance_slices(&mut buffers, num); + if num == 0 || buffers.is_empty() { + break; + } + } + match self.process_sections(&mut output_buffers) { + Ok(None) => Ok(()), + Ok(Some(missing)) => Err(Error::OutOfBounds(missing)), + Err(Error::OutOfBounds(_)) => Err(Error::SectionTooShort), + Err(err) => Err(err), + }?; + } else { + let total_size = self.sections.iter().map(|x| x.len).sum::<usize>(); + loop { + let to_skip = total_size - self.ready_section_data; + if to_skip == 0 { + break; + } + let available_codestream = box_parser.get_more_codestream(input)? as usize; + let to_skip = to_skip.min(available_codestream); + let skipped = if !box_parser.box_buffer.is_empty() { + box_parser.box_buffer.consume(to_skip) + } else { + input.skip(to_skip)? + }; + box_parser.consume_codestream(skipped as u64); + self.ready_section_data += skipped; + if skipped == 0 { + break; + } + } + if self.ready_section_data < total_size { + return Err(Error::OutOfBounds(total_size - self.ready_section_data)); + } else { + self.sections.clear(); + } + } + if self.sections.is_empty() { + // Go back to parsing a new frame header, if any. + self.process_without_output = false; + if regular_frame { + return Ok(()); + } + continue; + } + } else { + // Trying to read a frame or a file header. + assert!(self.frame.is_none()); + assert!(self.has_more_frames); + + let available_codestream = match box_parser.get_more_codestream(input) { + Err(Error::OutOfBounds(_)) => 0, + Ok(c) => c as usize, + Err(e) => return Err(e), + }; + let c = self.non_section_buf.refill( + |buf| { + if !box_parser.box_buffer.is_empty() { + Ok(box_parser.box_buffer.take(buf)) + } else { + input.read(buf) + } + }, + Some(available_codestream), + )?; + box_parser.consume_codestream(c as u64); + + self.process_non_section(decode_options)?; + + if self.decoder_state.is_some() && self.frame_header.is_none() { + // Return to caller if we found image info. + return Ok(()); + } + if self.frame.is_some() { + if self.has_visible_frame() { + // Return to caller if we found visible frame info. + return Ok(()); + } else { + self.process_without_output = true; + continue; + } + } + } + } + } +} diff --git a/third_party/rust/jxl/src/api/inner/codestream_parser/non_section.rs b/third_party/rust/jxl/src/api/inner/codestream_parser/non_section.rs @@ -0,0 +1,277 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +use std::io::IoSliceMut; + +use crate::{ + api::{ + Endianness, JxlBasicInfo, JxlBitDepth, JxlColorEncoding, JxlColorProfile, JxlColorType, + JxlDataFormat, JxlDecoderOptions, JxlExtraChannel, JxlPixelFormat, JxlTransferFunction, + inner::codestream_parser::SectionState, + }, + bit_reader::BitReader, + error::{Error, Result}, + frame::{DecoderState, Frame, Section}, + headers::{ + FileHeader, JxlHeader, + color_encoding::ColorSpace, + encodings::UnconditionalCoder, + frame_header::{FrameHeader, Toc, TocNonserialized}, + }, + icc::IncrementalIccReader, +}; + +use super::{CodestreamParser, SectionBuffer}; + +impl CodestreamParser { + pub(super) fn process_non_section(&mut self, decode_options: &JxlDecoderOptions) -> Result<()> { + if self.decoder_state.is_none() && self.file_header.is_none() { + // We don't have a file header yet. Try parsing that. + // TODO(veluca): make this incremental, as a file header might be multiple megabytes. + let mut br = BitReader::new(&self.non_section_buf); + br.skip_bits(self.non_section_bit_offset as usize)?; + let file_header = FileHeader::read(&mut br)?; + let data = &file_header.image_metadata; + self.animation = data.animation.clone(); + self.basic_info = Some(JxlBasicInfo { + size: ( + file_header.size.xsize() as usize, + file_header.size.ysize() as usize, + ), + bit_depth: if data.bit_depth.floating_point_sample() { + JxlBitDepth::Float { + bits_per_sample: data.bit_depth.bits_per_sample(), + exponent_bits_per_sample: data.bit_depth.exponent_bits_per_sample(), + } + } else { + JxlBitDepth::Int { + bits_per_sample: data.bit_depth.bits_per_sample(), + } + }, + orientation: data.orientation, + extra_channels: data + .extra_channel_info + .iter() + .map(|info| JxlExtraChannel { + ec_type: info.ec_type, + alpha_associated: info.alpha_associated(), + }) + .collect(), + animation: data + .animation + .as_ref() + .map(|anim| crate::api::JxlAnimation { + num_loops: anim.num_loops, + have_timecodes: anim.have_timecodes, + }), + }); + self.file_header = Some(file_header); + let bits = br.total_bits_read(); + self.non_section_buf.consume(bits / 8); + self.non_section_bit_offset = (bits % 8) as u8; + } + + if self.decoder_state.is_none() && self.embedded_color_profile.is_none() { + let file_header = self.file_header.as_ref().unwrap(); + // Parse (or extract from file header) the ICC profile. + let mut br = BitReader::new(&self.non_section_buf); + br.skip_bits(self.non_section_bit_offset as usize)?; + let embedded_color_profile = if file_header.image_metadata.color_encoding.want_icc { + if self.icc_parser.is_none() { + self.icc_parser = Some(IncrementalIccReader::new(&mut br)?); + } + let icc_parser = self.icc_parser.as_mut().unwrap(); + let mut bits = br.total_bits_read(); + for _ in 0..icc_parser.remaining() { + match icc_parser.read_one(&mut br) { + Ok(()) => bits = br.total_bits_read(), + Err(Error::OutOfBounds(c)) => { + self.non_section_buf.consume(bits / 8); + self.non_section_bit_offset = (bits % 8) as u8; + // Estimate >= one bit per remaining character to read. + return Err(Error::OutOfBounds(c + icc_parser.remaining() / 8)); + } + Err(e) => return Err(e), + } + } + self.non_section_buf.consume(bits / 8); + self.non_section_bit_offset = (bits % 8) as u8; + JxlColorProfile::Icc(self.icc_parser.take().unwrap().finalize()?) + } else { + JxlColorProfile::Simple(JxlColorEncoding::from_internal( + &file_header.image_metadata.color_encoding, + )?) + }; + let output_color_profile = if file_header.image_metadata.xyb_encoded { + let nonlinear_output_color_profile = match &embedded_color_profile { + JxlColorProfile::Icc(_) => JxlColorEncoding::srgb( + file_header.image_metadata.color_encoding.color_space == ColorSpace::Gray, + ), + JxlColorProfile::Simple(encoding) => encoding.clone(), + }; + JxlColorProfile::Simple(if decode_options.xyb_output_linear { + match nonlinear_output_color_profile { + JxlColorEncoding::RgbColorSpace { + white_point, + primaries, + transfer_function: _, + rendering_intent, + } => JxlColorEncoding::RgbColorSpace { + white_point, + primaries, + transfer_function: JxlTransferFunction::Linear, + rendering_intent, + }, + JxlColorEncoding::GrayscaleColorSpace { + white_point, + transfer_function: _, + rendering_intent, + } => JxlColorEncoding::GrayscaleColorSpace { + white_point, + transfer_function: JxlTransferFunction::Linear, + rendering_intent, + }, + JxlColorEncoding::XYB { .. } => unreachable!(), + } + } else { + nonlinear_output_color_profile + }) + } else { + embedded_color_profile.clone() + }; + self.embedded_color_profile = Some(embedded_color_profile); + self.output_color_profile = Some(output_color_profile); + self.pixel_format = Some(JxlPixelFormat { + color_type: if file_header.image_metadata.color_encoding.color_space + == ColorSpace::Gray + { + JxlColorType::Grayscale + } else { + JxlColorType::Rgb + }, + color_data_format: Some(JxlDataFormat::F32 { + endianness: Endianness::native(), + }), + extra_channel_format: vec![ + Some(JxlDataFormat::F32 { + endianness: Endianness::native() + }); + file_header.image_metadata.extra_channel_info.len() + ], + }); + + let mut br = BitReader::new(&self.non_section_buf); + br.skip_bits(self.non_section_bit_offset as usize)?; + br.jump_to_byte_boundary()?; + self.non_section_buf.consume(br.total_bits_read() / 8); + + // We now have image information. + let mut decoder_state = DecoderState::new(self.file_header.take().unwrap()); + decoder_state.xyb_output_linear = decode_options.xyb_output_linear; + decoder_state.render_spotcolors = decode_options.render_spot_colors; + self.decoder_state = Some(decoder_state); + // Reset bit offset to 0 since we've consumed everything up to a byte boundary + self.non_section_bit_offset = 0; + return Ok(()); + } + + let decoder_state = self.decoder_state.as_mut().unwrap(); + + if self.frame_header.is_none() { + // We don't have a frame header yet. Try parsing that. + // TODO(veluca): do we need to make this incremental? + let mut br = BitReader::new(&self.non_section_buf); + br.skip_bits(self.non_section_bit_offset as usize)?; + let mut frame_header = FrameHeader::read_unconditional( + &(), + &mut br, + &decoder_state.file_header.frame_header_nonserialized(), + )?; + frame_header.postprocess(&decoder_state.file_header.frame_header_nonserialized()); + self.frame_header = Some(frame_header); + let bits = br.total_bits_read(); + self.non_section_buf.consume(bits / 8); + self.non_section_bit_offset = (bits % 8) as u8; + } + + let mut br = BitReader::new(&self.non_section_buf); + br.skip_bits(self.non_section_bit_offset as usize)?; + let num_toc_entries = self.frame_header.as_ref().unwrap().num_toc_entries(); + let toc = Toc::read_unconditional( + &(), + &mut br, + &TocNonserialized { + num_entries: num_toc_entries as u32, + }, + )?; + br.jump_to_byte_boundary()?; + let frame = Frame::from_header_and_toc( + self.frame_header.take().unwrap(), + toc, + self.decoder_state.take().unwrap(), + )?; + let bits = br.total_bits_read(); + self.non_section_buf.consume(bits / 8); + self.non_section_bit_offset = (bits % 8) as u8; + + let mut sections: Vec<_> = frame + .toc() + .entries + .iter() + .map(|x| SectionBuffer { + len: *x as usize, + data: vec![], + section: Section::LfGlobal, // will be fixed later + }) + .collect(); + + let order = if frame.toc().permuted { + frame.toc().permutation.0.clone() + } else { + (0..sections.len() as u32).collect() + }; + + if sections.len() > 1 { + let base_sections = [Section::LfGlobal, Section::HfGlobal]; + let lf_sections = (0..frame.header().num_lf_groups()).map(|x| Section::Lf { group: x }); + let hf_sections = (0..frame.header().passes.num_passes).flat_map(|p| { + (0..frame.header().num_groups()).map(move |g| Section::Hf { + group: g, + pass: p as usize, + }) + }); + + for section in base_sections + .into_iter() + .chain(lf_sections) + .chain(hf_sections) + { + sections[order[frame.get_section_idx(section)] as usize].section = section; + } + } + + self.sections = sections.into_iter().collect(); + self.ready_section_data = 0; + + // Move data from the pre-section buffer into the sections. + for buf in self.sections.iter_mut() { + if self.non_section_buf.is_empty() { + break; + } + buf.data = vec![0; buf.len]; + self.ready_section_data += self + .non_section_buf + .take(&mut [IoSliceMut::new(&mut buf.data)]); + } + + self.section_state = + SectionState::new(frame.header().num_lf_groups(), frame.header().num_groups()); + assert!(self.available_sections.is_empty()); + + self.frame = Some(frame); + + Ok(()) + } +} diff --git a/third_party/rust/jxl/src/api/inner/codestream_parser/sections.rs b/third_party/rust/jxl/src/api/inner/codestream_parser/sections.rs @@ -0,0 +1,198 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +use crate::{ + api::{JxlColorType, JxlOutputBuffer}, + bit_reader::BitReader, + error::Result, + frame::Section, +}; + +use super::CodestreamParser; + +pub(super) struct SectionState { + lf_global_done: bool, + remaining_lf: usize, + hf_global_done: bool, + completed_passes: Vec<u8>, +} + +impl SectionState { + pub(super) fn new(num_lf_groups: usize, num_groups: usize) -> Self { + Self { + lf_global_done: false, + remaining_lf: num_lf_groups, + hf_global_done: false, + completed_passes: vec![0; num_groups], + } + } +} + +// No guarantees on the order of calls to f, or the order of retained elements in vec. +fn retain_by_value<T>(vec: &mut Vec<T>, mut f: impl FnMut(T) -> Option<T>) { + for pos in (0..vec.len()).rev() { + let element_to_test = vec.swap_remove(pos); + if let Some(v) = f(element_to_test) { + vec.push(v); + } + } +} + +impl CodestreamParser { + pub(super) fn process_sections( + &mut self, + output_buffers: &mut Option<&mut [JxlOutputBuffer<'_>]>, + ) -> Result<Option<usize>> { + // Dequeue ready sections. + while self + .sections + .front() + .is_some_and(|s| s.len <= self.ready_section_data) + { + let s = self.sections.pop_front().unwrap(); + self.ready_section_data -= s.len; + self.available_sections.push(s); + } + if self.available_sections.is_empty() { + return Ok(Some( + self.sections.front().unwrap().len - self.ready_section_data, + )); + } + let frame = self.frame.as_mut().unwrap(); + let frame_header = frame.header(); + if frame_header.num_groups() == 1 && frame_header.passes.num_passes == 1 { + // Single-group special case. + assert_eq!(self.available_sections.len(), 1); + assert!(self.sections.is_empty()); + let mut br = BitReader::new(&self.available_sections[0].data); + frame.decode_lf_global(&mut br)?; + frame.decode_lf_group(0, &mut br)?; + frame.decode_hf_global(&mut br)?; + frame.prepare_render_pipeline()?; + frame.finalize_lf()?; + frame.decode_hf_group(0, 0, &mut br)?; + self.available_sections.clear(); + } else { + let mut lf_global_section = None; + let mut lf_sections = vec![]; + let mut hf_global_section = None; + let mut sorted_sections_for_each_group = Vec::with_capacity(frame_header.num_groups()); + for _ in 0..frame_header.num_groups() { + sorted_sections_for_each_group.push(vec![]); + } + + loop { + let initial_sz = self.available_sections.len(); + retain_by_value(&mut self.available_sections, |sec| match sec.section { + Section::LfGlobal => { + lf_global_section = Some(sec); + self.section_state.lf_global_done = true; + None + } + Section::Lf { .. } => { + if !self.section_state.lf_global_done { + Some(sec) + } else { + lf_sections.push(sec); + self.section_state.remaining_lf -= 1; + None + } + } + Section::HfGlobal => { + if self.section_state.remaining_lf != 0 { + Some(sec) + } else { + hf_global_section = Some(sec); + self.section_state.hf_global_done = true; + None + } + } + Section::Hf { group, pass } => { + if !self.section_state.hf_global_done + && self.section_state.completed_passes[group] != pass as u8 + { + Some(sec) + } else { + sorted_sections_for_each_group[group].push(sec); + self.section_state.completed_passes[group] += 1; + None + } + } + }); + if self.available_sections.len() == initial_sz { + break; + } + } + + if let Some(lf_global) = lf_global_section { + frame.decode_lf_global(&mut BitReader::new(&lf_global.data))?; + } + + for lf_section in lf_sections { + let Section::Lf { group } = lf_section.section else { + unreachable!() + }; + frame.decode_lf_group(group, &mut BitReader::new(&lf_section.data))?; + } + + if let Some(hf_global) = hf_global_section { + frame.decode_hf_global(&mut BitReader::new(&hf_global.data))?; + frame.prepare_render_pipeline()?; + frame.finalize_lf()?; + } + + for g in sorted_sections_for_each_group { + // TODO(veluca): render all the available passes at once. + for sec in g { + let Section::Hf { group, pass } = sec.section else { + unreachable!() + }; + frame.decode_hf_group(group, pass, &mut BitReader::new(&sec.data))?; + } + } + } + // Frame is not yet complete. + if !self.sections.is_empty() { + return Ok(None); + } + assert!(self.available_sections.is_empty()); + + #[cfg(test)] + { + self.frame_callback.as_mut().map_or(Ok(()), |cb| { + cb(self.frame.as_ref().unwrap(), self.decoded_frames) + })?; + } + + let result = self.frame.take().unwrap().finalize()?; + + #[cfg(test)] + { + self.decoded_frames += 1; + } + + if let Some(state) = result.decoder_state { + self.decoder_state = Some(state); + } else { + self.has_more_frames = false; + } + // TODO(veluca): this code should be integrated in the render pipeline. + if let Some(channels) = result.channels + && let Some(bufs) = output_buffers + { + if self.pixel_format.as_ref().unwrap().color_type == JxlColorType::Grayscale { + for (buf, chan) in bufs.iter_mut().zip(channels.iter()) { + buf.write_from_f32(chan); + } + } else { + bufs[0].write_from_rgb_f32(&channels[0], &channels[1], &channels[2]); + for (buf, chan) in bufs.iter_mut().skip(1).zip(channels.iter().skip(3)) { + buf.write_from_f32(chan); + } + } + } + Ok(None) + } +} diff --git a/third_party/rust/jxl/src/api/inner/process.rs b/third_party/rust/jxl/src/api/inner/process.rs @@ -0,0 +1,120 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +use std::{ + io::IoSliceMut, + ops::{Deref, Range}, +}; + +use crate::error::Result; + +use crate::api::{JxlBitstreamInput, JxlDecoderInner, JxlOutputBuffer, ProcessingResult}; + +// General implementation strategy: +// - Anything that is not a section is read into a small buffer. +// - As soon as we know section sizes, data is read directly into sections. +// When the start of the populated range in `buf` goes past half of its length, +// the data in the buffer is moved back to the beginning. + +pub(super) struct SmallBuffer<const SIZE: usize> { + buf: [u8; SIZE], + range: Range<usize>, +} + +impl<const SIZE: usize> SmallBuffer<SIZE> { + pub(super) fn refill( + &mut self, + mut get_input: impl FnMut(&mut [IoSliceMut]) -> Result<usize, std::io::Error>, + max: Option<usize>, + ) -> Result<usize> { + let mut total = 0; + loop { + if self.range.start >= SIZE / 2 { + let start = self.range.start; + let len = self.range.len(); + let (pre, post) = self.buf.split_at_mut(start); + pre[0..len].copy_from_slice(&post[0..len]); + self.range.start -= start; + self.range.end -= start; + } + if self.range.len() >= SIZE / 2 { + break; + } + let stop = if let Some(max) = max { + (self.range.end + max.saturating_sub(total)).min(SIZE) + } else { + SIZE + }; + let num = get_input(&mut [IoSliceMut::new(&mut self.buf[self.range.end..stop])])?; + total += num; + self.range.end += num; + if num == 0 { + break; + } + } + Ok(total) + } + + pub(super) fn take(&mut self, mut buffers: &mut [IoSliceMut]) -> usize { + let mut num = 0; + while !self.range.is_empty() { + let Some((buf, rest)) = buffers.split_first_mut() else { + break; + }; + buffers = rest; + let len = self.range.len().min(buf.len()); + // Only copy 'len' bytes, not the entire range, to avoid panic when buf is smaller than range + buf[..len].copy_from_slice(&self.buf[self.range.start..self.range.start + len]); + self.range.start += len; + num += len; + } + num + } + + pub(super) fn consume(&mut self, amount: usize) -> usize { + let amount = amount.min(self.range.len()); + self.range.start += amount; + amount + } + + pub(super) fn new() -> Self { + Self { + buf: [0; SIZE], + range: 0..0, + } + } +} + +impl<const SIZE: usize> Deref for SmallBuffer<SIZE> { + type Target = [u8]; + fn deref(&self) -> &Self::Target { + &self.buf[self.range.clone()] + } +} + +impl JxlDecoderInner { + /// Process more of the input file. + /// This function will return when reaching the next decoding stage (i.e. finished decoding + /// file/frame header, or finished decoding a frame). + /// If called when decoding a frame with `None` for buffers, the frame will still be read, + /// but pixel data will not be produced. + pub fn process<In: JxlBitstreamInput>( + &mut self, + input: &mut In, + buffers: Option<&mut [JxlOutputBuffer]>, + ) -> Result<ProcessingResult<(), ()>> { + ProcessingResult::new(self.codestream_parser.process( + &mut self.box_parser, + input, + &self.options, + buffers, + )) + } + + /// Draws all the pixels we have data for. + pub fn flush_pixels(&mut self, _buffers: &mut [JxlOutputBuffer]) -> Result<()> { + todo!() + } +} diff --git a/third_party/rust/jxl/src/api/input.rs b/third_party/rust/jxl/src/api/input.rs @@ -0,0 +1,82 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +use std::io::{BufRead, BufReader, Error, IoSliceMut, Read, Seek, SeekFrom}; + +pub trait JxlBitstreamInput { + /// Returns an estimate bound of the total number of bytes that can be read via `read`. + /// Returning a too-low estimate here can impede parallelism. Returning a too-high + /// estimate can increase memory usage. + fn available_bytes(&mut self) -> Result<usize, Error>; + + /// Fills in `bufs` with more bytes, returning the number of bytes written. + /// Buffers are filled in order and to completion. + fn read(&mut self, bufs: &mut [IoSliceMut]) -> Result<usize, Error>; + + /// Skips up to `bytes` bytes of input. The provided implementation just uses `read`, but in + /// some cases this can be implemented faster. + /// Returns the number of bytes that were skipped. If this returns 0, it is assumed that no + /// more input is available. + fn skip(&mut self, bytes: usize) -> Result<usize, Error> { + let mut bytes = bytes; + const BUF_SIZE: usize = 1024; + let mut skip_buf = [0; BUF_SIZE]; + let mut skipped = 0; + while bytes > 0 { + let num = bytes.min(BUF_SIZE); + self.read(&mut [IoSliceMut::new(&mut skip_buf[..num])])?; + bytes -= num; + skipped += num; + } + Ok(skipped) + } + + /// Un-consumes read bytes. This will only be called at the end of a file stream, + /// to un-read potentially over-read bytes. If ensuring that data is not read past + /// the file end is not required, this method can safely be implemented as a no-op. + /// The provided implementation does nothing. + fn unconsume(&mut self, _count: usize) -> Result<(), Error> { + Ok(()) + } +} + +impl JxlBitstreamInput for &[u8] { + fn available_bytes(&mut self) -> Result<usize, Error> { + Ok(self.len()) + } + + fn read(&mut self, bufs: &mut [IoSliceMut]) -> Result<usize, Error> { + self.read_vectored(bufs) + } + + fn skip(&mut self, bytes: usize) -> Result<usize, Error> { + let num = bytes.min(self.len()); + self.consume(num); + Ok(num) + } +} + +impl<R: Read + Seek> JxlBitstreamInput for BufReader<R> { + fn available_bytes(&mut self) -> Result<usize, Error> { + let pos = self.stream_position()?; + let end = self.seek(SeekFrom::End(0))?; + self.seek(SeekFrom::Start(pos))?; + Ok(end.saturating_sub(pos) as usize) + } + + fn read(&mut self, bufs: &mut [IoSliceMut]) -> Result<usize, Error> { + self.read_vectored(bufs) + } + + fn skip(&mut self, bytes: usize) -> Result<usize, Error> { + let cur = self.stream_position()?; + self.seek(SeekFrom::Current(bytes as i64)) + .map(|x| x.saturating_sub(cur) as usize) + } + + fn unconsume(&mut self, count: usize) -> Result<(), Error> { + self.seek_relative(-(count as i64)) + } +} diff --git a/third_party/rust/jxl/src/api/options.rs b/third_party/rust/jxl/src/api/options.rs @@ -0,0 +1,42 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +pub enum JxlProgressiveMode { + /// Renders all pixels in every call to Process. + Eager, + /// Renders pixels once passes are completed. + Pass, + /// Renders pixels only once the final frame is ready. + FullFrame, +} + +#[non_exhaustive] +pub struct JxlDecoderOptions { + pub adjust_orientation: bool, + pub unpremultiply_alpha: bool, + pub render_spot_colors: bool, + pub coalescing: bool, + pub desired_intensity_target: Option<f32>, + pub skip_preview: bool, + pub progressive_mode: JxlProgressiveMode, + pub xyb_output_linear: bool, + pub enable_output: bool, +} + +impl Default for JxlDecoderOptions { + fn default() -> Self { + Self { + adjust_orientation: true, + unpremultiply_alpha: false, + render_spot_colors: true, + coalescing: true, + skip_preview: false, + desired_intensity_target: None, + progressive_mode: JxlProgressiveMode::Pass, + xyb_output_linear: true, + enable_output: true, + } + } +} diff --git a/third_party/rust/jxl/src/api/output.rs b/third_party/rust/jxl/src/api/output.rs @@ -0,0 +1,173 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +#![allow(unsafe_code)] + +use core::slice; +use std::{marker::PhantomData, mem::MaybeUninit, ops::Range}; + +use num_traits::ToBytes; + +use crate::image::Image; + +pub struct JxlOutputBuffer<'a> { + // Safety invariants: + // - uninit data is never written to `buf`. + // - `buf` is valid for writes for all bytes in the range + // `buf[i*bytes_between_rows..i*bytes_between_rows+bytes_per_row]` for all values of `i` + // from `0` to `num_rows-1`. + // - `self` has exclusive (write) access to the bytes in these ranges. + // - all the bytes in those ranges (and in between) are part of the same allocated object. + // - `num_rows > 0`, `bytes_per_row > 0`, `bytes_per_row <= bytes_between_rows`. + // - The computation `E = bytes_between_rows * (num_rows-1) + bytes_per_row` does not + // overflow and has a result that is at most `isize::MAX`. + buf: *mut MaybeUninit<u8>, + bytes_per_row: usize, + num_rows: usize, + bytes_between_rows: usize, + _ph: PhantomData<&'a mut u8>, +} + +impl<'a> JxlOutputBuffer<'a> { + fn check_vals(num_rows: usize, bytes_per_row: usize, bytes_between_rows: usize) { + assert!(num_rows > 0); + assert!(bytes_per_row > 0); + assert!(bytes_between_rows >= bytes_per_row); + assert!( + bytes_between_rows + .checked_mul(num_rows - 1) + .unwrap() + .checked_add(bytes_per_row) + .unwrap() + <= isize::MAX as usize + ); + } + + /// Creates a new JxlOutputBuffer from raw pointers. + /// It is guaranteed that `buf` will never be used to write uninitialized data. + /// + /// # Safety + /// - `buf` must be valid for writes for all bytes in the range + /// `buf[i*bytes_between_rows..i*bytes_between_rows+bytes_per_row]` for all values of `i` + /// from `0` to `num_rows-1`. + /// - The bytes in these ranges must not be accessed as long as the returned `Self` is in scope. + /// - All the bytes in those ranges (and in between) must be part of the same allocated object. + pub unsafe fn new_from_ptr( + buf: *mut MaybeUninit<u8>, + num_rows: usize, + bytes_per_row: usize, + bytes_between_rows: usize, + ) -> Self { + Self::check_vals(num_rows, bytes_per_row, bytes_between_rows); + // Safety: caller guarantees the buf-related requirements, and other safety invariants are + // checked in check_vals. + Self { + buf, + bytes_per_row, + bytes_between_rows, + num_rows, + _ph: PhantomData, + } + } + + pub fn new(buf: &'a mut [u8], num_rows: usize, bytes_per_row: usize) -> Self { + Self::new_uninit( + // Safety: `new_uninit` guarantees that no uninit data is ever written to the passed-in + // slice. Moreover, `T` and `MaybeUninit<T>` have the same memory layout. + unsafe { slice::from_raw_parts_mut(buf.as_mut_ptr().cast(), buf.len()) }, + num_rows, + bytes_per_row, + ) + } + + /// Creates a new JxlOutputBuffer from a slice of uninit data. + /// It is guaranteed that `buf` will never be used to write uninitalized data. + pub fn new_uninit( + buf: &'a mut [MaybeUninit<u8>], + num_rows: usize, + bytes_per_row: usize, + ) -> Self { + Self::check_vals(num_rows, bytes_per_row, bytes_per_row); + assert!(buf.len() >= bytes_per_row * num_rows); + // Safety note: buf-related requirements follow from this function mutably borrowing a + // slice, and other safety invariants are checked in check_vals. + Self { + buf: buf.as_mut_ptr(), + bytes_per_row, + bytes_between_rows: bytes_per_row, + num_rows, + _ph: PhantomData, + } + } + + /// # Safety + /// No uninit data must be written to the returned slice. + pub unsafe fn get(&mut self, row: usize, cols: Range<usize>) -> &mut [MaybeUninit<u8>] { + assert!(row < self.num_rows); + assert!(cols.start <= cols.end); + assert!(cols.end <= self.bytes_per_row); + let start = row * self.bytes_between_rows + cols.start; + // Safety: `start` is guaranteed to be <= isize::MAX, and `self.buf + start` is guaranteed + // to fit within the same allocated object, as per safety invariants of this struct. + // We checked above that `row` and `cols` satisfy the requirements to apply the safety + // invariant. + let start = unsafe { self.buf.add(start) }; + // Safety: due to the struct safety invariant, we know the entire slice is in a range of + // memory valid for writes. Moreover, the caller promises not to write uninitialized data + // in the returned slice. Finally, as we take self by mutable reference and `self` has + // exclusive access to the slices described in the safety invariant, we know aliasing + // rules will not be violated. + unsafe { slice::from_raw_parts_mut(start, cols.len()) } + } + + // TODO(veluca): the following methods should be removed, as data should be written + // incrementally. + pub(super) fn write_from_rgb_f32(&mut self, r: &Image<f32>, g: &Image<f32>, b: &Image<f32>) { + assert_eq!(r.size(), g.size()); + assert_eq!(r.size(), b.size()); + assert_eq!(self.bytes_per_row, r.size().0 * 12); + let (xsize, ysize) = r.size(); + assert_eq!(ysize, self.num_rows); + for y in 0..ysize { + let rrow = r.as_rect().row(y); + let grow = g.as_rect().row(y); + let brow = b.as_rect().row(y); + // Safety: uninit data is never written in `row`. + let row = unsafe { self.get(y, 0..12 * xsize) }; + let rgb = rrow + .iter() + .zip(grow.iter().zip(brow.iter())) + .map(|(r, (g, b))| { + let mut arr = [0; 12]; + arr[0..4].copy_from_slice(&r.to_ne_bytes()); + arr[4..8].copy_from_slice(&g.to_ne_bytes()); + arr[8..12].copy_from_slice(&b.to_ne_bytes()); + arr + }); + for (out, rgb) in row.chunks_exact_mut(12).zip(rgb) { + for i in 0..12 { + out[i].write(rgb[i]); + } + } + } + } + + pub(super) fn write_from_f32(&mut self, c: &Image<f32>) { + assert_eq!(self.bytes_per_row, c.size().0 * 4); + let (xsize, ysize) = c.size(); + assert_eq!(ysize, self.num_rows); + for y in 0..ysize { + let crow = c.as_rect().row(y); + // Safety: uninit data is never written in `row`. + let row = unsafe { self.get(y, 0..4 * xsize) }; + for (out, v) in row.chunks_exact_mut(4).zip(crow) { + let v = v.to_ne_bytes(); + for i in 0..4 { + out[i].write(v[i]); + } + } + } + } +} diff --git a/third_party/rust/jxl/src/api/signature.rs b/third_party/rust/jxl/src/api/signature.rs @@ -0,0 +1,162 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +use crate::{ + api::ProcessingResult, + error::{Error, Result}, +}; + +/// The magic bytes for a bare JPEG XL codestream. +pub(crate) const CODESTREAM_SIGNATURE: [u8; 2] = [0xff, 0x0a]; +/// The magic bytes for a file using the JPEG XL container format. +pub(crate) const CONTAINER_SIGNATURE: [u8; 12] = + [0, 0, 0, 0xc, b'J', b'X', b'L', b' ', 0xd, 0xa, 0x87, 0xa]; + +#[derive(Debug, PartialEq)] +pub enum JxlSignatureType { + Codestream, + Container, +} + +impl JxlSignatureType { + pub(crate) fn signature(&self) -> &[u8] { + match self { + JxlSignatureType::Container => &CONTAINER_SIGNATURE, + JxlSignatureType::Codestream => &CODESTREAM_SIGNATURE, + } + } +} + +pub(crate) fn check_signature_internal(file_prefix: &[u8]) -> Result<Option<JxlSignatureType>> { + let prefix_len = file_prefix.len(); + + for st in [JxlSignatureType::Codestream, JxlSignatureType::Container] { + let len = st.signature().len(); + // Determine the number of bytes to compare (the length of the shorter slice) + let len_to_check = prefix_len.min(len); + + if file_prefix[..len_to_check] == st.signature()[..len_to_check] { + // The prefix is a valid start. Now, is it complete? + return if prefix_len >= len { + Ok(Some(st)) + } else { + Err(Error::OutOfBounds(len - prefix_len)) + }; + } + } + // The prefix doesn't match the start of any known signature. + Ok(None) +} + +/// Checks if the given buffer starts with a valid JPEG XL signature. +/// +/// # Returns +/// +/// A `ProcessingResult` which is: +/// - `Complete(Some(_))` if a full container or codestream signature is found. +/// - `Complete(None)` if the prefix is definitively not a JXL signature. +/// - `NeedsMoreInput` if the prefix matches a signature but is too short. +pub fn check_signature(file_prefix: &[u8]) -> ProcessingResult<Option<JxlSignatureType>, ()> { + ProcessingResult::new(check_signature_internal(file_prefix)).unwrap() +} + +#[cfg(test)] +mod tests { + use crate::api::{ + CODESTREAM_SIGNATURE, CONTAINER_SIGNATURE, JxlSignatureType, ProcessingResult, + check_signature, + }; + + macro_rules! signature_test { + ($test_name:ident, $bytes:expr, Complete(Some($expected_type:expr))) => { + #[test] + fn $test_name() { + let result = check_signature($bytes); + match result { + ProcessingResult::Complete { result } => { + assert_eq!(result, Some($expected_type)); + } + _ => panic!("Expected Complete(Some(_)), but got {:?}", result), + } + } + }; + ($test_name:ident, $bytes:expr, Complete(None)) => { + #[test] + fn $test_name() { + let result = check_signature($bytes); + match result { + ProcessingResult::Complete { result } => { + assert_eq!(result, None); + } + _ => panic!("Expected Complete(None), but got {:?}", result), + } + } + }; + ($test_name:ident, $bytes:expr, NeedsMoreInput($expected_hint:expr)) => { + #[test] + fn $test_name() { + let result = check_signature($bytes); + match result { + ProcessingResult::NeedsMoreInput { size_hint, .. } => { + assert_eq!(size_hint, $expected_hint); + } + _ => panic!("Expected NeedsMoreInput, but got {:?}", result), + } + } + }; + } + + signature_test!( + full_container_sig, + &CONTAINER_SIGNATURE, + Complete(Some(JxlSignatureType::Container)) + ); + + signature_test!( + partial_container_sig, + &CONTAINER_SIGNATURE[..5], + NeedsMoreInput(CONTAINER_SIGNATURE.len() - 5) + ); + + signature_test!( + full_codestream_sig, + &CODESTREAM_SIGNATURE, + Complete(Some(JxlSignatureType::Codestream)) + ); + + signature_test!( + partial_codestream_sig, + &CODESTREAM_SIGNATURE[..1], + NeedsMoreInput(CODESTREAM_SIGNATURE.len() - 1) + ); + + signature_test!( + empty_prefix, + &[], + NeedsMoreInput(CODESTREAM_SIGNATURE.len()) + ); + + signature_test!(invalid_sig, &[0x12, 0x34, 0x56, 0x77], Complete(None)); + + signature_test!( + container_with_extra_data, + &{ + let mut data = CONTAINER_SIGNATURE.to_vec(); + data.extend_from_slice(&[0x11, 0x22, 0x33]); + data + }, + Complete(Some(JxlSignatureType::Container)) + ); + + signature_test!( + codestream_with_extra_data, + &{ + let mut data = CODESTREAM_SIGNATURE.to_vec(); + data.extend_from_slice(&[0x44, 0x55, 0x66]); + data + }, + Complete(Some(JxlSignatureType::Codestream)) + ); +} diff --git a/third_party/rust/jxl/src/api/test.rs b/third_party/rust/jxl/src/api/test.rs @@ -0,0 +1,67 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +macro_rules! create_output_buffers { + ($info:expr, $pixel_format:expr, $bufs:ident, $slices:ident) => { + use crate::api::{JxlColorType::Rgb, JxlOutputBuffer}; + use std::mem::MaybeUninit; + + let orientation = $info.orientation; + let (width, height) = $info.size; + + let (buffer_width, buffer_height) = if orientation.is_transposing() { + (height, width) + } else { + (width, height) + }; + + let num_channels = $pixel_format.color_type.samples_per_pixel(); + + let mut $bufs: Vec<Vec<MaybeUninit<u8>>> = Vec::new(); + + // For RGB images, first buffer holds interleaved RGB data + match $pixel_format.color_type == Rgb { + true => { + $bufs.push(vec![ + MaybeUninit::uninit(); + buffer_width * buffer_height * 12 + ]); + for _ in 3..num_channels { + $bufs.push(vec![ + MaybeUninit::uninit(); + buffer_width * buffer_height * 4 + ]); + } + } + false => { + // For grayscale or other formats, one buffer per channel + for _ in 0..num_channels { + $bufs.push(vec![ + MaybeUninit::uninit(); + buffer_width * buffer_height * 4 + ]); + } + } + } + + let mut $slices: Vec<JxlOutputBuffer> = $bufs + .iter_mut() + .enumerate() + .map(|(i, buffer)| { + let bytes_per_pixel = if i == 0 && $pixel_format.color_type == Rgb { + 12 // Interleaved RGB + } else { + 4 // Single channel + }; + JxlOutputBuffer::new_uninit( + buffer.as_mut_slice(), + buffer_height, + bytes_per_pixel * buffer_width, + ) + }) + .collect(); + }; +} +pub(crate) use create_output_buffers; diff --git a/third_party/rust/jxl/src/bit_reader.rs b/third_party/rust/jxl/src/bit_reader.rs @@ -0,0 +1,236 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +use std::fmt::Debug; + +use crate::{error::Error, util::tracing_wrappers::*}; +use byteorder::{ByteOrder, LittleEndian}; + +/// Reads bits from a sequence of bytes. +#[derive(Clone)] +pub struct BitReader<'a> { + data: &'a [u8], + bit_buf: u64, + bits_in_buf: usize, + total_bits_read: usize, +} + +impl Debug for BitReader<'_> { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!( + f, + "BitReader{{ data: [{} bytes], bit_buf: {:0width$b}, total_bits_read: {} }}", + self.data.len(), + self.bit_buf, + self.total_bits_read, + width = self.bits_in_buf + ) + } +} + +pub const MAX_BITS_PER_CALL: usize = 56; + +impl<'a> BitReader<'a> { + /// Constructs a BitReader for a given range of data. + pub fn new(data: &[u8]) -> BitReader<'_> { + BitReader { + data, + bit_buf: 0, + bits_in_buf: 0, + total_bits_read: 0, + } + } + + /// Reads `num` bits from the buffer without consuming them. + pub fn peek(&mut self, num: usize) -> u64 { + debug_assert!(num <= MAX_BITS_PER_CALL); + self.refill(); + self.bit_buf & ((1u64 << num) - 1) + } + + /// Advances by `num` bits. Similar to `skip_bits`, but bits must be in the buffer. + pub fn consume(&mut self, num: usize) -> Result<(), Error> { + if self.bits_in_buf < num { + return Err(Error::OutOfBounds((num - self.bits_in_buf).div_ceil(8))); + } + self.bit_buf >>= num; + self.bits_in_buf -= num; + self.total_bits_read = self.total_bits_read.wrapping_add(num); + Ok(()) + } + + /// Reads `num` bits from the buffer. + /// ``` + /// # use jxl::bit_reader::BitReader; + /// let mut br = BitReader::new(&[0, 1]); + /// assert_eq!(br.read(8)?, 0); + /// assert_eq!(br.read(4)?, 1); + /// assert_eq!(br.read(4)?, 0); + /// assert_eq!(br.total_bits_read(), 16); + /// assert!(br.read(1).is_err()); + /// # Ok::<(), jxl::error::Error>(()) + /// ``` + pub fn read(&mut self, num: usize) -> Result<u64, Error> { + let ret = self.peek(num); + self.consume(num)?; + Ok(ret) + } + + /// Returns the total number of bits that have been read or skipped. + pub fn total_bits_read(&self) -> usize { + self.total_bits_read + } + + /// Returns the total number of bits that can still be read or skipped. + pub fn total_bits_available(&self) -> usize { + self.data.len() * 8 + self.bits_in_buf + } + + /// Skips `num` bits. + /// ``` + /// # use jxl::bit_reader::BitReader; + /// let mut br = BitReader::new(&[0, 1]); + /// assert_eq!(br.read(8)?, 0); + /// br.skip_bits(4)?; + /// assert_eq!(br.total_bits_read(), 12); + /// # Ok::<(), jxl::error::Error>(()) + /// ``` + #[inline(never)] + pub fn skip_bits(&mut self, mut n: usize) -> Result<(), Error> { + // Check if we can skip within the current buffer + if let Some(next_remaining_bits) = self.bits_in_buf.checked_sub(n) { + self.total_bits_read += n; + self.bits_in_buf = next_remaining_bits; + self.bit_buf >>= n; + return Ok(()); + } + + // Adjust the number of bits to skip and reset the buffer + n -= self.bits_in_buf; + self.total_bits_read += self.bits_in_buf; + self.bit_buf = 0; + self.bits_in_buf = 0; + + // Check if the remaining bits to skip exceed the total bits in `data` + let bits_available = self.data.len() * 8; + if n > bits_available { + self.total_bits_read += bits_available; + return Err(Error::OutOfBounds(n - bits_available)); + } + + // Skip bytes directly in `data`, then handle leftover bits + self.total_bits_read += n / 8 * 8; + self.data = &self.data[n / 8..]; + n %= 8; + + // Refill the buffer and adjust for any remaining bits + self.refill(); + let to_consume = self.bits_in_buf.min(n); + // The bits loaded by refill() haven't been counted in total_bits_read yet, + // so we add (not subtract) the bits we're consuming. The original code + // incorrectly subtracted here, causing underflow when skip_bits was called + // on a fresh BitReader. + self.total_bits_read += to_consume; + n -= to_consume; + self.bit_buf >>= to_consume; + self.bits_in_buf -= to_consume; + if n > 0 { + Err(Error::OutOfBounds(n)) + } else { + Ok(()) + } + } + + /// Return the number of bits + pub fn bits_to_next_byte(&self) -> usize { + let byte_boundary = self.total_bits_read.div_ceil(8) * 8; + byte_boundary - self.total_bits_read + } + + /// Jumps to the next byte boundary. The skipped bytes have to be 0. + /// ``` + /// # use jxl::bit_reader::BitReader; + /// let mut br = BitReader::new(&[0, 1]); + /// assert_eq!(br.read(8)?, 0); + /// br.skip_bits(4)?; + /// br.jump_to_byte_boundary()?; + /// assert_eq!(br.total_bits_read(), 16); + /// # Ok::<(), jxl::error::Error>(()) + /// ``` + #[inline(never)] + pub fn jump_to_byte_boundary(&mut self) -> Result<(), Error> { + if self.read(self.bits_to_next_byte())? != 0 { + return Err(Error::NonZeroPadding); + } + Ok(()) + } + + fn refill(&mut self) { + // See Refill() in C++ code. + if self.data.len() >= 8 { + let bits = LittleEndian::read_u64(self.data); + self.bit_buf |= bits << self.bits_in_buf; + let read_bytes = (63 - self.bits_in_buf) >> 3; + self.bits_in_buf |= 56; + self.data = &self.data[read_bytes..]; + debug_assert!(56 <= self.bits_in_buf && self.bits_in_buf < 64); + } else { + self.refill_slow() + } + } + + #[inline(never)] + fn refill_slow(&mut self) { + while self.bits_in_buf < 56 { + if self.data.is_empty() { + return; + } + self.bit_buf |= (self.data[0] as u64) << self.bits_in_buf; + self.bits_in_buf += 8; + self.data = &self.data[1..]; + } + } + + /// Splits off a separate BitReader to handle the next `n` *full* bytes. + /// If `self` is not aligned to a byte boundary, it skips to the next byte boundary. + /// `self` is automatically advanced by `n` bytes. + pub fn split_at(&mut self, n: usize) -> Result<BitReader<'a>, Error> { + self.jump_to_byte_boundary()?; + let mut ret = Self { ..*self }; + self.skip_bits(n * 8)?; + let bytes_in_buf = ret.bits_in_buf / 8; + if n > bytes_in_buf { + // Prevent the returned bitreader from over-reading. + ret.data = &ret.data[..n - bytes_in_buf]; + } else { + ret.bits_in_buf = n * 8; + ret.bit_buf &= (1u64 << (n * 8)) - 1; + ret.data = &[]; + } + debug!(?n, ret=?ret); + Ok(ret) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_skip_bits_on_fresh_reader() { + // This test checks if skip_bits works correctly on a fresh BitReader + let data = [0x12, 0x34, 0x56, 0x78]; + let mut br = BitReader::new(&data); + + // Try to skip 1 bit on a fresh reader - this should work + br.skip_bits(1) + .expect("skip_bits should work on fresh reader"); + assert_eq!(br.total_bits_read(), 1); + + // Read the next 7 bits to complete the byte + let val = br.read(7).expect("read should work"); + assert_eq!(val, 0x12 >> 1); // Should get the lower 7 bits of 0x12 + } +} diff --git a/third_party/rust/jxl/src/color/mod.rs b/third_party/rust/jxl/src/color/mod.rs @@ -0,0 +1,6 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +pub mod tf; diff --git a/third_party/rust/jxl/src/color/tf.rs b/third_party/rust/jxl/src/color/tf.rs @@ -0,0 +1,601 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +use crate::util::eval_rational_poly; + +const SRGB_POWTABLE_UPPER: [u8; 16] = [ + 0x00, 0x0a, 0x19, 0x26, 0x32, 0x41, 0x4d, 0x5c, 0x68, 0x75, 0x83, 0x8f, 0xa0, 0xaa, 0xb9, 0xc6, +]; + +const SRGB_POWTABLE_LOWER: [u8; 16] = [ + 0x00, 0xb7, 0x04, 0x0d, 0xcb, 0xe7, 0x41, 0x68, 0x51, 0xd1, 0xeb, 0xf2, 0x00, 0xb7, 0x04, 0x0d, +]; + +/// Converts the linear samples with the sRGB transfer curve. +// Fast linear to sRGB conversion, ported from libjxl. Max error ~1.7e-4 +pub fn linear_to_srgb_fast(samples: &mut [f32]) { + for s in samples { + let v = s.to_bits() & 0x7fff_ffff; + let v_adj = f32::from_bits((v | 0x3e80_0000) & 0x3eff_ffff); + let pow = 0.059914046f32; + let pow = pow * v_adj - 0.10889456; + let pow = pow * v_adj + 0.107963754; + let pow = pow * v_adj + 0.018092343; + + // `mul` won't be used when `v` is small. + let idx = (v >> 23).wrapping_sub(118) as usize & 0xf; + let mul = 0x4000_0000 + | (u32::from(SRGB_POWTABLE_UPPER[idx]) << 18) + | (u32::from(SRGB_POWTABLE_LOWER[idx]) << 10); + + let v = f32::from_bits(v); + let small = v * 12.92; + let acc = pow * f32::from_bits(mul) - 0.055; + + *s = if v <= 0.0031308 { small } else { acc }.copysign(*s); + } +} + +/// Converts the linear samples with the sRGB transfer curve. +// Max error ~5e-7 +pub fn linear_to_srgb(samples: &mut [f32]) { + #[allow(clippy::excessive_precision)] + const P: [f32; 5] = [ + -5.135152395e-4, + 5.287254571e-3, + 3.903842876e-1, + 1.474205315, + 7.352629620e-1, + ]; + + #[allow(clippy::excessive_precision)] + const Q: [f32; 5] = [ + 1.004519624e-2, + 3.036675394e-1, + 1.340816930, + 9.258482155e-1, + 2.424867759e-2, + ]; + + for x in samples { + let a = x.abs(); + *x = if a <= 0.0031308 { + a * 12.92 + } else { + eval_rational_poly(a.sqrt(), P, Q) + } + .copysign(*x); + } +} + +/// Converts samples in sRGB transfer curve to linear. Inverse of `linear_to_srgb`. +pub fn srgb_to_linear(samples: &mut [f32]) { + #[allow(clippy::excessive_precision)] + const P: [f32; 5] = [ + 2.200248328e-4, + 1.043637593e-2, + 1.624820318e-1, + 7.961564959e-1, + 8.210152774e-1, + ]; + + #[allow(clippy::excessive_precision)] + const Q: [f32; 5] = [ + 2.631846970e-1, + 1.076976492, + 4.987528350e-1, + -5.512498495e-2, + 6.521209011e-3, + ]; + + for x in samples { + let a = x.abs(); + *x = if a <= 0.04045 { + a / 12.92 + } else { + eval_rational_poly(a, P, Q) + } + .copysign(*x); + } +} + +/// Converts the linear samples with the BT.709 transfer curve. +pub fn linear_to_bt709(samples: &mut [f32]) { + for s in samples { + let a = s.abs(); + *s = if a <= 0.018 { + a * 4.5 + } else { + crate::util::fast_powf(a, 0.45).mul_add(1.099, -0.099) + } + .copysign(*s); + } +} + +/// Converts samples in BT.709 transfer curve to linear. Inverse of `linear_to_bt709`. +pub fn bt709_to_linear(samples: &mut [f32]) { + for s in samples { + let a = s.abs(); + *s = if a <= 0.081 { + a / 4.5 + } else { + crate::util::fast_powf(a.mul_add(1.0 / 1.099, 0.099 / 1.099), 1.0 / 0.45) + } + .copysign(*s); + } +} + +const PQ_M1: f64 = 2610.0 / 16384.0; +const PQ_M2: f64 = (2523.0 / 4096.0) * 128.0; +const PQ_C1: f64 = 3424.0 / 4096.0; +const PQ_C2: f64 = (2413.0 / 4096.0) * 32.0; +const PQ_C3: f64 = (2392.0 / 4096.0) * 32.0; + +/// Converts linear sample to PQ signal using PQ inverse EOTF, where linear sample value of 1.0 +/// represents `intensity_target` display nits. +/// +/// This version uses original EOTF using double precision arithmetic internally. +pub fn linear_to_pq_precise(intensity_target: f32, samples: &mut [f32]) { + let mult = intensity_target as f64 * 10000f64.recip(); + + for s in samples { + if *s == 0.0 { + continue; + } + + let a = s.abs() as f64; + let xp = (a * mult).powf(PQ_M1); + let num = PQ_C1 + xp * PQ_C2; + let den = 1.0 + xp * PQ_C3; + let e = (num / den).powf(PQ_M2); + *s = (e as f32).copysign(*s); + } +} + +/// Converts PQ signal to linear sample using PQ EOTF, where linear sample value of 1.0 represents +/// `intensity_target` display nits. +/// +/// This version uses original EOTF using double precision arithmetic internally. +pub fn pq_to_linear_precise(intensity_target: f32, samples: &mut [f32]) { + let mult = 10000.0 / intensity_target as f64; + + for s in samples { + if *s == 0.0 { + continue; + } + + let a = s.abs() as f64; + let xp = a.powf(PQ_M2.recip()); + let num = (xp - PQ_C1).max(0.0); + let den = PQ_C2 - PQ_C3 * xp; + let y = (num / den).powf(PQ_M1.recip()); + *s = ((y * mult) as f32).copysign(*s); + } +} + +const PQ_EOTF_P: [f32; 5] = [ + 2.6297566e-4, + -6.235531e-3, + 7.386023e-1, + 2.6455317, + 5.500349e-1, +]; +const PQ_EOTF_Q: [f32; 5] = [ + 4.213501e2, + -4.2873682e2, + 1.7436467e2, + -3.3907887e1, + 2.6771877, +]; + +const PQ_INV_EOTF_P: [f32; 5] = [1.351392e-2, -1.095778, 5.522776e1, 1.492516e2, 4.838434e1]; +const PQ_INV_EOTF_Q: [f32; 5] = [1.012416, 2.016708e1, 9.26371e1, 1.120607e2, 2.590418e1]; +const PQ_INV_EOTF_P_SMALL: [f32; 5] = [ + 9.863406e-6, + 3.881234e-1, + 1.352821e2, + 6.889862e4, + -2.864824e5, +]; +const PQ_INV_EOTF_Q_SMALL: [f32; 5] = + [3.371868e1, 1.477719e3, 1.608477e4, -4.389884e4, -2.072546e5]; + +/// Converts linear sample to PQ signal using PQ inverse EOTF, where linear sample value of 1.0 +/// represents `intensity_target` display nits. +/// +/// This version uses approximate curve using rational polynomial. +// Max error: ~7e-7 at intensity_target = 10000 +pub fn linear_to_pq(intensity_target: f32, samples: &mut [f32]) { + let y_mult = intensity_target * 10000f32.recip(); + + for s in samples { + let a = s.abs(); + let a_scaled = a * y_mult; + let a_1_4 = a_scaled.sqrt().sqrt(); + + let y = if a < 1e-4 { + eval_rational_poly(a_1_4, PQ_INV_EOTF_P_SMALL, PQ_INV_EOTF_Q_SMALL) + } else { + eval_rational_poly(a_1_4, PQ_INV_EOTF_P, PQ_INV_EOTF_Q) + }; + + *s = y.copysign(*s); + } +} + +/// Converts PQ signal to linear sample using PQ EOTF, where linear sample value of 1.0 represents +/// `intensity_target` display nits. +/// +/// This version uses approximate curve using rational polynomial. +// Max error: ~3e-6 at intensity_target = 10000 +pub fn pq_to_linear(intensity_target: f32, samples: &mut [f32]) { + let y_mult = 10000.0 / intensity_target; + + for s in samples { + let a = s.abs(); + // a + a * a + let x = a.mul_add(a, a); + let y = eval_rational_poly(x, PQ_EOTF_P, PQ_EOTF_Q); + *s = (y * y_mult).copysign(*s); + } +} + +const HLG_A: f64 = 0.17883277; +const HLG_B: f64 = 1.0 - 4.0 * HLG_A; +const HLG_C: f64 = 0.5599107295; + +fn hlg_ootf_inner_precise(exp: f64, [lr, lg, lb]: [f32; 3], [sr, sg, sb]: [&mut [f32]; 3]) { + if exp.abs() < 0.1 { + return; + } + + let lr = lr as f64; + let lg = lg as f64; + let lb = lb as f64; + for ((r, g), b) in std::iter::zip(sr, sg).zip(sb) { + let dr = *r as f64; + let dg = *g as f64; + let db = *b as f64; + let mixed = dr.mul_add(lr, dg.mul_add(lg, db * lb)); + let mult = mixed.powf(exp); + *r = (dr * mult) as f32; + *g = (dg * mult) as f32; + *b = (db * mult) as f32; + } +} + +fn hlg_ootf_inner(exp: f32, [lr, lg, lb]: [f32; 3], [sr, sg, sb]: [&mut [f32]; 3]) { + if exp.abs() < 0.1 { + return; + } + + for ((r, g), b) in std::iter::zip(sr, sg).zip(sb) { + let mixed = r.mul_add(lr, g.mul_add(lg, *b * lb)); + let mult = crate::util::fast_powf(mixed, exp); + *r *= mult; + *g *= mult; + *b *= mult; + } +} + +/// Converts scene-referred linear samples to display-referred linear samples using HLG OOTF. +/// +/// This version uses double precision arithmetic internally. +pub fn hlg_scene_to_display_precise( + intensity_display: f32, + luminance_rgb: [f32; 3], + samples_rgb: [&mut [f32]; 3], +) { + let system_gamma = 1.2f64 * 1.111f64.powf((intensity_display as f64 / 1e3).log2()); + let gamma_sub_one = system_gamma - 1.0; + hlg_ootf_inner_precise(gamma_sub_one, luminance_rgb, samples_rgb); +} + +/// Converts display-referred linear samples to scene-referred linear samples using HLG inverse +/// OOTF. +/// +/// This version uses double precision arithmetic internally. +pub fn hlg_display_to_scene_precise( + intensity_display: f32, + luminance_rgb: [f32; 3], + samples_rgb: [&mut [f32]; 3], +) { + let system_gamma = 1.2f64 * 1.111f64.powf((intensity_display as f64 / 1e3).log2()); + let one_sub_gamma = 1.0 - system_gamma; + hlg_ootf_inner_precise(one_sub_gamma / system_gamma, luminance_rgb, samples_rgb); +} + +/// Converts scene-referred linear samples to display-referred linear samples using HLG OOTF. +/// +/// This version uses `fast_powf` to compute power function. +pub fn hlg_scene_to_display( + intensity_display: f32, + luminance_rgb: [f32; 3], + samples_rgb: [&mut [f32]; 3], +) { + let system_gamma = 1.2f32 * 1.111f32.powf((intensity_display / 1e3).log2()); + let gamma_sub_one = system_gamma - 1.0; + hlg_ootf_inner(gamma_sub_one, luminance_rgb, samples_rgb); +} + +/// Converts display-referred linear samples to scene-referred linear samples using HLG inverse +/// OOTF. +/// +/// This version uses `fast_powf` to compute power function. +pub fn hlg_display_to_scene( + intensity_display: f32, + luminance_rgb: [f32; 3], + samples_rgb: [&mut [f32]; 3], +) { + let system_gamma = 1.2f32 * 1.111f32.powf((intensity_display / 1e3).log2()); + let one_sub_gamma = 1.0 - system_gamma; + hlg_ootf_inner(one_sub_gamma / system_gamma, luminance_rgb, samples_rgb); +} + +/// Converts scene-referred linear sample to HLG signal. +/// +/// This version uses double precision arithmetic internally. +pub fn scene_to_hlg_precise(samples: &mut [f32]) { + for s in samples { + let a = s.abs() as f64; + let y = if a <= 1.0 / 12.0 { + (3.0 * a).sqrt() + } else { + // TODO(tirr-c): maybe use mul_add? + HLG_A * (12.0 * a - HLG_B).ln() + HLG_C + }; + *s = (y as f32).copysign(*s); + } +} + +/// Converts HLG signal to scene-referred linear sample. +/// +/// This version uses double precision arithmetic internally. +pub fn hlg_to_scene_precise(samples: &mut [f32]) { + for s in samples { + let a = s.abs() as f64; + let y = if a <= 0.5 { + a * a / 3.0 + } else { + (((a - HLG_C) / HLG_A).exp() + HLG_B) / 12.0 + }; + *s = (y as f32).copysign(*s); + } +} + +/// Converts scene-referred linear sample to HLG signal. +/// +/// This version uses `fast_log2f` to apply logarithmic function. +// Max error: ~5e-7 +pub fn scene_to_hlg(samples: &mut [f32]) { + for s in samples { + let a = s.abs(); + let y = if a <= 1.0 / 12.0 { + (3.0 * a).sqrt() + } else { + // TODO(tirr-c): maybe use mul_add? + let log = crate::util::fast_log2f(12.0 * a - HLG_B as f32); + // log2 x = ln x / ln 2, therefore ln x = (ln 2)(log2 x) + (HLG_A * std::f64::consts::LN_2) as f32 * log + HLG_C as f32 + }; + *s = y.copysign(*s); + } +} + +/// Converts HLG signal to scene-referred linear sample. +/// +/// This version uses `fast_pow2f` to apply logarithmic function. +// Max error: ~5e-6 +pub fn hlg_to_scene(samples: &mut [f32]) { + for s in samples { + let a = s.abs(); + let y = if a <= 0.5 { + a * a / 3.0 + } else { + const POW: f32 = (std::f64::consts::LOG2_E / HLG_A) as f32; + const ADD: f32 = (HLG_B / 12.0) as f32; + // TODO(OneDeuxTriSeiGo): replace raw constant with the below equation + // when std::f64::exp() can is available as a const fn. + // + // Equation: ((-HLG_B / HLG_A).exp() / 12.0) + // Constant: 0.003_639_807_079_052_639 + const MUL: f32 = 0.003_639_807; + + // TODO(OneDeuxTriSeiGo): maybe use mul_add? + crate::util::fast_pow2f(a * POW) * MUL + ADD + }; + *s = y.copysign(*s); + } +} + +#[cfg(test)] +mod test { + use test_log::test; + + use super::*; + use crate::util::test::assert_all_almost_abs_eq; + + fn arb_samples( + u: &mut arbtest::arbitrary::Unstructured, + ) -> arbtest::arbitrary::Result<Vec<f32>> { + const DENOM: u32 = 1 << 24; + + let mut samples = Vec::new(); + + // uniform distribution in [-1.0, 1.0] + while !u.is_empty() { + let a: u32 = u.int_in_range(0..=DENOM)?; + let signed: bool = u.arbitrary()?; + let x = a as f32 / DENOM as f32; + samples.push(if signed { -x } else { x }); + } + + Ok(samples) + } + + #[test] + fn srgb_roundtrip_arb() { + arbtest::arbtest(|u| { + let samples = arb_samples(u)?; + let mut output = samples.clone(); + + linear_to_srgb(&mut output); + srgb_to_linear(&mut output); + assert_all_almost_abs_eq(&output, &samples, 2e-6); + Ok(()) + }); + } + + #[test] + fn bt709_roundtrip_arb() { + arbtest::arbtest(|u| { + let samples = arb_samples(u)?; + let mut output = samples.clone(); + + linear_to_bt709(&mut output); + bt709_to_linear(&mut output); + assert_all_almost_abs_eq(&output, &samples, 5e-6); + Ok(()) + }); + } + + #[test] + fn linear_to_srgb_fast_arb() { + arbtest::arbtest(|u| { + let mut samples = arb_samples(u)?; + let mut fast = samples.clone(); + + linear_to_srgb(&mut samples); + linear_to_srgb_fast(&mut fast); + assert_all_almost_abs_eq(&samples, &fast, 1.7e-4); + Ok(()) + }); + } + + #[test] + fn linear_to_pq_arb() { + arbtest::arbtest(|u| { + let intensity_target = u.int_in_range(9900..=10100)? as f32; + let mut samples = arb_samples(u)?; + let mut precise = samples.clone(); + + linear_to_pq(intensity_target, &mut samples); + linear_to_pq_precise(intensity_target, &mut precise); + // Error seems to increase at intensity_target < 10000 + assert_all_almost_abs_eq(&samples, &precise, 8e-7); + Ok(()) + }); + } + + #[test] + fn pq_to_linear_arb() { + arbtest::arbtest(|u| { + let intensity_target = u.int_in_range(9900..=10100)? as f32; + let mut samples = arb_samples(u)?; + let mut precise = samples.clone(); + + pq_to_linear(intensity_target, &mut samples); + pq_to_linear_precise(intensity_target, &mut precise); + assert_all_almost_abs_eq(&samples, &precise, 3e-6); + Ok(()) + }); + } + + #[test] + fn hlg_ootf_arb() { + arbtest::arbtest(|u| { + let intensity_target = u.int_in_range(900..=1100)? as f32; + + let lr = 0.2 + u.int_in_range(0..=255)? as f32 / 255.0; + let lb = 0.2 + u.int_in_range(0..=255)? as f32 / 255.0; + let lg = 1.0 - lr - lb; + let luminance_rgb = [lr, lg, lb]; + + let r = u.int_in_range(0u32..=(1 << 24))? as f32 / (1 << 24) as f32; + let g = u.int_in_range(0u32..=(1 << 24))? as f32 / (1 << 24) as f32; + let b = u.int_in_range(0u32..=(1 << 24))? as f32 / (1 << 24) as f32; + + let mut fast_r = r; + let mut fast_g = g; + let mut fast_b = b; + let fast = [ + std::slice::from_mut(&mut fast_r), + std::slice::from_mut(&mut fast_g), + std::slice::from_mut(&mut fast_b), + ]; + hlg_display_to_scene(intensity_target, luminance_rgb, fast); + + let mut precise_r = r; + let mut precise_g = g; + let mut precise_b = b; + let precise = [ + std::slice::from_mut(&mut precise_r), + std::slice::from_mut(&mut precise_g), + std::slice::from_mut(&mut precise_b), + ]; + hlg_display_to_scene(intensity_target, luminance_rgb, precise); + + assert_all_almost_abs_eq( + &[fast_r, fast_g, fast_b], + &[precise_r, precise_g, precise_b], + 7.2e-7, + ); + + let mut fast_r = r; + let mut fast_g = g; + let mut fast_b = b; + let fast = [ + std::slice::from_mut(&mut fast_r), + std::slice::from_mut(&mut fast_g), + std::slice::from_mut(&mut fast_b), + ]; + hlg_scene_to_display(intensity_target, luminance_rgb, fast); + + let mut precise_r = r; + let mut precise_g = g; + let mut precise_b = b; + let precise = [ + std::slice::from_mut(&mut precise_r), + std::slice::from_mut(&mut precise_g), + std::slice::from_mut(&mut precise_b), + ]; + hlg_scene_to_display(intensity_target, luminance_rgb, precise); + + assert_all_almost_abs_eq( + &[fast_r, fast_g, fast_b], + &[precise_r, precise_g, precise_b], + 7.2e-7, + ); + + Ok(()) + }); + } + + #[test] + fn scene_to_hlg_arb() { + arbtest::arbtest(|u| { + let mut samples = arb_samples(u)?; + let mut precise = samples.clone(); + + scene_to_hlg(&mut samples); + scene_to_hlg_precise(&mut precise); + assert_all_almost_abs_eq(&samples, &precise, 5e-7); + Ok(()) + }); + } + + #[test] + fn hlg_to_scene_arb() { + arbtest::arbtest(|u| { + let mut samples = arb_samples(u)?; + let mut precise = samples.clone(); + + hlg_to_scene(&mut samples); + hlg_to_scene_precise(&mut precise); + assert_all_almost_abs_eq(&samples, &precise, 5e-6); + Ok(()) + }); + } +} diff --git a/third_party/rust/jxl/src/container/box_header.rs b/third_party/rust/jxl/src/container/box_header.rs @@ -0,0 +1,113 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. +// +// Originally written for jxl-oxide. + +use crate::error::Error; + +/// Box header used in JPEG XL containers. +#[derive(Debug, Clone)] +pub struct ContainerBoxHeader { + ty: ContainerBoxType, + box_size: Option<u64>, + is_last: bool, +} + +pub enum HeaderParseResult { + Done { + header: ContainerBoxHeader, + header_size: usize, + }, + NeedMoreData, +} + +impl ContainerBoxHeader { + pub(super) fn parse(buf: &[u8]) -> Result<HeaderParseResult, Error> { + let (tbox, box_size, header_size) = match *buf { + [ + 0, + 0, + 0, + 1, + t0, + t1, + t2, + t3, + s0, + s1, + s2, + s3, + s4, + s5, + s6, + s7, + .., + ] => { + let xlbox = u64::from_be_bytes([s0, s1, s2, s3, s4, s5, s6, s7]); + let tbox = ContainerBoxType([t0, t1, t2, t3]); + let xlbox = xlbox.checked_sub(16).ok_or(Error::InvalidBox)?; + (tbox, Some(xlbox), 16) + } + [s0, s1, s2, s3, t0, t1, t2, t3, ..] => { + let sbox = u32::from_be_bytes([s0, s1, s2, s3]); + let tbox = ContainerBoxType([t0, t1, t2, t3]); + let sbox = if sbox == 0 { + None + } else if let Some(sbox) = sbox.checked_sub(8) { + Some(sbox as u64) + } else { + return Err(Error::InvalidBox); + }; + (tbox, sbox, 8) + } + _ => return Ok(HeaderParseResult::NeedMoreData), + }; + let is_last = box_size.is_none(); + + let header = Self { + ty: tbox, + box_size, + is_last, + }; + Ok(HeaderParseResult::Done { + header, + header_size, + }) + } +} + +impl ContainerBoxHeader { + #[inline] + pub fn box_type(&self) -> ContainerBoxType { + self.ty + } + + #[inline] + pub fn box_size(&self) -> Option<u64> { + self.box_size + } + + #[inline] + pub fn is_last(&self) -> bool { + self.is_last + } +} + +#[derive(Debug, Copy, Clone, Eq, PartialEq, Hash)] +pub struct ContainerBoxType(pub [u8; 4]); + +impl ContainerBoxType { + pub const JXL: Self = Self(*b"JXL "); + pub const FILE_TYPE: Self = Self(*b"ftyp"); + pub const JXL_LEVEL: Self = Self(*b"jxll"); + pub const JUMBF: Self = Self(*b"jumb"); + pub const EXIF: Self = Self(*b"Exif"); + pub const XML: Self = Self(*b"xml "); + pub const BROTLI_COMPRESSED: Self = Self(*b"brob"); + pub const FRAME_INDEX: Self = Self(*b"jxli"); + pub const CODESTREAM: Self = Self(*b"jxlc"); + pub const PARTIAL_CODESTREAM: Self = Self(*b"jxlp"); + pub const JPEG_RECONSTRUCTION: Self = Self(*b"jbrd"); +} diff --git a/third_party/rust/jxl/src/container/mod.rs b/third_party/rust/jxl/src/container/mod.rs @@ -0,0 +1,187 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. +// +// Originally written for jxl-oxide. + +pub mod box_header; +pub mod parse; + +use box_header::*; +pub use parse::ParseEvent; +use parse::*; + +/// Container format parser. +#[derive(Debug, Default)] +pub struct ContainerParser { + state: DetectState, + jxlp_index_state: JxlpIndexState, + previous_consumed_bytes: usize, +} + +#[derive(Debug, Default)] +enum DetectState { + #[default] + WaitingSignature, + WaitingBoxHeader, + WaitingJxlpIndex(ContainerBoxHeader), + InAuxBox { + #[allow(unused)] + header: ContainerBoxHeader, + bytes_left: Option<usize>, + }, + InCodestream { + kind: BitstreamKind, + bytes_left: Option<usize>, + }, +} + +/// Structure of the decoded bitstream. +#[derive(Debug, Copy, Clone, Eq, PartialEq)] +pub enum BitstreamKind { + /// Decoder can't determine structure of the bitstream. + Unknown, + /// Bitstream is a direct JPEG XL codestream without box structure. + BareCodestream, + /// Bitstream is a JPEG XL container with box structure. + Container, + /// Bitstream is not a valid JPEG XL image. + Invalid, +} + +#[derive(Debug, Copy, Clone, Eq, PartialEq, Default)] +enum JxlpIndexState { + #[default] + Initial, + SingleJxlc, + Jxlp(u32), + JxlpFinished, +} + +impl ContainerParser { + pub fn new() -> Self { + Self::default() + } + + pub fn kind(&self) -> BitstreamKind { + match self.state { + DetectState::WaitingSignature => BitstreamKind::Unknown, + DetectState::WaitingBoxHeader + | DetectState::WaitingJxlpIndex(..) + | DetectState::InAuxBox { .. } => BitstreamKind::Container, + DetectState::InCodestream { kind, .. } => kind, + } + } + + /// Parses input buffer and generates parser events. + /// + /// The parser might not fully consume the buffer. Use [`previous_consumed_bytes`] to get how + /// many bytes are consumed. Bytes not consumed by the parser should be processed again. + /// + /// [`previous_consumed_bytes`]: ContainerParser::previous_consumed_bytes + pub fn process_bytes<'inner, 'buf>( + &'inner mut self, + input: &'buf [u8], + ) -> ParseEvents<'inner, 'buf> { + ParseEvents::new(self, input) + } + + /// Get how many bytes are consumed by the previous call to [`process_bytes`]. + /// + /// Bytes not consumed by the parser should be processed again. + /// + /// [`process_bytes`]: ContainerParser::process_bytes + pub fn previous_consumed_bytes(&self) -> usize { + self.previous_consumed_bytes + } +} + +#[cfg(test)] +impl ContainerParser { + pub(crate) fn collect_codestream(input: &[u8]) -> crate::error::Result<Vec<u8>> { + let mut parser = Self::new(); + let mut codestream = Vec::new(); + for event in parser.process_bytes(input) { + match event? { + ParseEvent::BitstreamKind(_) => {} + ParseEvent::Codestream(buf) => { + codestream.extend_from_slice(buf); + } + } + } + Ok(codestream) + } +} + +#[cfg(test)] +mod test { + use super::*; + use test_log::test; + + #[rustfmt::skip] + const HEADER: &[u8] = &[ + 0x00, 0x00, 0x00, 0x0c, b'J', b'X', b'L', b' ', 0x0d, 0x0a, 0x87, 0x0a, 0x00, 0x00, 0x00, 0x14, + b'f', b't', b'y', b'p', b'j', b'x', b'l', b' ', 0x00, 0x00, 0x00, 0x00, b'j', b'x', b'l', b' ', + ]; + + #[test] + fn parse_partial() { + arbtest::arbtest(|u| { + // Prepare arbitrary container format data with two jxlp boxes. + let total_len = u.arbitrary_len::<u8>()?; + let mut codestream0 = vec![0u8; total_len / 2]; + u.fill_buffer(&mut codestream0)?; + let mut codestream1 = vec![0u8; total_len - codestream0.len()]; + u.fill_buffer(&mut codestream1)?; + + let mut container = HEADER.to_vec(); + container.extend_from_slice(&(12 + codestream0.len() as u32).to_be_bytes()); + container.extend_from_slice(b"jxlp\x00\x00\x00\x00"); + container.extend_from_slice(&codestream0); + + container.extend_from_slice(&(12 + codestream1.len() as u32).to_be_bytes()); + container.extend_from_slice(b"jxlp\x80\x00\x00\x01"); + container.extend_from_slice(&codestream1); + + let mut expected = codestream0; + expected.extend(codestream1); + + // Create a list of arbitrary splits. + let mut tests = Vec::new(); + u.arbitrary_loop(Some(1), Some(10), |u| { + let split_at_idx = u.choose_index(container.len())?; + tests.push(container.split_at(split_at_idx)); + Ok(std::ops::ControlFlow::Continue(())) + })?; + + // Test if split index doesn't affect final codestream. + for (first, second) in tests { + let mut codestream = Vec::new(); + let mut parser = ContainerParser::new(); + + for event in parser.process_bytes(first) { + let event = event.unwrap(); + if let ParseEvent::Codestream(data) = event { + codestream.extend_from_slice(data); + } + } + + let consumed = parser.previous_consumed_bytes(); + let mut second_chunk = first[consumed..].to_vec(); + second_chunk.extend_from_slice(second); + + for event in parser.process_bytes(&second_chunk) { + let event = event.unwrap(); + if let ParseEvent::Codestream(data) = event { + codestream.extend_from_slice(data); + } + } + + assert_eq!(codestream, expected); + } + + Ok(()) + }); + } +} diff --git a/third_party/rust/jxl/src/container/parse.rs b/third_party/rust/jxl/src/container/parse.rs @@ -0,0 +1,273 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. +// +// Originally written for jxl-oxide. + +use super::{BitstreamKind, ContainerParser, DetectState, JxlpIndexState, box_header::*}; +use crate::{ + api::{CODESTREAM_SIGNATURE, CONTAINER_SIGNATURE}, + error::{Error, Result}, + util::tracing_wrappers::*, +}; + +/// Iterator that reads over a buffer and emits parser events. +pub struct ParseEvents<'inner, 'buf> { + inner: &'inner mut ContainerParser, + remaining_input: &'buf [u8], + finished: bool, +} + +impl<'inner, 'buf> ParseEvents<'inner, 'buf> { + pub(super) fn new(parser: &'inner mut ContainerParser, input: &'buf [u8]) -> Self { + parser.previous_consumed_bytes = 0; + Self { + inner: parser, + remaining_input: input, + finished: false, + } + } + + fn emit_single(&mut self) -> Result<Option<ParseEvent<'buf>>> { + let state = &mut self.inner.state; + let jxlp_index_state = &mut self.inner.jxlp_index_state; + let buf = &mut self.remaining_input; + + loop { + if buf.is_empty() { + self.finished = true; + return Ok(None); + } + + match state { + DetectState::WaitingSignature => { + if buf.starts_with(&CODESTREAM_SIGNATURE) { + info!("Codestream signature found"); + *state = DetectState::InCodestream { + kind: BitstreamKind::BareCodestream, + bytes_left: None, + }; + return Ok(Some(ParseEvent::BitstreamKind( + BitstreamKind::BareCodestream, + ))); + } else if buf.starts_with(&CONTAINER_SIGNATURE) { + info!("Container signature found"); + *state = DetectState::WaitingBoxHeader; + *buf = &buf[CONTAINER_SIGNATURE.len()..]; + return Ok(Some(ParseEvent::BitstreamKind(BitstreamKind::Container))); + } else if !CODESTREAM_SIGNATURE.starts_with(buf) + && !CONTAINER_SIGNATURE.starts_with(buf) + { + warn!(?buf, "Invalid signature"); + *state = DetectState::InCodestream { + kind: BitstreamKind::Invalid, + bytes_left: None, + }; + return Ok(Some(ParseEvent::BitstreamKind(BitstreamKind::Invalid))); + } else { + return Ok(None); + } + } + DetectState::WaitingBoxHeader => match ContainerBoxHeader::parse(buf)? { + HeaderParseResult::Done { + header, + header_size, + } => { + *buf = &buf[header_size..]; + let tbox = header.box_type(); + if tbox == ContainerBoxType::CODESTREAM { + match jxlp_index_state { + JxlpIndexState::Initial => { + *jxlp_index_state = JxlpIndexState::SingleJxlc; + } + JxlpIndexState::SingleJxlc => { + warn!("Duplicate jxlc box found"); + return Err(Error::InvalidBox); + } + JxlpIndexState::Jxlp(_) | JxlpIndexState::JxlpFinished => { + warn!("Found jxlc box instead of jxlp box"); + return Err(Error::InvalidBox); + } + } + + *state = DetectState::InCodestream { + kind: BitstreamKind::Container, + bytes_left: header.box_size().map(|x| x as usize), + }; + } else if tbox == ContainerBoxType::PARTIAL_CODESTREAM { + if let Some(box_size) = header.box_size() + && box_size < 4 + { + return Err(Error::InvalidBox); + } + + match jxlp_index_state { + JxlpIndexState::Initial => { + *jxlp_index_state = JxlpIndexState::Jxlp(0); + } + JxlpIndexState::Jxlp(index) => { + *index += 1; + } + JxlpIndexState::SingleJxlc => { + warn!("jxlp box found after jxlc box"); + return Err(Error::InvalidBox); + } + JxlpIndexState::JxlpFinished => { + warn!("found another jxlp box after the final one"); + return Err(Error::InvalidBox); + } + } + + *state = DetectState::WaitingJxlpIndex(header); + } else { + let bytes_left = header.box_size().map(|x| x as usize); + *state = DetectState::InAuxBox { header, bytes_left }; + } + } + HeaderParseResult::NeedMoreData => return Ok(None), + }, + DetectState::WaitingJxlpIndex(header) => { + let &[b0, b1, b2, b3, ..] = &**buf else { + return Ok(None); + }; + + let index = u32::from_be_bytes([b0, b1, b2, b3]); + *buf = &buf[4..]; + let is_last = index & 0x80000000 != 0; + let index = index & 0x7fffffff; + + match *jxlp_index_state { + JxlpIndexState::Jxlp(expected_index) if expected_index == index => { + if is_last { + *jxlp_index_state = JxlpIndexState::JxlpFinished; + } + } + #[allow(unused_variables)] + JxlpIndexState::Jxlp(expected_index) => { + warn!( + expected_index, + actual_index = index, + "Out-of-order jxlp box found", + ); + return Err(Error::InvalidBox); + } + state => { + unreachable!("invalid jxlp index state in WaitingJxlpIndex: {state:?}"); + } + } + + *state = DetectState::InCodestream { + kind: BitstreamKind::Container, + bytes_left: header.box_size().map(|x| x as usize - 4), + }; + } + DetectState::InCodestream { + bytes_left: None, .. + } => { + let payload = *buf; + *buf = &[]; + return Ok(Some(ParseEvent::Codestream(payload))); + } + DetectState::InCodestream { + bytes_left: Some(bytes_left), + .. + } => { + let payload = if buf.len() >= *bytes_left { + let (payload, remaining) = buf.split_at(*bytes_left); + *state = DetectState::WaitingBoxHeader; + *buf = remaining; + payload + } else { + let payload = *buf; + *bytes_left -= buf.len(); + *buf = &[]; + payload + }; + return Ok(Some(ParseEvent::Codestream(payload))); + } + DetectState::InAuxBox { + header: _, + bytes_left: None, + } => { + let _payload = *buf; + *buf = &[]; + // FIXME: emit auxiliary box event + } + DetectState::InAuxBox { + header: _, + bytes_left: Some(bytes_left), + } => { + let _payload = if buf.len() >= *bytes_left { + let (payload, remaining) = buf.split_at(*bytes_left); + *state = DetectState::WaitingBoxHeader; + *buf = remaining; + payload + } else { + let payload = *buf; + *bytes_left -= buf.len(); + *buf = &[]; + payload + }; + // FIXME: emit auxiliary box event + } + } + } + } +} + +impl std::fmt::Debug for ParseEvents<'_, '_> { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("ParseEvents") + .field("inner", &self.inner) + .field( + "remaining_input", + &format_args!("({} byte(s))", self.remaining_input.len()), + ) + .field("finished", &self.finished) + .finish() + } +} + +impl<'buf> Iterator for ParseEvents<'_, 'buf> { + type Item = Result<ParseEvent<'buf>>; + + fn next(&mut self) -> Option<Self::Item> { + if self.finished { + return None; + } + + let initial_buf = self.remaining_input; + let event = self.emit_single(); + + if event.is_err() { + self.finished = true; + } + + self.inner.previous_consumed_bytes += initial_buf.len() - self.remaining_input.len(); + event.transpose() + } +} + +/// Parser event emitted by [`ParseEvents`]. +pub enum ParseEvent<'buf> { + /// Bitstream structure is detected. + BitstreamKind(BitstreamKind), + /// Codestream data is read. + /// + /// Returned data may be partial. Complete codestream can be obtained by concatenating all data + /// of `Codestream` events. + Codestream(&'buf [u8]), +} + +impl std::fmt::Debug for ParseEvent<'_> { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::BitstreamKind(kind) => f.debug_tuple("BitstreamKind").field(kind).finish(), + Self::Codestream(buf) => f + .debug_tuple("Codestream") + .field(&format_args!("{} byte(s)", buf.len())) + .finish(), + } + } +} diff --git a/third_party/rust/jxl/src/entropy_coding/ans.rs b/third_party/rust/jxl/src/entropy_coding/ans.rs @@ -0,0 +1,489 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. +// +// Originally written for jxl-oxide. + +use crate::bit_reader::BitReader; +use crate::error::{Error, Result}; + +const LOG_SUM_PROBS: usize = 12; +const SUM_PROBS: u16 = 1 << LOG_SUM_PROBS; + +const RLE_MARKER_SYM: u16 = LOG_SUM_PROBS as u16 + 1; + +#[derive(Debug)] +struct AnsHistogram { + buckets: Vec<Bucket>, + log_bucket_size: usize, + bucket_mask: u32, + // For optimizing fast-lossless case. + single_symbol: Option<u32>, +} + +// log_alphabet_size <= 8 and log_bucket_size <= 7, so u8 is sufficient for symbols and cutoffs. +#[derive(Debug, Copy, Clone)] +#[repr(C)] +struct Bucket { + alias_symbol: u8, + alias_cutoff: u8, + dist: u16, + alias_offset: u16, + alias_dist_xor: u16, +} + +impl AnsHistogram { + fn decode_dist_two_symbols(br: &mut BitReader, dist: &mut [u16]) -> Result<usize> { + let table_size = dist.len(); + + let v0 = Self::read_u8(br)? as usize; + let v1 = Self::read_u8(br)? as usize; + if v0 == v1 { + return Err(Error::InvalidAnsHistogram); + } + + let alphabet_size = v0.max(v1) + 1; + if alphabet_size > table_size { + return Err(Error::InvalidAnsHistogram); + } + + let prob = br.read(LOG_SUM_PROBS)? as u16; + dist[v0] = prob; + dist[v1] = SUM_PROBS - prob; + + Ok(alphabet_size) + } + + fn decode_dist_single_symbol(br: &mut BitReader, dist: &mut [u16]) -> Result<usize> { + let table_size = dist.len(); + + let val = Self::read_u8(br)? as usize; + let alphabet_size = val + 1; + if alphabet_size > table_size { + return Err(Error::InvalidAnsHistogram); + } + + dist[val] = SUM_PROBS; + + Ok(alphabet_size) + } + + fn decode_dist_evenly_distributed(br: &mut BitReader, dist: &mut [u16]) -> Result<usize> { + let table_size = dist.len(); + + let alphabet_size = Self::read_u8(br)? as usize + 1; + if alphabet_size > table_size { + return Err(Error::InvalidAnsHistogram); + } + + let base = SUM_PROBS as usize / alphabet_size; + let remainder = SUM_PROBS as usize % alphabet_size; + dist[0..remainder].fill(base as u16 + 1); + dist[remainder..alphabet_size].fill(base as u16); + + Ok(alphabet_size) + } + + fn decode_dist_complex(br: &mut BitReader, dist: &mut [u16]) -> Result<usize> { + let table_size = dist.len(); + + let mut len = 0usize; + while len < 3 { + if br.read(1)? != 0 { + len += 1; + } else { + break; + } + } + + let shift = (br.read(len)? + (1 << len) - 1) as i16; + if shift > 13 { + return Err(Error::InvalidAnsHistogram); + } + + let alphabet_size = Self::read_u8(br)? as usize + 3; + if alphabet_size > table_size { + return Err(Error::InvalidAnsHistogram); + } + + // TODO(tirr-c): This could be an array of length `SUM_PROB / 4` (4 is from the minimum + // value of `repeat_count`). Change if using array is faster. + let mut repeat_ranges = Vec::new(); + let mut omit_data = None; + let mut idx = 0; + while idx < alphabet_size { + dist[idx] = Self::read_prefix(br)?; + if dist[idx] == RLE_MARKER_SYM { + let repeat_count = Self::read_u8(br)? as usize + 4; + if idx + repeat_count > alphabet_size { + return Err(Error::InvalidAnsHistogram); + } + repeat_ranges.push(idx..(idx + repeat_count)); + idx += repeat_count; + continue; + } + match &mut omit_data { + Some((log, pos)) => { + if dist[idx] > *log { + *log = dist[idx]; + *pos = idx; + } + } + data => { + *data = Some((dist[idx], idx)); + } + } + idx += 1; + } + let Some((_, omit_pos)) = omit_data else { + return Err(Error::InvalidAnsHistogram); + }; + if dist.get(omit_pos + 1) == Some(&RLE_MARKER_SYM) { + return Err(Error::InvalidAnsHistogram); + } + + let mut repeat_range_idx = 0usize; + let mut acc = 0; + let mut prev_dist = 0u16; + for (idx, code) in dist.iter_mut().enumerate() { + if repeat_range_idx < repeat_ranges.len() + && repeat_ranges[repeat_range_idx].start <= idx + { + if repeat_ranges[repeat_range_idx].end == idx { + repeat_range_idx += 1; + } else { + *code = prev_dist; + acc += *code; + // dist[omit_pos] > 0 + if acc >= SUM_PROBS { + return Err(Error::InvalidAnsHistogram); + } + continue; + } + } + + if *code == 0 { + prev_dist = 0; + continue; + } + if idx == omit_pos { + prev_dist = 0; + continue; + } + if *code > 1 { + let zeros = (*code - 1) as i16; + let bitcount = (shift - ((LOG_SUM_PROBS as i16 - zeros) >> 1)).clamp(0, zeros); + *code = (1 << zeros) + ((br.read(bitcount as usize)? as u16) << (zeros - bitcount)); + } + + prev_dist = *code; + acc += *code; + // dist[omit_pos] > 0 + if acc >= SUM_PROBS { + return Err(Error::InvalidAnsHistogram); + } + } + dist[omit_pos] = SUM_PROBS - acc; + + Ok(alphabet_size) + } + + fn build_alias_map(alphabet_size: usize, log_bucket_size: usize, dist: &[u16]) -> Vec<Bucket> { + #[derive(Debug)] + struct WorkingBucket { + dist: u16, + alias_symbol: u16, + alias_offset: u16, + alias_cutoff: u16, + } + + let bucket_size = 1u16 << log_bucket_size; + let mut buckets: Vec<_> = dist + .iter() + .enumerate() + .map(|(i, &dist)| WorkingBucket { + dist, + alias_symbol: if i < alphabet_size { i as u16 } else { 0 }, + alias_offset: 0, + alias_cutoff: dist, + }) + .collect(); + + let mut underfull = Vec::new(); + let mut overfull = Vec::new(); + for (idx, &WorkingBucket { dist, .. }) in buckets.iter().enumerate() { + match dist.cmp(&bucket_size) { + std::cmp::Ordering::Less => underfull.push(idx), + std::cmp::Ordering::Equal => {} + std::cmp::Ordering::Greater => overfull.push(idx), + } + } + while let (Some(o), Some(u)) = (overfull.pop(), underfull.pop()) { + let by = bucket_size - buckets[u].alias_cutoff; + buckets[o].alias_cutoff -= by; + buckets[u].alias_symbol = o as u16; + buckets[u].alias_offset = buckets[o].alias_cutoff; + match buckets[o].alias_cutoff.cmp(&bucket_size) { + std::cmp::Ordering::Less => underfull.push(o), + std::cmp::Ordering::Equal => {} + std::cmp::Ordering::Greater => overfull.push(o), + } + } + + // Assertion failure happens only if `dist` doesn't sum to `SUM_PROB`, which is checked + // before building alias map. + assert!(overfull.is_empty() && underfull.is_empty()); + + buckets + .iter() + .enumerate() + .map(|(idx, bucket)| { + if bucket.alias_cutoff == bucket_size { + Bucket { + dist: bucket.dist, + alias_symbol: idx as u8, + alias_offset: 0, + alias_cutoff: 0, + alias_dist_xor: 0, + } + } else { + Bucket { + dist: bucket.dist, + alias_symbol: bucket.alias_symbol as u8, + alias_offset: bucket.alias_offset - bucket.alias_cutoff, + alias_cutoff: bucket.alias_cutoff as u8, + alias_dist_xor: bucket.dist ^ buckets[bucket.alias_symbol as usize].dist, + } + } + }) + .collect() + } + + // log_alphabet_size: 5 + u(2) + pub fn decode(br: &mut BitReader, log_alpha_size: usize) -> Result<Self> { + debug_assert!((5..=8).contains(&log_alpha_size)); + let table_size = (1u16 << log_alpha_size) as usize; + // 4 <= log_bucket_size <= 7 + let log_bucket_size = LOG_SUM_PROBS - log_alpha_size; + let bucket_size = 1u16 << log_bucket_size; + let bucket_mask = bucket_size as u32 - 1; + + let mut dist = vec![0u16; table_size]; + let alphabet_size = if br.read(1)? != 0 { + if br.read(1)? != 0 { + Self::decode_dist_two_symbols(br, &mut dist)? + } else { + Self::decode_dist_single_symbol(br, &mut dist)? + } + } else if br.read(1)? != 0 { + Self::decode_dist_evenly_distributed(br, &mut dist)? + } else { + Self::decode_dist_complex(br, &mut dist)? + }; + + if let Some(single_sym_idx) = dist.iter().position(|&d| d == SUM_PROBS) { + let buckets = dist + .into_iter() + .enumerate() + .map(|(i, dist)| Bucket { + dist, + alias_symbol: single_sym_idx as u8, + alias_offset: bucket_size * i as u16, + alias_cutoff: 0, + alias_dist_xor: dist ^ SUM_PROBS, + }) + .collect(); + return Ok(Self { + buckets, + log_bucket_size, + bucket_mask, + single_symbol: Some(single_sym_idx as u32), + }); + } + + Ok(Self { + buckets: Self::build_alias_map(alphabet_size, log_bucket_size, &dist), + log_bucket_size, + bucket_mask, + single_symbol: None, + }) + } + + fn read_u8(bitstream: &mut BitReader) -> Result<u8> { + Ok(if bitstream.read(1)? != 0 { + let n = bitstream.read(3)?; + ((1 << n) + bitstream.read(n as usize)?) as u8 + } else { + 0 + }) + } + + fn read_prefix(br: &mut BitReader) -> Result<u16> { + // Prefix code lookup table. + #[rustfmt::skip] + const TABLE: [(u8, u8); 128] = [ + (10, 3), (12, 7), (7, 3), (3, 4), (6, 3), (8, 3), (9, 3), (5, 4), + (10, 3), ( 4, 4), (7, 3), (1, 4), (6, 3), (8, 3), (9, 3), (2, 4), + (10, 3), ( 0, 5), (7, 3), (3, 4), (6, 3), (8, 3), (9, 3), (5, 4), + (10, 3), ( 4, 4), (7, 3), (1, 4), (6, 3), (8, 3), (9, 3), (2, 4), + (10, 3), (11, 6), (7, 3), (3, 4), (6, 3), (8, 3), (9, 3), (5, 4), + (10, 3), ( 4, 4), (7, 3), (1, 4), (6, 3), (8, 3), (9, 3), (2, 4), + (10, 3), ( 0, 5), (7, 3), (3, 4), (6, 3), (8, 3), (9, 3), (5, 4), + (10, 3), ( 4, 4), (7, 3), (1, 4), (6, 3), (8, 3), (9, 3), (2, 4), + (10, 3), (13, 7), (7, 3), (3, 4), (6, 3), (8, 3), (9, 3), (5, 4), + (10, 3), ( 4, 4), (7, 3), (1, 4), (6, 3), (8, 3), (9, 3), (2, 4), + (10, 3), ( 0, 5), (7, 3), (3, 4), (6, 3), (8, 3), (9, 3), (5, 4), + (10, 3), ( 4, 4), (7, 3), (1, 4), (6, 3), (8, 3), (9, 3), (2, 4), + (10, 3), (11, 6), (7, 3), (3, 4), (6, 3), (8, 3), (9, 3), (5, 4), + (10, 3), ( 4, 4), (7, 3), (1, 4), (6, 3), (8, 3), (9, 3), (2, 4), + (10, 3), ( 0, 5), (7, 3), (3, 4), (6, 3), (8, 3), (9, 3), (5, 4), + (10, 3), ( 4, 4), (7, 3), (1, 4), (6, 3), (8, 3), (9, 3), (2, 4), + ]; + + let index = br.peek(7); + let (sym, bits) = TABLE[index as usize]; + br.consume(bits as usize)?; + Ok(sym as u16) + } +} + +impl AnsHistogram { + #[inline(always)] + pub fn read(&self, br: &mut BitReader, state: &mut u32) -> Result<u32> { + let idx = *state & 0xfff; + let i = (idx >> self.log_bucket_size) as usize; + let pos = idx & self.bucket_mask; + + let bucket = self.buckets[i]; + let alias_symbol = bucket.alias_symbol as usize; + let alias_cutoff = bucket.alias_cutoff as u32; + let dist = bucket.dist as u32; + + let map_to_alias = pos >= alias_cutoff; + let (offset, dist_xor) = if map_to_alias { + (bucket.alias_offset as u32, bucket.alias_dist_xor as u32) + } else { + (0, 0) + }; + + let dist = dist ^ dist_xor; + let symbol = if map_to_alias { alias_symbol } else { i }; + let offset = offset + pos; + + let next_state = (*state >> LOG_SUM_PROBS) * dist + offset; + let appended_state = (next_state << 16) | br.peek(16) as u32; + let select_appended = next_state < (1 << 16); + *state = if select_appended { + appended_state + } else { + next_state + }; + br.consume(if select_appended { 16 } else { 0 })?; + Ok(symbol as u32) + } + + // For optimizing fast-lossless case. + #[allow(unused)] + #[inline] + pub fn single_symbol(&self) -> Option<u32> { + self.single_symbol + } +} + +#[derive(Debug)] +pub struct AnsCodes { + histograms: Vec<AnsHistogram>, +} + +impl AnsCodes { + pub fn decode(num: usize, log_alpha_size: usize, br: &mut BitReader) -> Result<AnsCodes> { + let histograms = (0..num) + .map(|_| AnsHistogram::decode(br, log_alpha_size)) + .collect::<Result<_>>()?; + Ok(Self { histograms }) + } +} + +#[derive(Debug)] +pub struct AnsReader(u32); + +impl AnsReader { + /// Expected final ANS state. + const CHECKSUM: u32 = 0x130000; + + pub fn new_unused() -> Self { + Self(Self::CHECKSUM) + } + + pub fn init(br: &mut BitReader) -> Result<Self> { + let initial_state = br.read(32)? as u32; + Ok(Self(initial_state)) + } + + pub fn read(&mut self, codes: &AnsCodes, br: &mut BitReader, ctx: usize) -> Result<u32> { + codes.histograms[ctx].read(br, &mut self.0) + } + + pub fn check_final_state(self) -> Result<()> { + if self.0 == Self::CHECKSUM { + Ok(()) + } else { + Err(Error::AnsChecksumMismatch) + } + } +} + +#[cfg(test)] +mod test { + use super::*; + + fn validate_buckets(buckets: &[Bucket]) { + let dist_sum: u16 = buckets.iter().map(|bucket| bucket.dist).sum(); + assert_eq!(dist_sum, SUM_PROBS); + } + + #[test] + fn single_symbol() { + // Single symbol of 20 + let mut br = BitReader::new(&[0b00100101, 0b01]); + let histogram = AnsHistogram::decode(&mut br, 5).unwrap(); + validate_buckets(&histogram.buckets); + assert_eq!(histogram.buckets[20].dist, SUM_PROBS); + assert_eq!(histogram.single_symbol, Some(20)); + + // Single symbol of 32 (invalid) + let mut br = BitReader::new(&[0b00101101, 0b000]); + assert!(AnsHistogram::decode(&mut br, 5).is_err()); + } + + #[test] + fn two_symbols() { + // two symbols of 10 and 20, where the prob of symbol 10 is 256 + let mut br = BitReader::new(&[0b10011111, 0b10010010, 0b00000000, 0b00010]); + let histogram = AnsHistogram::decode(&mut br, 5).unwrap(); + validate_buckets(&histogram.buckets); + assert_eq!(histogram.buckets[10].dist, 256); + assert_eq!(histogram.buckets[20].dist, SUM_PROBS - 256); + } + + #[test] + fn decode_arb() { + arbtest::arbtest(|u| { + let total_len = u.arbitrary_len::<u8>()?; + let mut buf = vec![0u8; total_len]; + u.fill_buffer(&mut buf)?; + let mut br = BitReader::new(&buf); + + loop { + match AnsHistogram::decode(&mut br, 8) { + Ok(histogram) => validate_buckets(&histogram.buckets), + Err(Error::OutOfBounds(_)) => break, + Err(_) => {} + } + } + + Ok(()) + }); + } +} diff --git a/third_party/rust/jxl/src/entropy_coding/context_map.rs b/third_party/rust/jxl/src/entropy_coding/context_map.rs @@ -0,0 +1,76 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +use crate::bit_reader::BitReader; +use crate::error::Error; +use std::collections::HashSet; + +use crate::entropy_coding::decode::*; + +fn move_to_front(v: &mut [u8], index: u8) { + let value = v[index as usize]; + for i in (1..=index as usize).rev() { + v[i] = v[i - 1]; + } + v[0] = value; +} + +fn inverse_move_to_front(v: &mut [u8]) { + let mut mtf: [u8; 256] = std::array::from_fn(|x| x as u8); + for val in v.iter_mut() { + let index = *val; + *val = mtf[index as usize]; + if index != 0 { + move_to_front(&mut mtf, index); + } + } +} + +fn verify_context_map(ctx_map: &[u8]) -> Result<(), Error> { + let num_histograms = *ctx_map.iter().max().unwrap() as u32 + 1; + let distinct_histograms = ctx_map.iter().collect::<HashSet<_>>().len() as u32; + if distinct_histograms != num_histograms { + return Err(Error::InvalidContextMapHole( + num_histograms, + distinct_histograms, + )); + } + Ok(()) +} + +pub fn decode_context_map(num_contexts: usize, br: &mut BitReader) -> Result<Vec<u8>, Error> { + let is_simple = br.read(1)? != 0; + if is_simple { + let bits_per_entry = br.read(2)? as usize; + if bits_per_entry != 0 { + (0..num_contexts) + .map(|_| Ok(br.read(bits_per_entry)? as u8)) + .collect() + } else { + Ok(vec![0u8; num_contexts]) + } + } else { + let use_mtf = br.read(1)? != 0; + let histograms = Histograms::decode(1, br, /*allow_lz77=*/ num_contexts > 2)?; + let mut reader = SymbolReader::new(&histograms, br, None)?; + + let mut ctx_map: Vec<u8> = (0..num_contexts) + .map(|_| { + let mv = reader.read_unsigned(&histograms, br, 0usize)?; + if mv > u8::MAX as u32 { + Err(Error::InvalidContextMap(mv)) + } else { + Ok(mv as u8) + } + }) + .collect::<Result<_, _>>()?; + reader.check_final_state(&histograms)?; + if use_mtf { + inverse_move_to_front(&mut ctx_map[..]); + } + verify_context_map(&ctx_map[..])?; + Ok(ctx_map) + } +} diff --git a/third_party/rust/jxl/src/entropy_coding/decode.rs b/third_party/rust/jxl/src/entropy_coding/decode.rs @@ -0,0 +1,327 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +use jxl_macros::UnconditionalCoder; + +use crate::bit_reader::BitReader; +use crate::entropy_coding::ans::*; +use crate::entropy_coding::context_map::*; +use crate::entropy_coding::huffman::*; +use crate::entropy_coding::hybrid_uint::*; +use crate::error::{Error, Result}; +use crate::headers::encodings::*; +use crate::util::tracing_wrappers::*; + +pub fn decode_varint16(br: &mut BitReader) -> Result<u16> { + if br.read(1)? != 0 { + let nbits = br.read(4)? as usize; + if nbits == 0 { + Ok(1) + } else { + Ok((1 << nbits) + br.read(nbits)? as u16) + } + } else { + Ok(0) + } +} + +pub fn unpack_signed(unsigned: u32) -> i32 { + ((unsigned >> 1) ^ ((!unsigned) & 1).wrapping_sub(1)) as i32 +} + +#[derive(UnconditionalCoder, Debug)] +struct Lz77Params { + pub enabled: bool, + #[condition(enabled)] + #[coder(u2S(224, 512, 4096, Bits(15) + 8))] + pub min_symbol: Option<u32>, + #[condition(enabled)] + #[coder(u2S(3, 4, Bits(2) + 5, Bits(8) + 9))] + pub min_length: Option<u32>, +} + +#[derive(Debug)] +enum Codes { + Huffman(HuffmanCodes), + Ans(AnsCodes), +} + +#[derive(Debug)] +pub struct Histograms { + lz77_params: Lz77Params, + lz77_length_uint: Option<HybridUint>, + context_map: Vec<u8>, + // TODO(veluca): figure out why this is unused. + #[allow(dead_code)] + log_alpha_size: usize, + uint_configs: Vec<HybridUint>, + codes: Codes, +} + +#[derive(Debug)] +pub struct Lz77State { + min_symbol: u32, + min_length: u32, + dist_multiplier: u32, + window: Vec<u32>, + num_to_copy: u32, + copy_pos: u32, + num_decoded: u32, +} + +impl Lz77State { + const LOG_WINDOW_SIZE: u32 = 20; + const WINDOW_MASK: u32 = (1 << Self::LOG_WINDOW_SIZE) - 1; + + #[rustfmt::skip] + const SPECIAL_DISTANCES: [(i8, u8); 120] = [ + ( 0, 1), ( 1, 0), ( 1, 1), (-1, 1), ( 0, 2), ( 2, 0), ( 1, 2), (-1, 2), ( 2, 1), (-2, 1), + ( 2, 2), (-2, 2), ( 0, 3), ( 3, 0), ( 1, 3), (-1, 3), ( 3, 1), (-3, 1), ( 2, 3), (-2, 3), + ( 3, 2), (-3, 2), ( 0, 4), ( 4, 0), ( 1, 4), (-1, 4), ( 4, 1), (-4, 1), ( 3, 3), (-3, 3), + ( 2, 4), (-2, 4), ( 4, 2), (-4, 2), ( 0, 5), ( 3, 4), (-3, 4), ( 4, 3), (-4, 3), ( 5, 0), + ( 1, 5), (-1, 5), ( 5, 1), (-5, 1), ( 2, 5), (-2, 5), ( 5, 2), (-5, 2), ( 4, 4), (-4, 4), + ( 3, 5), (-3, 5), ( 5, 3), (-5, 3), ( 0, 6), ( 6, 0), ( 1, 6), (-1, 6), ( 6, 1), (-6, 1), + ( 2, 6), (-2, 6), ( 6, 2), (-6, 2), ( 4, 5), (-4, 5), ( 5, 4), (-5, 4), ( 3, 6), (-3, 6), + ( 6, 3), (-6, 3), ( 0, 7), ( 7, 0), ( 1, 7), (-1, 7), ( 5, 5), (-5, 5), ( 7, 1), (-7, 1), + ( 4, 6), (-4, 6), ( 6, 4), (-6, 4), ( 2, 7), (-2, 7), ( 7, 2), (-7, 2), ( 3, 7), (-3, 7), + ( 7, 3), (-7, 3), ( 5, 6), (-5, 6), ( 6, 5), (-6, 5), ( 8, 0), ( 4, 7), (-4, 7), ( 7, 4), + (-7, 4), ( 8, 1), ( 8, 2), ( 6, 6), (-6, 6), ( 8, 3), ( 5, 7), (-5, 7), ( 7, 5), (-7, 5), + ( 8, 4), ( 6, 7), (-6, 7), ( 7, 6), (-7, 6), ( 8, 5), ( 7, 7), (-7, 7), ( 8, 6), ( 8, 7), + ]; + + fn push_decoded_symbol(&mut self, token: u32) { + let offset = (self.num_decoded & Self::WINDOW_MASK) as usize; + if let Some(slot) = self.window.get_mut(offset) { + *slot = token; + } else { + debug_assert_eq!(self.window.len(), offset); + self.window.push(token); + } + self.num_decoded += 1; + } + + fn pull_symbol(&mut self) -> Option<u32> { + if let Some(next_num_to_copy) = self.num_to_copy.checked_sub(1) { + let sym = self.window[(self.copy_pos & Self::WINDOW_MASK) as usize]; + self.copy_pos += 1; + self.num_to_copy = next_num_to_copy; + Some(sym) + } else { + None + } + } +} + +#[derive(Debug)] +pub struct SymbolReader { + pub lz77_state: Option<Lz77State>, + pub ans_reader: AnsReader, +} + +impl SymbolReader { + pub fn new( + histograms: &Histograms, + br: &mut BitReader, + image_width: Option<usize>, + ) -> Result<SymbolReader> { + let ans_reader = if matches!(histograms.codes, Codes::Ans(_)) { + AnsReader::init(br)? + } else { + AnsReader::new_unused() + }; + + let lz77_state = if histograms.lz77_params.enabled { + Some(Lz77State { + min_symbol: histograms.lz77_params.min_symbol.unwrap(), + min_length: histograms.lz77_params.min_length.unwrap(), + dist_multiplier: image_width.unwrap_or(0) as u32, + window: Vec::new(), + num_to_copy: 0, + copy_pos: 0, + num_decoded: 0, + }) + } else { + None + }; + + Ok(SymbolReader { + lz77_state, + ans_reader, + }) + } + + pub fn read_unsigned( + &mut self, + histograms: &Histograms, + br: &mut BitReader, + context: usize, + ) -> Result<u32> { + let cluster = histograms.map_context_to_cluster(context); + if histograms.lz77_params.enabled { + let lz77_state = self.lz77_state.as_mut().unwrap(); + if let Some(sym) = lz77_state.pull_symbol() { + lz77_state.push_decoded_symbol(sym); + return Ok(sym); + } + let token = match &histograms.codes { + Codes::Huffman(hc) => hc.read(br, cluster)?, + Codes::Ans(ans) => self.ans_reader.read(ans, br, cluster)?, + }; + let Some(lz77_token) = token.checked_sub(lz77_state.min_symbol) else { + let sym = histograms.uint_configs[cluster].read(token, br)?; + lz77_state.push_decoded_symbol(sym); + return Ok(sym); + }; + if lz77_state.num_decoded == 0 { + return Err(Error::UnexpectedLz77Repeat); + } + + let num_to_copy = histograms + .lz77_length_uint + .as_ref() + .unwrap() + .read(lz77_token, br)?; + let Some(num_to_copy) = num_to_copy.checked_add(lz77_state.min_length) else { + warn!( + num_to_copy, + lz77_state.min_length, "LZ77 num_to_copy overflow" + ); + return Err(Error::ArithmeticOverflow); + }; + let lz_dist_cluster = *histograms.context_map.last().unwrap() as usize; + + let distance_sym = match &histograms.codes { + Codes::Huffman(hc) => hc.read(br, lz_dist_cluster)?, + Codes::Ans(ans) => self.ans_reader.read(ans, br, lz_dist_cluster)?, + }; + let distance_sym = histograms.uint_configs[lz_dist_cluster].read(distance_sym, br)?; + + let distance_sub_1 = if lz77_state.dist_multiplier == 0 { + distance_sym + } else if let Some(distance) = distance_sym.checked_sub(120) { + distance + } else { + let (offset, dist) = Lz77State::SPECIAL_DISTANCES[distance_sym as usize]; + let dist = (lz77_state.dist_multiplier * dist as u32) + .checked_add_signed(offset as i32 - 1); + dist.unwrap_or(0) + }; + + let distance = (((1 << 20) - 1).min(distance_sub_1) + 1).min(lz77_state.num_decoded); + lz77_state.copy_pos = lz77_state.num_decoded - distance; + + lz77_state.num_to_copy = num_to_copy; + let sym = lz77_state.pull_symbol().unwrap(); + lz77_state.push_decoded_symbol(sym); + Ok(sym) + } else { + let token = match &histograms.codes { + Codes::Huffman(hc) => hc.read(br, cluster)?, + Codes::Ans(ans) => self.ans_reader.read(ans, br, cluster)?, + }; + histograms.uint_configs[cluster].read(token, br) + } + } + + pub fn read_signed( + &mut self, + histograms: &Histograms, + br: &mut BitReader, + context: usize, + ) -> Result<i32> { + let unsigned = self.read_unsigned(histograms, br, context)?; + Ok(unpack_signed(unsigned)) + } + + pub fn check_final_state(self, histograms: &Histograms) -> Result<()> { + match &histograms.codes { + Codes::Huffman(_) => Ok(()), + Codes::Ans(_) => self.ans_reader.check_final_state(), + } + } +} + +impl Histograms { + pub fn decode(num_contexts: usize, br: &mut BitReader, allow_lz77: bool) -> Result<Histograms> { + let lz77_params = Lz77Params::read_unconditional(&(), br, &Empty {})?; + if !allow_lz77 && lz77_params.enabled { + return Err(Error::Lz77Disallowed); + } + let (num_contexts, lz77_length_uint) = if lz77_params.enabled { + ( + num_contexts + 1, + Some(HybridUint::decode(/*log_alpha_size=*/ 8, br)?), + ) + } else { + (num_contexts, None) + }; + + let context_map = if num_contexts > 1 { + decode_context_map(num_contexts, br)? + } else { + vec![0] + }; + assert_eq!(context_map.len(), num_contexts); + + let use_prefix_code = br.read(1)? != 0; + let log_alpha_size = if use_prefix_code { + HUFFMAN_MAX_BITS + } else { + br.read(2)? as usize + 5 + }; + let num_histograms = *context_map.iter().max().unwrap() + 1; + let uint_configs = ((0..num_histograms).map(|_| HybridUint::decode(log_alpha_size, br))) + .collect::<Result<_>>()?; + + let codes = if use_prefix_code { + Codes::Huffman(HuffmanCodes::decode(num_histograms as usize, br)?) + } else { + Codes::Ans(AnsCodes::decode( + num_histograms as usize, + log_alpha_size, + br, + )?) + }; + + Ok(Histograms { + lz77_params, + lz77_length_uint, + context_map, + log_alpha_size, + uint_configs, + codes, + }) + } + + pub fn map_context_to_cluster(&self, context: usize) -> usize { + self.context_map[context] as usize + } + + pub fn num_histograms(&self) -> usize { + *self.context_map.iter().max().unwrap() as usize + 1 + } +} + +#[cfg(test)] +impl Histograms { + /// Builds a decoder that reads an octet at a time and emits its bit-reversed value. + pub fn reverse_octet(num_contexts: usize) -> Self { + let d = HuffmanCodes::byte_histogram(); + let codes = Codes::Huffman(d); + let uint_configs = vec![HybridUint::new(8, 0, 0)]; + Self { + lz77_params: Lz77Params { + enabled: false, + min_symbol: None, + min_length: None, + }, + lz77_length_uint: None, + uint_configs, + log_alpha_size: 15, + context_map: vec![0u8; num_contexts], + codes, + } + } +} diff --git a/third_party/rust/jxl/src/entropy_coding/huffman.rs b/third_party/rust/jxl/src/entropy_coding/huffman.rs @@ -0,0 +1,531 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +use std::fmt::Debug; + +use crate::bit_reader::BitReader; +use crate::entropy_coding::decode::*; +use crate::error::{Error, Result}; +use crate::util::{CeilLog2, NewWithCapacity, tracing_wrappers::*}; + +pub const HUFFMAN_MAX_BITS: usize = 15; +const TABLE_BITS: usize = 8; +const TABLE_SIZE: usize = 1 << TABLE_BITS; +const CODE_LENGTHS_CODE: usize = 18; +const DEFAULT_CODE_LENGTH: u8 = 8; +const CODE_LENGTH_REPEAT_CODE: u8 = 16; + +#[derive(Clone, Copy)] +struct TableEntry { + bits: u8, + value: u16, +} + +impl Debug for TableEntry { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{}l{}", self.value, self.bits) + } +} + +#[derive(Debug)] +struct Table { + entries: Vec<TableEntry>, +} + +/* Returns reverse(reverse(key, len) + 1, len), where reverse(key, len) is the +bit-wise reversal of the len least significant bits of key. */ +fn get_next_key(key: u32, len: usize) -> u32 { + let mut step = 1 << (len - 1); + while key & step != 0 { + step >>= 1; + } + (key & (step.wrapping_sub(1))) + step +} + +/* Stores code in table[0], table[step], table[2*step], ..., table[end] */ +/* Assumes that end is an integer multiple of step */ +fn replicate_value(table: &mut [TableEntry], step: usize, value: TableEntry) { + for v in table.iter_mut().step_by(step) { + *v = value; + } +} + +/* Returns the table width of the next 2nd level table. count is the histogram +of bit lengths for the remaining symbols, len is the code length of the next +processed symbol */ +fn next_table_bit_size(count: &[u16], len: usize, root_bits: usize) -> usize { + let mut len = len; + let mut left = 1 << (len - root_bits); + while len < HUFFMAN_MAX_BITS { + if left <= count[len] { + break; + } + left -= count[len]; + len += 1; + left <<= 1; + } + len - root_bits +} + +impl Table { + fn decode_simple_table(al_size: usize, br: &mut BitReader) -> Result<Vec<TableEntry>> { + let max_bits = al_size.ceil_log2(); + let num_symbols = (br.read(2)? + 1) as usize; + let mut symbols = [0u16; 4]; + for symbol in symbols.iter_mut().take(num_symbols) { + let sym = br.read(max_bits)? as usize; + if sym >= al_size { + return Err(Error::InvalidHuffman); + } + *symbol = sym as u16; + } + if (0..num_symbols - 1).any(|i| symbols[..i].contains(&symbols[i + 1])) { + return Err(Error::InvalidHuffman); + } + + let special_4_symbols = if num_symbols == 4 { + br.read(1)? != 0 + } else { + false + }; + debug!(symbols = ?symbols[..num_symbols]); + match (num_symbols, special_4_symbols) { + (1, _) => Ok(vec![ + TableEntry { + bits: 0, + value: symbols[0] + }; + TABLE_SIZE + ]), + (2, _) => { + let mut ret = Vec::new_with_capacity(TABLE_SIZE)?; + symbols[0..2].sort_unstable(); + for _ in 0..(TABLE_SIZE >> 1) { + ret.push(TableEntry { + bits: 1, + value: symbols[0], + }); + ret.push(TableEntry { + bits: 1, + value: symbols[1], + }); + } + Ok(ret) + } + (3, _) => { + let mut ret = Vec::new_with_capacity(TABLE_SIZE)?; + symbols[1..3].sort_unstable(); + for _ in 0..(TABLE_SIZE >> 2) { + ret.push(TableEntry { + bits: 1, + value: symbols[0], + }); + ret.push(TableEntry { + bits: 2, + value: symbols[1], + }); + ret.push(TableEntry { + bits: 1, + value: symbols[0], + }); + ret.push(TableEntry { + bits: 2, + value: symbols[2], + }); + } + Ok(ret) + } + (4, false) => { + let mut ret = Vec::new_with_capacity(TABLE_SIZE)?; + symbols.sort_unstable(); + for _ in 0..(TABLE_SIZE >> 2) { + ret.push(TableEntry { + bits: 2, + value: symbols[0], + }); + ret.push(TableEntry { + bits: 2, + value: symbols[2], + }); + ret.push(TableEntry { + bits: 2, + value: symbols[1], + }); + ret.push(TableEntry { + bits: 2, + value: symbols[3], + }); + } + Ok(ret) + } + (4, true) => { + let mut ret = Vec::new_with_capacity(TABLE_SIZE)?; + symbols[2..4].sort_unstable(); + for _ in 0..(TABLE_SIZE >> 3) { + ret.push(TableEntry { + bits: 1, + value: symbols[0], + }); + ret.push(TableEntry { + bits: 2, + value: symbols[1], + }); + ret.push(TableEntry { + bits: 1, + value: symbols[0], + }); + ret.push(TableEntry { + bits: 3, + value: symbols[2], + }); + ret.push(TableEntry { + bits: 1, + value: symbols[0], + }); + ret.push(TableEntry { + bits: 2, + value: symbols[1], + }); + ret.push(TableEntry { + bits: 1, + value: symbols[0], + }); + ret.push(TableEntry { + bits: 3, + value: symbols[3], + }); + } + Ok(ret) + } + _ => unreachable!(), + } + } + + fn decode_huffman_code_lengths( + code_length_code_lengths: [u8; CODE_LENGTHS_CODE], + al_size: usize, + br: &mut BitReader, + ) -> Result<Vec<u8>> { + let table = Table::build(5, &code_length_code_lengths)?; + + let mut symbol = 0; + let mut prev_code_len = DEFAULT_CODE_LENGTH; + let mut repeat = 0u16; + let mut repeat_code_len = 0; + let mut space = 1 << 15; + + let mut code_lengths = vec![0u8; al_size]; + + while symbol < al_size && space > 0 { + let idx = br.peek(5) as usize; + br.consume(table[idx].bits as usize)?; + let code_len = table[idx].value as u8; + if code_len < CODE_LENGTH_REPEAT_CODE { + repeat = 0; + code_lengths[symbol] = code_len; + symbol += 1; + if code_len != 0 { + prev_code_len = code_len; + space -= 32768usize >> code_len; + } + } else { + let extra_bits = code_len - 14; + + let new_len = if code_len == CODE_LENGTH_REPEAT_CODE { + prev_code_len + } else { + 0 + }; + if repeat_code_len != new_len { + repeat = 0; + repeat_code_len = new_len; + } + let old_repeat = repeat; + if repeat > 0 { + repeat -= 2; + repeat <<= extra_bits; + } + repeat += br.read(extra_bits as usize)? as u16 + 3; + let repeat_delta = repeat - old_repeat; + if symbol + repeat_delta as usize > al_size { + return Err(Error::InvalidHuffman); + } + for i in 0..repeat_delta { + code_lengths[symbol + i as usize] = repeat_code_len; + } + symbol += repeat_delta as usize; + if repeat_code_len != 0 { + space -= (repeat_delta as usize) << (15 - repeat_code_len); + } + } + } + if space != 0 { + return Err(Error::InvalidHuffman); + } + Ok(code_lengths) + } + + #[instrument(level = "trace", ret, err)] + fn build(root_bits: usize, code_lengths: &[u8]) -> Result<Vec<TableEntry>> { + if code_lengths.len() > 1 << HUFFMAN_MAX_BITS { + return Err(Error::InvalidHuffman); + } + let mut counts = [0u16; HUFFMAN_MAX_BITS + 1]; + for &v in code_lengths.iter() { + counts[v as usize] += 1; + } + + /* symbols sorted by code length */ + let mut sorted = vec![0u16; code_lengths.len()]; + + /* offsets in sorted table for each length */ + let mut offset = [0; HUFFMAN_MAX_BITS + 1]; + let mut max_length = 1; + + /* generate offsets into sorted symbol table by code length */ + { + let mut sum = 0; + for len in 1..=HUFFMAN_MAX_BITS { + offset[len] = sum; + if counts[len] != 0 { + sum += counts[len]; + max_length = len; + } + } + } + + /* sort symbols by length, by symbol order within each length */ + for (symbol, len) in code_lengths.iter().enumerate() { + if *len != 0 { + sorted[offset[*len as usize] as usize] = symbol as u16; + offset[*len as usize] += 1; + } + } + + let mut table_bits = root_bits; + let mut table_size = 1 << table_bits; + let mut table_pos = 0; + let mut table = vec![TableEntry { bits: 0, value: 0 }; table_size]; + + /* special case code with only one value */ + if offset[HUFFMAN_MAX_BITS] == 1 { + for v in table.iter_mut() { + v.bits = 0; + v.value = sorted[0]; + } + return Ok(table); + } + + /* fill in root table */ + /* let's reduce the table size to a smaller size if possible, and */ + /* create the repetitions by memcpy if possible in the coming loop */ + if table_bits > max_length { + table_bits = max_length; + table_size = 1 << table_bits; + } + let mut key = 0u32; + let mut symbol = 0; + let mut bits = 1u8; + let mut step = 2; + loop { + loop { + if counts[bits as usize] == 0 { + break; + } + let value = sorted[symbol]; + symbol += 1; + replicate_value(&mut table[key as usize..], step, TableEntry { bits, value }); + key = get_next_key(key, bits as usize); + counts[bits as usize] -= 1; + } + step <<= 1; + bits += 1; + if bits as usize > table_bits { + break; + } + } + + /* if root_bits != table_bits we only created one fraction of the */ + /* table, and we need to replicate it now. */ + while table.len() != table_size { + for i in 0..table_size { + table[i + table_size] = table[i]; + } + table_size <<= 1; + } + trace!("table of length {}, table_size: {table_size}", table.len()); + + /* fill in 2nd level tables and add pointers to root table */ + let mask = (table.len() - 1) as u32; + let mut low = !0u32; + let mut step = 2; + for len in root_bits + 1..=max_length { + loop { + if counts[len] == 0 { + break; + } + if (key & mask) != low { + table_pos += table_size; + table_bits = next_table_bit_size(&counts, len, root_bits); + table_size = 1 << table_bits; + low = key & mask; + table[low as usize].bits = (table_bits + root_bits) as u8; + table[low as usize].value = (table_pos - low as usize) as u16; + if table.len() < table_pos + table_size { + table.resize(table_pos + table_size, TableEntry { bits: 0, value: 0 }); + } + } + counts[len] -= 1; + let bits = (len - root_bits) as u8; + let value = sorted[symbol]; + symbol += 1; + let pos = table_pos + (key as usize >> root_bits); + trace!( + "filling 2nd level table of len {len} starting at position {pos} ({table_pos} + {}) of {}", + key as usize >> root_bits, + table.len() + ); + replicate_value(&mut table[pos..], step, TableEntry { bits, value }); + key = get_next_key(key, len); + } + step <<= 1; + } + Ok(table) + } + + #[instrument(level = "trace", skip(br), ret, err)] + pub fn decode(al_size: usize, br: &mut BitReader) -> Result<Table> { + let entries = if al_size == 1 { + vec![TableEntry { bits: 0, value: 0 }; TABLE_SIZE] + } else { + assert!(al_size < 1 << HUFFMAN_MAX_BITS); + let simple_code_or_skip = br.read(2)? as usize; + if simple_code_or_skip == 1 { + Table::decode_simple_table(al_size, br)? + } else { + let mut code_length_code_lengths = [0u8; CODE_LENGTHS_CODE]; + let mut space = 32; + const STATIC_HUFF_BITS: [u8; 16] = [2, 2, 2, 3, 2, 2, 2, 4, 2, 2, 2, 3, 2, 2, 2, 4]; + const STATIC_HUFF_VALS: [u8; 16] = [0, 4, 3, 2, 0, 4, 3, 1, 0, 4, 3, 2, 0, 4, 3, 5]; + const CODE_LENGTH_CODE_ORDER: [u8; CODE_LENGTHS_CODE] = + [1, 2, 3, 4, 0, 5, 17, 6, 16, 7, 8, 9, 10, 11, 12, 13, 14, 15]; + let mut num_codes = 0; + for i in simple_code_or_skip..CODE_LENGTHS_CODE { + if space <= 0 { + break; + } + let idx = br.peek(4) as usize; + br.consume(STATIC_HUFF_BITS[idx] as usize)?; + let v = STATIC_HUFF_VALS[idx]; + code_length_code_lengths[CODE_LENGTH_CODE_ORDER[i] as usize] = v; + if v != 0 { + space -= 32 >> v; + num_codes += 1; + } + } + if num_codes != 1 && space != 0 { + return Err(Error::InvalidHuffman); + } + let code_lengths = + Table::decode_huffman_code_lengths(code_length_code_lengths, al_size, br)?; + debug!(?code_lengths); + Table::build(TABLE_BITS, &code_lengths)? + } + }; + Ok(Table { entries }) + } + + pub fn read(&self, br: &mut BitReader) -> Result<u32> { + let mut pos = br.peek(TABLE_BITS) as usize; + let mut n_bits = self.entries[pos].bits as usize; + if n_bits > TABLE_BITS { + br.consume(TABLE_BITS)?; + n_bits -= TABLE_BITS; + pos += self.entries[pos].value as usize; + pos += br.peek(n_bits) as usize; + } + br.consume(self.entries[pos].bits as usize)?; + Ok(self.entries[pos].value as u32) + } +} + +#[derive(Debug)] +pub struct HuffmanCodes { + tables: Vec<Table>, +} + +impl HuffmanCodes { + pub fn decode(num: usize, br: &mut BitReader) -> Result<HuffmanCodes> { + let alphabet_sizes: Vec<u16> = (0..num) + .map(|_| Ok(decode_varint16(br)? + 1)) + .collect::<Result<_>>()?; + let max = *alphabet_sizes.iter().max().unwrap(); + if max as usize > (1 << HUFFMAN_MAX_BITS) { + return Err(Error::AlphabetTooLargeHuff(max as usize)); + } + let tables = alphabet_sizes + .iter() + .map(|sz| Table::decode(*sz as usize, br)) + .collect::<Result<_>>()?; + Ok(HuffmanCodes { tables }) + } + pub fn read(&self, br: &mut BitReader, ctx: usize) -> Result<u32> { + self.tables[ctx].read(br) + } +} + +#[cfg(test)] +impl HuffmanCodes { + /// Builds Huffman histogram of 256 8-bit symbols. + pub(super) fn byte_histogram() -> HuffmanCodes { + let mut br = BitReader::new(&[0b11101111, 0b00111111, 0, 1, 0, 0b10100000, 0b0110]); + HuffmanCodes::decode(1, &mut br).unwrap() + } +} + +#[cfg(test)] +mod test { + use super::*; + use test_log::test; + + #[test] + fn byte_histogram() { + let codes = HuffmanCodes::byte_histogram(); + + let expected_arr = [8u8, 13, 21, 34, 55, 89, 144, 233]; + let bits = expected_arr.map(|v| v.reverse_bits()); + let mut br = BitReader::new(&bits); + + for expected in expected_arr { + assert_eq!(codes.read(&mut br, 0).unwrap(), expected as u32); + } + } + + #[test] + fn long_code() { + // This correctly sums to 4096 table entries + const CODE: [u8; 520] = [ + 3, 6, 7, 7, 7, 7, 7, 7, 7, 7, 8, 7, 8, 7, 8, 8, 8, 8, 8, 8, 8, 9, 9, 8, 9, 9, 9, 9, 9, + 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, + 9, 9, 9, 9, 9, 9, 9, 8, 9, 9, 8, 9, 9, 9, 9, 9, 8, 9, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, + 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, + 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 7, 7, 7, 7, 8, 8, 8, 8, 8, 8, 8, 8, 7, 8, 8, 8, 8, 8, + 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 9, 8, 9, 8, 8, 8, 8, 9, 9, 9, 9, 8, 8, 9, 9, 9, 9, 9, 9, + 9, 8, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, + 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 10, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 10, 9, 9, 9, 9, + 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 8, 4, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 6, 10, 7, 9, 9, 11, 12, 12, + ]; + assert!(Table::build(TABLE_BITS, &CODE).is_ok()); + } +} diff --git a/third_party/rust/jxl/src/entropy_coding/hybrid_uint.rs b/third_party/rust/jxl/src/entropy_coding/hybrid_uint.rs @@ -0,0 +1,80 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +use crate::bit_reader::BitReader; +use crate::error::Error; + +use crate::util::CeilLog2; + +#[derive(Debug)] +pub struct HybridUint { + split_token: u32, + split_exponent: u32, + msb_in_token: u32, + lsb_in_token: u32, +} + +impl HybridUint { + pub fn decode(log_alpha_size: usize, br: &mut BitReader) -> Result<HybridUint, Error> { + let split_exponent = br.read((log_alpha_size + 1).ceil_log2())? as u32; + let split_token = 1u32 << split_exponent; + let msb_in_token; + let lsb_in_token; + if split_exponent != log_alpha_size as u32 { + let nbits = (split_exponent + 1).ceil_log2() as usize; + msb_in_token = br.read(nbits)? as u32; + if msb_in_token > split_exponent { + return Err(Error::InvalidUintConfig(split_exponent, msb_in_token, None)); + } + let nbits = (split_exponent - msb_in_token + 1).ceil_log2() as usize; + lsb_in_token = br.read(nbits)? as u32; + } else { + msb_in_token = 0; + lsb_in_token = 0; + } + if lsb_in_token + msb_in_token > split_exponent { + return Err(Error::InvalidUintConfig( + split_exponent, + msb_in_token, + Some(lsb_in_token), + )); + } + Ok(HybridUint { + split_token, + split_exponent, + msb_in_token, + lsb_in_token, + }) + } + pub fn read(&self, symbol: u32, br: &mut BitReader) -> Result<u32, Error> { + if symbol < self.split_token { + return Ok(symbol); + } + let bits_in_token = self.lsb_in_token + self.msb_in_token; + let nbits = + self.split_exponent - bits_in_token + ((symbol - self.split_token) >> bits_in_token); + // To match the behaviour of libjxl, we limit nbits to 31. + if nbits > 31 { + return Err(Error::IntegerTooLarge(nbits)); + } + let low = symbol & ((1 << self.lsb_in_token) - 1); + let symbol_nolow = symbol >> self.lsb_in_token; + let bits = br.read(nbits as usize)? as u32; + let hi = (symbol_nolow & ((1 << self.msb_in_token) - 1)) | (1 << self.msb_in_token); + Ok((((hi << nbits) | bits) << self.lsb_in_token) | low) + } +} + +#[cfg(test)] +impl HybridUint { + pub fn new(split_exponent: u32, msb_in_token: u32, lsb_in_token: u32) -> Self { + Self { + split_token: 1 << split_exponent, + split_exponent, + msb_in_token, + lsb_in_token, + } + } +} diff --git a/third_party/rust/jxl/src/entropy_coding/mod.rs b/third_party/rust/jxl/src/entropy_coding/mod.rs @@ -0,0 +1,10 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +pub mod ans; +pub mod context_map; +pub mod decode; +pub mod huffman; +pub mod hybrid_uint; diff --git a/third_party/rust/jxl/src/error.rs b/third_party/rust/jxl/src/error.rs @@ -0,0 +1,253 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +use std::collections::TryReserveError; + +use thiserror::Error; + +use crate::{ + entropy_coding::huffman::HUFFMAN_MAX_BITS, features::spline::Point, image::DataTypeTag, +}; + +#[derive(Error, Debug)] +#[non_exhaustive] +pub enum Error { + #[error("Invalid raw quantization table")] + InvalidRawQuantTable, + #[error("Invalid distance band {0}: {1}")] + InvalidDistanceBand(usize, f32), + #[error("Invalid AFV bands")] + InvalidAFVBands, + #[error("Invalid quantization table weight: {0}")] + InvalidQuantizationTableWeight(f32), + #[error("Read out of bounds; size hint: {0}")] + OutOfBounds(usize), + #[error("Section is too short")] + SectionTooShort, + #[error("Non-zero padding bits")] + NonZeroPadding, + #[error("Invalid signature")] + InvalidSignature, + #[error("Invalid exponent_bits_per_sample: {0}")] + InvalidExponent(u32), + #[error("Invalid mantissa_bits: {0}")] + InvalidMantissa(i32), + #[error("Invalid bits_per_sample: {0}")] + InvalidBitsPerSample(u32), + #[error("Invalid enum value {0} for {1}")] + InvalidEnum(u32, String), + #[error("Value of dim_shift {0} is too large")] + DimShiftTooLarge(u32), + #[error("Float is NaN or Inf")] + FloatNaNOrInf, + #[error("Invalid gamma value: {0}")] + InvalidGamma(f32), + #[error("Invalid color encoding: no ICC and unknown TF / ColorSpace")] + InvalidColorEncoding, + #[error("Invalid color space: should be one of RGB, Gray or XYB")] + InvalidColorSpace, + #[error("Only perceptual rendering intent implemented for XYB ICC profile.")] + InvalidRenderingIntent, + #[error("Invalid intensity_target: {0}")] + InvalidIntensityTarget(f32), + #[error("Invalid min_nits: {0}")] + InvalidMinNits(f32), + #[error("Invalid linear_below {1}, relative_to_max_display is {0}")] + InvalidLinearBelow(bool, f32), + #[error("Overflow when computing a bitstream size")] + SizeOverflow, + #[error("Invalid ISOBMMF container")] + InvalidBox, + #[error("ICC is too large")] + IccTooLarge, + #[error("Invalid ICC stream: unexpected end of stream")] + IccEndOfStream, + #[error("Invalid ICC stream")] + InvalidIccStream, + #[error("Invalid HybridUintConfig: {0} {1} {2:?}")] + InvalidUintConfig(u32, u32, Option<u32>), + #[error("LZ77 enabled when explicitly disallowed")] + Lz77Disallowed, + #[error("LZ77 repeat symbol encountered without decoding any symbols")] + UnexpectedLz77Repeat, + #[error("Huffman alphabet too large: {0}, max is {max}", max = 1 << HUFFMAN_MAX_BITS)] + AlphabetTooLargeHuff(usize), + #[error("Invalid Huffman code")] + InvalidHuffman, + #[error("Invalid ANS histogram")] + InvalidAnsHistogram, + #[error("ANS stream checksum mismatch")] + AnsChecksumMismatch, + #[error("Integer too large: nbits {0} > 29")] + IntegerTooLarge(u32), + #[error("Invalid context map: context id {0} > 255")] + InvalidContextMap(u32), + #[error("Invalid context map: number of histogram {0}, number of distinct histograms {1}")] + InvalidContextMapHole(u32, u32), + #[error( + "Invalid permutation: skipped elements {skip} and encoded elements {end} don't fit in permutation of size {size}" + )] + InvalidPermutationSize { size: u32, skip: u32, end: u32 }, + #[error( + "Invalid permutation: Lehmer code {lehmer} out of bounds in permutation of size {size} at index {idx}" + )] + InvalidPermutationLehmerCode { size: u32, idx: u32, lehmer: u32 }, + #[error("Invalid quant encoding mode")] + InvalidQuantEncodingMode, + #[error("Invalid quant encoding with mode {mode} and required size {required_size}")] + InvalidQuantEncoding { mode: u8, required_size: usize }, + // FrameHeader format errors + #[error("Invalid extra channel upsampling: upsampling: {0} dim_shift: {1} ec_upsampling: {2}")] + InvalidEcUpsampling(u32, u32, u32), + #[error("Num_ds: {0} should be smaller than num_passes: {1}")] + NumPassesTooLarge(u32, u32), + #[error("Non-patch reference frame with a crop")] + NonPatchReferenceWithCrop, + #[error("Non-444 chroma subsampling is not allowed when adaptive DC smoothing is enabled")] + Non444ChromaSubsampling, + #[error("Non-444 chroma subsampling is not allowed for bigger than 8x8 transforms")] + InvalidBlockSizeForChromaSubsampling, + #[error("Out of memory: {0}")] + OutOfMemory(#[from] TryReserveError), + #[error("Image size too large: {0}x{1}")] + ImageSizeTooLarge(usize, usize), + #[error("Invalid image size: {0}x{1}")] + InvalidImageSize(usize, usize), + #[error("Rect out of bounds: {0}x{1}+{2}+{3} rect in {4}x{5} view")] + RectOutOfBounds(usize, usize, usize, usize, usize, usize), + // Generic arithmetic overflow. Prefer using other errors if possible. + #[error("Arithmetic overflow")] + ArithmeticOverflow, + #[error("Empty frame sequence")] + NoFrames, + #[error( + "Pipeline channel type mismatch: stage {0} channel {1}, expected {2:?} but found {3:?}" + )] + PipelineChannelTypeMismatch(String, usize, DataTypeTag, DataTypeTag), + #[error("Pipeline has a stage ({0}) with a shift after an expand stage")] + PipelineShiftAfterExpand(String), + #[error("Channel {0} was not used in the render pipeline")] + PipelineChannelUnused(usize), + #[error("Trying to copy rects of different size, src: {0}x{1} dst {2}x{3}")] + CopyOfDifferentSize(usize, usize, usize, usize), + #[error("LF quantization factor is too small: {0}")] + LfQuantFactorTooSmall(f32), + #[error("HF quantization factor is too small: {0}")] + HfQuantFactorTooSmall(f32), + #[error("Invalid modular mode predictor: {0}")] + InvalidPredictor(u32), + #[error("Invalid modular mode property: {0}")] + InvalidProperty(u32), + #[error("Invalid alpha channel for blending: {0}, limit is {1}")] + PatchesInvalidAlphaChannel(usize, usize), + #[error("Invalid patch blend mode: {0}, limit is {1}")] + PatchesInvalidBlendMode(u8, u8), + #[error("Invalid Patch: negative {0}-coordinate: {1} base {0}, {2} delta {0}")] + PatchesInvalidDelta(String, usize, i32), + #[error( + "Invalid position specified in reference frame in {0}-coordinate: {0}0 + {0}size = {1} + {2} > {3} = reference_frame {0}size" + )] + PatchesInvalidPosition(String, usize, usize, usize), + #[error("Patches invalid reference frame at index {0}")] + PatchesInvalidReference(usize), + #[error("Invalid Patch {0}: at {1} + {2} > {3}")] + PatchesOutOfBounds(String, usize, usize, usize), + #[error("Patches cannot use frames saved post color transforms")] + PatchesPostColorTransform(), + #[error("Too many {0}: {1}, limit is {2}")] + PatchesTooMany(String, usize, usize), + #[error("Reference too large: {0}, limit is {1}")] + PatchesRefTooLarge(usize, usize), + #[error("Point list is empty")] + PointListEmpty, + #[error("Too large area for spline: {0}, limit is {1}")] + SplinesAreaTooLarge(u64, u64), + #[error("Too large manhattan_distance reached: {0}, limit is {1}")] + SplinesDistanceTooLarge(u64, u64), + #[error("Too many splines: {0}, limit is {1}")] + SplinesTooMany(u32, u32), + #[error("Spline has adjacent coinciding control points: point[{0}]: {1:?}, point[{2}]: {3:?}")] + SplineAdjacentCoincidingControlPoints(usize, Point, usize, Point), + #[error("Too many control points for splines: {0}, limit is {1}")] + SplinesTooManyControlPoints(u32, u32), + #[error( + "Spline point outside valid bounds: coordinates: {0:?}, out of bounds: {1}, bounds: {2:?}" + )] + SplinesPointOutOfRange(Point, i32, std::ops::Range<i32>), + #[error("Spline coordinates out of bounds: {0}, limit is {1}")] + SplinesCoordinatesLimit(i32, i32), + #[error("Spline delta-delta is out of bounds: {0}, limit is {1}")] + SplinesDeltaLimit(i64, i64), + #[error("Modular tree too large: {0}, limit is {1}")] + TreeTooLarge(usize, usize), + #[error("Modular tree too tall: {0}, limit is {1}")] + TreeTooTall(usize, usize), + #[error("Modular tree multiplier too large: {0}, limit is {1}")] + TreeMultiplierTooLarge(u32, u32), + #[error("Modular tree multiplier too large: {0}, multiplier log is {1}")] + TreeMultiplierBitsTooLarge(u32, u32), + #[error( + "Modular tree splits on property {0} at value {1}, which is outside the possible range of [{2}, {3}]" + )] + TreeSplitOnEmptyRange(u8, i32, i32, i32), + #[error("Modular stream requested a global tree but there isn't one")] + NoGlobalTree, + #[error("Invalid transform id")] + InvalidTransformId, + #[error("Invalid RCT type {0}")] + InvalidRCT(u32), + #[error("Invalid channel range: {0}..{1}, {2} total channels")] + InvalidChannelRange(usize, usize, usize), + #[error("Invalid transform: mixing different channels (different shape or different shift)")] + MixingDifferentChannels, + #[error("Invalid transform: squeezing meta-channels needs an in-place transform")] + MetaSqueezeRequiresInPlace, + #[error("Invalid transform: too many squeezes (shift > 30)")] + TooManySqueezes, + #[error("Invalid BlockConextMap: too big: num_lf_context: {0}, num_qf_thresholds: {1}")] + BlockContextMapSizeTooBig(usize, usize), + #[error("Invalid BlockConextMap: too many distinct contexts.")] + TooManyBlockContexts, + #[error("Base color correlation out of range.")] + BaseColorCorrelationOutOfRange, + #[error("Invalid EPF sharpness param {0}")] + InvalidEpfValue(i32), + #[error("Invalid VarDCT transform type {0}")] + InvalidVarDCTTransform(usize), + #[error("Invalid VarDCT transform map")] + InvalidVarDCTTransformMap, + #[error("VarDCT transform overflows HF group")] + HFBlockOutOfBounds, + #[error("Invalid AC: nonzeros {0} is too large for {1} 8x8 blocks")] + InvalidNumNonZeros(usize, usize), + #[error("Invalid AC: {0} nonzeros after decoding block")] + EndOfBlockResidualNonZeros(usize), + #[error("Unknown transfer function for ICC profile")] + TransferFunctionUnknown, + #[error("Attempting to write out of Bounds when writing ICC")] + IccWriteOutOfBounds, + #[error("Invalid tag string when writing ICC: {0}")] + IccInvalidTagString(String), + #[error("Invalid text for ICC MLuc string, not ascii: {0}")] + IccMlucTextNotAscii(String), + #[error("ICC value is out of range / NaN: {0}")] + IccValueOutOfRangeS15Fixed16(f32), + #[error("Y value is too small: {0}")] + IccInvalidWhitePointY(f32), + #[error("{2}: wx: {0}, wy: {1}")] + IccInvalidWhitePoint(f32, f32, String), + #[error("Determinant is zero or too small, matrix is close to singular: |det| = {0}.")] + MatrixInversionFailed(f64), + #[error("Unsupported transfer function when writing ICC")] + IccUnsupportedTransferFunction, + #[error("Table size too large when writing ICC: {0}")] + IccTableSizeExceeded(usize), + #[error("Invalid CMS configuration: requested ICC but no CMS is configured")] + ICCOutputNoCMS, + #[error("I/O error: {0}")] + IOError(#[from] std::io::Error), +} + +pub type Result<T, E = Error> = std::result::Result<T, E>; diff --git a/third_party/rust/jxl/src/features/blending.rs b/third_party/rust/jxl/src/features/blending.rs @@ -0,0 +1,781 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +use crate::headers::extra_channels::{ExtraChannel, ExtraChannelInfo}; + +use super::patches::{PatchBlendMode, PatchBlending}; + +#[inline] +fn maybe_clamp(v: f32, clamp: bool) -> f32 { + if clamp { v.clamp(0.0, 1.0) } else { v } +} + +pub fn perform_blending<T: AsRef<[f32]>, V: AsMut<[f32]>>( + bg: &mut [V], + fg: &[T], + color_blending: &PatchBlending, + ec_blending: &[PatchBlending], + extra_channel_info: &[ExtraChannelInfo], +) { + let has_alpha = extra_channel_info + .iter() + .any(|info| info.ec_type == ExtraChannel::Alpha); + let num_ec = extra_channel_info.len(); + let xsize = bg[0].as_mut().len(); + + let mut tmp = vec![vec![0.0f32; xsize]; 3 + num_ec]; + + for i in 0..num_ec { + let alpha = ec_blending[i].alpha_channel; + let clamp = ec_blending[i].clamp; + let alpha_associated = extra_channel_info[alpha].alpha_associated(); + + match ec_blending[i].mode { + PatchBlendMode::Add => { + for x in 0..xsize { + tmp[3 + i][x] = bg[3 + i].as_mut()[x] + fg[3 + i].as_ref()[x]; + } + } + PatchBlendMode::BlendAbove => { + if i == alpha { + for x in 0..xsize { + let fa = maybe_clamp(fg[3 + alpha].as_ref()[x], clamp); + tmp[3 + i][x] = 1.0 - (1.0 - fa) * (1.0 - bg[3 + i].as_mut()[x]); + } + } else if alpha_associated { + for x in 0..xsize { + let fa = maybe_clamp(fg[3 + alpha].as_ref()[x], clamp); + tmp[3 + i][x] = fg[3 + i].as_ref()[x] + bg[3 + i].as_mut()[x] * (1.0 - fa); + } + } else { + for x in 0..xsize { + let fa = maybe_clamp(fg[3 + alpha].as_ref()[x], clamp); + let new_a = 1.0 - (1.0 - fa) * (1.0 - bg[3 + alpha].as_mut()[x]); + let rnew_a = if new_a > 0.0 { 1.0 / new_a } else { 0.0 }; + tmp[3 + i][x] = (fg[3 + i].as_ref()[x] * fa + + bg[3 + i].as_mut()[x] * bg[3 + alpha].as_mut()[x] * (1.0 - fa)) + * rnew_a; + } + } + } + PatchBlendMode::BlendBelow => { + if i == alpha { + for x in 0..xsize { + let ba = maybe_clamp(bg[3 + alpha].as_mut()[x], clamp); + tmp[3 + i][x] = 1.0 - (1.0 - ba) * (1.0 - fg[3 + i].as_ref()[x]); + } + } else if alpha_associated { + for x in 0..xsize { + let ba = maybe_clamp(bg[3 + alpha].as_mut()[x], clamp); + tmp[3 + i][x] = bg[3 + i].as_mut()[x] + fg[3 + i].as_ref()[x] * (1.0 - ba); + } + } else { + for x in 0..xsize { + let ba = maybe_clamp(bg[3 + alpha].as_mut()[x], clamp); + let new_a = 1.0 - (1.0 - ba) * (1.0 - fg[3 + alpha].as_ref()[x]); + let rnew_a = if new_a > 0.0 { 1.0 / new_a } else { 0.0 }; + tmp[3 + i][x] = (bg[3 + i].as_mut()[x] * ba + + fg[3 + i].as_ref()[x] * fg[3 + alpha].as_ref()[x] * (1.0 - ba)) + * rnew_a; + } + } + } + PatchBlendMode::AlphaWeightedAddAbove => { + if i == alpha { + tmp[3 + i].copy_from_slice(bg[3 + i].as_mut()); + } else if clamp { + for x in 0..xsize { + tmp[3 + i][x] = bg[3 + i].as_mut()[x] + + fg[3 + i].as_ref()[x] * fg[3 + alpha].as_ref()[x].clamp(0.0, 1.0); + } + } else { + for x in 0..xsize { + tmp[3 + i][x] = bg[3 + i].as_mut()[x] + + fg[3 + i].as_ref()[x] * fg[3 + alpha].as_ref()[x]; + } + } + } + PatchBlendMode::AlphaWeightedAddBelow => { + if i == alpha { + tmp[3 + i].copy_from_slice(fg[3 + i].as_ref()); + } else if clamp { + for x in 0..xsize { + tmp[3 + i][x] = fg[3 + i].as_ref()[x] + + bg[3 + i].as_mut()[x] * bg[3 + alpha].as_mut()[x].clamp(0.0, 1.0); + } + } else { + for x in 0..xsize { + tmp[3 + i][x] = fg[3 + i].as_ref()[x] + + bg[3 + i].as_mut()[x] * bg[3 + alpha].as_mut()[x]; + } + } + } + PatchBlendMode::Mul => { + if clamp { + for x in 0..xsize { + tmp[3 + i][x] = + bg[3 + i].as_mut()[x] * fg[3 + i].as_ref()[x].clamp(0.0, 1.0); + } + } else { + for x in 0..xsize { + tmp[3 + i][x] = bg[3 + i].as_mut()[x] * fg[3 + i].as_ref()[x]; + } + } + } + PatchBlendMode::Replace => { + tmp[3 + i].copy_from_slice(fg[3 + i].as_ref()); + } + PatchBlendMode::None => { + tmp[3 + i].copy_from_slice(bg[3 + i].as_mut()); + } + } + } + + let alpha = color_blending.alpha_channel; + let clamp = color_blending.clamp; + + match color_blending.mode { + PatchBlendMode::Add => { + for c in 0..3 { + for x in 0..xsize { + tmp[c][x] = bg[c].as_mut()[x] + fg[c].as_ref()[x]; + } + } + } + PatchBlendMode::AlphaWeightedAddAbove => { + for c in 0..3 { + if !has_alpha { + for x in 0..xsize { + tmp[c][x] = bg[c].as_mut()[x] + fg[c].as_ref()[x]; + } + } else if clamp { + for x in 0..xsize { + tmp[c][x] = bg[c].as_mut()[x] + + fg[c].as_ref()[x] * fg[3 + alpha].as_ref()[x].clamp(0.0, 1.0); + } + } else { + for x in 0..xsize { + tmp[c][x] = + bg[c].as_mut()[x] + fg[c].as_ref()[x] * fg[3 + alpha].as_ref()[x]; + } + } + } + } + PatchBlendMode::AlphaWeightedAddBelow => { + for c in 0..3 { + if !has_alpha { + for x in 0..xsize { + tmp[c][x] = bg[c].as_mut()[x] + fg[c].as_ref()[x]; + } + } else if clamp { + for x in 0..xsize { + tmp[c][x] = fg[c].as_ref()[x] + + bg[c].as_mut()[x] * bg[3 + alpha].as_mut()[x].clamp(0.0, 1.0); + } + } else { + for x in 0..xsize { + tmp[c][x] = + fg[c].as_ref()[x] + bg[c].as_mut()[x] * bg[3 + alpha].as_mut()[x]; + } + } + } + } + PatchBlendMode::BlendAbove => { + if !has_alpha { + for c in 0..3 { + tmp[c].copy_from_slice(fg[c].as_ref()); + } + } else if extra_channel_info[alpha].alpha_associated() { + for x in 0..xsize { + let fa = maybe_clamp(fg[3 + alpha].as_ref()[x], clamp); + for c in 0..3 { + tmp[c][x] = fg[c].as_ref()[x] + bg[c].as_mut()[x] * (1.0 - fa); + } + tmp[3 + alpha][x] = 1.0 - (1.0 - fa) * (1.0 - bg[3 + alpha].as_mut()[x]); + } + } else { + for x in 0..xsize { + let fa = maybe_clamp(fg[3 + alpha].as_ref()[x], clamp); + let new_a = 1.0 - (1.0 - fa) * (1.0 - bg[3 + alpha].as_mut()[x]); + let rnew_a = if new_a > 0.0 { 1.0 / new_a } else { 0.0 }; + for c in 0..3 { + tmp[c][x] = (fg[c].as_ref()[x] * fa + + bg[c].as_mut()[x] * bg[3 + alpha].as_mut()[x] * (1.0 - fa)) + * rnew_a; + } + tmp[3 + alpha][x] = new_a; + } + } + } + PatchBlendMode::BlendBelow => { + if !has_alpha { + for c in 0..3 { + tmp[c].copy_from_slice(bg[c].as_mut()); + } + } else if extra_channel_info[alpha].alpha_associated() { + for x in 0..xsize { + let ba = maybe_clamp(bg[3 + alpha].as_mut()[x], clamp); + for c in 0..3 { + tmp[c][x] = bg[c].as_mut()[x] + fg[c].as_ref()[x] * (1.0 - ba); + } + tmp[3 + alpha][x] = 1.0 - (1.0 - ba) * (1.0 - fg[3 + alpha].as_ref()[x]); + } + } else { + for x in 0..xsize { + let ba = maybe_clamp(bg[3 + alpha].as_mut()[x], clamp); + let new_a = 1.0 - (1.0 - ba) * (1.0 - fg[3 + alpha].as_ref()[x]); + let rnew_a = if new_a > 0.0 { 1.0 / new_a } else { 0.0 }; + for c in 0..3 { + tmp[c][x] = (bg[c].as_mut()[x] * ba + + fg[c].as_ref()[x] * fg[3 + alpha].as_ref()[x] * (1.0 - ba)) + * rnew_a; + } + tmp[3 + alpha][x] = new_a; + } + } + } + PatchBlendMode::Mul => { + for c in 0..3 { + for x in 0..xsize { + tmp[c][x] = bg[c].as_mut()[x] * maybe_clamp(fg[c].as_ref()[x], clamp); + } + } + } + PatchBlendMode::Replace => { + for c in 0..3 { + tmp[c].copy_from_slice(fg[c].as_ref()); + } + } + PatchBlendMode::None => { + for c in 0..3 { + tmp[c].copy_from_slice(bg[c].as_mut()); + } + } + } + for i in 0..(3 + num_ec) { + bg[i].as_mut().copy_from_slice(&tmp[i]); + } +} + +#[cfg(test)] +mod tests { + fn clamp(x: f32) -> f32 { + x.clamp(0.0, 1.0) + } + + mod perform_blending_tests { + use super::{super::*, *}; + use crate::{headers::bit_depth::BitDepth, util::test::assert_all_almost_abs_eq}; + use test_log::test; + + const ABS_DELTA: f32 = 1e-6; + + // Helper for expected value calculations based on C++ logic + + // Alpha compositing formula: Ao = As + Ab * (1 - As) + // Used for kBlend modes for the alpha channel itself. + fn expected_alpha_blend(fg_a: f32, bg_a: f32) -> f32 { + fg_a + bg_a * (1.0 - fg_a) + } + + // Color compositing for kBlend, premultiplied alpha: Co = Cs_premult + Cb_premult * (1 - As) + fn expected_color_blend_premultiplied(c_fg: f32, c_bg: f32, fg_a: f32) -> f32 { + c_fg + c_bg * (1.0 - fg_a) + } + + // Color compositing for kBlend, non-premultiplied alpha: Co = (Cs * As + Cb * Ab * (1 - As)) / Ao_blend + fn expected_color_blend_non_premultiplied( + c_fg: f32, + fg_a: f32, // Foreground color and its alpha + c_bg: f32, + bg_a: f32, // Background color and its alpha + alpha_blend_out: f32, // The resulting alpha from expected_alpha_blend(fg_a, bg_a) + ) -> f32 { + if alpha_blend_out.abs() < ABS_DELTA { + // Avoid division by zero + 0.0 + } else { + (c_fg * fg_a + c_bg * bg_a * (1.0 - fg_a)) / alpha_blend_out + } + } + + // For kAlphaWeightedAdd modes: Co = Cb + Cs * As + fn expected_alpha_weighted_add(c_bg: f32, c_fg: f32, fg_a: f32) -> f32 { + c_bg + c_fg * fg_a + } + + // For kMul mode: Co = Cb * Cs + fn expected_mul_blend(c_bg: f32, c_fg: f32) -> f32 { + c_bg * c_fg + } + + #[test] + fn test_color_replace_fg_over_bg() { + let mut bg_r = [0.1]; + let mut bg_g = [0.2]; + let mut bg_b = [0.3]; + let fg_r = [0.7]; + let fg_g = [0.8]; + let fg_b = [0.9]; + + let mut bg_channels: [&mut [f32]; 3] = [&mut bg_r, &mut bg_g, &mut bg_b]; + let fg_channels: [&[f32]; 3] = [&fg_r, &fg_g, &fg_b]; + + let color_blending = PatchBlending { + mode: PatchBlendMode::Replace, + alpha_channel: 0, // Not used for Replace + clamp: false, + }; + + let ec_blending: [PatchBlending; 0] = []; + let extra_channel_info: [ExtraChannelInfo; 0] = []; + + perform_blending( + &mut bg_channels, + &fg_channels, + &color_blending, + &ec_blending, + &extra_channel_info, + ); + + // Expected: output color is fg color + assert_all_almost_abs_eq(&bg_r, &fg_r, ABS_DELTA); + assert_all_almost_abs_eq(&bg_g, &fg_g, ABS_DELTA); + assert_all_almost_abs_eq(&bg_b, &fg_b, ABS_DELTA); + } + + #[test] + fn test_color_add() { + let mut bg_r = [0.1]; + let mut bg_g = [0.2]; + let mut bg_b = [0.3]; + let fg_r = [0.7]; + let fg_g = [0.6]; + let fg_b = [0.5]; + let expected_r = [bg_r[0] + fg_r[0]]; + let expected_g = [bg_g[0] + fg_g[0]]; + let expected_b = [bg_b[0] + fg_b[0]]; + + let mut bg_channels: [&mut [f32]; 3] = [&mut bg_r, &mut bg_g, &mut bg_b]; + let fg_channels: [&[f32]; 3] = [&fg_r, &fg_g, &fg_b]; + + let color_blending = PatchBlending { + mode: PatchBlendMode::Add, + alpha_channel: 0, // Not used + clamp: false, + }; + let ec_blending: [PatchBlending; 0] = []; + let extra_channel_info: [ExtraChannelInfo; 0] = []; + + perform_blending( + &mut bg_channels, + &fg_channels, + &color_blending, + &ec_blending, + &extra_channel_info, + ); + + assert_all_almost_abs_eq(&bg_r, &expected_r, ABS_DELTA); + assert_all_almost_abs_eq(&bg_g, &expected_g, ABS_DELTA); + assert_all_almost_abs_eq(&bg_b, &expected_b, ABS_DELTA); + } + + #[test] + fn test_color_blend_above_premultiplied_alpha() { + // BG: R=0.1, G=0.2, B=0.3, A=0.8 (premultiplied) + // FG: R=0.4, G=0.3, B=0.2, A=0.5 (premultiplied) + let mut bg_r = [0.1]; + let mut bg_g = [0.2]; + let mut bg_b = [0.3]; + let mut bg_a = [0.8]; + let fg_r = [0.4]; + let fg_g = [0.3]; + let fg_b = [0.2]; + let fg_a = [0.5]; + let fga = fg_a[0]; // Not clamped + let bga = bg_a[0]; + + // Expected alpha: Ao = Afg + Abg * (1 - Afg) + let expected_a_val = expected_alpha_blend(fga, bga); + // Expected color: Co = Cfg_premult + Cbg_premult * (1 - Afg) + let expected_r_val = expected_color_blend_premultiplied(fg_r[0], bg_r[0], fga); + let expected_g_val = expected_color_blend_premultiplied(fg_g[0], bg_g[0], fga); + let expected_b_val = expected_color_blend_premultiplied(fg_b[0], bg_b[0], fga); + + let mut bg_channels: [&mut [f32]; 4] = [&mut bg_r, &mut bg_g, &mut bg_b, &mut bg_a]; + let fg_channels: [&[f32]; 4] = [&fg_r, &fg_g, &fg_b, &fg_a]; + + let color_blending = PatchBlending { + mode: PatchBlendMode::BlendAbove, + alpha_channel: 0, // EC0 is the alpha channel + clamp: false, + }; + // EC0 (alpha) blending rule. + // For BlendAbove color mode, the alpha channel itself is also blended using source-over. + // So this ec_blending rule for alpha will be effectively overridden by color blending's alpha calculation. + let ec_blending = [PatchBlending { + mode: PatchBlendMode::Replace, // Arbitrary, will be overwritten by color blend + alpha_channel: 0, + clamp: false, + }]; + let extra_channel_info = [ExtraChannelInfo::new( + false, + ExtraChannel::Alpha, + BitDepth::f32(), // Assuming f32 + 0, + "alpha".to_string(), + true, // alpha_associated = true (premultiplied) + None, + None, + )]; + + perform_blending( + &mut bg_channels, + &fg_channels, + &color_blending, + &ec_blending, + &extra_channel_info, + ); + + assert_all_almost_abs_eq(&bg_a, &[expected_a_val], ABS_DELTA); + assert_all_almost_abs_eq(&bg_r, &[expected_r_val], ABS_DELTA); + assert_all_almost_abs_eq(&bg_g, &[expected_g_val], ABS_DELTA); + assert_all_almost_abs_eq(&bg_b, &[expected_b_val], ABS_DELTA); + } + + #[test] + fn test_color_blend_above_non_premultiplied_alpha() { + // BG: R=0.1, G=0.2, B=0.3 (unpremult), A=0.8 + // FG: R=0.7, G=0.6, B=0.5 (unpremult), A=0.5 + let mut bg_r = [0.1]; + let mut bg_g = [0.2]; + let mut bg_b = [0.3]; + let mut bg_a = [0.8]; + let fg_r = [0.7]; + let fg_g = [0.6]; + let fg_b = [0.5]; + let fg_a = [0.5]; + let fga = fg_a[0]; + let bga = bg_a[0]; + + // Expected alpha: Ao = Afg + Abg * (1 - Afg) + let expected_a_val = expected_alpha_blend(fga, bga); + // Expected color: Co = (Cfg_unpremult * Afg + Cbg_unpremult * Abg * (1 - Afg)) / Ao_blend + let expected_r_val = + expected_color_blend_non_premultiplied(fg_r[0], fga, bg_r[0], bga, expected_a_val); + let expected_g_val = + expected_color_blend_non_premultiplied(fg_g[0], fga, bg_g[0], bga, expected_a_val); + let expected_b_val = + expected_color_blend_non_premultiplied(fg_b[0], fga, bg_b[0], bga, expected_a_val); + + let mut bg_channels: [&mut [f32]; 4] = [&mut bg_r, &mut bg_g, &mut bg_b, &mut bg_a]; + let fg_channels: [&[f32]; 4] = [&fg_r, &fg_g, &fg_b, &fg_a]; + + let color_blending = PatchBlending { + mode: PatchBlendMode::BlendAbove, + alpha_channel: 0, // EC0 + clamp: false, + }; + let ec_blending = [PatchBlending { + // This will be overwritten for the alpha channel by color blending + mode: PatchBlendMode::Replace, + alpha_channel: 0, + clamp: false, + }]; + let extra_channel_info = [ExtraChannelInfo::new( + false, + ExtraChannel::Alpha, + BitDepth::f32(), + 0, + "alpha".to_string(), + false, // alpha_associated = false (non-premultiplied) + None, + None, + )]; + + perform_blending( + &mut bg_channels, + &fg_channels, + &color_blending, + &ec_blending, + &extra_channel_info, + ); + + assert_all_almost_abs_eq(&bg_a, &[expected_a_val], ABS_DELTA); + assert_all_almost_abs_eq(&bg_r, &[expected_r_val], ABS_DELTA); + assert_all_almost_abs_eq(&bg_g, &[expected_g_val], ABS_DELTA); + assert_all_almost_abs_eq(&bg_b, &[expected_b_val], ABS_DELTA); + } + + #[test] + fn test_color_alpha_weighted_add_above() { + let mut bg_r = [0.1]; + let mut bg_g = [0.2]; + let mut bg_b = [0.3]; + let mut bg_a = [0.8]; // bg alpha used by ec_blending + let fg_r = [0.7]; + let fg_g = [0.6]; + let fg_b = [0.5]; + let fg_a = [0.5]; // fg alpha used for weighting + let fga_for_weighting = fg_a[0]; // Not clamped as color_blending.clamp is false + + // Expected color: Co = Cbg + Cfg * Afg_for_weighting + let expected_r_val = expected_alpha_weighted_add(bg_r[0], fg_r[0], fga_for_weighting); + let expected_g_val = expected_alpha_weighted_add(bg_g[0], fg_g[0], fga_for_weighting); + let expected_b_val = expected_alpha_weighted_add(bg_b[0], fg_b[0], fga_for_weighting); + + // Expected alpha (EC0): Blended according to ec_blending[0]. + // Mode is BlendAbove, alpha_channel is 0 (itself). + // C++: PerformAlphaBlending(bg[3+0], bg[3+0], fg[3+0], fg[3+0], tmp.Row(3+0), ...) + // This means it's the "alpha channel blends itself" case: Ao = Afg + Abg * (1 - Afg) + // fg_alpha_for_ec0 = fg_a[0], bg_alpha_for_ec0 = bg_a[0]. ec_blending[0].clamp is false. + let expected_a_val = expected_alpha_blend(fg_a[0], bg_a[0]); + + let mut bg_channels: [&mut [f32]; 4] = [&mut bg_r, &mut bg_g, &mut bg_b, &mut bg_a]; + let fg_channels: [&[f32]; 4] = [&fg_r, &fg_g, &fg_b, &fg_a]; + + let color_blending = PatchBlending { + mode: PatchBlendMode::AlphaWeightedAddAbove, + alpha_channel: 0, // EC0 is alpha + clamp: false, + }; + // For AlphaWeightedAdd color mode, the alpha channel (EC0) value is determined by its ec_blending rule. + // Let's make EC0 blend itself using BlendAbove mode. + let ec_blending = [PatchBlending { + mode: PatchBlendMode::BlendAbove, // Alpha channel EC0 blends itself + alpha_channel: 0, // using itself as alpha reference + clamp: false, + }]; + let extra_channel_info = [ExtraChannelInfo::new( + false, + ExtraChannel::Alpha, + BitDepth::f32(), + 0, + "alpha".to_string(), + true, // alpha_associated (doesn't directly affect AWA color, but affects EC0's own blend) + None, + None, + )]; + + perform_blending( + &mut bg_channels, + &fg_channels, + &color_blending, + &ec_blending, + &extra_channel_info, + ); + + assert_all_almost_abs_eq(&bg_r, &[expected_r_val], ABS_DELTA); + assert_all_almost_abs_eq(&bg_g, &[expected_g_val], ABS_DELTA); + assert_all_almost_abs_eq(&bg_b, &[expected_b_val], ABS_DELTA); + assert_all_almost_abs_eq(&bg_a, &[expected_a_val], ABS_DELTA); + } + + #[test] + fn test_color_mul_with_clamp() { + let mut bg_r = [0.5]; + let mut bg_g = [0.8]; + let mut bg_b = [1.0]; + let fg_r = [1.5]; + let fg_g = [-0.2]; + let fg_b = [0.5]; // fg values will be clamped + let expected_r = [expected_mul_blend(bg_r[0], clamp(fg_r[0]))]; // 0.5 * 1.0 = 0.5 + let expected_g = [expected_mul_blend(bg_g[0], clamp(fg_g[0]))]; // 0.8 * 0.0 = 0.0 + let expected_b = [expected_mul_blend(bg_b[0], clamp(fg_b[0]))]; // 1.0 * 0.5 = 0.5 + + let mut bg_channels: [&mut [f32]; 3] = [&mut bg_r, &mut bg_g, &mut bg_b]; + let fg_channels: [&[f32]; 3] = [&fg_r, &fg_g, &fg_b]; + + let color_blending = PatchBlending { + mode: PatchBlendMode::Mul, + alpha_channel: 0, // Not used + clamp: true, // Clamp fg values + }; + let ec_blending: [PatchBlending; 0] = []; + let extra_channel_info: [ExtraChannelInfo; 0] = []; + + perform_blending( + &mut bg_channels, + &fg_channels, + &color_blending, + &ec_blending, + &extra_channel_info, + ); + + assert_all_almost_abs_eq(&bg_r, &expected_r, ABS_DELTA); + assert_all_almost_abs_eq(&bg_g, &expected_g, ABS_DELTA); + assert_all_almost_abs_eq(&bg_b, &expected_b, ABS_DELTA); + } + + #[test] + fn test_ec_blend_data_with_separate_alpha_premultiplied() { + // Color: Replace FG over BG (to keep it simple) + // EC0: Data channel + // EC1: Alpha channel for EC0 + let mut bg_r = [0.1]; + let mut bg_g = [0.1]; + let mut bg_b = [0.1]; + let mut bg_ec0 = [0.2]; + let mut bg_ec1_alpha = [0.9]; // EC1 is alpha for EC0 + + let fg_r = [0.5]; + let fg_g = [0.5]; + let fg_b = [0.5]; + let fg_ec0 = [0.6]; + let fg_ec1_alpha = [0.4]; + + // EC1 (Alpha channel for EC0) blending: BlendAbove, uses itself as alpha. + // fg_alpha = fg_ec1_alpha[0], bg_alpha = bg_ec1_alpha[0] + let expected_out_ec1_alpha = expected_alpha_blend(fg_ec1_alpha[0], bg_ec1_alpha[0]); + + // EC0 (Data channel) blending: BlendAbove, uses EC1 as alpha. + // fg_alpha_for_ec0 = fg_ec1_alpha[0] (not clamped as ec_blending[0].clamp is false) + // is_premultiplied = extra_channel_info[ec_blending[0].alpha_channel (is 1)].alpha_associated = true. + // Formula: out = fg_data + bg_data * (1.f - fg_alpha_of_data) + let expected_out_ec0 = + expected_color_blend_premultiplied(fg_ec0[0], bg_ec0[0], fg_ec1_alpha[0]); + + let mut bg_channels: [&mut [f32]; 5] = [ + &mut bg_r, + &mut bg_g, + &mut bg_b, + &mut bg_ec0, + &mut bg_ec1_alpha, + ]; + let fg_channels: [&[f32]; 5] = [&fg_r, &fg_g, &fg_b, &fg_ec0, &fg_ec1_alpha]; + + let color_blending = PatchBlending { + // Simple color replace + mode: PatchBlendMode::Replace, + alpha_channel: 0, + clamp: false, + }; + + let ec_blending = [ + PatchBlending { + // EC0 (data) uses EC1 as alpha + mode: PatchBlendMode::BlendAbove, + alpha_channel: 1, // EC1 is alpha for EC0 + clamp: false, + }, + PatchBlending { + // EC1 (alpha) blends itself + mode: PatchBlendMode::BlendAbove, + alpha_channel: 1, // EC1 uses itself as alpha + clamp: false, + }, + ]; + let extra_channel_info = [ + ExtraChannelInfo::new( + false, + ExtraChannel::Unknown, + BitDepth::f32(), + 0, + "ec0".to_string(), + false, + None, + None, + ), // EC0 data + ExtraChannelInfo::new( + false, + ExtraChannel::Alpha, + BitDepth::f32(), + 0, + "alpha_for_ec0".to_string(), + true, // EC1 is premultiplied alpha + None, + None, + ), // EC1 alpha + ]; + + perform_blending( + &mut bg_channels, + &fg_channels, + &color_blending, + &ec_blending, + &extra_channel_info, + ); + + // Expected Color (Replace) + assert_all_almost_abs_eq(&bg_r, &fg_r, ABS_DELTA); + assert_all_almost_abs_eq(&bg_g, &fg_g, ABS_DELTA); + assert_all_almost_abs_eq(&bg_b, &fg_b, ABS_DELTA); + + assert_all_almost_abs_eq(&bg_ec1_alpha, &[expected_out_ec1_alpha], ABS_DELTA); + + assert_all_almost_abs_eq(&bg_ec0, &[expected_out_ec0], ABS_DELTA); + } + + #[test] + fn test_no_alpha_channel_blend_above_falls_back_to_copy_fg() { + let mut bg_r = [0.1]; + let mut bg_g = [0.2]; + let mut bg_b = [0.3]; + let fg_r = [0.7]; + let fg_g = [0.8]; + let fg_b = [0.9]; + + let mut bg_channels: [&mut [f32]; 3] = [&mut bg_r, &mut bg_g, &mut bg_b]; + let fg_channels: [&[f32]; 3] = [&fg_r, &fg_g, &fg_b]; + + let color_blending = PatchBlending { + mode: PatchBlendMode::BlendAbove, + alpha_channel: 0, // Irrelevant as no alpha EIs + clamp: false, + }; + + let ec_blending: [PatchBlending; 0] = []; + // No ExtraChannelInfo means has_alpha will be false. + let extra_channel_info: [ExtraChannelInfo; 0] = []; + + perform_blending( + &mut bg_channels, + &fg_channels, + &color_blending, + &ec_blending, + &extra_channel_info, + ); + + // Expected: output color is fg color due to fallback + assert_all_almost_abs_eq(&bg_r, &fg_r, ABS_DELTA); + assert_all_almost_abs_eq(&bg_g, &fg_g, ABS_DELTA); + assert_all_almost_abs_eq(&bg_b, &fg_b, ABS_DELTA); + } + + #[test] + fn test_empty_pixels() { + let mut bg_r: [f32; 0] = []; + let mut bg_g: [f32; 0] = []; + let mut bg_b: [f32; 0] = []; + let fg_r: [f32; 0] = []; + let fg_g: [f32; 0] = []; + let fg_b: [f32; 0] = []; + + let mut bg_channels: [&mut [f32]; 3] = [&mut bg_r, &mut bg_g, &mut bg_b]; + let fg_channels: [&[f32]; 3] = [&fg_r, &fg_g, &fg_b]; + + let color_blending = PatchBlending { + mode: PatchBlendMode::Replace, + alpha_channel: 0, + clamp: false, + }; + let ec_blending: [PatchBlending; 0] = []; + let extra_channel_info: [ExtraChannelInfo; 0] = []; + + perform_blending( + &mut bg_channels, + &fg_channels, + &color_blending, + &ec_blending, + &extra_channel_info, + ); + + // Expect output slices to also be empty and unchanged. + assert_eq!(bg_r.len(), 0); + assert_eq!(bg_g.len(), 0); + assert_eq!(bg_b.len(), 0); + } + } +} diff --git a/third_party/rust/jxl/src/features/epf.rs b/third_party/rust/jxl/src/features/epf.rs @@ -0,0 +1,74 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +use crate::{ + SIGMA_PADDING, + error::Result, + frame::{HfMetadata, LfGlobalState, transform_map::*}, + headers::frame_header::{Encoding, FrameHeader}, + image::Image, +}; + +pub fn create_sigma_image( + frame_header: &FrameHeader, + lf_global: &LfGlobalState, + hf_meta: &Option<HfMetadata>, +) -> Result<Image<f32>> { + let size_blocks = frame_header.size_blocks(); + let rf = &frame_header.restoration_filter; + let sigma_xsize = size_blocks.0 + 2 * SIGMA_PADDING; + let sigma_ysize = size_blocks.1 + 2 * SIGMA_PADDING; + let mut sigma_image = Image::<f32>::new((sigma_xsize, sigma_ysize))?; + #[allow(clippy::excessive_precision)] + const INV_SIGMA_NUM: f32 = -1.1715728752538099024; + if frame_header.encoding == Encoding::VarDCT { + let hf_meta = hf_meta.as_ref().unwrap(); + let raw_quant_map = hf_meta.raw_quant_map.as_rect(); + let transform_map = hf_meta.transform_map.as_rect(); + let quant_params = lf_global.quant_params.as_ref().unwrap(); + let quant_scale = 1.0 / quant_params.inv_global_scale(); + let epf_map = hf_meta.epf_map.as_rect(); + let mut sigma_rect = sigma_image.as_rect_mut(); + for by in 0..size_blocks.1 { + let sby = SIGMA_PADDING + by; + for bx in 0..size_blocks.0 { + let sbx = SIGMA_PADDING + bx; + let raw_quant = raw_quant_map.row(by)[bx]; + let raw_transform_id = transform_map.row(by)[bx]; + let transform_id = raw_transform_id & 127; + let is_first_block = raw_transform_id >= 128; + if !is_first_block { + continue; + } + let transform_type = HfTransformType::from_usize(transform_id as usize)?; + let cx = covered_blocks_x(transform_type) as usize; + let cy = covered_blocks_y(transform_type) as usize; + let sigma_quant = + rf.epf_quant_mul / (quant_scale * raw_quant as f32 * INV_SIGMA_NUM); + for iy in 0..cy { + for ix in 0..cx { + let sharpness = epf_map.row(by + iy)[bx + ix] as usize; + let sigma = (sigma_quant * rf.epf_sharp_lut[sharpness]).min(-1e-4); + sigma_rect.row(sby + iy)[sbx + ix] = 1.0 / sigma; + } + } + } + sigma_rect.row(sby)[SIGMA_PADDING - 1] = sigma_rect.row(sby)[SIGMA_PADDING]; + sigma_rect.row(sby)[SIGMA_PADDING + size_blocks.0] = + sigma_rect.row(sby)[SIGMA_PADDING + size_blocks.0 - 1]; + } + for bx in 0..sigma_xsize { + sigma_rect.row(SIGMA_PADDING - 1)[bx] = sigma_rect.row(SIGMA_PADDING)[bx]; + sigma_rect.row(SIGMA_PADDING + size_blocks.1)[bx] = + sigma_rect.row(SIGMA_PADDING + size_blocks.1 - 1)[bx]; + } + } else { + // TODO(szabadka): Instead of allocating an image, return an enum with image and f32 + // variants. + let sigma = INV_SIGMA_NUM / rf.epf_sigma_for_modular; + sigma_image.fill(sigma); + } + Ok(sigma_image) +} diff --git a/third_party/rust/jxl/src/features/mod.rs b/third_party/rust/jxl/src/features/mod.rs @@ -0,0 +1,10 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +pub mod blending; +pub mod epf; +pub mod noise; +pub mod patches; +pub mod spline; diff --git a/third_party/rust/jxl/src/features/noise.rs b/third_party/rust/jxl/src/features/noise.rs @@ -0,0 +1,40 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +use crate::{bit_reader::BitReader, error::Result}; +#[derive(Debug, PartialEq, Default, Clone, Copy)] +pub struct Noise { + pub lut: [f32; 8], +} + +impl Noise { + pub fn read(br: &mut BitReader) -> Result<Noise> { + let mut noise = Noise::default(); + for l in &mut noise.lut { + *l = (br.read(10)? as f32) / ((1 << 10) as f32); + } + Ok(noise) + } + pub fn strength(&self, vx: f32) -> f32 { + let k_scale = (self.lut.len() - 2) as f32; + let scaled_vx = f32::max(0.0, vx * k_scale); + let pre_floor_x = scaled_vx.floor(); + let pre_frac_x = scaled_vx - pre_floor_x; + let floor_x = if scaled_vx >= k_scale + 1.0 { + k_scale + } else { + pre_floor_x + }; + let frac_x = if scaled_vx >= k_scale + 1.0 { + 1.0 + } else { + pre_frac_x + }; + let floor_x_int = floor_x as usize; + let low = self.lut[floor_x_int]; + let hi = self.lut[floor_x_int + 1]; + ((hi - low) * frac_x + low).clamp(0.0, 1.0) + } +} diff --git a/third_party/rust/jxl/src/features/patches.rs b/third_party/rust/jxl/src/features/patches.rs @@ -0,0 +1,1960 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +use num_derive::FromPrimitive; +use num_traits::FromPrimitive; + +use crate::{ + bit_reader::BitReader, + entropy_coding::decode::Histograms, + entropy_coding::decode::SymbolReader, + error::{Error, Result}, + features::blending::perform_blending, + frame::{DecoderState, ReferenceFrame}, + headers::extra_channels::ExtraChannelInfo, + util::{NewWithCapacity, slice, tracing_wrappers::*}, +}; + +// Context numbers as specified in Section C.4.5, Listing C.2: +#[derive(Debug, PartialEq, Eq, Clone, Copy)] +#[repr(usize)] +pub enum PatchContext { + NumRefPatch = 0, + ReferenceFrame = 1, + PatchSize = 2, + PatchReferencePosition = 3, + PatchPosition = 4, + PatchBlendMode = 5, + PatchOffset = 6, + PatchCount = 7, + PatchAlphaChannel = 8, + PatchClamp = 9, +} + +impl PatchContext { + const NUM: usize = 10; +} + +/// Blend modes +#[derive(Debug, PartialEq, Eq, Clone, Copy, FromPrimitive)] +#[repr(u8)] +pub enum PatchBlendMode { + // The new values are the old ones. Useful to skip some channels. + None = 0, + // The new values (in the crop) replace the old ones: sample = new + Replace = 1, + // The new values (in the crop) get added to the old ones: sample = old + new + Add = 2, + // The new values (in the crop) get multiplied by the old ones: + // sample = old * new + // This blend mode is only supported if BlendColorSpace is kEncoded. The + // range of the new value matters for multiplication purposes, and its + // nominal range of 0..1 is computed the same way as this is done for the + // alpha values in kBlend and kAlphaWeightedAdd. + Mul = 3, + // The new values (in the crop) replace the old ones if alpha>0: + // For first alpha channel: + // alpha = old + new * (1 - old) + // For other channels if !alpha_associated: + // sample = ((1 - new_alpha) * old * old_alpha + new_alpha * new) / alpha + // For other channels if alpha_associated: + // sample = (1 - new_alpha) * old + new + // The alpha formula applies to the alpha used for the division in the other + // channels formula, and applies to the alpha channel itself if its + // blend_channel value matches itself. + // If using kBlendAbove, new is the patch and old is the original image; if + // using kBlendBelow, the meaning is inverted. + BlendAbove = 4, + BlendBelow = 5, + // The new values (in the crop) are added to the old ones if alpha>0: + // For first alpha channel: sample = sample = old + new * (1 - old) + // For other channels: sample = old + alpha * new + AlphaWeightedAddAbove = 6, + AlphaWeightedAddBelow = 7, +} + +impl PatchBlendMode { + pub const NUM_BLEND_MODES: u8 = 8; + + pub fn uses_alpha(self) -> bool { + matches!( + self, + Self::BlendAbove + | Self::BlendBelow + | Self::AlphaWeightedAddAbove + | Self::AlphaWeightedAddBelow + ) + } + + pub fn uses_clamp(self) -> bool { + self.uses_alpha() || self == Self::Mul + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct PatchBlending { + pub mode: PatchBlendMode, + pub alpha_channel: usize, + pub clamp: bool, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct PatchReferencePosition { + // Not using `ref` like in the spec here, because it is a keyword. + reference: usize, + x0: usize, + y0: usize, + xsize: usize, + ysize: usize, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct PatchPosition { + x: usize, + y: usize, + ref_pos_idx: usize, +} + +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +struct PatchTreeNode { + left_child: isize, + right_child: isize, + y_center: usize, + start: usize, + num: usize, +} + +#[derive(Clone, Debug, Default, PartialEq, Eq)] +pub struct PatchesDictionary { + pub positions: Vec<PatchPosition>, + pub ref_positions: Vec<PatchReferencePosition>, + blendings: Vec<PatchBlending>, + blendings_stride: usize, + patch_tree: Vec<PatchTreeNode>, + // Number of patches for each row. + num_patches: Vec<usize>, + sorted_patches_y0: Vec<(usize, usize)>, + sorted_patches_y1: Vec<(usize, usize)>, +} + +impl PatchesDictionary { + fn compute_patch_tree(&mut self) -> Result<()> { + #[derive(Debug, Clone, Copy)] + struct PatchInterval { + idx: usize, + y0: usize, + y1: usize, + } + + self.patch_tree.clear(); + self.num_patches.clear(); + self.sorted_patches_y0.clear(); + self.sorted_patches_y1.clear(); + + if self.positions.is_empty() { + return Ok(()); + } + + // Create a y-interval for each patch. + let mut intervals: Vec<PatchInterval> = Vec::new_with_capacity(self.positions.len())?; + for (i, pos) in self.positions.iter().enumerate() { + let ref_pos = self.ref_positions[pos.ref_pos_idx]; + if ref_pos.xsize > 0 && ref_pos.ysize > 0 { + intervals.push(PatchInterval { + idx: i, + y0: pos.y, + y1: pos.y + self.ref_positions[pos.ref_pos_idx].ysize, + }); + } + } + + let intervals_len = intervals.len(); + let sort_by_y0 = |intervals: &mut Vec<PatchInterval>, start: usize, end: usize| { + intervals[start..end].sort_unstable_by_key(|i| i.y0); + }; + let sort_by_y1 = |intervals: &mut Vec<PatchInterval>, start: usize, end: usize| { + intervals[start..end].sort_unstable_by_key(|i| i.y1); + }; + + // Count the number of patches for each row. + sort_by_y1(&mut intervals, 0, intervals_len); + self.num_patches + .resize(intervals.last().map_or(0, |iv| iv.y1), 0); //Safe last() + for iv in &intervals { + for y in iv.y0..iv.y1 { + self.num_patches[y] += 1; + } + } + + let root = PatchTreeNode { + start: 0, + num: intervals.len(), + ..Default::default() + }; + self.patch_tree.push(root); + + let mut next = 0; + while next < self.patch_tree.len() { + let node = &mut self.patch_tree[next]; // Borrow mutably *before* accessing fields + let start = node.start; + let end = node.start + node.num; + + // Choose the y_center for this node to be the median of interval starts. + sort_by_y0(&mut intervals, start, end); + let middle_idx = start + node.num / 2; + node.y_center = intervals[middle_idx].y0; + + // Divide the intervals in [start, end) into three groups: + let mut right_start = middle_idx; + while right_start < end && intervals[right_start].y0 == node.y_center { + right_start += 1; + } + + sort_by_y1(&mut intervals, start, right_start); + let mut left_end = right_start; + while left_end > start && intervals[left_end - 1].y1 > node.y_center { + left_end -= 1; + } + + // Fill in sorted_patches_y0_ and sorted_patches_y1_ for the current node. + node.num = right_start - left_end; + node.start = self.sorted_patches_y0.len(); + + self.sorted_patches_y1 + .try_reserve(right_start.saturating_sub(left_end))?; + self.sorted_patches_y0 + .try_reserve(right_start.saturating_sub(left_end))?; + for i in (left_end..right_start).rev() { + self.sorted_patches_y1 + .push((intervals[i].y1, intervals[i].idx)); + } + sort_by_y0(&mut intervals, left_end, right_start); + for interval in intervals.iter().take(right_start).skip(left_end) { + self.sorted_patches_y0.push((interval.y0, interval.idx)); + } + + // Create the left and right nodes (if not empty). + // We modify left_child/right_child on the *original* node in patch_tree, + // so we have to do the assignment *before* we push the new nodes. + self.patch_tree[next].left_child = -1; + self.patch_tree[next].right_child = -1; + + if left_end > start { + let mut left = PatchTreeNode::default(); + left.start = start; + left.num = left_end - left.start; + self.patch_tree[next].left_child = self.patch_tree.len() as isize; + self.patch_tree.try_reserve(1)?; + self.patch_tree.push(left); + } + if right_start < end { + let mut right = PatchTreeNode::default(); + right.start = right_start; + right.num = end - right.start; + self.patch_tree[next].right_child = self.patch_tree.len() as isize; + self.patch_tree.try_reserve(1)?; + self.patch_tree.push(right); + } + + next += 1; + } + Ok(()) + } + + #[instrument(level = "debug", skip(br), ret, err)] + pub fn read( + br: &mut BitReader, + xsize: usize, + ysize: usize, + num_extra_channels: usize, + reference_frames: &[Option<ReferenceFrame>], + ) -> Result<PatchesDictionary> { + let blendings_stride = num_extra_channels + 1; + let patches_histograms = Histograms::decode(PatchContext::NUM, br, true)?; + let mut patches_reader = SymbolReader::new(&patches_histograms, br, None)?; + let num_ref_patch = patches_reader.read_unsigned( + &patches_histograms, + br, + PatchContext::NumRefPatch as usize, + )? as usize; + let num_pixels = xsize * ysize; + let max_ref_patches = 1024 + num_pixels / 4; + let max_patches = max_ref_patches * 4; + let max_blending_infos = max_patches * 4; + if num_ref_patch > max_ref_patches { + return Err(Error::PatchesTooMany( + "reference patches".to_string(), + num_ref_patch, + max_ref_patches, + )); + } + let mut total_patches = 0; + let mut next_size = 1; + let mut positions: Vec<PatchPosition> = Vec::new(); + let mut blendings = Vec::new(); + let mut ref_positions = Vec::new_with_capacity(num_ref_patch)?; + for _ in 0..num_ref_patch { + let reference = patches_reader.read_unsigned( + &patches_histograms, + br, + PatchContext::ReferenceFrame as usize, + )? as usize; + if reference >= DecoderState::MAX_STORED_FRAMES { + return Err(Error::PatchesRefTooLarge( + reference, + DecoderState::MAX_STORED_FRAMES, + )); + } + + let x0 = patches_reader.read_unsigned( + &patches_histograms, + br, + PatchContext::PatchReferencePosition as usize, + )? as usize; + let y0 = patches_reader.read_unsigned( + &patches_histograms, + br, + PatchContext::PatchReferencePosition as usize, + )? as usize; + let ref_pos_xsize = patches_reader.read_unsigned( + &patches_histograms, + br, + PatchContext::PatchSize as usize, + )? as usize + + 1; + let ref_pos_ysize = patches_reader.read_unsigned( + &patches_histograms, + br, + PatchContext::PatchSize as usize, + )? as usize + + 1; + let reference_frame = &reference_frames[reference]; + // TODO(firsching): make sure this check is correct in the presence of downsampled extra channels (also in libjxl). + match reference_frame { + None => return Err(Error::PatchesInvalidReference(reference)), + Some(reference) => { + if !reference.saved_before_color_transform { + return Err(Error::PatchesPostColorTransform()); + } + if x0 + ref_pos_xsize > reference.frame[0].size().0 { + return Err(Error::PatchesInvalidPosition( + "x".to_string(), + x0, + ref_pos_xsize, + reference.frame[0].size().0, + )); + } + if y0 + ref_pos_ysize > reference.frame[0].size().1 { + return Err(Error::PatchesInvalidPosition( + "y".to_string(), + y0, + ref_pos_ysize, + reference.frame[0].size().1, + )); + } + } + } + + let id_count = patches_reader.read_unsigned( + &patches_histograms, + br, + PatchContext::PatchCount as usize, + )? as usize + + 1; + if id_count > max_patches + 1 { + return Err(Error::PatchesTooMany( + "patches".to_string(), + id_count, + max_patches, + )); + } + total_patches += id_count; + + if total_patches > max_patches { + return Err(Error::PatchesTooMany( + "patches".to_string(), + total_patches, + max_patches, + )); + } + + if next_size < total_patches { + next_size *= 2; + next_size = std::cmp::min(next_size, max_patches); + } + if next_size * blendings_stride > max_blending_infos { + return Err(Error::PatchesTooMany( + "blending_info".to_string(), + total_patches, + max_patches, + )); + } + positions.try_reserve(next_size.saturating_sub(positions.len()))?; + blendings.try_reserve( + (next_size * PatchBlendMode::NUM_BLEND_MODES as usize) + .saturating_sub(blendings.len()), + )?; + + for i in 0..id_count { + let mut pos = PatchPosition { + x: 0, + y: 0, + ref_pos_idx: ref_positions.len(), + }; + if i == 0 { + // Read initial position + pos.x = patches_reader.read_unsigned( + &patches_histograms, + br, + PatchContext::PatchPosition as usize, + )? as usize; + pos.y = patches_reader.read_unsigned( + &patches_histograms, + br, + PatchContext::PatchPosition as usize, + )? as usize; + } else { + // Read offsets and calculate new position + let delta_x = patches_reader.read_signed( + &patches_histograms, + br, + PatchContext::PatchOffset as usize, + )?; + if delta_x < 0 && (-delta_x as usize) > positions.last().unwrap().x { + return Err(Error::PatchesInvalidDelta( + "x".to_string(), + positions.last().unwrap().x, + delta_x, + )); + } + pos.x = (positions.last().unwrap().x as i32 + delta_x) as usize; + + let delta_y = patches_reader.read_signed( + &patches_histograms, + br, + PatchContext::PatchOffset as usize, + )?; + if delta_y < 0 && (-delta_y as usize) > positions.last().unwrap().y { + return Err(Error::PatchesInvalidDelta( + "y".to_string(), + positions.last().unwrap().y, + delta_y, + )); + } + pos.y = (positions.last().unwrap().y as i32 + delta_y) as usize; + } + + if pos.x + ref_pos_xsize > xsize { + return Err(Error::PatchesOutOfBounds( + "x".to_string(), + pos.x, + ref_pos_xsize, + xsize, + )); + } + if pos.y + ref_pos_ysize > ysize { + return Err(Error::PatchesOutOfBounds( + "y".to_string(), + pos.y, + ref_pos_ysize, + ysize, + )); + } + + for _ in 0..blendings_stride { + let mut alpha_channel = 0; + let mut clamp = false; + let maybe_blend_mode = patches_reader.read_unsigned( + &patches_histograms, + br, + PatchContext::PatchBlendMode as usize, + )? as u8; + let blend_mode = match PatchBlendMode::from_u8(maybe_blend_mode) { + None => { + return Err(Error::PatchesInvalidBlendMode( + maybe_blend_mode, + PatchBlendMode::NUM_BLEND_MODES, + )); + } + Some(blend_mode) => blend_mode, + }; + + if PatchBlendMode::uses_alpha(blend_mode) && blendings_stride > 2 { + alpha_channel = patches_reader.read_unsigned( + &patches_histograms, + br, + PatchContext::PatchAlphaChannel as usize, + )? as usize; + if alpha_channel >= num_extra_channels { + return Err(Error::PatchesInvalidAlphaChannel( + alpha_channel, + num_extra_channels, + )); + } + } + + if PatchBlendMode::uses_clamp(blend_mode) { + clamp = patches_reader.read_unsigned( + &patches_histograms, + br, + PatchContext::PatchClamp as usize, + )? != 0; + } + blendings.push(PatchBlending { + mode: blend_mode, + alpha_channel, + clamp, + }); + } + positions.push(pos); + } + + ref_positions.push(PatchReferencePosition { + reference, + x0, + y0, + xsize: ref_pos_xsize, + ysize: ref_pos_ysize, + }) + } + + let mut patches_dict = PatchesDictionary { + positions, + blendings, + ref_positions, + blendings_stride, + num_patches: vec![], + sorted_patches_y0: vec![], + sorted_patches_y1: vec![], + patch_tree: vec![], + }; + patches_dict.compute_patch_tree()?; + Ok(patches_dict) + } + + pub fn set_patches_for_row(&self, y: usize, patches_for_row_result: &mut Vec<usize>) { + patches_for_row_result.clear(); + if self.num_patches.len() <= y || self.num_patches[y] == 0 { + return; + } + + let mut tree_idx: isize = 0; + loop { + if tree_idx == -1 { + break; + } + + // Safe access using get() and unwrap_or(). No need for the assert. + let node = self.patch_tree.get(tree_idx as usize).unwrap_or_else(|| { + // TODO(firsching): Handle panic differently? + panic!("Invalid tree_idx: {tree_idx}"); + }); + + if y <= node.y_center { + for i in 0..node.num { + let p = self.sorted_patches_y0[node.start + i]; + if y < p.0 { + break; + } + patches_for_row_result.push(p.1); + } + tree_idx = if y < node.y_center { + node.left_child + } else { + -1 + }; + } else { + for i in 0..node.num { + let p = self.sorted_patches_y1[node.start + i]; + if y >= p.0 { + break; + } + patches_for_row_result.push(p.1); + } + tree_idx = node.right_child; + } + } + + // Ensure that the relative order of patches is preserved. + patches_for_row_result.sort(); + } + + pub fn add_one_row( + &self, + row: &mut [&mut [f32]], + row_pos: (usize, usize), + xsize: usize, + extra_channel_info: &[ExtraChannelInfo], + reference_frames: &[Option<ReferenceFrame>], + patches_for_row_result: &mut Vec<usize>, + ) { + // TODO(zond): Allocate a buffer for this when building the stage instead of when executing it. + let mut out = row + .iter_mut() + .map(|s| &mut s[..xsize]) + .collect::<Vec<&mut [f32]>>(); + let num_ec = extra_channel_info.len(); + assert!(num_ec + 1 == self.blendings_stride); + let dummy_fg = vec![0f32]; + let mut fg = vec![dummy_fg.as_slice(); 3 + num_ec]; + self.set_patches_for_row(row_pos.1, &mut *patches_for_row_result); + for pos_idx in patches_for_row_result.iter() { + let pos = &self.positions[*pos_idx]; + assert!(row_pos.1 >= pos.y); // assert patch starts at or before current row + if pos.x >= row_pos.0 + out[0].len() { + // if patch starts before end of current chunk, continue + continue; + } + + let ref_pos = &self.ref_positions[pos.ref_pos_idx]; + assert!(pos.y + ref_pos.ysize > row_pos.1); // assert patch ends after current row + if pos.x + ref_pos.xsize < row_pos.0 { + // if patch ends before current chunk, continue + continue; + } + + let (ref_x0, out_x0, ref_xsize) = if pos.x < row_pos.0 { + // if patch starts before current chunk + // crop the first part of the patch and use the first part of the chunk + ( + ref_pos.x0 + row_pos.0 - pos.x, + 0, + ref_pos.xsize + pos.x - row_pos.0, + ) + } else { + // otherwise + // use the first part of the patch and crop the first part of the chunk + (ref_pos.x0, pos.x - row_pos.0, ref_pos.xsize) + }; + let (ref_x1, out_x1) = if out[0].len() - out_x0 < ref_xsize { + // if rest of chunk is smaller than patch + // crop the last part of the patch and use the last part of the chunk + (ref_x0 + out[0].len() - out_x0, out[0].len()) + } else { + // otherwise + // use the last part of the patch and crop the last part of the chunk + (ref_x0 + ref_xsize, out_x0 + ref_xsize) + }; + let ref_pos_y = ref_pos.y0 + row_pos.1 - pos.y; + + for (c, fg_ptr) in fg.iter_mut().enumerate().take(3) { + *fg_ptr = &(reference_frames[ref_pos.reference].as_ref().unwrap().frame[c] + .as_rect() + .row(ref_pos_y)[ref_x0..ref_x1]); + } + for i in 0..num_ec { + fg[3 + i] = &(reference_frames[ref_pos.reference].as_ref().unwrap().frame[3 + i] + .as_rect() + .row(ref_pos_y)[ref_x0..ref_x1]); + } + + let blending_idx = pos_idx * self.blendings_stride; + perform_blending( + &mut slice!(&mut out, .., out_x0..out_x1), + &fg, + &self.blendings[blending_idx], + &self.blendings[blending_idx + 1..], + extra_channel_info, + ); + } + } +} + +#[cfg(test)] +mod tests { + + mod read_patches_tests { + use super::super::*; + use test_log::test; + + #[test] + fn read_single_patch_dict() -> Result<()> { + let mut br = BitReader::new(&[0x12, 0x4a, 0x8c, 0x63, 0x13, 0x01, 0xa6, 0x53, 0x01]); + let got_dict = PatchesDictionary::read( + &mut br, + 1024, + 1024, + 0, + &[Some(ReferenceFrame::blank(1024, 1024, 1, true).unwrap())], + )?; + let want_dict = PatchesDictionary { + positions: vec![PatchPosition { + x: 10, + y: 20, + ref_pos_idx: 0, + }], + ref_positions: vec![PatchReferencePosition { + reference: 0, + x0: 0, + y0: 0, + xsize: 1, + ysize: 1, + }], + blendings: vec![PatchBlending { + mode: PatchBlendMode::Add, + alpha_channel: 0, + clamp: false, + }], + blendings_stride: 1, + num_patches: vec![ + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, + ], + patch_tree: vec![PatchTreeNode { + left_child: -1, + right_child: -1, + y_center: 20, + start: 0, + num: 1, + }], + sorted_patches_y0: vec![(20, 0)], + sorted_patches_y1: vec![(21, 0)], + }; + assert_eq!(got_dict, want_dict); + Ok(()) + } + + #[test] + fn read_multi_patch_dict() -> Result<()> { + let mut br = BitReader::new(&[ + 0x12, 0xc6, 0x26, 0x3f, 0x08, 0x4e, 0xb6, 0x0d, 0xf2, 0xde, 0xb6, 0x6d, + ]); + let got_dict = PatchesDictionary::read( + &mut br, + 1024, + 1024, + 2, + &[Some(ReferenceFrame::blank(1024, 1024, 1, true).unwrap())], + )?; + let want_dict = PatchesDictionary { + positions: vec![ + PatchPosition { + x: 0, + y: 0, + ref_pos_idx: 0, + }, + PatchPosition { + x: 5, + y: 5, + ref_pos_idx: 1, + }, + ], + ref_positions: vec![ + PatchReferencePosition { + reference: 0, + x0: 0, + y0: 0, + xsize: 2, + ysize: 1, + }, + PatchReferencePosition { + reference: 0, + x0: 0, + y0: 0, + xsize: 1, + ysize: 2, + }, + ], + blendings: vec![ + PatchBlending { + mode: PatchBlendMode::BlendAbove, + alpha_channel: 1, + clamp: false, + }, + PatchBlending { + mode: PatchBlendMode::Mul, + alpha_channel: 0, + clamp: true, + }, + PatchBlending { + mode: PatchBlendMode::Mul, + alpha_channel: 0, + clamp: true, + }, + PatchBlending { + mode: PatchBlendMode::Mul, + alpha_channel: 0, + clamp: true, + }, + PatchBlending { + mode: PatchBlendMode::Mul, + alpha_channel: 0, + clamp: true, + }, + PatchBlending { + mode: PatchBlendMode::Mul, + alpha_channel: 0, + clamp: true, + }, + ], + blendings_stride: 3, + num_patches: vec![1, 0, 0, 0, 0, 1, 1], + patch_tree: vec![ + PatchTreeNode { + left_child: 1, + right_child: -1, + y_center: 5, + start: 0, + num: 1, + }, + PatchTreeNode { + left_child: -1, + right_child: -1, + y_center: 0, + start: 1, + num: 1, + }, + ], + sorted_patches_y0: vec![(5, 1), (0, 0)], + sorted_patches_y1: vec![(7, 1), (1, 0)], + }; + assert_eq!(got_dict, want_dict); + Ok(()) + } + + #[test] + fn read_large_patch_dict() -> Result<()> { + let mut br = BitReader::new(&[ + 0x12, 0x4e, 0x50, 0x76, 0xeb, 0x41, 0x0d, 0x7e, 0xe5, 0x8e, 0xd2, 0x5d, 0x01, + ]); + let got_dict = PatchesDictionary::read( + &mut br, + 1024, + 1024, + 1, + &[Some(ReferenceFrame::blank(1024, 1024, 1, true).unwrap())], + )?; + let want_dict = PatchesDictionary { + positions: vec![PatchPosition { + x: 2, + y: 3, + ref_pos_idx: 0, + }], + ref_positions: vec![PatchReferencePosition { + reference: 0, + x0: 0, + y0: 0, + xsize: 300, + ysize: 200, + }], + blendings: vec![ + PatchBlending { + mode: PatchBlendMode::AlphaWeightedAddBelow, + alpha_channel: 0, + clamp: false, + }, + PatchBlending { + mode: PatchBlendMode::Mul, + alpha_channel: 0, + clamp: false, + }, + ], + blendings_stride: 2, + num_patches: vec![ + 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, + 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, + 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, + 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, + 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, + 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, + 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, + 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, + ], + patch_tree: vec![PatchTreeNode { + left_child: -1, + right_child: -1, + y_center: 3, + start: 0, + num: 1, + }], + sorted_patches_y0: vec![(3, 0)], + sorted_patches_y1: vec![(203, 0)], + }; + assert_eq!(got_dict, want_dict); + Ok(()) + } + + #[test] + fn read_clamped_patch_dict() -> Result<()> { + let mut br = BitReader::new(&[0x12, 0xc6, 0x26, 0x1f, 0x70, 0xce, 0x06]); + let got_dict = PatchesDictionary::read( + &mut br, + 1024, + 1024, + 0, + &[Some(ReferenceFrame::blank(1024, 1024, 1, true).unwrap())], + )?; + let want_dict = PatchesDictionary { + positions: vec![PatchPosition { + x: 4, + y: 4, + ref_pos_idx: 0, + }], + ref_positions: vec![PatchReferencePosition { + reference: 0, + x0: 0, + y0: 0, + xsize: 1, + ysize: 1, + }], + blendings: vec![PatchBlending { + mode: PatchBlendMode::Mul, + alpha_channel: 0, + clamp: true, + }], + blendings_stride: 1, + num_patches: vec![0, 0, 0, 0, 1], + patch_tree: vec![PatchTreeNode { + left_child: -1, + right_child: -1, + y_center: 4, + start: 0, + num: 1, + }], + sorted_patches_y0: vec![(4, 0)], + sorted_patches_y1: vec![(5, 0)], + }; + assert_eq!(got_dict, want_dict); + Ok(()) + } + + #[test] + fn read_dup_patch_dict() -> Result<()> { + let mut br = BitReader::new(&[0x12, 0x0a, 0x8d, 0x88, 0x03, 0x31, 0xd7, 0x35]); + let got_dict = PatchesDictionary::read( + &mut br, + 1024, + 1024, + 0, + &[Some(ReferenceFrame::blank(1024, 1024, 1, true).unwrap())], + )?; + let want_dict = PatchesDictionary { + positions: vec![ + PatchPosition { + x: 0, + y: 0, + ref_pos_idx: 0, + }, + PatchPosition { + x: 5, + y: 5, + ref_pos_idx: 0, + }, + ], + ref_positions: vec![PatchReferencePosition { + reference: 0, + x0: 0, + y0: 0, + xsize: 1, + ysize: 1, + }], + blendings: vec![ + PatchBlending { + mode: PatchBlendMode::Add, + alpha_channel: 0, + clamp: false, + }, + PatchBlending { + mode: PatchBlendMode::Add, + alpha_channel: 0, + clamp: false, + }, + ], + blendings_stride: 1, + num_patches: vec![1, 0, 0, 0, 0, 1], + patch_tree: vec![ + PatchTreeNode { + left_child: 1, + right_child: -1, + y_center: 5, + start: 0, + num: 1, + }, + PatchTreeNode { + left_child: -1, + right_child: -1, + y_center: 0, + start: 1, + num: 1, + }, + ], + sorted_patches_y0: vec![(5, 1), (0, 0)], + sorted_patches_y1: vec![(6, 1), (1, 0)], + }; + assert_eq!(got_dict, want_dict); + Ok(()) + } + } + + mod set_patches_for_row_tests { + use super::super::*; + use test_log::test; + + // Helper to create a PatchesDictionary for tests + fn create_dictionary( + positions: Vec<PatchPosition>, + ref_positions: Vec<PatchReferencePosition>, + ) -> PatchesDictionary { + // Using default/empty blendings for these tests as they don't affect get_patches_for_row + let mut dict = PatchesDictionary { + positions, + ref_positions, + ..Default::default() + }; + dict.compute_patch_tree().unwrap(); + dict + } + + #[test] + fn test_no_patches() { + let dict = create_dictionary(vec![], vec![]); + let mut patches_for_row_result = vec![]; + dict.set_patches_for_row(0, &mut patches_for_row_result); + assert_eq!(patches_for_row_result, vec![] as Vec<usize>); + dict.set_patches_for_row(10, &mut patches_for_row_result); + assert_eq!(patches_for_row_result, vec![] as Vec<usize>); + } + + #[test] + fn test_single_patch_hit() { + let ref_positions = vec![PatchReferencePosition { + reference: 0, + x0: 0, + y0: 0, + xsize: 10, + ysize: 5, + }]; + let positions = vec![PatchPosition { + x: 0, + y: 10, + ref_pos_idx: 0, + }]; + let dict = create_dictionary(positions, ref_positions); + let mut patches_for_row_result = vec![]; + + // Patch covers rows 10, 11, 12, 13, 14 + dict.set_patches_for_row(10, &mut patches_for_row_result); + assert_eq!(patches_for_row_result, vec![0]); // First row of patch + dict.set_patches_for_row(12, &mut patches_for_row_result); + assert_eq!(patches_for_row_result, vec![0]); // Middle row of patch + dict.set_patches_for_row(14, &mut patches_for_row_result); + assert_eq!(patches_for_row_result, vec![0]); // Last row of patch + } + + #[test] + fn test_single_patch_miss() { + let ref_positions = vec![PatchReferencePosition { + reference: 0, + x0: 0, + y0: 0, + xsize: 10, + ysize: 5, + }]; // Covers y=10 to y=14 + let positions = vec![PatchPosition { + x: 0, + y: 10, + ref_pos_idx: 0, + }]; + let dict = create_dictionary(positions, ref_positions); + let mut patches_for_row_result = vec![]; + + dict.set_patches_for_row(9, &mut patches_for_row_result); + assert_eq!(patches_for_row_result, vec![] as Vec<usize>); // Row before patch + dict.set_patches_for_row(15, &mut patches_for_row_result); + assert_eq!(patches_for_row_result, vec![] as Vec<usize>); // Row after patch + } + + #[test] + fn test_single_patch_height_one() { + let ref_positions = vec![PatchReferencePosition { + reference: 0, + x0: 0, + y0: 0, + xsize: 10, + ysize: 1, + }]; // Covers y=5 only + let positions = vec![PatchPosition { + x: 0, + y: 5, + ref_pos_idx: 0, + }]; + let dict = create_dictionary(positions, ref_positions); + let mut patches_for_row_result = vec![]; + + dict.set_patches_for_row(4, &mut patches_for_row_result); + assert_eq!(patches_for_row_result, vec![] as Vec<usize>); + dict.set_patches_for_row(5, &mut patches_for_row_result); + assert_eq!(patches_for_row_result, vec![0]); + dict.set_patches_for_row(6, &mut patches_for_row_result); + assert_eq!(patches_for_row_result, vec![] as Vec<usize>); + } + + #[test] + fn test_multiple_patches_non_overlapping() { + let ref_positions = vec![ + PatchReferencePosition { + reference: 0, + x0: 0, + y0: 0, + xsize: 10, + ysize: 3, + }, // Patch 0: rows 5,6,7 + PatchReferencePosition { + reference: 0, + x0: 0, + y0: 0, + xsize: 10, + ysize: 2, + }, // Patch 1: rows 10,11 + ]; + let positions = vec![ + PatchPosition { + x: 0, + y: 5, + ref_pos_idx: 0, + }, + PatchPosition { + x: 0, + y: 10, + ref_pos_idx: 1, + }, + ]; + let dict = create_dictionary(positions, ref_positions); + let mut patches_for_row_result = vec![]; + + dict.set_patches_for_row(4, &mut patches_for_row_result); + assert_eq!(patches_for_row_result, vec![] as Vec<usize>); + dict.set_patches_for_row(5, &mut patches_for_row_result); + assert_eq!(patches_for_row_result, vec![0]); + dict.set_patches_for_row(7, &mut patches_for_row_result); + assert_eq!(patches_for_row_result, vec![0]); + dict.set_patches_for_row(8, &mut patches_for_row_result); + assert_eq!(patches_for_row_result, vec![] as Vec<usize>); // Between patches + dict.set_patches_for_row(9, &mut patches_for_row_result); + assert_eq!(patches_for_row_result, vec![] as Vec<usize>); // Between patches + dict.set_patches_for_row(10, &mut patches_for_row_result); + assert_eq!(patches_for_row_result, vec![1]); + dict.set_patches_for_row(11, &mut patches_for_row_result); + assert_eq!(patches_for_row_result, vec![1]); + dict.set_patches_for_row(12, &mut patches_for_row_result); + assert_eq!(patches_for_row_result, vec![] as Vec<usize>); + } + + #[test] + fn test_multiple_patches_overlapping() { + let ref_positions = vec![ + PatchReferencePosition { + reference: 0, + x0: 0, + y0: 0, + xsize: 10, + ysize: 5, + }, // Patch 0: rows 10-14 + PatchReferencePosition { + reference: 0, + x0: 0, + y0: 0, + xsize: 10, + ysize: 4, + }, // Patch 1: rows 12-15 + ]; + let positions = vec![ + PatchPosition { + x: 0, + y: 10, + ref_pos_idx: 0, + }, // idx 0 + PatchPosition { + x: 0, + y: 12, + ref_pos_idx: 1, + }, // idx 1 + ]; + let dict = create_dictionary(positions, ref_positions); + let mut patches_for_row_result = vec![]; + + dict.set_patches_for_row(10, &mut patches_for_row_result); + assert_eq!(patches_for_row_result, vec![0]); // Only patch 0 + dict.set_patches_for_row(11, &mut patches_for_row_result); + assert_eq!(patches_for_row_result, vec![0]); // Only patch 0 + dict.set_patches_for_row(12, &mut patches_for_row_result); + assert_eq!(patches_for_row_result, vec![0, 1]); // Both patches (sorted indices) + dict.set_patches_for_row(13, &mut patches_for_row_result); + assert_eq!(patches_for_row_result, vec![0, 1]); // Both patches + dict.set_patches_for_row(14, &mut patches_for_row_result); + assert_eq!(patches_for_row_result, vec![0, 1]); // Patch 0 ends, Patch 1 continues + dict.set_patches_for_row(15, &mut patches_for_row_result); + assert_eq!(patches_for_row_result, vec![1]); // Only patch 1 + dict.set_patches_for_row(16, &mut patches_for_row_result); + assert_eq!(patches_for_row_result, vec![] as Vec<usize>); + } + + #[test] + fn test_multiple_patches_adjacent() { + let ref_positions = vec![ + PatchReferencePosition { + reference: 0, + x0: 0, + y0: 0, + xsize: 10, + ysize: 2, + }, // Patch 0: rows 5,6 + PatchReferencePosition { + reference: 0, + x0: 0, + y0: 0, + xsize: 10, + ysize: 3, + }, // Patch 1: rows 7,8,9 + ]; + let positions = vec![ + PatchPosition { + x: 0, + y: 5, + ref_pos_idx: 0, + }, + PatchPosition { + x: 0, + y: 7, + ref_pos_idx: 1, + }, + ]; + let dict = create_dictionary(positions, ref_positions); + let mut patches_for_row_result = vec![]; + + dict.set_patches_for_row(4, &mut patches_for_row_result); + assert_eq!(patches_for_row_result, vec![] as Vec<usize>); + dict.set_patches_for_row(5, &mut patches_for_row_result); + assert_eq!(patches_for_row_result, vec![0]); + dict.set_patches_for_row(6, &mut patches_for_row_result); + assert_eq!(patches_for_row_result, vec![0]); + dict.set_patches_for_row(7, &mut patches_for_row_result); + assert_eq!(patches_for_row_result, vec![1]); // Patch 0 ends, Patch 1 starts + dict.set_patches_for_row(8, &mut patches_for_row_result); + assert_eq!(patches_for_row_result, vec![1]); + dict.set_patches_for_row(9, &mut patches_for_row_result); + assert_eq!(patches_for_row_result, vec![1]); + dict.set_patches_for_row(10, &mut patches_for_row_result); + assert_eq!(patches_for_row_result, vec![] as Vec<usize>); + } + + #[test] + fn test_multiple_patches_same_start_different_heights() { + let ref_positions = vec![ + PatchReferencePosition { + reference: 0, + x0: 0, + y0: 0, + xsize: 10, + ysize: 2, + }, // Patch 0: rows 3,4 + PatchReferencePosition { + reference: 0, + x0: 0, + y0: 0, + xsize: 10, + ysize: 4, + }, // Patch 1: rows 3,4,5,6 + ]; + let positions = vec![ + PatchPosition { + x: 0, + y: 3, + ref_pos_idx: 0, + }, // idx 0 + PatchPosition { + x: 0, + y: 3, + ref_pos_idx: 1, + }, // idx 1 + ]; + let dict = create_dictionary(positions, ref_positions); + let mut patches_for_row_result = vec![]; + + dict.set_patches_for_row(2, &mut patches_for_row_result); + assert_eq!(patches_for_row_result, vec![] as Vec<usize>); + dict.set_patches_for_row(3, &mut patches_for_row_result); + assert_eq!(patches_for_row_result, vec![0, 1]); // Both cover + dict.set_patches_for_row(4, &mut patches_for_row_result); + assert_eq!(patches_for_row_result, vec![0, 1]); // Both cover + dict.set_patches_for_row(5, &mut patches_for_row_result); + assert_eq!(patches_for_row_result, vec![1]); // Only patch 1 (longer) + dict.set_patches_for_row(6, &mut patches_for_row_result); + assert_eq!(patches_for_row_result, vec![1]); // Only patch 1 + dict.set_patches_for_row(7, &mut patches_for_row_result); + assert_eq!(patches_for_row_result, vec![] as Vec<usize>); + } + + #[test] + fn test_patches_out_of_order_definition() { + // Define patches in a non-sorted order of their y positions + // get_patches_for_row should still return sorted indices if multiple apply. + let ref_positions = vec![ + PatchReferencePosition { + reference: 0, + x0: 0, + y0: 0, + xsize: 5, + ysize: 3, + }, // Patch 0 (idx 0): rows 10,11,12 + PatchReferencePosition { + reference: 0, + x0: 0, + y0: 0, + xsize: 5, + ysize: 3, + }, // Patch 1 (idx 1): rows 5,6,7 + PatchReferencePosition { + reference: 0, + x0: 0, + y0: 0, + xsize: 5, + ysize: 3, + }, // Patch 2 (idx 2): rows 10,11,12 (overlaps with 0) + ]; + let positions = vec![ + PatchPosition { + x: 0, + y: 10, + ref_pos_idx: 0, + }, // Patch 0 + PatchPosition { + x: 0, + y: 5, + ref_pos_idx: 1, + }, // Patch 1 + PatchPosition { + x: 0, + y: 10, + ref_pos_idx: 2, + }, // Patch 2 + ]; + let dict = create_dictionary(positions, ref_positions); + let mut patches_for_row_result = vec![]; + + dict.set_patches_for_row(4, &mut patches_for_row_result); + assert_eq!(patches_for_row_result, vec![] as Vec<usize>); + dict.set_patches_for_row(5, &mut patches_for_row_result); + assert_eq!(patches_for_row_result, vec![1]); + dict.set_patches_for_row(6, &mut patches_for_row_result); + assert_eq!(patches_for_row_result, vec![1]); + dict.set_patches_for_row(7, &mut patches_for_row_result); + assert_eq!(patches_for_row_result, vec![1]); + dict.set_patches_for_row(8, &mut patches_for_row_result); + assert_eq!(patches_for_row_result, vec![] as Vec<usize>); + dict.set_patches_for_row(9, &mut patches_for_row_result); + assert_eq!(patches_for_row_result, vec![] as Vec<usize>); + dict.set_patches_for_row(10, &mut patches_for_row_result); + assert_eq!(patches_for_row_result, vec![0, 2]); // Patches 0 and 2, indices sorted + dict.set_patches_for_row(11, &mut patches_for_row_result); + assert_eq!(patches_for_row_result, vec![0, 2]); + dict.set_patches_for_row(12, &mut patches_for_row_result); + assert_eq!(patches_for_row_result, vec![0, 2]); + dict.set_patches_for_row(13, &mut patches_for_row_result); + assert_eq!(patches_for_row_result, vec![] as Vec<usize>); + } + } + + mod add_one_row_tests { + use super::super::*; + use crate::{ + headers::{bit_depth::BitDepth, extra_channels::ExtraChannel}, + image::Image, + util::test::assert_all_almost_abs_eq, + }; + + const MAX_ABS_DELTA: f32 = 1e-6; // Adjusted for typical f32 comparisons + + fn create_reference_frame( + width: usize, + height: usize, + channel_data: Vec<Vec<f32>>, + ) -> Result<Option<ReferenceFrame>> { + let mut frame_channels = Vec::new(); + for data_vec in channel_data { + assert_eq!( + data_vec.len(), + width * height, + "Channel data length mismatch" + ); + let img = Image::new_with_data((width, height), data_vec); + frame_channels.push(img); + } + Ok(Some(ReferenceFrame { + frame: frame_channels, + saved_before_color_transform: true, + })) + } + + #[test] + fn test_add_one_row_simple_replace() -> Result<()> { + let xsize = 10; + let num_base_channels = 3; // R, G, B + let const_val = 1.0; + + let ref_frames = vec![create_reference_frame( + xsize, + 1, + vec![vec![const_val; xsize]; num_base_channels], + )?]; + let extra_channel_info: Vec<ExtraChannelInfo> = Vec::new(); + + let ref_positions = vec![PatchReferencePosition { + reference: 0, // Points to main_ref_frame + x0: 2, + y0: 0, + xsize: 3, + ysize: 1, + }]; + let positions = vec![PatchPosition { + x: 2, + y: 0, + ref_pos_idx: 0, + }]; + let blendings = vec![PatchBlending { + mode: PatchBlendMode::Replace, + alpha_channel: 0, + clamp: false, // Clamping set to false + }]; + let mut patches_dict = PatchesDictionary { + positions, + ref_positions, + blendings, + blendings_stride: 1 + extra_channel_info.len(), + patch_tree: Vec::new(), + num_patches: Vec::new(), + sorted_patches_y0: Vec::new(), + sorted_patches_y1: Vec::new(), + }; + patches_dict.compute_patch_tree()?; + + let mut r_data: Vec<f32> = vec![0.0; xsize]; + let mut g_data: Vec<f32> = vec![0.0; xsize]; + let mut b_data: Vec<f32> = vec![0.0; xsize]; + let mut row_slices: Vec<&mut [f32]> = vec![&mut r_data, &mut g_data, &mut b_data]; + + let expected_r = vec![0.0, 0.0, 1.0, 1.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0]; + + patches_dict.add_one_row( + &mut row_slices, + (0, 0), + xsize, + &extra_channel_info, + &ref_frames, // Pass the Vec<ReferenceFrame> + &mut vec![], + ); + + assert_all_almost_abs_eq(&r_data, &expected_r, MAX_ABS_DELTA); + assert_all_almost_abs_eq(&g_data, &expected_r, MAX_ABS_DELTA); + assert_all_almost_abs_eq(&b_data, &expected_r, MAX_ABS_DELTA); + Ok(()) + } + + #[test] + fn test_add_one_row_simple_add() -> Result<()> { + let xsize = 10; + let y_coord = 0; + let num_base_channels = 3; + let const_val = 0.2; + + let ref_frames = vec![create_reference_frame( + xsize, + 1, + vec![vec![const_val; xsize]; num_base_channels], + )?]; + let extra_channel_info: Vec<ExtraChannelInfo> = Vec::new(); + + let ref_positions = vec![PatchReferencePosition { + reference: 0, + x0: 2, + y0: 0, + xsize: 3, + ysize: 1, + }]; + let positions = vec![PatchPosition { + x: 2, + y: 0, + ref_pos_idx: 0, + }]; + let blendings = vec![PatchBlending { + mode: PatchBlendMode::Add, + alpha_channel: 0, + clamp: false, // Clamping set to false + }]; + let mut patches_dict = PatchesDictionary { + positions, + ref_positions, + blendings, + blendings_stride: 1 + extra_channel_info.len(), + patch_tree: Vec::new(), + num_patches: Vec::new(), + sorted_patches_y0: Vec::new(), + sorted_patches_y1: Vec::new(), + }; + patches_dict.compute_patch_tree()?; + + let mut r_data: Vec<f32> = vec![0.5; xsize]; + let mut g_data: Vec<f32> = vec![0.5; xsize]; + let mut b_data: Vec<f32> = vec![0.5; xsize]; + let mut row_slices: Vec<&mut [f32]> = vec![&mut r_data, &mut g_data, &mut b_data]; + + let mut expected_r: Vec<f32> = vec![0.5; xsize]; + for r in expected_r.iter_mut().take(5).skip(2) { + *r = 0.5 + 0.2 + } + + patches_dict.add_one_row( + &mut row_slices, + (0, y_coord), + xsize, + &extra_channel_info, + &ref_frames, + &mut vec![], + ); + + assert_all_almost_abs_eq(&r_data, &expected_r, MAX_ABS_DELTA); + assert_all_almost_abs_eq(&g_data, &expected_r, MAX_ABS_DELTA); + assert_all_almost_abs_eq(&b_data, &expected_r, MAX_ABS_DELTA); + Ok(()) + } + + #[test] + fn test_add_one_row_overlapping_replace() -> Result<()> { + let xsize = 10; + let y_coord = 0; + let num_base_channels = 3; + + let main_ref_frame1 = + create_reference_frame(xsize, 1, vec![vec![1.0; xsize]; num_base_channels])?; + let main_ref_frame2 = + create_reference_frame(xsize, 1, vec![vec![2.0; xsize]; num_base_channels])?; + + let ref_frames = vec![main_ref_frame1, main_ref_frame2]; + let extra_channel_info: Vec<ExtraChannelInfo> = Vec::new(); + + let ref_positions = vec![ + PatchReferencePosition { + reference: 0, // Points to main_ref_frame1 + x0: 0, + y0: 0, + xsize: 4, + ysize: 1, + }, + PatchReferencePosition { + reference: 1, // Points to main_ref_frame2 + x0: 0, + y0: 0, + xsize: 3, + ysize: 1, + }, + ]; + let positions = vec![ + PatchPosition { + x: 2, + y: 0, + ref_pos_idx: 0, + }, // P1: canvas [2..6] with 1.0 + PatchPosition { + x: 4, + y: 0, + ref_pos_idx: 1, + }, // P2: canvas [4..7] with 2.0 + ]; + let blendings = vec![ + PatchBlending { + mode: PatchBlendMode::Replace, + alpha_channel: 0, + clamp: false, // Clamping set to false + }, // For P1 + PatchBlending { + mode: PatchBlendMode::Replace, + alpha_channel: 0, + clamp: false, // Clamping set to false + }, // For P2 + ]; + let mut patches_dict = PatchesDictionary { + positions, + ref_positions, + blendings, + blendings_stride: 1 + extra_channel_info.len(), + patch_tree: Vec::new(), + num_patches: Vec::new(), + sorted_patches_y0: Vec::new(), + sorted_patches_y1: Vec::new(), + }; + patches_dict.compute_patch_tree()?; + + let mut r_data: Vec<f32> = vec![0.0; xsize]; + let mut g_data: Vec<f32> = vec![0.0; xsize]; + let mut b_data: Vec<f32> = vec![0.0; xsize]; + let mut row_slices: Vec<&mut [f32]> = vec![&mut r_data, &mut g_data, &mut b_data]; + + let expected_r: Vec<f32> = vec![0.0, 0.0, 1.0, 1.0, 2.0, 2.0, 2.0, 0.0, 0.0, 0.0]; + + patches_dict.add_one_row( + &mut row_slices, + (0, y_coord), + xsize, + &extra_channel_info, + &ref_frames, + &mut vec![], + ); + + assert_all_almost_abs_eq(&r_data, &expected_r, MAX_ABS_DELTA); + assert_all_almost_abs_eq(&g_data, &expected_r, MAX_ABS_DELTA); + assert_all_almost_abs_eq(&b_data, &expected_r, MAX_ABS_DELTA); + Ok(()) + } + + #[test] + fn test_add_one_row_blend_above_ec_alpha_non_associated() -> Result<()> { + let xsize = 1; + let y_coord = 0; + + let initial_color_val = 0.1; + let initial_ec0_alpha = 0.4; + let ref_color_val = 0.8; + let ref_ec0_alpha_val = 0.5; + + let ec_info = vec![ExtraChannelInfo::new( + true, + ExtraChannel::Alpha, + BitDepth::f32(), + 0, + "AlphaEC".to_string(), + false, // alpha_associated = false + None, + None, + )]; + + let main_ref_frame_data = vec![ + vec![ref_color_val; xsize], // R + vec![ref_color_val; xsize], // G + vec![ref_color_val; xsize], // B + vec![ref_ec0_alpha_val; xsize], // EC0 (Alpha) + ]; + let main_ref_frame = create_reference_frame(xsize, 1, main_ref_frame_data)?; + + let ref_frames = vec![main_ref_frame]; + + let ref_positions = vec![PatchReferencePosition { + reference: 0, + x0: 0, + y0: 0, + xsize: 1, + ysize: 1, + }]; + let positions = vec![PatchPosition { + x: 0, + y: 0, + ref_pos_idx: 0, + }]; + let blendings = vec![ + PatchBlending { + mode: PatchBlendMode::BlendAbove, + alpha_channel: 0, // Alpha for color is EC0 + clamp: false, // Clamping set to false + }, // Color + PatchBlending { + mode: PatchBlendMode::BlendAbove, + alpha_channel: 0, // Alpha for EC0 is EC0 itself + clamp: false, // Clamping set to false + }, // EC0 + ]; + let mut patches_dict = PatchesDictionary { + positions, + ref_positions, + blendings, + blendings_stride: 1 + ec_info.len(), // Color + 1 EC + patch_tree: Vec::new(), + num_patches: Vec::new(), + sorted_patches_y0: Vec::new(), + sorted_patches_y1: Vec::new(), + }; + patches_dict.compute_patch_tree()?; + + let mut r_data = vec![initial_color_val; xsize]; + let mut g_data = vec![initial_color_val; xsize]; + let mut b_data = vec![initial_color_val; xsize]; + let mut ec0_data = vec![initial_ec0_alpha; xsize]; + let mut row_slices: Vec<&mut [f32]> = + vec![&mut r_data, &mut g_data, &mut b_data, &mut ec0_data]; + + // Calculations based on C++ logic for non-associated alpha: + // OutputAlpha = OldAlpha + PatchAlpha * (1 - OldAlpha) + // OutputColor = (OldColor * OldAlpha * (1 - PatchAlpha) + PatchColor * PatchAlpha) / OutputAlpha + // (If OutputAlpha is very small, OutputColor is 0) + let canvas_alpha_val = initial_ec0_alpha; // old_alpha + let patch_alpha_val = ref_ec0_alpha_val; // ref_alpha + + let expected_ec0 = canvas_alpha_val + patch_alpha_val * (1.0 - canvas_alpha_val); + + let canvas_color_val = initial_color_val; // old_color + let patch_color_val = ref_color_val; // ref_color (straight) + + let expected_color = if expected_ec0.abs() < 1e-5 { + // Threshold similar to kSmallAlpha + 0.0 + } else { + (canvas_color_val * canvas_alpha_val * (1.0 - patch_alpha_val) + + patch_color_val * patch_alpha_val) + / expected_ec0 + }; + + patches_dict.add_one_row( + &mut row_slices, + (0, y_coord), + xsize, + &ec_info, + &ref_frames, + &mut vec![], + ); + + assert_all_almost_abs_eq(&r_data, &vec![expected_color], MAX_ABS_DELTA); + assert_all_almost_abs_eq(&g_data, &vec![expected_color], MAX_ABS_DELTA); + assert_all_almost_abs_eq(&b_data, &vec![expected_color], MAX_ABS_DELTA); + assert_all_almost_abs_eq(&ec0_data, &vec![expected_ec0], MAX_ABS_DELTA); + Ok(()) + } + + #[test] + fn test_add_one_row_blend_above_ec_alpha_associated() -> Result<()> { + let xsize = 1; + let y_coord = 0; + + let initial_color_val = 0.1; // Canvas color, assumed associated + let initial_ec0_alpha = 0.4; // Canvas alpha + let ref_color_val = 0.8; // Patch color, straight (non-associated) + let ref_ec0_alpha_val = 0.5; // Patch alpha + + let ec_info = vec![ExtraChannelInfo::new( + true, + ExtraChannel::Alpha, + BitDepth::f32(), + 0, + "AlphaEC".to_string(), + true, // alpha_associated = true + None, + None, + )]; + + let main_ref_frame_data = vec![ + vec![ref_color_val; xsize], // R + vec![ref_color_val; xsize], // G + vec![ref_color_val; xsize], // B + vec![ref_ec0_alpha_val; xsize], // EC0 (Alpha) + ]; + let main_ref_frame = create_reference_frame(xsize, 1, main_ref_frame_data)?; + + let ref_frames = vec![main_ref_frame]; + + let ref_positions = vec![PatchReferencePosition { + reference: 0, + x0: 0, + y0: 0, + xsize: 1, + ysize: 1, + }]; + let positions = vec![PatchPosition { + x: 0, + y: 0, + ref_pos_idx: 0, + }]; + let blendings = vec![ + PatchBlending { + mode: PatchBlendMode::BlendAbove, + alpha_channel: 0, + clamp: false, // Clamping set to false + }, // Color + PatchBlending { + mode: PatchBlendMode::BlendAbove, + alpha_channel: 0, + clamp: false, // Clamping set to false + }, // EC0 + ]; + let mut patches_dict = PatchesDictionary { + positions, + ref_positions, + blendings, + blendings_stride: 1 + ec_info.len(), + patch_tree: Vec::new(), + num_patches: Vec::new(), + sorted_patches_y0: Vec::new(), + sorted_patches_y1: Vec::new(), + }; + patches_dict.compute_patch_tree()?; + + let mut r_data = vec![initial_color_val; xsize]; + let mut g_data = vec![initial_color_val; xsize]; + let mut b_data = vec![initial_color_val; xsize]; + let mut ec0_data = vec![initial_ec0_alpha; xsize]; + let mut row_slices: Vec<&mut [f32]> = + vec![&mut r_data, &mut g_data, &mut b_data, &mut ec0_data]; + + let expected_ec0 = ref_ec0_alpha_val + initial_ec0_alpha * (1.0 - ref_ec0_alpha_val); + + let expected_color = ref_color_val + initial_color_val * (1.0 - ref_ec0_alpha_val); + + patches_dict.add_one_row( + &mut row_slices, + (0, y_coord), + xsize, + &ec_info, + &ref_frames, + &mut vec![], + ); + + assert_all_almost_abs_eq(&ec0_data, &vec![expected_ec0], MAX_ABS_DELTA); + assert_all_almost_abs_eq(&r_data, &vec![expected_color], MAX_ABS_DELTA); + assert_all_almost_abs_eq(&g_data, &vec![expected_color], MAX_ABS_DELTA); + assert_all_almost_abs_eq(&b_data, &vec![expected_color], MAX_ABS_DELTA); + Ok(()) + } + + #[test] + fn test_add_one_row_mul_blend() -> Result<()> { + let xsize = 2; + let y_coord = 0; + let num_base_channels = 3; + + let initial_vals = vec![0.5, 2.0]; + let ref_vals = vec![0.8, 0.7]; + + let main_ref_channel_data: Vec<Vec<f32>> = (0..num_base_channels) + .map(|_| ref_vals.clone()) // Each color channel gets ref_vals + .collect(); + let main_ref_frame = create_reference_frame(xsize, 1, main_ref_channel_data)?; + + let dummy_channel_data: Vec<Vec<f32>> = + (0..num_base_channels).map(|_| vec![0.0; xsize]).collect(); + let dummy_ref_frame1 = create_reference_frame(xsize, 1, dummy_channel_data.clone())?; + let dummy_ref_frame2 = create_reference_frame(xsize, 1, dummy_channel_data)?; + + let ref_frames = vec![main_ref_frame, dummy_ref_frame1, dummy_ref_frame2]; + let extra_channel_info: Vec<ExtraChannelInfo> = Vec::new(); + + let ref_positions = vec![PatchReferencePosition { + reference: 0, + x0: 0, + y0: 0, + xsize, + ysize: 1, + }]; + let positions = vec![PatchPosition { + x: 0, + y: 0, + ref_pos_idx: 0, + }]; + + // Test Mul (always without clamp as per new instruction) + let blendings = vec![PatchBlending { + mode: PatchBlendMode::Mul, + alpha_channel: 0, + clamp: false, // Clamping set to false + }]; + let mut dict = PatchesDictionary { + positions: positions.clone(), + ref_positions: ref_positions.clone(), + blendings, + blendings_stride: 1, // Only color channels + patch_tree: Vec::new(), + num_patches: Vec::new(), + sorted_patches_y0: Vec::new(), + sorted_patches_y1: Vec::new(), + }; + dict.compute_patch_tree()?; + + let mut r_data = initial_vals.clone(); + let mut g_data = initial_vals.clone(); + let mut b_data = initial_vals.clone(); + let mut slices: Vec<&mut [f32]> = vec![&mut r_data, &mut g_data, &mut b_data]; + dict.add_one_row( + &mut slices, + (0, y_coord), + xsize, + &extra_channel_info, + &ref_frames, + &mut vec![], + ); + + let expected_vals = vec![0.5 * 0.8, 2.0 * 0.7]; // [0.4, 1.4] + assert_all_almost_abs_eq(&r_data, &expected_vals, MAX_ABS_DELTA); + assert_all_almost_abs_eq(&g_data, &expected_vals, MAX_ABS_DELTA); + assert_all_almost_abs_eq(&b_data, &expected_vals, MAX_ABS_DELTA); + + Ok(()) + } + + #[test] + fn test_add_one_row_none_blend() -> Result<()> { + let xsize = 5; + let y_coord = 0; + let num_base_channels = 3; + let const_val = 100.0; + + let main_channel_data: Vec<Vec<f32>> = (0..num_base_channels) + .map(|_| vec![const_val; xsize]) + .collect(); + let main_ref_frame = create_reference_frame(xsize, 1, main_channel_data)?; + + let dummy_channel_data: Vec<Vec<f32>> = + (0..num_base_channels).map(|_| vec![0.0; xsize]).collect(); + let dummy_ref_frame1 = create_reference_frame(xsize, 1, dummy_channel_data.clone())?; + let dummy_ref_frame2 = create_reference_frame(xsize, 1, dummy_channel_data)?; + + let ref_frames = vec![main_ref_frame, dummy_ref_frame1, dummy_ref_frame2]; + let extra_channel_info: Vec<ExtraChannelInfo> = Vec::new(); + + let ref_positions = vec![PatchReferencePosition { + reference: 0, + x0: 0, + y0: 0, + xsize: 3, + ysize: 1, + }]; + let positions = vec![PatchPosition { + x: 1, + y: 0, + ref_pos_idx: 0, + }]; + let blendings = vec![PatchBlending { + mode: PatchBlendMode::None, + alpha_channel: 0, + clamp: false, // Clamping set to false + }]; + + let mut patches_dict = PatchesDictionary { + positions, + ref_positions, + blendings, + blendings_stride: 1, // Only color channels + patch_tree: Vec::new(), + num_patches: Vec::new(), + sorted_patches_y0: Vec::new(), + sorted_patches_y1: Vec::new(), + }; + patches_dict.compute_patch_tree()?; + + let initial_data: Vec<f32> = (0..xsize).map(|i| i as f32 * 0.1 + 0.05).collect(); + let mut r_data = initial_data.clone(); + let mut g_data = initial_data.clone(); + let mut b_data = initial_data.clone(); + let mut row_slices: Vec<&mut [f32]> = vec![&mut r_data, &mut g_data, &mut b_data]; + + patches_dict.add_one_row( + &mut row_slices, + (0, y_coord), + xsize, + &extra_channel_info, + &ref_frames, + &mut vec![], + ); + + assert_all_almost_abs_eq(&r_data, &initial_data, MAX_ABS_DELTA); + assert_all_almost_abs_eq(&g_data, &initial_data, MAX_ABS_DELTA); + assert_all_almost_abs_eq(&b_data, &initial_data, MAX_ABS_DELTA); + Ok(()) + } + } +} diff --git a/third_party/rust/jxl/src/features/spline.rs b/third_party/rust/jxl/src/features/spline.rs @@ -0,0 +1,1884 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +use std::{ + f32::consts::{FRAC_1_SQRT_2, PI, SQRT_2}, + iter::{self, zip}, + ops, +}; + +use crate::{ + bit_reader::BitReader, + entropy_coding::decode::{Histograms, SymbolReader, unpack_signed}, + error::{Error, Result}, + frame::color_correlation_map::ColorCorrelationParams, + util::{CeilLog2, NewWithCapacity, fast_cos, fast_erff, tracing_wrappers::*}, +}; +const MAX_NUM_CONTROL_POINTS: u32 = 1 << 20; +const MAX_NUM_CONTROL_POINTS_PER_PIXEL_RATIO: u32 = 2; +const DELTA_LIMIT: i64 = 1 << 30; +const SPLINE_POS_LIMIT: i32 = 1 << 23; + +const QUANTIZATION_ADJUSTMENT_CONTEXT: usize = 0; +const STARTING_POSITION_CONTEXT: usize = 1; +const NUM_SPLINES_CONTEXT: usize = 2; +const NUM_CONTROL_POINTS_CONTEXT: usize = 3; +const CONTROL_POINTS_CONTEXT: usize = 4; +const DCT_CONTEXT: usize = 5; +const NUM_SPLINE_CONTEXTS: usize = 6; +const DESIRED_RENDERING_DISTANCE: f32 = 1.0; + +#[derive(Debug, Clone, Copy, Default)] +pub struct Point { + pub x: f32, + pub y: f32, +} + +impl Point { + fn new(x: f32, y: f32) -> Self { + Point { x, y } + } + fn abs(&self) -> f32 { + self.x.hypot(self.y) + } +} + +impl PartialEq for Point { + fn eq(&self, other: &Self) -> bool { + (self.x - other.x).abs() < 1e-3 && (self.y - other.y).abs() < 1e-3 + } +} + +impl ops::Add<Point> for Point { + type Output = Point; + fn add(self, rhs: Point) -> Point { + Point { + x: self.x + rhs.x, + y: self.y + rhs.y, + } + } +} + +impl ops::Sub<Point> for Point { + type Output = Point; + fn sub(self, rhs: Point) -> Point { + Point { + x: self.x - rhs.x, + y: self.y - rhs.y, + } + } +} + +impl ops::Mul<f32> for Point { + type Output = Point; + fn mul(self, rhs: f32) -> Point { + Point { + x: self.x * rhs, + y: self.y * rhs, + } + } +} + +impl ops::Div<f32> for Point { + type Output = Point; + fn div(self, rhs: f32) -> Point { + let inv = 1.0 / rhs; + Point { + x: self.x * inv, + y: self.y * inv, + } + } +} + +#[derive(Default, Debug)] +pub struct Spline { + control_points: Vec<Point>, + // X, Y, B. + color_dct: [Dct32; 3], + // Splines are drawn by normalized Gaussian splatting. This controls the + // Gaussian's parameter along the spline. + sigma_dct: Dct32, + // The estimated area in pixels covered by the spline. + estimated_area_reached: u64, +} + +impl Spline { + pub fn validate_adjacent_point_coincidence(&self) -> Result<()> { + if let Some(((index, p0), p1)) = zip( + self.control_points + .iter() + .take(self.control_points.len() - 1) + .enumerate(), + self.control_points.iter().skip(1), + ) + .find(|((_, p0), p1)| **p0 == **p1) + { + return Err(Error::SplineAdjacentCoincidingControlPoints( + index, + *p0, + index + 1, + *p1, + )); + } + Ok(()) + } +} + +#[derive(Debug, Default, Clone)] +pub struct QuantizedSpline { + // Double delta-encoded. + pub control_points: Vec<(i64, i64)>, + pub color_dct: [[i32; 32]; 3], + pub sigma_dct: [i32; 32], +} + +fn inv_adjusted_quant(adjustment: i32) -> f32 { + if adjustment >= 0 { + 1.0 / (1.0 + 0.125 * adjustment as f32) + } else { + 1.0 - 0.125 * adjustment as f32 + } +} + +fn validate_spline_point_pos<T: num_traits::ToPrimitive>(x: T, y: T) -> Result<()> { + let xi = x.to_i32().unwrap(); + let yi = y.to_i32().unwrap(); + let ok_range = -(1i32 << 23)..(1i32 << 23); + if !ok_range.contains(&xi) { + return Err(Error::SplinesPointOutOfRange( + Point { + x: xi as f32, + y: yi as f32, + }, + xi, + ok_range, + )); + } + if !ok_range.contains(&yi) { + return Err(Error::SplinesPointOutOfRange( + Point { + x: xi as f32, + y: yi as f32, + }, + yi, + ok_range, + )); + } + Ok(()) +} + +const CHANNEL_WEIGHT: [f32; 4] = [0.0042, 0.075, 0.07, 0.3333]; + +fn area_limit(image_size: u64) -> u64 { + (1024 * image_size + (1u64 << 32)).min(1u64 << 42) +} + +impl QuantizedSpline { + #[instrument(level = "debug", skip(br), ret, err)] + pub fn read( + br: &mut BitReader, + splines_histograms: &Histograms, + splines_reader: &mut SymbolReader, + max_control_points: u32, + total_num_control_points: &mut u32, + ) -> Result<QuantizedSpline> { + let num_control_points = + splines_reader.read_unsigned(splines_histograms, br, NUM_CONTROL_POINTS_CONTEXT)?; + *total_num_control_points += num_control_points; + if *total_num_control_points > max_control_points { + return Err(Error::SplinesTooManyControlPoints( + *total_num_control_points, + max_control_points, + )); + } + let mut control_points = Vec::new_with_capacity(num_control_points as usize)?; + for _ in 0..num_control_points { + let x = + splines_reader.read_signed(splines_histograms, br, CONTROL_POINTS_CONTEXT)? as i64; + let y = + splines_reader.read_signed(splines_histograms, br, CONTROL_POINTS_CONTEXT)? as i64; + control_points.push((x, y)); + // Add check that double deltas are not outrageous (not in spec). + let max_delta_delta = x.abs().max(y.abs()); + if max_delta_delta >= DELTA_LIMIT { + return Err(Error::SplinesDeltaLimit(max_delta_delta, DELTA_LIMIT)); + } + } + // Decode DCTs and populate the QuantizedSpline struct + let mut color_dct = [[0; 32]; 3]; + let mut sigma_dct = [0; 32]; + + let mut decode_dct = |dct: &mut [i32; 32]| -> Result<()> { + for value in dct.iter_mut() { + *value = splines_reader.read_signed(splines_histograms, br, DCT_CONTEXT)?; + } + Ok(()) + }; + + for channel in &mut color_dct { + decode_dct(channel)?; + } + decode_dct(&mut sigma_dct)?; + + Ok(QuantizedSpline { + control_points, + color_dct, + sigma_dct, + }) + } + + pub fn dequantize( + &self, + starting_point: &Point, + quantization_adjustment: i32, + y_to_x: f32, + y_to_b: f32, + image_size: u64, + ) -> Result<Spline> { + let area_limit = area_limit(image_size); + + let mut result = Spline { + control_points: Vec::new_with_capacity(self.control_points.len() + 1)?, + ..Default::default() + }; + + let px = starting_point.x.round(); + let py = starting_point.y.round(); + validate_spline_point_pos(px, py)?; + + let mut current_x = px as i32; + let mut current_y = py as i32; + result + .control_points + .push(Point::new(current_x as f32, current_y as f32)); + + let mut current_delta_x = 0i32; + let mut current_delta_y = 0i32; + let mut manhattan_distance = 0u64; + + for &(dx, dy) in &self.control_points { + current_delta_x += dx as i32; + current_delta_y += dy as i32; + validate_spline_point_pos(current_delta_x, current_delta_y)?; + + manhattan_distance += + current_delta_x.unsigned_abs() as u64 + current_delta_y.unsigned_abs() as u64; + + if manhattan_distance > area_limit { + return Err(Error::SplinesDistanceTooLarge( + manhattan_distance, + area_limit, + )); + } + + current_x += current_delta_x; + current_y += current_delta_y; + validate_spline_point_pos(current_x, current_y)?; + + result + .control_points + .push(Point::new(current_x as f32, current_y as f32)); + } + + let inv_quant = inv_adjusted_quant(quantization_adjustment); + + for (c, weight) in CHANNEL_WEIGHT.iter().enumerate().take(3) { + for i in 0..32 { + let inv_dct_factor = if i == 0 { FRAC_1_SQRT_2 } else { 1.0 }; + result.color_dct[c].0[i] = + self.color_dct[c][i] as f32 * inv_dct_factor * weight * inv_quant; + } + } + + for i in 0..32 { + result.color_dct[0].0[i] += y_to_x * result.color_dct[1].0[i]; + result.color_dct[2].0[i] += y_to_b * result.color_dct[1].0[i]; + } + + let mut width_estimate = 0; + let mut color = [0u64; 3]; + + for (c, color_val) in color.iter_mut().enumerate() { + for i in 0..32 { + *color_val += (inv_quant * self.color_dct[c][i].abs() as f32).ceil() as u64; + } + } + + color[0] += y_to_x.abs().ceil() as u64 * color[1]; + color[2] += y_to_b.abs().ceil() as u64 * color[1]; + + let max_color = color[0].max(color[1]).max(color[2]); + let logcolor = 1u64.max((1u64 + max_color).ceil_log2()); + + let weight_limit = + (((area_limit as f32 / logcolor as f32) / manhattan_distance.max(1) as f32).sqrt()) + .ceil(); + + for i in 0..32 { + let inv_dct_factor = if i == 0 { FRAC_1_SQRT_2 } else { 1.0 }; + result.sigma_dct.0[i] = + self.sigma_dct[i] as f32 * inv_dct_factor * CHANNEL_WEIGHT[3] * inv_quant; + + let weight_f = (inv_quant * self.sigma_dct[i].abs() as f32).ceil(); + let weight = weight_limit.min(weight_f.max(1.0)) as u64; + width_estimate += weight * weight * logcolor; + } + + result.estimated_area_reached = width_estimate * manhattan_distance; + + Ok(result) + } +} + +#[derive(Debug, Clone, Copy, Default)] +struct SplineSegment { + center_x: f32, + center_y: f32, + maximum_distance: f32, + inv_sigma: f32, + sigma_over_4_times_intensity: f32, + color: [f32; 3], +} + +#[derive(Debug, Default, Clone)] +pub struct Splines { + pub quantization_adjustment: i32, + pub splines: Vec<QuantizedSpline>, + pub starting_points: Vec<Point>, + segments: Vec<SplineSegment>, + segment_indices: Vec<usize>, + segment_y_start: Vec<u64>, +} + +fn draw_centripetal_catmull_rom_spline(points: &[Point]) -> Result<Vec<Point>> { + if points.is_empty() { + return Ok(vec![]); + } + if points.len() == 1 { + return Ok(vec![points[0]]); + } + const NUM_POINTS: usize = 16; + // Create a view of points with one prepended and one appended point. + let extended_points = iter::once(points[0] + (points[0] - points[1])) + .chain(points.iter().cloned()) + .chain(iter::once( + points[points.len() - 1] + (points[points.len() - 1] - points[points.len() - 2]), + )); + // Pair each point with the sqrt of the distance to the next point. + let points_and_deltas = extended_points + .chain(iter::once(Point::default())) + .scan(Point::default(), |previous, p| { + let result = Some((*previous, (p - *previous).abs().sqrt())); + *previous = p; + result + }) + .skip(1); + // Window the points with a [Point; 4] window. + let windowed_points = points_and_deltas + .scan([(Point::default(), 0.0); 4], |window, p| { + (window[0], window[1], window[2], window[3]) = + (window[1], window[2], window[3], (p.0, p.1)); + Some([window[0], window[1], window[2], window[3]]) + }) + .skip(3); + // Create the points necessary per window, and flatten the result. + let result = windowed_points + .flat_map(|p| { + let mut window_result = [Point::default(); NUM_POINTS]; + window_result[0] = p[1].0; + let mut t = [0.0; 4]; + for k in 0..3 { + // TODO(from libjxl): Restrict d[k] with reasonable limit and spec it. + t[k + 1] = t[k] + p[k].1; + } + for (i, window_point) in window_result.iter_mut().enumerate().skip(1) { + let tt = p[0].1 + ((i as f32) / (NUM_POINTS as f32)) * p[1].1; + let mut a = [Point::default(); 3]; + for k in 0..3 { + // TODO(from libjxl): Reciprocal multiplication would be faster. + a[k] = p[k].0 + (p[k + 1].0 - p[k].0) * ((tt - t[k]) / p[k].1); + } + let mut b = [Point::default(); 2]; + for k in 0..2 { + b[k] = a[k] + (a[k + 1] - a[k]) * ((tt - t[k]) / (p[k].1 + p[k + 1].1)); + } + *window_point = b[0] + (b[1] - b[0]) * ((tt - t[1]) / p[1].1); + } + window_result + }) + .chain(iter::once(points[points.len() - 1])) + .collect(); + Ok(result) +} + +fn for_each_equally_spaced_point<F: FnMut(Point, f32)>( + points: &[Point], + desired_distance: f32, + mut f: F, +) { + if points.is_empty() { + return; + } + let mut accumulated_distance = 0.0; + f(points[0], desired_distance); + if points.len() == 1 { + return; + } + for index in 0..(points.len() - 1) { + let mut current = points[index]; + let next = points[index + 1]; + let segment = next - current; + let segment_length = segment.abs(); + let unit_step = segment / segment_length; + if accumulated_distance + segment_length >= desired_distance { + current = current + unit_step * (desired_distance - accumulated_distance); + f(current, desired_distance); + accumulated_distance -= desired_distance; + } + accumulated_distance += segment_length; + while accumulated_distance >= desired_distance { + current = current + unit_step * desired_distance; + f(current, desired_distance); + accumulated_distance -= desired_distance; + } + } + f(points[points.len() - 1], accumulated_distance); +} + +#[derive(Default, Clone, Copy, Debug)] +struct Dct32([f32; 32]); + +impl Dct32 { + fn continuous_idct(&self, t: f32) -> f32 { + const MULTIPLIERS: [f32; 32] = [ + PI / 32.0 * 0.0, + PI / 32.0 * 1.0, + PI / 32.0 * 2.0, + PI / 32.0 * 3.0, + PI / 32.0 * 4.0, + PI / 32.0 * 5.0, + PI / 32.0 * 6.0, + PI / 32.0 * 7.0, + PI / 32.0 * 8.0, + PI / 32.0 * 9.0, + PI / 32.0 * 10.0, + PI / 32.0 * 11.0, + PI / 32.0 * 12.0, + PI / 32.0 * 13.0, + PI / 32.0 * 14.0, + PI / 32.0 * 15.0, + PI / 32.0 * 16.0, + PI / 32.0 * 17.0, + PI / 32.0 * 18.0, + PI / 32.0 * 19.0, + PI / 32.0 * 20.0, + PI / 32.0 * 21.0, + PI / 32.0 * 22.0, + PI / 32.0 * 23.0, + PI / 32.0 * 24.0, + PI / 32.0 * 25.0, + PI / 32.0 * 26.0, + PI / 32.0 * 27.0, + PI / 32.0 * 28.0, + PI / 32.0 * 29.0, + PI / 32.0 * 30.0, + PI / 32.0 * 31.0, + ]; + let tandhalf = t + 0.5; + zip(MULTIPLIERS.iter(), self.0.iter()) + .map(|(multiplier, coeff)| SQRT_2 * coeff * fast_cos(multiplier * tandhalf)) + .sum() + } +} + +impl Splines { + #[cfg(test)] + pub fn create( + quantization_adjustment: i32, + splines: Vec<QuantizedSpline>, + starting_points: Vec<Point>, + ) -> Splines { + Splines { + quantization_adjustment, + splines, + starting_points, + segments: vec![], + segment_indices: vec![], + segment_y_start: vec![], + } + } + pub fn draw_segments(&self, row: &mut [&mut [f32]], row_pos: (usize, usize), xsize: usize) { + let first_segment_index_pos = self.segment_y_start[row_pos.1]; + let last_segment_index_pos = self.segment_y_start[row_pos.1 + 1]; + for segment_index_pos in first_segment_index_pos..last_segment_index_pos { + self.draw_segment( + row, + row_pos, + xsize, + &self.segments[self.segment_indices[segment_index_pos as usize]], + ); + } + } + fn draw_segment( + &self, + row: &mut [&mut [f32]], + row_pos: (usize, usize), + xsize: usize, + segment: &SplineSegment, + ) { + let (x0, y) = row_pos; + let x1 = x0 + xsize; + let clamped_x0 = x0.max((segment.center_x - segment.maximum_distance).round() as usize); + // one-past-the-end + let clamped_x1 = x1.min((segment.center_x + segment.maximum_distance).round() as usize + 1); + for x in clamped_x0..clamped_x1 { + self.draw_segment_at(row, (x, y), x0, segment); + } + } + fn draw_segment_at( + &self, + row: &mut [&mut [f32]], + pixel_pos: (usize, usize), + row_x0: usize, + segment: &SplineSegment, + ) { + let (x, y) = pixel_pos; + let inv_sigma = segment.inv_sigma; + let half = 0.5f32; + let one_over_2s2 = 0.353_553_38_f32; + let sigma_over_4_times_intensity = segment.sigma_over_4_times_intensity; + let dx = x as f32 - segment.center_x; + let dy = y as f32 - segment.center_y; + let sqd = dx * dx + dy * dy; + let distance = sqd.sqrt(); + let one_dimensional_factor = fast_erff((distance * half + one_over_2s2) * inv_sigma) + - fast_erff((distance * half - one_over_2s2) * inv_sigma); + let local_intensity = + sigma_over_4_times_intensity * one_dimensional_factor * one_dimensional_factor; + for (channel_index, row) in row.iter_mut().enumerate() { + let cm = segment.color[channel_index]; + let inp = row[x - row_x0]; + row[x - row_x0] = cm * local_intensity + inp; + } + } + + fn add_segment( + &mut self, + center: &Point, + intensity: f32, + color: [f32; 3], + sigma: f32, + segments_by_y: &mut Vec<(u64, usize)>, + ) { + if sigma.is_infinite() + || sigma == 0.0 + || (1.0 / sigma).is_infinite() + || intensity.is_infinite() + { + return; + } + // TODO(zond): Use 3 if not JXL_HIGH_PRECISION + const DISTANCE_EXP: f32 = 5.0; + let max_color = [0.01, color[0], color[1], color[2]] + .iter() + .map(|chan| (chan * intensity).abs()) + .max_by(|a, b| a.total_cmp(b)) + .unwrap(); + let max_distance = + (-2.0 * sigma * sigma * (0.1f32.ln() * DISTANCE_EXP - max_color.ln())).sqrt(); + let segment = SplineSegment { + center_x: center.x, + center_y: center.y, + color, + inv_sigma: 1.0 / sigma, + sigma_over_4_times_intensity: 0.25 * sigma * intensity, + maximum_distance: max_distance, + }; + let y0 = (center.y - max_distance).round() as i64; + let y1 = (center.y + max_distance).round() as i64 + 1; + for y in 0.max(y0)..y1 { + segments_by_y.push((y as u64, self.segments.len())); + } + self.segments.push(segment); + } + + fn add_segments_from_points( + &mut self, + spline: &Spline, + points_to_draw: &[(Point, f32)], + length: f32, + desired_distance: f32, + segments_by_y: &mut Vec<(u64, usize)>, + ) { + let inv_length = 1.0 / length; + for (point_index, (point, multiplier)) in points_to_draw.iter().enumerate() { + let progress = (point_index as f32 * desired_distance * inv_length).min(1.0); + let mut color = [0.0; 3]; + for (index, coeffs) in spline.color_dct.iter().enumerate() { + color[index] = coeffs.continuous_idct((32.0 - 1.0) * progress); + } + let sigma = spline.sigma_dct.continuous_idct((32.0 - 1.0) * progress); + self.add_segment(point, *multiplier, color, sigma, segments_by_y); + } + } + + pub fn initialize_draw_cache( + &mut self, + image_xsize: u64, + image_ysize: u64, + color_correlation_params: &ColorCorrelationParams, + ) -> Result<()> { + let mut total_estimated_area_reached = 0u64; + let mut splines = Vec::new(); + let area_limit = area_limit(image_xsize * image_ysize); + for (index, qspline) in self.splines.iter().enumerate() { + let spline = qspline.dequantize( + &self.starting_points[index], + self.quantization_adjustment, + color_correlation_params.y_to_x_lf(), + color_correlation_params.y_to_b_lf(), + image_xsize * image_ysize, + )?; + total_estimated_area_reached += spline.estimated_area_reached; + if total_estimated_area_reached > area_limit { + return Err(Error::SplinesAreaTooLarge( + total_estimated_area_reached, + area_limit, + )); + } + spline.validate_adjacent_point_coincidence()?; + splines.push(spline); + } + + if total_estimated_area_reached + > (8 * image_xsize * image_ysize + (1u64 << 25)).min(1u64 << 30) + { + warn!( + "Large total_estimated_area_reached, expect slower decoding:{}", + total_estimated_area_reached + ); + } + + let mut segments_by_y = Vec::new(); + + self.segments.clear(); + for spline in splines { + let mut points_to_draw = Vec::<(Point, f32)>::new(); + let intermediate_points = draw_centripetal_catmull_rom_spline(&spline.control_points)?; + for_each_equally_spaced_point( + &intermediate_points, + DESIRED_RENDERING_DISTANCE, + |p, d| points_to_draw.push((p, d)), + ); + let length = (points_to_draw.len() - 2) as f32 * DESIRED_RENDERING_DISTANCE + + points_to_draw[points_to_draw.len() - 1].1; + if length <= 0.0 { + continue; + } + self.add_segments_from_points( + &spline, + &points_to_draw, + length, + DESIRED_RENDERING_DISTANCE, + &mut segments_by_y, + ); + } + + // TODO(from libjxl): Consider linear sorting here. + segments_by_y.sort_by_key(|segment| segment.0); + + self.segment_indices.clear(); + self.segment_indices.try_reserve(segments_by_y.len())?; + self.segment_indices.resize(segments_by_y.len(), 0); + + self.segment_y_start.clear(); + self.segment_y_start.try_reserve(image_ysize as usize + 1)?; + self.segment_y_start.resize(image_ysize as usize + 1, 0); + + for (i, segment) in segments_by_y.iter().enumerate() { + self.segment_indices[i] = segment.1; + let y = segment.0; + if y < image_ysize { + self.segment_y_start[y as usize + 1] += 1; + } + } + for y in 0..image_ysize { + self.segment_y_start[y as usize + 1] += self.segment_y_start[y as usize]; + } + Ok(()) + } + + #[instrument(level = "debug", skip(br), ret, err)] + pub fn read(br: &mut BitReader, num_pixels: u32) -> Result<Splines> { + trace!(pos = br.total_bits_read()); + let splines_histograms = Histograms::decode(NUM_SPLINE_CONTEXTS, br, true)?; + let mut splines_reader = SymbolReader::new(&splines_histograms, br, None)?; + let num_splines = + 1 + splines_reader.read_unsigned(&splines_histograms, br, NUM_SPLINES_CONTEXT)?; + let max_control_points = + MAX_NUM_CONTROL_POINTS.min(num_pixels / MAX_NUM_CONTROL_POINTS_PER_PIXEL_RATIO); + if num_splines > max_control_points { + return Err(Error::SplinesTooMany(num_splines, max_control_points)); + } + + let mut starting_points = Vec::new(); + let mut last_x = 0; + let mut last_y = 0; + for i in 0..num_splines { + let unsigned_x = + splines_reader.read_unsigned(&splines_histograms, br, STARTING_POSITION_CONTEXT)?; + let unsigned_y = + splines_reader.read_unsigned(&splines_histograms, br, STARTING_POSITION_CONTEXT)?; + + let (x, y) = if i != 0 { + ( + unpack_signed(unsigned_x) + last_x, + unpack_signed(unsigned_y) + last_y, + ) + } else { + (unsigned_x as i32, unsigned_y as i32) + }; + // It is not in spec, but reasonable limit to avoid overflows. + let max_coordinate = x.abs().max(y.abs()); + if max_coordinate >= SPLINE_POS_LIMIT { + return Err(Error::SplinesCoordinatesLimit( + max_coordinate, + SPLINE_POS_LIMIT, + )); + } + + starting_points.push(Point { + x: x as f32, + y: y as f32, + }); + + last_x = x; + last_y = y; + } + + let quantization_adjustment = + splines_reader.read_signed(&splines_histograms, br, QUANTIZATION_ADJUSTMENT_CONTEXT)?; + + let mut splines = Vec::new(); + let mut num_control_points = 0u32; + for _ in 0..num_splines { + splines.push(QuantizedSpline::read( + br, + &splines_histograms, + &mut splines_reader, + max_control_points, + &mut num_control_points, + )?); + } + splines_reader.check_final_state(&splines_histograms)?; + Ok(Splines { + quantization_adjustment, + splines, + starting_points, + ..Splines::default() + }) + } +} + +#[cfg(test)] +#[allow(clippy::excessive_precision)] +mod test_splines { + use std::{f32::consts::SQRT_2, iter::zip}; + use test_log::test; + + use crate::{ + error::{Error, Result}, + features::spline::SplineSegment, + frame::color_correlation_map::ColorCorrelationParams, + util::test::{assert_all_almost_abs_eq, assert_almost_abs_eq, assert_almost_eq}, + }; + + use super::{ + DESIRED_RENDERING_DISTANCE, Dct32, Point, QuantizedSpline, Spline, Splines, + draw_centripetal_catmull_rom_spline, for_each_equally_spaced_point, + }; + + #[test] + fn dequantize() -> Result<(), Error> { + // Golden data generated by libjxl. + let quantized_and_dequantized = [ + ( + QuantizedSpline { + control_points: vec![ + (109, 105), + (-247, -261), + (168, 427), + (-46, -360), + (-61, 181), + ], + color_dct: [ + [ + 12223, 9452, 5524, 16071, 1048, 17024, 14833, 7690, 21952, 2405, 2571, + 2190, 1452, 2500, 18833, 1667, 5857, 21619, 1310, 20000, 10429, 11667, + 7976, 18786, 12976, 18548, 14786, 12238, 8667, 3405, 19929, 8429, + ], + [ + 177, 712, 127, 999, 969, 356, 105, 12, 1132, 309, 353, 415, 1213, 156, + 988, 524, 316, 1100, 64, 36, 816, 1285, 183, 889, 839, 1099, 79, 1316, + 287, 105, 689, 841, + ], + [ + 780, -201, -38, -695, -563, -293, -88, 1400, -357, 520, 979, 431, -118, + 590, -971, -127, 157, 206, 1266, 204, -320, -223, 704, -687, -276, + -716, 787, -1121, 40, 292, 249, -10, + ], + ], + sigma_dct: [ + 139, 65, 133, 5, 137, 272, 88, 178, 71, 256, 254, 82, 126, 252, 152, 53, + 281, 15, 8, 209, 285, 156, 73, 56, 36, 287, 86, 244, 270, 94, 224, 156, + ], + }, + Spline { + control_points: vec![ + Point { x: 109.0, y: 54.0 }, + Point { x: 218.0, y: 159.0 }, + Point { x: 80.0, y: 3.0 }, + Point { x: 110.0, y: 274.0 }, + Point { x: 94.0, y: 185.0 }, + Point { x: 17.0, y: 277.0 }, + ], + color_dct: [ + Dct32([ + 36.300457, + 39.69839859, + 23.20079994, + 67.49819946, + 4.401599884, + 71.50080109, + 62.29859924, + 32.29800034, + 92.19839478, + 10.10099983, + 10.79819965, + 9.197999954, + 6.098399639, + 10.5, + 79.09859467, + 7.001399517, + 24.59939957, + 90.79979706, + 5.501999855, + 84.0, + 43.80179977, + 49.00139999, + 33.49919891, + 78.90119934, + 54.49919891, + 77.90159607, + 62.10119629, + 51.39959717, + 36.40139771, + 14.30099964, + 83.70179749, + 35.40179825, + ]), + Dct32([ + 9.386842728, + 53.40000153, + 9.525000572, + 74.92500305, + 72.67500305, + 26.70000076, + 7.875000477, + 0.9000000358, + 84.90000153, + 23.17500114, + 26.47500038, + 31.12500191, + 90.9750061, + 11.70000076, + 74.1000061, + 39.30000305, + 23.70000076, + 82.5, + 4.800000191, + 2.700000048, + 61.20000076, + 96.37500763, + 13.72500038, + 66.67500305, + 62.92500305, + 82.42500305, + 5.925000191, + 98.70000458, + 21.52500153, + 7.875000477, + 51.67500305, + 63.07500076, + ]), + Dct32([ + 47.99487305, + 39.33000183, + 6.865000725, + 26.27500153, + 33.2650032, + 6.190000534, + 1.715000629, + 98.90000153, + 59.91000366, + 59.57500458, + 95.00499725, + 61.29500198, + 82.71500397, + 53.0, + 6.130004883, + 30.41000366, + 34.69000244, + 96.91999817, + 93.4200058, + 16.97999954, + 38.80000305, + 80.76500702, + 63.00499725, + 18.5850029, + 43.60500336, + 32.30500412, + 61.01499939, + 20.23000336, + 24.32500076, + 28.31500053, + 69.10500336, + 62.375, + ]), + ], + sigma_dct: Dct32([ + 32.75933838, + 21.66449928, + 44.32889938, + 1.666499972, + 45.66209793, + 90.6576004, + 29.33039856, + 59.32740021, + 23.66429901, + 85.32479858, + 84.6581955, + 27.33059883, + 41.99580002, + 83.99160004, + 50.66159821, + 17.66489983, + 93.65729523, + 4.999499798, + 2.666399956, + 69.65969849, + 94.9905014, + 51.99480057, + 24.33090019, + 18.66479874, + 11.99880028, + 95.65709686, + 28.66379929, + 81.32519531, + 89.99099731, + 31.3302002, + 74.65919495, + 51.99480057, + ]), + estimated_area_reached: 19843491681, + }, + ), + ( + QuantizedSpline { + control_points: vec![ + (24, -32), + (-178, -7), + (226, 151), + (121, -172), + (-184, 39), + (-201, -182), + (301, 404), + ], + color_dct: [ + [ + 5051, 6881, 5238, 1571, 9952, 19762, 2048, 13524, 16405, 2310, 1286, + 4714, 16857, 21429, 12500, 15524, 1857, 5595, 6286, 17190, 15405, + 20738, 310, 16071, 10952, 16286, 15571, 8452, 6929, 3095, 9905, 5690, + ], + [ + 899, 1059, 836, 388, 1291, 247, 235, 203, 1073, 747, 1283, 799, 356, + 1281, 1231, 561, 477, 720, 309, 733, 1013, 477, 779, 1183, 32, 1041, + 1275, 367, 88, 1047, 321, 931, + ], + [ + -78, 244, -883, 943, -682, 752, 107, 262, -75, 557, -202, -575, -231, + -731, -605, 732, 682, 650, 592, -14, -1035, 913, -188, -95, 286, -574, + -509, 67, 86, -1056, 592, 380, + ], + ], + sigma_dct: [ + 308, 8, 125, 7, 119, 237, 209, 60, 277, 215, 126, 186, 90, 148, 211, 136, + 188, 142, 140, 124, 272, 140, 274, 165, 24, 209, 76, 254, 185, 83, 11, 141, + ], + }, + Spline { + control_points: vec![ + Point { x: 172.0, y: 309.0 }, + Point { x: 196.0, y: 277.0 }, + Point { x: 42.0, y: 238.0 }, + Point { x: 114.0, y: 350.0 }, + Point { x: 307.0, y: 290.0 }, + Point { x: 316.0, y: 269.0 }, + Point { x: 124.0, y: 66.0 }, + Point { x: 233.0, y: 267.0 }, + ], + color_dct: [ + Dct32([ + 15.00070381, + 28.90019989, + 21.99959946, + 6.598199844, + 41.79839706, + 83.00039673, + 8.601599693, + 56.80079651, + 68.90100098, + 9.701999664, + 5.401199818, + 19.79879951, + 70.79940033, + 90.00180054, + 52.5, + 65.20079803, + 7.799399853, + 23.49899864, + 26.40119934, + 72.19799805, + 64.7009964, + 87.09959412, + 1.301999927, + 67.49819946, + 45.99839783, + 68.40119934, + 65.39820099, + 35.49839783, + 29.10179901, + 12.9989996, + 41.60099792, + 23.89799881, + ]), + Dct32([ + 47.67667389, + 79.42500305, + 62.70000076, + 29.10000038, + 96.82500458, + 18.52500153, + 17.625, + 15.22500038, + 80.4750061, + 56.02500153, + 96.2250061, + 59.92500305, + 26.70000076, + 96.07500458, + 92.32500458, + 42.07500076, + 35.77500153, + 54.00000381, + 23.17500114, + 54.97500229, + 75.9750061, + 35.77500153, + 58.42500305, + 88.7250061, + 2.400000095, + 78.07500458, + 95.625, + 27.52500153, + 6.600000381, + 78.52500153, + 24.07500076, + 69.82500458, + ]), + Dct32([ + 43.81587219, + 96.50500488, + 0.8899993896, + 95.11000061, + 49.0850029, + 71.16500092, + 25.11499977, + 33.56500244, + 75.2250061, + 95.01499939, + 82.08500671, + 19.67500305, + 10.53000069, + 44.90500259, + 49.9750061, + 93.31500244, + 83.51499939, + 99.5, + 64.61499786, + 53.99500275, + 3.525009155, + 99.68499756, + 45.2650032, + 82.07500458, + 22.42000008, + 37.89500427, + 59.99499893, + 32.21500015, + 12.62000084, + 4.605003357, + 65.51499939, + 96.42500305, + ]), + ], + sigma_dct: Dct32([ + 72.58903503, + 2.666399956, + 41.66249847, + 2.333099842, + 39.66270065, + 78.99209595, + 69.65969849, + 19.99799919, + 92.32409668, + 71.65950012, + 41.99580002, + 61.9937973, + 29.99699974, + 49.32839966, + 70.32630157, + 45.3288002, + 62.66040039, + 47.32859802, + 46.66199875, + 41.32920074, + 90.6576004, + 46.66199875, + 91.32419586, + 54.99449921, + 7.999199867, + 69.65969849, + 25.3307991, + 84.6581955, + 61.66049957, + 27.66390038, + 3.66629982, + 46.99530029, + ]), + estimated_area_reached: 25829781306, + }, + ), + ( + QuantizedSpline { + control_points: vec![ + (157, -89), + (-244, 41), + (-58, 168), + (429, -185), + (-361, 198), + (230, -269), + (-416, 203), + (167, 65), + (460, -344), + ], + color_dct: [ + [ + 5691, 15429, 1000, 2524, 5595, 4048, 18881, 1357, 14381, 3952, 22595, + 15167, 20857, 2500, 905, 14548, 5452, 19500, 19143, 9643, 10929, 6048, + 9476, 7143, 11952, 21524, 6643, 22310, 15500, 11476, 5310, 10452, + ], + [ + 470, 880, 47, 1203, 1295, 211, 475, 8, 907, 528, 325, 1145, 769, 1035, + 633, 905, 57, 72, 1216, 780, 1, 696, 47, 637, 843, 580, 1144, 477, 669, + 479, 256, 643, + ], + [ + 1169, -301, 1041, -725, -43, -22, 774, 134, -822, 499, 456, -287, -713, + -776, 76, 449, 750, 580, -207, -643, 956, -426, 377, -64, 101, -250, + -164, 259, 169, -240, 430, -22, + ], + ], + sigma_dct: [ + 354, 5, 75, 56, 140, 226, 84, 187, 151, 70, 257, 288, 137, 99, 100, 159, + 79, 176, 59, 210, 278, 68, 171, 65, 230, 263, 69, 199, 107, 107, 170, 202, + ], + }, + Spline { + control_points: vec![ + Point { x: 100.0, y: 186.0 }, + Point { x: 257.0, y: 97.0 }, + Point { x: 170.0, y: 49.0 }, + Point { x: 25.0, y: 169.0 }, + Point { x: 309.0, y: 104.0 }, + Point { x: 232.0, y: 237.0 }, + Point { x: 385.0, y: 101.0 }, + Point { x: 122.0, y: 168.0 }, + Point { x: 26.0, y: 300.0 }, + Point { x: 390.0, y: 88.0 }, + ], + color_dct: [ + Dct32([ + 16.90140724, + 64.80179596, + 4.199999809, + 10.60079956, + 23.49899864, + 17.00160027, + 79.30019379, + 5.699399948, + 60.40019608, + 16.59840012, + 94.89899445, + 63.70139694, + 87.59939575, + 10.5, + 3.80099988, + 61.10159683, + 22.89839935, + 81.8999939, + 80.40059662, + 40.50059891, + 45.90179825, + 25.40159988, + 39.79919815, + 30.00059891, + 50.19839859, + 90.40079498, + 27.90059853, + 93.70199585, + 65.09999847, + 48.19919968, + 22.30200005, + 43.89839935, + ]), + Dct32([ + 24.92551422, + 66.0, + 3.525000095, + 90.2250061, + 97.12500763, + 15.82500076, + 35.625, + 0.6000000238, + 68.02500153, + 39.60000229, + 24.37500191, + 85.875, + 57.67500305, + 77.625, + 47.47500229, + 67.875, + 4.275000095, + 5.400000095, + 91.20000458, + 58.50000381, + 0.07500000298, + 52.20000076, + 3.525000095, + 47.77500153, + 63.22500229, + 43.5, + 85.80000305, + 35.77500153, + 50.17500305, + 35.92500305, + 19.20000076, + 48.22500229, + ]), + Dct32([ + 82.78805542, + 44.93000031, + 76.39500427, + 39.4750061, + 94.11500549, + 14.2850008, + 89.80500031, + 9.980000496, + 10.48500061, + 74.52999878, + 56.29500198, + 65.78500366, + 7.765003204, + 23.30500031, + 52.79500198, + 99.30500031, + 56.77500153, + 46.0, + 76.71000671, + 13.49000549, + 66.99499512, + 22.38000107, + 29.91499901, + 43.29500198, + 70.2950058, + 26.0, + 74.31999969, + 53.90499878, + 62.00500488, + 19.12500381, + 49.30000305, + 46.68500137, + ]), + ], + sigma_dct: Dct32([ + 83.43025208, + 1.666499972, + 24.99749947, + 18.66479874, + 46.66199875, + 75.32579803, + 27.99720001, + 62.32709885, + 50.32830048, + 23.33099937, + 85.65809631, + 95.99040222, + 45.66209793, + 32.99670029, + 33.32999802, + 52.99469757, + 26.33069992, + 58.66079712, + 19.66469955, + 69.99299622, + 92.65740204, + 22.6644001, + 56.99430084, + 21.66449928, + 76.65899658, + 87.65789795, + 22.99769974, + 66.3266983, + 35.6631012, + 35.6631012, + 56.6609993, + 67.32659912, + ]), + estimated_area_reached: 47263284396, + }, + ), + ]; + for (quantized, want_dequantized) in quantized_and_dequantized { + let got_dequantized = quantized.dequantize( + &want_dequantized.control_points[0], + 0, + 0.0, + 1.0, + 2u64 << 30, + )?; + assert_eq!( + got_dequantized.control_points.len(), + want_dequantized.control_points.len() + ); + assert_all_almost_abs_eq( + got_dequantized + .control_points + .iter() + .map(|p| p.x) + .collect::<Vec<f32>>(), + want_dequantized + .control_points + .iter() + .map(|p| p.x) + .collect::<Vec<f32>>(), + 1e-6, + ); + assert_all_almost_abs_eq( + got_dequantized + .control_points + .iter() + .map(|p| p.y) + .collect::<Vec<f32>>(), + want_dequantized + .control_points + .iter() + .map(|p| p.y) + .collect::<Vec<f32>>(), + 1e-6, + ); + for index in 0..got_dequantized.color_dct.len() { + assert_all_almost_abs_eq( + got_dequantized.color_dct[index].0, + want_dequantized.color_dct[index].0, + 1e-4, + ); + } + assert_all_almost_abs_eq( + got_dequantized.sigma_dct.0, + want_dequantized.sigma_dct.0, + 1e-4, + ); + assert_eq!( + got_dequantized.estimated_area_reached, + want_dequantized.estimated_area_reached, + ); + } + Ok(()) + } + + #[test] + fn centripetal_catmull_rom_spline() -> Result<(), Error> { + let control_points = vec![Point { x: 1.0, y: 2.0 }, Point { x: 4.0, y: 3.0 }]; + let want_result = [ + Point { x: 1.0, y: 2.0 }, + Point { + x: 1.187500119, + y: 2.0625, + }, + Point { x: 1.375, y: 2.125 }, + Point { + x: 1.562499881, + y: 2.1875, + }, + Point { + x: 1.750000119, + y: 2.25, + }, + Point { + x: 1.9375, + y: 2.3125, + }, + Point { x: 2.125, y: 2.375 }, + Point { + x: 2.312500238, + y: 2.4375, + }, + Point { + x: 2.500000238, + y: 2.5, + }, + Point { + x: 2.6875, + y: 2.5625, + }, + Point { + x: 2.875000477, + y: 2.625, + }, + Point { + x: 3.062499762, + y: 2.6875, + }, + Point { x: 3.25, y: 2.75 }, + Point { + x: 3.4375, + y: 2.8125, + }, + Point { + x: 3.624999762, + y: 2.875, + }, + Point { + x: 3.812500238, + y: 2.9375, + }, + Point { x: 4.0, y: 3.0 }, + ]; + let got_result = draw_centripetal_catmull_rom_spline(&control_points)?; + assert_all_almost_abs_eq( + got_result.iter().map(|p| p.x).collect::<Vec<f32>>(), + want_result.iter().map(|p| p.x).collect::<Vec<f32>>(), + 1e-10, + ); + Ok(()) + } + + #[test] + fn equally_spaced_points() -> Result<(), Error> { + let desired_rendering_distance = 10.0f32; + let segments = [ + Point { x: 0.0, y: 0.0 }, + Point { x: 5.0, y: 0.0 }, + Point { x: 35.0, y: 0.0 }, + Point { x: 35.0, y: 10.0 }, + ]; + let want_results = [ + (Point { x: 0.0, y: 0.0 }, desired_rendering_distance), + (Point { x: 10.0, y: 0.0 }, desired_rendering_distance), + (Point { x: 20.0, y: 0.0 }, desired_rendering_distance), + (Point { x: 30.0, y: 0.0 }, desired_rendering_distance), + (Point { x: 35.0, y: 5.0 }, desired_rendering_distance), + (Point { x: 35.0, y: 10.0 }, 5.0f32), + ]; + let mut got_results = Vec::<(Point, f32)>::new(); + for_each_equally_spaced_point(&segments, desired_rendering_distance, |p, d| { + got_results.push((p, d)) + }); + assert_all_almost_abs_eq( + got_results.iter().map(|(p, _)| p.x).collect::<Vec<f32>>(), + want_results.iter().map(|(p, _)| p.x).collect::<Vec<f32>>(), + 1e-9, + ); + assert_all_almost_abs_eq( + got_results.iter().map(|(p, _)| p.y).collect::<Vec<f32>>(), + want_results.iter().map(|(p, _)| p.y).collect::<Vec<f32>>(), + 1e-9, + ); + assert_all_almost_abs_eq( + got_results.iter().map(|(_, d)| *d).collect::<Vec<f32>>(), + want_results.iter().map(|(_, d)| *d).collect::<Vec<f32>>(), + 1e-9, + ); + Ok(()) + } + + #[test] + fn dct32() -> Result<(), Error> { + let mut dct = Dct32::default(); + for (i, coeff) in dct.0.iter_mut().enumerate() { + *coeff = 0.05f32 * i as f32; + } + // Golden numbers come from libjxl. + let want_out = [ + 16.7353153229, + -18.6041717529, + 7.9931735992, + -7.1250801086, + 4.6699867249, + -4.3367614746, + 3.2450540066, + -3.0694460869, + 2.4446771145, + -2.3350939751, + 1.9243829250, + -1.8484034538, + 1.5531382561, + -1.4964176416, + 1.2701368332, + -1.2254891396, + 1.0434474945, + -1.0067725182, + 0.8544843197, + -0.8232427835, + 0.6916543841, + -0.6642799377, + 0.5473306179, + -0.5226536393, + 0.4161090851, + -0.3933961987, + 0.2940555215, + -0.2726306915, + 0.1781132221, + -0.1574717760, + 0.0656886101, + -0.0454511642, + ]; + for (t, want) in want_out.iter().enumerate() { + let got_out = dct.continuous_idct(t as f32); + assert_almost_abs_eq(got_out, *want, 1e-4); + } + Ok(()) + } + + fn verify_segment_almost_equal(seg1: &SplineSegment, seg2: &SplineSegment) { + assert_almost_eq(seg1.center_x, seg2.center_x, 1e-2, 1e-4); + assert_almost_eq(seg1.center_y, seg2.center_y, 1e-2, 1e-4); + for (got, want) in zip(seg1.color.iter(), seg2.color.iter()) { + assert_almost_eq(*got, *want, 1e-2, 1e-4); + } + assert_almost_eq(seg1.inv_sigma, seg2.inv_sigma, 1e-2, 1e-4); + assert_almost_eq(seg1.maximum_distance, seg2.maximum_distance, 1e-2, 1e-4); + assert_almost_eq( + seg1.sigma_over_4_times_intensity, + seg2.sigma_over_4_times_intensity, + 1e-2, + 1e-4, + ); + } + + #[test] + fn spline_segments_add_segment() -> Result<(), Error> { + let mut splines = Splines::default(); + let mut segments_by_y = Vec::<(u64, usize)>::new(); + + splines.add_segment( + &Point { x: 10.0, y: 20.0 }, + 0.5, + [0.5, 0.6, 0.7], + 0.8, + &mut segments_by_y, + ); + // Golden numbers come from libjxl. + let want_segment = SplineSegment { + center_x: 10.0, + center_y: 20.0, + color: [0.5, 0.6, 0.7], + inv_sigma: 1.25, + maximum_distance: 3.65961, + sigma_over_4_times_intensity: 0.1, + }; + assert_eq!(splines.segments.len(), 1); + verify_segment_almost_equal(&splines.segments[0], &want_segment); + let want_segments_by_y = [ + (16, 0), + (17, 0), + (18, 0), + (19, 0), + (20, 0), + (21, 0), + (22, 0), + (23, 0), + (24, 0), + ]; + for (got, want) in zip(segments_by_y.iter(), want_segments_by_y.iter()) { + assert_eq!(got.0, want.0); + assert_eq!(got.1, want.1); + } + Ok(()) + } + + #[test] + fn spline_segments_add_segments_from_points() -> Result<(), Error> { + let mut splines = Splines::default(); + let mut segments_by_y = Vec::<(u64, usize)>::new(); + let mut color_dct = [Dct32::default(); 3]; + for (channel_index, channel_dct) in color_dct.iter_mut().enumerate() { + for (coeff_index, coeff) in channel_dct.0.iter_mut().enumerate() { + *coeff = 0.1 * channel_index as f32 + 0.05 * coeff_index as f32; + } + } + let mut sigma_dct = Dct32::default(); + for (coeff_index, coeff) in sigma_dct.0.iter_mut().enumerate() { + *coeff = 0.06 * coeff_index as f32; + } + let spline = Spline { + control_points: vec![], + color_dct, + sigma_dct, + estimated_area_reached: 0, + }; + let points_to_draw = vec![ + (Point { x: 10.0, y: 20.0 }, 1.0), + (Point { x: 11.0, y: 21.0 }, 1.0), + (Point { x: 12.0, y: 21.0 }, 1.0), + ]; + splines.add_segments_from_points( + &spline, + &points_to_draw, + SQRT_2 + 1.0, + DESIRED_RENDERING_DISTANCE, + &mut segments_by_y, + ); + // Golden numbers come from libjxl. + let want_segments = [ + SplineSegment { + center_x: 10.0, + center_y: 20.0, + color: [16.73531532, 19.68646049, 22.63760757], + inv_sigma: 0.04979490861, + maximum_distance: 108.6400299, + sigma_over_4_times_intensity: 5.020593643, + }, + SplineSegment { + center_x: 11.0, + center_y: 21.0, + color: [-0.8199231625, -0.7960500717, -0.7721766233], + inv_sigma: -1.016355753, + maximum_distance: 4.680418015, + sigma_over_4_times_intensity: -0.2459768653, + }, + SplineSegment { + center_x: 12.0, + center_y: 21.0, + color: [-0.7767754197, -0.7544237971, -0.7320720553], + inv_sigma: -1.072811365, + maximum_distance: 4.423510075, + sigma_over_4_times_intensity: -0.2330325693, + }, + ]; + assert_eq!(splines.segments.len(), want_segments.len()); + for (got, want) in zip(splines.segments.iter(), want_segments.iter()) { + verify_segment_almost_equal(got, want); + } + let want_segments_by_y: Vec<(u64, usize)> = (0..=129) + .map(|c| (c, 0)) + .chain((16..=26).map(|c| (c, 1))) + .chain((17..=25).map(|c| (c, 2))) + .collect(); + for (got, want) in zip(segments_by_y.iter(), want_segments_by_y.iter()) { + assert_eq!(got.0, want.0); + assert_eq!(got.1, want.1); + } + Ok(()) + } + + #[test] + fn init_draw_cache() -> Result<(), Error> { + let mut splines = Splines { + splines: vec![ + QuantizedSpline { + control_points: vec![ + (109, 105), + (-247, -261), + (168, 427), + (-46, -360), + (-61, 181), + ], + color_dct: [ + [ + 12223, 9452, 5524, 16071, 1048, 17024, 14833, 7690, 21952, 2405, 2571, + 2190, 1452, 2500, 18833, 1667, 5857, 21619, 1310, 20000, 10429, 11667, + 7976, 18786, 12976, 18548, 14786, 12238, 8667, 3405, 19929, 8429, + ], + [ + 177, 712, 127, 999, 969, 356, 105, 12, 1132, 309, 353, 415, 1213, 156, + 988, 524, 316, 1100, 64, 36, 816, 1285, 183, 889, 839, 1099, 79, 1316, + 287, 105, 689, 841, + ], + [ + 780, -201, -38, -695, -563, -293, -88, 1400, -357, 520, 979, 431, -118, + 590, -971, -127, 157, 206, 1266, 204, -320, -223, 704, -687, -276, + -716, 787, -1121, 40, 292, 249, -10, + ], + ], + sigma_dct: [ + 139, 65, 133, 5, 137, 272, 88, 178, 71, 256, 254, 82, 126, 252, 152, 53, + 281, 15, 8, 209, 285, 156, 73, 56, 36, 287, 86, 244, 270, 94, 224, 156, + ], + }, + QuantizedSpline { + control_points: vec![ + (24, -32), + (-178, -7), + (226, 151), + (121, -172), + (-184, 39), + (-201, -182), + (301, 404), + ], + color_dct: [ + [ + 5051, 6881, 5238, 1571, 9952, 19762, 2048, 13524, 16405, 2310, 1286, + 4714, 16857, 21429, 12500, 15524, 1857, 5595, 6286, 17190, 15405, + 20738, 310, 16071, 10952, 16286, 15571, 8452, 6929, 3095, 9905, 5690, + ], + [ + 899, 1059, 836, 388, 1291, 247, 235, 203, 1073, 747, 1283, 799, 356, + 1281, 1231, 561, 477, 720, 309, 733, 1013, 477, 779, 1183, 32, 1041, + 1275, 367, 88, 1047, 321, 931, + ], + [ + -78, 244, -883, 943, -682, 752, 107, 262, -75, 557, -202, -575, -231, + -731, -605, 732, 682, 650, 592, -14, -1035, 913, -188, -95, 286, -574, + -509, 67, 86, -1056, 592, 380, + ], + ], + sigma_dct: [ + 308, 8, 125, 7, 119, 237, 209, 60, 277, 215, 126, 186, 90, 148, 211, 136, + 188, 142, 140, 124, 272, 140, 274, 165, 24, 209, 76, 254, 185, 83, 11, 141, + ], + }, + ], + starting_points: vec![Point { x: 10.0, y: 20.0 }, Point { x: 5.0, y: 40.0 }], + ..Default::default() + }; + splines.initialize_draw_cache( + 1 << 15, + 1 << 15, + &ColorCorrelationParams { + color_factor: 1, + base_correlation_x: 0.0, + base_correlation_b: 0.0, + ytox_lf: 0, + ytob_lf: 0, + }, + )?; + assert_eq!(splines.segments.len(), 1940); + let want_segments_sample = [ + ( + 22, + SplineSegment { + center_x: 25.77652359, + center_y: 35.33295059, + color: [-524.996582, -509.9048462, 43.3883667], + inv_sigma: -0.00197347207, + maximum_distance: 3021.377197, + sigma_over_4_times_intensity: -126.6802902, + }, + ), + ( + 474, + SplineSegment { + center_x: -16.45600891, + center_y: 78.81845856, + color: [-117.6707535, -133.5515594, 343.5632629], + inv_sigma: -0.002631845651, + maximum_distance: 2238.376221, + sigma_over_4_times_intensity: -94.9903717, + }, + ), + ( + 835, + SplineSegment { + center_x: -71.93701172, + center_y: 230.0635529, + color: [44.79507446, 298.9411621, -395.3574524], + inv_sigma: 0.01869126037, + maximum_distance: 316.4499207, + sigma_over_4_times_intensity: 13.3752346, + }, + ), + ( + 1066, + SplineSegment { + center_x: -126.2593002, + center_y: -22.97857094, + color: [-136.4196625, 194.757019, -98.18778992], + inv_sigma: 0.007531851064, + maximum_distance: 769.2540283, + sigma_over_4_times_intensity: 33.19237137, + }, + ), + ( + 1328, + SplineSegment { + center_x: 73.70871735, + center_y: 56.31413269, + color: [-13.44394779, 162.6139221, 93.78419495], + inv_sigma: 0.003664178308, + maximum_distance: 1572.710327, + sigma_over_4_times_intensity: 68.2281189, + }, + ), + ( + 1545, + SplineSegment { + center_x: 77.48892975, + center_y: -92.33877563, + color: [-220.6807556, 66.13040924, -32.26184082], + inv_sigma: 0.03166157752, + maximum_distance: 183.6748352, + sigma_over_4_times_intensity: 7.89600563, + }, + ), + ( + 1774, + SplineSegment { + center_x: -16.43594933, + center_y: -144.8626556, + color: [57.31535339, -46.36843109, 92.14952087], + inv_sigma: -0.01524505392, + maximum_distance: 371.4827271, + sigma_over_4_times_intensity: -16.39876175, + }, + ), + ( + 1929, + SplineSegment { + center_x: 61.19338608, + center_y: -10.70717049, + color: [-69.78807068, 300.6082458, -476.5135803], + inv_sigma: 0.003229281865, + maximum_distance: 1841.37854, + sigma_over_4_times_intensity: 77.41659546, + }, + ), + ]; + for (index, segment) in want_segments_sample { + verify_segment_almost_equal(&segment, &splines.segments[index]); + } + Ok(()) + } +} diff --git a/third_party/rust/jxl/src/frame.rs b/third_party/rust/jxl/src/frame.rs @@ -0,0 +1,1235 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +use crate::{ + GROUP_DIM, + bit_reader::BitReader, + entropy_coding::decode::Histograms, + error::Result, + features::{ + epf::create_sigma_image, noise::Noise, patches::PatchesDictionary, spline::Splines, + }, + headers::{ + FileHeader, Orientation, + color_encoding::ColorSpace, + encodings::UnconditionalCoder, + extra_channels::{ExtraChannel, ExtraChannelInfo}, + frame_header::{Encoding, FrameHeader, Toc, TocNonserialized}, + permutation::Permutation, + }, + image::Image, + render::{ + RenderPipeline, RenderPipelineBuilder, SaveStage, SaveStageType, SimpleRenderPipeline, + SimpleRenderPipelineBuilder, stages::*, + }, + util::{CeilLog2, Xorshift128Plus, tracing_wrappers::*}, +}; +use adaptive_lf_smoothing::adaptive_lf_smoothing; +use block_context_map::BlockContextMap; +use coeff_order::decode_coeff_orders; +use color_correlation_map::ColorCorrelationParams; +use group::decode_vardct_group; +use modular::{FullModularImage, ModularStreamId, Tree}; +use modular::{decode_hf_metadata, decode_vardct_lf}; +use quant_weights::DequantMatrices; +use quantizer::LfQuantFactors; +use quantizer::QuantizerParams; +use transform_map::*; + +use std::sync::Arc; + +mod adaptive_lf_smoothing; +mod block_context_map; +mod coeff_order; +pub mod color_correlation_map; +mod group; +pub mod modular; +mod quant_weights; +pub mod quantizer; +pub mod transform_map; + +#[derive(Debug, PartialEq, Eq, Clone, Copy)] +pub enum Section { + LfGlobal, + Lf { group: usize }, + HfGlobal, + Hf { group: usize, pass: usize }, +} + +pub struct LfGlobalState { + patches: Option<PatchesDictionary>, + splines: Option<Splines>, + noise: Option<Noise>, + lf_quant: LfQuantFactors, + pub quant_params: Option<QuantizerParams>, + block_context_map: Option<BlockContextMap>, + color_correlation_params: Option<ColorCorrelationParams>, + tree: Option<Tree>, + modular_global: FullModularImage, +} + +pub struct PassState { + coeff_orders: Vec<Permutation>, + histograms: Histograms, +} + +pub struct HfGlobalState { + num_histograms: u32, + passes: Vec<PassState>, + dequant_matrices: DequantMatrices, + hf_coefficients: Option<(Image<i32>, Image<i32>, Image<i32>)>, +} + +#[derive(Clone, Debug)] +pub struct ReferenceFrame { + pub frame: Vec<Image<f32>>, + pub saved_before_color_transform: bool, +} + +impl ReferenceFrame { + #[cfg(test)] + pub fn blank( + width: usize, + height: usize, + num_channels: usize, + saved_before_color_transform: bool, + ) -> Result<Self> { + let frame = (0..num_channels) + .map(|_| Image::new_constant((width, height), 0.0)) + .collect::<Result<_>>()?; + Ok(Self { + frame, + saved_before_color_transform, + }) + } +} + +#[derive(Debug)] +pub struct DecoderState { + pub(super) file_header: FileHeader, + pub(super) reference_frames: [Option<ReferenceFrame>; Self::MAX_STORED_FRAMES], + pub(super) lf_frames: [Option<[Image<f32>; 3]>; 4], + pub xyb_output_linear: bool, + pub enable_output: bool, + pub render_spotcolors: bool, +} + +impl DecoderState { + pub const MAX_STORED_FRAMES: usize = 4; + + pub fn new(file_header: FileHeader) -> Self { + Self { + file_header, + reference_frames: [None, None, None, None], + lf_frames: [None, None, None, None], + xyb_output_linear: true, + enable_output: true, + render_spotcolors: true, + } + } + + pub fn extra_channel_info(&self) -> &Vec<ExtraChannelInfo> { + &self.file_header.image_metadata.extra_channel_info + } + + pub fn reference_frame(&self, i: usize) -> Option<&ReferenceFrame> { + assert!(i < Self::MAX_STORED_FRAMES); + self.reference_frames[i].as_ref() + } +} + +pub struct HfMetadata { + ytox_map: Image<i8>, + ytob_map: Image<i8>, + pub raw_quant_map: Image<i32>, + pub transform_map: Image<u8>, + pub epf_map: Image<u8>, + used_hf_types: u32, +} + +pub struct Frame { + header: FrameHeader, + toc: Toc, + modular_color_channels: usize, + lf_global: Option<LfGlobalState>, + hf_global: Option<HfGlobalState>, + lf_image: Option<[Image<f32>; 3]>, + quant_lf: Image<u8>, + hf_meta: Option<HfMetadata>, + decoder_state: DecoderState, + render_pipeline: Option<SimpleRenderPipeline>, +} + +pub struct FrameOutput { + pub decoder_state: Option<DecoderState>, + pub channels: Option<Vec<Image<f32>>>, +} + +impl Frame { + pub fn new(br: &mut BitReader, decoder_state: DecoderState) -> Result<Self> { + let mut frame_header = FrameHeader::read_unconditional( + &(), + br, + &decoder_state.file_header.frame_header_nonserialized(), + )?; + frame_header.postprocess(&decoder_state.file_header.frame_header_nonserialized()); + let num_toc_entries = frame_header.num_toc_entries(); + let toc = Toc::read_unconditional( + &(), + br, + &TocNonserialized { + num_entries: num_toc_entries as u32, + }, + ) + .unwrap(); + br.jump_to_byte_boundary()?; + Self::from_header_and_toc(frame_header, toc, decoder_state) + } + + pub fn from_header_and_toc( + frame_header: FrameHeader, + toc: Toc, + decoder_state: DecoderState, + ) -> Result<Self> { + let image_metadata = &decoder_state.file_header.image_metadata; + let is_modular_gray = !frame_header.do_ycbcr + && !image_metadata.xyb_encoded + && image_metadata.color_encoding.color_space == ColorSpace::Gray; + let modular_color_channels = if frame_header.encoding == Encoding::VarDCT { + 0 + } else if is_modular_gray { + 1 + } else { + 3 + }; + let size_blocks = frame_header.size_blocks(); + let lf_image = if frame_header.encoding == Encoding::VarDCT { + if frame_header.has_lf_frame() { + decoder_state.lf_frames[frame_header.lf_level as usize].clone() + } else { + Some([ + Image::new(size_blocks)?, + Image::new(size_blocks)?, + Image::new(size_blocks)?, + ]) + } + } else { + None + }; + let quant_lf = Image::new(size_blocks)?; + let size_color_tiles = (size_blocks.0.div_ceil(8), size_blocks.1.div_ceil(8)); + let hf_meta = if frame_header.encoding == Encoding::VarDCT { + Some(HfMetadata { + ytox_map: Image::new(size_color_tiles)?, + ytob_map: Image::new(size_color_tiles)?, + raw_quant_map: Image::new(size_blocks)?, + transform_map: Image::new_with_default( + size_blocks, + HfTransformType::INVALID_TRANSFORM, + )?, + epf_map: Image::new(size_blocks)?, + used_hf_types: 0, + }) + } else { + None + }; + Ok(Self { + header: frame_header, + modular_color_channels, + toc, + lf_global: None, + hf_global: None, + lf_image, + quant_lf, + hf_meta, + decoder_state, + render_pipeline: None, + }) + } + + pub fn toc(&self) -> &Toc { + &self.toc + } + + pub fn header(&self) -> &FrameHeader { + &self.header + } + + pub fn total_bytes_in_toc(&self) -> usize { + self.toc.entries.iter().map(|x| *x as usize).sum() + } + + /// Given a bit reader pointing at the end of the TOC, returns a vector of `BitReader`s, each + /// of which reads a specific section. + pub fn sections<'a>(&self, br: &'a mut BitReader) -> Result<Vec<BitReader<'a>>> { + debug!(toc = ?self.toc); + let ret = self + .toc + .entries + .iter() + .scan(br, |br, count| Some(br.split_at(*count as usize))) + .collect::<Result<Vec<_>>>()?; + if !self.toc.permuted { + return Ok(ret); + } + let mut inv_perm = vec![0; ret.len()]; + for (i, pos) in self.toc.permutation.iter().enumerate() { + inv_perm[*pos as usize] = i; + } + let mut shuffled_ret = ret.clone(); + for (br, pos) in ret.into_iter().zip(inv_perm.into_iter()) { + shuffled_ret[pos] = br; + } + Ok(shuffled_ret) + } + + #[instrument(level = "debug", skip(self), ret)] + pub fn get_section_idx(&self, section: Section) -> usize { + if self.header.num_toc_entries() == 1 { + 0 + } else { + match section { + Section::LfGlobal => 0, + Section::Lf { group } => 1 + group, + Section::HfGlobal => self.header.num_lf_groups() + 1, + Section::Hf { group, pass } => { + 2 + self.header.num_lf_groups() + self.header.num_groups() * pass + group + } + } + } + } + + #[instrument(level = "debug", skip_all)] + pub fn decode_lf_global(&mut self, br: &mut BitReader) -> Result<()> { + debug!(section_size = br.total_bits_available()); + assert!(self.lf_global.is_none()); + trace!(pos = br.total_bits_read()); + + let patches = if self.header.has_patches() { + info!("decoding patches"); + Some(PatchesDictionary::read( + br, + self.header.width as usize, + self.header.height as usize, + self.decoder_state.extra_channel_info().len(), + &self.decoder_state.reference_frames, + )?) + } else { + None + }; + + let splines = if self.header.has_splines() { + info!("decoding splines"); + Some(Splines::read(br, self.header.width * self.header.height)?) + } else { + None + }; + + let noise = if self.header.has_noise() { + info!("decoding noise"); + Some(Noise::read(br)?) + } else { + None + }; + + let lf_quant = LfQuantFactors::new(br)?; + debug!(?lf_quant); + + let quant_params = if self.header.encoding == Encoding::VarDCT { + info!("decoding VarDCT quantizer params"); + Some(QuantizerParams::read(br)?) + } else { + None + }; + debug!(?quant_params); + + let block_context_map = if self.header.encoding == Encoding::VarDCT { + info!("decoding block context map"); + Some(BlockContextMap::read(br)?) + } else { + None + }; + debug!(?block_context_map); + + let color_correlation_params = if self.header.encoding == Encoding::VarDCT { + info!("decoding color correlation params"); + Some(ColorCorrelationParams::read(br)?) + } else { + None + }; + debug!(?color_correlation_params); + + let tree = if br.read(1)? == 1 { + let size_limit = (1024 + + self.header.width as usize + * self.header.height as usize + * (self.modular_color_channels + + self.decoder_state.extra_channel_info().len()) + / 16) + .min(1 << 22); + Some(Tree::read(br, size_limit)?) + } else { + None + }; + + let modular_global = FullModularImage::read( + &self.header, + &self.decoder_state.file_header.image_metadata, + self.modular_color_channels, + &tree, + br, + )?; + + self.lf_global = Some(LfGlobalState { + patches, + splines, + noise, + lf_quant, + quant_params, + block_context_map, + color_correlation_params, + tree, + modular_global, + }); + + Ok(()) + } + + #[instrument(level = "debug", skip(self, br))] + pub fn decode_lf_group(&mut self, group: usize, br: &mut BitReader) -> Result<()> { + debug!(section_size = br.total_bits_available()); + let lf_global = self.lf_global.as_mut().unwrap(); + if self.header.encoding == Encoding::VarDCT && !self.header.has_lf_frame() { + info!("decoding VarDCT LF with group id {}", group); + decode_vardct_lf( + group, + &self.header, + &self.decoder_state.file_header.image_metadata, + &lf_global.tree, + lf_global.color_correlation_params.as_ref().unwrap(), + lf_global.quant_params.as_ref().unwrap(), + &lf_global.lf_quant, + lf_global.block_context_map.as_ref().unwrap(), + self.lf_image.as_mut().unwrap(), + &mut self.quant_lf, + br, + )?; + } + lf_global.modular_global.read_stream( + ModularStreamId::ModularLF(group), + &self.header, + &lf_global.tree, + br, + )?; + if self.header.encoding == Encoding::VarDCT { + info!("decoding HF metadata with group id {}", group); + let hf_meta = self.hf_meta.as_mut().unwrap(); + decode_hf_metadata( + group, + &self.header, + &self.decoder_state.file_header.image_metadata, + &lf_global.tree, + hf_meta, + br, + )?; + } + Ok(()) + } + + #[instrument(level = "debug", skip_all)] + pub fn decode_hf_global(&mut self, br: &mut BitReader) -> Result<()> { + debug!(section_size = br.total_bits_available()); + if self.header.encoding == Encoding::Modular { + return Ok(()); + } + let lf_global = self.lf_global.as_mut().unwrap(); + let mut dequant_matrices = DequantMatrices::decode(&self.header, lf_global, br)?; + dequant_matrices.ensure_computed(self.hf_meta.as_ref().unwrap().used_hf_types)?; + let block_context_map = lf_global.block_context_map.as_mut().unwrap(); + let num_histo_bits = self.header.num_groups().ceil_log2(); + let num_histograms: u32 = br.read(num_histo_bits)? as u32 + 1; + info!( + "Processing HFGlobal section with {} passes and {} histograms", + self.header.passes.num_passes, num_histograms + ); + let mut passes: Vec<PassState> = vec![]; + #[allow(unused_variables)] + for i in 0..self.header.passes.num_passes as usize { + let used_orders = match br.read(2)? { + 0 => 0x5f, + 1 => 0x13, + 2 => 0, + _ => br.read(coeff_order::NUM_ORDERS)?, + } as u32; + debug!(used_orders); + let coeff_orders = decode_coeff_orders(used_orders, br)?; + assert_eq!(coeff_orders.len(), 3 * coeff_order::NUM_ORDERS); + let num_contexts = num_histograms as usize * block_context_map.num_ac_contexts(); + info!( + "Deconding histograms for pass {} with {} contexts", + i, num_contexts + ); + let histograms = Histograms::decode(num_contexts, br, true)?; + debug!("Found {} histograms", histograms.num_histograms()); + passes.push(PassState { + coeff_orders, + histograms, + }); + } + let hf_coefficients = if passes.len() <= 1 { + None + } else { + let xs = GROUP_DIM * GROUP_DIM; + let ys = self.header.num_groups(); + Some(( + Image::new((xs, ys))?, + Image::new((xs, ys))?, + Image::new((xs, ys))?, + )) + }; + self.hf_global = Some(HfGlobalState { + num_histograms, + passes, + dequant_matrices, + hf_coefficients, + }); + Ok(()) + } + + pub fn build_render_pipeline( + decoder_state: &DecoderState, + frame_header: &FrameHeader, + lf_global: &LfGlobalState, + epf_sigma: &Option<Arc<Image<f32>>>, + ) -> Result<SimpleRenderPipeline> { + let num_channels = frame_header.num_extra_channels as usize + 3; + let num_temp_channels = if frame_header.has_noise() { 3 } else { 0 }; + let metadata = &decoder_state.file_header.image_metadata; + let mut pipeline = SimpleRenderPipelineBuilder::new( + num_channels + num_temp_channels, + frame_header.size_upsampled(), + frame_header.upsampling.ilog2() as usize, + frame_header.log_group_dim(), + ); + if frame_header.encoding == Encoding::Modular { + if decoder_state.file_header.image_metadata.xyb_encoded { + pipeline = + pipeline.add_stage(ConvertModularXYBToF32Stage::new(0, &lf_global.lf_quant))? + } else { + for i in 0..3 { + pipeline = + pipeline.add_stage(ConvertModularToF32Stage::new(i, metadata.bit_depth))?; + } + } + } + for i in 3..num_channels { + pipeline = pipeline.add_stage(ConvertModularToF32Stage::new(i, metadata.bit_depth))?; + } + + for c in 0..3 { + if frame_header.hshift(c) != 0 { + pipeline = pipeline.add_stage(HorizontalChromaUpsample::new(c))?; + } + if frame_header.vshift(c) != 0 { + pipeline = pipeline.add_stage(VerticalChromaUpsample::new(c))?; + } + } + + let filters = &frame_header.restoration_filter; + if filters.gab { + pipeline = pipeline + .add_stage(GaborishStage::new( + 0, + filters.gab_x_weight1, + filters.gab_x_weight2, + ))? + .add_stage(GaborishStage::new( + 1, + filters.gab_y_weight1, + filters.gab_y_weight2, + ))? + .add_stage(GaborishStage::new( + 2, + filters.gab_b_weight1, + filters.gab_b_weight2, + ))?; + } + + let rf = &frame_header.restoration_filter; + if rf.epf_iters >= 3 { + pipeline = pipeline.add_stage(Epf0Stage::new( + rf.epf_pass0_sigma_scale, + rf.epf_border_sad_mul, + rf.epf_channel_scale, + epf_sigma.as_ref().unwrap().clone(), + ))? + } + if rf.epf_iters >= 1 { + pipeline = pipeline.add_stage(Epf1Stage::new( + 1.0, + rf.epf_border_sad_mul, + rf.epf_channel_scale, + epf_sigma.as_ref().unwrap().clone(), + ))? + } + if rf.epf_iters >= 2 { + pipeline = pipeline.add_stage(Epf2Stage::new( + rf.epf_pass2_sigma_scale, + rf.epf_border_sad_mul, + rf.epf_channel_scale, + epf_sigma.as_ref().unwrap().clone(), + ))? + } + + let late_ec_upsample = frame_header.upsampling > 1 + && frame_header + .ec_upsampling + .iter() + .all(|x| *x == frame_header.upsampling); + + if !late_ec_upsample { + let transform_data = &decoder_state.file_header.transform_data; + for (ec, ec_up) in frame_header.ec_upsampling.iter().enumerate() { + if *ec_up > 1 { + pipeline = match *ec_up { + 2 => pipeline.add_stage(Upsample2x::new(transform_data, 3 + ec)), + 4 => pipeline.add_stage(Upsample4x::new(transform_data, 3 + ec)), + 8 => pipeline.add_stage(Upsample8x::new(transform_data, 3 + ec)), + _ => unreachable!(), + }?; + } + } + } + + if frame_header.has_patches() { + // TODO(szabadka): Avoid cloning everything. + pipeline = pipeline.add_stage(PatchesStage { + patches: lf_global.patches.clone().unwrap(), + extra_channels: metadata.extra_channel_info.clone(), + decoder_state: Arc::new(decoder_state.reference_frames.to_vec()), + })? + } + + if frame_header.has_splines() { + pipeline = pipeline.add_stage(SplinesStage::new( + lf_global.splines.clone().unwrap(), + frame_header.size(), + &lf_global.color_correlation_params.unwrap_or_default(), + ))? + } + + if frame_header.upsampling > 1 { + let transform_data = &decoder_state.file_header.transform_data; + let nb_channels = if late_ec_upsample { + 3 + frame_header.ec_upsampling.len() + } else { + 3 + }; + for c in 0..nb_channels { + pipeline = match frame_header.upsampling { + 2 => pipeline.add_stage(Upsample2x::new(transform_data, c)), + 4 => pipeline.add_stage(Upsample4x::new(transform_data, c)), + 8 => pipeline.add_stage(Upsample8x::new(transform_data, c)), + _ => unreachable!(), + }?; + } + } + + if frame_header.has_noise() { + pipeline = pipeline + .add_stage(ConvolveNoiseStage::new(num_channels))? + .add_stage(ConvolveNoiseStage::new(num_channels + 1))? + .add_stage(ConvolveNoiseStage::new(num_channels + 2))? + .add_stage(AddNoiseStage::new( + *lf_global.noise.as_ref().unwrap(), + lf_global.color_correlation_params.unwrap_or_default(), + num_channels, + ))?; + } + if frame_header.lf_level != 0 { + for i in 0..3 { + pipeline = pipeline.add_save_stage(SaveStage::<f32>::new( + SaveStageType::Lf, + i, + frame_header.size_upsampled(), + 1.0, + Orientation::Identity, + )?)?; + } + } + if frame_header.can_be_referenced && frame_header.save_before_ct { + for i in 0..num_channels { + pipeline = pipeline.add_save_stage(SaveStage::<f32>::new( + SaveStageType::Reference, + i, + frame_header.size_upsampled(), + 1.0, + Orientation::Identity, + )?)?; + } + } + + let mut linear = false; + let output_color_info = OutputColorInfo::from_header(&decoder_state.file_header)?; + if frame_header.do_ycbcr { + pipeline = pipeline.add_stage(YcbcrToRgbStage::new(0))?; + } else if decoder_state.file_header.image_metadata.xyb_encoded { + pipeline = pipeline.add_stage(XybStage::new(0, output_color_info.clone()))?; + if decoder_state.xyb_output_linear { + linear = true; + } else { + pipeline = + pipeline.add_stage(FromLinearStage::new(0, output_color_info.tf.clone()))?; + } + } + + if frame_header.needs_blending() { + if linear { + pipeline = + pipeline.add_stage(FromLinearStage::new(0, output_color_info.tf.clone()))?; + linear = false; + } + pipeline = pipeline.add_stage(BlendingStage::new( + frame_header, + &decoder_state.file_header, + &decoder_state.reference_frames, + )?)?; + pipeline = pipeline.add_stage(ExtendToImageDimensionsStage::new( + frame_header, + &decoder_state.file_header, + &decoder_state.reference_frames, + )?)?; + } + let image_size = &decoder_state.file_header.size; + let image_size = (image_size.xsize() as usize, image_size.ysize() as usize); + + if frame_header.can_be_referenced && !frame_header.save_before_ct { + if linear { + pipeline = + pipeline.add_stage(FromLinearStage::new(0, output_color_info.tf.clone()))?; + linear = false; + } + for i in 0..num_channels { + pipeline = pipeline.add_save_stage(SaveStage::<f32>::new( + SaveStageType::Reference, + i, + image_size, + 1.0, + Orientation::Identity, + )?)?; + } + } + + if decoder_state.render_spotcolors { + for (i, info) in decoder_state + .file_header + .image_metadata + .extra_channel_info + .iter() + .enumerate() + { + if info.ec_type == ExtraChannel::SpotColor { + pipeline = + pipeline.add_stage(SpotColorStage::new(i, info.spot_color.unwrap()))?; + } + } + } + + if frame_header.is_visible() { + let color_space = decoder_state + .file_header + .image_metadata + .color_encoding + .color_space; + let num_color_channels = if color_space == ColorSpace::Gray { + 1 + } else { + 3 + }; + for i in (0..num_color_channels).chain(3..num_channels) { + if decoder_state.render_spotcolors + && i > 3 + && decoder_state.file_header.image_metadata.extra_channel_info[i - 3].ec_type + == ExtraChannel::SpotColor + { + continue; + } + if decoder_state.file_header.image_metadata.xyb_encoded + && decoder_state.xyb_output_linear + && !linear + { + pipeline = + pipeline.add_stage(ToLinearStage::new(0, output_color_info.tf.clone()))?; + linear = true; + } + pipeline = pipeline.add_save_stage(SaveStage::<f32>::new( + SaveStageType::Output, + i, + image_size, + 255.0, + metadata.orientation, + )?)?; + } + } + pipeline.build() + } + + pub fn prepare_render_pipeline(&mut self) -> Result<()> { + let lf_global = self.lf_global.as_mut().unwrap(); + let epf_sigma = if self.header.restoration_filter.epf_iters > 0 { + let sigma_image = create_sigma_image(&self.header, lf_global, &self.hf_meta)?; + Some(Arc::new(sigma_image)) + } else { + None + }; + + let render_pipeline = + Self::build_render_pipeline(&self.decoder_state, &self.header, lf_global, &epf_sigma)?; + self.render_pipeline = Some(render_pipeline); + if self.decoder_state.enable_output { + lf_global.modular_global.process_output( + 0, + 0, + &self.header, + self.render_pipeline.as_mut().unwrap(), + )?; + for group in 0..self.header.num_lf_groups() { + lf_global.modular_global.process_output( + 1, + group, + &self.header, + self.render_pipeline.as_mut().unwrap(), + )?; + } + } + Ok(()) + } + + pub fn finalize_lf(&mut self) -> Result<()> { + if self.header.should_do_adaptive_lf_smoothing() { + let lf_global = self.lf_global.as_mut().unwrap(); + let lf_quant = &lf_global.lf_quant; + let inv_quant_lf = lf_global.quant_params.as_mut().unwrap().inv_quant_lf(); + adaptive_lf_smoothing( + [ + inv_quant_lf * lf_quant.quant_factors[0], + inv_quant_lf * lf_quant.quant_factors[1], + inv_quant_lf * lf_quant.quant_factors[2], + ], + self.lf_image.as_mut().unwrap(), + ) + } else { + Ok(()) + } + } + + #[instrument(level = "debug", skip(self, br))] + pub fn decode_hf_group(&mut self, group: usize, pass: usize, br: &mut BitReader) -> Result<()> { + debug!(section_size = br.total_bits_available()); + if self.header.has_noise() { + // TODO(sboukortt): consider making this a dedicated stage + let num_channels = self.header.num_extra_channels as usize + 3; + + let group_dim = self.header.group_dim() as u32; + let xsize_groups = self.header.size_groups().0; + let gx = (group % xsize_groups) as u32; + let gy = (group / xsize_groups) as u32; + // TODO(sboukortt): test upsampling+noise + let upsampling = self.header.upsampling; + let x0 = gx * upsampling * group_dim; + let y0 = gy * upsampling * group_dim; + // TODO(sboukortt): actual frame indices for the first two + let mut rng = Xorshift128Plus::new_with_seeds(1, 0, x0, y0); + let bits_to_float = |bits: u32| f32::from_bits((bits >> 9) | 0x3F800000); + let rp = self.render_pipeline.as_mut().unwrap(); + for i in 0..3 { + let mut buf = rp.get_buffer_for_group(num_channels + i, group)?; + let mut rect = buf.as_rect_mut(); + let (xsize, ysize) = rect.size(); + const FLOATS_PER_BATCH: usize = + Xorshift128Plus::N * size_of::<u64>() / size_of::<f32>(); + let mut batch = [0u64; Xorshift128Plus::N]; + + for y in 0..ysize { + let row = rect.row(y); + for batch_index in 0..xsize.div_ceil(FLOATS_PER_BATCH) { + rng.fill(&mut batch); + let batch_size = + (xsize - batch_index * FLOATS_PER_BATCH).min(FLOATS_PER_BATCH); + for i in 0..batch_size { + let x = FLOATS_PER_BATCH * batch_index + i; + let k = i / 2; + let high_bytes = i % 2 != 0; + let bits = if high_bytes { + ((batch[k] & 0xFFFFFFFF00000000) >> 32) as u32 + } else { + (batch[k] & 0xFFFFFFFF) as u32 + }; + row[x] = bits_to_float(bits); + } + } + } + + rp.set_buffer_for_group(num_channels + i, group, 1, buf); + } + } + let lf_global = self.lf_global.as_mut().unwrap(); + if self.header.encoding == Encoding::VarDCT { + info!("Decoding VarDCT group {group}, pass {pass}"); + let hf_global = self.hf_global.as_mut().unwrap(); + let hf_meta = self.hf_meta.as_mut().unwrap(); + let pixels = decode_vardct_group( + group, + pass, + &self.header, + lf_global, + hf_global, + hf_meta, + &self.lf_image, + &self.quant_lf, + &self + .decoder_state + .file_header + .transform_data + .opsin_inverse_matrix + .quant_biases, + br, + )?; + if self.decoder_state.enable_output + && pass + 1 == self.header.passes.num_passes as usize + { + for (c, p) in pixels.into_iter().enumerate() { + self.render_pipeline + .as_mut() + .unwrap() + .set_buffer_for_group(c, group, 1, p); + } + } + } + lf_global.modular_global.read_stream( + ModularStreamId::ModularHF { group, pass }, + &self.header, + &lf_global.tree, + br, + )?; + if self.decoder_state.enable_output { + lf_global.modular_global.process_output( + 2 + pass, + group, + &self.header, + self.render_pipeline.as_mut().unwrap(), + )?; + // TODO(veluca): pass actual output buffers. + self.render_pipeline.as_mut().unwrap().do_render(&mut [])?; + } + Ok(()) + } + + pub fn finalize(mut self) -> Result<FrameOutput> { + let mut output_frame_data = Vec::<Image<f32>>::new(); + let mut reference_frame_data = Vec::<Image<f32>>::new(); + let mut lf_frame_data = [ + Image::<f32>::new((0, 0))?, + Image::<f32>::new((0, 0))?, + Image::<f32>::new((0, 0))?, + ]; + + let mut lf_chan = 0; + if let Some(render_pipeline) = self.render_pipeline { + for stage in render_pipeline + .into_stages() + .into_iter() + .filter_map(|x| x.downcast::<SaveStage<f32>>().ok()) + { + match stage.stage_type { + SaveStageType::Output => { + output_frame_data.push(stage.into_buffer()); + } + SaveStageType::Reference => { + reference_frame_data.push(stage.into_buffer()); + } + SaveStageType::Lf => { + lf_frame_data[lf_chan] = stage.into_buffer(); + lf_chan += 1; + } + } + } + } + + if self.header.can_be_referenced { + info!("Saving frame in slot {}", self.header.save_as_reference); + self.decoder_state.reference_frames[self.header.save_as_reference as usize] = + Some(ReferenceFrame { + frame: reference_frame_data, + saved_before_color_transform: self.header.save_before_ct, + }); + } + + if self.header.lf_level != 0 { + self.decoder_state.lf_frames[(self.header.lf_level - 1) as usize] = Some(lf_frame_data); + } + let channels = if self.header.is_visible() { + Some(output_frame_data) + } else { + None + }; + let decoder_state = if self.header.is_last { + None + } else { + Some(self.decoder_state) + }; + let frame_output = FrameOutput { + decoder_state, + channels, + }; + Ok(frame_output) + } +} + +#[cfg(test)] +mod test { + use crate::api::{JxlDecoderOptions, test::create_output_buffers}; + use std::panic; + + use crate::{ + api::{JxlDecoder, ProcessingResult, states}, + error::Error, + features::spline::Point, + util::test::assert_almost_abs_eq, + }; + use test_log::test; + + use super::Frame; + + #[allow(clippy::type_complexity)] + fn read_frames( + mut input: &[u8], + callback: Box<dyn FnMut(&Frame, usize) -> Result<(), Error>>, + ) -> Result<usize, Error> { + let options = JxlDecoderOptions::default(); + let mut decoded_frames; + let mut initialized_decoder = + JxlDecoder::<states::Initialized>::new(options).with_frame_callback(callback); + + let mut decoder_with_image_info = loop { + let process_result = initialized_decoder.process(&mut input); + match process_result.unwrap() { + ProcessingResult::Complete { result } => break result, + ProcessingResult::NeedsMoreInput { fallback, .. } => { + if input.is_empty() { + panic!("Unexpected end of input while reading image info"); + } + initialized_decoder = fallback; + } + } + }; + + let basic_info = decoder_with_image_info.basic_info().clone(); + let pixel_format = decoder_with_image_info.current_pixel_format().clone(); + + loop { + let mut decoder_with_frame_info = loop { + let process_result = decoder_with_image_info.process(&mut input); + match process_result.unwrap() { + ProcessingResult::Complete { result } => break result, + ProcessingResult::NeedsMoreInput { fallback, .. } => { + if input.is_empty() { + panic!("Unexpected end of input while reading frame info"); + } + decoder_with_image_info = fallback; + } + } + }; + + create_output_buffers!(basic_info, pixel_format, output_buffers, output_slices); + + decoder_with_image_info = loop { + let process_result = + decoder_with_frame_info.process(&mut input, &mut output_slices); + match process_result.unwrap() { + ProcessingResult::Complete { result } => break result, + ProcessingResult::NeedsMoreInput { fallback, .. } => { + if input.is_empty() { + panic!("Unexpected end of input while decoding frame"); + } + decoder_with_frame_info = fallback; + } + } + }; + + decoded_frames = decoder_with_image_info.decoded_frames(); + + if !decoder_with_image_info.has_more_frames() { + break; + } + } + + Ok(decoded_frames) + } + + #[test] + fn splines() -> Result<(), Error> { + let verify_frame = move |frame: &Frame, _| { + let lf_global = frame.lf_global.as_ref().unwrap(); + let splines = lf_global.splines.as_ref().unwrap(); + assert_eq!(splines.quantization_adjustment, 0); + let expected_starting_points = [Point { x: 9.0, y: 54.0 }].to_vec(); + assert_eq!(splines.starting_points, expected_starting_points); + assert_eq!(splines.splines.len(), 1); + let spline = splines.splines[0].clone(); + let expected_control_points = [ + (109, 105), + (-130, -261), + (-66, 193), + (227, -52), + (-170, 290), + ] + .to_vec(); + assert_eq!(spline.control_points.clone(), expected_control_points); + + const EXPECTED_COLOR_DCT: [[i32; 32]; 3] = [ + { + let mut row = [0; 32]; + row[0] = 168; + row[1] = 119; + row + }, + { + let mut row = [0; 32]; + row[0] = 9; + row[2] = 7; + row + }, + { + let mut row = [0; 32]; + row[0] = -10; + row[1] = 7; + row + }, + ]; + assert_eq!(spline.color_dct, EXPECTED_COLOR_DCT); + const EXPECTED_SIGMA_DCT: [i32; 32] = { + let mut dct = [0; 32]; + dct[0] = 4; + dct[7] = 2; + dct + }; + assert_eq!(spline.sigma_dct, EXPECTED_SIGMA_DCT); + Ok(()) + }; + assert_eq!( + read_frames( + include_bytes!("../resources/test/splines.jxl"), + Box::new(verify_frame), + )?, + 1 + ); + Ok(()) + } + + #[test] + fn noise() -> Result<(), Error> { + let verify_frame = |frame: &Frame, _| { + let lf_global = frame.lf_global.as_ref().unwrap(); + let noise = lf_global.noise.as_ref().unwrap(); + let want_noise = [ + 0.000000, 0.000977, 0.002930, 0.003906, 0.005859, 0.006836, 0.008789, 0.010742, + ]; + for (index, noise_param) in want_noise.iter().enumerate() { + assert_almost_abs_eq(noise.lut[index], *noise_param, 1e-6); + } + Ok(()) + }; + assert_eq!( + read_frames( + include_bytes!("../resources/test/8x8_noise.jxl"), + Box::new(verify_frame), + )?, + 1 + ); + Ok(()) + } + + #[test] + fn patches() -> Result<(), Error> { + let verify_frame = |frame: &Frame, frame_index| { + if frame_index == 0 { + assert!(!frame.header().has_patches()); + assert!(frame.header().can_be_referenced); + } else if frame_index == 1 { + assert!(frame.header().has_patches()); + assert!(!frame.header().can_be_referenced); + } + Ok(()) + }; + assert_eq!( + read_frames( + include_bytes!("../resources/test/grayscale_patches_modular.jxl"), + Box::new(verify_frame), + )?, + 2 + ); + Ok(()) + } + + #[test] + fn multiple_lf_420() -> Result<(), Error> { + let verify_frame = |frame: &Frame, _| { + assert!(frame.header().is420()); + let Some(lf_image) = &frame.lf_image else { + panic!("no lf_image"); + }; + for y in 0..146 { + let sample_cb_row = lf_image[0].as_rect().row(y); + let sample_cr_row = lf_image[2].as_rect().row(y); + for x in 0..146 { + let sample_cb = sample_cb_row[x]; + let sample_cr = sample_cr_row[x]; + let no_chroma = sample_cb == 0.0 && sample_cr == 0.0; + if y < 128 || x < 128 { + assert!(!no_chroma); + } else { + assert!(no_chroma); + } + } + } + Ok(()) + }; + read_frames( + include_bytes!("../resources/test/multiple_lf_420.jxl"), + Box::new(verify_frame), + )?; + Ok(()) + } + + #[test] + fn xyb_grayscale_patches() -> Result<(), Error> { + let verify_frame = |frame: &Frame, frame_index| { + if frame_index == 0 { + assert_eq!( + frame.header.frame_type, + crate::headers::frame_header::FrameType::ReferenceOnly, + ); + assert_eq!( + frame.header.encoding, + crate::headers::frame_header::Encoding::Modular, + ); + assert_eq!(frame.modular_color_channels, 3); + } else { + assert!(frame.header.has_patches()); + assert_eq!(frame.modular_color_channels, 0); + } + Ok(()) + }; + assert_eq!( + read_frames( + include_bytes!("../resources/test/grayscale_patches_var_dct.jxl"), + Box::new(verify_frame), + )?, + 2 + ); + Ok(()) + } +} diff --git a/third_party/rust/jxl/src/frame/adaptive_lf_smoothing.rs b/third_party/rust/jxl/src/frame/adaptive_lf_smoothing.rs @@ -0,0 +1,98 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +use crate::error::Result; +use crate::image::Image; +use num_traits::abs; + +#[allow(clippy::excessive_precision)] +const W_SIDE: f32 = 0.20345139757231578; +#[allow(clippy::excessive_precision)] +const W_CORNER: f32 = 0.0334829185968739; +const W_CENTER: f32 = 1.0 - 4.0 * (W_SIDE + W_CORNER); + +fn compute_pixel_channel( + dc_factor: f32, + gap: f32, + x: usize, + row_top: &[f32], + row: &[f32], + row_bottom: &[f32], +) -> (f32, f32, f32) { + let tl = row_top[x - 1]; + let tc = row_top[x]; + let tr = row_top[x + 1]; + let ml = row[x - 1]; + let mc = row[x]; + let mr = row[x + 1]; + let bl = row_bottom[x - 1]; + let bc = row_bottom[x]; + let br = row_bottom[x + 1]; + let corner = tl + tr + bl + br; + let side = ml + mr + tc + bc; + let sm = corner * W_CORNER + side * W_SIDE + mc * W_CENTER; + (mc, sm, gap.max(abs((mc - sm) / dc_factor))) +} + +pub fn adaptive_lf_smoothing(lf_factors: [f32; 3], lf_image: &mut [Image<f32>; 3]) -> Result<()> { + let xsize = lf_image[0].size().0; + let ysize = lf_image[0].size().1; + if ysize <= 2 || xsize <= 2 { + return Ok(()); + } + let mut smoothed: [Image<f32>; 3] = [ + Image::<f32>::new((xsize, ysize))?, + Image::<f32>::new((xsize, ysize))?, + Image::<f32>::new((xsize, ysize))?, + ]; + for c in 0..3 { + for y in [0, ysize - 1] { + smoothed[c] + .as_rect_mut() + .row(y) + .copy_from_slice(lf_image[c].as_rect().row(y)); + } + } + for y in 1..ysize - 1 { + for x in [0, xsize - 1] { + for c in 0..3 { + smoothed[c].as_rect_mut().row(y)[x] = lf_image[c].as_rect().row(y)[x]; + } + } + for x in 1..xsize - 1 { + let gap = 0.5; + let (mc_x, sm_x, gap) = compute_pixel_channel( + lf_factors[0], + gap, + x, + lf_image[0].as_rect().row(y - 1), + lf_image[0].as_rect().row(y), + lf_image[0].as_rect().row(y + 1), + ); + let (mc_y, sm_y, gap) = compute_pixel_channel( + lf_factors[1], + gap, + x, + lf_image[1].as_rect().row(y - 1), + lf_image[1].as_rect().row(y), + lf_image[1].as_rect().row(y + 1), + ); + let (mc_b, sm_b, gap) = compute_pixel_channel( + lf_factors[2], + gap, + x, + lf_image[2].as_rect().row(y - 1), + lf_image[2].as_rect().row(y), + lf_image[2].as_rect().row(y + 1), + ); + let factor = (3.0 - 4.0 * gap).max(0.0); + smoothed[0].as_rect_mut().row(y)[x] = (sm_x - mc_x) * factor + mc_x; + smoothed[1].as_rect_mut().row(y)[x] = (sm_y - mc_y) * factor + mc_y; + smoothed[2].as_rect_mut().row(y)[x] = (sm_b - mc_b) * factor + mc_b; + } + } + *lf_image = smoothed; + Ok(()) +} diff --git a/third_party/rust/jxl/src/frame/block_context_map.rs b/third_party/rust/jxl/src/frame/block_context_map.rs @@ -0,0 +1,147 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +use crate::{ + bit_reader::BitReader, + entropy_coding::context_map::*, + entropy_coding::decode::unpack_signed, + error::{Error, Result}, + frame::coeff_order::NUM_ORDERS, +}; + +pub const NON_ZERO_BUCKETS: usize = 37; +pub const ZERO_DENSITY_CONTEXT_COUNT: usize = 458; + +pub const COEFF_FREQ_CONTEXT: [usize; 64] = [ + 0xBAD, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 15, 16, 16, 17, 17, 18, 18, 19, + 19, 20, 20, 21, 21, 22, 22, 23, 23, 23, 23, 24, 24, 24, 24, 25, 25, 25, 25, 26, 26, 26, 26, 27, + 27, 27, 27, 28, 28, 28, 28, 29, 29, 29, 29, 30, 30, 30, 30, +]; + +pub const COEFF_NUM_NONZERO_CONTEXT: [usize; 64] = [ + 0xBAD, 0, 31, 62, 62, 93, 93, 93, 93, 123, 123, 123, 123, 152, 152, 152, 152, 152, 152, 152, + 152, 180, 180, 180, 180, 180, 180, 180, 180, 180, 180, 180, 180, 206, 206, 206, 206, 206, 206, + 206, 206, 206, 206, 206, 206, 206, 206, 206, 206, 206, 206, 206, 206, 206, 206, 206, 206, 206, + 206, 206, 206, 206, 206, 206, +]; + +pub fn zero_density_context( + nonzeros_left: usize, + k: usize, + num_blocks: usize, + prev: usize, +) -> usize { + let nonzeros_left_norm = nonzeros_left.div_ceil(num_blocks); + let k_norm = k / num_blocks; + assert!((1..64).contains(&k_norm)); + assert!((1..64).contains(&nonzeros_left_norm)); + (COEFF_NUM_NONZERO_CONTEXT[nonzeros_left_norm] + COEFF_FREQ_CONTEXT[k_norm]) * 2 + prev +} + +#[derive(Debug)] +pub struct BlockContextMap { + pub lf_thresholds: [Vec<i32>; 3], + pub qf_thresholds: Vec<u32>, + pub context_map: Vec<u8>, + pub num_lf_contexts: usize, + pub num_contexts: usize, +} + +impl BlockContextMap { + pub fn num_ac_contexts(&self) -> usize { + self.num_contexts * (NON_ZERO_BUCKETS + ZERO_DENSITY_CONTEXT_COUNT) + } + pub fn read(br: &mut BitReader) -> Result<BlockContextMap, Error> { + if br.read(1)? == 1 { + Ok(BlockContextMap { + lf_thresholds: [vec![], vec![], vec![]], + qf_thresholds: vec![], + context_map: vec![ + 0, 1, 2, 2, 3, 3, 4, 5, 6, 6, 6, 6, 6, // + 7, 8, 9, 9, 10, 11, 12, 13, 14, 14, 14, 14, 14, // + 7, 8, 9, 9, 10, 11, 12, 13, 14, 14, 14, 14, 14, // + ], + num_lf_contexts: 1, + num_contexts: 15, + }) + } else { + let mut num_lf_contexts: usize = 1; + let mut lf_thresholds: [Vec<i32>; 3] = [vec![], vec![], vec![]]; + for thr in lf_thresholds.iter_mut() { + let num_lf_thresholds = br.read(4)? as usize; + let mut v: Vec<i32> = vec![0; num_lf_thresholds]; + for val in v.iter_mut() { + let uval = match br.read(2)? { + 0 => br.read(4)?, + 1 => br.read(8)? + 16, + 2 => br.read(16)? + 272, + _ => br.read(32)? + 65808, + }; + *val = unpack_signed(uval as u32) + } + *thr = v; + num_lf_contexts *= num_lf_thresholds + 1; + } + let num_qf_thresholds = br.read(4)? as usize; + let mut qf_thresholds: Vec<u32> = vec![0; num_qf_thresholds]; + for val in qf_thresholds.iter_mut() { + *val = match br.read(2)? { + 0 => br.read(2)?, + 1 => br.read(3)? + 4, + 2 => br.read(5)? + 12, + _ => br.read(8)? + 44, + } as u32 + + 1; + } + if num_lf_contexts * (num_qf_thresholds + 1) > 64 { + return Err(Error::BlockContextMapSizeTooBig( + num_lf_contexts, + num_qf_thresholds, + )); + } + let context_map_size = 3 * NUM_ORDERS * num_lf_contexts * (num_qf_thresholds + 1); + let context_map = decode_context_map(context_map_size, br)?; + assert_eq!(context_map.len(), context_map_size); + let num_contexts = *context_map.iter().max().unwrap() as usize + 1; + if num_contexts > 16 { + Err(Error::TooManyBlockContexts) + } else { + Ok(BlockContextMap { + lf_thresholds, + qf_thresholds, + context_map, + num_lf_contexts, + num_contexts, + }) + } + } + } + pub fn block_context(&self, lf_idx: usize, qf: u32, shape_id: usize, c: usize) -> usize { + let mut qf_idx: usize = 0; + for t in &self.qf_thresholds { + if qf > *t { + qf_idx += 1; + } + } + let mut idx = if c < 2 { c ^ 1 } else { 2 }; + idx = idx * NUM_ORDERS + shape_id; + idx = idx * (self.qf_thresholds.len() + 1) + qf_idx; + idx = idx * self.num_lf_contexts + lf_idx; + self.context_map[idx] as usize + } + pub fn nonzero_context(&self, nonzeros: usize, block_context: usize) -> usize { + let context: usize = if nonzeros < 8 { + nonzeros + } else if nonzeros < 64 { + 4 + nonzeros / 2 + } else { + 36 + }; + context * self.num_contexts + block_context + } + pub fn zero_density_context_offset(&self, block_context: usize) -> usize { + self.num_contexts * NON_ZERO_BUCKETS + ZERO_DENSITY_CONTEXT_COUNT * block_context + } +} diff --git a/third_party/rust/jxl/src/frame/coeff_order.rs b/third_party/rust/jxl/src/frame/coeff_order.rs @@ -0,0 +1,121 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +use crate::{ + BLOCK_DIM, BLOCK_SIZE, + bit_reader::BitReader, + entropy_coding::decode::SymbolReader, + error::Result, + frame::Histograms, + frame::transform_map::*, + headers::permutation::Permutation, + util::{CeilLog2, tracing_wrappers::*}, +}; + +use std::mem; + +pub const NUM_ORDERS: usize = 13; + +pub const TRANSFORM_TYPE_LUT: [HfTransformType; NUM_ORDERS] = [ + HfTransformType::DCT, + HfTransformType::IDENTITY, // a.k.a. "Hornuss" + HfTransformType::DCT16X16, + HfTransformType::DCT32X32, + HfTransformType::DCT8X16, + HfTransformType::DCT8X32, + HfTransformType::DCT16X32, + HfTransformType::DCT64X64, + HfTransformType::DCT32X64, + HfTransformType::DCT128X128, + HfTransformType::DCT64X128, + HfTransformType::DCT256X256, + HfTransformType::DCT128X256, +]; + +pub const NUM_PERMUTATION_CONTEXTS: usize = 8; + +pub fn natural_coeff_order(transform: HfTransformType) -> Vec<u32> { + let cx = covered_blocks_x(transform) as usize; + let cy = covered_blocks_y(transform) as usize; + let xsize: usize = cx * BLOCK_DIM; + assert!(cx >= cy); + // We compute the zigzag order for a cx x cx block, then discard all the + // lines that are not multiple of the ratio between cx and cy. + let xs = cx / cy; + let xsm = xs - 1; + let xss = xs.ceil_log2(); + let mut out: Vec<u32> = vec![0; cx * cy * BLOCK_SIZE]; + // First half of the block + let mut cur = cx * cy; + for i in 0..xsize { + for j in 0..(i + 1) { + let mut x = j; + let mut y = i - j; + if i % 2 != 0 { + mem::swap(&mut x, &mut y); + } + if (y & xsm) != 0 { + continue; + } + y >>= xss; + let val; + if x < cx && y < cy { + val = y * cx + x; + } else { + val = cur; + cur += 1; + } + out[val] = (y * xsize + x) as u32; + } + } + // Second half + for ir in 1..xsize { + let ip = xsize - ir; + let i = ip - 1; + for j in 0..(i + 1) { + let mut x = xsize - 1 - (i - j); + let mut y = xsize - 1 - j; + if i % 2 != 0 { + mem::swap(&mut x, &mut y); + } + if (y & xsm) != 0 { + continue; + } + y >>= xss; + let val = cur; + cur += 1; + out[val] = (y * xsize + x) as u32; + } + } + out +} + +pub fn decode_coeff_orders(used_orders: u32, br: &mut BitReader) -> Result<Vec<Permutation>> { + // TODO(szabadka): Compute natural coefficient orders only for those transform that are used. + let all_component_orders = 3 * NUM_ORDERS; + let mut permutations: Vec<Permutation> = (0..all_component_orders) + .map(|o| Permutation(natural_coeff_order(TRANSFORM_TYPE_LUT[o / 3]))) + .collect(); + if used_orders == 0 { + return Ok(permutations); + } + let histograms = Histograms::decode(NUM_PERMUTATION_CONTEXTS, br, true)?; + let mut reader = SymbolReader::new(&histograms, br, None)?; + for (ord, transform_type) in TRANSFORM_TYPE_LUT.iter().enumerate() { + if used_orders & (1 << ord) == 0 { + continue; + } + debug!(?transform_type); + let num_blocks = covered_blocks_x(*transform_type) * covered_blocks_y(*transform_type); + for c in 0..3 { + let size = num_blocks * BLOCK_SIZE as u32; + let permutation = Permutation::decode(size, num_blocks, &histograms, br, &mut reader)?; + let index = 3 * ord + c; + permutations[index].compose(&permutation); + } + } + reader.check_final_state(&histograms)?; + Ok(permutations) +} diff --git a/third_party/rust/jxl/src/frame/color_correlation_map.rs b/third_party/rust/jxl/src/frame/color_correlation_map.rs @@ -0,0 +1,91 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +use crate::{ + BLOCK_DIM, + bit_reader::BitReader, + error::{Error, Result}, +}; +use std::default::Default; + +pub const COLOR_TILE_DIM: usize = 64; + +const _: () = assert!(COLOR_TILE_DIM % BLOCK_DIM == 0); + +pub const COLOR_TILE_DIM_IN_BLOCKS: usize = COLOR_TILE_DIM / BLOCK_DIM; + +pub const DEFAULT_COLOR_FACTOR: u32 = 84; + +#[derive(Debug, Copy, Clone)] +pub struct ColorCorrelationParams { + pub color_factor: u32, + pub base_correlation_x: f32, + pub base_correlation_b: f32, + pub ytox_lf: i32, + pub ytob_lf: i32, +} + +impl Default for ColorCorrelationParams { + fn default() -> Self { + Self { + color_factor: DEFAULT_COLOR_FACTOR, + base_correlation_x: 0.0, + base_correlation_b: 1.0, + ytox_lf: 0, + ytob_lf: 0, + } + } +} + +impl ColorCorrelationParams { + pub fn read(br: &mut BitReader) -> Result<ColorCorrelationParams, Error> { + if br.read(1)? == 1 { + Ok(Self::default()) + } else { + let color_factor = match br.read(2)? { + 0 => DEFAULT_COLOR_FACTOR, + 1 => 256, + 2 => (br.read(8)? + 2) as u32, + _ => (br.read(16)? + 258) as u32, + }; + use half::f16; + let val_x = f16::from_bits(br.read(16)? as u16); + let val_b = f16::from_bits(br.read(16)? as u16); + if !val_x.is_finite() || !val_b.is_finite() { + return Err(Error::FloatNaNOrInf); + } + let base_correlation_x = val_x.to_f32(); + let base_correlation_b = val_b.to_f32(); + if base_correlation_x > 4.0 || base_correlation_b > 4.0 { + return Err(Error::BaseColorCorrelationOutOfRange); + } + let ytox_lf = br.read(8)? as i32 - 128; + let ytob_lf = br.read(8)? as i32 - 128; + Ok(ColorCorrelationParams { + color_factor, + base_correlation_x, + base_correlation_b, + ytox_lf, + ytob_lf, + }) + } + } + + pub fn y_to_x(&self, factor: i32) -> f32 { + self.base_correlation_x + (factor as f32) / (self.color_factor as f32) + } + + pub fn y_to_x_lf(&self) -> f32 { + self.y_to_x(self.ytox_lf) + } + + pub fn y_to_b(&self, factor: i32) -> f32 { + self.base_correlation_b + (factor as f32) / (self.color_factor as f32) + } + + pub fn y_to_b_lf(&self) -> f32 { + self.y_to_b(self.ytob_lf) + } +} diff --git a/third_party/rust/jxl/src/frame/group.rs b/third_party/rust/jxl/src/frame/group.rs @@ -0,0 +1,645 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +use num_traits::Float; + +use crate::{ + BLOCK_DIM, BLOCK_SIZE, GROUP_DIM, + bit_reader::BitReader, + entropy_coding::decode::SymbolReader, + error::{Error, Result}, + frame::{ + HfGlobalState, HfMetadata, LfGlobalState, block_context_map::*, + color_correlation_map::COLOR_TILE_DIM_IN_BLOCKS, quant_weights::DequantMatrices, + transform_map::*, + }, + headers::frame_header::FrameHeader, + image::{Image, ImageRect, Rect}, + util::{CeilLog2, tracing_wrappers::*}, + var_dct::{ + dct::{DCT1D, DCT1DImpl, compute_scaled_dct}, + dct_scales::{DctResampleScales, HasDctResampleScales, dct_total_resample_scale}, + transform::*, + }, +}; + +// Computes the lowest-frequency ROWSxCOLS-sized square in output, which is a +// DCT_ROWS*DCT_COLS-sized DCT block, by doing a ROWS*COLS DCT on the input +// block. +fn reinterpreting_dct< + const DCT_ROWS: usize, + const DCT_COLS: usize, + const ROWS: usize, + const COLS: usize, +>( + input: &ImageRect<f32>, + output: &mut [f32], + output_stride: usize, + block: &mut [f32], +) where + DctResampleScales<ROWS, DCT_ROWS>: HasDctResampleScales<ROWS>, + DctResampleScales<COLS, DCT_COLS>: HasDctResampleScales<COLS>, + DCT1DImpl<ROWS>: DCT1D, + DCT1DImpl<COLS>: DCT1D, +{ + let mut dct_input = [[0.0; COLS]; ROWS]; + #[allow(clippy::needless_range_loop)] + for y in 0..ROWS { + dct_input[y].copy_from_slice(&input.row(y)[0..COLS]); + } + compute_scaled_dct::<ROWS, COLS>(dct_input, block); + if ROWS < COLS { + for y in 0..ROWS { + for x in 0..COLS { + output[y * output_stride + x] = block[y * COLS + x] + * dct_total_resample_scale::<ROWS, DCT_ROWS>(y) + * dct_total_resample_scale::<COLS, DCT_COLS>(x); + } + } + } else { + for y in 0..COLS { + for x in 0..ROWS { + output[y * output_stride + x] = block[y * ROWS + x] + * dct_total_resample_scale::<COLS, DCT_COLS>(y) + * dct_total_resample_scale::<ROWS, DCT_ROWS>(x); + } + } + } +} + +fn lowest_frequencies_from_lf( + hf_type: HfTransformType, + lf: &ImageRect<f32>, + llf: &mut [f32], + scratch: &mut [f32], +) { + match hf_type { + HfTransformType::DCT16X8 => { + reinterpreting_dct::< + /*DCT_ROWS=*/ { 2 * BLOCK_DIM }, + /*DCT_COLS=*/ BLOCK_DIM, + /*ROWS=*/ 2, + /*COLS=*/ 1, + >(lf, llf, 2 * BLOCK_DIM, scratch); + } + HfTransformType::DCT8X16 => { + reinterpreting_dct::< + /*DCT_ROWS=*/ BLOCK_DIM, + /*DCT_COLS=*/ { 2 * BLOCK_DIM }, + /*ROWS=*/ 1, + /*COLS=*/ 2, + >(lf, llf, 2 * BLOCK_DIM, scratch); + } + HfTransformType::DCT16X16 => { + reinterpreting_dct::< + /*DCT_ROWS=*/ { 2 * BLOCK_DIM }, + /*DCT_COLS=*/ { 2 * BLOCK_DIM }, + /*ROWS=*/ 2, + /*COLS=*/ 2, + >(lf, llf, 2 * BLOCK_DIM, scratch); + } + HfTransformType::DCT32X8 => { + reinterpreting_dct::< + /*DCT_ROWS=*/ { 4 * BLOCK_DIM }, + /*DCT_COLS=*/ BLOCK_DIM, + /*ROWS=*/ 4, + /*COLS=*/ 1, + >(lf, llf, 4 * BLOCK_DIM, scratch); + } + HfTransformType::DCT8X32 => { + reinterpreting_dct::< + /*DCT_ROWS=*/ BLOCK_DIM, + /*DCT_COLS=*/ { 4 * BLOCK_DIM }, + /*ROWS=*/ 1, + /*COLS=*/ 4, + >(lf, llf, 4 * BLOCK_DIM, scratch); + } + HfTransformType::DCT32X16 => { + reinterpreting_dct::< + /*DCT_ROWS=*/ { 4 * BLOCK_DIM }, + /*DCT_COLS=*/ { 2 * BLOCK_DIM }, + /*ROWS=*/ 4, + /*COLS=*/ 2, + >(lf, llf, 4 * BLOCK_DIM, scratch); + } + HfTransformType::DCT16X32 => { + reinterpreting_dct::< + /*DCT_ROWS=*/ { 2 * BLOCK_DIM }, + /*DCT_COLS=*/ { 4 * BLOCK_DIM }, + /*ROWS=*/ 2, + /*COLS=*/ 4, + >(lf, llf, 4 * BLOCK_DIM, scratch); + } + HfTransformType::DCT32X32 => { + reinterpreting_dct::< + /*DCT_ROWS=*/ { 4 * BLOCK_DIM }, + /*DCT_COLS=*/ { 4 * BLOCK_DIM }, + /*ROWS=*/ 4, + /*COLS=*/ 4, + >(lf, llf, 4 * BLOCK_DIM, scratch); + } + HfTransformType::DCT64X32 => { + reinterpreting_dct::< + /*DCT_ROWS=*/ { 8 * BLOCK_DIM }, + /*DCT_COLS=*/ { 4 * BLOCK_DIM }, + /*ROWS=*/ 8, + /*COLS=*/ 4, + >(lf, llf, 8 * BLOCK_DIM, scratch); + } + HfTransformType::DCT32X64 => { + reinterpreting_dct::< + /*DCT_ROWS=*/ { 4 * BLOCK_DIM }, + /*DCT_COLS=*/ { 8 * BLOCK_DIM }, + /*ROWS=*/ 4, + /*COLS=*/ 8, + >(lf, llf, 8 * BLOCK_DIM, scratch); + } + HfTransformType::DCT64X64 => { + reinterpreting_dct::< + /*DCT_ROWS=*/ { 8 * BLOCK_DIM }, + /*DCT_COLS=*/ { 8 * BLOCK_DIM }, + /*ROWS=*/ 8, + /*COLS=*/ 8, + >(lf, llf, 8 * BLOCK_DIM, scratch); + } + HfTransformType::DCT128X64 => { + reinterpreting_dct::< + /*DCT_ROWS=*/ { 16 * BLOCK_DIM }, + /*DCT_COLS=*/ { 8 * BLOCK_DIM }, + /*ROWS=*/ 16, + /*COLS=*/ 8, + >(lf, llf, 16 * BLOCK_DIM, scratch); + } + HfTransformType::DCT64X128 => { + reinterpreting_dct::< + /*DCT_ROWS=*/ { 8 * BLOCK_DIM }, + /*DCT_COLS=*/ { 16 * BLOCK_DIM }, + /*ROWS=*/ 8, + /*COLS=*/ 16, + >(lf, llf, 16 * BLOCK_DIM, scratch); + } + HfTransformType::DCT128X128 => { + reinterpreting_dct::< + /*DCT_ROWS=*/ { 16 * BLOCK_DIM }, + /*DCT_COLS=*/ { 16 * BLOCK_DIM }, + /*ROWS=*/ 16, + /*COLS=*/ 16, + >(lf, llf, 16 * BLOCK_DIM, scratch); + } + HfTransformType::DCT256X128 => { + reinterpreting_dct::< + /*DCT_ROWS=*/ { 32 * BLOCK_DIM }, + /*DCT_COLS=*/ { 16 * BLOCK_DIM }, + /*ROWS=*/ 32, + /*COLS=*/ 16, + >(lf, llf, 32 * BLOCK_DIM, scratch); + } + HfTransformType::DCT128X256 => { + reinterpreting_dct::< + /*DCT_ROWS=*/ { 16 * BLOCK_DIM }, + /*DCT_COLS=*/ { 32 * BLOCK_DIM }, + /*ROWS=*/ 16, + /*COLS=*/ 32, + >(lf, llf, 32 * BLOCK_DIM, scratch); + } + HfTransformType::DCT256X256 => { + reinterpreting_dct::< + /*DCT_ROWS=*/ { 32 * BLOCK_DIM }, + /*DCT_COLS=*/ { 32 * BLOCK_DIM }, + /*ROWS=*/ 32, + /*COLS=*/ 32, + >(lf, llf, 32 * BLOCK_DIM, scratch); + } + HfTransformType::DCT + | HfTransformType::DCT2X2 + | HfTransformType::DCT4X4 + | HfTransformType::DCT4X8 + | HfTransformType::DCT8X4 + | HfTransformType::AFV0 + | HfTransformType::AFV1 + | HfTransformType::AFV2 + | HfTransformType::AFV3 + | HfTransformType::IDENTITY => { + llf[0] = lf.row(0)[0]; + } + } +} + +fn predict_num_nonzeros(nzeros_map: &Image<u32>, bx: usize, by: usize) -> usize { + if bx == 0 { + if by == 0 { + 32 + } else { + nzeros_map.as_rect().row(by - 1)[0] as usize + } + } else if by == 0 { + nzeros_map.as_rect().row(by)[bx - 1] as usize + } else { + (nzeros_map.as_rect().row(by - 1)[bx] + nzeros_map.as_rect().row(by)[bx - 1]).div_ceil(2) + as usize + } +} + +fn adjust_quant_bias(c: usize, quant_i: i32, biases: &[f32; 4]) -> f32 { + match quant_i { + 0 => 0.0, + 1 => biases[c], + -1 => -biases[c], + _ => { + let quant = quant_i as f32; + quant - biases[3] / quant + } + } +} + +#[allow(clippy::too_many_arguments)] +fn dequant_lane( + scaled_dequant_x: f32, + scaled_dequant_y: f32, + scaled_dequant_b: f32, + dequant_matrices: &[f32], + size: usize, + k: usize, + x_cc_mul: f32, + b_cc_mul: f32, + biases: &[f32; 4], + qblock: &[&[i32]; 3], + block: &mut [Vec<f32>; 3], +) { + let x_mul = dequant_matrices[k] * scaled_dequant_x; + let y_mul = dequant_matrices[size + k] * scaled_dequant_y; + let b_mul = dequant_matrices[2 * size + k] * scaled_dequant_b; + + let quantized_x = qblock[0][k]; + let quantized_y = qblock[1][k]; + let quantized_b = qblock[2][k]; + + let dequant_x_cc = adjust_quant_bias(0, quantized_x, biases) * x_mul; + let dequant_y = adjust_quant_bias(1, quantized_y, biases) * y_mul; + let dequant_b_cc = adjust_quant_bias(2, quantized_b, biases) * b_mul; + + let dequant_x = x_cc_mul * dequant_y + dequant_x_cc; + let dequant_b = b_cc_mul * dequant_y + dequant_b_cc; + block[0][k] = dequant_x; + block[1][k] = dequant_y; + block[2][k] = dequant_b; +} + +#[allow(clippy::too_many_arguments)] +fn dequant_block( + hf_type: HfTransformType, + inv_global_scale: f32, + quant: u32, + x_dm_multiplier: f32, + b_dm_multiplier: f32, + x_cc_mul: f32, + b_cc_mul: f32, + size: usize, + dequant_matrices: &DequantMatrices, + covered_blocks: usize, + lf: &Option<[ImageRect<f32>; 3]>, + biases: &[f32; 4], + qblock: &[&[i32]; 3], + block: &mut [Vec<f32>; 3], + scratch: &mut [f32], +) { + let scaled_dequant_y = inv_global_scale / (quant as f32); + + let scaled_dequant_x = scaled_dequant_y * x_dm_multiplier; + let scaled_dequant_b = scaled_dequant_y * b_dm_multiplier; + + let matrices = dequant_matrices.matrix(hf_type, 0); + + for k in 0..covered_blocks * BLOCK_SIZE { + dequant_lane( + scaled_dequant_x, + scaled_dequant_y, + scaled_dequant_b, + matrices, + size, + k, + x_cc_mul, + b_cc_mul, + biases, + qblock, + block, + ); + } + if let Some(lf) = lf.as_ref() { + for c in 0..3 { + lowest_frequencies_from_lf(hf_type, &lf[c], &mut block[c], scratch); + } + } +} + +#[allow(clippy::too_many_arguments)] +#[allow(clippy::type_complexity)] +pub fn decode_vardct_group( + group: usize, + pass: usize, + frame_header: &FrameHeader, + lf_global: &mut LfGlobalState, + hf_global: &mut HfGlobalState, + hf_meta: &HfMetadata, + lf_image: &Option<[Image<f32>; 3]>, + quant_lf: &Image<u8>, + quant_biases: &[f32; 4], + br: &mut BitReader, +) -> Result<[Image<f32>; 3], Error> { + let x_dm_multiplier = (1.0 / (1.25)).powf(frame_header.x_qm_scale as f32 - 2.0); + let b_dm_multiplier = (1.0 / (1.25)).powf(frame_header.b_qm_scale as f32 - 2.0); + + let num_histo_bits = hf_global.num_histograms.ceil_log2(); + let histogram_index: usize = br.read(num_histo_bits as usize)? as usize; + debug!(?histogram_index); + let mut reader = SymbolReader::new(&hf_global.passes[pass].histograms, br, None)?; + let block_group_rect = frame_header.block_group_rect(group); + let group_size = ( + block_group_rect.size.0 * BLOCK_DIM, + block_group_rect.size.1 * BLOCK_DIM, + ); + let mut pixels: [Image<f32>; 3] = [ + Image::new(( + group_size.0 >> frame_header.hshift(0), + group_size.1 >> frame_header.vshift(0), + ))?, + Image::new(( + group_size.0 >> frame_header.hshift(1), + group_size.1 >> frame_header.vshift(1), + ))?, + Image::new(( + group_size.0 >> frame_header.hshift(2), + group_size.1 >> frame_header.vshift(2), + ))?, + ]; + debug!(?block_group_rect); + let max_block_size = HfTransformType::VALUES + .iter() + .filter(|&transform_type| (hf_meta.used_hf_types & (1 << *transform_type as u32)) != 0) + .map(|&transform_type| { + BLOCK_SIZE + * covered_blocks_x(transform_type) as usize + * covered_blocks_y(transform_type) as usize + }) + .max() + .unwrap_or(0); + let mut scratch = vec![0.0; max_block_size]; + let color_correlation_params = lf_global.color_correlation_params.as_ref().unwrap(); + let cmap_rect = Rect { + origin: ( + block_group_rect.origin.0 / COLOR_TILE_DIM_IN_BLOCKS, + block_group_rect.origin.1 / COLOR_TILE_DIM_IN_BLOCKS, + ), + size: ( + block_group_rect.size.0.div_ceil(COLOR_TILE_DIM_IN_BLOCKS), + block_group_rect.size.1.div_ceil(COLOR_TILE_DIM_IN_BLOCKS), + ), + }; + let quant_params = lf_global.quant_params.as_ref().unwrap(); + let inv_global_scale = quant_params.inv_global_scale(); + let ytox_map = hf_meta.ytox_map.as_rect(); + let ytox_map_rect = ytox_map.rect(cmap_rect)?; + let ytob_map = hf_meta.ytob_map.as_rect(); + let ytob_map_rect = ytob_map.rect(cmap_rect)?; + let transform_map = hf_meta.transform_map.as_rect(); + let transform_map_rect = transform_map.rect(block_group_rect)?; + let raw_quant_map = hf_meta.raw_quant_map.as_rect(); + let raw_quant_map_rect = raw_quant_map.rect(block_group_rect)?; + let mut num_nzeros: [Image<u32>; 3] = [ + Image::new(( + block_group_rect.size.0 >> frame_header.hshift(0), + block_group_rect.size.1 >> frame_header.vshift(0), + ))?, + Image::new(( + block_group_rect.size.0 >> frame_header.hshift(1), + block_group_rect.size.1 >> frame_header.vshift(1), + ))?, + Image::new(( + block_group_rect.size.0 >> frame_header.hshift(2), + block_group_rect.size.1 >> frame_header.vshift(2), + ))?, + ]; + let quant_lf_rect = quant_lf.as_rect().rect(block_group_rect)?; + let block_context_map = lf_global.block_context_map.as_mut().unwrap(); + let context_offset = histogram_index * block_context_map.num_ac_contexts(); + let mut coeffs_storage; + let mut hf_coefficients_rects; + let coeffs = match hf_global.hf_coefficients.as_mut() { + Some(hf_coefficients) => { + hf_coefficients_rects = ( + hf_coefficients.0.as_rect_mut(), + hf_coefficients.1.as_rect_mut(), + hf_coefficients.2.as_rect_mut(), + ); + [ + hf_coefficients_rects.0.row(group), + hf_coefficients_rects.1.row(group), + hf_coefficients_rects.2.row(group), + ] + } + None => { + coeffs_storage = vec![0; 3 * GROUP_DIM * GROUP_DIM]; + let (coeffs_x, coeffs_y_b) = coeffs_storage.split_at_mut(GROUP_DIM * GROUP_DIM); + let (coeffs_y, coeffs_b) = coeffs_y_b.split_at_mut(GROUP_DIM * GROUP_DIM); + [coeffs_x, coeffs_y, coeffs_b] + } + }; + let shift_for_pass = if pass < frame_header.passes.shift.len() { + frame_header.passes.shift[pass] + } else { + 0 + }; + let mut coeffs_offset = 0; + let mut transform_buffer: [Vec<f32>; 3] = [ + vec![0.0; MAX_COEFF_AREA], + vec![0.0; MAX_COEFF_AREA], + vec![0.0; MAX_COEFF_AREA], + ]; + + let hshift = [ + frame_header.hshift(0), + frame_header.hshift(1), + frame_header.hshift(2), + ]; + let vshift = [ + frame_header.vshift(0), + frame_header.vshift(1), + frame_header.vshift(2), + ]; + let lf = match lf_image.as_ref() { + None => None, + Some(lf_planes) => { + let r: [Rect; 3] = core::array::from_fn(|i| Rect { + origin: ( + block_group_rect.origin.0 >> hshift[i], + block_group_rect.origin.1 >> vshift[i], + ), + size: ( + block_group_rect.size.0 >> hshift[i], + block_group_rect.size.1 >> vshift[i], + ), + }); + + let [lf_x, lf_y, lf_b] = lf_planes.each_ref(); + Some([ + lf_x.as_rect().rect(r[0])?, + lf_y.as_rect().rect(r[1])?, + lf_b.as_rect().rect(r[2])?, + ]) + } + }; + for by in 0..block_group_rect.size.1 { + let sby = [by >> vshift[0], by >> vshift[1], by >> vshift[2]]; + let ty = by / COLOR_TILE_DIM_IN_BLOCKS; + + let row_cmap_x = ytox_map_rect.row(ty); + let row_cmap_b = ytob_map_rect.row(ty); + + for bx in 0..block_group_rect.size.0 { + let sbx = [bx >> hshift[0], bx >> hshift[1], bx >> hshift[2]]; + let tx = bx / COLOR_TILE_DIM_IN_BLOCKS; + let x_cc_mul = color_correlation_params.y_to_x(row_cmap_x[tx] as i32); + let b_cc_mul = color_correlation_params.y_to_b(row_cmap_b[tx] as i32); + let raw_quant = raw_quant_map_rect.row(by)[bx] as u32; + let quant_lf = quant_lf_rect.row(by)[bx] as usize; + let raw_transform_id = transform_map_rect.row(by)[bx]; + let transform_id = raw_transform_id & 127; + let is_first_block = raw_transform_id >= 128; + if !is_first_block { + continue; + } + let lf_rects = match lf.as_ref() { + None => None, + Some(lf) => { + let [lf_x, lf_y, lf_b] = lf.each_ref(); + Some([ + lf_x.rect(Rect { + origin: (sbx[0], sby[0]), + size: (lf_x.size().0 - sbx[0], lf_x.size().1 - sby[0]), + })?, + lf_y.rect(Rect { + origin: (sbx[1], sby[1]), + size: (lf_y.size().0 - sbx[1], lf_y.size().1 - sby[1]), + })?, + lf_b.rect(Rect { + origin: (sbx[2], sby[2]), + size: (lf_b.size().0 - sbx[2], lf_b.size().1 - sby[2]), + })?, + ]) + } + }; + + let transform_type = HfTransformType::from_usize(transform_id as usize)?; + let cx = covered_blocks_x(transform_type) as usize; + let cy = covered_blocks_y(transform_type) as usize; + let shape_id = block_shape_id(transform_type) as usize; + let block_size = (cx * BLOCK_DIM, cy * BLOCK_DIM); + let block_rect = Rect { + origin: (bx * BLOCK_DIM, by * BLOCK_DIM), + size: block_size, + }; + let num_blocks = cx * cy; + let num_coeffs = num_blocks * BLOCK_SIZE; + for c in [1, 0, 2] { + if (sbx[c] << hshift[c]) != bx || (sby[c] << vshift[c] != by) { + continue; + } + trace!( + "Decoding block ({},{}) channel {} with {}x{} block transform {} (shape id {})", + sbx[c], sby[c], c, cx, cy, transform_id, shape_id + ); + let predicted_nzeros = predict_num_nonzeros(&num_nzeros[c], sbx[c], sby[c]); + let block_context = + block_context_map.block_context(quant_lf, raw_quant, shape_id, c); + let nonzero_context = block_context_map + .nonzero_context(predicted_nzeros, block_context) + + context_offset; + let mut nonzeros = + reader.read_unsigned(&hf_global.passes[pass].histograms, br, nonzero_context)? + as usize; + trace!( + "block ({},{},{c}) predicted_nzeros: {predicted_nzeros} \ + nzero_ctx: {nonzero_context} (offset: {context_offset}) \ + nzeros: {nonzeros}", + sbx[c], sby[c] + ); + if nonzeros + num_blocks > num_coeffs { + return Err(Error::InvalidNumNonZeros(nonzeros, num_blocks)); + } + for iy in 0..cy { + for ix in 0..cx { + num_nzeros[c].as_rect_mut().row(sby[c] + iy)[sbx[c] + ix] = + nonzeros.div_ceil(num_blocks) as u32; + } + } + let histo_offset = + block_context_map.zero_density_context_offset(block_context) + context_offset; + let mut prev = if nonzeros > num_coeffs / 16 { 0 } else { 1 }; + for k in num_blocks..num_coeffs { + if nonzeros == 0 { + break; + } + let ctx = histo_offset + zero_density_context(nonzeros, k, num_blocks, prev); + let coeff = reader.read_signed(&hf_global.passes[pass].histograms, br, ctx)? + << shift_for_pass; + prev = if coeff != 0 { 1 } else { 0 }; + nonzeros -= prev; + let coeff_index = + hf_global.passes[pass].coeff_orders[shape_id * 3 + c][k] as usize; + coeffs[c][coeffs_offset + coeff_index] += coeff; + } + if nonzeros != 0 { + return Err(Error::EndOfBlockResidualNonZeros(nonzeros)); + } + } + let qblock = [ + &coeffs[0][coeffs_offset..], + &coeffs[1][coeffs_offset..], + &coeffs[2][coeffs_offset..], + ]; + dequant_block( + transform_type, + inv_global_scale, + raw_quant, + x_dm_multiplier, + b_dm_multiplier, + x_cc_mul, + b_cc_mul, + num_coeffs, + &hf_global.dequant_matrices, + num_blocks, + &lf_rects, + quant_biases, + &qblock, + &mut transform_buffer, + &mut scratch, + ); + for c in [1, 0, 2] { + if (sbx[c] << hshift[c]) != bx || (sby[c] << vshift[c] != by) { + continue; + } + transform_to_pixels(transform_type, &mut transform_buffer[c])?; + let mut output = pixels[c].as_rect_mut(); + let downsampled_rect = Rect { + origin: ( + block_rect.origin.0 >> hshift[c], + block_rect.origin.1 >> vshift[c], + ), + size: block_rect.size, + }; + let mut output_rect = output.rect(downsampled_rect)?; + for i in 0..downsampled_rect.size.1 { + let offset = i * downsampled_rect.size.0; + output_rect.row(i).copy_from_slice( + &transform_buffer[c][offset..offset + downsampled_rect.size.0], + ); + } + } + coeffs_offset += num_coeffs; + } + } + reader.check_final_state(&hf_global.passes[pass].histograms)?; + Ok(pixels) +} diff --git a/third_party/rust/jxl/src/frame/modular.rs b/third_party/rust/jxl/src/frame/modular.rs @@ -0,0 +1,842 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +use std::{cell::RefCell, cmp::min, fmt::Debug}; + +use crate::{ + bit_reader::BitReader, + error::{Error, Result}, + frame::{ + ColorCorrelationParams, HfMetadata, + block_context_map::BlockContextMap, + quantizer::{self, LfQuantFactors, QuantizerParams}, + transform_map::*, + }, + headers::{ + ImageMetadata, JxlHeader, bit_depth::BitDepth, frame_header::FrameHeader, + modular::GroupHeader, + }, + image::{Image, Rect}, + render::{RenderPipeline, SimpleRenderPipeline}, + util::{CeilLog2, tracing_wrappers::*}, +}; + +mod borrowed_buffers; +pub(crate) mod decode; +mod predict; +mod transforms; +mod tree; + +use borrowed_buffers::with_buffers; +pub use decode::ModularStreamId; +use decode::decode_modular_subbitstream; +pub use predict::Predictor; +use transforms::{TransformStepChunk, make_grids}; +pub use tree::Tree; + +#[derive(Clone, PartialEq, Eq, Copy)] +struct ChannelInfo { + // The index of the output channel in the render pipeline, or -1 for non-output channels. + output_channel_idx: isize, + // width, height + size: (usize, usize), + shift: Option<(usize, usize)>, // None for meta-channels + bit_depth: BitDepth, +} + +impl Debug for ChannelInfo { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{}x{}", self.size.0, self.size.1)?; + if let Some(shift) = self.shift { + write!(f, "(shift {},{})", shift.0, shift.1)?; + } else { + write!(f, "(meta)")?; + } + write!(f, "{:?}", self.bit_depth)?; + if self.output_channel_idx >= 0 { + write!(f, "(output channel {})", self.output_channel_idx)?; + } + Ok(()) + } +} + +impl ChannelInfo { + fn is_meta(&self) -> bool { + self.shift.is_none() + } + + fn is_meta_or_small(&self, group_dim: usize) -> bool { + self.is_meta() || (self.size.0 <= group_dim && self.size.1 <= group_dim) + } + + fn is_shift_in_range(&self, min: usize, max: usize) -> bool { + assert!(min <= max); + self.shift.is_some_and(|(a, b)| { + let shift = a.min(b); + min <= shift && shift <= max + }) + } + + fn is_equivalent(&self, other: &ChannelInfo) -> bool { + self.size == other.size && self.shift == other.shift && self.bit_depth == other.bit_depth + } +} + +#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Copy)] +enum ModularGridKind { + // Single big channel. + None, + // 2048x2048 image-pixels (if modular_group_shift == 1). + Lf, + // 256x256 image-pixels (if modular_group_shift == 1). + Hf, +} + +impl ModularGridKind { + fn grid_dim(&self, frame_header: &FrameHeader, shift: (usize, usize)) -> (usize, usize) { + let group_dim = match self { + ModularGridKind::None => 0, + ModularGridKind::Lf => frame_header.lf_group_dim(), + ModularGridKind::Hf => frame_header.group_dim(), + }; + (group_dim >> shift.0, group_dim >> shift.1) + } + fn grid_shape(&self, frame_header: &FrameHeader) -> (usize, usize) { + match self { + ModularGridKind::None => (1, 1), + ModularGridKind::Lf => frame_header.size_lf_groups(), + ModularGridKind::Hf => frame_header.size_groups(), + } + } +} + +// All the information on a specific buffer needed by Modular decoding. +#[derive(Debug)] +pub(crate) struct ModularChannel { + // Actual pixel buffer. + pub data: Image<i32>, + // Holds additional information such as the weighted predictor's error channel's last row for + // the transform chunk that produced this buffer. + auxiliary_data: Option<Image<i32>>, + // Shift of the channel (None if this is a meta-channel). + shift: Option<(usize, usize)>, + bit_depth: BitDepth, +} + +impl ModularChannel { + pub fn new(size: (usize, usize), bit_depth: BitDepth) -> Result<Self> { + Self::new_with_shift(size, Some((0, 0)), bit_depth) + } + + fn new_with_shift( + size: (usize, usize), + shift: Option<(usize, usize)>, + bit_depth: BitDepth, + ) -> Result<Self> { + Ok(ModularChannel { + data: Image::new(size)?, + auxiliary_data: None, + shift, + bit_depth, + }) + } + + fn try_clone(&self) -> Result<Self> { + Ok(ModularChannel { + data: self.data.try_clone()?, + auxiliary_data: self + .auxiliary_data + .as_ref() + .map(Image::try_clone) + .transpose()?, + shift: self.shift, + bit_depth: self.bit_depth, + }) + } + + fn channel_info(&self) -> ChannelInfo { + ChannelInfo { + output_channel_idx: -1, + size: self.data.size(), + shift: self.shift, + bit_depth: self.bit_depth, + } + } +} + +// Note: this type uses interior mutability to get mutable references to multiple buffers at once. +// In principle, this is not needed, but the overhead should be minimal so using `unsafe` here is +// probably not worth it. +#[derive(Debug)] +struct ModularBuffer { + data: RefCell<Option<ModularChannel>>, + // Number of times this buffer will be used, *including* when it is used for output. + remaining_uses: usize, + used_by_transforms: Vec<usize>, + size: (usize, usize), +} + +impl ModularBuffer { + // Gives out a copy of the buffer + auxiliary buffer, marking the buffer as used. + // If this was the last usage of the buffer, does not actually copy the buffer. + fn get_buffer(&mut self) -> Result<ModularChannel> { + self.remaining_uses = self.remaining_uses.checked_sub(1).unwrap(); + if self.remaining_uses == 0 { + Ok(self.data.borrow_mut().take().unwrap()) + } else { + Ok(self + .data + .borrow() + .as_ref() + .map(ModularChannel::try_clone) + .transpose()? + .unwrap()) + } + } + + fn mark_used(&mut self) { + self.remaining_uses = self.remaining_uses.checked_sub(1).unwrap(); + if self.remaining_uses == 0 { + *self.data.borrow_mut() = None; + } + } +} + +#[derive(Debug)] +struct ModularBufferInfo { + info: ChannelInfo, + // The index of coded channel in the bit-stream, or -1 for non-coded channels. + coded_channel_id: isize, + #[cfg_attr(not(feature = "tracing"), allow(dead_code))] + description: String, + grid_kind: ModularGridKind, + grid_shape: (usize, usize), + buffer_grid: Vec<ModularBuffer>, +} + +impl ModularBufferInfo { + fn get_grid_idx( + &self, + output_grid_kind: ModularGridKind, + output_grid_pos: (usize, usize), + ) -> usize { + let grid_pos = match (output_grid_kind, self.grid_kind) { + (_, ModularGridKind::None) => (0, 0), + (ModularGridKind::Lf, ModularGridKind::Lf) + | (ModularGridKind::Hf, ModularGridKind::Hf) => output_grid_pos, + (ModularGridKind::Hf, ModularGridKind::Lf) => { + (output_grid_pos.0 / 8, output_grid_pos.1 / 8) + } + _ => unreachable!("invalid combination of output grid kind and buffer grid kind"), + }; + self.grid_shape.0 * grid_pos.1 + grid_pos.0 + } + fn get_grid_rect( + &self, + frame_header: &FrameHeader, + output_grid_kind: ModularGridKind, + output_grid_pos: (usize, usize), + ) -> Rect { + let chan_size = self.info.size; + if output_grid_kind == ModularGridKind::None { + assert_eq!(self.grid_kind, output_grid_kind); + return Rect { + origin: (0, 0), + size: chan_size, + }; + } + let shift = self.info.shift.unwrap(); + let grid_dim = output_grid_kind.grid_dim(frame_header, shift); + let bx = output_grid_pos.0 * grid_dim.0; + let by = output_grid_pos.1 * grid_dim.1; + let size = ( + (chan_size.0 - bx).min(grid_dim.0), + (chan_size.1 - by).min(grid_dim.1), + ); + let origin = match (output_grid_kind, self.grid_kind) { + (ModularGridKind::Lf, ModularGridKind::Lf) + | (ModularGridKind::Hf, ModularGridKind::Hf) => (0, 0), + (_, ModularGridKind::None) => (bx, by), + (ModularGridKind::Hf, ModularGridKind::Lf) => { + let lf_grid_dim = self.grid_kind.grid_dim(frame_header, shift); + (bx % lf_grid_dim.0, by % lf_grid_dim.1) + } + _ => unreachable!("invalid combination of output grid kind and buffer grid kind"), + }; + Rect { origin, size } + } +} + +/// A modular image is a sequence of channels to which one or more transforms might have been +/// applied. We represent a modular image as a list of buffers, some of which are coded in the +/// bitstream; other buffers are obtained as the output of one of the transformation steps. +/// Some buffers are marked as `output`: those are the buffers corresponding to the pre-transform +/// image channels. +/// The buffers are internally divided in grids, matching the sizes of the groups they are coded +/// in (with appropriate shifts), or the size of the data produced by applying the appropriate +/// transforms to each of the groups in the input of the transforms. +#[derive(Debug)] +pub struct FullModularImage { + buffer_info: Vec<ModularBufferInfo>, + transform_steps: Vec<TransformStepChunk>, + // List of buffer indices of the channels of the modular image encoded in each kind of section. + // In order, LfGlobal, LfGroup, HfGroup(pass 0), ..., HfGroup(last pass). + section_buffer_indices: Vec<Vec<usize>>, + modular_color_channels: usize, +} + +impl FullModularImage { + #[instrument(level = "debug", skip_all)] + pub fn read( + frame_header: &FrameHeader, + image_metadata: &ImageMetadata, + modular_color_channels: usize, + global_tree: &Option<Tree>, + br: &mut BitReader, + ) -> Result<Self> { + let mut channels = vec![]; + for c in 0..modular_color_channels { + let shift = (frame_header.hshift(c), frame_header.vshift(c)); + let size = frame_header.size(); + channels.push(ChannelInfo { + output_channel_idx: c as isize, + size: (size.0.div_ceil(1 << shift.0), size.1.div_ceil(1 << shift.1)), + shift: Some(shift), + bit_depth: image_metadata.bit_depth, + }); + } + + for (idx, ecups) in frame_header.ec_upsampling.iter().enumerate() { + let shift_ec = ecups.ceil_log2(); + let shift_color = frame_header.upsampling.ceil_log2(); + let shift = shift_ec + .checked_sub(shift_color) + .expect("ec_upsampling >= upsampling should be checked in frame header") + as usize; + let size = frame_header.size_upsampled(); + let size = ( + size.0.div_ceil(*ecups as usize), + size.1.div_ceil(*ecups as usize), + ); + channels.push(ChannelInfo { + output_channel_idx: 3 + idx as isize, + size, + shift: Some((shift, shift)), + bit_depth: image_metadata.bit_depth, + }); + } + + #[cfg(feature = "tracing")] + for (i, ch) in channels.iter().enumerate() { + trace!("Modular channel {i}: {ch:?}"); + } + + if channels.is_empty() { + return Ok(Self { + buffer_info: vec![], + transform_steps: vec![], + section_buffer_indices: vec![vec![]; 2 + frame_header.passes.num_passes as usize], + modular_color_channels, + }); + } + + trace!("reading modular header"); + let header = GroupHeader::read(br)?; + + let (mut buffer_info, transform_steps) = + transforms::apply::meta_apply_transforms(&channels, &header)?; + + // Assign each (channel, group) pair present in the bitstream to the section in which it + // will be decoded. + let mut section_buffer_indices: Vec<Vec<usize>> = vec![]; + + let mut sorted_buffers: Vec<_> = buffer_info + .iter() + .enumerate() + .filter_map(|(i, b)| { + if b.coded_channel_id >= 0 { + Some((b.coded_channel_id, i)) + } else { + None + } + }) + .collect(); + + sorted_buffers.sort_by_key(|x| x.0); + + section_buffer_indices.push( + sorted_buffers + .iter() + .take_while(|x| { + buffer_info[x.1] + .info + .is_meta_or_small(frame_header.group_dim()) + }) + .map(|x| x.1) + .collect(), + ); + + section_buffer_indices.push( + sorted_buffers + .iter() + .skip_while(|x| { + buffer_info[x.1] + .info + .is_meta_or_small(frame_header.group_dim()) + }) + .filter(|x| buffer_info[x.1].info.is_shift_in_range(3, usize::MAX)) + .map(|x| x.1) + .collect(), + ); + + for pass in 0..frame_header.passes.num_passes as usize { + let (min_shift, max_shift) = frame_header.passes.downsampling_bracket(pass); + section_buffer_indices.push( + sorted_buffers + .iter() + .skip_while(|x| { + buffer_info[x.1] + .info + .is_meta_or_small(frame_header.group_dim()) + }) + .filter(|x| { + buffer_info[x.1] + .info + .is_shift_in_range(min_shift, max_shift) + }) + .map(|x| x.1) + .collect(), + ); + } + + // Ensure that the channel list in each group is sorted by actual channel ID. + for list in section_buffer_indices.iter_mut() { + list.sort_by_key(|x| buffer_info[*x].coded_channel_id); + } + + trace!(?section_buffer_indices); + #[cfg(feature = "tracing")] + for (section, indices) in section_buffer_indices.iter().enumerate() { + let section_name = match section { + 0 => "LF global".to_string(), + 1 => "LF groups".to_string(), + _ => format!("HF groups, pass {}", section - 2), + }; + trace!("Coded modular channels in {section_name}"); + for i in indices { + let bi = &buffer_info[*i]; + trace!( + "Channel {i} {:?} coded id: {}", + bi.info, bi.coded_channel_id + ); + } + } + + let transform_steps = make_grids( + frame_header, + transform_steps, + §ion_buffer_indices, + &mut buffer_info, + ); + + #[cfg(feature = "tracing")] + for (i, bi) in buffer_info.iter().enumerate() { + trace!( + "Channel {i} {:?} coded_id: {} '{}' {:?} grid {:?}", + bi.info, bi.coded_channel_id, bi.description, bi.grid_kind, bi.grid_shape + ); + for (pos, buf) in bi.buffer_grid.iter().enumerate() { + trace!( + "Channel {i} grid {pos} ({}, {}) size: {:?}, uses: {}, used_by: {:?}", + pos % bi.grid_shape.0, + pos / bi.grid_shape.0, + buf.size, + buf.remaining_uses, + buf.used_by_transforms + ); + } + } + + #[cfg(feature = "tracing")] + for (i, ts) in transform_steps.iter().enumerate() { + trace!("Transform {i}: {ts:?}"); + } + + with_buffers(&buffer_info, §ion_buffer_indices[0], 0, |bufs| { + decode_modular_subbitstream( + bufs, + ModularStreamId::GlobalData.get_id(frame_header), + Some(header), + global_tree, + br, + ) + })?; + + Ok(FullModularImage { + buffer_info, + transform_steps, + section_buffer_indices, + modular_color_channels, + }) + } + + #[allow(clippy::type_complexity)] + #[instrument(level = "debug", skip(self, frame_header, global_tree, br), ret)] + pub fn read_stream( + &mut self, + stream: ModularStreamId, + frame_header: &FrameHeader, + global_tree: &Option<Tree>, + br: &mut BitReader, + ) -> Result<()> { + if self.buffer_info.is_empty() { + info!("No modular channels to decode"); + return Ok(()); + } + let (section_id, grid) = match stream { + ModularStreamId::ModularLF(group) => (1, group), + ModularStreamId::ModularHF { pass, group } => (2 + pass, group), + _ => { + unreachable!( + "read_stream should only be used for streams that are part of the main Modular image" + ); + } + }; + + with_buffers( + &self.buffer_info, + &self.section_buffer_indices[section_id], + grid, + |bufs| { + decode_modular_subbitstream( + bufs, + stream.get_id(frame_header), + None, + global_tree, + br, + ) + }, + )?; + Ok(()) + } + + pub fn process_output( + &mut self, + section_id: usize, + grid: usize, + frame_header: &FrameHeader, + render_pipeline: &mut SimpleRenderPipeline, + ) -> Result<()> { + let mut maybe_output = |bi: &mut ModularBufferInfo, grid: usize| -> Result<()> { + if bi.info.output_channel_idx >= 0 { + let chan = bi.info.output_channel_idx as usize; + debug!("Rendering channel {chan:?}, grid position {grid}"); + let buf = bi.buffer_grid[grid].get_buffer()?; + // TODO(veluca): figure out what to do with passes here. + if chan == 0 && self.modular_color_channels == 1 { + for i in 0..2 { + render_pipeline.set_buffer_for_group( + i, + grid, + 1, + buf.data.as_rect().to_image()?, + ); + } + render_pipeline.set_buffer_for_group(2, grid, 1, buf.data); + } else { + render_pipeline.set_buffer_for_group(chan, grid, 1, buf.data); + } + } + Ok(()) + }; + + let mut new_ready_transform_chunks = vec![]; + for buf in self.section_buffer_indices[section_id].iter().copied() { + maybe_output(&mut self.buffer_info[buf], grid)?; + let new_chunks = self.buffer_info[buf].buffer_grid[grid] + .used_by_transforms + .to_vec(); + trace!("Buffer {buf} grid position {grid} used by chunks {new_chunks:?}"); + new_ready_transform_chunks.extend(new_chunks); + } + + trace!(?new_ready_transform_chunks); + + while let Some(tfm) = new_ready_transform_chunks.pop() { + trace!("tfm = {tfm} chunk = {:?}", self.transform_steps[tfm]); + for (new_buf, new_grid) in + self.transform_steps[tfm].dep_ready(frame_header, &mut self.buffer_info)? + { + maybe_output(&mut self.buffer_info[new_buf], new_grid)?; + let new_chunks = self.buffer_info[new_buf].buffer_grid[new_grid] + .used_by_transforms + .to_vec(); + trace!("Buffer {new_buf} grid position {new_grid} used by chunks {new_chunks:?}"); + new_ready_transform_chunks.extend(new_chunks); + } + } + + Ok(()) + } +} + +#[allow(clippy::too_many_arguments)] +fn dequant_lf( + r: Rect, + lf: &mut [Image<f32>; 3], + quant_lf: &mut Image<u8>, + input: [&Image<i32>; 3], + color_correlation_params: &ColorCorrelationParams, + quant_params: &QuantizerParams, + lf_quant: &LfQuantFactors, + mul: f32, + frame_header: &FrameHeader, + bctx: &BlockContextMap, +) -> Result<()> { + let inv_quant_lf = (quantizer::GLOBAL_SCALE_DENOM as f32) + / (quant_params.global_scale as f32 * quant_params.quant_lf as f32); + let lf_factors = lf_quant.quant_factors.map(|factor| factor * inv_quant_lf); + + let lf_rects = lf.each_mut().map(|lf| lf.as_rect_mut()); + + if frame_header.is444() { + let [mut lf0, mut lf1, mut lf2] = lf_rects; + let mut lf_rects = (lf0.rect(r)?, lf1.rect(r)?, lf2.rect(r)?); + + let fac_x = lf_factors[0] * mul; + let fac_y = lf_factors[1] * mul; + let fac_b = lf_factors[2] * mul; + let cfl_fac_x = color_correlation_params.y_to_x_lf(); + let cfl_fac_b = color_correlation_params.y_to_b_lf(); + for y in 0..r.size.1 { + let quant_row_x = input[1].as_rect().row(y); + let quant_row_y = input[0].as_rect().row(y); + let quant_row_b = input[2].as_rect().row(y); + let dec_row_x = lf_rects.0.row(y); + let dec_row_y = lf_rects.1.row(y); + let dec_row_b = lf_rects.2.row(y); + for x in 0..r.size.0 { + let in_x = quant_row_x[x] as f32 * fac_x; + let in_y = quant_row_y[x] as f32 * fac_y; + let in_b = quant_row_b[x] as f32 * fac_b; + dec_row_y[x] = in_y; + dec_row_x[x] = in_y * cfl_fac_x + in_x; + dec_row_b[x] = in_y * cfl_fac_b + in_b; + } + } + } else { + for (c, mut lf_rect) in lf_rects.into_iter().enumerate() { + let rect = Rect { + origin: ( + r.origin.0 >> frame_header.hshift(c), + r.origin.1 >> frame_header.vshift(c), + ), + size: ( + r.size.0 >> frame_header.hshift(c), + r.size.1 >> frame_header.vshift(c), + ), + }; + let mut lf_rect = lf_rect.rect(rect)?; + let fac = lf_factors[c] * mul; + let ch = input[if c < 2 { c ^ 1 } else { c }]; + for y in 0..rect.size.1 { + let quant_row = ch.as_rect().row(y); + let row = lf_rect.row(y); + for x in 0..rect.size.0 { + row[x] = quant_row[x] as f32 * fac; + } + } + } + } + let mut quant_lf_as_rect = quant_lf.as_rect_mut(); + let mut quant_lf_rect = quant_lf_as_rect.rect(r)?; + if bctx.num_lf_contexts <= 1 { + for y in 0..r.size.1 { + quant_lf_rect.row(y).fill(0); + } + } else { + for y in 0..r.size.1 { + let qlf_row_val = quant_lf_rect.row(y); + let quant_row_x = input[1].as_rect().row(y >> frame_header.vshift(0)); + let quant_row_y = input[0].as_rect().row(y >> frame_header.vshift(1)); + let quant_row_b = input[2].as_rect().row(y >> frame_header.vshift(2)); + for x in 0..r.size.0 { + let bucket_x = bctx.lf_thresholds[0] + .iter() + .filter(|&t| quant_row_x[x >> frame_header.hshift(0)] > *t) + .count(); + let bucket_y = bctx.lf_thresholds[1] + .iter() + .filter(|&t| quant_row_y[x >> frame_header.hshift(1)] > *t) + .count(); + let bucket_b = bctx.lf_thresholds[2] + .iter() + .filter(|&t| quant_row_b[x >> frame_header.hshift(2)] > *t) + .count(); + let mut bucket = bucket_x; + bucket *= bctx.lf_thresholds[2].len() + 1; + bucket += bucket_b; + bucket *= bctx.lf_thresholds[1].len() + 1; + bucket += bucket_y; + qlf_row_val[x] = bucket as u8; + } + } + } + Ok(()) +} + +#[allow(clippy::too_many_arguments)] +pub fn decode_vardct_lf( + group: usize, + frame_header: &FrameHeader, + image_metadata: &ImageMetadata, + global_tree: &Option<Tree>, + color_correlation_params: &ColorCorrelationParams, + quant_params: &QuantizerParams, + lf_quant: &LfQuantFactors, + bctx: &BlockContextMap, + lf_image: &mut [Image<f32>; 3], + quant_lf: &mut Image<u8>, + br: &mut BitReader, +) -> Result<()> { + let extra_precision = br.read(2)?; + debug!(?extra_precision); + let mul = 1.0 / (1 << extra_precision) as f32; + let stream_id = ModularStreamId::VarDCTLF(group).get_id(frame_header); + debug!(?stream_id); + let r = frame_header.lf_group_rect(group); + debug!(?r); + let shrink_rect = |size: (usize, usize), c| { + ( + size.0 >> frame_header.hshift(c), + size.1 >> frame_header.vshift(c), + ) + }; + let mut buffers = [ + ModularChannel::new(shrink_rect(r.size, 1), image_metadata.bit_depth)?, + ModularChannel::new(shrink_rect(r.size, 0), image_metadata.bit_depth)?, + ModularChannel::new(shrink_rect(r.size, 2), image_metadata.bit_depth)?, + ]; + decode_modular_subbitstream( + buffers.iter_mut().collect(), + stream_id, + None, + global_tree, + br, + )?; + dequant_lf( + r, + lf_image, + quant_lf, + [&buffers[0].data, &buffers[1].data, &buffers[2].data], + color_correlation_params, + quant_params, + lf_quant, + mul, + frame_header, + bctx, + ) +} + +pub fn decode_hf_metadata( + group: usize, + frame_header: &FrameHeader, + image_metadata: &ImageMetadata, + global_tree: &Option<Tree>, + hf_meta: &mut HfMetadata, + br: &mut BitReader, +) -> Result<()> { + let stream_id = ModularStreamId::LFMeta(group).get_id(frame_header); + debug!(?stream_id); + let r = frame_header.lf_group_rect(group); + debug!(?r); + let upper_bound = r.size.0 * r.size.1; + let count_num_bits = upper_bound.ceil_log2(); + let count: usize = br.read(count_num_bits)? as usize + 1; + debug!(?count); + let cr = Rect { + origin: (r.origin.0 >> 3, r.origin.1 >> 3), + size: (r.size.0.div_ceil(8), r.size.1.div_ceil(8)), + }; + let mut buffers = [ + ModularChannel::new_with_shift(cr.size, Some((3, 3)), image_metadata.bit_depth)?, + ModularChannel::new_with_shift(cr.size, Some((3, 3)), image_metadata.bit_depth)?, + ModularChannel::new((count, 2), image_metadata.bit_depth)?, + ModularChannel::new(r.size, image_metadata.bit_depth)?, + ]; + decode_modular_subbitstream( + buffers.iter_mut().collect(), + stream_id, + None, + global_tree, + br, + )?; + let ytox_image = buffers[0].data.as_rect(); + let ytob_image = buffers[1].data.as_rect(); + let mut ytox_map = hf_meta.ytox_map.as_rect_mut(); + let mut ytob_map = hf_meta.ytob_map.as_rect_mut(); + let mut ytox_map_rect = ytox_map.rect(cr)?; + let mut ytob_map_rect = ytob_map.rect(cr)?; + let i8min: i32 = i8::MIN.into(); + let i8max: i32 = i8::MAX.into(); + for y in 0..cr.size.1 { + for x in 0..cr.size.0 { + ytox_map_rect.row(y)[x] = ytox_image.row(y)[x].clamp(i8min, i8max) as i8; + ytob_map_rect.row(y)[x] = ytob_image.row(y)[x].clamp(i8min, i8max) as i8; + } + } + let transform_image = buffers[2].data.as_rect(); + let epf_image = buffers[3].data.as_rect(); + let mut transform_map = hf_meta.transform_map.as_rect_mut(); + let mut transform_map_rect = transform_map.rect(r)?; + let mut raw_quant_map = hf_meta.raw_quant_map.as_rect_mut(); + let mut raw_quant_map_rect = raw_quant_map.rect(r)?; + let mut epf_map = hf_meta.epf_map.as_rect_mut(); + let mut epf_map_rect = epf_map.rect(r)?; + let mut num: usize = 0; + let mut used_hf_types: u32 = 0; + for y in 0..r.size.1 { + for x in 0..r.size.0 { + let epf_val = epf_image.row(y)[x]; + if !(0..8).contains(&epf_val) { + return Err(Error::InvalidEpfValue(epf_val)); + } + epf_map_rect.row(y)[x] = epf_val as u8; + if transform_map_rect.row(y)[x] != HfTransformType::INVALID_TRANSFORM { + continue; + } + if num >= count { + return Err(Error::InvalidVarDCTTransformMap); + } + let raw_transform = transform_image.row(0)[num]; + let raw_quant = 1 + transform_image.row(1)[num].clamp(0, 255); + used_hf_types |= 1 << raw_transform; + let transform_type = HfTransformType::from_usize(raw_transform as usize)?; + let cx = covered_blocks_x(transform_type) as usize; + let cy = covered_blocks_y(transform_type) as usize; + if (cx > 1 || cy > 1) && !frame_header.is444() { + return Err(Error::InvalidBlockSizeForChromaSubsampling); + } + let next_group = ((x / 32 + 1) * 32, (y / 32 + 1) * 32); + if x + cx > min(r.size.0, next_group.0) || y + cy > min(r.size.1, next_group.1) { + return Err(Error::HFBlockOutOfBounds); + } + let transform_id = raw_transform as u8; + for iy in 0..cy { + for ix in 0..cx { + transform_map_rect.row(y + iy)[x + ix] = if iy == 0 && ix == 0 { + transform_id + 128 // Set highest bit to signal first block. + } else { + transform_id + }; + raw_quant_map_rect.row(y + iy)[x + ix] = raw_quant; + } + } + num += 1; + } + } + hf_meta.used_hf_types |= used_hf_types; + Ok(()) +} diff --git a/third_party/rust/jxl/src/frame/modular/borrowed_buffers.rs b/third_party/rust/jxl/src/frame/modular/borrowed_buffers.rs @@ -0,0 +1,36 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +use std::{cell::RefMut, ops::DerefMut}; + +use crate::{error::Result, image::Image}; + +use super::{ModularBufferInfo, ModularChannel}; + +pub fn with_buffers<T>( + buffers: &[ModularBufferInfo], + indices: &[usize], + grid: usize, + f: impl FnOnce(Vec<&mut ModularChannel>) -> Result<T>, +) -> Result<T> { + let mut bufs = vec![]; + for i in indices { + // Allocate buffers if they are not present. + let buf = &buffers[*i]; + let b = &buf.buffer_grid[grid]; + let mut data = b.data.borrow_mut(); + if data.is_none() { + *data = Some(ModularChannel { + data: Image::new(b.size)?, + auxiliary_data: None, + shift: buf.info.shift, + bit_depth: buf.info.bit_depth, + }); + } + + bufs.push(RefMut::map(data, |x| x.as_mut().unwrap())); + } + f(bufs.iter_mut().map(|x| x.deref_mut()).collect()) +} diff --git a/third_party/rust/jxl/src/frame/modular/decode.rs b/third_party/rust/jxl/src/frame/modular/decode.rs @@ -0,0 +1,215 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +use crate::{ + bit_reader::BitReader, + entropy_coding::decode::SymbolReader, + error::{Error, Result}, + frame::quantizer::NUM_QUANT_TABLES, + headers::{JxlHeader, frame_header::FrameHeader, modular::GroupHeader}, + image::Image, + util::tracing_wrappers::*, +}; + +use super::{ + ModularChannel, Tree, + predict::{WeightedPredictorState, clamped_gradient}, + transforms::apply::meta_apply_local_transforms, + tree::NUM_NONREF_PROPERTIES, +}; + +use num_traits::abs; +use std::cmp::max; + +#[derive(Debug)] +pub enum ModularStreamId { + GlobalData, + VarDCTLF(usize), + ModularLF(usize), + LFMeta(usize), + QuantTable(usize), + ModularHF { pass: usize, group: usize }, +} + +impl ModularStreamId { + pub fn get_id(&self, frame_header: &FrameHeader) -> usize { + match self { + Self::GlobalData => 0, + Self::VarDCTLF(g) => 1 + g, + Self::ModularLF(g) => 1 + frame_header.num_lf_groups() + g, + Self::LFMeta(g) => 1 + frame_header.num_lf_groups() * 2 + g, + Self::QuantTable(q) => 1 + frame_header.num_lf_groups() * 3 + q, + Self::ModularHF { pass, group } => { + 1 + frame_header.num_lf_groups() * 3 + + NUM_QUANT_TABLES + + frame_header.num_groups() * *pass + + *group + } + } + } +} + +fn precompute_references( + buffers: &mut [&mut ModularChannel], + chan: usize, + y: usize, + references: &mut Image<i32>, +) { + references.as_rect_mut().apply(|_, v: &mut i32| *v = 0); + let mut offset = 0; + let num_extra_props = references.size().0; + for i in 0..chan { + if offset >= num_extra_props { + break; + } + let j = chan - i - 1; + if buffers[j].data.size() != buffers[chan].data.size() + || buffers[j].shift != buffers[chan].shift + { + continue; + } + let mut refs = references.as_rect_mut(); + let ref_chan = buffers[j].data.as_rect(); + for x in 0..buffers[chan].data.size().0 { + let v = ref_chan.row(y)[x]; + refs.row(x)[offset] = abs(v); + refs.row(x)[offset + 1] = v; + let vleft = if x > 0 { ref_chan.row(y)[x - 1] } else { 0 }; + let vtop = if y > 0 { ref_chan.row(y - 1)[x] } else { vleft }; + let vtopleft = if x > 0 && y > 0 { + ref_chan.row(y - 1)[x - 1] + } else { + vleft + }; + let vpredicted = clamped_gradient(vleft as i64, vtop as i64, vtopleft as i64); + refs.row(x)[offset + 2] = abs(v as i64 - vpredicted) as i32; + refs.row(x)[offset + 3] = (v as i64 - vpredicted) as i32; + } + offset += 4; + } +} + +#[allow(clippy::too_many_arguments)] +#[instrument(level = "debug", skip(buffers, reader, tree, br))] +fn decode_modular_channel( + buffers: &mut [&mut ModularChannel], + chan: usize, + stream_id: usize, + header: &GroupHeader, + tree: &Tree, + reader: &mut SymbolReader, + br: &mut BitReader, +) -> Result<()> { + debug!("reading channel"); + let size = buffers[chan].data.size(); + let mut wp_state = WeightedPredictorState::new(&header.wp_header, size.0); + let mut num_ref_props = + max(0, tree.max_property() as i32 - NUM_NONREF_PROPERTIES as i32) as usize; + num_ref_props = num_ref_props.div_ceil(4) * 4; + let mut references = Image::<i32>::new((num_ref_props, size.0))?; + let num_properties = NUM_NONREF_PROPERTIES + num_ref_props; + let make_pixel = + |dec: i32, mul: u32, guess: i64| -> i32 { (guess + (mul as i64) * (dec as i64)) as i32 }; + for y in 0..size.1 { + precompute_references(buffers, chan, y, &mut references); + let mut property_buffer: Vec<i32> = vec![0; num_properties]; + property_buffer[0] = chan as i32; + property_buffer[1] = stream_id as i32; + for x in 0..size.0 { + let prediction_result = tree.predict( + buffers, + chan, + &mut wp_state, + x, + y, + &references, + &mut property_buffer, + ); + let dec = + reader.read_signed(&tree.histograms, br, prediction_result.context as usize)?; + let val = make_pixel(dec, prediction_result.multiplier, prediction_result.guess); + buffers[chan].data.as_rect_mut().row(y)[x] = val; + trace!(y, x, val, dec, ?property_buffer, ?prediction_result); + wp_state.update_errors(val, (x, y), size.0); + } + } + + Ok(()) +} + +// This function will decode a header and apply local transforms if a header is not given. +// The intended use of passing a header is for the DcGlobal section. +pub fn decode_modular_subbitstream( + buffers: Vec<&mut ModularChannel>, + stream_id: usize, + header: Option<GroupHeader>, + global_tree: &Option<Tree>, + br: &mut BitReader, +) -> Result<()> { + // Skip decoding if all grids are zero-sized. + let is_empty = buffers.iter().all(|buffer| { + let size = buffer.data.size(); + size.0 == 0 || size.1 == 0 + }); + if is_empty { + return Ok(()); + } + let mut transform_steps = vec![]; + let mut buffer_storage = vec![]; + + let buffers = buffers.into_iter().collect::<Vec<_>>(); + let (header, mut buffers) = match header { + Some(h) => (h, buffers), + None => { + let h = GroupHeader::read(br)?; + if !h.transforms.is_empty() { + // Note: reassigning to `buffers` here convinces the borrow checker that the borrow of + // `buffer_storage` ought to outlive `buffers[..]`'s lifetime, which obviously breaks + // applying transforms later. + let new_bufs; + (new_bufs, transform_steps) = + meta_apply_local_transforms(buffers, &mut buffer_storage, &h)?; + (h, new_bufs) + } else { + (h, buffers) + } + } + }; + + if header.use_global_tree && global_tree.is_none() { + return Err(Error::NoGlobalTree); + } + let local_tree = if !header.use_global_tree { + Some(Tree::read(br, 1024)?) + } else { + None + }; + let tree = if header.use_global_tree { + global_tree.as_ref().unwrap() + } else { + local_tree.as_ref().unwrap() + }; + + let image_width = buffers + .iter() + .map(|info| info.channel_info().size.0) + .max() + .unwrap_or(0); + let mut reader = SymbolReader::new(&tree.histograms, br, Some(image_width))?; + + for i in 0..buffers.len() { + decode_modular_channel(&mut buffers, i, stream_id, &header, tree, &mut reader, br)?; + } + + reader.check_final_state(&tree.histograms)?; + + drop(buffers); + + for step in transform_steps.iter().rev() { + step.local_apply(&mut buffer_storage)?; + } + + Ok(()) +} diff --git a/third_party/rust/jxl/src/frame/modular/predict.rs b/third_party/rust/jxl/src/frame/modular/predict.rs @@ -0,0 +1,516 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +use std::array::from_fn; + +use crate::{ + error::{Error, Result}, + headers::modular::WeightedHeader, + image::{Image, ImageRect}, + util::floor_log2_nonzero, +}; +use num_derive::FromPrimitive; +use num_traits::FromPrimitive; + +#[repr(u8)] +#[derive(Debug, FromPrimitive, Clone, Copy, PartialEq, Eq)] +pub enum Predictor { + Zero = 0, + West = 1, + North = 2, + AverageWestAndNorth = 3, + Select = 4, + Gradient = 5, + Weighted = 6, + NorthEast = 7, + NorthWest = 8, + WestWest = 9, + AverageWestAndNorthWest = 10, + AverageNorthAndNorthWest = 11, + AverageNorthAndNorthEast = 12, + AverageAll = 13, +} + +impl Predictor { + pub fn requires_full_row(&self) -> bool { + matches!( + self, + Predictor::Weighted + | Predictor::NorthEast + | Predictor::AverageNorthAndNorthEast + | Predictor::AverageAll + ) + } +} + +impl TryFrom<u32> for Predictor { + type Error = Error; + + fn try_from(value: u32) -> Result<Self> { + Self::from_u32(value).ok_or(Error::InvalidPredictor(value)) + } +} + +pub struct PredictionData { + pub left: i32, + pub top: i32, + pub toptop: i32, + pub topleft: i32, + pub topright: i32, + pub leftleft: i32, + pub toprightright: i32, +} + +impl PredictionData { + pub fn get(rect: ImageRect<i32>, x: usize, y: usize) -> Self { + let left = if x > 0 { + rect.row(y)[x - 1] + } else if y > 0 { + rect.row(y - 1)[0] + } else { + 0 + }; + let top = if y > 0 { rect.row(y - 1)[x] } else { left }; + let topleft = if x > 0 && y > 0 { + rect.row(y - 1)[x - 1] + } else { + left + }; + let topright = if x + 1 < rect.size().0 && y > 0 { + rect.row(y - 1)[x + 1] + } else { + top + }; + let leftleft = if x > 1 { rect.row(y)[x - 2] } else { left }; + let toptop = if y > 1 { rect.row(y - 2)[x] } else { top }; + let toprightright = if x + 2 < rect.size().0 && y > 0 { + rect.row(y - 1)[x + 2] + } else { + topright + }; + Self { + left, + top, + toptop, + topleft, + topright, + leftleft, + toprightright, + } + } + + #[allow(clippy::too_many_arguments)] + pub fn get_with_neighbors( + rect: ImageRect<i32>, + rect_left: Option<ImageRect<i32>>, + rect_top: Option<ImageRect<i32>>, + rect_top_left: Option<ImageRect<i32>>, + rect_right: Option<ImageRect<i32>>, + rect_top_right: Option<ImageRect<i32>>, + x: usize, + y: usize, + xsize: usize, + ysize: usize, + ) -> Self { + let left = if x > 0 { + rect.row(y)[x - 1] + } else if let Some(l) = rect_left { + l.row(y)[xsize - 1] + } else if y > 0 { + rect.row(y - 1)[0] + } else if let Some(t) = rect_top { + t.row(ysize - 1)[0] + } else { + 0 + }; + let top = if y > 0 { + rect.row(y - 1)[x] + } else if let Some(t) = rect_top { + t.row(ysize - 1)[x] + } else { + left + }; + let topleft = if x > 0 { + if y > 0 { + rect.row(y - 1)[x - 1] + } else if let Some(t) = rect_top { + t.row(ysize - 1)[x - 1] + } else { + left + } + } else if y > 0 { + if let Some(l) = rect_left { + l.row(y - 1)[xsize - 1] + } else { + left + } + } else if let Some(tl) = rect_top_left { + tl.row(ysize - 1)[xsize - 1] + } else { + left + }; + let topright = if x + 1 < rect.size().0 { + if y > 0 { + rect.row(y - 1)[x + 1] + } else if let Some(t) = rect_top { + t.row(ysize - 1)[x + 1] + } else { + top + } + } else if y > 0 { + if let Some(r) = rect_right { + r.row(y - 1)[0] + } else { + top + } + } else if let Some(tr) = rect_top_right { + tr.row(ysize - 1)[0] + } else { + top + }; + let leftleft = if x > 1 { + rect.row(y)[x - 2] + } else if let Some(l) = rect_left { + l.row(y)[xsize + x - 2] + } else { + left + }; + let toptop = if y > 1 { + rect.row(y - 2)[x] + } else if let Some(t) = rect_top { + t.row(ysize + y - 2)[x] + } else { + top + }; + let toprightright = if x + 2 < rect.size().0 { + if y > 0 { + rect.row(y - 1)[x + 2] + } else if let Some(t) = rect_top { + t.row(ysize - 1)[x + 2] + } else { + topright + } + } else if y > 0 { + if let Some(r) = rect_right { + r.row(y - 1)[x + 2 - rect.size().0] + } else { + topright + } + } else if let Some(tr) = rect_top_right { + tr.row(ysize - 1)[x + 2 - rect.size().0] + } else { + topright + }; + Self { + left, + top, + toptop, + topleft, + topright, + leftleft, + toprightright, + } + } +} + +pub fn clamped_gradient(left: i64, top: i64, topleft: i64) -> i64 { + // Same code/logic as libjxl. + let min = left.min(top); + let max = left.max(top); + let grad = left + top - topleft; + let grad_clamp_max = if topleft < min { max } else { grad }; + if topleft > max { min } else { grad_clamp_max } +} + +impl Predictor { + pub const NUM_PREDICTORS: u32 = Predictor::AverageAll as u32 + 1; + + pub fn predict_one( + &self, + PredictionData { + left, + top, + toptop, + topleft, + topright, + leftleft, + toprightright, + }: PredictionData, + wp_pred: i64, + ) -> i64 { + match self { + Predictor::Zero => 0, + Predictor::West => left as i64, + Predictor::North => top as i64, + Predictor::Select => Self::select(left as i64, top as i64, topleft as i64), + Predictor::Gradient => clamped_gradient(left as i64, top as i64, topleft as i64), + Predictor::Weighted => wp_pred, + Predictor::WestWest => leftleft as i64, + Predictor::NorthEast => topright as i64, + Predictor::NorthWest => topleft as i64, + Predictor::AverageWestAndNorth => (top as i64 + left as i64) / 2, + Predictor::AverageWestAndNorthWest => (left as i64 + topleft as i64) / 2, + Predictor::AverageNorthAndNorthWest => (top as i64 + topleft as i64) / 2, + Predictor::AverageNorthAndNorthEast => (top as i64 + topright as i64) / 2, + Predictor::AverageAll => { + (6 * top as i64 - 2 * toptop as i64 + + 7 * left as i64 + + leftleft as i64 + + toprightright as i64 + + 3 * topright as i64 + + 8) + / 16 + } + } + } + + fn select(left: i64, top: i64, topleft: i64) -> i64 { + let p = left + top - topleft; + if (p - left).abs() < (p - top).abs() { + left + } else { + top + } + } +} + +const NUM_PREDICTORS: usize = 4; +const PRED_EXTRA_BITS: i64 = 3; +const PREDICTION_ROUND: i64 = ((1 << PRED_EXTRA_BITS) >> 1) - 1; +// Allows to approximate division by a number from 1 to 64. +// for (int i = 0; i < 64; i++) divlookup[i] = (1 << 24) / (i + 1); +const DIVLOOKUP: [u32; 64] = [ + 16777216, 8388608, 5592405, 4194304, 3355443, 2796202, 2396745, 2097152, 1864135, 1677721, + 1525201, 1398101, 1290555, 1198372, 1118481, 1048576, 986895, 932067, 883011, 838860, 798915, + 762600, 729444, 699050, 671088, 645277, 621378, 599186, 578524, 559240, 541200, 524288, 508400, + 493447, 479349, 466033, 453438, 441505, 430185, 419430, 409200, 399457, 390167, 381300, 372827, + 364722, 356962, 349525, 342392, 335544, 328965, 322638, 316551, 310689, 305040, 299593, 294337, + 289262, 284359, 279620, 275036, 270600, 266305, 262144, +]; + +fn add_bits(x: i32) -> i64 { + (x as i64) << PRED_EXTRA_BITS +} + +fn error_weight(x: u32, maxweight: u32) -> u32 { + let shift = floor_log2_nonzero(x + 1) as i32 - 5; + if shift < 0 { + 4u32 + maxweight * DIVLOOKUP[x as usize] + } else { + 4u32 + ((maxweight * DIVLOOKUP[x as usize >> shift]) >> shift) + } +} + +fn weighted_average(pixels: &[i64; NUM_PREDICTORS], weights: &mut [u32; NUM_PREDICTORS]) -> i64 { + let log_weight = floor_log2_nonzero(weights.iter().fold(0u32, |sum, el| sum + *el)); + let weight_sum = weights.iter_mut().fold(0, |sum, el| { + *el >>= log_weight - 4; + sum + *el + }); + let sum = weights + .iter() + .enumerate() + .fold(((weight_sum >> 1) - 1) as i64, |sum, (i, weight)| { + sum + pixels[i] * *weight as i64 + }); + (sum * DIVLOOKUP[(weight_sum - 1) as usize] as i64) >> 24 +} + +#[derive(Debug)] +pub struct WeightedPredictorState<'a> { + prediction: [i64; NUM_PREDICTORS], + pred: i64, + pred_errors: [Vec<u32>; NUM_PREDICTORS], + error: Vec<i32>, + wp_header: &'a WeightedHeader, +} + +impl<'a> WeightedPredictorState<'a> { + pub fn new(wp_header: &'a WeightedHeader, xsize: usize) -> WeightedPredictorState<'a> { + let num_errors = (xsize + 2) * 2; + WeightedPredictorState { + prediction: [0; NUM_PREDICTORS], + pred: 0, + pred_errors: from_fn(|_| vec![0; num_errors]), + error: vec![0; num_errors], + wp_header, + } + } + + pub fn save_state(&self, wp_image: &mut Image<i32>, xsize: usize) { + wp_image + .as_rect_mut() + .row(0) + .copy_from_slice(&self.error[xsize + 2..]); + } + + pub fn restore_state(&mut self, wp_image: &Image<i32>, xsize: usize) { + self.error[xsize + 2..].copy_from_slice(wp_image.as_rect().row(0)); + } + + pub fn update_errors(&mut self, correct_val: i32, pos: (usize, usize), xsize: usize) { + let (cur_row, prev_row) = if pos.1 & 1 != 0 { + (0, xsize + 2) + } else { + (xsize + 2, 0) + }; + let val = add_bits(correct_val); + self.error[cur_row + pos.0] = (self.pred - val) as i32; + for (i, pred_err) in self.pred_errors.iter_mut().enumerate() { + let err = + (((self.prediction[i] - val).abs() + PREDICTION_ROUND) >> PRED_EXTRA_BITS) as u32; + pred_err[cur_row + pos.0] = err; + let idx = prev_row + pos.0 + 1; + pred_err[idx] = pred_err[idx].wrapping_add(err); + } + } + + pub fn predict_and_property( + &mut self, + pos: (usize, usize), + xsize: usize, + data: &PredictionData, + ) -> (i64, i32) { + let (cur_row, prev_row) = if pos.1 & 1 != 0 { + (0, xsize + 2) + } else { + (xsize + 2, 0) + }; + let pos_n = prev_row + pos.0; + let pos_ne = if pos.0 < xsize - 1 { pos_n + 1 } else { pos_n }; + let pos_nw = if pos.0 > 0 { pos_n - 1 } else { pos_n }; + let mut weights = [0u32; NUM_PREDICTORS]; + for (i, weight) in weights.iter_mut().enumerate() { + *weight = error_weight( + self.pred_errors[i][pos_n] + .wrapping_add(self.pred_errors[i][pos_ne]) + .wrapping_add(self.pred_errors[i][pos_nw]), + self.wp_header.w(i).unwrap(), + ); + } + let n = add_bits(data.top); + let w = add_bits(data.left); + let ne = add_bits(data.topright); + let nw = add_bits(data.topleft); + let nn = add_bits(data.toptop); + + let te_w = if pos.0 == 0 { + 0 + } else { + self.error[cur_row + pos.0 - 1] as i64 + }; + let te_n = self.error[pos_n] as i64; + let te_nw = self.error[pos_nw] as i64; + let sum_wn = te_n + te_w; + let te_ne = self.error[pos_ne] as i64; + + let mut p = te_w; + if te_n.abs() > p.abs() { + p = te_n; + } + if te_nw.abs() > p.abs() { + p = te_nw; + } + if te_ne.abs() > p.abs() { + p = te_ne; + } + + self.prediction[0] = w + ne - n; + self.prediction[1] = n - (((sum_wn + te_ne) * self.wp_header.p1c as i64) >> 5); + self.prediction[2] = w - (((sum_wn + te_nw) * self.wp_header.p2c as i64) >> 5); + self.prediction[3] = n + - ((te_nw * (self.wp_header.p3ca as i64) + + (te_n * (self.wp_header.p3cb as i64)) + + (te_ne * (self.wp_header.p3cc as i64)) + + ((nn - n) * (self.wp_header.p3cd as i64)) + + ((nw - w) * (self.wp_header.p3ce as i64))) + >> 5); + + self.pred = weighted_average(&self.prediction, &mut weights); + + if ((te_n ^ te_w) | (te_n ^ te_nw)) <= 0 { + let mx = w.max(ne.max(n)); + let mn = w.min(ne.min(n)); + self.pred = mn.max(mx.min(self.pred)); + } + ((self.pred + PREDICTION_ROUND) >> PRED_EXTRA_BITS, p as i32) + } +} + +#[cfg(test)] +mod tests { + use crate::headers::modular::{GroupHeader, WeightedHeader}; + + use super::{PredictionData, WeightedPredictorState}; + + struct SimpleRandom { + out: i64, + } + + impl SimpleRandom { + fn new() -> SimpleRandom { + SimpleRandom { out: 1 } + } + fn next(&mut self) -> i64 { + self.out = self.out * 48271 % 0x7fffffff; + self.out + } + } + + fn step( + rng: &mut SimpleRandom, + state: &mut WeightedPredictorState, + xsize: usize, + ysize: usize, + ) -> (i64, i32) { + let pos = (rng.next() as usize % xsize, rng.next() as usize % ysize); + let res = state.predict_and_property( + pos, + xsize, + &PredictionData { + top: rng.next() as i32 % 256, + left: rng.next() as i32 % 256, + topright: rng.next() as i32 % 256, + topleft: rng.next() as i32 % 256, + toptop: rng.next() as i32 % 256, + leftleft: 0, + toprightright: 0, + }, + ); + state.update_errors((rng.next() % 256) as i32, pos, xsize); + res + } + + #[test] + fn predict_and_update_errors() { + let mut rng = SimpleRandom::new(); + let header = GroupHeader { + use_global_tree: false, + wp_header: WeightedHeader { + all_default: true, + p1c: rng.next() as u32 % 32, + p2c: rng.next() as u32 % 32, + p3ca: rng.next() as u32 % 32, + p3cb: rng.next() as u32 % 32, + p3cc: rng.next() as u32 % 32, + p3cd: rng.next() as u32 % 32, + p3ce: rng.next() as u32 % 32, + w0: rng.next() as u32 % 16, + w1: rng.next() as u32 % 16, + w2: rng.next() as u32 % 16, + w3: rng.next() as u32 % 16, + }, + transforms: Vec::new(), + }; + let xsize = 8; + let ysize = 8; + let mut state = WeightedPredictorState::new(&header.wp_header, xsize); + // The golden number results are generated by using the libjxl predictor with the same input numbers. + assert_eq!(step(&mut rng, &mut state, xsize, ysize), (135i64, 0i32)); + assert_eq!(step(&mut rng, &mut state, xsize, ysize), (110i64, -60i32)); + assert_eq!(step(&mut rng, &mut state, xsize, ysize), (165i64, 0i32)); + assert_eq!(step(&mut rng, &mut state, xsize, ysize), (153i64, -60i32)); + } +} diff --git a/third_party/rust/jxl/src/frame/modular/transforms.rs b/third_party/rust/jxl/src/frame/modular/transforms.rs @@ -0,0 +1,376 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +use std::cell::RefCell; + +use apply::TransformStep; +pub use apply::TransformStepChunk; +use num_derive::FromPrimitive; + +use crate::frame::modular::ModularBuffer; +use crate::headers::frame_header::FrameHeader; +use crate::util::tracing_wrappers::*; + +use super::{ModularBufferInfo, ModularGridKind, Predictor}; + +pub(super) mod apply; +mod palette; +mod rct; +mod squeeze; + +#[derive(Debug, FromPrimitive, PartialEq, Clone, Copy)] +pub enum RctPermutation { + Rgb = 0, + Gbr = 1, + Brg = 2, + Rbg = 3, + Grb = 4, + Bgr = 5, +} + +#[derive(Debug, FromPrimitive, PartialEq, Clone, Copy)] +pub enum RctOp { + Noop = 0, + AddFirstToThird = 1, + AddFirstToSecond = 2, + AddFirstToSecondAndThird = 3, + AddAvgToSecond = 4, + AddFirstToThirdAndAvgToSecond = 5, + YCoCg = 6, +} + +#[instrument(level = "trace", skip_all, ret)] +pub fn make_grids( + frame_header: &FrameHeader, + transform_steps: Vec<TransformStep>, + section_buffer_indices: &[Vec<usize>], + buffer_info: &mut Vec<ModularBufferInfo>, +) -> Vec<TransformStepChunk> { + // Initialize grid sizes, starting from coded channels. + for i in section_buffer_indices[1].iter() { + buffer_info[*i].grid_kind = ModularGridKind::Lf; + } + for buffer_indices in section_buffer_indices.iter().skip(2) { + for i in buffer_indices.iter() { + buffer_info[*i].grid_kind = ModularGridKind::Hf; + } + } + + trace!(?buffer_info, "post set grid kind for coded channels"); + + // Transforms can be un-applied in the opposite order they appear with in the array, + // so we can use that information to propagate grid kinds. + + for step in transform_steps.iter().rev() { + match step { + TransformStep::Rct { + buf_in, buf_out, .. + } => { + let grid_in = buffer_info[buf_in[0]].grid_kind; + for i in 0..3 { + assert_eq!(grid_in, buffer_info[buf_in[i]].grid_kind); + } + for i in 0..3 { + buffer_info[buf_out[i]].grid_kind = grid_in; + } + } + TransformStep::Palette { + buf_in, buf_out, .. + } => { + for buf in buf_out.iter() { + buffer_info[*buf].grid_kind = buffer_info[*buf_in].grid_kind; + } + } + TransformStep::HSqueeze { buf_in, buf_out } + | TransformStep::VSqueeze { buf_in, buf_out } => { + buffer_info[*buf_out].grid_kind = buffer_info[buf_in[0]] + .grid_kind + .max(buffer_info[buf_in[1]].grid_kind); + } + } + } + + // Set grid shapes. + for buf in buffer_info.iter_mut() { + buf.grid_shape = buf.grid_kind.grid_shape(frame_header); + } + + trace!(?buffer_info, "post propagate grid kind"); + + let get_grid_indices = |shape: (usize, usize)| { + (0..shape.1).flat_map(move |y| (0..shape.0).map(move |x| (x as isize, y as isize))) + }; + + // Create grids. + for g in buffer_info.iter_mut() { + let is_output = g.info.output_channel_idx >= 0; + g.buffer_grid = get_grid_indices(g.grid_shape) + .map(|(x, y)| ModularBuffer { + data: RefCell::new(None), + remaining_uses: if is_output { 1 } else { 0 }, + used_by_transforms: vec![], + size: g + .get_grid_rect(frame_header, g.grid_kind, (x as usize, y as usize)) + .size, + }) + .collect(); + } + + trace!(?buffer_info, "with grids"); + + let add_transform_step = + |transform: &TransformStep, + grid_pos: (isize, isize), + grid_transform_steps: &mut Vec<TransformStepChunk>| { + let ts = grid_transform_steps.len(); + grid_transform_steps.push(TransformStepChunk { + step: transform.clone(), + grid_pos: (grid_pos.0 as usize, grid_pos.1 as usize), + incomplete_deps: 0, + }); + ts + }; + + let add_grid_use = |ts: usize, + input_buffer_idx: usize, + output_grid_kind: ModularGridKind, + output_grid_shape: (usize, usize), + output_grid_pos: (isize, isize), + grid_transform_steps: &mut Vec<TransformStepChunk>, + buffer_info: &mut Vec<ModularBufferInfo>| { + let output_grid_size = (output_grid_shape.0 as isize, output_grid_shape.1 as isize); + if output_grid_pos.0 < 0 + || output_grid_pos.0 >= output_grid_size.0 + || output_grid_pos.1 < 0 + || output_grid_pos.1 >= output_grid_size.1 + { + // Skip adding uses of non-existent grid positions. + return; + } + let output_grid_pos = (output_grid_pos.0 as usize, output_grid_pos.1 as usize); + let input_grid_pos = + buffer_info[input_buffer_idx].get_grid_idx(output_grid_kind, output_grid_pos); + if !buffer_info[input_buffer_idx].buffer_grid[input_grid_pos] + .used_by_transforms + .contains(&ts) + { + buffer_info[input_buffer_idx].buffer_grid[input_grid_pos].remaining_uses += 1; + buffer_info[input_buffer_idx].buffer_grid[input_grid_pos] + .used_by_transforms + .push(ts); + grid_transform_steps[ts].incomplete_deps += 1; + } + }; + + // Add grid-ed transforms. + let mut grid_transform_steps = vec![]; + + for transform in transform_steps { + match &transform { + TransformStep::Rct { + buf_in, buf_out, .. + } => { + // Easy case: we just depend on the 3 input buffers in the same location. + let out_kind = buffer_info[buf_out[0]].grid_kind; + let out_shape = buffer_info[buf_out[0]].grid_shape; + for (x, y) in get_grid_indices(out_shape) { + let ts = add_transform_step(&transform, (x, y), &mut grid_transform_steps); + for bin in buf_in { + add_grid_use( + ts, + *bin, + out_kind, + out_shape, + (x, y), + &mut grid_transform_steps, + buffer_info, + ); + } + } + } + TransformStep::Palette { + buf_in, + buf_pal, + buf_out, + predictor, + .. + } if predictor.requires_full_row() => { + // Delta palette with AverageAll or Weighted. Those are special, because we can + // only make progress one full image row at a time (since we need decoded values + // from the previous row or two rows). + let out_kind = buffer_info[buf_out[0]].grid_kind; + let out_shape = buffer_info[buf_out[0]].grid_shape; + let mut ts = 0; + for (x, y) in get_grid_indices(out_shape) { + if x == 0 { + ts = add_transform_step(&transform, (x, y), &mut grid_transform_steps); + add_grid_use( + ts, + *buf_pal, + out_kind, + out_shape, + (x, y), + &mut grid_transform_steps, + buffer_info, + ); + } + add_grid_use( + ts, + *buf_in, + out_kind, + out_shape, + (x, y), + &mut grid_transform_steps, + buffer_info, + ); + for out in buf_out.iter() { + add_grid_use( + ts, + *out, + out_kind, + out_shape, + (x, y - 1), + &mut grid_transform_steps, + buffer_info, + ); + } + } + } + TransformStep::Palette { + buf_in, + buf_pal, + buf_out, + predictor, + .. + } => { + // Maybe-delta palette: we depend on the palette and the input buffer in the same + // location. We may also depend on other grid positions in the output buffer, + // according to the used predictor. + let out_kind = buffer_info[buf_out[0]].grid_kind; + let out_shape = buffer_info[buf_out[0]].grid_shape; + for (x, y) in get_grid_indices(out_shape) { + let ts = add_transform_step(&transform, (x, y), &mut grid_transform_steps); + add_grid_use( + ts, + *buf_pal, + out_kind, + out_shape, + (x, y), + &mut grid_transform_steps, + buffer_info, + ); + add_grid_use( + ts, + *buf_in, + out_kind, + out_shape, + (x, y), + &mut grid_transform_steps, + buffer_info, + ); + let offsets = match predictor { + Predictor::Zero => [].as_slice(), + _ => &[(0, -1), (-1, 0), (-1, -1)], + }; + for (dx, dy) in offsets { + for out in buf_out.iter() { + add_grid_use( + ts, + *out, + out_kind, + out_shape, + (x + dx, y + dy), + &mut grid_transform_steps, + buffer_info, + ); + } + } + } + } + TransformStep::HSqueeze { buf_in, buf_out } => { + let out_kind = buffer_info[*buf_out].grid_kind; + let out_shape = buffer_info[*buf_out].grid_shape; + for (x, y) in get_grid_indices(out_shape) { + let ts = add_transform_step(&transform, (x, y), &mut grid_transform_steps); + // Average and residuals from the same position + for bin in buf_in { + add_grid_use( + ts, + *bin, + out_kind, + out_shape, + (x, y), + &mut grid_transform_steps, + buffer_info, + ); + } + // Next average + add_grid_use( + ts, + buf_in[0], + out_kind, + out_shape, + (x + 1, y), + &mut grid_transform_steps, + buffer_info, + ); + // Previous decoded + add_grid_use( + ts, + *buf_out, + out_kind, + out_shape, + (x - 1, y), + &mut grid_transform_steps, + buffer_info, + ); + } + } + TransformStep::VSqueeze { buf_in, buf_out } => { + let out_kind = buffer_info[*buf_out].grid_kind; + let out_shape = buffer_info[*buf_out].grid_shape; + for (x, y) in get_grid_indices(out_shape) { + let ts = add_transform_step(&transform, (x, y), &mut grid_transform_steps); + // Average and residuals from the same position + for bin in buf_in { + add_grid_use( + ts, + *bin, + out_kind, + out_shape, + (x, y), + &mut grid_transform_steps, + buffer_info, + ); + } + // Next average + add_grid_use( + ts, + buf_in[0], + out_kind, + out_shape, + (x, y + 1), + &mut grid_transform_steps, + buffer_info, + ); + // Previous decoded + add_grid_use( + ts, + *buf_out, + out_kind, + out_shape, + (x, y - 1), + &mut grid_transform_steps, + buffer_info, + ); + } + } + } + } + + trace!(?grid_transform_steps, ?buffer_info); + + grid_transform_steps +} diff --git a/third_party/rust/jxl/src/frame/modular/transforms/apply.rs b/third_party/rust/jxl/src/frame/modular/transforms/apply.rs @@ -0,0 +1,972 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +use std::{cell::Ref, fmt::Debug}; + +use num_traits::FromPrimitive; + +use crate::{ + error::{Error, Result}, + frame::modular::{ + ChannelInfo, ModularBufferInfo, ModularChannel, ModularGridKind, Predictor, + borrowed_buffers::with_buffers, + }, + headers::{self, frame_header::FrameHeader, modular::TransformId, modular::WeightedHeader}, + util::tracing_wrappers::*, +}; +use std::cell::RefMut; +use std::ops::Deref; +use std::ops::DerefMut; + +use super::{RctOp, RctPermutation}; + +#[derive(Debug, Clone)] +pub enum TransformStep { + Rct { + buf_in: [usize; 3], + buf_out: [usize; 3], + op: RctOp, + perm: RctPermutation, + }, + Palette { + buf_in: usize, + buf_pal: usize, + buf_out: Vec<usize>, + num_colors: usize, + num_deltas: usize, + predictor: Predictor, + wp_header: WeightedHeader, + }, + HSqueeze { + buf_in: [usize; 2], + buf_out: usize, + }, + VSqueeze { + buf_in: [usize; 2], + buf_out: usize, + }, +} + +#[derive(Debug)] +pub struct TransformStepChunk { + pub(super) step: TransformStep, + // Grid position this transform should produce. + // Note that this is a lie for Palette with AverageAll or Weighted, as the transform with + // position (0, y) will produce the entire row of blocks (*, y) (and there will be no + // transforms with position (x, y) with x > 0). + pub(super) grid_pos: (usize, usize), + // Number of inputs that are not yet available. + pub(super) incomplete_deps: usize, +} + +impl TransformStepChunk { + // Marks that one dependency of this transform is ready, and potentially runs the transform, + // returning the new buffers that are now ready. + #[instrument(level = "trace", skip_all)] + pub fn dep_ready( + &mut self, + frame_header: &FrameHeader, + buffers: &mut [ModularBufferInfo], + ) -> Result<Vec<(usize, usize)>> { + self.incomplete_deps = self.incomplete_deps.checked_sub(1).unwrap(); + if self.incomplete_deps > 0 { + trace!( + "skipping transform chunk because incomplete_deps = {}", + self.incomplete_deps + ); + return Ok(vec![]); + } + let buf_out: &[usize] = match &self.step { + TransformStep::Rct { buf_out, .. } => buf_out, + TransformStep::Palette { buf_out, .. } => buf_out, + TransformStep::HSqueeze { buf_out, .. } | TransformStep::VSqueeze { buf_out, .. } => { + &[*buf_out] + } + }; + + let out_grid_kind = buffers[buf_out[0]].grid_kind; + let out_grid = buffers[buf_out[0]].get_grid_idx(out_grid_kind, self.grid_pos); + let out_size = buffers[buf_out[0]].info.size; + for bo in buf_out { + assert_eq!(out_grid_kind, buffers[*bo].grid_kind); + assert_eq!(out_size, buffers[*bo].info.size); + } + + match &self.step { + TransformStep::Rct { + buf_in, + buf_out, + op, + perm, + } => { + for i in 0..3 { + assert_eq!(out_grid_kind, buffers[buf_in[i]].grid_kind); + assert_eq!(out_size, buffers[buf_in[i]].info.size); + // Optimistically move the buffers to the output if possible. + // If not, creates buffers in the output that are a copy of the input buffers. + // This should be rare. + *buffers[buf_out[i]].buffer_grid[out_grid].data.borrow_mut() = + Some(buffers[buf_in[i]].buffer_grid[out_grid].get_buffer()?); + } + with_buffers(buffers, buf_out, out_grid, |mut bufs| { + super::rct::do_rct_step(&mut bufs, *op, *perm); + Ok(()) + })?; + Ok(buf_out.iter().map(|x| (*x, out_grid)).collect()) + } + TransformStep::Palette { + buf_in, + buf_pal, + buf_out, + .. + } if buffers[*buf_in].info.size.0 == 0 => { + // Nothing to do, just bookkeeping. + buffers[*buf_in].buffer_grid[out_grid].mark_used(); + buffers[*buf_pal].buffer_grid[0].mark_used(); + with_buffers(buffers, buf_out, out_grid, |_| Ok(()))?; + Ok(buf_out.iter().map(|x| (*x, out_grid)).collect()) + } + TransformStep::Palette { + buf_in, + buf_pal, + buf_out, + num_colors, + num_deltas, + predictor, + .. + } if !predictor.requires_full_row() => { + assert_eq!(out_grid_kind, buffers[*buf_in].grid_kind); + assert_eq!(out_size, buffers[*buf_in].info.size); + + { + let img_in = + Ref::map(buffers[*buf_in].buffer_grid[out_grid].data.borrow(), |x| { + x.as_ref().unwrap() + }); + let img_pal = Ref::map(buffers[*buf_pal].buffer_grid[0].data.borrow(), |x| { + x.as_ref().unwrap() + }); + // Ensure that the output buffers are present. + // TODO(szabadka): Extend the callback to support many grid points. + with_buffers(buffers, buf_out, out_grid, |_| Ok(()))?; + let grid_shape = buffers[buf_out[0]].grid_shape; + let grid_x = out_grid % grid_shape.0; + let grid_y = out_grid / grid_shape.0; + let border = if *predictor == Predictor::Zero { 0 } else { 1 }; + let grid_x0 = grid_x.saturating_sub(border); + let grid_y0 = grid_y.saturating_sub(border); + let grid_x1 = grid_x + 1; + let grid_y1 = grid_y + 1; + let mut out_bufs = vec![]; + for i in buf_out { + for gy in grid_y0..grid_y1 { + for gx in grid_x0..grid_x1 { + let grid = gy * grid_shape.0 + gx; + let buf = &buffers[*i]; + let b = &buf.buffer_grid[grid]; + let data = b.data.borrow_mut(); + out_bufs.push(RefMut::map(data, |x| x.as_mut().unwrap())); + } + } + } + let mut out_buf_refs: Vec<&mut ModularChannel> = + out_bufs.iter_mut().map(|x| x.deref_mut()).collect(); + super::palette::do_palette_step_one_group( + &img_in, + &img_pal, + &mut out_buf_refs, + grid_x - grid_x0, + grid_y - grid_y0, + grid_x1 - grid_x0, + grid_y1 - grid_y0, + *num_colors, + *num_deltas, + *predictor, + ); + } + buffers[*buf_in].buffer_grid[out_grid].mark_used(); + buffers[*buf_pal].buffer_grid[0].mark_used(); + Ok(buf_out.iter().map(|x| (*x, out_grid)).collect()) + } + TransformStep::Palette { + buf_in, + buf_pal, + buf_out, + num_colors, + num_deltas, + predictor, + wp_header, + } => { + assert_eq!(out_grid_kind, buffers[*buf_in].grid_kind); + assert_eq!(out_size, buffers[*buf_in].info.size); + let mut generated_chunks = Vec::<(usize, usize)>::new(); + let grid_shape = buffers[buf_out[0]].grid_shape; + { + assert_eq!(out_grid % grid_shape.0, 0); + let grid_y = out_grid / grid_shape.0; + let grid_y0 = grid_y.saturating_sub(1); + let grid_y1 = grid_y + 1; + let mut in_bufs = vec![]; + for grid_x in 0..grid_shape.0 { + let grid = grid_y * grid_shape.0 + grid_x; + in_bufs.push(Ref::map( + buffers[*buf_in].buffer_grid[grid].data.borrow(), + |x| x.as_ref().unwrap(), + )); + // Ensure that the output buffers are present. + // TODO(szabadka): Extend the callback to support many grid points. + with_buffers(buffers, buf_out, out_grid + grid_x, |_| Ok(()))?; + } + let in_buf_refs: Vec<&ModularChannel> = + in_bufs.iter().map(|x| x.deref()).collect(); + let img_pal = Ref::map(buffers[*buf_pal].buffer_grid[0].data.borrow(), |x| { + x.as_ref().unwrap() + }); + let mut out_bufs = vec![]; + for i in buf_out { + for grid_y in grid_y0..grid_y1 { + for grid_x in 0..grid_shape.0 { + let grid = grid_y * grid_shape.0 + grid_x; + let buf = &buffers[*i]; + let b = &buf.buffer_grid[grid]; + let data = b.data.borrow_mut(); + out_bufs.push(RefMut::map(data, |x| x.as_mut().unwrap())); + } + } + } + let mut out_buf_refs: Vec<&mut ModularChannel> = + out_bufs.iter_mut().map(|x| x.deref_mut()).collect(); + super::palette::do_palette_step_group_row( + &in_buf_refs, + &img_pal, + &mut out_buf_refs, + grid_y - grid_y0, + grid_shape.0, + *num_colors, + *num_deltas, + *predictor, + wp_header, + )?; + } + buffers[*buf_pal].buffer_grid[0].mark_used(); + for grid_x in 0..grid_shape.0 { + buffers[*buf_in].buffer_grid[out_grid + grid_x].mark_used(); + for buf in buf_out { + generated_chunks.push((*buf, out_grid + grid_x)); + } + } + Ok(generated_chunks) + } + TransformStep::HSqueeze { buf_in, buf_out } => { + let buf_avg = &buffers[buf_in[0]]; + let buf_res = &buffers[buf_in[1]]; + let in_grid = buf_avg.get_grid_idx(out_grid_kind, self.grid_pos); + assert_eq!(out_grid_kind, buf_res.grid_kind); + { + trace!( + "HSqueeze {:?} -> {:?}, grid {out_grid} grid pos {:?}", + buf_in, buf_out, self.grid_pos + ); + let (gx, gy) = self.grid_pos; + let in_avg = Ref::map(buf_avg.buffer_grid[in_grid].data.borrow(), |x| { + x.as_ref().unwrap() + }); + let has_next = gx + 1 < buffers[*buf_out].grid_shape.0; + let gx_next = if has_next { gx + 1 } else { gx }; + let next_avg_grid = buf_avg.get_grid_idx(out_grid_kind, (gx_next, gy)); + let in_next_avg = + Ref::map(buf_avg.buffer_grid[next_avg_grid].data.borrow(), |x| { + x.as_ref().unwrap() + }); + let in_next_avg_rect = if has_next { + Some(in_next_avg.data.as_rect().rect(buf_avg.get_grid_rect( + frame_header, + out_grid_kind, + (gx_next, gy), + ))?) + } else { + None + }; + let in_res = Ref::map(buf_res.buffer_grid[out_grid].data.borrow(), |x| { + x.as_ref().unwrap() + }); + let out_prev = if gx == 0 { + None + } else { + let prev_out_grid = + buffers[*buf_out].get_grid_idx(out_grid_kind, (gx - 1, gy)); + Some(Ref::map( + buffers[*buf_out].buffer_grid[prev_out_grid].data.borrow(), + |x| x.as_ref().unwrap(), + )) + }; + + with_buffers(buffers, &[*buf_out], out_grid, |mut bufs| { + super::squeeze::do_hsqueeze_step( + &in_avg.data.as_rect().rect(buf_avg.get_grid_rect( + frame_header, + out_grid_kind, + (gx, gy), + ))?, + &in_res.data.as_rect().rect(buf_res.get_grid_rect( + frame_header, + out_grid_kind, + (gx, gy), + ))?, + &in_next_avg_rect, + &out_prev, + &mut bufs, + ); + Ok(()) + })?; + } + buffers[buf_in[0]].buffer_grid[in_grid].mark_used(); + buffers[buf_in[1]].buffer_grid[out_grid].mark_used(); + Ok(vec![(*buf_out, out_grid)]) + } + TransformStep::VSqueeze { buf_in, buf_out } => { + let buf_avg = &buffers[buf_in[0]]; + let buf_res = &buffers[buf_in[1]]; + let in_grid = buf_avg.get_grid_idx(out_grid_kind, self.grid_pos); + assert_eq!(out_grid_kind, buf_res.grid_kind); + { + trace!( + "VSqueeze {:?} -> {:?} grid: {out_grid:?} grid pos: {:?}", + buf_in, buf_out, self.grid_pos + ); + let (gx, gy) = self.grid_pos; + let in_avg = Ref::map(buf_avg.buffer_grid[in_grid].data.borrow(), |x| { + x.as_ref().unwrap() + }); + let has_next = gy + 1 < buffers[*buf_out].grid_shape.1; + let gy_next = if has_next { gy + 1 } else { gy }; + let next_avg_grid = buf_avg.get_grid_idx(out_grid_kind, (gx, gy_next)); + let in_next_avg = + Ref::map(buf_avg.buffer_grid[next_avg_grid].data.borrow(), |x| { + x.as_ref().unwrap() + }); + let in_next_avg_rect = if has_next { + Some(in_next_avg.data.as_rect().rect(buf_avg.get_grid_rect( + frame_header, + out_grid_kind, + (gx, gy_next), + ))?) + } else { + None + }; + let in_res = Ref::map(buf_res.buffer_grid[out_grid].data.borrow(), |x| { + x.as_ref().unwrap() + }); + let out_prev = if gy == 0 { + None + } else { + let prev_out_grid = + buffers[*buf_out].get_grid_idx(out_grid_kind, (gx, gy - 1)); + Some(Ref::map( + buffers[*buf_out].buffer_grid[prev_out_grid].data.borrow(), + |x| x.as_ref().unwrap(), + )) + }; + + with_buffers(buffers, &[*buf_out], out_grid, |mut bufs| { + super::squeeze::do_vsqueeze_step( + &in_avg.data.as_rect().rect(buf_avg.get_grid_rect( + frame_header, + out_grid_kind, + (gx, gy), + ))?, + &in_res.data.as_rect().rect(buf_res.get_grid_rect( + frame_header, + out_grid_kind, + (gx, gy), + ))?, + &in_next_avg_rect, + &out_prev, + &mut bufs, + ); + Ok(()) + })?; + } + buffers[buf_in[0]].buffer_grid[in_grid].mark_used(); + buffers[buf_in[1]].buffer_grid[out_grid].mark_used(); + Ok(vec![(*buf_out, out_grid)]) + } + } + } +} + +#[instrument(level = "trace", err)] +fn check_equal_channels( + channels: &[(usize, ChannelInfo)], + first_chan: usize, + num: usize, +) -> Result<()> { + if first_chan + num > channels.len() { + return Err(Error::InvalidChannelRange( + first_chan, + first_chan + num, + channels.len(), + )); + } + for inc in 1..num { + if !channels[first_chan] + .1 + .is_equivalent(&channels[first_chan + inc].1) + { + return Err(Error::MixingDifferentChannels); + } + } + Ok(()) +} + +fn meta_apply_single_transform( + transform: &headers::modular::Transform, + header: &headers::modular::GroupHeader, + channels: &mut Vec<(usize, ChannelInfo)>, + transform_steps: &mut Vec<TransformStep>, + mut add_transform_buffer: impl FnMut(ChannelInfo, String) -> usize, +) -> Result<()> { + match transform.id { + TransformId::Rct => { + let begin_channel = transform.begin_channel as usize; + let op = RctOp::from_u32(transform.rct_type % 7).unwrap(); + let perm = RctPermutation::from_u32(transform.rct_type / 7) + .expect("header decoding should ensure rct_type < 42"); + check_equal_channels(channels, begin_channel, 3)?; + let mut buf_in = [0; 3]; + let buf_out = [ + channels[begin_channel].0, + channels[begin_channel + 1].0, + channels[begin_channel + 2].0, + ]; + for i in 0..3 { + let c = &mut channels[begin_channel + i]; + let mut info = c.1; + info.output_channel_idx = -1; + c.0 = add_transform_buffer( + info, + format!( + "RCT (op {op:?} perm {perm:?}) starting at channel {begin_channel}, \ + input {i}" + ), + ); + buf_in[i] = c.0; + } + transform_steps.push(TransformStep::Rct { + buf_out, + buf_in, + op, + perm, + }); + trace!("applied RCT: {channels:?}"); + } + TransformId::Squeeze => { + let steps = if transform.squeezes.is_empty() { + super::squeeze::default_squeeze(channels) + } else { + transform.squeezes.clone() + }; + for step in steps { + super::squeeze::check_squeeze_params(channels, &step)?; + let in_place = step.in_place; + let horizontal = step.horizontal; + let begin_channel = step.begin_channel as usize; + let num_channels = step.num_channels as usize; + let end_channel = begin_channel + num_channels; + let new_chan_offset = if in_place { + end_channel + } else { + channels.len() + }; + for ic in 0..num_channels { + let chan = &channels[begin_channel + ic].1; + let new_shift = if let Some(shift) = chan.shift { + if shift.0 > 30 || shift.1 > 30 { + return Err(Error::TooManySqueezes); + } + if horizontal { + Some((shift.0 + 1, shift.1)) + } else { + Some((shift.0, shift.1 + 1)) + } + } else { + None + }; + let w = chan.size.0; + let h = chan.size.1; + let (new_size_0, new_size_1) = if horizontal { + ((w.div_ceil(2), h), (w - w.div_ceil(2), h)) + } else { + ((w, h.div_ceil(2)), (w, h - h.div_ceil(2))) + }; + let new_0 = ChannelInfo { + output_channel_idx: -1, + shift: new_shift, + size: new_size_0, + bit_depth: chan.bit_depth, + }; + let buf_0 = add_transform_buffer( + new_0, + format!("Squeezed channel, original channel {}", begin_channel + ic), + ); + let new_1 = ChannelInfo { + output_channel_idx: -1, + shift: new_shift, + size: new_size_1, + bit_depth: chan.bit_depth, + }; + let buf_1 = add_transform_buffer( + new_1, + format!("Squeeze residual, original channel {}", begin_channel + ic), + ); + if horizontal { + transform_steps.push(TransformStep::HSqueeze { + buf_in: [buf_0, buf_1], + buf_out: channels[begin_channel + ic].0, + }); + } else { + transform_steps.push(TransformStep::VSqueeze { + buf_in: [buf_0, buf_1], + buf_out: channels[begin_channel + ic].0, + }); + } + channels[begin_channel + ic] = (buf_0, new_0); + channels.insert(new_chan_offset + ic, (buf_1, new_1)); + trace!("applied squeeze: {channels:?}"); + } + } + } + TransformId::Palette => { + let begin_channel = transform.begin_channel as usize; + let num_channels = transform.num_channels as usize; + let num_colors = transform.num_colors as usize; + let num_deltas = transform.num_deltas as usize; + let pred = Predictor::from_u32(transform.predictor_id) + .expect("header decoding should ensure a valid predictor"); + check_equal_channels(channels, begin_channel, num_channels)?; + // We already checked the bit_depth for all channels from `begin_channel` is + // equal in the line above. + let bit_depth = channels[begin_channel].1.bit_depth; + let pchan_info = ChannelInfo { + output_channel_idx: -1, + shift: None, + size: (num_colors + num_deltas, num_channels), + bit_depth, + }; + let pchan = add_transform_buffer( + pchan_info, + format!( + "Palette for palette transform starting at channel {begin_channel} with \ + {num_channels} channels" + ), + ); + let mut inchan_info = channels[begin_channel].1; + inchan_info.output_channel_idx = -1; + let inchan = add_transform_buffer( + inchan_info, + format!( + "Pixel data for palette transform starting at channel {begin_channel} with \ + {num_channels} channels", + ), + ); + transform_steps.push(TransformStep::Palette { + buf_in: inchan, + buf_pal: pchan, + buf_out: channels[begin_channel..(begin_channel + num_channels)] + .iter() + .map(|x| x.0) + .collect(), + num_colors, + num_deltas, + predictor: pred, + wp_header: header.wp_header.clone(), + }); + channels.drain(begin_channel + 1..begin_channel + num_channels); + channels[begin_channel].0 = inchan; + channels.insert(0, (pchan, pchan_info)); + trace!("applied palette: {channels:?}"); + } + TransformId::Invalid => { + unreachable!("header decoding for invalid transforms should fail"); + } + } + Ok(()) +} + +#[instrument(level = "trace", ret)] +pub fn meta_apply_transforms( + channels: &[ChannelInfo], + header: &headers::modular::GroupHeader, +) -> Result<(Vec<ModularBufferInfo>, Vec<TransformStep>)> { + let mut buffer_info = vec![]; + let mut transform_steps = vec![]; + // (buffer id, channel info) + let mut channels: Vec<_> = channels.iter().cloned().enumerate().collect(); + + // First, add all the pre-transform channels to the buffer list. + for chan in channels.iter() { + buffer_info.push(ModularBufferInfo { + info: chan.1, + coded_channel_id: -1, + description: format!( + "Input channel {}, size {}x{}", + chan.0, chan.1.size.0, chan.1.size.1 + ), + // To be filled by make_grids. + grid_kind: ModularGridKind::None, + grid_shape: (0, 0), + buffer_grid: vec![], + }); + } + + let mut add_transform_buffer = |info, description| { + buffer_info.push(ModularBufferInfo { + info, + coded_channel_id: -1, + description, + // To be filled by make_grids. + grid_kind: ModularGridKind::None, + grid_shape: (0, 0), + buffer_grid: vec![], + }); + buffer_info.len() - 1 + }; + + // Apply transforms to the channel list. + for transform in &header.transforms { + meta_apply_single_transform( + transform, + header, + &mut channels, + &mut transform_steps, + &mut add_transform_buffer, + )?; + } + + // All the channels left over at the end of applying transforms are the channels that are + // actually coded. + for (chid, chan) in channels.iter().enumerate() { + buffer_info[chan.0].coded_channel_id = chid as isize; + } + + #[cfg(feature = "tracing")] + for (i, transform) in transform_steps.iter().enumerate() { + trace!("Transform step {i}: {transform:?}"); + } + + Ok((buffer_info, transform_steps)) +} + +#[derive(Debug)] +pub enum LocalTransformBuffer<'a> { + // This channel has been consumed by some transform. + Empty, + // This channel has not been written to yet. + Placeholder(ChannelInfo), + // Temporary, locally-allocated channel. + Owned(ModularChannel), + // Channel belonging to the global image. + Borrowed(&'a mut ModularChannel), +} + +impl LocalTransformBuffer<'_> { + fn channel_info(&self) -> ChannelInfo { + match self { + LocalTransformBuffer::Empty => unreachable!("an empty buffer has no channel info"), + LocalTransformBuffer::Owned(m) => m.channel_info(), + LocalTransformBuffer::Placeholder(c) => *c, + LocalTransformBuffer::Borrowed(m) => m.channel_info(), + } + } + + fn borrow_mut(&mut self) -> &mut ModularChannel { + match self { + LocalTransformBuffer::Owned(m) => m, + LocalTransformBuffer::Borrowed(m) => m, + LocalTransformBuffer::Empty => unreachable!("tried to borrow an empty channel"), + LocalTransformBuffer::Placeholder(_) => { + unreachable!("tried to borrow a placeholder channel") + } + } + } + + fn take(&mut self) -> Self { + assert!(!matches!(self, LocalTransformBuffer::Empty)); + let mut r = LocalTransformBuffer::Empty; + std::mem::swap(self, &mut r); + r + } + + fn allocate_if_needed(&mut self) -> Result<()> { + if let LocalTransformBuffer::Placeholder(c) = self { + *self = LocalTransformBuffer::Owned(ModularChannel::new_with_shift( + c.size, + c.shift, + c.bit_depth, + )?); + } + Ok(()) + } +} + +#[instrument(level = "trace", ret)] +pub fn meta_apply_local_transforms<'a, 'b>( + channels_in: Vec<&'a mut ModularChannel>, + buffer_storage: &'b mut Vec<LocalTransformBuffer<'a>>, + header: &headers::modular::GroupHeader, +) -> Result<(Vec<&'b mut ModularChannel>, Vec<TransformStep>)> { + let mut transform_steps = vec![]; + + // (buffer id, channel info) + let mut channels: Vec<_> = channels_in + .iter() + .map(|x| x.channel_info()) + .enumerate() + .collect(); + + debug!(?channels, "initial channels"); + + // First, add all the pre-transform channels to the buffer list. + buffer_storage.extend(channels_in.into_iter().map(LocalTransformBuffer::Borrowed)); + + #[allow(unused_variables)] + let mut add_transform_buffer = |info, description| { + trace!(description, ?info, "adding channel buffer"); + buffer_storage.push(LocalTransformBuffer::Placeholder(info)); + buffer_storage.len() - 1 + }; + + // Apply transforms to the channel list. + for transform in &header.transforms { + meta_apply_single_transform( + transform, + header, + &mut channels, + &mut transform_steps, + &mut add_transform_buffer, + )?; + } + + debug!(?channels, ?buffer_storage, "channels after transforms"); + debug!(?transform_steps); + + // Ensure that the buffer indices in `channels` appear in increasing order, by reordering them + // if necessary. + if !channels.iter().map(|x| x.0).is_sorted() { + let mut buf_new_position: Vec<_> = channels.iter().map(|x| x.0).collect(); + buf_new_position.sort(); + let buf_tmp: Vec<_> = channels + .iter() + .map(|x| { + let mut b = LocalTransformBuffer::Empty; + std::mem::swap(&mut b, &mut buffer_storage[x.0]); + b + }) + .collect(); + + let mut buf_remap: Vec<_> = (0..buffer_storage.len()).collect(); + + for (new_pos, (ch_info, buf)) in buf_new_position + .iter() + .cloned() + .zip(channels.iter_mut().zip(buf_tmp.into_iter())) + { + assert!(matches!( + buffer_storage[new_pos], + LocalTransformBuffer::Empty + )); + buf_remap[ch_info.0] = new_pos; + buffer_storage[new_pos] = buf; + ch_info.0 = new_pos; + } + + for step in transform_steps.iter_mut() { + use std::iter::once; + match step { + TransformStep::Rct { + buf_in, buf_out, .. + } => { + for b in buf_in.iter_mut().chain(buf_out.iter_mut()) { + *b = buf_remap[*b]; + } + } + TransformStep::Palette { + buf_in, + buf_pal, + buf_out, + .. + } => { + for b in once(buf_in).chain(once(buf_pal)).chain(buf_out.iter_mut()) { + *b = buf_remap[*b]; + } + } + TransformStep::HSqueeze { buf_in, buf_out } + | TransformStep::VSqueeze { buf_in, buf_out } => { + for b in once(buf_out).chain(buf_in.iter_mut()) { + *b = buf_remap[*b]; + } + } + } + } + } + + debug!(?channels, ?buffer_storage, "sorted channels"); + + debug!(?transform_steps); + + // Since RCT steps will try to transfer buffers from the source channels to the destination + // channels, make sure we do the reverse transformation here (to have the caller-provided + // buffers be used for writing temporary data). + for ts in transform_steps.iter() { + if let TransformStep::Rct { + buf_in, buf_out, .. + } = ts + { + for c in 0..3 { + assert_eq!( + buffer_storage[buf_in[c]].channel_info(), + buffer_storage[buf_out[c]].channel_info() + ); + assert!(matches!( + buffer_storage[buf_in[c]], + LocalTransformBuffer::Placeholder(_) + )); + buffer_storage.swap(buf_in[c], buf_out[c]); + } + } + } + + debug!(?channels, ?buffer_storage, "RCT-adjusted channels"); + + // Allocate all the coded channels if they aren't yet. + for (buf, _) in channels.iter() { + buffer_storage[*buf].allocate_if_needed()?; + } + + debug!(?channels, ?buffer_storage, "allocated buffers"); + + // Extract references to to-be-decoded buffers. + let mut coded_buffers = Vec::with_capacity(channels.len()); + let mut buffer_tail = &mut buffer_storage[..]; + let mut last_buffer = None; + for (buf, _) in channels { + let offset = if let Some(lb) = last_buffer { + buf.checked_sub(lb).unwrap() + } else { + buf + 1 + }; + let cur_buf; + (cur_buf, buffer_tail) = buffer_tail.split_at_mut(offset); + coded_buffers.push(cur_buf.last_mut().unwrap().borrow_mut()); + last_buffer = Some(buf); + } + + Ok((coded_buffers, transform_steps)) +} + +impl TransformStep { + // Marks that one dependency of this transform is ready, and potentially runs the transform, + // returning the new buffers that are now ready. + pub fn local_apply(&self, buffers: &mut [LocalTransformBuffer]) -> Result<()> { + match self { + TransformStep::Rct { + buf_in, + buf_out, + op, + perm, + } => { + for i in 0..3 { + assert_eq!( + buffers[buf_in[i]].channel_info(), + buffers[buf_out[i]].channel_info() + ); + } + let [mut a, mut b, mut c] = [ + buffers[buf_in[0]].take(), + buffers[buf_in[1]].take(), + buffers[buf_in[2]].take(), + ]; + { + let mut bufs = [a.borrow_mut(), b.borrow_mut(), c.borrow_mut()]; + super::rct::do_rct_step(&mut bufs, *op, *perm); + } + buffers[buf_out[0]] = a; + buffers[buf_out[1]] = b; + buffers[buf_out[2]] = c; + } + TransformStep::Palette { + buf_in, + buf_pal, + buf_out, + num_colors, + num_deltas, + predictor, + wp_header, + } => { + for b in buf_out.iter() { + assert_eq!( + buffers[*b].channel_info().size, + buffers[*buf_in].channel_info().size + ); + buffers[*b].allocate_if_needed()?; + } + let mut img_in = buffers[*buf_in].take(); + let mut img_pal = buffers[*buf_pal].take(); + let mut out_bufs: Vec<_> = buf_out.iter().map(|x| buffers[*x].take()).collect(); + { + let mut bufs: Vec<_> = out_bufs.iter_mut().map(|x| x.borrow_mut()).collect(); + super::palette::do_palette_step_general( + img_in.borrow_mut(), + img_pal.borrow_mut(), + &mut bufs, + *num_colors, + *num_deltas, + *predictor, + wp_header, + ); + } + for (pos, buf) in buf_out.iter().zip(out_bufs.into_iter()) { + buffers[*pos] = buf; + } + } + TransformStep::HSqueeze { buf_in, buf_out } => { + buffers[*buf_out].allocate_if_needed()?; + let mut out_buf = buffers[*buf_out].take(); + let mut in_avg = buffers[buf_in[0]].take(); + let mut in_res = buffers[buf_in[1]].take(); + { + let mut bufs: Vec<_> = vec![out_buf.borrow_mut()]; + super::squeeze::do_hsqueeze_step( + &in_avg.borrow_mut().data.as_rect(), + &in_res.borrow_mut().data.as_rect(), + &None, + &None, + &mut bufs, + ); + } + buffers[*buf_out] = out_buf; + } + TransformStep::VSqueeze { buf_in, buf_out } => { + buffers[*buf_out].allocate_if_needed()?; + let mut out_buf = buffers[*buf_out].take(); + let mut in_avg = buffers[buf_in[0]].take(); + let mut in_res = buffers[buf_in[1]].take(); + { + let mut bufs: Vec<_> = vec![out_buf.borrow_mut()]; + super::squeeze::do_vsqueeze_step( + &in_avg.borrow_mut().data.as_rect(), + &in_res.borrow_mut().data.as_rect(), + &None, + &None, + &mut bufs, + ); + } + buffers[*buf_out] = out_buf; + } + }; + + Ok(()) + } +} diff --git a/third_party/rust/jxl/src/frame/modular/transforms/palette.rs b/third_party/rust/jxl/src/frame/modular/transforms/palette.rs @@ -0,0 +1,452 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +use crate::{ + error::Result, + frame::modular::{ + ModularChannel, Predictor, + predict::{PredictionData, WeightedPredictorState}, + }, + headers::modular::WeightedHeader, + image::{Image, ImageRect}, +}; + +const RGB_CHANNELS: usize = 3; + +// 5x5x5 color cube for the larger cube. +const LARGE_CUBE: usize = 5; + +// Smaller interleaved color cube to fill the holes of the larger cube. +const SMALL_CUBE: usize = 4; +const SMALL_CUBE_BITS: usize = 2; +// SMALL_CUBE ** 3 +const LARGE_CUBE_OFFSET: usize = SMALL_CUBE * SMALL_CUBE * SMALL_CUBE; + +fn scale<const DENOM: usize>(value: usize, bit_depth: usize) -> i32 { + // return (value * ((1 << bit_depth) - 1)) / DENOM; + // We only call this function with SMALL_CUBE or LARGE_CUBE - 1 as DENOM, + // allowing us to avoid a division here. + const { + assert!(DENOM == 4, "denom must be 4"); + } + ((value * ((1 << bit_depth) - 1)) >> 2) as i32 +} + +// The purpose of this function is solely to extend the interpretation of +// palette indices to implicit values. If index < nb_deltas, indicating that the +// result is a delta palette entry, it is the responsibility of the caller to +// treat it as such. +fn get_palette_value( + palette: &ImageRect<i32>, + index: isize, + c: usize, + palette_size: usize, + bit_depth: usize, +) -> i32 { + if index < 0 { + const DELTA_PALETTE: [[i32; 3]; 72] = [ + [0, 0, 0], + [4, 4, 4], + [11, 0, 0], + [0, 0, -13], + [0, -12, 0], + [-10, -10, -10], + [-18, -18, -18], + [-27, -27, -27], + [-18, -18, 0], + [0, 0, -32], + [-32, 0, 0], + [-37, -37, -37], + [0, -32, -32], + [24, 24, 45], + [50, 50, 50], + [-45, -24, -24], + [-24, -45, -45], + [0, -24, -24], + [-34, -34, 0], + [-24, 0, -24], + [-45, -45, -24], + [64, 64, 64], + [-32, 0, -32], + [0, -32, 0], + [-32, 0, 32], + [-24, -45, -24], + [45, 24, 45], + [24, -24, -45], + [-45, -24, 24], + [80, 80, 80], + [64, 0, 0], + [0, 0, -64], + [0, -64, -64], + [-24, -24, 45], + [96, 96, 96], + [64, 64, 0], + [45, -24, -24], + [34, -34, 0], + [112, 112, 112], + [24, -45, -45], + [45, 45, -24], + [0, -32, 32], + [24, -24, 45], + [0, 96, 96], + [45, -24, 24], + [24, -45, -24], + [-24, -45, 24], + [0, -64, 0], + [96, 0, 0], + [128, 128, 128], + [64, 0, 64], + [144, 144, 144], + [96, 96, 0], + [-36, -36, 36], + [45, -24, -45], + [45, -45, -24], + [0, 0, -96], + [0, 128, 128], + [0, 96, 0], + [45, 24, -45], + [-128, 0, 0], + [24, -45, 24], + [-45, 24, -45], + [64, 0, -64], + [64, -64, -64], + [96, 0, 96], + [45, -45, 24], + [24, 45, -45], + [64, 64, -64], + [128, 128, 0], + [0, 0, -128], + [-24, 45, -45], + ]; + if c >= RGB_CHANNELS { + return 0; + } + // Do not open the brackets, otherwise INT32_MIN negation could overflow. + let mut index = -(index + 1) as usize; + index %= 1 + 2 * (DELTA_PALETTE.len() - 1); + const MULTIPLIER: [i32; 2] = [-1, 1]; + let mut result = DELTA_PALETTE[(index + 1) >> 1][c] * MULTIPLIER[index & 1]; + if bit_depth > 8 { + result *= 1 << (bit_depth - 8); + } + result + } else { + let mut index = index as usize; + if palette_size <= index && index < palette_size + LARGE_CUBE_OFFSET { + if c >= RGB_CHANNELS { + return 0; + } + index -= palette_size; + index >>= c * SMALL_CUBE_BITS; + scale::<SMALL_CUBE>(index % SMALL_CUBE, bit_depth) + + (1 << (0.max(bit_depth as isize - 3))) + } else if palette_size + LARGE_CUBE_OFFSET <= index { + if c >= RGB_CHANNELS { + return 0; + } + index -= palette_size + LARGE_CUBE_OFFSET; + // TODO(eustas): should we take care of ambiguity created by + // index >= LARGE_CUBE ** 3 ? + match c { + 0 => (), + 1 => { + index /= LARGE_CUBE; + } + 2 => { + index /= LARGE_CUBE * LARGE_CUBE; + } + _ => (), + } + scale::<{ LARGE_CUBE - 1 }>(index % LARGE_CUBE, bit_depth) + } else { + palette.row(c)[index] + } + } +} + +pub fn do_palette_step_general( + buf_in: &ModularChannel, + buf_pal: &ModularChannel, + buf_out: &mut [&mut ModularChannel], + num_colors: usize, + num_deltas: usize, + predictor: Predictor, + wp_header: &WeightedHeader, +) { + let (w, h) = buf_in.data.size(); + let palette = buf_pal.data.as_rect(); + let bit_depth = buf_in.bit_depth.bits_per_sample().min(24) as usize; + + if w == 0 { + // Nothing to do. + // Avoid touching "empty" channels with non-zero height. + } else if num_deltas == 0 && predictor == Predictor::Zero { + for (chan_index, out) in buf_out.iter_mut().enumerate() { + for y in 0..h { + for x in 0..w { + let index = buf_in.data.as_rect().row(y)[x]; + let palette_value = get_palette_value( + &palette, + index as isize, + /*c=*/ chan_index, + /*palette_size=*/ num_colors, + /*bit_depth=*/ bit_depth, + ); + out.data.as_rect_mut().row(y)[x] = palette_value; + } + } + } + } else if predictor == Predictor::Weighted { + let w = buf_in.data.size().0; + for (chan_index, out) in buf_out.iter_mut().enumerate() { + let mut wp_state = WeightedPredictorState::new(wp_header, w); + for y in 0..h { + let idx = buf_in.data.as_rect().row(y); + for (x, &index) in idx.iter().enumerate() { + let palette_entry = get_palette_value( + &palette, + index as isize, + /*c=*/ chan_index, + /*palette_size=*/ num_colors + num_deltas, + /*bit_depth=*/ bit_depth, + ); + let val = if index < num_deltas as i32 { + let prediction_data = PredictionData::get(out.data.as_rect(), x, y); + let (wp_pred, _) = + wp_state.predict_and_property((x, y), w, &prediction_data); + let pred = predictor.predict_one(prediction_data, wp_pred); + (pred + palette_entry as i64) as i32 + } else { + palette_entry + }; + out.data.as_rect_mut().row(y)[x] = val; + wp_state.update_errors(val, (x, y), w); + } + } + } + } else { + for (chan_index, out) in buf_out.iter_mut().enumerate() { + for y in 0..h { + let idx = buf_in.data.as_rect().row(y); + for (x, &index) in idx.iter().enumerate() { + let palette_entry = get_palette_value( + &palette, + index as isize, + /*c=*/ chan_index, + /*palette_size=*/ num_colors + num_deltas, + /*bit_depth=*/ bit_depth, + ); + let val = if index < num_deltas as i32 { + let pred = predictor.predict_one( + PredictionData::get(out.data.as_rect(), x, y), + /*wp_pred=*/ 0, + ); + (pred + palette_entry as i64) as i32 + } else { + palette_entry + }; + out.data.as_rect_mut().row(y)[x] = val; + } + } + } + } +} + +#[allow(clippy::too_many_arguments)] +fn get_prediction_data( + buf: &mut [&mut ModularChannel], + idx: usize, + grid_x: usize, + grid_y: usize, + grid_xsize: usize, + x: usize, + y: usize, + xsize: usize, + ysize: usize, +) -> PredictionData { + PredictionData::get_with_neighbors( + buf[idx].data.as_rect(), + if grid_x > 0 { + Some(buf[idx - 1].data.as_rect()) + } else { + None + }, + if grid_y > 0 { + Some(buf[idx - grid_xsize].data.as_rect()) + } else { + None + }, + if grid_x > 0 && grid_y > 0 { + Some(buf[idx - grid_xsize - 1].data.as_rect()) + } else { + None + }, + if grid_x + 1 < grid_xsize { + Some(buf[idx + 1].data.as_rect()) + } else { + None + }, + if grid_x + 1 < grid_xsize && grid_y > 0 { + Some(buf[idx - grid_xsize + 1].data.as_rect()) + } else { + None + }, + x, + y, + xsize, + ysize, + ) +} + +#[allow(clippy::too_many_arguments)] +pub fn do_palette_step_one_group( + buf_in: &ModularChannel, + buf_pal: &ModularChannel, + buf_out: &mut [&mut ModularChannel], + grid_x: usize, + grid_y: usize, + grid_xsize: usize, + grid_ysize: usize, + num_colors: usize, + num_deltas: usize, + predictor: Predictor, +) { + let h = buf_in.data.size().1; + let palette = buf_pal.data.as_rect(); + let bit_depth = buf_in.bit_depth.bits_per_sample().min(24) as usize; + let num_c = buf_out.len() / (grid_xsize * grid_ysize); + let (xsize, ysize) = buf_out[0].data.size(); + + for c in 0..num_c { + for y in 0..h { + let index_img = buf_in.data.as_rect().row(y); + let out_idx = c * grid_ysize * grid_xsize + grid_y * grid_xsize + grid_x; + for (x, &index) in index_img.iter().enumerate() { + let palette_entry = get_palette_value( + &palette, + index as isize, + c, + /*palette_size=*/ num_colors + num_deltas, + /*bit_depth=*/ bit_depth, + ); + let val = if index < num_deltas as i32 { + let pred = predictor.predict_one( + get_prediction_data( + buf_out, out_idx, grid_x, grid_y, grid_xsize, x, y, xsize, ysize, + ), + /*wp_pred=*/ 0, + ); + (pred + palette_entry as i64) as i32 + } else { + palette_entry + }; + buf_out[out_idx].data.as_rect_mut().row(y)[x] = val; + } + } + } +} + +#[allow(clippy::too_many_arguments)] +pub fn do_palette_step_group_row( + buf_in: &[&ModularChannel], + buf_pal: &ModularChannel, + buf_out: &mut [&mut ModularChannel], + grid_y: usize, + grid_xsize: usize, + num_colors: usize, + num_deltas: usize, + predictor: Predictor, + wp_header: &WeightedHeader, +) -> Result<()> { + let palette = buf_pal.data.as_rect(); + let h = buf_in[0].data.size().1; + let bit_depth = buf_in[0].bit_depth.bits_per_sample().min(24) as usize; + let grid_ysize = grid_y + 1; + let num_c = buf_out.len() / (grid_xsize * grid_ysize); + let total_w = buf_out[0..grid_xsize] + .iter() + .map(|buf| buf.data.size().0) + .sum(); + let (xsize, ysize) = buf_out[0].data.size(); + + if predictor == Predictor::Weighted { + for c in 0..num_c { + let mut wp_state = WeightedPredictorState::new(wp_header, total_w); + let out_row_idx = c * grid_ysize * grid_xsize + grid_y * grid_xsize; + if grid_y > 0 { + let prev_row_idx = out_row_idx - grid_y * grid_xsize; + wp_state.restore_state( + buf_out[prev_row_idx].auxiliary_data.as_ref().unwrap(), + total_w, + ); + } + for y in 0..h { + for (grid_x, index_buf) in buf_in.iter().enumerate().take(grid_xsize) { + let index_img = index_buf.data.as_rect().row(y); + let out_idx = out_row_idx + grid_x; + for (x, &index) in index_img.iter().enumerate() { + let palette_entry = get_palette_value( + &palette, + index as isize, + c, + /*palette_size=*/ num_colors + num_deltas, + /*bit_depth=*/ bit_depth, + ); + let val = if index < num_deltas as i32 { + let prediction_data = get_prediction_data( + buf_out, out_idx, grid_x, grid_y, grid_xsize, x, y, xsize, ysize, + ); + let (pred, _) = wp_state.predict_and_property( + (grid_x * xsize + x, y & 1), + total_w, + &prediction_data, + ); + (pred + palette_entry as i64) as i32 + } else { + palette_entry + }; + buf_out[out_idx].data.as_rect_mut().row(y)[x] = val; + wp_state.update_errors(val, (grid_x * xsize + x, y & 1), total_w); + } + } + } + let mut wp_image = Image::<i32>::new((total_w + 2, 1))?; + wp_state.save_state(&mut wp_image, total_w); + buf_out[out_row_idx].auxiliary_data = Some(wp_image); + } + } else { + for c in 0..num_c { + for y in 0..h { + for (grid_x, index_buf) in buf_in.iter().enumerate().take(grid_xsize) { + let index_img = index_buf.data.as_rect().row(y); + let out_idx = c * grid_ysize * grid_xsize + grid_y * grid_xsize + grid_x; + for (x, &index) in index_img.iter().enumerate() { + let palette_entry = get_palette_value( + &palette, + index as isize, + c, + /*palette_size=*/ num_colors + num_deltas, + /*bit_depth=*/ bit_depth, + ); + let val = if index < num_deltas as i32 { + let pred = predictor.predict_one( + get_prediction_data( + buf_out, out_idx, grid_x, grid_y, grid_xsize, x, y, xsize, + ysize, + ), + /*wp_pred=*/ 0, + ); + (pred + palette_entry as i64) as i32 + } else { + palette_entry + }; + buf_out[out_idx].data.as_rect_mut().row(y)[x] = val; + } + } + } + } + } + Ok(()) +} diff --git a/third_party/rust/jxl/src/frame/modular/transforms/rct.rs b/third_party/rust/jxl/src/frame/modular/transforms/rct.rs @@ -0,0 +1,92 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +use crate::{ + frame::modular::{ + ModularChannel, + transforms::{RctOp, RctPermutation}, + }, + util::tracing_wrappers::*, +}; + +// Applies a RCT in-place to the given buffers. +#[instrument(level = "debug", skip(buffers), ret)] +pub fn do_rct_step(buffers: &mut [&mut ModularChannel], op: RctOp, perm: RctPermutation) { + let size = buffers[0].data.size(); + + let [r, g, b] = buffers else { + unreachable!("incorrect buffer count for RCT"); + }; + + let buffers = [r, g, b]; + + 'rct: { + let apply_rct: fn(i32, i32, i32) -> (i32, i32, i32) = match op { + RctOp::Noop => break 'rct, + RctOp::YCoCg => |y, co, cg| { + let y = y.wrapping_sub(cg >> 1); + let g = cg.wrapping_add(y); + let y = y.wrapping_sub(co >> 1); + let r = y.wrapping_add(co); + (r, g, y) + }, + RctOp::AddFirstToThird => |v0, v1, v2| (v0, v1, v2.wrapping_add(v0)), + RctOp::AddFirstToSecond => |v0, v1, v2| (v0, v1.wrapping_add(v0), v2), + RctOp::AddFirstToSecondAndThird => { + |v0, v1, v2| (v0, v1.wrapping_add(v0), v2.wrapping_add(v0)) + } + RctOp::AddAvgToSecond => { + |v0, v1, v2| (v0, v1.wrapping_add((v0.wrapping_add(v2)) >> 1), v2) + } + RctOp::AddFirstToThirdAndAvgToSecond => |v0, v1, v2| { + let v2 = v0.wrapping_add(v2); + (v0, v1.wrapping_add((v0.wrapping_add(v2)) >> 1), v2) + }, + }; + + for pos_y in 0..size.1 { + for pos_x in 0..size.0 { + let [v0, v1, v2] = [0, 1, 2].map(|x| buffers[x].data.as_rect().row(pos_y)[pos_x]); + let (w0, w1, w2) = apply_rct(v0, v1, v2); + for (i, p) in [w0, w1, w2].iter().enumerate() { + buffers[i].data.as_rect_mut().row(pos_y)[pos_x] = *p; + } + } + } + } + + let [r, g, b] = buffers; + + // Note: Gbr and Brg use the *inverse* permutation compared to libjxl, because we *first* write + // to the buffers and then permute them, while in libjxl the buffers to be written to are + // permuted first. + // The same is true for Rbg/Grb/Bgr, but since those are involutions it doesn't change + // anything. + match perm { + RctPermutation::Rgb => {} + RctPermutation::Gbr => { + // out[1, 2, 0] = in[0, 1, 2] + std::mem::swap(&mut g.data, &mut b.data); // [1, 0, 2] + std::mem::swap(&mut r.data, &mut g.data); + } + RctPermutation::Brg => { + // out[2, 0, 1] = in[0, 1, 2] + std::mem::swap(&mut r.data, &mut b.data); // [1, 0, 2] + std::mem::swap(&mut r.data, &mut g.data); + } + RctPermutation::Rbg => { + // out[0, 2, 1] = in[0, 1, 2] + std::mem::swap(&mut b.data, &mut g.data); + } + RctPermutation::Grb => { + // out[1, 0, 2] = in[0, 1, 2] + std::mem::swap(&mut r.data, &mut g.data); + } + RctPermutation::Bgr => { + // out[2, 1, 0] = in[0, 1, 2] + std::mem::swap(&mut r.data, &mut b.data); + } + } +} diff --git a/third_party/rust/jxl/src/frame/modular/transforms/squeeze.rs b/third_party/rust/jxl/src/frame/modular/transforms/squeeze.rs @@ -0,0 +1,272 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +use std::cell::Ref; + +use crate::{ + error::{Error, Result}, + frame::modular::{ChannelInfo, ModularChannel}, + headers::modular::SqueezeParams, + image::ImageRect, +}; + +use crate::util::tracing_wrappers::*; + +#[instrument(level = "trace", err)] +pub fn check_squeeze_params( + channels: &[(usize, ChannelInfo)], + params: &SqueezeParams, +) -> Result<()> { + let end_channel = (params.begin_channel + params.num_channels) as usize; + if end_channel > channels.len() { + return Err(Error::InvalidChannelRange( + params.begin_channel as usize, + params.num_channels as usize, + channels.len(), + )); + } + if channels[params.begin_channel as usize].1.is_meta() != channels[end_channel - 1].1.is_meta() + { + return Err(Error::MixingDifferentChannels); + } + if channels[params.begin_channel as usize].1.is_meta() && !params.in_place { + return Err(Error::MetaSqueezeRequiresInPlace); + } + Ok(()) +} + +pub fn default_squeeze(data_channel_info: &[(usize, ChannelInfo)]) -> Vec<SqueezeParams> { + let num_meta_channels = data_channel_info + .iter() + .take_while(|x| x.1.is_meta()) + .count(); + + let mut w = data_channel_info[num_meta_channels].1.size.0; + let mut h = data_channel_info[num_meta_channels].1.size.1; + let nc = data_channel_info.len() - num_meta_channels; + + let mut params = vec![]; + + if nc > 2 && data_channel_info[1].1.size == (w, h) { + // 420 previews + let sp = SqueezeParams { + horizontal: true, + in_place: false, + begin_channel: num_meta_channels as u32 + 1, + num_channels: 2, + }; + params.push(sp); + params.push(SqueezeParams { + horizontal: false, + ..sp + }); + } + + const MAX_FIRST_PREVIEW_SIZE: usize = 8; + + let sp = SqueezeParams { + begin_channel: num_meta_channels as u32, + num_channels: nc as u32, + in_place: true, + horizontal: false, + }; + + // vertical first on tall images + if w <= h && h > MAX_FIRST_PREVIEW_SIZE { + params.push(SqueezeParams { + horizontal: false, + ..sp + }); + h = h.div_ceil(2); + } + while w > MAX_FIRST_PREVIEW_SIZE || h > MAX_FIRST_PREVIEW_SIZE { + if w > MAX_FIRST_PREVIEW_SIZE { + params.push(SqueezeParams { + horizontal: true, + ..sp + }); + w = w.div_ceil(2); + } + if h > MAX_FIRST_PREVIEW_SIZE { + params.push(SqueezeParams { + horizontal: false, + ..sp + }); + h = h.div_ceil(2); + } + } + + params +} + +#[inline(always)] +fn smooth_tendency(b: i64, a: i64, n: i64) -> i64 { + let mut diff = 0; + if b >= a && a >= n { + diff = (4 * b - 3 * n - a + 6) / 12; + // 2c = a<<1 + diff - diff&1 <= 2b so diff - diff&1 <= 2b - 2a + // 2d = a<<1 - diff - diff&1 >= 2n so diff + diff&1 <= 2a - 2n + if diff - (diff & 1) > 2 * (b - a) { + diff = 2 * (b - a) + 1; + } + if diff + (diff & 1) > 2 * (a - n) { + diff = 2 * (a - n); + } + } else if b <= a && a <= n { + diff = (4 * b - 3 * n - a - 6) / 12; + // 2c = a<<1 + diff + diff&1 >= 2b so diff + diff&1 >= 2b - 2a + // 2d = a<<1 - diff + diff&1 <= 2n so diff - diff&1 >= 2a - 2n + if diff + (diff & 1) < 2 * (b - a) { + diff = 2 * (b - a) - 1; + } + if diff - (diff & 1) < 2 * (a - n) { + diff = 2 * (a - n); + } + } + diff +} + +#[inline(always)] +fn unsqueeze(avg: i32, res: i32, next_avg: i32, prev: i32) -> (i32, i32) { + let tendency = smooth_tendency(prev as i64, avg as i64, next_avg as i64); + let diff = (res as i64) + tendency; + let a = (avg as i64) + (diff / 2); + let b = a - diff; + (a as i32, b as i32) +} + +pub fn do_hsqueeze_step( + in_avg: &ImageRect<'_, i32>, + in_res: &ImageRect<'_, i32>, + in_next_avg: &Option<ImageRect<'_, i32>>, + out_prev: &Option<Ref<'_, ModularChannel>>, + buffers: &mut [&mut ModularChannel], +) { + trace!("hsqueeze step in_avg: {in_avg:?} in_res: {in_res:?} in_next_avg: {in_next_avg:?}"); + let out = buffers.first_mut().unwrap(); + // Shortcut: guarantees that row is at least 1px in the main loop + if out.data.size().0 == 0 { + return; + } + let (w, h) = in_res.size(); + // Another shortcut: when output row has just 1px + if w == 0 { + for y in 0..h { + out.data.as_rect_mut().row(y)[0] = in_avg.row(y)[0]; + } + return; + } + // Otherwise: 2 or more in in row + + debug_assert!(w >= 1); + let has_tail = out.data.size().0 & 1 == 1; + if has_tail { + debug_assert!(in_avg.size().0 == w + 1); + debug_assert!(out.data.size().0 == 2 * w + 1); + } + + for y in 0..h { + let avg_row = in_avg.row(y); + let res_row = in_res.row(y); + let mut prev_b = match out_prev { + None => avg_row[0], + Some(mc) => mc.data.as_rect().row(y)[mc.data.size().0 - 1], + }; + // Guarantee that `avg_row[x + 1]` is available. + let x_end = if has_tail { w } else { w - 1 }; + for x in 0..x_end { + let (a, b) = unsqueeze(avg_row[x], res_row[x], avg_row[x + 1], prev_b); + out.data.as_rect_mut().row(y)[2 * x] = a; + out.data.as_rect_mut().row(y)[2 * x + 1] = b; + prev_b = b; + } + if !has_tail { + let last_avg = match in_next_avg { + None => avg_row[w - 1], + Some(mc) => mc.row(y)[0], + }; + let (a, b) = unsqueeze(avg_row[w - 1], res_row[w - 1], last_avg, prev_b); + out.data.as_rect_mut().row(y)[2 * w - 2] = a; + out.data.as_rect_mut().row(y)[2 * w - 1] = b; + } else { + // 1 last pixel + out.data.as_rect_mut().row(y)[2 * w] = in_avg.row(y)[w]; + } + } +} + +pub fn do_vsqueeze_step( + in_avg: &ImageRect<'_, i32>, + in_res: &ImageRect<'_, i32>, + in_next_avg: &Option<ImageRect<'_, i32>>, + out_prev: &Option<Ref<'_, ModularChannel>>, + buffers: &mut [&mut ModularChannel], +) { + trace!("vsqueeze step in_avg: {in_avg:?} in_res: {in_res:?} in_next_avg: {in_next_avg:?}"); + let mut out = buffers.first_mut().unwrap().data.as_rect_mut(); + // Shortcut: guarantees that there at least 1 output row + if out.size().1 == 0 { + return; + } + let (w, h) = in_res.size(); + // Another shortcut: when there is one output row + if h == 0 { + out.row(0).copy_from_slice(in_avg.row(0)); + return; + } + // Otherwise: 2 or more rows + + debug_assert!(h > 0); // i.e. h - 1 >= 0 + let has_tail = out.size().1 & 1 == 1; + if has_tail { + debug_assert!(in_avg.size().1 == h + 1); + debug_assert!(out.size().1 == 2 * h + 1); + } + + { + let prev_b_row = match out_prev { + None => in_avg.row(0), + Some(mc) => mc.data.as_rect().row(mc.data.size().1 - 1), + }; + let avg_row = in_avg.row(0); + let res_row = in_res.row(0); + let avg_row_next = if !has_tail && (h == 1) { + debug_assert!(in_next_avg.is_none()); + in_avg.row(0) + } else { + in_avg.row(1) + }; + for x in 0..w { + let (a, b) = unsqueeze(avg_row[x], res_row[x], avg_row_next[x], prev_b_row[x]); + out.row(0)[x] = a; + out.row(1)[x] = b; + } + } + for y in 1..h { + let avg_row = in_avg.row(y); + let res_row = in_res.row(y); + let avg_row_next = if has_tail || y < h - 1 { + in_avg.row(y + 1) + } else { + match in_next_avg { + None => avg_row, + Some(mc) => mc.row(0), + } + }; + for x in 0..w { + let (a, b) = unsqueeze( + avg_row[x], + res_row[x], + avg_row_next[x], + out.row(2 * y - 1)[x], + ); + out.row(2 * y)[x] = a; + out.row(2 * y + 1)[x] = b; + } + } + if has_tail { + out.row(2 * h).copy_from_slice(in_avg.row(h)); + } +} diff --git a/third_party/rust/jxl/src/frame/modular/tree.rs b/third_party/rust/jxl/src/frame/modular/tree.rs @@ -0,0 +1,314 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +use std::fmt::Debug; + +use super::{ModularChannel, Predictor, predict::WeightedPredictorState}; +use crate::{ + bit_reader::BitReader, + entropy_coding::decode::Histograms, + entropy_coding::decode::SymbolReader, + error::{Error, Result}, + frame::modular::predict::PredictionData, + image::Image, + util::{NewWithCapacity, tracing_wrappers::*}, +}; + +#[derive(Debug)] +pub enum TreeNode { + Split { + property: u8, + val: i32, + left: u32, + right: u32, + }, + Leaf { + predictor: Predictor, + offset: i32, + multiplier: u32, + id: u32, + }, +} + +pub struct Tree { + pub nodes: Vec<TreeNode>, + pub histograms: Histograms, +} + +impl Debug for Tree { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "Tree[{:?}]", self.nodes) + } +} + +#[derive(Debug)] +pub struct PredictionResult { + pub guess: i64, + pub multiplier: u32, + pub context: u32, +} + +pub const NUM_NONREF_PROPERTIES: usize = 16; +pub const PROPERTIES_PER_PREVCHAN: usize = 4; + +const SPLIT_VAL_CONTEXT: usize = 0; +const PROPERTY_CONTEXT: usize = 1; +const PREDICTOR_CONTEXT: usize = 2; +const OFFSET_CONTEXT: usize = 3; +const MULTIPLIER_LOG_CONTEXT: usize = 4; +const MULTIPLIER_BITS_CONTEXT: usize = 5; +const NUM_TREE_CONTEXTS: usize = 6; + +impl Tree { + #[instrument(level = "debug", skip(br), err)] + pub fn read(br: &mut BitReader, size_limit: usize) -> Result<Tree> { + assert!(size_limit <= u32::MAX as usize); + trace!(pos = br.total_bits_read()); + let tree_histograms = Histograms::decode(NUM_TREE_CONTEXTS, br, true)?; + let mut tree_reader = SymbolReader::new(&tree_histograms, br, None)?; + // TODO(veluca): consider early-exiting for trees known to be infinite. + let mut tree: Vec<TreeNode> = vec![]; + let mut to_decode = 1; + let mut leaf_id = 0; + let mut max_property = 0; + while to_decode > 0 { + if tree.len() > size_limit { + return Err(Error::TreeTooLarge(tree.len(), size_limit)); + } + if tree.len() >= tree.capacity() { + tree.try_reserve(tree.len() * 2 + 1)?; + } + to_decode -= 1; + let property = tree_reader.read_unsigned(&tree_histograms, br, PROPERTY_CONTEXT)?; + trace!(property); + if let Some(property) = property.checked_sub(1) { + // inner node. + if property > 255 { + return Err(Error::InvalidProperty(property)); + } + max_property = max_property.max(property); + let splitval = tree_reader.read_signed(&tree_histograms, br, SPLIT_VAL_CONTEXT)?; + let left_child = (tree.len() + to_decode + 1) as u32; + let node = TreeNode::Split { + property: property as u8, + val: splitval, + left: left_child, + right: left_child + 1, + }; + trace!("split node {:?}", node); + to_decode += 2; + tree.push(node); + } else { + let predictor = Predictor::try_from(tree_reader.read_unsigned( + &tree_histograms, + br, + PREDICTOR_CONTEXT, + )?)?; + let offset = tree_reader.read_signed(&tree_histograms, br, OFFSET_CONTEXT)?; + let mul_log = + tree_reader.read_unsigned(&tree_histograms, br, MULTIPLIER_LOG_CONTEXT)?; + if mul_log >= 31 { + return Err(Error::TreeMultiplierTooLarge(mul_log, 31)); + } + let mul_bits = + tree_reader.read_unsigned(&tree_histograms, br, MULTIPLIER_BITS_CONTEXT)?; + let multiplier = (mul_bits as u64 + 1) << mul_log; + if multiplier > (u32::MAX as u64) { + return Err(Error::TreeMultiplierBitsTooLarge(mul_bits, mul_log)); + } + let node = TreeNode::Leaf { + predictor, + offset, + id: leaf_id, + multiplier: multiplier as u32, + }; + leaf_id += 1; + trace!("leaf node {:?}", node); + tree.push(node); + } + } + tree_reader.check_final_state(&tree_histograms)?; + + let num_properties = max_property as usize + 1; + let mut property_ranges = Vec::new_with_capacity(num_properties * tree.len())?; + property_ranges.resize(num_properties * tree.len(), (i32::MIN, i32::MAX)); + let mut height = Vec::new_with_capacity(tree.len())?; + height.resize(tree.len(), 0); + for i in 0..tree.len() { + const HEIGHT_LIMIT: usize = 2048; + if height[i] > HEIGHT_LIMIT { + return Err(Error::TreeTooLarge(height[i], HEIGHT_LIMIT)); + } + if let TreeNode::Split { + property, + val, + left, + right, + } = tree[i] + { + height[left as usize] = height[i] + 1; + height[right as usize] = height[i] + 1; + for p in 0..num_properties { + if p == property as usize { + let (l, u) = property_ranges[i * num_properties + p]; + if l > val || u <= val { + return Err(Error::TreeSplitOnEmptyRange(p as u8, val, l, u)); + } + trace!( + "splitting at node {i} on property {p}, range [{l}, {u}] at position {val}" + ); + property_ranges[left as usize * num_properties + p] = (val + 1, u); + property_ranges[right as usize * num_properties + p] = (l, val); + } else { + property_ranges[left as usize * num_properties + p] = + property_ranges[i * num_properties + p]; + property_ranges[right as usize * num_properties + p] = + property_ranges[i * num_properties + p]; + } + } + } else { + #[cfg(feature = "tracing")] + { + for p in 0..num_properties { + let (l, u) = property_ranges[i * num_properties + p]; + trace!("final range at node {i} property {p}: [{l}, {u}]"); + } + } + } + } + + let histograms = Histograms::decode(tree.len().div_ceil(2), br, true)?; + + Ok(Tree { + nodes: tree, + histograms, + }) + } + + pub fn max_property(&self) -> usize { + self.nodes + .iter() + .map(|x| match x { + TreeNode::Leaf { .. } => 0, + TreeNode::Split { property, .. } => *property, + }) + .max() + .unwrap() as usize + } + + pub fn num_prev_channels(&self) -> usize { + self.max_property().saturating_sub(NUM_NONREF_PROPERTIES) / PROPERTIES_PER_PREVCHAN + } + + // Note: `property_buffer` is passed as input because this implementation relies on having the + // previous values available for computing the local gradient property. + // Also, the first two properties (the static properties) should be already set by the caller. + // All other properties should be 0 on the first call in a row. + #[instrument(level = "trace", skip(buffers), ret)] + #[allow(clippy::too_many_arguments)] + pub(super) fn predict( + &self, + buffers: &mut [&mut ModularChannel], + index: usize, + wp_state: &mut WeightedPredictorState, + x: usize, + y: usize, + references: &Image<i32>, + property_buffer: &mut [i32], + ) -> PredictionResult { + let img = &buffers[index].data; + let prediction_data = PredictionData::get(img.as_rect(), x, y); + let PredictionData { + left, + top, + toptop, + topleft, + topright, + leftleft, + toprightright: _toprightright, + } = prediction_data; + + trace!( + left, + top, topleft, topright, leftleft, toptop, _toprightright + ); + + // Position + property_buffer[2] = y as i32; + property_buffer[3] = x as i32; + + // Neighbours + property_buffer[4] = top.abs(); + property_buffer[5] = left.abs(); + property_buffer[6] = top; + property_buffer[7] = left; + + // Local gradient + property_buffer[8] = left.wrapping_sub(property_buffer[9]); + property_buffer[9] = left.wrapping_add(top).wrapping_sub(topleft); + + // FFV1 context properties + property_buffer[10] = left.wrapping_sub(topleft); + property_buffer[11] = topleft.wrapping_sub(top); + property_buffer[12] = top.wrapping_sub(topright); + property_buffer[13] = top.wrapping_sub(toptop); + property_buffer[14] = left.wrapping_sub(leftleft); + + // Weighted predictor property. + let (wp_pred, property) = + wp_state.predict_and_property((x, y), img.size().0, &prediction_data); + property_buffer[15] = property; + + // Reference properties. + let num_refs = references.size().0; + let ref_properties = &mut property_buffer[NUM_NONREF_PROPERTIES..]; + ref_properties[..num_refs].copy_from_slice(&references.as_rect().row(x)[..num_refs]); + + trace!(?property_buffer, "new properties"); + + let mut tree_node = 0; + while let TreeNode::Split { + property, + val, + left, + right, + } = self.nodes[tree_node] + { + if property_buffer[property as usize] > val { + trace!( + "left at node {tree_node} [{} > {val}]", + property_buffer[property as usize] + ); + tree_node = left as usize; + } else { + trace!( + "right at node {tree_node} [{} <= {val}]", + property_buffer[property as usize] + ); + tree_node = right as usize; + } + } + + trace!(leaf = ?self.nodes[tree_node]); + + let TreeNode::Leaf { + predictor, + offset, + multiplier, + id, + } = self.nodes[tree_node] + else { + unreachable!(); + }; + + let pred = predictor.predict_one(prediction_data, wp_pred); + + PredictionResult { + guess: pred + offset as i64, + multiplier, + context: id, + } + } +} diff --git a/third_party/rust/jxl/src/frame/quant_weights.rs b/third_party/rust/jxl/src/frame/quant_weights.rs @@ -0,0 +1,3140 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +use std::{f32::consts::SQRT_2, sync::OnceLock}; + +use half::f16; + +use crate::{ + BLOCK_DIM, BLOCK_SIZE, + bit_reader::BitReader, + error::{ + Error::{ + HfQuantFactorTooSmall, InvalidDistanceBand, InvalidQuantEncoding, + InvalidQuantEncodingMode, InvalidQuantizationTableWeight, InvalidRawQuantTable, + }, + Result, + }, + frame::{ + LfGlobalState, + modular::{ModularChannel, ModularStreamId, decode::decode_modular_subbitstream}, + transform_map::HfTransformType, + }, + headers::{bit_depth::BitDepth, frame_header::FrameHeader}, +}; + +pub const INV_LF_QUANT: [f32; 3] = [4096.0, 512.0, 256.0]; + +pub const LF_QUANT: [f32; 3] = [ + 1.0 / INV_LF_QUANT[0], + 1.0 / INV_LF_QUANT[1], + 1.0 / INV_LF_QUANT[2], +]; + +const ALMOST_ZERO: f32 = 1e-8; + +const LOG2_NUM_QUANT_MODES: usize = 3; + +#[derive(Debug)] +pub struct DctQuantWeightParams { + params: [[f32; Self::MAX_DISTANCE_BANDS]; 3], + num_bands: usize, +} +impl DctQuantWeightParams { + const LOG2_MAX_DISTANCE_BANDS: usize = 4; + const MAX_DISTANCE_BANDS: usize = 1 + (1 << Self::LOG2_MAX_DISTANCE_BANDS); + + pub fn from_array<const N: usize>(values: &[[f32; N]; 3]) -> Self { + let mut result = Self { + params: [[0.0; Self::MAX_DISTANCE_BANDS]; 3], + num_bands: N, + }; + for (params, values) in result.params.iter_mut().zip(values) { + params[..values.len()].copy_from_slice(values); + } + result + } + + pub fn decode(br: &mut BitReader) -> Result<Self> { + let num_bands = br.read(Self::LOG2_MAX_DISTANCE_BANDS)? as usize + 1; + let mut params = [[0.0; Self::MAX_DISTANCE_BANDS]; 3]; + for row in params.iter_mut() { + for item in row[..num_bands].iter_mut() { + *item = f32::from(f16::from_bits(br.read(16)? as u16)); + } + if row[0] < ALMOST_ZERO { + return Err(HfQuantFactorTooSmall(row[0])); + } + row[0] *= 64.0; + } + Ok(DctQuantWeightParams { params, num_bands }) + } +} + +#[allow(clippy::large_enum_variant)] +#[derive(Debug)] +pub enum QuantEncoding { + Library, + // a.k.a. "Hornuss" + Identity { + xyb_weights: [[f32; 3]; 3], + }, + Dct2 { + xyb_weights: [[f32; 6]; 3], + }, + Dct4 { + params: DctQuantWeightParams, + xyb_mul: [[f32; 2]; 3], + }, + Dct4x8 { + params: DctQuantWeightParams, + xyb_mul: [f32; 3], + }, + Afv { + params4x8: DctQuantWeightParams, + params4x4: DctQuantWeightParams, + weights: [[f32; 9]; 3], + }, + Dct { + params: DctQuantWeightParams, + }, + Raw { + qtable: Vec<i32>, + qtable_den: f32, + }, +} + +impl QuantEncoding { + // TODO(veluca): figure out if this should actually be unused. + #[allow(dead_code)] + pub fn raw_from_qtable(qtable: Vec<i32>, shift: i32) -> Self { + Self::Raw { + qtable, + qtable_den: (1 << shift) as f32 * (1.0 / (8.0 * 255.0)), + } + } + + pub fn decode( + mut required_size_x: usize, + mut required_size_y: usize, + index: usize, + header: &FrameHeader, + lf_global: &LfGlobalState, + br: &mut BitReader, + ) -> Result<Self> { + let required_size = required_size_x * required_size_y; + required_size_x *= BLOCK_DIM; + required_size_y *= BLOCK_DIM; + let mode = br.read(LOG2_NUM_QUANT_MODES)? as u8; + match mode { + 0 => Ok(Self::Library), + 1 => { + if required_size != 1 { + return Err(InvalidQuantEncoding { + mode, + required_size, + }); + } + let mut xyb_weights = [[0.0; 3]; 3]; + for row in xyb_weights.iter_mut() { + for item in row.iter_mut() { + *item = f32::from(f16::from_bits(br.read(16)? as u16)); + if item.abs() < ALMOST_ZERO { + return Err(HfQuantFactorTooSmall(*item)); + } + *item *= 64.0; + } + } + Ok(Self::Identity { xyb_weights }) + } + 2 => { + if required_size != 1 { + return Err(InvalidQuantEncoding { + mode, + required_size, + }); + } + let mut xyb_weights = [[0.0; 6]; 3]; + for row in xyb_weights.iter_mut() { + for item in row.iter_mut() { + *item = f32::from(f16::from_bits(br.read(16)? as u16)); + if item.abs() < ALMOST_ZERO { + return Err(HfQuantFactorTooSmall(*item)); + } + *item *= 64.0; + } + } + Ok(Self::Dct2 { xyb_weights }) + } + 3 => { + if required_size != 1 { + return Err(InvalidQuantEncoding { + mode, + required_size, + }); + } + let mut xyb_mul = [[0.0; 2]; 3]; + for row in xyb_mul.iter_mut() { + for item in row.iter_mut() { + *item = f32::from(f16::from_bits(br.read(16)? as u16)); + if item.abs() < ALMOST_ZERO { + return Err(HfQuantFactorTooSmall(*item)); + } + } + } + let params = DctQuantWeightParams::decode(br)?; + Ok(Self::Dct4 { params, xyb_mul }) + } + 4 => { + if required_size != 1 { + return Err(InvalidQuantEncoding { + mode, + required_size, + }); + } + let mut xyb_mul = [0.0; 3]; + for item in xyb_mul.iter_mut() { + *item = f32::from(f16::from_bits(br.read(16)? as u16)); + if item.abs() < ALMOST_ZERO { + return Err(HfQuantFactorTooSmall(*item)); + } + } + let params = DctQuantWeightParams::decode(br)?; + Ok(Self::Dct4x8 { params, xyb_mul }) + } + 5 => { + if required_size != 1 { + return Err(InvalidQuantEncoding { + mode, + required_size, + }); + } + let mut weights = [[0.0; 9]; 3]; + for row in weights.iter_mut() { + for item in row.iter_mut() { + *item = f32::from(f16::from_bits(br.read(16)? as u16)); + } + for item in row[0..6].iter_mut() { + *item *= 64.0; + } + } + let params4x8 = DctQuantWeightParams::decode(br)?; + let params4x4 = DctQuantWeightParams::decode(br)?; + Ok(Self::Afv { + params4x8, + params4x4, + weights, + }) + } + 6 => { + let params = DctQuantWeightParams::decode(br)?; + Ok(Self::Dct { params }) + } + 7 => { + let qtable_den = f32::from(f16::from_bits(br.read(16)? as u16)); + if qtable_den < ALMOST_ZERO { + // qtable[] values are already checked for <= 0 so the denominator may not be negative. + return Err(InvalidRawQuantTable); + } + let bit_depth = BitDepth::integer_samples(8); + let mut image = [ + ModularChannel::new((required_size_x, required_size_y), bit_depth)?, + ModularChannel::new((required_size_x, required_size_y), bit_depth)?, + ModularChannel::new((required_size_x, required_size_y), bit_depth)?, + ]; + let stream_id = ModularStreamId::QuantTable(index).get_id(header); + decode_modular_subbitstream( + image.iter_mut().collect(), + stream_id, + None, + &lf_global.tree, + br, + )?; + let mut qtable = Vec::with_capacity(required_size_x * required_size_y * 3); + for channel in image.iter_mut() { + for entry in channel.data.as_rect().iter() { + qtable.push(entry); + if entry <= 0 { + return Err(InvalidRawQuantTable); + } + } + } + Ok(Self::Raw { qtable, qtable_den }) + } + _ => Err(InvalidQuantEncoding { + mode, + required_size, + }), + } + } +} + +#[derive(Clone, Copy, Debug)] +enum QuantTable { + // Update QuantTable::VALUES when changing this! + Dct, + Identity, + Dct2x2, + Dct4x4, + Dct16x16, + Dct32x32, + // Dct16x8 + Dct8x16, + // Dct32x8 + Dct8x32, + // Dct32x16 + Dct16x32, + Dct4x8, + // Dct8x4 + Afv0, + // Afv1 + // Afv2 + // Afv3 + Dct64x64, + // Dct64x32, + Dct32x64, + Dct128x128, + // Dct128x64, + Dct64x128, + Dct256x256, + // Dct256x128, + Dct128x256, +} + +impl QuantTable { + pub const CARDINALITY: usize = Self::VALUES.len(); + pub const VALUES: [QuantTable; 17] = [ + QuantTable::Dct, + QuantTable::Identity, + QuantTable::Dct2x2, + QuantTable::Dct4x4, + QuantTable::Dct16x16, + QuantTable::Dct32x32, + // QuantTable::Dct16x8 + QuantTable::Dct8x16, + // QuantTable::Dct32x8 + QuantTable::Dct8x32, + // QuantTable::Dct32x16 + QuantTable::Dct16x32, + QuantTable::Dct4x8, + // QuantTable::Dct8x4 + QuantTable::Afv0, + // QuantTable::Afv1 + // QuantTable::Afv2 + // QuantTable::Afv3 + QuantTable::Dct64x64, + // QuantTable::Dct64x32, + QuantTable::Dct32x64, + QuantTable::Dct128x128, + // QuantTable::Dct128x64, + QuantTable::Dct64x128, + QuantTable::Dct256x256, + // QuantTable::Dct256x128, + QuantTable::Dct128x256, + ]; + pub fn from_usize(idx: usize) -> Result<QuantTable> { + match QuantTable::VALUES.get(idx) { + Some(table) => Ok(*table), + None => Err(InvalidQuantEncodingMode), + } + } + fn for_strategy(strategy: HfTransformType) -> QuantTable { + match strategy { + HfTransformType::DCT => QuantTable::Dct, + HfTransformType::IDENTITY => QuantTable::Identity, + HfTransformType::DCT2X2 => QuantTable::Dct2x2, + HfTransformType::DCT4X4 => QuantTable::Dct4x4, + HfTransformType::DCT16X16 => QuantTable::Dct16x16, + HfTransformType::DCT32X32 => QuantTable::Dct32x32, + HfTransformType::DCT16X8 | HfTransformType::DCT8X16 => QuantTable::Dct8x16, + HfTransformType::DCT32X8 | HfTransformType::DCT8X32 => QuantTable::Dct8x32, + HfTransformType::DCT32X16 | HfTransformType::DCT16X32 => QuantTable::Dct16x32, + HfTransformType::DCT4X8 | HfTransformType::DCT8X4 => QuantTable::Dct4x8, + HfTransformType::AFV0 + | HfTransformType::AFV1 + | HfTransformType::AFV2 + | HfTransformType::AFV3 => QuantTable::Afv0, + HfTransformType::DCT64X64 => QuantTable::Dct64x64, + HfTransformType::DCT64X32 | HfTransformType::DCT32X64 => QuantTable::Dct32x64, + HfTransformType::DCT128X128 => QuantTable::Dct128x128, + HfTransformType::DCT128X64 | HfTransformType::DCT64X128 => QuantTable::Dct64x128, + HfTransformType::DCT256X256 => QuantTable::Dct256x256, + HfTransformType::DCT256X128 | HfTransformType::DCT128X256 => QuantTable::Dct128x256, + } + } +} + +pub struct DequantMatrices { + computed_mask: u32, + table: Vec<f32>, + inv_table: Vec<f32>, + table_offsets: [usize; HfTransformType::CARDINALITY * 3], + encodings: Vec<QuantEncoding>, +} + +#[allow(clippy::excessive_precision)] +impl DequantMatrices { + fn dct() -> QuantEncoding { + QuantEncoding::Dct { + params: DctQuantWeightParams::from_array(&[ + [3150.0, 0.0, -0.4, -0.4, -0.4, -2.0], + [560.0, 0.0, -0.3, -0.3, -0.3, -0.3], + [512.0, -2.0, -1.0, 0.0, -1.0, -2.0], + ]), + } + } + fn id() -> QuantEncoding { + QuantEncoding::Identity { + xyb_weights: [ + [280.0, 3160.0, 3160.0], + [60.0, 864.0, 864.0], + [18.0, 200.0, 200.0], + ], + } + } + fn dct2x2() -> QuantEncoding { + QuantEncoding::Dct2 { + xyb_weights: [ + [3840.0, 2560.0, 1280.0, 640.0, 480.0, 300.0], + [960.0, 640.0, 320.0, 180.0, 140.0, 120.0], + [640.0, 320.0, 128.0, 64.0, 32.0, 16.0], + ], + } + } + fn dct4x4() -> QuantEncoding { + QuantEncoding::Dct4 { + params: DctQuantWeightParams::from_array(&[ + [2200.0, 0.0, 0.0, 0.0], + [392.0, 0.0, 0.0, 0.0], + [112.0, -0.25, -0.25, -0.5], + ]), + xyb_mul: [[1.0, 1.0], [1.0, 1.0], [1.0, 1.0]], + } + } + fn dct16x16() -> QuantEncoding { + QuantEncoding::Dct { + params: DctQuantWeightParams::from_array(&[ + [ + 8996.8725711814115328, + -1.3000777393353804, + -0.49424529824571225, + -0.439093774457103443, + -0.6350101832695744, + -0.90177264050827612, + -1.6162099239887414, + ], + [ + 3191.48366296844234752, + -0.67424582104194355, + -0.80745813428471001, + -0.44925837484843441, + -0.35865440981033403, + -0.31322389111877305, + -0.37615025315725483, + ], + [ + 1157.50408145487200256, + -2.0531423165804414, + -1.4, + -0.50687130033378396, + -0.42708730624733904, + -1.4856834539296244, + -4.9209142884401604, + ], + ]), + } + } + fn dct32x32() -> QuantEncoding { + QuantEncoding::Dct { + params: DctQuantWeightParams::from_array(&[ + [ + 15718.40830982518931456, + -1.025, + -0.98, + -0.9012, + -0.4, + -0.48819395464, + -0.421064, + -0.27, + ], + [ + 7305.7636810695983104, + -0.8041958212306401, + -0.7633036457487539, + -0.55660379990111464, + -0.49785304658857626, + -0.43699592683512467, + -0.40180866526242109, + -0.27321683125358037, + ], + [ + 3803.53173721215041536, + -3.060733579805728, + -2.0413270132490346, + -2.0235650159727417, + -0.5495389509954993, + -0.4, + -0.4, + -0.3, + ], + ]), + } + } + + // dct16x8 + fn dct8x16() -> QuantEncoding { + QuantEncoding::Dct { + params: DctQuantWeightParams::from_array(&[ + [7240.7734393502, -0.7, -0.7, -0.2, -0.2, -0.2, -0.5], + [1448.15468787004, -0.5, -0.5, -0.5, -0.2, -0.2, -0.2], + [506.854140754517, -1.4, -0.2, -0.5, -0.5, -1.5, -3.6], + ]), + } + } + + // dct32x8 + fn dct8x32() -> QuantEncoding { + QuantEncoding::Dct { + params: DctQuantWeightParams::from_array(&[ + [ + 16283.2494710648897, + -1.7812845336559429, + -1.6309059012653515, + -1.0382179034313539, + -0.85, + -0.7, + -0.9, + -1.2360638576849587, + ], + [ + 5089.15750884921511936, + -0.320049391452786891, + -0.35362849922161446, + -0.30340000000000003, + -0.61, + -0.5, + -0.5, + -0.6, + ], + [ + 3397.77603275308720128, + -0.321327362693153371, + -0.34507619223117997, + -0.70340000000000003, + -0.9, + -1.0, + -1.0, + -1.1754605576265209, + ], + ]), + } + } + + // dct32x16 + fn dct16x32() -> QuantEncoding { + QuantEncoding::Dct { + params: DctQuantWeightParams::from_array(&[ + [ + 13844.97076442300573, + -0.97113799999999995, + -0.658, + -0.42026, + -0.22712, + -0.2206, + -0.226, + -0.6, + ], + [ + 4798.964084220744293, + -0.61125308982767057, + -0.83770786552491361, + -0.79014862079498627, + -0.2692727459704829, + -0.38272769465388551, + -0.22924222653091453, + -0.20719098826199578, + ], + [ + 1807.236946760964614, + -1.2, + -1.2, + -0.7, + -0.7, + -0.7, + -0.4, + -0.5, + ], + ]), + } + } + + // dct8x4 + fn dct4x8() -> QuantEncoding { + QuantEncoding::Dct4x8 { + params: DctQuantWeightParams::from_array(&[ + [ + 2198.050556016380522, + -0.96269623020744692, + -0.76194253026666783, + -0.6551140670773547, + ], + [ + 764.3655248643528689, + -0.92630200888366945, + -0.9675229603596517, + -0.27845290869168118, + ], + [ + 527.107573587542228, + -1.4594385811273854, + -1.450082094097871593, + -1.5843722511996204, + ], + ]), + xyb_mul: [1.0, 1.0, 1.0], + } + } + // AFV + fn afv0() -> QuantEncoding { + let QuantEncoding::Dct4x8 { + params: params4x8, .. + } = Self::dct4x8() + else { + unreachable!(); + }; + let QuantEncoding::Dct4 { + params: params4x4, .. + } = Self::dct4x4() + else { + unreachable!() + }; + QuantEncoding::Afv { + params4x8, + params4x4, + weights: [ + [ + 3072.0, 3072.0, // 4x4/4x8 DC tendency. + 256.0, 256.0, 256.0, // AFV corner. + 414.0, 0.0, 0.0, 0.0, // AFV high freqs. + ], + [ + 1024.0, 1024.0, // 4x4/4x8 DC tendency. + 50.0, 50.0, 50.0, // AFV corner. + 58.0, 0.0, 0.0, 0.0, // AFV high freqs. + ], + [ + 384.0, 384.0, // 4x4/4x8 DC tendency. + 12.0, 12.0, 12.0, // AFV corner. + 22.0, -0.25, -0.25, -0.25, // AFV high freqs. + ], + ], + } + } + + fn dct64x64() -> QuantEncoding { + QuantEncoding::Dct { + params: DctQuantWeightParams::from_array(&[ + [ + 0.9 * 26629.073922049845, + -1.025, + -0.78, + -0.65012, + -0.19041574084286472, + -0.20819395464, + -0.421064, + -0.32733845535848671, + ], + [ + 0.9 * 9311.3238710010046, + -0.3041958212306401, + -0.3633036457487539, + -0.35660379990111464, + -0.3443074455424403, + -0.33699592683512467, + -0.30180866526242109, + -0.27321683125358037, + ], + [ + 0.9 * 4992.2486445538634, + -1.2, + -1.2, + -0.8, + -0.7, + -0.7, + -0.4, + -0.5, + ], + ]), + } + } + + // dct64x32 + fn dct32x64() -> QuantEncoding { + QuantEncoding::Dct { + params: DctQuantWeightParams::from_array(&[ + [ + 0.65 * 23629.073922049845, + -1.025, + -0.78, + -0.65012, + -0.19041574084286472, + -0.20819395464, + -0.421064, + -0.32733845535848671, + ], + [ + 0.65 * 8611.3238710010046, + -0.3041958212306401, + -0.3633036457487539, + -0.35660379990111464, + -0.3443074455424403, + -0.33699592683512467, + -0.30180866526242109, + -0.27321683125358037, + ], + [ + 0.65 * 4492.2486445538634, + -1.2, + -1.2, + -0.8, + -0.7, + -0.7, + -0.4, + -0.5, + ], + ]), + } + } + fn dct128x128() -> QuantEncoding { + QuantEncoding::Dct { + params: DctQuantWeightParams::from_array(&[ + [ + 1.8 * 26629.073922049845, + -1.025, + -0.78, + -0.65012, + -0.19041574084286472, + -0.20819395464, + -0.421064, + -0.32733845535848671, + ], + [ + 1.8 * 9311.3238710010046, + -0.3041958212306401, + -0.3633036457487539, + -0.35660379990111464, + -0.3443074455424403, + -0.33699592683512467, + -0.30180866526242109, + -0.27321683125358037, + ], + [ + 1.8 * 4992.2486445538634, + -1.2, + -1.2, + -0.8, + -0.7, + -0.7, + -0.4, + -0.5, + ], + ]), + } + } + + // dct128x64 + fn dct64x128() -> QuantEncoding { + QuantEncoding::Dct { + params: DctQuantWeightParams::from_array(&[ + [ + 1.3 * 23629.073922049845, + -1.025, + -0.78, + -0.65012, + -0.19041574084286472, + -0.20819395464, + -0.421064, + -0.32733845535848671, + ], + [ + 1.3 * 8611.3238710010046, + -0.3041958212306401, + -0.3633036457487539, + -0.35660379990111464, + -0.3443074455424403, + -0.33699592683512467, + -0.30180866526242109, + -0.27321683125358037, + ], + [ + 1.3 * 4492.2486445538634, + -1.2, + -1.2, + -0.8, + -0.7, + -0.7, + -0.4, + -0.5, + ], + ]), + } + } + fn dct256x256() -> QuantEncoding { + QuantEncoding::Dct { + params: DctQuantWeightParams::from_array(&[ + [ + 3.6 * 26629.073922049845, + -1.025, + -0.78, + -0.65012, + -0.19041574084286472, + -0.20819395464, + -0.421064, + -0.32733845535848671, + ], + [ + 3.6 * 9311.3238710010046, + -0.3041958212306401, + -0.3633036457487539, + -0.35660379990111464, + -0.3443074455424403, + -0.33699592683512467, + -0.30180866526242109, + -0.27321683125358037, + ], + [ + 3.6 * 4992.2486445538634, + -1.2, + -1.2, + -0.8, + -0.7, + -0.7, + -0.4, + -0.5, + ], + ]), + } + } + + // dct256x128 + fn dct128x256() -> QuantEncoding { + QuantEncoding::Dct { + params: DctQuantWeightParams::from_array(&[ + [ + 2.6 * 23629.073922049845, + -1.025, + -0.78, + -0.65012, + -0.19041574084286472, + -0.20819395464, + -0.421064, + -0.32733845535848671, + ], + [ + 2.6 * 8611.3238710010046, + -0.3041958212306401, + -0.3633036457487539, + -0.35660379990111464, + -0.3443074455424403, + -0.33699592683512467, + -0.30180866526242109, + -0.27321683125358037, + ], + [ + 2.6 * 4492.2486445538634, + -1.2, + -1.2, + -0.8, + -0.7, + -0.7, + -0.4, + -0.5, + ], + ]), + } + } + + pub fn library() -> &'static [QuantEncoding; QuantTable::CARDINALITY] { + static QUANTS: OnceLock<[QuantEncoding; QuantTable::CARDINALITY]> = OnceLock::new(); + QUANTS.get_or_init(|| { + [ + DequantMatrices::dct(), + DequantMatrices::id(), + DequantMatrices::dct2x2(), + DequantMatrices::dct4x4(), + DequantMatrices::dct16x16(), + DequantMatrices::dct32x32(), + DequantMatrices::dct8x16(), + DequantMatrices::dct8x32(), + DequantMatrices::dct16x32(), + DequantMatrices::dct4x8(), + DequantMatrices::afv0(), + DequantMatrices::dct64x64(), + DequantMatrices::dct32x64(), + // Same default for large transforms (128+) as for 64x* transforms. + DequantMatrices::dct128x128(), + DequantMatrices::dct64x128(), + DequantMatrices::dct256x256(), + DequantMatrices::dct128x256(), + ] + }) + } + + pub fn matrix(&self, quant_kind: HfTransformType, c: usize) -> &[f32] { + assert_ne!((1 << quant_kind as u32) & self.computed_mask, 0); + &self.table[self.table_offsets[quant_kind as usize * 3 + c]..] + } + + // TODO(veluca): figure out if this should actually be unused. + #[allow(dead_code)] + pub fn inv_matrix(&self, quant_kind: HfTransformType, c: usize) -> &[f32] { + assert_ne!((1 << quant_kind as u32) & self.computed_mask, 0); + &self.inv_table[self.table_offsets[quant_kind as usize * 3 + c]..] + } + + pub fn decode( + header: &FrameHeader, + lf_global: &LfGlobalState, + br: &mut BitReader, + ) -> Result<Self> { + let all_default = br.read(1)? == 1; + let mut encodings = Vec::with_capacity(QuantTable::CARDINALITY); + if all_default { + for _ in 0..QuantTable::CARDINALITY { + encodings.push(QuantEncoding::Library) + } + } else { + for (i, (&required_size_x, required_size_y)) in Self::REQUIRED_SIZE_X + .iter() + .zip(Self::REQUIRED_SIZE_Y) + .enumerate() + { + encodings.push(QuantEncoding::decode( + required_size_x, + required_size_y, + i, + header, + lf_global, + br, + )?); + } + } + Ok(Self { + computed_mask: 0, + table: vec![0.0; Self::TOTAL_TABLE_SIZE], + inv_table: vec![0.0; Self::TOTAL_TABLE_SIZE], + table_offsets: [0; HfTransformType::CARDINALITY * 3], + encodings, + }) + } + + pub const REQUIRED_SIZE_X: [usize; QuantTable::CARDINALITY] = + [1, 1, 1, 1, 2, 4, 1, 1, 2, 1, 1, 8, 4, 16, 8, 32, 16]; + + pub const REQUIRED_SIZE_Y: [usize; QuantTable::CARDINALITY] = + [1, 1, 1, 1, 2, 4, 2, 4, 4, 1, 1, 8, 8, 16, 16, 32, 32]; + + pub const SUM_REQUIRED_X_Y: usize = 2056; + + pub const TOTAL_TABLE_SIZE: usize = Self::SUM_REQUIRED_X_Y * BLOCK_SIZE * 3; + + pub fn ensure_computed(&mut self, acs_mask: u32) -> Result<()> { + let mut offsets = [0usize; QuantTable::CARDINALITY * 3]; + let mut pos = 0usize; + for i in 0..QuantTable::CARDINALITY { + let num = DequantMatrices::REQUIRED_SIZE_X[i] + * DequantMatrices::REQUIRED_SIZE_Y[i] + * BLOCK_SIZE; + for c in 0..3 { + offsets[3 * i + c] = pos + c * num; + } + pos += 3 * num; + } + for i in 0..HfTransformType::CARDINALITY { + for c in 0..3 { + self.table_offsets[i * 3 + c] = offsets + [QuantTable::for_strategy(HfTransformType::from_usize(i)?) as usize * 3 + c]; + } + } + let mut kind_mask = 0u32; + for i in 0..HfTransformType::CARDINALITY { + if acs_mask & (1u32 << i) != 0 { + kind_mask |= 1u32 << QuantTable::for_strategy(HfTransformType::VALUES[i]) as u32; + } + } + let mut computed_kind_mask = 0u32; + for i in 0..HfTransformType::CARDINALITY { + if self.computed_mask & (1u32 << i) != 0 { + computed_kind_mask |= + 1u32 << QuantTable::for_strategy(HfTransformType::VALUES[i]) as u32; + } + } + for table in 0..QuantTable::CARDINALITY { + if (1u32 << table) & computed_kind_mask != 0 { + continue; + } + if (1u32 << table) & !kind_mask != 0 { + continue; + } + match self.encodings[table] { + QuantEncoding::Library => { + self.compute_quant_table(true, table, offsets[table * 3])? + } + _ => self.compute_quant_table(false, table, offsets[table * 3])?, + }; + } + self.computed_mask |= acs_mask; + Ok(()) + } + fn compute_quant_table( + &mut self, + library: bool, + table_num: usize, + offset: usize, + ) -> Result<usize> { + let encoding = if library { + &DequantMatrices::library()[table_num] + } else { + &self.encodings[table_num] + }; + let quant_table_idx = QuantTable::from_usize(table_num)? as usize; + let wrows = 8 * DequantMatrices::REQUIRED_SIZE_X[quant_table_idx]; + let wcols = 8 * DequantMatrices::REQUIRED_SIZE_Y[quant_table_idx]; + let num = wrows * wcols; + let mut weights = vec![0f32; 3 * num]; + match encoding { + QuantEncoding::Library => { + // Library and copy quant encoding should get replaced by the actual + // parameters by the caller. + return Err(InvalidQuantEncodingMode); + } + QuantEncoding::Identity { xyb_weights } => { + for c in 0..3 { + for i in 0..64 { + weights[64 * c + i] = xyb_weights[c][0]; + } + weights[64 * c + 1] = xyb_weights[c][1]; + weights[64 * c + 8] = xyb_weights[c][1]; + weights[64 * c + 9] = xyb_weights[c][2]; + } + } + QuantEncoding::Dct2 { xyb_weights } => { + for (c, xyb_weight) in xyb_weights.iter().enumerate() { + let start = c * 64; + weights[start] = 0xBAD as f32; + weights[start + 1] = xyb_weight[0]; + weights[start + 8] = xyb_weight[0]; + weights[start + 9] = xyb_weight[1]; + for y in 0..2 { + for x in 0..2 { + weights[start + y * 8 + x + 2] = xyb_weight[2]; + weights[start + (y + 2) * 8 + x] = xyb_weight[2]; + } + } + for y in 0..2 { + for x in 0..2 { + weights[start + (y + 2) * 8 + x + 2] = xyb_weight[3]; + } + } + for y in 0..4 { + for x in 0..4 { + weights[start + y * 8 + x + 4] = xyb_weight[4]; + weights[start + (y + 4) * 8 + x] = xyb_weight[4]; + } + } + for y in 0..4 { + for x in 0..4 { + weights[start + (y + 4) * 8 + x + 4] = xyb_weight[5]; + } + } + } + } + QuantEncoding::Dct4 { params, xyb_mul } => { + let mut weights4x4 = [0f32; 3 * 4 * 4]; + get_quant_weights(4, 4, params, &mut weights4x4)?; + for c in 0..3 { + for y in 0..BLOCK_DIM { + for x in 0..BLOCK_DIM { + weights[c * num + y * BLOCK_DIM + x] = + weights4x4[c * 16 + (y / 2) * 4 + (x / 2)]; + } + } + weights[c * num + 1] /= xyb_mul[c][0]; + weights[c * num + BLOCK_DIM] /= xyb_mul[c][0]; + weights[c * num + BLOCK_DIM + 1] /= xyb_mul[c][1]; + } + } + QuantEncoding::Dct4x8 { params, xyb_mul } => { + let mut weights4x8 = [0f32; 3 * 4 * 8]; + get_quant_weights(4, 8, params, &mut weights4x8)?; + for c in 0..3 { + for y in 0..BLOCK_DIM { + for x in 0..BLOCK_DIM { + weights[c * num + y * BLOCK_DIM + x] = + weights4x8[c * 32 + (y / 2) * 8 + x]; + } + } + weights[c * num + BLOCK_DIM] /= xyb_mul[c]; + } + } + QuantEncoding::Dct { params } => { + get_quant_weights(wrows, wcols, params, &mut weights)?; + } + QuantEncoding::Raw { qtable, qtable_den } => { + if qtable.len() != 3 * num { + return Err(InvalidRawQuantTable); + } + for i in 0..3 * num { + weights[i] = 1f32 / (qtable_den * qtable[i] as f32); + } + } + QuantEncoding::Afv { + params4x8, + params4x4, + weights: afv_weights, + } => { + const FREQS: [f32; 16] = [ + 0xBAD as f32, + 0xBAD as f32, + 0.8517778890324296, + 5.37778436506804, + 0xBAD as f32, + 0xBAD as f32, + 4.734747904497923, + 5.449245381693219, + 1.6598270267479331, + 4f32, + 7.275749096817861, + 10.423227632456525, + 2.662932286148962, + 7.630657783650829, + 8.962388608184032, + 12.97166202570235, + ]; + let mut weights4x8 = [0f32; 3 * 4 * 8]; + get_quant_weights(4, 8, params4x8, &mut weights4x8)?; + let mut weights4x4 = [0f32; 3 * 4 * 4]; + get_quant_weights(4, 4, params4x4, &mut weights4x4)?; + const LO: f32 = 0.8517778890324296; + const HI: f32 = 12.97166202570235f32 - LO + 1e-6f32; + for c in 0..3 { + let mut bands = [0f32; 4]; + bands[0] = afv_weights[c][5]; + if bands[0] < ALMOST_ZERO { + return Err(InvalidDistanceBand(0, bands[0])); + } + for i in 1..4 { + bands[i] = bands[i - 1] * mult(afv_weights[c][i + 5]); + if bands[i] < ALMOST_ZERO { + return Err(InvalidDistanceBand(i, bands[i])); + } + } + + { + let start = c * 64; + weights[start] = 1f32; + let mut set = |x, y, val| { + weights[start + y * 8 + x] = val; + }; + set(0, 1, afv_weights[c][0]); + set(1, 0, afv_weights[c][1]); + set(0, 2, afv_weights[c][2]); + set(2, 0, afv_weights[c][3]); + set(2, 2, afv_weights[c][4]); + + for y in 0..4 { + for x in 0..4 { + if x < 2 && y < 2 { + continue; + } + let val = interpolate(FREQS[y * 4 + x] - LO, HI, &bands); + set(2 * x, 2 * y, val); + } + } + } + + for y in 0..BLOCK_DIM / 2 { + for x in 0..BLOCK_DIM { + if x == 0 && y == 0 { + continue; + } + weights[c * num + (2 * y + 1) * BLOCK_DIM + x] = + weights4x8[c * 32 + y * 8 + x]; + } + } + + for y in 0..BLOCK_DIM / 2 { + for x in 0..BLOCK_DIM / 2 { + if x == 0 && y == 0 { + continue; + } + weights[c * num + (2 * y) * BLOCK_DIM + 2 * x + 1] = + weights4x4[c * 16 + y * 4 + x]; + } + } + } + } + } + for (i, weight) in weights.iter().enumerate() { + if !(ALMOST_ZERO..=1.0 / ALMOST_ZERO).contains(weight) { + return Err(InvalidQuantizationTableWeight(*weight)); + } + self.table[offset + i] = 1f32 / weight; + self.inv_table[offset + i] = *weight; + } + let (xs, ys) = coefficient_layout( + DequantMatrices::REQUIRED_SIZE_X[quant_table_idx], + DequantMatrices::REQUIRED_SIZE_Y[quant_table_idx], + ); + for c in 0..3 { + for y in 0..ys { + for x in 0..xs { + self.inv_table[offset + c * ys * xs * BLOCK_SIZE + y * BLOCK_DIM * xs + x] = + 0f32; + } + } + } + Ok(0) + } +} + +fn coefficient_layout(rows: usize, cols: usize) -> (usize, usize) { + ( + if rows < cols { rows } else { cols }, + if rows < cols { cols } else { rows }, + ) +} + +fn get_quant_weights( + rows: usize, + cols: usize, + distance_bands: &DctQuantWeightParams, + out: &mut [f32], +) -> Result<()> { + for c in 0..3 { + let mut bands = [0f32; DctQuantWeightParams::MAX_DISTANCE_BANDS]; + bands[0] = distance_bands.params[c][0]; + if bands[0] < ALMOST_ZERO { + return Err(InvalidDistanceBand(0, bands[0])); + } + for i in 1..distance_bands.num_bands { + bands[i] = bands[i - 1] * mult(distance_bands.params[c][i]); + if bands[i] < ALMOST_ZERO { + return Err(InvalidDistanceBand(i, bands[i])); + } + } + let scale = (distance_bands.num_bands - 1) as f32 / (SQRT_2 + 1e-6); + let rcpcol = scale / (cols - 1) as f32; + let rcprow = scale / (rows - 1) as f32; + for y in 0..rows { + let dy = y as f32 * rcprow; + let dy2 = dy * dy; + for x in 0..cols { + let dx = x as f32 * rcpcol; + let scaled_distance = (dx * dx + dy2).sqrt(); + let weight = if distance_bands.num_bands == 1 { + bands[0] + } else { + interpolate_vec(scaled_distance, &bands) + }; + out[c * cols * rows + y * cols + x] = weight; + } + } + } + Ok(()) +} + +fn interpolate_vec(scaled_pos: f32, array: &[f32]) -> f32 { + let idxf32 = scaled_pos.floor(); + let frac = scaled_pos - idxf32; + let idx = idxf32 as usize; + let a = array[idx]; + let b = array[1..][idx]; + (b / a).powf(frac) * a +} + +fn interpolate(pos: f32, max: f32, array: &[f32]) -> f32 { + let scaled_pos = pos * (array.len() - 1) as f32 / max; + let idx = scaled_pos as usize; + let a = array[idx]; + let b = array[idx + 1]; + a * (b / a).powf(scaled_pos - idx as f32) +} + +fn mult(v: f32) -> f32 { + if v > 0f32 { + 1f32 + v + } else { + 1f32 / (1f32 - v) + } +} + +#[cfg(test)] +mod test { + use super::*; + use crate::error::Result; + use crate::frame::quant_weights::DequantMatrices; + use crate::util::test::{assert_all_almost_abs_eq, assert_almost_abs_eq, assert_almost_eq}; + + #[test] + fn check_required_x_y() { + assert_eq!( + DequantMatrices::SUM_REQUIRED_X_Y, + DequantMatrices::REQUIRED_SIZE_X + .iter() + .zip(DequantMatrices::REQUIRED_SIZE_Y) + .map(|(&x, y)| x * y) + .sum() + ); + } + + #[test] + fn check_dequant_matrix_correctness() -> Result<()> { + let mut matrices = DequantMatrices { + computed_mask: 0, + table: vec![0.0; DequantMatrices::TOTAL_TABLE_SIZE], + inv_table: vec![0.0; DequantMatrices::TOTAL_TABLE_SIZE], + table_offsets: [0; HfTransformType::CARDINALITY * 3], + encodings: (0..QuantTable::CARDINALITY) + .map(|_| QuantEncoding::Library) + .collect(), + }; + matrices.ensure_computed(!0)?; + + // Golden data produced by libjxl. + let target_offsets: [usize; 81] = [ + 0, 64, 128, 192, 256, 320, 384, 448, 512, 576, 640, 704, 768, 1024, 1280, 1536, 2560, + 3584, 4608, 4736, 4864, 4608, 4736, 4864, 4992, 5248, 5504, 4992, 5248, 5504, 5760, + 6272, 6784, 5760, 6272, 6784, 7296, 7360, 7424, 7296, 7360, 7424, 7488, 7552, 7616, + 7488, 7552, 7616, 7488, 7552, 7616, 7488, 7552, 7616, 7680, 11776, 15872, 19968, 22016, + 24064, 19968, 22016, 24064, 26112, 42496, 58880, 75264, 83456, 91648, 75264, 83456, + 91648, 99840, 165376, 230912, 296448, 329216, 361984, 296448, 329216, 361984, + ]; + assert_all_almost_abs_eq(matrices.table_offsets, target_offsets, 0); + + // Golden data produced by libjxl. + let target_table = [ + 0.000317f32, + 0.000629f32, + 0.000457f32, + 0.000367f32, + 0.000378f32, + 0.000709f32, + 0.000593f32, + 0.000566f32, + 0.000629f32, + 0.001192f32, + 0.000943f32, + 0.001786f32, + 0.003042f32, + 0.002372f32, + 0.001998f32, + 0.002044f32, + 0.003341f32, + 0.002907f32, + 0.002804f32, + 0.003042f32, + 0.004229f32, + 0.003998f32, + 0.001953f32, + 0.011969f32, + 0.011719f32, + 0.007886f32, + 0.008374f32, + 0.015337f32, + 0.011719f32, + 0.011719f32, + 0.011969f32, + 0.032080f32, + 0.025368f32, + 0.003571f32, + 0.003571f32, + 0.003571f32, + 0.003571f32, + 0.003571f32, + 0.003571f32, + 0.003571f32, + 0.003571f32, + 0.003571f32, + 0.003571f32, + 0.003571f32, + 0.016667f32, + 0.016667f32, + 0.016667f32, + 0.016667f32, + 0.016667f32, + 0.016667f32, + 0.016667f32, + 0.016667f32, + 0.016667f32, + 0.016667f32, + 0.016667f32, + 0.055556f32, + 0.055556f32, + 0.055556f32, + 0.055556f32, + 0.055556f32, + 0.055556f32, + 0.055556f32, + 0.055556f32, + 0.055556f32, + 0.055556f32, + 0.055556f32, + 0.000335f32, + 0.002083f32, + 0.002083f32, + 0.001563f32, + 0.000781f32, + 0.002083f32, + 0.003333f32, + 0.002083f32, + 0.002083f32, + 0.003333f32, + 0.003333f32, + 0.000335f32, + 0.007143f32, + 0.007143f32, + 0.005556f32, + 0.003125f32, + 0.007143f32, + 0.008333f32, + 0.007143f32, + 0.007143f32, + 0.008333f32, + 0.008333f32, + 0.000335f32, + 0.031250f32, + 0.031250f32, + 0.015625f32, + 0.007812f32, + 0.031250f32, + 0.062500f32, + 0.031250f32, + 0.031250f32, + 0.062500f32, + 0.062500f32, + 0.000455f32, + 0.000455f32, + 0.000455f32, + 0.000455f32, + 0.000455f32, + 0.000455f32, + 0.000455f32, + 0.000455f32, + 0.000455f32, + 0.000455f32, + 0.000455f32, + 0.002551f32, + 0.002551f32, + 0.002551f32, + 0.002551f32, + 0.002551f32, + 0.002551f32, + 0.002551f32, + 0.002551f32, + 0.002551f32, + 0.002551f32, + 0.002551f32, + 0.008929f32, + 0.014654f32, + 0.012241f32, + 0.011161f32, + 0.010455f32, + 0.015352f32, + 0.013951f32, + 0.012706f32, + 0.014654f32, + 0.020926f32, + 0.017433f32, + 0.000111f32, + 0.000469f32, + 0.000258f32, + 0.000640f32, + 0.000388f32, + 0.001007f32, + 0.000566f32, + 0.001880f32, + 0.000946f32, + 0.000886f32, + 0.001880f32, + 0.000313f32, + 0.001168f32, + 0.000531f32, + 0.001511f32, + 0.000962f32, + 0.001959f32, + 0.001399f32, + 0.002531f32, + 0.001908f32, + 0.001850f32, + 0.002531f32, + 0.000864f32, + 0.007969f32, + 0.002684f32, + 0.010653f32, + 0.006434f32, + 0.015981f32, + 0.009743f32, + 0.040354f32, + 0.014631f32, + 0.013468f32, + 0.040354f32, + 0.000064f32, + 0.000135f32, + 0.000279f32, + 0.000521f32, + 0.000760f32, + 0.001145f32, + 0.000502f32, + 0.000647f32, + 0.000911f32, + 0.001286f32, + 0.001685f32, + 0.000137f32, + 0.000257f32, + 0.000464f32, + 0.000739f32, + 0.001126f32, + 0.001645f32, + 0.000706f32, + 0.000959f32, + 0.001327f32, + 0.001839f32, + 0.002404f32, + 0.000263f32, + 0.001155f32, + 0.003800f32, + 0.010779f32, + 0.016740f32, + 0.024003f32, + 0.010258f32, + 0.014299f32, + 0.019509f32, + 0.026824f32, + 0.035546f32, + 0.000138f32, + 0.000515f32, + 0.000425f32, + 0.000333f32, + 0.000362f32, + 0.000559f32, + 0.000507f32, + 0.000500f32, + 0.000538f32, + 0.000686f32, + 0.000666f32, + 0.000691f32, + 0.002504f32, + 0.001785f32, + 0.001353f32, + 0.001443f32, + 0.002721f32, + 0.002469f32, + 0.002432f32, + 0.002617f32, + 0.003340f32, + 0.003241f32, + 0.001973f32, + 0.010000f32, + 0.006529f32, + 0.005339f32, + 0.005497f32, + 0.012033f32, + 0.009689f32, + 0.009374f32, + 0.011033f32, + 0.031220f32, + 0.026814f32, + 0.000138f32, + 0.000515f32, + 0.000425f32, + 0.000333f32, + 0.000362f32, + 0.000559f32, + 0.000507f32, + 0.000500f32, + 0.000538f32, + 0.000686f32, + 0.000666f32, + 0.000691f32, + 0.002504f32, + 0.001785f32, + 0.001353f32, + 0.001443f32, + 0.002721f32, + 0.002469f32, + 0.002432f32, + 0.002617f32, + 0.003340f32, + 0.003241f32, + 0.001973f32, + 0.010000f32, + 0.006529f32, + 0.005339f32, + 0.005497f32, + 0.012033f32, + 0.009689f32, + 0.009374f32, + 0.011033f32, + 0.031220f32, + 0.026814f32, + 0.000061f32, + 0.001686f32, + 0.000890f32, + 0.000539f32, + 0.000524f32, + 0.003058f32, + 0.002221f32, + 0.001956f32, + 0.002130f32, + 0.002809f32, + 0.007926f32, + 0.000196f32, + 0.000734f32, + 0.000453f32, + 0.000376f32, + 0.000372f32, + 0.001148f32, + 0.000906f32, + 0.000822f32, + 0.000877f32, + 0.001084f32, + 0.002058f32, + 0.000294f32, + 0.001684f32, + 0.000872f32, + 0.000599f32, + 0.000587f32, + 0.003612f32, + 0.002411f32, + 0.002042f32, + 0.002282f32, + 0.003276f32, + 0.009683f32, + 0.000061f32, + 0.001686f32, + 0.000890f32, + 0.000539f32, + 0.000524f32, + 0.003058f32, + 0.002221f32, + 0.001956f32, + 0.002130f32, + 0.002809f32, + 0.007926f32, + 0.000196f32, + 0.000734f32, + 0.000453f32, + 0.000376f32, + 0.000372f32, + 0.001148f32, + 0.000906f32, + 0.000822f32, + 0.000877f32, + 0.001084f32, + 0.002058f32, + 0.000294f32, + 0.001684f32, + 0.000872f32, + 0.000599f32, + 0.000587f32, + 0.003612f32, + 0.002411f32, + 0.002042f32, + 0.002282f32, + 0.003276f32, + 0.009683f32, + 0.000072f32, + 0.000339f32, + 0.000172f32, + 0.000429f32, + 0.000308f32, + 0.000552f32, + 0.000422f32, + 0.000388f32, + 0.000557f32, + 0.000496f32, + 0.000935f32, + 0.000208f32, + 0.001118f32, + 0.000422f32, + 0.001498f32, + 0.000958f32, + 0.002133f32, + 0.001464f32, + 0.001310f32, + 0.002154f32, + 0.001903f32, + 0.002817f32, + 0.000553f32, + 0.004679f32, + 0.001639f32, + 0.008626f32, + 0.003998f32, + 0.015372f32, + 0.008305f32, + 0.006659f32, + 0.015623f32, + 0.012761f32, + 0.026404f32, + 0.000072f32, + 0.000339f32, + 0.000172f32, + 0.000429f32, + 0.000308f32, + 0.000552f32, + 0.000422f32, + 0.000388f32, + 0.000557f32, + 0.000496f32, + 0.000935f32, + 0.000208f32, + 0.001118f32, + 0.000422f32, + 0.001498f32, + 0.000958f32, + 0.002133f32, + 0.001464f32, + 0.001310f32, + 0.002154f32, + 0.001903f32, + 0.002817f32, + 0.000553f32, + 0.004679f32, + 0.001639f32, + 0.008626f32, + 0.003998f32, + 0.015372f32, + 0.008305f32, + 0.006659f32, + 0.015623f32, + 0.012761f32, + 0.026404f32, + 0.000455f32, + 0.001419f32, + 0.001007f32, + 0.000853f32, + 0.000733f32, + 0.001530f32, + 0.001456f32, + 0.001211f32, + 0.001672f32, + 0.002347f32, + 0.001967f32, + 0.001308f32, + 0.004385f32, + 0.002909f32, + 0.002409f32, + 0.002080f32, + 0.004796f32, + 0.004518f32, + 0.003629f32, + 0.005108f32, + 0.006026f32, + 0.005529f32, + 0.001897f32, + 0.009714f32, + 0.005643f32, + 0.004386f32, + 0.003585f32, + 0.010940f32, + 0.010108f32, + 0.007561f32, + 0.012828f32, + 0.024294f32, + 0.017414f32, + 0.000455f32, + 0.001419f32, + 0.001007f32, + 0.000853f32, + 0.000733f32, + 0.001530f32, + 0.001456f32, + 0.001211f32, + 0.001672f32, + 0.002347f32, + 0.001967f32, + 0.001308f32, + 0.004385f32, + 0.002909f32, + 0.002409f32, + 0.002080f32, + 0.004796f32, + 0.004518f32, + 0.003629f32, + 0.005108f32, + 0.006026f32, + 0.005529f32, + 0.001897f32, + 0.009714f32, + 0.005643f32, + 0.004386f32, + 0.003585f32, + 0.010940f32, + 0.010108f32, + 0.007561f32, + 0.012828f32, + 0.024294f32, + 0.017414f32, + 1.000000f32, + 0.002415f32, + 0.001007f32, + 0.003906f32, + 0.000733f32, + 0.001530f32, + 0.002415f32, + 0.001211f32, + 0.002415f32, + 0.002415f32, + 0.001967f32, + 1.000000f32, + 0.017241f32, + 0.002909f32, + 0.020000f32, + 0.002080f32, + 0.004796f32, + 0.017241f32, + 0.003629f32, + 0.017241f32, + 0.017241f32, + 0.005529f32, + 1.000000f32, + 0.058364f32, + 0.005643f32, + 0.083333f32, + 0.003585f32, + 0.010940f32, + 0.064815f32, + 0.007561f32, + 0.050237f32, + 0.088778f32, + 0.017414f32, + 1.000000f32, + 0.002415f32, + 0.001007f32, + 0.003906f32, + 0.000733f32, + 0.001530f32, + 0.002415f32, + 0.001211f32, + 0.002415f32, + 0.002415f32, + 0.001967f32, + 1.000000f32, + 0.017241f32, + 0.002909f32, + 0.020000f32, + 0.002080f32, + 0.004796f32, + 0.017241f32, + 0.003629f32, + 0.017241f32, + 0.017241f32, + 0.005529f32, + 1.000000f32, + 0.058364f32, + 0.005643f32, + 0.083333f32, + 0.003585f32, + 0.010940f32, + 0.064815f32, + 0.007561f32, + 0.050237f32, + 0.088778f32, + 0.017414f32, + 1.000000f32, + 0.002415f32, + 0.001007f32, + 0.003906f32, + 0.000733f32, + 0.001530f32, + 0.002415f32, + 0.001211f32, + 0.002415f32, + 0.002415f32, + 0.001967f32, + 1.000000f32, + 0.017241f32, + 0.002909f32, + 0.020000f32, + 0.002080f32, + 0.004796f32, + 0.017241f32, + 0.003629f32, + 0.017241f32, + 0.017241f32, + 0.005529f32, + 1.000000f32, + 0.058364f32, + 0.005643f32, + 0.083333f32, + 0.003585f32, + 0.010940f32, + 0.064815f32, + 0.007561f32, + 0.050237f32, + 0.088778f32, + 0.017414f32, + 1.000000f32, + 0.002415f32, + 0.001007f32, + 0.003906f32, + 0.000733f32, + 0.001530f32, + 0.002415f32, + 0.001211f32, + 0.002415f32, + 0.002415f32, + 0.001967f32, + 1.000000f32, + 0.017241f32, + 0.002909f32, + 0.020000f32, + 0.002080f32, + 0.004796f32, + 0.017241f32, + 0.003629f32, + 0.017241f32, + 0.017241f32, + 0.005529f32, + 1.000000f32, + 0.058364f32, + 0.005643f32, + 0.083333f32, + 0.003585f32, + 0.010940f32, + 0.064815f32, + 0.007561f32, + 0.050237f32, + 0.088778f32, + 0.017414f32, + 0.000042f32, + 0.000152f32, + 0.000298f32, + 0.000128f32, + 0.000268f32, + 0.000407f32, + 0.000268f32, + 0.000364f32, + 0.000299f32, + 0.000380f32, + 0.000623f32, + 0.000119f32, + 0.000213f32, + 0.000391f32, + 0.000195f32, + 0.000328f32, + 0.000571f32, + 0.000329f32, + 0.000525f32, + 0.000393f32, + 0.000542f32, + 0.000803f32, + 0.000223f32, + 0.001090f32, + 0.003367f32, + 0.000867f32, + 0.002454f32, + 0.006359f32, + 0.002462f32, + 0.005715f32, + 0.003396f32, + 0.005943f32, + 0.010539f32, + 0.000065f32, + 0.000136f32, + 0.000249f32, + 0.000399f32, + 0.000481f32, + 0.000616f32, + 0.000394f32, + 0.000449f32, + 0.000528f32, + 0.000700f32, + 0.000944f32, + 0.000179f32, + 0.000237f32, + 0.000329f32, + 0.000453f32, + 0.000619f32, + 0.000836f32, + 0.000444f32, + 0.000554f32, + 0.000714f32, + 0.000920f32, + 0.001172f32, + 0.000342f32, + 0.000788f32, + 0.001774f32, + 0.003270f32, + 0.005731f32, + 0.009499f32, + 0.003143f32, + 0.004679f32, + 0.007422f32, + 0.010736f32, + 0.015538f32, + 0.000065f32, + 0.000136f32, + 0.000249f32, + 0.000399f32, + 0.000481f32, + 0.000616f32, + 0.000394f32, + 0.000449f32, + 0.000528f32, + 0.000700f32, + 0.000944f32, + 0.000179f32, + 0.000237f32, + 0.000329f32, + 0.000453f32, + 0.000619f32, + 0.000836f32, + 0.000444f32, + 0.000554f32, + 0.000714f32, + 0.000920f32, + 0.001172f32, + 0.000342f32, + 0.000788f32, + 0.001774f32, + 0.003270f32, + 0.005731f32, + 0.009499f32, + 0.003143f32, + 0.004679f32, + 0.007422f32, + 0.010736f32, + 0.015538f32, + 0.000021f32, + 0.000148f32, + 0.000127f32, + 0.000094f32, + 0.000083f32, + 0.000212f32, + 0.000175f32, + 0.000163f32, + 0.000159f32, + 0.000164f32, + 0.000329f32, + 0.000060f32, + 0.000194f32, + 0.000149f32, + 0.000122f32, + 0.000113f32, + 0.000294f32, + 0.000251f32, + 0.000224f32, + 0.000217f32, + 0.000228f32, + 0.000420f32, + 0.000111f32, + 0.001651f32, + 0.001032f32, + 0.000701f32, + 0.000605f32, + 0.003305f32, + 0.002650f32, + 0.002162f32, + 0.002031f32, + 0.002222f32, + 0.005691f32, + 0.000033f32, + 0.000120f32, + 0.000234f32, + 0.000104f32, + 0.000213f32, + 0.000334f32, + 0.000214f32, + 0.000303f32, + 0.000236f32, + 0.000315f32, + 0.000521f32, + 0.000089f32, + 0.000161f32, + 0.000297f32, + 0.000148f32, + 0.000254f32, + 0.000444f32, + 0.000255f32, + 0.000412f32, + 0.000299f32, + 0.000424f32, + 0.000638f32, + 0.000171f32, + 0.000850f32, + 0.002654f32, + 0.000698f32, + 0.002002f32, + 0.005130f32, + 0.002014f32, + 0.004672f32, + 0.002695f32, + 0.004847f32, + 0.008953f32, + 0.000033f32, + 0.000120f32, + 0.000234f32, + 0.000104f32, + 0.000213f32, + 0.000334f32, + 0.000214f32, + 0.000303f32, + 0.000236f32, + 0.000315f32, + 0.000521f32, + 0.000089f32, + 0.000161f32, + 0.000297f32, + 0.000148f32, + 0.000254f32, + 0.000444f32, + 0.000255f32, + 0.000412f32, + 0.000299f32, + 0.000424f32, + 0.000638f32, + 0.000171f32, + 0.000850f32, + 0.002654f32, + 0.000698f32, + 0.002002f32, + 0.005130f32, + 0.002014f32, + 0.004672f32, + 0.002695f32, + 0.004847f32, + 0.008953f32, + 0.000010f32, + 0.000062f32, + 0.000026f32, + 0.000077f32, + 0.000055f32, + 0.000106f32, + 0.000076f32, + 0.000069f32, + 0.000108f32, + 0.000087f32, + 0.000165f32, + 0.000030f32, + 0.000072f32, + 0.000044f32, + 0.000103f32, + 0.000067f32, + 0.000147f32, + 0.000101f32, + 0.000086f32, + 0.000149f32, + 0.000124f32, + 0.000211f32, + 0.000056f32, + 0.000487f32, + 0.000166f32, + 0.000920f32, + 0.000424f32, + 0.001655f32, + 0.000897f32, + 0.000664f32, + 0.001683f32, + 0.001291f32, + 0.002862f32, + 0.000016f32, + 0.000115f32, + 0.000099f32, + 0.000073f32, + 0.000065f32, + 0.000164f32, + 0.000136f32, + 0.000127f32, + 0.000124f32, + 0.000128f32, + 0.000256f32, + 0.000045f32, + 0.000144f32, + 0.000111f32, + 0.000091f32, + 0.000084f32, + 0.000219f32, + 0.000187f32, + 0.000168f32, + 0.000162f32, + 0.000171f32, + 0.000314f32, + 0.000086f32, + 0.001260f32, + 0.000790f32, + 0.000537f32, + 0.000465f32, + 0.002528f32, + 0.002026f32, + 0.001657f32, + 0.001560f32, + 0.001709f32, + 0.004355f32, + 0.000016f32, + 0.000115f32, + 0.000099f32, + 0.000073f32, + 0.000065f32, + 0.000164f32, + 0.000136f32, + 0.000127f32, + 0.000124f32, + 0.000128f32, + 0.000256f32, + 0.000045f32, + 0.000144f32, + 0.000111f32, + 0.000091f32, + 0.000084f32, + 0.000219f32, + 0.000187f32, + 0.000168f32, + 0.000162f32, + 0.000171f32, + 0.000314f32, + 0.000086f32, + 0.001260f32, + 0.000790f32, + 0.000537f32, + 0.000465f32, + 0.002528f32, + 0.002026f32, + 0.001657f32, + 0.001560f32, + 0.001709f32, + 0.004355f32, + ]; + let mut target_table_index = 0; + for i in 0..HfTransformType::CARDINALITY { + let qt_idx = QuantTable::for_strategy(HfTransformType::from_usize(i)?) as usize; + let size = DequantMatrices::REQUIRED_SIZE_X[qt_idx] + * DequantMatrices::REQUIRED_SIZE_Y[qt_idx] + * BLOCK_SIZE; + for c in 0..3 { + let start = matrices.table_offsets[3 * i + c]; + for j in (start..start + size).step_by(size / 10) { + assert_almost_abs_eq(matrices.table[j], target_table[target_table_index], 1e-5); + target_table_index += 1; + } + } + } + // Golden data produced by libjxl. + let target_inv_table = [ + 0.000000f32, + 1_590.757_8_f32, + 2_188.414_6_f32, + 2_726.993_f32, + 2_648.627_7_f32, + 1_410.374_f32, + 1_686.279_4_f32, + 1_765.964_1_f32, + 1_590.757_8_f32, + 838.702_15_f32, + 1_060.592_9_f32, + 0.000000f32, + 328.723_8_f32, + 421.547_42_f32, + 500.443_73_f32, + 489.194_1_f32, + 299.277_44_f32, + 344.016_4_f32, + 356.627_56_f32, + 328.723_8_f32, + 236.484_7_f32, + 250.119_8_f32, + 0.000000f32, + 83.550804f32, + 85.333_32_f32, + 126.804_84_f32, + 119.412384f32, + 65.203_81_f32, + 85.333_23_f32, + 85.333244f32, + 83.550804f32, + 31.172384f32, + 39.419487f32, + 0.000000f32, + 280.000000f32, + 280.000000f32, + 280.000000f32, + 280.000000f32, + 280.000000f32, + 280.000000f32, + 280.000000f32, + 280.000000f32, + 280.000000f32, + 280.000000f32, + 0.000000f32, + 60.000000f32, + 60.000000f32, + 60.000000f32, + 60.000000f32, + 60.000000f32, + 60.000000f32, + 60.000000f32, + 60.000000f32, + 60.000000f32, + 60.000000f32, + 0.000000f32, + 18.000000f32, + 18.000000f32, + 18.000000f32, + 18.000000f32, + 18.000000f32, + 18.000000f32, + 18.000000f32, + 18.000000f32, + 18.000000f32, + 18.000000f32, + 0.000000f32, + 480.000000f32, + 480.000000f32, + 640.000000f32, + 1280.000000f32, + 480.000000f32, + 300.000000f32, + 480.000000f32, + 480.000000f32, + 300.000000f32, + 300.000000f32, + 0.000000f32, + 140.000000f32, + 140.000000f32, + 180.000000f32, + 320.000000f32, + 140.000000f32, + 120.000000f32, + 140.000000f32, + 140.000000f32, + 120.000000f32, + 120.000000f32, + 0.000000f32, + 32.000000f32, + 32.000000f32, + 64.000000f32, + 128.000000f32, + 32.000000f32, + 16.000000f32, + 32.000000f32, + 32.000000f32, + 16.000000f32, + 16.000000f32, + 0.000000f32, + 2_199.999_5_f32, + 2_199.998_8_f32, + 2_199.997_f32, + 2_199.997_8_f32, + 2_199.999_3_f32, + 2_199.997_f32, + 2_199.998_3_f32, + 2_199.999_5_f32, + 2_199.997_f32, + 2_199.998_3_f32, + 0.000000f32, + 391.999_94_f32, + 391.999_76_f32, + 391.999_45_f32, + 391.999_63_f32, + 391.999_85_f32, + 391.999_45_f32, + 391.999_7_f32, + 391.999_94_f32, + 391.999_45_f32, + 391.999_7_f32, + 0.000000f32, + 68.239_35_f32, + 81.689606f32, + 89.600_1_f32, + 95.651_69_f32, + 65.137_18_f32, + 71.680_09_f32, + 78.702_8_f32, + 68.239_35_f32, + 47.786777f32, + 57.363_38_f32, + 0.000000f32, + 2_134.025_f32, + 3_880.564_5_f32, + 1_561.426_1_f32, + 2_580.273_2_f32, + 993.463_3_f32, + 1_766.659_2_f32, + 531.866_15_f32, + 1_057.315_1_f32, + 1_129.141_7_f32, + 531.866_15_f32, + 0.000000f32, + 856.371_03_f32, + 1_884.007_1_f32, + 661.633_f32, + 1_039.256_2_f32, + 510.514_92_f32, + 714.580_6_f32, + 395.167_5_f32, + 524.174_9_f32, + 540.578_f32, + 395.167_5_f32, + 0.000000f32, + 125.492_86_f32, + 372.602_72_f32, + 93.867_99_f32, + 155.421_52_f32, + 62.573_67_f32, + 102.638_92_f32, + 24.780632f32, + 68.345_99_f32, + 74.248_6_f32, + 24.780632f32, + 0.000000f32, + 7_394.222_7_f32, + 3_578.031_f32, + 1_919.217_f32, + 1_315.414_7_f32, + 873.481_14_f32, + 1_993.637_6_f32, + 1_544.639_8_f32, + 1_097.804_9_f32, + 777.788_7_f32, + 593.406_25_f32, + 0.000000f32, + 3_889.284_f32, + 2_156.401_9_f32, + 1_353.482_9_f32, + 888.446_17_f32, + 607.867_1_f32, + 1_416.747_f32, + 1_042.852_7_f32, + 753.371_15_f32, + 543.717_3_f32, + 415.901_12_f32, + 0.000000f32, + 865.445_74_f32, + 263.145_42_f32, + 92.775_6_f32, + 59.736_86_f32, + 41.660606f32, + 97.485_29_f32, + 69.935_29_f32, + 51.259308f32, + 37.279_94_f32, + 28.132723f32, + 0.000000f32, + 1_943.121_3_f32, + 2_353.786_9_f32, + 3_003.803_7_f32, + 2_759.090_6_f32, + 1_787.990_7_f32, + 1_970.900_4_f32, + 2_000.403_2_f32, + 1_859.104_f32, + 1_456.708_3_f32, + 1_501.485_4_f32, + 0.000000f32, + 399.333_1_f32, + 560.170_4_f32, + 739.322_3_f32, + 692.840_9_f32, + 367.452_f32, + 405.042_f32, + 411.105_13_f32, + 382.066_6_f32, + 299.369_78_f32, + 308.571_96_f32, + 0.000000f32, + 100.000015f32, + 153.171_57_f32, + 187.309_95_f32, + 181.919_97_f32, + 83.107_42_f32, + 103.207_18_f32, + 106.674_48_f32, + 90.637794f32, + 32.030476f32, + 37.294_44_f32, + 0.000000f32, + 1_943.121_3_f32, + 2_353.786_9_f32, + 3_003.803_7_f32, + 2_759.090_6_f32, + 1_787.990_7_f32, + 1_970.900_4_f32, + 2_000.403_2_f32, + 1_859.104_f32, + 1_456.708_3_f32, + 1_501.485_4_f32, + 0.000000f32, + 399.333_1_f32, + 560.170_4_f32, + 739.322_3_f32, + 692.840_9_f32, + 367.452_f32, + 405.042_f32, + 411.105_13_f32, + 382.066_6_f32, + 299.369_78_f32, + 308.571_96_f32, + 0.000000f32, + 100.000015f32, + 153.171_57_f32, + 187.309_95_f32, + 181.919_97_f32, + 83.107_42_f32, + 103.207_18_f32, + 106.674_48_f32, + 90.637794f32, + 32.030476f32, + 37.294_44_f32, + 0.000000f32, + 593.167_f32, + 1_123.533_1_f32, + 1_855.866_1_f32, + 1_908.905_2_f32, + 326.994_14_f32, + 450.258_58_f32, + 511.279_08_f32, + 469.570_34_f32, + 356.047_f32, + 126.164_13_f32, + 0.000000f32, + 1_362.585_8_f32, + 2_208.564_f32, + 2_662.055_2_f32, + 2_690.115_7_f32, + 871.264_95_f32, + 1_103.732_8_f32, + 1_216.298_7_f32, + 1_139.726_f32, + 922.482_9_f32, + 485.887_88_f32, + 0.000000f32, + 593.843_3_f32, + 1_146.650_3_f32, + 1_669.026_5_f32, + 1_704.578_1_f32, + 276.873_93_f32, + 414.831_3_f32, + 489.748_35_f32, + 438.224_15_f32, + 305.274_23_f32, + 103.268776f32, + 0.000000f32, + 593.167_f32, + 1_123.533_1_f32, + 1_855.866_1_f32, + 1_908.905_2_f32, + 326.994_14_f32, + 450.258_58_f32, + 511.279_08_f32, + 469.570_34_f32, + 356.047_f32, + 126.164_13_f32, + 0.000000f32, + 1_362.585_8_f32, + 2_208.564_f32, + 2_662.055_2_f32, + 2_690.115_7_f32, + 871.264_95_f32, + 1_103.732_8_f32, + 1_216.298_7_f32, + 1_139.726_f32, + 922.482_9_f32, + 485.887_88_f32, + 0.000000f32, + 593.843_3_f32, + 1_146.650_3_f32, + 1_669.026_5_f32, + 1_704.578_1_f32, + 276.873_93_f32, + 414.831_3_f32, + 489.748_35_f32, + 438.224_15_f32, + 305.274_23_f32, + 103.268776f32, + 0.000000f32, + 2_951.45_f32, + 5_803.090_3_f32, + 2_333.720_7_f32, + 3_250.279_3_f32, + 1_812.437_7_f32, + 2_367.215_3_f32, + 2_575.901_6_f32, + 1_794.716_f32, + 2_014.430_4_f32, + 1_070.065_2_f32, + 0.000000f32, + 894.280_7_f32, + 2_366.964_4_f32, + 667.591_43_f32, + 1_044.054_9_f32, + 468.918_46_f32, + 683.237_6_f32, + 763.157_1_f32, + 464.274_26_f32, + 525.577_7_f32, + 355.034_94_f32, + 0.000000f32, + 213.711_44_f32, + 609.947_6_f32, + 115.928_83_f32, + 250.111_22_f32, + 65.055_44_f32, + 120.410_9_f32, + 150.171_69_f32, + 64.008_35_f32, + 78.361_82_f32, + 37.872444f32, + 0.000000f32, + 2_951.45_f32, + 5_803.090_3_f32, + 2_333.720_7_f32, + 3_250.279_3_f32, + 1_812.437_7_f32, + 2_367.215_3_f32, + 2_575.901_6_f32, + 1_794.716_f32, + 2_014.430_4_f32, + 1_070.065_2_f32, + 0.000000f32, + 894.280_7_f32, + 2_366.964_4_f32, + 667.591_43_f32, + 1_044.054_9_f32, + 468.918_46_f32, + 683.237_6_f32, + 763.157_1_f32, + 464.274_26_f32, + 525.577_7_f32, + 355.034_94_f32, + 0.000000f32, + 213.711_44_f32, + 609.947_6_f32, + 115.928_83_f32, + 250.111_22_f32, + 65.055_44_f32, + 120.410_9_f32, + 150.171_69_f32, + 64.008_35_f32, + 78.361_82_f32, + 37.872444f32, + 0.000000f32, + 704.524_2_f32, + 993.091_86_f32, + 1_173.002_2_f32, + 1_364.454_1_f32, + 653.527_9_f32, + 687.044_7_f32, + 825.446_53_f32, + 597.922_5_f32, + 426.046_05_f32, + 508.395_4_f32, + 0.000000f32, + 228.070_65_f32, + 343.725_7_f32, + 415.080_72_f32, + 480.806_3_f32, + 208.487_34_f32, + 221.326_13_f32, + 275.591_58_f32, + 195.755_52_f32, + 165.941_71_f32, + 180.871_83_f32, + 0.000000f32, + 102.945_56_f32, + 177.209_15_f32, + 227.986_3_f32, + 278.956_88_f32, + 91.407_4_f32, + 98.934_02_f32, + 132.264_72_f32, + 77.957184f32, + 41.162014f32, + 57.426_7_f32, + 0.000000f32, + 704.524_2_f32, + 993.091_86_f32, + 1_173.002_2_f32, + 1_364.454_1_f32, + 653.527_9_f32, + 687.044_7_f32, + 825.446_53_f32, + 597.922_5_f32, + 426.046_05_f32, + 508.395_4_f32, + 0.000000f32, + 228.070_65_f32, + 343.725_7_f32, + 415.080_72_f32, + 480.806_3_f32, + 208.487_34_f32, + 221.326_13_f32, + 275.591_58_f32, + 195.755_52_f32, + 165.941_71_f32, + 180.871_83_f32, + 0.000000f32, + 102.945_56_f32, + 177.209_15_f32, + 227.986_3_f32, + 278.956_88_f32, + 91.407_4_f32, + 98.934_02_f32, + 132.264_72_f32, + 77.957184f32, + 41.162014f32, + 57.426_7_f32, + 0.000000f32, + 413.999_94_f32, + 993.091_86_f32, + 256.000000f32, + 1_364.454_1_f32, + 653.527_9_f32, + 413.999_66_f32, + 825.446_53_f32, + 413.999_73_f32, + 413.999_42_f32, + 508.395_4_f32, + 0.000000f32, + 57.999_99_f32, + 343.725_7_f32, + 50.000000f32, + 480.806_3_f32, + 208.487_34_f32, + 57.999954f32, + 275.591_58_f32, + 57.999_96_f32, + 57.999_92_f32, + 180.871_83_f32, + 0.000000f32, + 17.133793f32, + 177.209_15_f32, + 12.000000f32, + 278.956_88_f32, + 91.407_4_f32, + 15.428568f32, + 132.264_72_f32, + 19.905685f32, + 11.264012f32, + 57.426_7_f32, + 0.000000f32, + 413.999_94_f32, + 993.091_86_f32, + 256.000000f32, + 1_364.454_1_f32, + 653.527_9_f32, + 413.999_66_f32, + 825.446_53_f32, + 413.999_73_f32, + 413.999_42_f32, + 508.395_4_f32, + 0.000000f32, + 57.999_99_f32, + 343.725_7_f32, + 50.000000f32, + 480.806_3_f32, + 208.487_34_f32, + 57.999954f32, + 275.591_58_f32, + 57.999_96_f32, + 57.999_92_f32, + 180.871_83_f32, + 0.000000f32, + 17.133793f32, + 177.209_15_f32, + 12.000000f32, + 278.956_88_f32, + 91.407_4_f32, + 15.428568f32, + 132.264_72_f32, + 19.905685f32, + 11.264012f32, + 57.426_7_f32, + 0.000000f32, + 413.999_94_f32, + 993.091_86_f32, + 256.000000f32, + 1_364.454_1_f32, + 653.527_9_f32, + 413.999_66_f32, + 825.446_53_f32, + 413.999_73_f32, + 413.999_42_f32, + 508.395_4_f32, + 0.000000f32, + 57.999_99_f32, + 343.725_7_f32, + 50.000000f32, + 480.806_3_f32, + 208.487_34_f32, + 57.999954f32, + 275.591_58_f32, + 57.999_96_f32, + 57.999_92_f32, + 180.871_83_f32, + 0.000000f32, + 17.133793f32, + 177.209_15_f32, + 12.000000f32, + 278.956_88_f32, + 91.407_4_f32, + 15.428568f32, + 132.264_72_f32, + 19.905685f32, + 11.264012f32, + 57.426_7_f32, + 0.000000f32, + 413.999_94_f32, + 993.091_86_f32, + 256.000000f32, + 1_364.454_1_f32, + 653.527_9_f32, + 413.999_66_f32, + 825.446_53_f32, + 413.999_73_f32, + 413.999_42_f32, + 508.395_4_f32, + 0.000000f32, + 57.999_99_f32, + 343.725_7_f32, + 50.000000f32, + 480.806_3_f32, + 208.487_34_f32, + 57.999954f32, + 275.591_58_f32, + 57.999_96_f32, + 57.999_92_f32, + 180.871_83_f32, + 0.000000f32, + 17.133793f32, + 177.209_15_f32, + 12.000000f32, + 278.956_88_f32, + 91.407_4_f32, + 15.428568f32, + 132.264_72_f32, + 19.905685f32, + 11.264012f32, + 57.426_7_f32, + 0.000000f32, + 6_582.816_4_f32, + 3_359.388_7_f32, + 7_791.875_5_f32, + 3_729.601_3_f32, + 2_454.834_f32, + 3_725.529_f32, + 2_744.766_4_f32, + 3_349.231_4_f32, + 2_634.738_8_f32, + 1_604.219_2_f32, + 0.000000f32, + 4_684.622_6_f32, + 2_554.650_4_f32, + 5_132.672_f32, + 3_046.984_1_f32, + 1_750.527_1_f32, + 3_041.338_1_f32, + 1_903.525_6_f32, + 2_542.798_3_f32, + 1_845.962_4_f32, + 1_245.448_2_f32, + 0.000000f32, + 917.482_67_f32, + 297.010_62_f32, + 1_153.165_4_f32, + 407.573_73_f32, + 157.246_46_f32, + 406.220_34_f32, + 174.986_19_f32, + 294.497_8_f32, + 168.263_95_f32, + 94.887_52_f32, + 0.000000f32, + 7_337.233_f32, + 4_022.487_3_f32, + 2_505.753_7_f32, + 2_076.849_f32, + 1_622.844_2_f32, + 2_538.46_f32, + 2_227.385_5_f32, + 1_893.945_9_f32, + 1_428.085_7_f32, + 1_059.229_1_f32, + 0.000000f32, + 4_215.989_f32, + 3_039.568_4_f32, + 2_205.074_5_f32, + 1_614.652_6_f32, + 1_196.811_4_f32, + 2_254.153_8_f32, + 1_805.539_f32, + 1_401.511_f32, + 1_087.307_5_f32, + 853.323_6_f32, + 0.000000f32, + 1_268.412_f32, + 563.855_9_f32, + 305.841_95_f32, + 174.499_47_f32, + 105.278_36_f32, + 318.157_75_f32, + 213.698_53_f32, + 134.729_1_f32, + 93.148544f32, + 64.359_23_f32, + 0.000000f32, + 7_337.233_f32, + 4_022.487_3_f32, + 2_505.753_7_f32, + 2_076.849_f32, + 1_622.844_2_f32, + 2_538.46_f32, + 2_227.385_5_f32, + 1_893.945_9_f32, + 1_428.085_7_f32, + 1_059.229_1_f32, + 0.000000f32, + 4_215.989_f32, + 3_039.568_4_f32, + 2_205.074_5_f32, + 1_614.652_6_f32, + 1_196.811_4_f32, + 2_254.153_8_f32, + 1_805.539_f32, + 1_401.511_f32, + 1_087.307_5_f32, + 853.323_6_f32, + 0.000000f32, + 1_268.412_f32, + 563.855_9_f32, + 305.841_95_f32, + 174.499_47_f32, + 105.278_36_f32, + 318.157_75_f32, + 213.698_53_f32, + 134.729_1_f32, + 93.148544f32, + 64.359_23_f32, + 0.000000f32, + 6_766.112_f32, + 7_894.433_6_f32, + 10_627.11_f32, + 12_049.799_f32, + 4_716.169_f32, + 5_715.242_f32, + 6_145.953_6_f32, + 6_284.098_6_f32, + 6_085.547_f32, + 3_040.497_3_f32, + 0.000000f32, + 5_164.680_7_f32, + 6_709.773_f32, + 8_223.506_f32, + 8_877.354_5_f32, + 3_396.971_2_f32, + 3_985.426_5_f32, + 4_455.855_5_f32, + 4_610.580_6_f32, + 4_388.779_f32, + 2_379.244_9_f32, + 0.000000f32, + 605.837_65_f32, + 968.753_9_f32, + 1_427.095_5_f32, + 1_653.824_2_f32, + 302.614_8_f32, + 377.299_22_f32, + 462.613_9_f32, + 492.384_12_f32, + 449.969_45_f32, + 175.714_2_f32, + 0.000000f32, + 8_341.217_f32, + 4_268.698_f32, + 9_660.034_f32, + 4_689.038_f32, + 2_994.779_5_f32, + 4_679.943_f32, + 3_301.692_6_f32, + 4_245.340_3_f32, + 3_177.615_5_f32, + 1_918.587_6_f32, + 0.000000f32, + 6_214.485_4_f32, + 3_367.615_5_f32, + 6_734.941_4_f32, + 3_939.316_2_f32, + 2_253.354_2_f32, + 3_926.354_7_f32, + 2_424.556_6_f32, + 3_339.36_f32, + 2_355.843_5_f32, + 1_568.312_6_f32, + 0.000000f32, + 1_176.600_6_f32, + 376.791_66_f32, + 1_432.170_2_f32, + 499.566_1_f32, + 194.944_96_f32, + 496.622_07_f32, + 214.034_21_f32, + 371.035_58_f32, + 206.326_42_f32, + 111.690674f32, + 0.000000f32, + 8_341.217_f32, + 4_268.698_f32, + 9_660.034_f32, + 4_689.038_f32, + 2_994.779_5_f32, + 4_679.943_f32, + 3_301.692_6_f32, + 4_245.340_3_f32, + 3_177.615_5_f32, + 1_918.587_6_f32, + 0.000000f32, + 6_214.485_4_f32, + 3_367.615_5_f32, + 6_734.941_4_f32, + 3_939.316_2_f32, + 2_253.354_2_f32, + 3_926.354_7_f32, + 2_424.556_6_f32, + 3_339.36_f32, + 2_355.843_5_f32, + 1_568.312_6_f32, + 0.000000f32, + 1_176.600_6_f32, + 376.791_66_f32, + 1_432.170_2_f32, + 499.566_1_f32, + 194.944_96_f32, + 496.622_07_f32, + 214.034_21_f32, + 371.035_58_f32, + 206.326_42_f32, + 111.690674f32, + 0.000000f32, + 16_091.596_f32, + 37_886.605_f32, + 13_018.401_f32, + 18_061.066_f32, + 9_417.376_f32, + 13_138.253_f32, + 14_536.568_f32, + 9_251.921_f32, + 11_539.108_f32, + 6_057.114_3_f32, + 0.000000f32, + 13_859.232_f32, + 22_801.957_f32, + 9_733.233_f32, + 14_894.771_f32, + 6_785.853_f32, + 9_871.179_f32, + 11_663.131_f32, + 6_696.171_4_f32, + 8_087.468_3_f32, + 4_742.546_f32, + 0.000000f32, + 2_052.832_f32, + 6_023.987_f32, + 1_086.971_4_f32, + 2_357.806_6_f32, + 604.310_36_f32, + 1_115.282_2_f32, + 1_506.556_9_f32, + 594.140_56_f32, + 774.890_87_f32, + 349.454_04_f32, + 0.000000f32, + 8_696.038_f32, + 10_137.892_f32, + 13_662.467_f32, + 15_456.473_f32, + 6_081.469_f32, + 7_342.178_f32, + 7_888.129_4_f32, + 8_059.175_3_f32, + 7_800.867_f32, + 3_911.675_f32, + 0.000000f32, + 6_930.828_6_f32, + 8_992.589_f32, + 11_005.801_f32, + 11_864.501_f32, + 4_558.516_6_f32, + 5_342.787_6_f32, + 5_964.875_5_f32, + 6_164.648_f32, + 5_863.846_7_f32, + 3_188.496_8_f32, + 0.000000f32, + 793.949_34_f32, + 1_266.555_7_f32, + 1_861.544_6_f32, + 2_151.571_5_f32, + 395.616_76_f32, + 493.578_8_f32, + 603.603_2_f32, + 641.048_8_f32, + 585.055_24_f32, + 229.617_22_f32, + 0.000000f32, + 8_696.038_f32, + 10_137.892_f32, + 13_662.467_f32, + 15_456.473_f32, + 6_081.469_f32, + 7_342.178_f32, + 7_888.129_4_f32, + 8_059.175_3_f32, + 7_800.867_f32, + 3_911.675_f32, + 0.000000f32, + 6_930.828_6_f32, + 8_992.589_f32, + 11_005.801_f32, + 11_864.501_f32, + 4_558.516_6_f32, + 5_342.787_6_f32, + 5_964.875_5_f32, + 6_164.648_f32, + 5_863.846_7_f32, + 3_188.496_8_f32, + 0.000000f32, + 793.949_34_f32, + 1_266.555_7_f32, + 1_861.544_6_f32, + 2_151.571_5_f32, + 395.616_76_f32, + 493.578_8_f32, + 603.603_2_f32, + 641.048_8_f32, + 585.055_24_f32, + 229.617_22_f32, + ]; + let mut target_inv_table_index = 0; + for i in 0..HfTransformType::CARDINALITY { + let qt_idx = QuantTable::for_strategy(HfTransformType::from_usize(i)?) as usize; + let size = DequantMatrices::REQUIRED_SIZE_X[qt_idx] + * DequantMatrices::REQUIRED_SIZE_Y[qt_idx] + * BLOCK_SIZE; + for c in 0..3 { + let start = matrices.table_offsets[3 * i + c]; + for j in (start..start + size).step_by(size / 10) { + assert_almost_eq( + matrices.inv_table[j], + target_inv_table[target_inv_table_index], + 1f32, + 1e-3, + ); + target_inv_table_index += 1; + } + } + } + Ok(()) + } +} diff --git a/third_party/rust/jxl/src/frame/quantizer.rs b/third_party/rust/jxl/src/frame/quantizer.rs @@ -0,0 +1,76 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +use crate::{ + bit_reader::BitReader, + error::{Error, Result}, + frame::quant_weights, + headers::encodings::{Empty, UnconditionalCoder}, +}; + +pub const NUM_QUANT_TABLES: usize = 17; +pub const GLOBAL_SCALE_DENOM: usize = 1 << 16; + +#[derive(Debug)] +pub struct LfQuantFactors { + pub quant_factors: [f32; 3], + pub inv_quant_factors: [f32; 3], +} + +impl LfQuantFactors { + pub fn new(br: &mut BitReader) -> Result<LfQuantFactors> { + let mut quant_factors = [0.0f32; 3]; + if br.read(1)? == 1 { + quant_factors = quant_weights::LF_QUANT; + } else { + for qf in quant_factors.iter_mut() { + *qf = f32::read_unconditional(&(), br, &Empty {})? / 128.0; + if *qf < 1e-8 { + return Err(Error::LfQuantFactorTooSmall(*qf)); + } + } + } + + let inv_quant_factors = quant_factors.map(f32::recip); + + Ok(LfQuantFactors { + quant_factors, + inv_quant_factors, + }) + } +} + +#[derive(Debug)] +pub struct QuantizerParams { + pub global_scale: u32, + pub quant_lf: u32, +} + +impl QuantizerParams { + pub fn read(br: &mut BitReader) -> Result<QuantizerParams> { + let global_scale = match br.read(2)? { + 0 => br.read(11)? + 1, + 1 => br.read(11)? + 2049, + 2 => br.read(12)? + 4097, + _ => br.read(16)? + 8193, + }; + let quant_lf = match br.read(2)? { + 0 => 16, + 1 => br.read(5)? + 1, + 2 => br.read(8)? + 1, + _ => br.read(16)? + 1, + }; + Ok(QuantizerParams { + global_scale: global_scale as u32, + quant_lf: quant_lf as u32, + }) + } + pub fn inv_global_scale(&self) -> f32 { + GLOBAL_SCALE_DENOM as f32 / self.global_scale as f32 + } + pub fn inv_quant_lf(&self) -> f32 { + self.inv_global_scale() / self.quant_lf as f32 + } +} diff --git a/third_party/rust/jxl/src/frame/transform_map.rs b/third_party/rust/jxl/src/frame/transform_map.rs @@ -0,0 +1,122 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +use crate::BLOCK_DIM; +use crate::error::{Error::InvalidVarDCTTransform, Result}; + +pub const MAX_COEFF_BLOCKS: usize = 32; +pub const MAX_BLOCK_DIM: usize = BLOCK_DIM * MAX_COEFF_BLOCKS; +pub const MAX_COEFF_AREA: usize = MAX_BLOCK_DIM * MAX_BLOCK_DIM; + +#[allow(clippy::upper_case_acronyms)] +#[derive(Copy, Clone, Debug, PartialEq)] +pub enum HfTransformType { + // Update HfTransformType::VALUES when changing this! + // Regular block size DCT + DCT = 0, + // Encode pixels without transforming + // a.k.a "Hornuss" + IDENTITY = 1, + // Use 2-by-2 DCT + DCT2X2 = 2, + // Use 4-by-4 DCT + DCT4X4 = 3, + // Use 16-by-16 DCT + DCT16X16 = 4, + // Use 32-by-32 DCT + DCT32X32 = 5, + // Use 16-by-8 DCT + DCT16X8 = 6, + // Use 8-by-16 DCT + DCT8X16 = 7, + // Use 32-by-8 DCT + DCT32X8 = 8, + // Use 8-by-32 DCT + DCT8X32 = 9, + // Use 32-by-16 DCT + DCT32X16 = 10, + // Use 16-by-32 DCT + DCT16X32 = 11, + // 4x8 and 8x4 DCT + DCT4X8 = 12, + DCT8X4 = 13, + // Corner-DCT. + AFV0 = 14, + AFV1 = 15, + AFV2 = 16, + AFV3 = 17, + // Larger DCTs + DCT64X64 = 18, + DCT64X32 = 19, + DCT32X64 = 20, + // No transforms smaller than 64x64 are allowed below. + DCT128X128 = 21, + DCT128X64 = 22, + DCT64X128 = 23, + DCT256X256 = 24, + DCT256X128 = 25, + DCT128X256 = 26, +} + +impl HfTransformType { + pub const INVALID_TRANSFORM: u8 = Self::CARDINALITY as u8; + pub const CARDINALITY: usize = Self::VALUES.len(); + pub const VALUES: [HfTransformType; 27] = [ + HfTransformType::DCT, + HfTransformType::IDENTITY, + HfTransformType::DCT2X2, + HfTransformType::DCT4X4, + HfTransformType::DCT16X16, + HfTransformType::DCT32X32, + HfTransformType::DCT16X8, + HfTransformType::DCT8X16, + HfTransformType::DCT32X8, + HfTransformType::DCT8X32, + HfTransformType::DCT32X16, + HfTransformType::DCT16X32, + HfTransformType::DCT4X8, + HfTransformType::DCT8X4, + HfTransformType::AFV0, + HfTransformType::AFV1, + HfTransformType::AFV2, + HfTransformType::AFV3, + HfTransformType::DCT64X64, + HfTransformType::DCT64X32, + HfTransformType::DCT32X64, + HfTransformType::DCT128X128, + HfTransformType::DCT128X64, + HfTransformType::DCT64X128, + HfTransformType::DCT256X256, + HfTransformType::DCT256X128, + HfTransformType::DCT128X256, + ]; + pub fn from_usize(idx: usize) -> Result<HfTransformType> { + match HfTransformType::VALUES.get(idx) { + Some(transform) => Ok(*transform), + None => Err(InvalidVarDCTTransform(idx)), + } + } +} + +pub fn covered_blocks_x(transform: HfTransformType) -> u32 { + let lut: [u32; HfTransformType::CARDINALITY] = [ + 1, 1, 1, 1, 2, 4, 1, 2, 1, 4, 2, 4, 1, 1, 1, 1, 1, 1, 8, 4, 8, 16, 8, 16, 32, 16, 32, + ]; + lut[transform as usize] +} + +pub fn covered_blocks_y(transform: HfTransformType) -> u32 { + let lut: [u32; HfTransformType::CARDINALITY] = [ + 1, 1, 1, 1, 2, 4, 2, 1, 4, 1, 4, 2, 1, 1, 1, 1, 1, 1, 8, 8, 4, 16, 16, 8, 32, 32, 16, + ]; + lut[transform as usize] +} + +pub fn block_shape_id(transform: HfTransformType) -> u32 { + let lut: [u32; HfTransformType::CARDINALITY] = [ + 0, 1, 1, 1, 2, 3, 4, 4, 5, 5, 6, 6, 1, 1, 1, 1, 1, 1, 7, 8, 8, 9, 10, 10, 11, 12, 12, + ]; + lut[transform as usize] +} diff --git a/third_party/rust/jxl/src/headers/bit_depth.rs b/third_party/rust/jxl/src/headers/bit_depth.rs @@ -0,0 +1,96 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +use crate::{bit_reader::BitReader, error::Error, headers::encodings::*}; +use jxl_macros::UnconditionalCoder; + +use std::fmt::Debug; + +#[derive(UnconditionalCoder, Clone, Copy, PartialEq, Eq)] +#[validate] +pub struct BitDepth { + #[default(false)] + floating_point_sample: bool, + #[select_coder(floating_point_sample)] + #[coder_true(u2S(32, 16, 24, Bits(6)+1))] + #[coder_false(u2S(8, 10, 12, Bits(6)+1))] + #[default(8)] + bits_per_sample: u32, + #[condition(floating_point_sample)] + #[default(0)] + #[coder(Bits(4)+1)] + exponent_bits_per_sample: u32, +} + +impl Debug for BitDepth { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + if self.floating_point_sample { + match (self.bits_per_sample, self.exponent_bits_per_sample) { + (32, 8) => { + write!(f, "F32") + } + (16, 5) => { + write!(f, "F16") + } + _ => { + write!( + f, + "FloatE{}M{}", + self.exponent_bits_per_sample, + self.bits_per_sample - self.exponent_bits_per_sample - 1 + ) + } + } + } else { + write!(f, "U{}", self.bits_per_sample) + } + } +} + +impl BitDepth { + pub fn integer_samples(bits_per_sample: u32) -> BitDepth { + BitDepth { + floating_point_sample: false, + bits_per_sample, + exponent_bits_per_sample: 0, + } + } + #[cfg(test)] + pub fn f32() -> BitDepth { + BitDepth { + floating_point_sample: true, + bits_per_sample: 32, + exponent_bits_per_sample: 8, + } + } + pub fn bits_per_sample(&self) -> u32 { + self.bits_per_sample + } + pub fn exponent_bits_per_sample(&self) -> u32 { + self.exponent_bits_per_sample + } + pub fn floating_point_sample(&self) -> bool { + self.floating_point_sample + } + fn check(&self, _: &Empty) -> Result<(), Error> { + if self.floating_point_sample { + if self.exponent_bits_per_sample < 2 || self.exponent_bits_per_sample > 8 { + Err(Error::InvalidExponent(self.exponent_bits_per_sample)) + } else { + let mantissa_bits = + self.bits_per_sample as i32 - self.exponent_bits_per_sample as i32 - 1; + if !(2..=23).contains(&mantissa_bits) { + Err(Error::InvalidMantissa(mantissa_bits)) + } else { + Ok(()) + } + } + } else if self.bits_per_sample > 31 { + Err(Error::InvalidBitsPerSample(self.bits_per_sample)) + } else { + Ok(()) + } + } +} diff --git a/third_party/rust/jxl/src/headers/color_encoding.rs b/third_party/rust/jxl/src/headers/color_encoding.rs @@ -0,0 +1,204 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +use crate::{bit_reader::BitReader, error::Error, headers::encodings::*}; +use jxl_macros::UnconditionalCoder; +use num_derive::FromPrimitive; +use std::fmt; + +#[allow(clippy::upper_case_acronyms)] +#[derive(UnconditionalCoder, Copy, Clone, PartialEq, Debug, FromPrimitive)] +pub enum ColorSpace { + RGB, + Gray, + XYB, + Unknown, +} + +impl fmt::Display for ColorSpace { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!( + f, + "{}", + match self { + ColorSpace::RGB => "RGB", + ColorSpace::Gray => "Gra", + ColorSpace::XYB => "XYB", + ColorSpace::Unknown => "CS?", + } + ) + } +} + +#[allow(clippy::upper_case_acronyms)] +#[derive(UnconditionalCoder, Copy, Clone, PartialEq, Debug, FromPrimitive)] +pub enum WhitePoint { + D65 = 1, + Custom = 2, + E = 10, + DCI = 11, +} + +#[allow(clippy::upper_case_acronyms)] +#[derive(UnconditionalCoder, Copy, Clone, PartialEq, Debug, FromPrimitive)] +pub enum Primaries { + SRGB = 1, + Custom = 2, + BT2100 = 9, + P3 = 11, +} + +#[allow(clippy::upper_case_acronyms)] +#[derive(UnconditionalCoder, Copy, Clone, PartialEq, Debug, FromPrimitive)] +pub enum TransferFunction { + BT709 = 1, + Unknown = 2, + Linear = 8, + SRGB = 13, + PQ = 16, + DCI = 17, + HLG = 18, +} + +#[derive(UnconditionalCoder, Copy, Clone, PartialEq, Debug, FromPrimitive)] +pub enum RenderingIntent { + Perceptual = 0, + Relative, + Saturation, + Absolute, +} + +impl fmt::Display for RenderingIntent { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!( + f, + "{}", + match self { + RenderingIntent::Perceptual => "Per", + RenderingIntent::Relative => "Rel", + RenderingIntent::Saturation => "Sat", + RenderingIntent::Absolute => "Abs", + } + ) + } +} + +#[derive(UnconditionalCoder, Debug, Clone)] +pub struct CustomXY { + #[default(0)] + #[coder(u2S(Bits(19), Bits(19) + 524288, Bits(20) + 1048576, Bits(21) + 2097152))] + pub x: i32, + #[default(0)] + #[coder(u2S(Bits(19), Bits(19) + 524288, Bits(20) + 1048576, Bits(21) + 2097152))] + pub y: i32, +} + +impl CustomXY { + /// Converts the stored scaled integer coordinates to f32 (x, y) values. + pub fn as_f32_coords(&self) -> (f32, f32) { + (self.x as f32 / 1_000_000.0, self.y as f32 / 1_000_000.0) + } + + pub fn from_f32_coords(x: f32, y: f32) -> Self { + Self { + x: (x * 1_000_000.0).round() as i32, + y: (y * 1_000_000.0).round() as i32, + } + } +} + +pub struct CustomTransferFunctionNonserialized { + color_space: ColorSpace, +} + +#[derive(UnconditionalCoder, Debug, Clone)] +#[nonserialized(CustomTransferFunctionNonserialized)] +#[validate] +pub struct CustomTransferFunction { + #[condition(nonserialized.color_space != ColorSpace::XYB)] + #[default(false)] + pub have_gamma: bool, + #[condition(have_gamma)] + #[default(3333333)] // XYB gamma + #[coder(Bits(24))] + pub gamma: u32, + #[condition(!have_gamma && nonserialized.color_space != ColorSpace::XYB)] + #[default(TransferFunction::SRGB)] + pub transfer_function: TransferFunction, +} + +impl CustomTransferFunction { + #[cfg(test)] + pub fn empty() -> CustomTransferFunction { + CustomTransferFunction { + have_gamma: false, + gamma: 0, + transfer_function: TransferFunction::Unknown, + } + } + pub fn gamma(&self) -> f32 { + assert!(self.have_gamma); + self.gamma as f32 * 0.0000001 + } + + pub fn check(&self, _: &CustomTransferFunctionNonserialized) -> Result<(), Error> { + if self.have_gamma { + let gamma = self.gamma(); + if gamma > 1.0 || gamma * 8192.0 < 1.0 { + Err(Error::InvalidGamma(gamma)) + } else { + Ok(()) + } + } else { + Ok(()) + } + } +} + +#[derive(UnconditionalCoder, Debug, Clone)] +#[validate] +pub struct ColorEncoding { + // all_default is never read. + #[allow(dead_code)] + #[all_default] + all_default: bool, + #[default(false)] + pub want_icc: bool, + #[default(ColorSpace::RGB)] + pub color_space: ColorSpace, + #[condition(!want_icc && color_space != ColorSpace::XYB)] + #[default(WhitePoint::D65)] + pub white_point: WhitePoint, + // TODO(veluca): can this be merged in the enum?! + #[condition(white_point == WhitePoint::Custom)] + #[default(CustomXY::default(&field_nonserialized))] + pub white: CustomXY, + #[condition(!want_icc && color_space != ColorSpace::XYB && color_space != ColorSpace::Gray)] + #[default(Primaries::SRGB)] + pub primaries: Primaries, + #[condition(primaries == Primaries::Custom)] + #[default([CustomXY::default(&field_nonserialized), CustomXY::default(&field_nonserialized), CustomXY::default(&field_nonserialized)])] + pub custom_primaries: [CustomXY; 3], + #[condition(!want_icc)] + #[default(CustomTransferFunction::default(&field_nonserialized))] + #[nonserialized(color_space: color_space)] + pub tf: CustomTransferFunction, + #[condition(!want_icc)] + #[default(RenderingIntent::Relative)] + pub rendering_intent: RenderingIntent, +} + +impl ColorEncoding { + pub fn check(&self, _: &Empty) -> Result<(), Error> { + if !self.want_icc + && (self.color_space == ColorSpace::Unknown + || self.tf.transfer_function == TransferFunction::Unknown) + { + Err(Error::InvalidColorEncoding) + } else { + Ok(()) + } + } +} diff --git a/third_party/rust/jxl/src/headers/encodings.rs b/third_party/rust/jxl/src/headers/encodings.rs @@ -0,0 +1,423 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +use super::{frame_header::PermutationNonserialized, permutation::Permutation}; +use crate::{ + bit_reader::BitReader, + entropy_coding::decode::{Histograms, SymbolReader, unpack_signed}, + error::Error, +}; + +pub enum U32 { + Bits(usize), + BitsOffset { n: usize, off: u32 }, + Val(u32), +} + +impl U32 { + pub fn read(&self, br: &mut BitReader) -> Result<u32, Error> { + match *self { + U32::Bits(n) => Ok(br.read(n)? as u32), + U32::BitsOffset { n, off } => Ok(br.read(n)? as u32 + off), + U32::Val(val) => Ok(val), + } + } +} + +pub enum U32Coder { + Direct(U32), + Select(U32, U32, U32, U32), +} + +#[derive(Default)] +pub struct Empty {} + +pub trait UnconditionalCoder<Config> +where + Self: Sized, +{ + type Nonserialized; + fn read_unconditional( + config: &Config, + br: &mut BitReader, + nonserialized: &Self::Nonserialized, + ) -> Result<Self, Error>; +} + +impl UnconditionalCoder<()> for bool { + type Nonserialized = Empty; + fn read_unconditional( + _: &(), + br: &mut BitReader, + _: &Self::Nonserialized, + ) -> Result<bool, Error> { + Ok(br.read(1)? != 0) + } +} + +impl UnconditionalCoder<()> for f32 { + type Nonserialized = Empty; + fn read_unconditional( + _: &(), + br: &mut BitReader, + _: &Self::Nonserialized, + ) -> Result<f32, Error> { + use half::f16; + let ret = f16::from_bits(br.read(16)? as u16); + if !ret.is_finite() { + Err(Error::FloatNaNOrInf) + } else { + Ok(ret.to_f32()) + } + } +} + +impl UnconditionalCoder<U32Coder> for u32 { + type Nonserialized = Empty; + fn read_unconditional( + config: &U32Coder, + br: &mut BitReader, + _: &Self::Nonserialized, + ) -> Result<u32, Error> { + match config { + U32Coder::Direct(u) => u.read(br), + U32Coder::Select(u0, u1, u2, u3) => { + let selector = br.read(2)?; + match selector { + 0 => u0.read(br), + 1 => u1.read(br), + 2 => u2.read(br), + _ => u3.read(br), + } + } + } + } +} + +impl UnconditionalCoder<U32Coder> for i32 { + type Nonserialized = Empty; + fn read_unconditional( + config: &U32Coder, + br: &mut BitReader, + nonserialized: &Self::Nonserialized, + ) -> Result<i32, Error> { + let u = u32::read_unconditional(config, br, nonserialized)?; + Ok(unpack_signed(u)) + } +} + +impl UnconditionalCoder<()> for u64 { + type Nonserialized = Empty; + fn read_unconditional( + _: &(), + br: &mut BitReader, + _: &Self::Nonserialized, + ) -> Result<u64, Error> { + match br.read(2)? { + 0 => Ok(0), + 1 => Ok(1 + br.read(4)?), + 2 => Ok(17 + br.read(8)?), + _ => { + let mut result: u64 = br.read(12)?; + let mut shift = 12; + while br.read(1)? == 1 { + if shift >= 60 { + assert_eq!(shift, 60); + return Ok(result | (br.read(4)? << shift)); + } + result |= br.read(8)? << shift; + shift += 8; + } + Ok(result) + } + } + } +} + +impl UnconditionalCoder<()> for String { + type Nonserialized = Empty; + fn read_unconditional( + _: &(), + br: &mut BitReader, + nonserialized: &Self::Nonserialized, + ) -> Result<String, Error> { + let len = u32::read_unconditional( + &U32Coder::Select( + U32::Val(0), + U32::Bits(4), + U32::BitsOffset { n: 5, off: 16 }, + U32::BitsOffset { n: 10, off: 48 }, + ), + br, + nonserialized, + )?; + let mut ret = String::new(); + ret.reserve(len as usize); + for _ in 0..len { + ret.push(br.read(8)? as u8 as char); + } + Ok(ret) + } +} + +impl UnconditionalCoder<()> for Permutation { + type Nonserialized = PermutationNonserialized; + fn read_unconditional( + _: &(), + br: &mut BitReader, + nonserialized: &Self::Nonserialized, + ) -> Result<Permutation, Error> { + // TODO: This is quadratic when incrementally parsing byte by byte, + // we might want to find a better way of reading the permutation. + let ret = if nonserialized.permuted { + let size = nonserialized.num_entries; + let num_contexts = 8; + let histograms = Histograms::decode(num_contexts, br, /*allow_lz77=*/ true)?; + let mut reader = SymbolReader::new(&histograms, br, None)?; + Permutation::decode(size, 0, &histograms, br, &mut reader) + } else { + Ok(Permutation::default()) + }; + br.jump_to_byte_boundary()?; + ret + } +} + +impl<T: UnconditionalCoder<Config>, Config, const N: usize> UnconditionalCoder<Config> for [T; N] { + type Nonserialized = T::Nonserialized; + fn read_unconditional( + config: &Config, + br: &mut BitReader, + nonserialized: &Self::Nonserialized, + ) -> Result<[T; N], Error> { + use array_init::try_array_init; + try_array_init(|_| T::read_unconditional(config, br, nonserialized)) + } +} + +pub struct VectorCoder<T: Sized> { + pub size_coder: U32Coder, + pub value_coder: T, +} + +impl<Config, T: UnconditionalCoder<Config>> UnconditionalCoder<VectorCoder<Config>> for Vec<T> { + type Nonserialized = T::Nonserialized; + fn read_unconditional( + config: &VectorCoder<Config>, + br: &mut BitReader, + nonserialized: &Self::Nonserialized, + ) -> Result<Vec<T>, Error> { + let len = u32::read_unconditional(&config.size_coder, br, &Empty {})?; + let mut ret: Vec<T> = Vec::new(); + ret.reserve_exact(len as usize); + for _ in 0..len { + ret.push(T::read_unconditional( + &config.value_coder, + br, + nonserialized, + )?); + } + Ok(ret) + } +} + +pub struct SelectCoder<T: Sized> { + pub use_true: bool, + pub coder_true: T, + pub coder_false: T, +} + +// Marker trait to avoid conflicting declarations for [T; N]. +pub trait Selectable {} +impl Selectable for u32 {} + +impl<Config, T: UnconditionalCoder<Config> + Selectable> UnconditionalCoder<SelectCoder<Config>> + for T +{ + type Nonserialized = <T as UnconditionalCoder<Config>>::Nonserialized; + fn read_unconditional( + config: &SelectCoder<Config>, + br: &mut BitReader, + nonserialized: &Self::Nonserialized, + ) -> Result<T, Error> { + if config.use_true { + T::read_unconditional(&config.coder_true, br, nonserialized) + } else { + T::read_unconditional(&config.coder_false, br, nonserialized) + } + } +} + +pub trait ConditionalCoder<Config> +where + Self: Sized, +{ + type Nonserialized; + fn read_conditional( + config: &Config, + condition: bool, + br: &mut BitReader, + nonserialized: &Self::Nonserialized, + ) -> Result<Self, Error>; +} + +impl<Config, T: UnconditionalCoder<Config>> ConditionalCoder<Config> for Option<T> { + type Nonserialized = T::Nonserialized; + fn read_conditional( + config: &Config, + condition: bool, + br: &mut BitReader, + nonserialized: &Self::Nonserialized, + ) -> Result<Option<T>, Error> { + if condition { + Ok(Some(T::read_unconditional(config, br, nonserialized)?)) + } else { + Ok(None) + } + } +} + +impl ConditionalCoder<()> for String { + type Nonserialized = Empty; + fn read_conditional( + _: &(), + condition: bool, + br: &mut BitReader, + nonserialized: &Empty, + ) -> Result<String, Error> { + if condition { + String::read_unconditional(&(), br, nonserialized) + } else { + Ok(String::new()) + } + } +} + +impl<Config, T: UnconditionalCoder<Config>> ConditionalCoder<VectorCoder<Config>> for Vec<T> { + type Nonserialized = T::Nonserialized; + fn read_conditional( + config: &VectorCoder<Config>, + condition: bool, + br: &mut BitReader, + nonserialized: &Self::Nonserialized, + ) -> Result<Vec<T>, Error> { + if condition { + Vec::read_unconditional(config, br, nonserialized) + } else { + Ok(Vec::new()) + } + } +} + +pub trait DefaultedElementCoder<Config, T> +where + Self: Sized, +{ + type Nonserialized; + fn read_defaulted_element( + config: &Config, + condition: bool, + default: T, + br: &mut BitReader, + nonserialized: &Self::Nonserialized, + ) -> Result<Self, Error>; +} + +impl<Config, T> DefaultedElementCoder<VectorCoder<Config>, T> for Vec<T> +where + T: UnconditionalCoder<Config> + Clone, +{ + type Nonserialized = T::Nonserialized; + + fn read_defaulted_element( + config: &VectorCoder<Config>, + condition: bool, + default: T, + br: &mut BitReader, + nonserialized: &Self::Nonserialized, + ) -> Result<Self, Error> { + let len = u32::read_unconditional(&config.size_coder, br, &Empty {})?; + if condition { + let mut ret: Vec<T> = Vec::new(); + ret.reserve_exact(len as usize); + for _ in 0..len { + ret.push(T::read_unconditional( + &config.value_coder, + br, + nonserialized, + )?); + } + Ok(ret) + } else { + Ok(vec![default; len as usize]) + } + } +} + +pub trait DefaultedCoder<Config> +where + Self: Sized, +{ + type Nonserialized; + fn read_defaulted( + config: &Config, + condition: bool, + default: Self, + br: &mut BitReader, + nonserialized: &Self::Nonserialized, + ) -> Result<Self, Error>; +} + +impl<Config, T: UnconditionalCoder<Config>> DefaultedCoder<Config> for T { + type Nonserialized = T::Nonserialized; + fn read_defaulted( + config: &Config, + condition: bool, + default: Self, + br: &mut BitReader, + nonserialized: &Self::Nonserialized, + ) -> Result<T, Error> { + if condition { + Ok(T::read_unconditional(config, br, nonserialized)?) + } else { + Ok(default) + } + } +} + +// TODO(veluca93): this will likely need to be implemented differently if +// there are extensions. +#[derive(Debug, PartialEq, Default, Clone)] +pub struct Extensions {} + +impl UnconditionalCoder<()> for Extensions { + type Nonserialized = Empty; + fn read_unconditional( + _: &(), + br: &mut BitReader, + _: &Self::Nonserialized, + ) -> Result<Extensions, Error> { + let selector = u64::read_unconditional(&(), br, &Empty {})?; + let mut total_size: u64 = 0; + for i in 0..64 { + if (selector & (1u64 << i)) != 0 { + let size = u64::read_unconditional(&(), br, &Empty {})?; + let sum = total_size.checked_add(size); + if let Some(s) = sum { + total_size = s; + } else { + return Err(Error::SizeOverflow); + } + } + } + let total_size = usize::try_from(total_size); + if let Ok(ts) = total_size { + br.skip_bits(ts)?; + } else { + return Err(Error::SizeOverflow); + } + Ok(Extensions {}) + } +} diff --git a/third_party/rust/jxl/src/headers/extra_channels.rs b/third_party/rust/jxl/src/headers/extra_channels.rs @@ -0,0 +1,99 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +use crate::{ + bit_reader::BitReader, + error::Error, + headers::{bit_depth::BitDepth, encodings::*}, +}; +use jxl_macros::UnconditionalCoder; +use num_derive::FromPrimitive; + +#[allow(clippy::upper_case_acronyms)] +#[derive(UnconditionalCoder, Copy, Clone, PartialEq, Debug, FromPrimitive, Eq)] +pub enum ExtraChannel { + Alpha, + Depth, + SpotColor, + SelectionMask, + Black, + CFA, + Thermal, + Reserved0, + Reserved1, + Reserved2, + Reserved3, + Reserved4, + Reserved5, + Reserved6, + Reserved7, + Unknown, + Optional, +} + +// TODO(veluca): figure out if these fields should be unused. +#[allow(dead_code)] +#[derive(UnconditionalCoder, Debug, Clone)] +#[validate] +pub struct ExtraChannelInfo { + #[all_default] + all_default: bool, + #[default(ExtraChannel::Alpha)] + pub ec_type: ExtraChannel, + #[default(BitDepth::default(&field_nonserialized))] + bit_depth: BitDepth, + #[coder(u2S(0, 3, 4, Bits(3) + 1))] + #[default(0)] + dim_shift: u32, + name: String, + // TODO(veluca93): if using Option<bool>, this is None when all_default. + #[condition(ec_type == ExtraChannel::Alpha)] + #[default(false)] + alpha_associated: bool, + #[condition(ec_type == ExtraChannel::SpotColor)] + pub spot_color: Option<[f32; 4]>, + #[condition(ec_type == ExtraChannel::CFA)] + #[coder(u2S(1, Bits(2), Bits(4) + 3, Bits(8) + 19))] + cfa_channel: Option<u32>, +} + +impl ExtraChannelInfo { + #[cfg(test)] + #[allow(clippy::too_many_arguments)] + pub fn new( + all_default: bool, + ec_type: ExtraChannel, + bit_depth: BitDepth, + dim_shift: u32, + name: String, + alpha_associated: bool, + spot_color: Option<[f32; 4]>, + cfa_channel: Option<u32>, + ) -> ExtraChannelInfo { + ExtraChannelInfo { + all_default, + ec_type, + bit_depth, + dim_shift, + name, + alpha_associated, + spot_color, + cfa_channel, + } + } + pub fn dim_shift(&self) -> u32 { + self.dim_shift + } + pub fn alpha_associated(&self) -> bool { + self.alpha_associated + } + fn check(&self, _: &Empty) -> Result<(), Error> { + if self.dim_shift > 3 { + Err(Error::DimShiftTooLarge(self.dim_shift)) + } else { + Ok(()) + } + } +} diff --git a/third_party/rust/jxl/src/headers/frame_header.rs b/third_party/rust/jxl/src/headers/frame_header.rs @@ -0,0 +1,805 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +#![allow(clippy::excessive_precision)] + +use crate::{ + BLOCK_DIM, GROUP_DIM, + bit_reader::BitReader, + error::Error, + headers::{encodings::*, extra_channels::ExtraChannelInfo}, + image::Rect, + util::FloorLog2, +}; + +use jxl_macros::UnconditionalCoder; +use num_derive::FromPrimitive; +use std::cmp::min; + +use super::{Animation, permutation::Permutation}; + +#[derive(UnconditionalCoder, Copy, Clone, PartialEq, Debug, FromPrimitive)] +pub enum FrameType { + RegularFrame = 0, + LFFrame = 1, + ReferenceOnly = 2, + SkipProgressive = 3, +} + +#[derive(UnconditionalCoder, Copy, Clone, PartialEq, Debug, FromPrimitive)] +pub enum Encoding { + VarDCT = 0, + Modular = 1, +} + +struct Flags; + +impl Flags { + pub const ENABLE_NOISE: u64 = 1; + pub const ENABLE_PATCHES: u64 = 2; + pub const ENABLE_SPLINES: u64 = 0x10; + pub const USE_LF_FRAME: u64 = 0x20; + pub const SKIP_ADAPTIVE_LF_SMOOTHING: u64 = 0x80; +} + +#[derive(UnconditionalCoder, Debug, PartialEq)] +pub struct Passes { + #[coder(u2S(1, 2, 3, Bits(3) + 4))] + #[default(1)] + pub num_passes: u32, + + #[coder(u2S(0, 1, 2, Bits(1) + 3))] + #[default(0)] + #[condition(num_passes != 1)] + num_ds: u32, + + #[size_coder(explicit(num_passes - 1))] + #[coder(Bits(2))] + #[default_element(0)] + #[condition(num_passes != 1)] + pub shift: Vec<u32>, + + #[size_coder(explicit(num_ds))] + #[coder(u2S(1, 2, 4, 8))] + #[default_element(1)] + #[condition(num_passes != 1)] + downsample: Vec<u32>, + + #[size_coder(explicit(num_ds))] + #[coder(u2S(0, 1, 2, Bits(3)))] + #[default_element(0)] + #[condition(num_passes != 1)] + last_pass: Vec<u32>, +} + +impl Passes { + pub fn downsampling_bracket(&self, pass: usize) -> (usize, usize) { + let mut max_shift = 2; + let mut min_shift = 3; + for i in 0..pass + 1 { + for j in 0..self.num_ds as usize { + min_shift = self.downsample[j].floor_log2(); + } + if i + 1 == self.num_passes as usize { + min_shift = 0 + } + if i != pass { + max_shift = min_shift - 1; + } + } + (min_shift as usize, max_shift as usize) + } +} + +#[derive(UnconditionalCoder, Copy, Clone, PartialEq, Debug, FromPrimitive)] +pub enum BlendingMode { + Replace = 0, + Add = 1, + Blend = 2, + AlphaWeightedAdd = 3, + Mul = 4, +} + +#[derive(Default)] +pub struct BlendingInfoNonserialized { + num_extra_channels: u32, + full_frame: bool, +} + +#[derive(UnconditionalCoder, Debug, PartialEq, Clone)] +#[nonserialized(BlendingInfoNonserialized)] +pub struct BlendingInfo { + #[coder(u2S(0, 1, 2, Bits(2) + 3))] + #[default(BlendingMode::Replace)] + pub mode: BlendingMode, + + /* Spec: "Let multi_extra be true if and only if and the number of extra channels is at least two." + libjxl condition is num_extra_channels > 0 */ + #[coder(u2S(0, 1, 2, Bits(3) + 3))] + #[default(0)] + #[condition(nonserialized.num_extra_channels > 0 && + (mode == BlendingMode::Blend || mode == BlendingMode::AlphaWeightedAdd))] + pub alpha_channel: u32, + + #[default(false)] + #[condition(nonserialized.num_extra_channels > 0 && + (mode == BlendingMode::Blend || mode == BlendingMode::AlphaWeightedAdd || mode == BlendingMode::Mul))] + pub clamp: bool, + + #[coder(u2S(0, 1, 2, 3))] + #[default(0)] + // This condition is called `resets_canvas` in the spec + #[condition(!(nonserialized.full_frame && mode == BlendingMode::Replace))] + pub source: u32, +} + +pub struct RestorationFilterNonserialized { + encoding: Encoding, +} + +#[derive(UnconditionalCoder, Debug, PartialEq)] +#[nonserialized(RestorationFilterNonserialized)] +pub struct RestorationFilter { + #[all_default] + all_default: bool, + + #[default(true)] + pub gab: bool, + + #[default(false)] + #[condition(gab)] + gab_custom: bool, + + #[default(0.115169525)] + #[condition(gab_custom)] + pub gab_x_weight1: f32, + + #[default(0.061248592)] + #[condition(gab_custom)] + pub gab_x_weight2: f32, + + #[default(0.115169525)] + #[condition(gab_custom)] + pub gab_y_weight1: f32, + + #[default(0.061248592)] + #[condition(gab_custom)] + pub gab_y_weight2: f32, + + #[default(0.115169525)] + #[condition(gab_custom)] + pub gab_b_weight1: f32, + + #[default(0.061248592)] + #[condition(gab_custom)] + pub gab_b_weight2: f32, + + #[coder(Bits(2))] + #[default(2)] + pub epf_iters: u32, + + #[default(false)] + #[condition(epf_iters > 0 && nonserialized.encoding == Encoding::VarDCT)] + epf_sharp_custom: bool, + + #[default([0.0, 1.0 / 7.0, 2.0 / 7.0, 3.0 / 7.0, 4.0 / 7.0, 5.0 / 7.0, 6.0 / 7.0, 1.0])] + #[condition(epf_sharp_custom)] + pub epf_sharp_lut: [f32; 8], + + #[default(false)] + #[condition(epf_iters > 0)] + epf_weight_custom: bool, + + #[default([40.0, 5.0, 3.5])] + #[condition(epf_weight_custom)] + pub epf_channel_scale: [f32; 3], + + #[default(0.45)] + #[condition(epf_weight_custom)] + epf_pass1_zeroflush: f32, + + #[default(0.6)] + #[condition(epf_weight_custom)] + epf_pass2_zeroflush: f32, + + #[default(false)] + #[condition(epf_iters > 0)] + epf_sigma_custom: bool, + + #[default(0.46)] + #[condition(epf_sigma_custom && nonserialized.encoding == Encoding::VarDCT)] + pub epf_quant_mul: f32, + + #[default(0.9)] + #[condition(epf_sigma_custom)] + pub epf_pass0_sigma_scale: f32, + + #[default(6.5)] + #[condition(epf_sigma_custom)] + pub epf_pass2_sigma_scale: f32, + + #[default(2.0 / 3.0)] + #[condition(epf_sigma_custom)] + pub epf_border_sad_mul: f32, + + #[default(1.0)] + #[condition(epf_iters > 0 && nonserialized.encoding == Encoding::Modular)] + pub epf_sigma_for_modular: f32, + + #[default(Extensions::default())] + extensions: Extensions, +} + +pub struct TocNonserialized { + pub num_entries: u32, +} + +pub struct PermutationNonserialized { + pub num_entries: u32, + pub permuted: bool, +} + +#[derive(UnconditionalCoder, Debug, PartialEq)] +#[nonserialized(TocNonserialized)] +pub struct Toc { + #[default(false)] + pub permuted: bool, + + // Here we don't use `condition(permuted)`, because `jump_to_byte_boundary` needs to be executed in both cases + #[default(Permutation::default())] + #[nonserialized(num_entries: nonserialized.num_entries, permuted: permuted)] + pub permutation: Permutation, + + #[coder(u2S(Bits(10), Bits(14) + 1024, Bits(22) + 17408, Bits(30) + 4211712))] + #[size_coder(explicit(nonserialized.num_entries))] + pub entries: Vec<u32>, +} + +pub struct FrameHeaderNonserialized { + pub xyb_encoded: bool, + pub num_extra_channels: u32, + pub extra_channel_info: Vec<ExtraChannelInfo>, + pub have_animation: bool, + pub have_timecode: bool, + pub img_width: u32, + pub img_height: u32, +} + +const H_SHIFT: [usize; 4] = [0, 1, 1, 0]; +const V_SHIFT: [usize; 4] = [0, 1, 0, 1]; + +fn compute_jpeg_shift(jpeg_upsampling: &[u32], shift_table: &[usize]) -> u32 { + jpeg_upsampling + .iter() + .map(|&ch| shift_table[ch as usize]) + .max() + .unwrap_or(0) as u32 +} + +#[derive(UnconditionalCoder, Debug, PartialEq)] +#[nonserialized(FrameHeaderNonserialized)] +#[aligned] +#[validate] +pub struct FrameHeader { + #[all_default] + all_default: bool, + + #[coder(Bits(2))] + #[default(FrameType::RegularFrame)] + pub frame_type: FrameType, + + #[coder(Bits(1))] + #[default(Encoding::VarDCT)] + pub encoding: Encoding, + + #[default(0)] + flags: u64, + + #[default(false)] + #[condition(!nonserialized.xyb_encoded)] + pub do_ycbcr: bool, + + #[coder(Bits(2))] + #[default([0, 0, 0])] + #[condition(do_ycbcr && flags & Flags::USE_LF_FRAME == 0)] + jpeg_upsampling: [u32; 3], + + #[coder(u2S(1, 2, 4, 8))] + #[default(1)] + #[condition(flags & Flags::USE_LF_FRAME == 0)] + pub upsampling: u32, + + #[size_coder(explicit(nonserialized.num_extra_channels))] + #[coder(u2S(1, 2, 4, 8))] + #[default_element(1)] + #[condition(flags & Flags::USE_LF_FRAME == 0)] + pub ec_upsampling: Vec<u32>, + + #[coder(Bits(2))] + #[default(1)] + #[condition(encoding == Encoding::Modular)] + group_size_shift: u32, + + #[coder(Bits(3))] + #[default(3)] + #[condition(encoding == Encoding::VarDCT && nonserialized.xyb_encoded)] + pub x_qm_scale: u32, + + #[coder(Bits(3))] + #[default(2)] + #[condition(encoding == Encoding::VarDCT && nonserialized.xyb_encoded)] + pub b_qm_scale: u32, + + #[condition(frame_type != FrameType::ReferenceOnly)] + #[default(Passes::default(&field_nonserialized))] + pub passes: Passes, + + #[coder(u2S(1, 2, 3, 4))] + #[default(0)] + #[condition(frame_type == FrameType::LFFrame)] + pub lf_level: u32, + + #[default(false)] + #[condition(frame_type != FrameType::LFFrame)] + have_crop: bool, + + #[coder(u2S(Bits(8), Bits(11) + 256, Bits(14) + 2304, Bits(30) + 18688))] + #[default(0)] + #[condition(have_crop && frame_type != FrameType::ReferenceOnly)] + pub x0: i32, + + #[coder(u2S(Bits(8), Bits(11) + 256, Bits(14) + 2304, Bits(30) + 18688))] + #[default(0)] + #[condition(have_crop && frame_type != FrameType::ReferenceOnly)] + pub y0: i32, + + #[coder(u2S(Bits(8), Bits(11) + 256, Bits(14) + 2304, Bits(30) + 18688))] + #[default(0)] + #[condition(have_crop)] + frame_width: u32, + + #[coder(u2S(Bits(8), Bits(11) + 256, Bits(14) + 2304, Bits(30) + 18688))] + #[default(0)] + #[condition(have_crop)] + frame_height: u32, + + // The following 2 fields are not actually serialized, but just used as variables to help with + // defining later conditions. + #[default(x0 <= 0 && y0 <= 0 && (frame_width as i64 + x0 as i64) >= nonserialized.img_width as i64 && + (frame_height as i64 + y0 as i64) >= nonserialized.img_height as i64)] + #[condition(false)] + completely_covers: bool, + + #[default(!have_crop || completely_covers)] + #[condition(false)] + full_frame: bool, + + /* "normal_frame" denotes the condition !all_default + && (frame_type == kRegularFrame || frame_type == kSkipProgressive) */ + #[default(BlendingInfo::default(&field_nonserialized))] + #[condition(frame_type == FrameType::RegularFrame || frame_type == FrameType::SkipProgressive)] + #[nonserialized(num_extra_channels : nonserialized.num_extra_channels, full_frame : full_frame)] + pub blending_info: BlendingInfo, + + #[size_coder(explicit(nonserialized.num_extra_channels))] + #[condition(frame_type == FrameType::RegularFrame || frame_type == FrameType::SkipProgressive)] + #[default_element(BlendingInfo::default(&field_nonserialized))] + #[nonserialized(num_extra_channels : nonserialized.num_extra_channels, full_frame: full_frame)] + pub ec_blending_info: Vec<BlendingInfo>, + + #[coder(u2S(0, 1, Bits(8), Bits(32)))] + #[default(0)] + #[condition((frame_type == FrameType::RegularFrame || + frame_type == FrameType::SkipProgressive) && nonserialized.have_animation)] + pub duration: u32, + + #[coder(Bits(32))] + #[default(0)] + #[condition((frame_type == FrameType::RegularFrame || + frame_type == FrameType::SkipProgressive) && nonserialized.have_timecode)] + timecode: u32, + + #[default(frame_type == FrameType::RegularFrame)] + #[condition(frame_type == FrameType::RegularFrame || frame_type == FrameType::SkipProgressive)] + pub is_last: bool, + + #[coder(Bits(2))] + #[default(0)] + #[condition(frame_type != FrameType::LFFrame && !is_last)] + pub save_as_reference: u32, + + // The following 2 fields are not actually serialized, but just used as variables to help with + // defining later conditions. + #[default(!is_last && frame_type != FrameType::LFFrame && (duration == 0 || save_as_reference != 0))] + #[condition(false)] + pub can_be_referenced: bool, + + #[default(can_be_referenced && blending_info.mode == BlendingMode::Replace && full_frame && + (frame_type == FrameType::RegularFrame || frame_type == FrameType::SkipProgressive))] + #[condition(false)] + save_before_ct_def_false: bool, + + #[default(frame_type == FrameType::LFFrame)] + #[condition(frame_type == FrameType::ReferenceOnly || save_before_ct_def_false)] + pub save_before_ct: bool, + + pub name: String, + + #[default(RestorationFilter::default(&field_nonserialized))] + #[nonserialized(encoding : encoding)] + pub restoration_filter: RestorationFilter, + + #[default(Extensions::default())] + extensions: Extensions, + + // The following fields are not actually serialized, but just used as variables for + // implementing the methods below. + #[coder(Bits(0))] + #[default(if frame_width == 0 { nonserialized.img_width } else { frame_width })] + #[condition(false)] + pub width: u32, + + #[coder(Bits(0))] + #[default(if frame_height == 0 { nonserialized.img_height } else { frame_height })] + #[condition(false)] + pub height: u32, + + #[coder(Bits(0))] + #[default(compute_jpeg_shift(&jpeg_upsampling, &H_SHIFT))] + #[condition(false)] + pub maxhs: u32, + + #[coder(Bits(0))] + #[default(compute_jpeg_shift(&jpeg_upsampling, &V_SHIFT))] + #[condition(false)] + pub maxvs: u32, + + #[coder(Bits(0))] + #[default(nonserialized.num_extra_channels)] + #[condition(false)] + pub num_extra_channels: u32, +} + +impl FrameHeader { + pub fn log_group_dim(&self) -> usize { + (GROUP_DIM.ilog2() - 1 + self.group_size_shift) as usize + } + pub fn group_dim(&self) -> usize { + 1 << self.log_group_dim() + } + pub fn lf_group_dim(&self) -> usize { + self.group_dim() * BLOCK_DIM + } + + pub fn num_groups(&self) -> usize { + self.size_groups().0 * self.size_groups().1 + } + + pub fn num_lf_groups(&self) -> usize { + self.size_lf_groups().0 * self.size_lf_groups().1 + } + + pub fn num_toc_entries(&self) -> usize { + let num_groups = self.num_groups(); + let num_dc_groups = self.num_lf_groups(); + + if num_groups == 1 && self.passes.num_passes == 1 { + 1 + } else { + 2 + num_dc_groups + num_groups * self.passes.num_passes as usize + } + } + + pub fn duration(&self, animation: &Animation) -> f64 { + (self.duration as f64) * 1000.0 * (animation.tps_denominator as f64) + / (animation.tps_numerator as f64) + } + + pub fn has_patches(&self) -> bool { + self.flags & Flags::ENABLE_PATCHES != 0 + } + + pub fn has_noise(&self) -> bool { + self.flags & Flags::ENABLE_NOISE != 0 + } + + pub fn has_splines(&self) -> bool { + self.flags & Flags::ENABLE_SPLINES != 0 + } + pub fn has_lf_frame(&self) -> bool { + self.flags & Flags::USE_LF_FRAME != 0 + } + pub fn should_do_adaptive_lf_smoothing(&self) -> bool { + self.flags & Flags::SKIP_ADAPTIVE_LF_SMOOTHING == 0 + && !self.has_lf_frame() + && self.encoding == Encoding::VarDCT + } + pub fn raw_hshift(&self, c: usize) -> usize { + H_SHIFT[self.jpeg_upsampling[c] as usize] + } + pub fn hshift(&self, c: usize) -> usize { + (self.maxhs as usize) - self.raw_hshift(c) + } + pub fn raw_vshift(&self, c: usize) -> usize { + V_SHIFT[self.jpeg_upsampling[c] as usize] + } + pub fn vshift(&self, c: usize) -> usize { + (self.maxvs as usize) - self.raw_vshift(c) + } + pub fn is444(&self) -> bool { + self.hshift(0) == 0 && self.vshift(0) == 0 && // Cb + self.hshift(2) == 0 && self.vshift(2) == 0 && // Cr + self.hshift(1) == 0 && self.vshift(1) == 0 // Y + } + pub fn is420(&self) -> bool { + self.hshift(0) == 1 && self.vshift(0) == 1 && // Cb + self.hshift(2) == 1 && self.vshift(2) == 1 && // Cr + self.hshift(1) == 0 && self.vshift(1) == 0 // Y + } + pub fn is422(&self) -> bool { + self.hshift(0) == 1 && self.vshift(0) == 0 && // Cb + self.hshift(2) == 1 && self.vshift(2) == 0 && // Cr + self.hshift(1) == 0 && self.vshift(1) == 0 // Y + } + pub fn is440(&self) -> bool { + self.hshift(0) == 0 && self.vshift(0) == 1 && // Cb + self.hshift(2) == 0 && self.vshift(2) == 1 && // Cr + self.hshift(1) == 0 && self.vshift(1) == 0 // Y + } + + pub fn is_visible(&self) -> bool { + (self.is_last || self.duration > 0) + && (self.frame_type == FrameType::RegularFrame + || self.frame_type == FrameType::SkipProgressive) + } + + pub fn needs_blending(&self) -> bool { + if !(self.frame_type == FrameType::RegularFrame + || self.frame_type == FrameType::SkipProgressive) + { + return false; + } + let replace_all = self.blending_info.mode == BlendingMode::Replace + && self + .ec_blending_info + .iter() + .all(|x| x.mode == BlendingMode::Replace); + self.have_crop || !replace_all + } + + /// The dimensions of this frame, as coded in the codestream, excluding padding pixels. + pub fn size(&self) -> (usize, usize) { + let (width, height) = self.size_upsampled(); + ( + width.div_ceil(self.upsampling as usize), + height.div_ceil(self.upsampling as usize), + ) + } + + /// The dimensions of this frame, as coded in the codestream, in 8x8 blocks. + pub fn size_blocks(&self) -> (usize, usize) { + ( + self.size().0.div_ceil(BLOCK_DIM << self.maxhs) << self.maxhs, + self.size().1.div_ceil(BLOCK_DIM << self.maxvs) << self.maxvs, + ) + } + + /// The dimensions of this frame, as coded in the codestream but including padding pixels. + pub fn size_padded(&self) -> (usize, usize) { + if self.encoding == Encoding::Modular { + self.size() + } else { + ( + self.size_blocks().0 * BLOCK_DIM, + self.size_blocks().1 * BLOCK_DIM, + ) + } + } + + /// The dimensions of this frame, after upsampling. + pub fn size_upsampled(&self) -> (usize, usize) { + ( + self.width.div_ceil(1 << (3 * self.lf_level)) as usize, + self.height.div_ceil(1 << (3 * self.lf_level)) as usize, + ) + } + + pub fn size_padded_upsampled(&self) -> (usize, usize) { + let (xsize, ysize) = self.size_padded(); + ( + xsize * self.upsampling as usize, + ysize * self.upsampling as usize, + ) + } + + /// The dimensions of this frame, in groups. + pub fn size_groups(&self) -> (usize, usize) { + ( + self.size().0.div_ceil(self.group_dim()), + self.size().1.div_ceil(self.group_dim()), + ) + } + + /// The dimensions of this frame, in LF groups. + pub fn size_lf_groups(&self) -> (usize, usize) { + ( + self.size_blocks().0.div_ceil(self.group_dim()), + self.size_blocks().1.div_ceil(self.group_dim()), + ) + } + + pub fn block_group_rect(&self, group: usize) -> Rect { + let group_dims = self.size_groups(); + let block_dims = self.size_blocks(); + let group_dim_in_blocks = self.group_dim() >> 3; + let gx = group % group_dims.0; + let gy = group / group_dims.0; + let origin = (gx * group_dim_in_blocks, gy * group_dim_in_blocks); + let size = ( + min( + block_dims.0.checked_sub(origin.0).unwrap(), + group_dim_in_blocks, + ), + min( + block_dims.1.checked_sub(origin.1).unwrap(), + group_dim_in_blocks, + ), + ); + Rect { origin, size } + } + + pub fn lf_group_rect(&self, group: usize) -> Rect { + let lf_dims = self.size_lf_groups(); + let block_dims = self.size_blocks(); + let gx = group % lf_dims.0; + let gy = group / lf_dims.0; + let origin = (gx * self.group_dim(), gy * self.group_dim()); + let size = ( + min( + block_dims.0.checked_sub(origin.0).unwrap(), + self.group_dim(), + ), + min( + block_dims.1.checked_sub(origin.1).unwrap(), + self.group_dim(), + ), + ); + Rect { origin, size } + } + + pub fn postprocess(&mut self, nonserialized: &FrameHeaderNonserialized) { + if self.upsampling > 1 { + for i in 0..nonserialized.extra_channel_info.len() { + let dim_shift = nonserialized.extra_channel_info[i].dim_shift(); + self.ec_upsampling[i] <<= dim_shift; + } + } + if self.encoding != Encoding::VarDCT || !nonserialized.xyb_encoded { + self.x_qm_scale = 2; + } + } + + fn check(&self, nonserialized: &FrameHeaderNonserialized) -> Result<(), Error> { + if self.upsampling > 1 + && let Some((info, upsampling)) = nonserialized + .extra_channel_info + .iter() + .zip(&self.ec_upsampling) + .find(|(info, ec_upsampling)| { + ((*ec_upsampling << info.dim_shift()) < self.upsampling) + || (**ec_upsampling > 8) + }) + { + return Err(Error::InvalidEcUpsampling( + self.upsampling, + info.dim_shift(), + *upsampling, + )); + } + + if self.passes.num_ds >= self.passes.num_passes { + return Err(Error::NumPassesTooLarge( + self.passes.num_ds, + self.passes.num_passes, + )); + } + + if !self.save_before_ct && !self.full_frame && self.frame_type == FrameType::ReferenceOnly { + return Err(Error::NonPatchReferenceWithCrop); + } + if !self.is444() + && ((self.flags & Flags::SKIP_ADAPTIVE_LF_SMOOTHING) == 0) + && self.encoding == Encoding::VarDCT + { + return Err(Error::Non444ChromaSubsampling); + } + Ok(()) + } +} + +#[cfg(test)] +mod test_frame_header { + use super::*; + use crate::util::test::read_headers_and_toc; + use test_log::test; + + #[test] + fn test_basic() { + let (_, frame_header, toc) = + read_headers_and_toc(include_bytes!("../../resources/test/basic.jxl")).unwrap(); + assert_eq!(frame_header.frame_type, FrameType::RegularFrame); + assert_eq!(frame_header.encoding, Encoding::VarDCT); + assert_eq!(frame_header.flags, 0); + assert_eq!(frame_header.upsampling, 1); + assert_eq!(frame_header.x_qm_scale, 2); + assert_eq!(frame_header.b_qm_scale, 2); + assert!(!frame_header.have_crop); + assert!(!frame_header.save_before_ct); + assert_eq!(frame_header.name, String::from("")); + assert_eq!(frame_header.restoration_filter.epf_iters, 1); + assert_eq!( + toc, + Toc { + permuted: false, + permutation: Permutation::default(), + entries: [53].to_vec(), + } + ) + } + + #[test] + fn test_extra_channel() { + let frame_header = + read_headers_and_toc(include_bytes!("../../resources/test/extra_channels.jxl")) + .unwrap() + .1; + assert_eq!(frame_header.frame_type, FrameType::RegularFrame); + assert_eq!(frame_header.encoding, Encoding::Modular); + assert_eq!(frame_header.flags, 0); + assert_eq!(frame_header.upsampling, 1); + assert_eq!(frame_header.ec_upsampling, vec![1]); + // libjxl x_qm_scale = 2, but condition is false (should be 3 according to the draft) + // Doesn't actually matter since this is modular mode and the value doesn't get used. + assert_eq!(frame_header.x_qm_scale, 3); + assert_eq!(frame_header.b_qm_scale, 2); + assert!(!frame_header.have_crop); + assert!(!frame_header.save_before_ct); + assert_eq!(frame_header.name, String::from("")); + assert_eq!(frame_header.restoration_filter.epf_iters, 0); + assert!(!frame_header.restoration_filter.gab); + } + + #[test] + fn test_has_permutation() { + let (_, frame_header, toc) = + read_headers_and_toc(include_bytes!("../../resources/test/has_permutation.jxl")) + .unwrap(); + assert_eq!(frame_header.frame_type, FrameType::RegularFrame); + assert_eq!(frame_header.encoding, Encoding::VarDCT); + assert_eq!(frame_header.flags, 0); + assert_eq!(frame_header.upsampling, 1); + assert_eq!(frame_header.x_qm_scale, 3); + assert_eq!(frame_header.b_qm_scale, 2); + assert!(!frame_header.have_crop); + assert!(!frame_header.save_before_ct); + assert_eq!(frame_header.name, String::from("")); + assert_eq!(frame_header.restoration_filter.epf_iters, 1); + assert_eq!( + toc, + Toc { + permuted: true, + permutation: Permutation(vec![ + 0u32, 1, 42, 48, 2, 3, 4, 5, 6, 7, 8, 9, 43, 10, 11, 12, 13, 14, 15, 16, 17, + 44, 18, 19, 20, 21, 22, 23, 24, 25, 45, 26, 27, 28, 29, 30, 31, 32, 33, 46, 34, + 35, 36, 37, 38, 39, 40, 41, 47, + ]), + entries: vec![ + 155, 992, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 9, 9, 9, 9, 9, 9, 9, + 9, 9, 9, 9, 9, 9, 9, 9, 9, 5, 5, 5, 5, 5, 5, 5, 5, 697, 5, 5, 5, 5, 5, 60, + ], + }, + ) + } +} diff --git a/third_party/rust/jxl/src/headers/image_metadata.rs b/third_party/rust/jxl/src/headers/image_metadata.rs @@ -0,0 +1,156 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +use crate::{ + bit_reader::BitReader, + error::Error, + headers::{bit_depth::*, color_encoding::*, encodings::*, extra_channels::*, size::*}, +}; +use jxl_macros::UnconditionalCoder; +use num_derive::FromPrimitive; + +#[derive(Debug, Default, Clone)] +pub struct Signature; + +impl Signature { + pub fn new() -> Signature { + Signature {} + } +} + +impl crate::headers::encodings::UnconditionalCoder<()> for Signature { + type Nonserialized = Empty; + fn read_unconditional(_: &(), br: &mut BitReader, _: &Empty) -> Result<Signature, Error> { + let sig1 = br.read(8)? as u8; + let sig2 = br.read(8)? as u8; + if (sig1, sig2) != (0xff, 0x0a) { + Err(Error::InvalidSignature) + } else { + Ok(Signature {}) + } + } +} + +#[derive(UnconditionalCoder, Copy, Clone, PartialEq, Debug, FromPrimitive)] +pub enum Orientation { + Identity = 1, + FlipHorizontal = 2, + Rotate180 = 3, + FlipVertical = 4, + Transpose = 5, + Rotate90 = 6, + AntiTranspose = 7, + Rotate270 = 8, +} + +impl Orientation { + pub fn is_transposing(self) -> bool { + matches!( + self, + Orientation::Transpose + | Orientation::AntiTranspose + | Orientation::Rotate90 + | Orientation::Rotate270 + ) + } +} + +#[derive(UnconditionalCoder, Debug, Clone)] +pub struct Animation { + #[coder(u2S(100, 1000, Bits(10) + 1, Bits(30) + 1))] + pub tps_numerator: u32, + #[coder(u2S(1, 1001, Bits(8) + 1, Bits(10) + 1))] + pub tps_denominator: u32, + #[coder(u2S(0, Bits(3), Bits(16), Bits(32)))] + pub num_loops: u32, + pub have_timecodes: bool, +} + +#[derive(UnconditionalCoder, Debug, Clone)] +#[validate] +pub struct ToneMapping { + #[all_default] + pub all_default: bool, + #[default(255.0)] + pub intensity_target: f32, + #[default(0.0)] + pub min_nits: f32, + #[default(false)] + pub relative_to_max_display: bool, + #[default(0.0)] + pub linear_below: f32, +} + +impl ToneMapping { + #[cfg(test)] + pub fn empty() -> ToneMapping { + ToneMapping { + all_default: false, + intensity_target: 0f32, + min_nits: 0f32, + relative_to_max_display: false, + linear_below: 0f32, + } + } + pub fn check(&self, _: &Empty) -> Result<(), Error> { + if self.intensity_target <= 0.0 { + Err(Error::InvalidIntensityTarget(self.intensity_target)) + } else if self.min_nits < 0.0 || self.min_nits > self.intensity_target { + Err(Error::InvalidMinNits(self.min_nits)) + } else if self.linear_below < 0.0 + || (self.relative_to_max_display && self.linear_below > 1.0) + { + Err(Error::InvalidLinearBelow( + self.relative_to_max_display, + self.linear_below, + )) + } else { + Ok(()) + } + } +} + +// TODO(veluca): figure out if these fields should be unused. +#[allow(dead_code)] +#[derive(UnconditionalCoder, Debug, Clone)] +pub struct ImageMetadata { + #[all_default] + all_default: bool, + #[default(false)] + extra_fields: bool, + #[condition(extra_fields)] + #[default(Orientation::Identity)] + #[coder(Bits(3) + 1)] + pub orientation: Orientation, + #[condition(extra_fields)] + #[default(false)] + have_intrinsic_size: bool, // TODO(veluca93): fold have_ fields in Option. + #[condition(have_intrinsic_size)] + pub intrinsic_size: Option<Size>, + #[condition(extra_fields)] + #[default(false)] + have_preview: bool, + #[condition(have_preview)] + pub preview: Option<Preview>, + #[condition(extra_fields)] + #[default(false)] + have_animation: bool, + #[condition(have_animation)] + pub animation: Option<Animation>, + #[default(BitDepth::default(&field_nonserialized))] + pub bit_depth: BitDepth, + #[default(true)] + pub modular_16bit_sufficient: bool, + #[size_coder(implicit(u2S(0, 1, Bits(4) + 2, Bits(12) + 1)))] + pub extra_channel_info: Vec<ExtraChannelInfo>, + #[default(true)] + pub xyb_encoded: bool, + #[default(ColorEncoding::default(&field_nonserialized))] + pub color_encoding: ColorEncoding, + #[condition(extra_fields)] + #[default(ToneMapping::default(&field_nonserialized))] + pub tone_mapping: ToneMapping, + extensions: Option<Extensions>, +} diff --git a/third_party/rust/jxl/src/headers/mod.rs b/third_party/rust/jxl/src/headers/mod.rs @@ -0,0 +1,68 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +pub mod bit_depth; +pub mod color_encoding; +pub mod encodings; +pub mod extra_channels; +pub mod frame_header; +pub mod image_metadata; +pub mod modular; +pub mod permutation; +pub mod size; +pub mod transform_data; + +use crate::{bit_reader::BitReader, error::Error, headers::encodings::*}; +use frame_header::FrameHeaderNonserialized; +use jxl_macros::UnconditionalCoder; + +pub use image_metadata::*; +pub use size::Size; +pub use transform_data::*; + +#[derive(UnconditionalCoder, Debug, Clone)] +pub struct FileHeader { + #[allow(dead_code)] + signature: Signature, + pub size: Size, + pub image_metadata: ImageMetadata, + #[nonserialized(xyb_encoded : image_metadata.xyb_encoded)] + pub transform_data: CustomTransformData, +} + +pub trait JxlHeader +where + Self: Sized, +{ + fn read(br: &mut BitReader) -> Result<Self, Error>; +} + +impl<T> JxlHeader for T +where + T: UnconditionalCoder<()>, + T::Nonserialized: Default, +{ + fn read(br: &mut BitReader) -> Result<Self, Error> { + Self::read_unconditional(&(), br, &T::Nonserialized::default()) + } +} + +impl FileHeader { + pub fn frame_header_nonserialized(&self) -> FrameHeaderNonserialized { + let have_timecode = match self.image_metadata.animation { + Some(ref animation) => animation.have_timecodes, + None => false, + }; + FrameHeaderNonserialized { + xyb_encoded: self.image_metadata.xyb_encoded, + num_extra_channels: self.image_metadata.extra_channel_info.len() as u32, + extra_channel_info: self.image_metadata.extra_channel_info.clone(), + have_animation: self.image_metadata.animation.is_some(), + have_timecode, + img_width: self.size.xsize(), + img_height: self.size.ysize(), + } + } +} diff --git a/third_party/rust/jxl/src/headers/modular.rs b/third_party/rust/jxl/src/headers/modular.rs @@ -0,0 +1,166 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +use crate::{ + bit_reader::BitReader, + error::{Error, Result}, + frame::modular::Predictor, + headers::encodings::*, +}; +use jxl_macros::UnconditionalCoder; +use num_derive::FromPrimitive; + +use super::encodings; + +#[derive(UnconditionalCoder, Debug, PartialEq, Clone)] +pub struct WeightedHeader { + #[all_default] + pub all_default: bool, + + #[coder(Bits(5))] + #[default(16)] + pub p1c: u32, + + #[coder(Bits(5))] + #[default(10)] + pub p2c: u32, + + #[coder(Bits(5))] + #[default(7)] + pub p3ca: u32, + + #[coder(Bits(5))] + #[default(7)] + pub p3cb: u32, + + #[coder(Bits(5))] + #[default(7)] + pub p3cc: u32, + + #[coder(Bits(5))] + #[default(0)] + pub p3cd: u32, + + #[coder(Bits(5))] + #[default(0)] + pub p3ce: u32, + + #[coder(Bits(4))] + #[default(0xd)] + pub w0: u32, + + #[coder(Bits(4))] + #[default(0xc)] + pub w1: u32, + + #[coder(Bits(4))] + #[default(0xc)] + pub w2: u32, + + #[coder(Bits(4))] + #[default(0xc)] + pub w3: u32, +} + +impl WeightedHeader { + pub fn w(&self, i: usize) -> Result<u32> { + match i { + 0 => Ok(self.w0), + 1 => Ok(self.w1), + 2 => Ok(self.w2), + 3 => Ok(self.w3), + _ => unreachable!( + "WeightedHeader::w called with an out-of-bounds index: {}. + This indicates a logical error in the calling code, which should ensure 'i' is within 0..=3.", + i), + } + } +} + +#[derive(UnconditionalCoder, Debug, PartialEq, Clone, Copy)] +pub struct SqueezeParams { + pub horizontal: bool, + pub in_place: bool, + #[coder(u2S(Bits(3), Bits(6) + 8, Bits(10) + 72, Bits(13) + 1096))] + pub begin_channel: u32, + #[coder(u2S(1, 2, 3, Bits(4) + 4))] + pub num_channels: u32, +} + +#[derive(UnconditionalCoder, Copy, Clone, PartialEq, Debug, FromPrimitive)] +pub enum TransformId { + Rct = 0, + Palette = 1, + Squeeze = 2, + Invalid = 3, +} + +#[derive(UnconditionalCoder, Debug, PartialEq)] +#[validate] +pub struct Transform { + #[coder(Bits(2))] + pub id: TransformId, + + #[condition(id == TransformId::Rct || id == TransformId::Palette)] + #[coder(u2S(Bits(3), Bits(6) + 8, Bits(10) + 72, Bits(13) + 1096))] + #[default(0)] + pub begin_channel: u32, + + #[condition(id == TransformId::Rct)] + #[coder(u2S(6, Bits(2), Bits(4) + 2, Bits(6) + 10))] + #[default(6)] + pub rct_type: u32, + + #[condition(id == TransformId::Palette)] + #[coder(u2S(1, 3, 4, Bits(13) + 1))] + #[default(3)] + pub num_channels: u32, + + #[condition(id == TransformId::Palette)] + #[coder(u2S(Bits(8), Bits(10) + 256, Bits(12) + 1280, Bits(16)+5376))] + #[default(256)] + pub num_colors: u32, + + #[condition(id == TransformId::Palette)] + #[coder(u2S(0, Bits(8)+1, Bits(10) + 257, Bits(16)+1281))] + #[default(0)] + pub num_deltas: u32, + + #[condition(id == TransformId::Palette)] + #[coder(Bits(4))] + #[default(0)] + pub predictor_id: u32, + + #[condition(id == TransformId::Squeeze)] + #[size_coder(implicit(u2S(0, Bits(4) + 1, Bits(6) + 9, Bits(8) + 41)))] + #[default(Vec::new())] + pub squeezes: Vec<SqueezeParams>, +} + +impl Transform { + fn check(&self, _: &encodings::Empty) -> Result<()> { + if self.id == TransformId::Invalid { + return Err(Error::InvalidTransformId); + } + + if self.rct_type >= 42 { + return Err(Error::InvalidRCT(self.rct_type)); + } + + if self.predictor_id >= Predictor::NUM_PREDICTORS { + return Err(Error::InvalidPredictor(self.predictor_id)); + } + + Ok(()) + } +} + +#[derive(UnconditionalCoder, Debug, PartialEq)] +pub struct GroupHeader { + pub use_global_tree: bool, + pub wp_header: WeightedHeader, + #[size_coder(implicit(u2S(0, 1, Bits(4) + 2, Bits(8) + 18)))] + pub transforms: Vec<Transform>, +} diff --git a/third_party/rust/jxl/src/headers/permutation.rs b/third_party/rust/jxl/src/headers/permutation.rs @@ -0,0 +1,340 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +use crate::bit_reader::BitReader; +use crate::entropy_coding::decode::Histograms; +use crate::entropy_coding::decode::SymbolReader; +use crate::error::{Error, Result}; +use crate::util::{CeilLog2, NewWithCapacity, tracing_wrappers::instrument, value_of_lowest_1_bit}; + +#[derive(Debug, PartialEq, Default)] +pub struct Permutation(pub Vec<u32>); + +impl std::ops::Deref for Permutation { + type Target = [u32]; + + fn deref(&self) -> &[u32] { + &self.0 + } +} + +impl Permutation { + /// Decode a permutation from entropy-coded stream. + pub fn decode( + size: u32, + skip: u32, + histograms: &Histograms, + br: &mut BitReader, + entropy_reader: &mut SymbolReader, + ) -> Result<Self> { + let end = entropy_reader.read_unsigned(histograms, br, get_context(size))?; + Self::decode_inner(size, skip, end, |ctx| { + entropy_reader.read_unsigned(histograms, br, ctx) + }) + } + + fn decode_inner( + size: u32, + skip: u32, + end: u32, + mut read: impl FnMut(usize) -> Result<u32>, + ) -> Result<Self> { + if end > size - skip { + return Err(Error::InvalidPermutationSize { size, skip, end }); + } + + let mut lehmer = Vec::new_with_capacity(end as usize)?; + + let mut prev_val = 0u32; + for idx in skip..(skip + end) { + let val = read(get_context(prev_val))?; + if val >= size - idx { + return Err(Error::InvalidPermutationLehmerCode { + size, + idx, + lehmer: val, + }); + } + lehmer.push(val); + prev_val = val; + } + + // Initialize the full permutation vector with skipped elements intact + let mut permutation = Vec::new_with_capacity((size - skip) as usize)?; + permutation.extend(0..size); + + // Decode the Lehmer code into the slice starting at `skip` + let permuted_slice = decode_lehmer_code(&lehmer, &permutation[skip as usize..])?; + + // Replace the target slice in `permutation` + permutation[skip as usize..].copy_from_slice(&permuted_slice); + + // Ensure the permutation has the correct size + assert_eq!(permutation.len(), size as usize); + + Ok(Self(permutation)) + } + + pub fn compose(&mut self, other: &Permutation) { + assert_eq!(self.0.len(), other.0.len()); + let mut tmp: Vec<u32> = vec![0; self.0.len()]; + for (i, val) in tmp.iter_mut().enumerate().take(self.0.len()) { + *val = self.0[other.0[i] as usize] + } + self.0.copy_from_slice(&tmp[..]); + } +} + +// Decodes the Lehmer code in `code` and returns the permuted slice. +#[instrument(level = "debug", ret, err)] +fn decode_lehmer_code(code: &[u32], permutation_slice: &[u32]) -> Result<Vec<u32>> { + let n = permutation_slice.len(); + if n == 0 { + return Err(Error::InvalidPermutationLehmerCode { + size: 0, + idx: 0, + lehmer: 0, + }); + } + + let mut permuted = Vec::new_with_capacity(n)?; + permuted.extend_from_slice(permutation_slice); + + let padded_n = (n as u32).next_power_of_two() as usize; + + // Allocate temp array inside the function + let mut temp = Vec::new_with_capacity(padded_n)?; + temp.extend((0..padded_n as u32).map(|x| value_of_lowest_1_bit(x + 1))); + + for (i, permuted_item) in permuted.iter_mut().enumerate() { + let code_i = *code.get(i).unwrap_or(&0); + + // Adjust the maximum allowed value for code_i + if code_i as usize > n - i - 1 { + return Err(Error::InvalidPermutationLehmerCode { + size: n as u32, + idx: i as u32, + lehmer: code_i, + }); + } + + let mut rank = code_i + 1; + + // Extract i-th unused element via implicit order-statistics tree. + let mut bit = padded_n; + let mut next = 0usize; + while bit != 0 { + let cand = next + bit; + if cand == 0 || cand > padded_n { + return Err(Error::InvalidPermutationLehmerCode { + size: n as u32, + idx: i as u32, + lehmer: code_i, + }); + } + bit >>= 1; + if temp[cand - 1] < rank { + next = cand; + rank -= temp[cand - 1]; + } + } + + *permuted_item = permutation_slice[next]; + + next += 1; + while next <= padded_n { + temp[next - 1] -= 1; + next += value_of_lowest_1_bit(next as u32) as usize; + } + } + + Ok(permuted) +} + +// Decodes the Lehmer code in `code` and returns the permuted vector. +#[cfg(test)] +fn decode_lehmer_code_naive(code: &[u32], permutation_slice: &[u32]) -> Result<Vec<u32>> { + let n = code.len(); + if n == 0 { + return Err(Error::InvalidPermutationLehmerCode { + size: 0, + idx: 0, + lehmer: 0, + }); + } + + // Ensure permutation_slice has sufficient length + if permutation_slice.len() < n { + return Err(Error::InvalidPermutationLehmerCode { + size: n as u32, + idx: 0, + lehmer: 0, + }); + } + + // Create temp array with values from permutation_slice + let mut temp = permutation_slice.to_vec(); + let mut permuted = Vec::new_with_capacity(n)?; + + // Iterate over the Lehmer code + for (i, &idx) in code.iter().enumerate() { + if idx as usize >= temp.len() { + return Err(Error::InvalidPermutationLehmerCode { + size: n as u32, + idx: i as u32, + lehmer: idx, + }); + } + + // Assign temp[idx] to permuted vector + permuted.push(temp.remove(idx as usize)); + } + + // Append any remaining elements from temp to permuted + permuted.extend(temp); + + Ok(permuted) +} + +fn get_context(x: u32) -> usize { + (x + 1).ceil_log2().min(7) as usize +} + +#[cfg(test)] +mod test { + use super::*; + use arbtest::arbitrary::{self, Arbitrary, Unstructured}; + use core::assert_eq; + use test_log::test; + + #[test] + fn generate_permutation_arbtest() { + arbtest::arbtest(|u| { + let input = PermutationInput::arbitrary(u)?; + + let permutation_slice = input.permutation.as_slice(); + + let perm1 = decode_lehmer_code(&input.code, permutation_slice); + let perm2 = decode_lehmer_code_naive(&input.code, permutation_slice); + + assert_eq!( + perm1.map_err(|x| x.to_string()), + perm2.map_err(|x| x.to_string()) + ); + Ok(()) + }); + } + + #[derive(Debug)] + struct PermutationInput { + code: Vec<u32>, + permutation: Vec<u32>, + } + + impl<'a> Arbitrary<'a> for PermutationInput { + fn arbitrary(u: &mut Unstructured<'a>) -> Result<Self, arbitrary::Error> { + // Generate a reasonable size to prevent tests from taking too long + let size_lehmer = u.int_in_range(1..=1000)?; + + let mut lehmer: Vec<u32> = Vec::with_capacity(size_lehmer as usize); + for i in 0..size_lehmer { + let max_val = size_lehmer - i - 1; + let val = if max_val > 0 { + u.int_in_range(0..=max_val)? + } else { + 0 + }; + lehmer.push(val); + } + + let mut permutation = Vec::new(); + let size_permutation = u.int_in_range(size_lehmer..=1000)?; + permutation.extend(0..size_permutation); + + let num_of_swaps = u.int_in_range(0..=100)?; + for _ in 0..num_of_swaps { + // Randomly swap two positions + let pos1 = u.int_in_range(0..=size_permutation - 1)?; + let pos2 = u.int_in_range(0..=size_permutation - 1)?; + permutation.swap(pos1 as usize, pos2 as usize); + } + + Ok(PermutationInput { + code: lehmer, + permutation, + }) + } + } + + #[test] + fn simple() { + // Lehmer code: [1, 1, 2, 3, 3, 6, 0, 1] + let code = vec![1u32, 1, 2, 3, 3, 6, 0, 1]; + let skip = 4; + let size = 16; + + let permutation_slice: Vec<u32> = (skip..size).collect(); + + let permuted = decode_lehmer_code(&code, &permutation_slice).unwrap(); + let permuted_naive = decode_lehmer_code_naive(&code, &permutation_slice).unwrap(); + + let mut permutation = Vec::with_capacity(size as usize); + permutation.extend(0..skip); // Add skipped elements + permutation.extend(permuted.iter()); + let expected_permutation = vec![0, 1, 2, 3, 5, 6, 8, 10, 11, 15, 4, 9, 7, 12, 13, 14]; + + assert_eq!(permutation, expected_permutation); + assert_eq!(permuted, permuted_naive); + } + + #[test] + fn decode_lehmer_compare_different_length() -> Result<(), Box<dyn std::error::Error>> { + // Lehmer code: [1, 1, 2, 3, 3, 6, 0, 1] + let code = vec![1u32, 1, 2, 3, 3, 6, 0, 1]; + let skip = 4; + let size = 16; + + let permutation_slice: Vec<u32> = (skip..size).collect(); + + let permuted_optimized = decode_lehmer_code(&code, &permutation_slice)?; + let permuted_naive = decode_lehmer_code_naive(&code, &permutation_slice)?; + + let expected_permuted = vec![5u32, 6, 8, 10, 11, 15, 4, 9, 7, 12, 13, 14]; + + assert_eq!(permuted_optimized, expected_permuted); + assert_eq!(permuted_naive, expected_permuted); + assert_eq!(permuted_optimized, permuted_naive); + + Ok(()) + } + + #[test] + fn decode_lehmer_compare_same_length() -> Result<(), Box<dyn std::error::Error>> { + // Lehmer code: [2, 3, 0, 0, 0] + let code = vec![2u32, 3, 0, 0, 0]; + let n = code.len(); + let permutation_slice: Vec<u32> = (0..n as u32).collect(); + + let permuted_optimized = decode_lehmer_code(&code, &permutation_slice)?; + let permuted_naive = decode_lehmer_code_naive(&code, &permutation_slice)?; + + let expected_permutation = vec![2u32, 4, 0, 1, 3]; + + assert_eq!(permuted_optimized, expected_permutation); + assert_eq!(permuted_naive, expected_permutation); + assert_eq!(permuted_optimized, permuted_naive); + + Ok(()) + } + + #[test] + fn lehmer_out_of_bounds() { + let code = vec![4]; + let permutation_slice: Vec<u32> = (4..8).collect(); + + let result = decode_lehmer_code(&code, &permutation_slice); + assert!(result.is_err()); + } +} diff --git a/third_party/rust/jxl/src/headers/size.rs b/third_party/rust/jxl/src/headers/size.rs @@ -0,0 +1,112 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +use crate::{bit_reader::BitReader, error::Error, headers::encodings::*}; +use jxl_macros::UnconditionalCoder; +use num_derive::FromPrimitive; + +#[derive(UnconditionalCoder, Copy, Clone, PartialEq, Debug, FromPrimitive, Default)] +enum AspectRatio { + #[default] + Unknown = 0, + Ratio1Over1 = 1, + Ratio12Over10 = 2, + Ratio4Over3 = 3, + Ratio3Over2 = 4, + Ratio16Over9 = 5, + Ratio5Over4 = 6, + Ratio2Over1 = 7, +} + +#[derive(UnconditionalCoder, Debug, Clone, Default)] +pub struct Size { + small: bool, + #[condition(small)] + #[coder(Bits(5) + 1)] + ysize_div8: Option<u32>, + #[condition(!small)] + #[coder(1 + u2S(Bits(9), Bits(13), Bits(18), Bits(30)))] + ysize: Option<u32>, + #[coder(Bits(3))] + ratio: AspectRatio, + #[condition(small && ratio == AspectRatio::Unknown)] + #[coder(Bits(5) + 1)] + xsize_div8: Option<u32>, + #[condition(!small && ratio == AspectRatio::Unknown)] + #[coder(1 + u2S(Bits(9), Bits(13), Bits(18), Bits(30)))] + xsize: Option<u32>, +} + +#[derive(UnconditionalCoder, Debug, Clone)] +pub struct Preview { + div8: bool, + #[condition(div8)] + #[coder(u2S(16, 32, Bits(5) + 1, Bits(9) + 33))] + ysize_div8: Option<u32>, + #[condition(!div8)] + #[coder(1 + u2S(Bits(6), Bits(8) + 64, Bits(10) + 320, Bits(12) + 1344))] + ysize: Option<u32>, + #[coder(Bits(3))] + ratio: AspectRatio, + #[condition(div8 && ratio == AspectRatio::Unknown)] + #[coder(u2S(16, 32, Bits(5) + 1, Bits(9) + 33))] + xsize_div8: Option<u32>, + #[condition(!div8 && ratio == AspectRatio::Unknown)] + #[coder(1 + u2S(Bits(6), Bits(8) + 64, Bits(10) + 320, Bits(12) + 1344))] + xsize: Option<u32>, +} + +fn map_aspect_ratio<T: Fn() -> u32>(ysize: u32, ratio: AspectRatio, fallback: T) -> u32 { + match ratio { + AspectRatio::Unknown => fallback(), + AspectRatio::Ratio1Over1 => ysize, + AspectRatio::Ratio12Over10 => (ysize as u64 * 12 / 10) as u32, + AspectRatio::Ratio4Over3 => (ysize as u64 * 4 / 3) as u32, + AspectRatio::Ratio3Over2 => (ysize as u64 * 3 / 2) as u32, + AspectRatio::Ratio16Over9 => (ysize as u64 * 16 / 9) as u32, + AspectRatio::Ratio5Over4 => (ysize as u64 * 5 / 4) as u32, + AspectRatio::Ratio2Over1 => ysize * 2, + } +} + +impl Size { + pub fn ysize(&self) -> u32 { + if self.small { + self.ysize_div8.unwrap() * 8 + } else { + self.ysize.unwrap() + } + } + + pub fn xsize(&self) -> u32 { + map_aspect_ratio(self.ysize(), self.ratio, /* fallback */ || { + if self.small { + self.xsize_div8.unwrap() * 8 + } else { + self.xsize.unwrap() + } + }) + } +} + +impl Preview { + pub fn ysize(&self) -> u32 { + if self.div8 { + self.ysize_div8.unwrap() * 8 + } else { + self.ysize.unwrap() + } + } + + pub fn xsize(&self) -> u32 { + map_aspect_ratio(self.ysize(), self.ratio, /* fallback */ || { + if self.div8 { + self.xsize_div8.unwrap() * 8 + } else { + self.xsize.unwrap() + } + }) + } +} diff --git a/third_party/rust/jxl/src/headers/transform_data.rs b/third_party/rust/jxl/src/headers/transform_data.rs @@ -0,0 +1,342 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +#![allow(clippy::excessive_precision)] + +use crate::{bit_reader::BitReader, error::Error, headers::encodings::*}; +use jxl_macros::UnconditionalCoder; +#[derive(Default)] +pub struct CustomTransformDataNonserialized { + pub xyb_encoded: bool, +} + +#[derive(UnconditionalCoder, Debug, Clone)] +pub struct OpsinInverseMatrix { + // We never read the all_default field. + #[allow(dead_code)] + #[all_default] + all_default: bool, + #[default([11.031566901960783, -9.866943921568629, -0.16462299647058826, + -3.254147380392157, 4.418770392156863, -0.16462299647058826, + -3.6588512862745097, 2.7129230470588235, 1.9459282392156863])] + pub inverse_matrix: [f32; 9], + #[default([-0.0037930732552754493; 3])] + pub opsin_biases: [f32; 3], + #[default([1.0 - 0.05465007330715401, 1.0 - 0.07005449891748593, 1.0 - 0.049935103337343655, 0.145])] + pub quant_biases: [f32; 4], +} + +const DEFAULT_KERN_2: [f32; 15] = [ + -0.01716200, + -0.03452303, + -0.04022174, + -0.02921014, + -0.00624645, + 0.14111091, + 0.28896755, + 0.00278718, + -0.01610267, + 0.56661550, + 0.03777607, + -0.01986694, + -0.03144731, + -0.01185068, + -0.00213539, +]; + +const DEFAULT_KERN_4: [f32; 55] = [ + -0.02419067, + -0.03491987, + -0.03693351, + -0.03094285, + -0.00529785, + -0.01663432, + -0.03556863, + -0.03888905, + -0.03516850, + -0.00989469, + 0.23651958, + 0.33392945, + -0.01073543, + -0.01313181, + -0.03556694, + 0.13048175, + 0.40103025, + 0.03951150, + -0.02077584, + 0.46914198, + -0.00209270, + -0.01484589, + -0.04064806, + 0.18942530, + 0.56279892, + 0.06674400, + -0.02335494, + -0.03551682, + -0.00754830, + -0.02267919, + -0.02363578, + 0.00315804, + -0.03399098, + -0.01359519, + -0.00091653, + -0.00335467, + -0.01163294, + -0.01610294, + -0.00974088, + -0.00191622, + -0.01095446, + -0.03198464, + -0.04455121, + -0.02799790, + -0.00645912, + 0.06390599, + 0.22963888, + 0.00630981, + -0.01897349, + 0.67537268, + 0.08483369, + -0.02534994, + -0.02205197, + -0.01667999, + -0.00384443, +]; + +const DEFAULT_KERN_8: [f32; 210] = [ + -0.02928613, + -0.03706353, + -0.03783812, + -0.03324558, + -0.00447632, + -0.02519406, + -0.03752601, + -0.03901508, + -0.03663285, + -0.00646649, + -0.02066407, + -0.03838633, + -0.04002101, + -0.03900035, + -0.00901973, + -0.01626393, + -0.03954148, + -0.04046620, + -0.03979621, + -0.01224485, + 0.29895328, + 0.35757708, + -0.02447552, + -0.01081748, + -0.04314594, + 0.23903219, + 0.41119301, + -0.00573046, + -0.01450239, + -0.04246845, + 0.17567618, + 0.45220643, + 0.02287757, + -0.01936783, + -0.03583255, + 0.11572472, + 0.47416733, + 0.06284440, + -0.02685066, + 0.42720050, + -0.02248939, + -0.01155273, + -0.04562755, + 0.28689496, + 0.49093869, + -0.00007891, + -0.01545926, + -0.04562659, + 0.21238920, + 0.53980934, + 0.03369474, + -0.02070211, + -0.03866988, + 0.14229550, + 0.56593398, + 0.08045181, + -0.02888298, + -0.03680918, + -0.00542229, + -0.02920477, + -0.02788574, + -0.02118180, + -0.03942402, + -0.00775547, + -0.02433614, + -0.03193943, + -0.02030828, + -0.04044014, + -0.01074016, + -0.01930822, + -0.03620399, + -0.01974125, + -0.03919545, + -0.01456093, + -0.00045072, + -0.00360110, + -0.01020207, + -0.01231907, + -0.00638988, + -0.00071592, + -0.00279122, + -0.00957115, + -0.01288327, + -0.00730937, + -0.00107783, + -0.00210156, + -0.00890705, + -0.01317668, + -0.00813895, + -0.00153491, + -0.02128481, + -0.04173044, + -0.04831487, + -0.03293190, + -0.00525260, + -0.01720322, + -0.04052736, + -0.05045706, + -0.03607317, + -0.00738030, + -0.01341764, + -0.03965629, + -0.05151616, + -0.03814886, + -0.01005819, + 0.18968273, + 0.33063684, + -0.01300105, + -0.01372950, + -0.04017465, + 0.13727832, + 0.36402234, + 0.01027890, + -0.01832107, + -0.03365072, + 0.08734506, + 0.38194295, + 0.04338228, + -0.02525993, + 0.56408126, + 0.00458352, + -0.01648227, + -0.04887868, + 0.24585519, + 0.62026135, + 0.04314807, + -0.02213737, + -0.04158014, + 0.16637289, + 0.65027023, + 0.09621636, + -0.03101388, + -0.04082742, + -0.00904519, + -0.02790922, + -0.02117818, + 0.00798662, + -0.03995711, + -0.01243427, + -0.02231705, + -0.02946266, + 0.00992055, + -0.03600283, + -0.01684920, + -0.00111684, + -0.00411204, + -0.01297130, + -0.01723725, + -0.01022545, + -0.00165306, + -0.00313110, + -0.01218016, + -0.01763266, + -0.01125620, + -0.00231663, + -0.01374149, + -0.03797620, + -0.05142937, + -0.03117307, + -0.00581914, + -0.01064003, + -0.03608089, + -0.05272168, + -0.03375670, + -0.00795586, + 0.09628104, + 0.27129991, + -0.00353779, + -0.01734151, + -0.03153981, + 0.05686230, + 0.28500998, + 0.02230594, + -0.02374955, + 0.68214326, + 0.05018048, + -0.02320852, + -0.04383616, + 0.18459474, + 0.71517975, + 0.10805613, + -0.03263677, + -0.03637639, + -0.01394373, + -0.02511203, + -0.01728636, + 0.05407331, + -0.02867568, + -0.01893131, + -0.00240854, + -0.00446511, + -0.01636187, + -0.02377053, + -0.01522848, + -0.00333334, + -0.00819975, + -0.02964169, + -0.04499287, + -0.02745350, + -0.00612408, + 0.02727416, + 0.19446600, + 0.00159832, + -0.02232473, + 0.74982506, + 0.11452620, + -0.03348048, + -0.01605681, + -0.02070339, + -0.00458223, +]; + +// TODO(firsching): remove once we use this! +#[allow(dead_code)] +#[derive(UnconditionalCoder, Debug, Clone)] +#[nonserialized(CustomTransformDataNonserialized)] +pub struct CustomTransformData { + #[all_default] + all_default: bool, + #[condition(nonserialized.xyb_encoded)] + #[default(OpsinInverseMatrix::default(&field_nonserialized))] + pub opsin_inverse_matrix: OpsinInverseMatrix, + #[default(0)] + #[coder(Bits(3))] + custom_weight_mask: u32, + #[condition((custom_weight_mask & 1) != 0)] + #[default(DEFAULT_KERN_2)] + pub weights2: [f32; 15], + #[condition((custom_weight_mask & 2) != 0)] + #[default(DEFAULT_KERN_4)] + pub weights4: [f32; 55], + #[condition((custom_weight_mask & 4) != 0)] + #[default(DEFAULT_KERN_8)] + pub weights8: [f32; 210], +} diff --git a/third_party/rust/jxl/src/icc.rs b/third_party/rust/jxl/src/icc.rs @@ -0,0 +1,202 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +use std::io::Cursor; + +use byteorder::{BigEndian, ReadBytesExt, WriteBytesExt}; + +use crate::bit_reader::*; +use crate::entropy_coding::decode::Histograms; +use crate::entropy_coding::decode::SymbolReader; +use crate::error::{Error, Result}; +use crate::headers::encodings::*; +use crate::util::NewWithCapacity; +use crate::util::tracing_wrappers::warn; + +mod header; +mod stream; +mod tag; + +use header::read_header; +use stream::{IccStream, read_varint_from_reader}; +use tag::{read_single_command, read_tag_list}; + +const ICC_CONTEXTS: usize = 41; +const ICC_HEADER_SIZE: u64 = 128; + +fn read_icc_inner(stream: &mut IccStream) -> Result<Vec<u8>, Error> { + let output_size = stream.read_varint()?; + let commands_size = stream.read_varint()?; + if stream.bytes_read().saturating_add(commands_size) > stream.len() { + return Err(Error::InvalidIccStream); + } + + // Simple check to avoid allocating too large buffer. + if output_size > (1 << 28) { + return Err(Error::IccTooLarge); + } + + if output_size + 65536 < stream.len() { + return Err(Error::IccTooLarge); + } + + // Extract command stream first. + let commands = stream.read_to_vec_exact(commands_size as usize)?; + let mut commands_stream = Cursor::new(commands); + // `stream` contains data stream from here. + let data_stream = stream; + + // Decode ICC profile header. + let mut decoded_profile = read_header(data_stream, output_size)?; + if output_size <= ICC_HEADER_SIZE { + return Ok(decoded_profile); + } + + // Convert to slice writer to prevent buffer from growing. + // `read_header` above returns buffer with capacity of `output_size`, so this doesn't realloc. + debug_assert_eq!(decoded_profile.capacity(), output_size as usize); + decoded_profile.resize(output_size as usize, 0); + let mut decoded_profile_writer = Cursor::new(&mut *decoded_profile); + decoded_profile_writer.set_position(ICC_HEADER_SIZE); + + // Decode tag list. + let v = read_varint_from_reader(&mut commands_stream)?; + if let Some(num_tags) = v.checked_sub(1) { + if (output_size - ICC_HEADER_SIZE) / 12 < num_tags { + warn!(output_size, num_tags, "num_tags too large"); + return Err(Error::InvalidIccStream); + } + + let num_tags = num_tags as u32; + decoded_profile_writer + .write_u32::<BigEndian>(num_tags) + .map_err(|_| Error::InvalidIccStream)?; + + read_tag_list( + data_stream, + &mut commands_stream, + &mut decoded_profile_writer, + num_tags, + output_size, + )?; + } + + // Decode tag data. + // Will not enter the loop if end of stream was reached while decoding tag list. + while let Ok(command) = commands_stream.read_u8() { + read_single_command( + data_stream, + &mut commands_stream, + &mut decoded_profile_writer, + command, + )?; + } + + // Validate output size. + let actual_len = decoded_profile_writer.position(); + if actual_len != output_size { + warn!(output_size, actual_len, "ICC profile size mismatch"); + return Err(Error::InvalidIccStream); + } + + Ok(decoded_profile) +} + +/// Struct to incrementally decode an ICC profile. +pub struct IncrementalIccReader { + histograms: Histograms, + reader: SymbolReader, + out_buf: Vec<u8>, + len: usize, + // [prev, prev_prev] + prev_bytes: [u8; 2], +} + +impl IncrementalIccReader { + pub fn new(br: &mut BitReader) -> Result<Self> { + let len = u64::read_unconditional(&(), br, &Empty {})?; + if len > 1u64 << 20 { + return Err(Error::IccTooLarge); + } + + let len = len as usize; + + let histograms = Histograms::decode(ICC_CONTEXTS, br, true)?; + let reader = SymbolReader::new(&histograms, br, None)?; + Ok(Self { + histograms, + reader, + len, + out_buf: Vec::new_with_capacity(len)?, + prev_bytes: [0, 0], + }) + } + + fn get_icc_ctx(&self) -> u32 { + if self.out_buf.len() <= ICC_HEADER_SIZE as usize { + return 0; + } + + let [b1, b2] = self.prev_bytes; + + let p1 = match b1 { + b'a'..=b'z' | b'A'..=b'Z' => 0, + b'0'..=b'9' | b'.' | b',' => 1, + 0..=1 => 2 + b1 as u32, + 2..=15 => 4, + 241..=254 => 5, + 255 => 6, + _ => 7, + }; + let p2 = match b2 { + b'a'..=b'z' | b'A'..=b'Z' => 0, + b'0'..=b'9' | b'.' | b',' => 1, + 0..=15 => 2, + 241..=255 => 3, + _ => 4, + }; + + 1 + p1 + 8 * p2 + } + + pub fn num_coded_bytes(&self) -> usize { + self.len + } + + pub fn remaining(&self) -> usize { + assert!(self.num_coded_bytes() >= self.out_buf.len()); + self.num_coded_bytes() - self.out_buf.len() + } + + pub fn read_one(&mut self, br: &mut BitReader) -> Result<()> { + let ctx = self.get_icc_ctx() as usize; + let sym = self.reader.read_unsigned(&self.histograms, br, ctx)?; + if sym >= 256 { + warn!(sym, "Invalid symbol in ICC stream"); + return Err(Error::InvalidIccStream); + } + let b = sym as u8; + self.out_buf.push(b); + self.prev_bytes = [b, self.prev_bytes[0]]; + + Ok(()) + } + + pub fn read_all(&mut self, br: &mut BitReader) -> Result<()> { + for _ in self.out_buf.len()..self.num_coded_bytes() { + self.read_one(br)? + } + Ok(()) + } + + pub fn finalize(self) -> Result<Vec<u8>> { + assert_eq!(self.num_coded_bytes(), self.out_buf.len()); + self.reader.check_final_state(&self.histograms)?; + let mut stream = IccStream::new(self.out_buf); + let profile = read_icc_inner(&mut stream)?; + stream.finalize()?; + Ok(profile) + } +} diff --git a/third_party/rust/jxl/src/icc/header.rs b/third_party/rust/jxl/src/icc/header.rs @@ -0,0 +1,51 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +use crate::{error::Result, util::NewWithCapacity}; + +use super::{ICC_HEADER_SIZE, IccStream}; + +fn predict_header(idx: usize, output_size: u32, header: &[u8]) -> u8 { + match idx { + 0..=3 => output_size.to_be_bytes()[idx], + 8 => 4, + 12..=23 => b"mntrRGB XYZ "[idx - 12], + 36..=39 => b"acsp"[idx - 36], + // APPL + 41 | 42 if header[40] == b'A' => b'P', + 43 if header[40] == b'A' => b'L', + // MSFT + 41 if header[40] == b'M' => b'S', + 42 if header[40] == b'M' => b'F', + 43 if header[40] == b'M' => b'T', + // SGI_ + 42 if header[40] == b'S' && header[41] == b'G' => b'I', + 43 if header[40] == b'S' && header[41] == b'G' => b' ', + // SUNW + 42 if header[40] == b'S' && header[41] == b'U' => b'N', + 43 if header[40] == b'S' && header[41] == b'U' => b'W', + 70 => 246, + 71 => 214, + 73 => 1, + 78 => 211, + 79 => 45, + 80..=83 => header[4 + idx - 80], + _ => 0, + } +} + +pub(super) fn read_header(data_stream: &mut IccStream, output_size: u64) -> Result<Vec<u8>> { + let header_size = output_size.min(ICC_HEADER_SIZE); + let header_data = data_stream.read_to_vec_exact(header_size as usize)?; + + let mut profile = Vec::new_with_capacity(output_size as usize)?; + + for (idx, &e) in header_data.iter().enumerate() { + let p = predict_header(idx, output_size as u32, &header_data); + profile.push(p.wrapping_add(e)); + } + + Ok(profile) +} diff --git a/third_party/rust/jxl/src/icc/stream.rs b/third_party/rust/jxl/src/icc/stream.rs @@ -0,0 +1,118 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +use std::io::{Read, Write}; + +use byteorder::{ReadBytesExt, WriteBytesExt}; + +use crate::error::{Error, Result}; +use crate::util::NewWithCapacity; +use crate::util::tracing_wrappers::{instrument, warn}; + +fn read_varint(mut read_one: impl FnMut() -> Result<u8>) -> Result<u64> { + let mut value = 0u64; + let mut shift = 0; + while shift < 63 { + let b = read_one()?; + value |= ((b & 0x7f) as u64) << shift; + if b & 0x80 == 0 { + break; + } + shift += 7; + } + Ok(value) +} + +pub(super) fn read_varint_from_reader(stream: &mut impl Read) -> Result<u64> { + read_varint(|| stream.read_u8().map_err(|_| Error::IccEndOfStream)) +} + +pub(super) struct IccStream { + command_stream: Vec<u8>, + bytes_read: u64, +} + +impl IccStream { + pub(super) fn new(command_stream: Vec<u8>) -> Self { + Self { + command_stream, + bytes_read: 0, + } + } + + pub fn len(&self) -> u64 { + self.command_stream.len() as u64 + } + + pub fn bytes_read(&self) -> u64 { + self.bytes_read + } + + pub fn remaining_bytes(&self) -> u64 { + self.len() - self.bytes_read + } + + fn read_one(&mut self) -> Result<u8> { + if self.remaining_bytes() == 0 { + return Err(Error::IccEndOfStream); + } + self.bytes_read += 1; + Ok(self.command_stream[self.bytes_read as usize - 1]) + } + + pub fn read_exact(&mut self, buf: &mut [u8]) -> Result<()> { + if buf.len() > self.remaining_bytes() as usize { + return Err(Error::IccEndOfStream); + } + + for b in buf { + *b = self.read_one()?; + } + + Ok(()) + } + + pub fn read_to_vec_exact(&mut self, len: usize) -> Result<Vec<u8>> { + if len > self.remaining_bytes() as usize { + return Err(Error::IccEndOfStream); + } + + let mut out = Vec::new_with_capacity(len)?; + + for _ in 0..len { + out.push(self.read_one()?); + } + + Ok(out) + } + + pub fn read_varint(&mut self) -> Result<u64> { + read_varint(|| self.read_one()) + } + + pub fn copy_bytes(&mut self, writer: &mut impl Write, len: usize) -> Result<()> { + if len > self.remaining_bytes() as usize { + return Err(Error::IccEndOfStream); + } + + for _ in 0..len { + let b = self.read_one()?; + writer.write_u8(b).map_err(|_| Error::InvalidIccStream)?; + } + + Ok(()) + } + + #[instrument(skip_all, err)] + pub fn finalize(self) -> Result<()> { + // Check if all bytes are read + if self.bytes_read == self.command_stream.len() as u64 { + Ok(()) + } else { + warn!("ICC stream is not fully consumed"); + Err(Error::InvalidIccStream) + } + } +} diff --git a/third_party/rust/jxl/src/icc/tag.rs b/third_party/rust/jxl/src/icc/tag.rs @@ -0,0 +1,242 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +use std::io::{Cursor, Write}; + +use byteorder::{BigEndian, ReadBytesExt, WriteBytesExt}; + +use crate::error::{Error, Result}; +use crate::util::NewWithCapacity; +use crate::util::tracing_wrappers::warn; + +use super::{ICC_HEADER_SIZE, IccStream, read_varint_from_reader}; + +const COMMON_TAGS: [&[u8; 4]; 19] = [ + b"rTRC", b"rXYZ", b"cprt", b"wtpt", b"bkpt", b"rXYZ", b"gXYZ", b"bXYZ", b"kXYZ", b"rTRC", + b"gTRC", b"bTRC", b"kTRC", b"chad", b"desc", b"chrm", b"dmnd", b"dmdd", b"lumi", +]; + +const COMMON_DATA: [&[u8; 4]; 8] = [ + b"XYZ ", b"desc", b"text", b"mluc", b"para", b"curv", b"sf32", b"gbd ", +]; + +pub(super) fn read_tag_list( + data_stream: &mut IccStream, + commands_stream: &mut Cursor<Vec<u8>>, + decoded_profile: &mut Cursor<&mut [u8]>, + num_tags: u32, + output_size: u64, +) -> Result<()> { + let mut prev_tagstart = num_tags * 12 + ICC_HEADER_SIZE as u32; + let mut prev_tagsize = 0u32; + + loop { + let Ok(command) = commands_stream.read_u8() else { + // End of commands stream + return Ok(()); + }; + + let tagcode = command & 63; + let tag = match tagcode { + 0 => break, + 1 => { + let mut tag = [0u8; 4]; + data_stream.read_exact(&mut tag)?; + tag + } + 2..=20 => *COMMON_TAGS[(tagcode - 2) as usize], + _ => return Err(Error::InvalidIccStream), + }; + + let tagstart = if command & 64 == 0 { + prev_tagstart + prev_tagsize + } else { + read_varint_from_reader(commands_stream)? as u32 + }; + let tagsize = match &tag { + _ if command & 128 != 0 => read_varint_from_reader(commands_stream)? as u32, + b"rXYZ" | b"gXYZ" | b"bXYZ" | b"kXYZ" | b"wtpt" | b"bkpt" | b"lumi" => 20, + _ => prev_tagsize, + }; + if (tagstart as u64 + tagsize as u64) > output_size { + warn!(output_size, tagstart, tagsize, "tag size overflow"); + return Err(Error::InvalidIccStream); + } + + prev_tagstart = tagstart; + prev_tagsize = tagsize; + + let write_result = (|| -> std::io::Result<()> { + decoded_profile.write_all(&tag)?; + decoded_profile.write_u32::<BigEndian>(tagstart)?; + decoded_profile.write_u32::<BigEndian>(tagsize)?; + if tagcode == 2 { + decoded_profile.write_all(b"gTRC")?; + decoded_profile.write_u32::<BigEndian>(tagstart)?; + decoded_profile.write_u32::<BigEndian>(tagsize)?; + decoded_profile.write_all(b"bTRC")?; + decoded_profile.write_u32::<BigEndian>(tagstart)?; + decoded_profile.write_u32::<BigEndian>(tagsize)?; + } else if tagcode == 3 { + decoded_profile.write_all(b"gXYZ")?; + decoded_profile.write_u32::<BigEndian>(tagstart + tagsize)?; + decoded_profile.write_u32::<BigEndian>(tagsize)?; + decoded_profile.write_all(b"bXYZ")?; + decoded_profile.write_u32::<BigEndian>(tagstart + tagsize * 2)?; + decoded_profile.write_u32::<BigEndian>(tagsize)?; + } + Ok(()) + })(); + write_result.map_err(|_| Error::InvalidIccStream)?; + } + + Ok(()) +} + +fn shuffle_w2(bytes: &[u8]) -> Result<Vec<u8>> { + let len = bytes.len(); + let mut out = Vec::new_with_capacity(len)?; + + let height = len / 2; + let odd = len % 2; + for idx in 0..height { + out.push(bytes[idx]); + out.push(bytes[idx + height + odd]); + } + if odd != 0 { + out.push(bytes[height]); + } + Ok(out) +} + +fn shuffle_w4(bytes: &[u8]) -> Result<Vec<u8>> { + let len = bytes.len(); + let mut out = Vec::new_with_capacity(len)?; + + let step = len / 4; + let wide_count = len % 4; + for idx in 0..step { + let mut base = idx; + for _ in 0..wide_count { + out.push(bytes[base]); + base += step + 1; + } + for _ in wide_count..4 { + out.push(bytes[base]); + base += step; + } + } + for idx in 1..=wide_count { + out.push(bytes[(step + 1) * idx - 1]); + } + Ok(out) +} + +pub(super) fn read_single_command( + data_stream: &mut IccStream, + commands_stream: &mut Cursor<Vec<u8>>, + decoded_profile: &mut Cursor<&mut [u8]>, + command: u8, +) -> Result<()> { + use std::num::Wrapping; + + match command { + 1 => { + let num = read_varint_from_reader(commands_stream)? as usize; + data_stream.copy_bytes(decoded_profile, num)?; + } + 2 | 3 => { + let num = read_varint_from_reader(commands_stream)? as usize; + let bytes = data_stream.read_to_vec_exact(num)?; + let bytes = if command == 2 { + shuffle_w2(&bytes)? + } else { + shuffle_w4(&bytes)? + }; + decoded_profile + .write_all(&bytes) + .map_err(|_| Error::InvalidIccStream)?; + } + 4 => { + let flags = commands_stream + .read_u8() + .map_err(|_| Error::InvalidIccStream)?; + let width = ((flags & 3) + 1) as usize; + let order = (flags >> 2) & 3; + if width == 3 || order == 3 { + return Err(Error::InvalidIccStream); + } + + let stride = if (flags & 16) == 0 { + width + } else { + let stride = read_varint_from_reader(commands_stream)? as usize; + if stride < width { + return Err(Error::InvalidIccStream); + } + stride + }; + if stride.saturating_mul(4) >= decoded_profile.position() as usize { + return Err(Error::InvalidIccStream); + } + + let num = read_varint_from_reader(commands_stream)? as usize; + let bytes = data_stream.read_to_vec_exact(num)?; + let bytes = match width { + 1 => bytes, + 2 => shuffle_w2(&bytes)?, + 4 => shuffle_w4(&bytes)?, + _ => unreachable!(), + }; + + for i in (0..num).step_by(width) { + let base_position = decoded_profile.position() as usize; + let buffer_slice = decoded_profile.get_ref(); + + let mut prev = [Wrapping(0u32); 3]; + for (j, p) in prev[..=order as usize].iter_mut().enumerate() { + let offset = base_position - (stride * (j + 1)); + let mut p_bytes = [0u8; 4]; + p_bytes[(4 - width)..].copy_from_slice(&buffer_slice[offset..offset + width]); + p.0 = u32::from_be_bytes(p_bytes); + } + + let p = match order { + 0 => prev[0], + 1 => Wrapping(2) * prev[0] - prev[1], + 2 => Wrapping(3) * (prev[0] - prev[1]) + prev[2], + _ => unreachable!(), + }; + + decoded_profile.set_position(base_position as u64); + for j in 0..width.min(num - i) { + let val = Wrapping(bytes[i + j] as u32) + (p >> (8 * (width - 1 - j))); + decoded_profile + .write_u8(val.0 as u8) + .map_err(|_| Error::InvalidIccStream)?; + } + } + } + 10 => { + let mut bytes = [0u8; 20]; + bytes[..4].copy_from_slice(b"XYZ "); + data_stream.read_exact(&mut bytes[8..])?; + decoded_profile + .write_all(&bytes) + .map_err(|_| Error::InvalidIccStream)?; + } + 16..=23 => { + decoded_profile + .write_all(COMMON_DATA[command as usize - 16]) + .and_then(|_| decoded_profile.write_all(&[0u8; 4])) + .map_err(|_| Error::InvalidIccStream)?; + } + _ => { + return Err(Error::InvalidIccStream); + } + } + + Ok(()) +} diff --git a/third_party/rust/jxl/src/image.rs b/third_party/rust/jxl/src/image.rs @@ -0,0 +1,610 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +use std::{fmt::Debug, ops::Mul}; + +use crate::{ + error::{Error, Result}, + util::{ShiftRightCeil, tracing_wrappers::*}, +}; + +mod private { + pub trait Sealed {} +} + +#[derive(Copy, Clone, PartialEq, Eq, Debug)] +pub enum DataTypeTag { + U8, + U16, + U32, + F32, + I8, + I16, + I32, + F16, + F64, +} + +pub trait ImageDataType: + private::Sealed + Copy + Default + 'static + Debug + PartialEq + Mul<Self, Output = Self> +{ + /// ID of this data type. Different types *must* have different values. + const DATA_TYPE_ID: DataTypeTag; + + fn from_f64(f: f64) -> Self; + fn to_f64(self) -> f64; + #[cfg(test)] + fn random<R: rand::Rng>(rng: &mut R) -> Self; +} + +#[cfg(test)] +macro_rules! type_min { + (f32) => { + 0.0f32 + }; + ($ty: ty) => { + <$ty>::MIN + }; +} + +#[cfg(test)] +macro_rules! type_max { + (f32) => { + 1.0f32 + }; + ($ty: ty) => { + <$ty>::MAX + }; +} + +macro_rules! impl_image_data_type { + ($ty: ident, $id: ident) => { + impl private::Sealed for $ty {} + impl ImageDataType for $ty { + const DATA_TYPE_ID: DataTypeTag = DataTypeTag::$id; + fn from_f64(f: f64) -> $ty { + f as $ty + } + fn to_f64(self) -> f64 { + self as f64 + } + #[cfg(test)] + fn random<R: rand::Rng>(rng: &mut R) -> Self { + use rand::distributions::{Distribution, Uniform}; + let min = type_min!($ty); + let max = type_max!($ty); + Uniform::new_inclusive(min, max).sample(rng) + } + } + }; +} + +impl_image_data_type!(u8, U8); +impl_image_data_type!(u16, U16); +impl_image_data_type!(u32, U32); +impl_image_data_type!(f32, F32); +impl_image_data_type!(i8, I8); +impl_image_data_type!(i16, I16); +impl_image_data_type!(i32, I32); + +impl private::Sealed for half::f16 {} +impl ImageDataType for half::f16 { + const DATA_TYPE_ID: DataTypeTag = DataTypeTag::F16; + fn from_f64(f: f64) -> half::f16 { + half::f16::from_f64(f) + } + fn to_f64(self) -> f64 { + <half::f16>::to_f64(self) + } + #[cfg(test)] + fn random<R: rand::Rng>(rng: &mut R) -> Self { + use rand::distributions::{Distribution, Uniform}; + Self::from_f64(Uniform::new(0.0f32, 1.0f32).sample(rng) as f64) + } +} + +// Meant to be used by the simple render pipeline and in general for +// testing purposes. +impl_image_data_type!(f64, F64); + +#[derive(Clone)] +pub struct Image<T: ImageDataType> { + // width, height + size: (usize, usize), + data: Vec<T>, +} + +impl<T: ImageDataType> Debug for Image<T> { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{:?} {}x{}", T::DATA_TYPE_ID, self.size.0, self.size.1,) + } +} + +#[derive(Clone, Copy, Debug)] +pub struct Rect { + pub origin: (usize, usize), + // width, height + pub size: (usize, usize), +} + +impl Rect { + pub fn is_within(&self, size: (usize, usize)) -> Result<()> { + if self + .origin + .0 + .checked_add(self.size.0) + .ok_or(Error::ArithmeticOverflow)? + > size.0 + || self + .origin + .1 + .checked_add(self.size.1) + .ok_or(Error::ArithmeticOverflow)? + > size.1 + { + Err(Error::RectOutOfBounds( + self.size.0, + self.size.1, + self.origin.0, + self.origin.1, + size.0, + size.1, + )) + } else { + Ok(()) + } + } +} + +#[derive(Clone, Copy)] +pub struct ImageRect<'a, T: ImageDataType> { + pub rect: Rect, + image: &'a Image<T>, +} + +pub struct ImageRectMut<'a, T: ImageDataType> { + rect: Rect, + image: &'a mut Image<T>, +} + +impl<T: ImageDataType> Debug for ImageRect<'_, T> { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!( + f, + "{:?} {}x{}+{}+{}", + T::DATA_TYPE_ID, + self.rect.size.0, + self.rect.size.1, + self.rect.origin.0, + self.rect.origin.1 + ) + } +} + +impl<T: ImageDataType> Debug for ImageRectMut<'_, T> { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!( + f, + "mut {:?} {}x{}+{}+{}", + T::DATA_TYPE_ID, + self.rect.size.0, + self.rect.size.1, + self.rect.origin.0, + self.rect.origin.1 + ) + } +} + +impl<T: ImageDataType> Image<T> { + #[instrument(err)] + pub fn new_with_default(size: (usize, usize), default: T) -> Result<Image<T>> { + let (xsize, ysize) = size; + // These limits let us not worry about overflows. + if xsize as u64 >= i64::MAX as u64 / 4 || ysize as u64 >= i64::MAX as u64 / 4 { + return Err(Error::ImageSizeTooLarge(xsize, ysize)); + } + let total_size = xsize + .checked_mul(ysize) + .ok_or(Error::ImageSizeTooLarge(xsize, ysize))?; + // To simplify modular transform logic, we allow empty images, because some modular + // meta-images can have 0 xsize or ysize (e.g. delta-palette, reference property image). + debug!("trying to allocate image"); + let mut data = vec![]; + data.try_reserve_exact(total_size)?; + data.resize(total_size, default); + Ok(Image { + size: (xsize, ysize), + data, + }) + } + + pub fn new(size: (usize, usize)) -> Result<Image<T>> { + Self::new_with_default(size, T::default()) + } + + #[cfg(test)] + pub fn new_with_data(size: (usize, usize), data: Vec<T>) -> Image<T> { + Image { size, data } + } + + #[cfg(test)] + pub fn new_random<R: rand::Rng>(size: (usize, usize), rng: &mut R) -> Result<Image<T>> { + let mut img = Self::new(size)?; + img.data.iter_mut().for_each(|x| *x = T::random(rng)); + Ok(img) + } + + #[cfg(test)] + pub fn new_range(size: (usize, usize), start: f32, step: f32) -> Result<Image<T>> { + let mut img = Self::new(size)?; + img.data + .iter_mut() + .enumerate() + .for_each(|(index, x)| *x = T::from_f64((start + step * index as f32) as f64)); + Ok(img) + } + + #[cfg(test)] + pub fn new_constant(size: (usize, usize), val: T) -> Result<Image<T>> { + let mut img = Self::new(size)?; + img.data.iter_mut().for_each(|x| *x = val); + Ok(img) + } + + pub fn size(&self) -> (usize, usize) { + self.size + } + + pub fn fill(&mut self, v: T) { + let size = self.size; + let mut rect = self.as_rect_mut(); + for y in 0..size.1 { + for x in 0..size.0 { + rect.row(y)[x] = v; + } + } + } + + pub fn group_rect(&self, group_id: usize, log_group_size: (usize, usize)) -> ImageRect<'_, T> { + let xgroups = self.size.0.shrc(log_group_size.0); + let group = (group_id % xgroups, group_id / xgroups); + let origin = (group.0 << log_group_size.0, group.1 << log_group_size.1); + let size = ( + (self.size.0 - origin.0).min(1 << log_group_size.0), + (self.size.1 - origin.1).min(1 << log_group_size.1), + ); + trace!( + "making rect {}x{}+{}+{} for group {group_id} in image of size {:?}, log group sizes {:?}", + size.0, size.1, origin.0, origin.1, self.size, log_group_size + ); + self.as_rect().rect(Rect { origin, size }).unwrap() + } + + pub fn group_rect_mut( + &mut self, + group_id: usize, + log_group_size: usize, + ) -> ImageRectMut<'_, T> { + let xgroups = self.size.0.shrc(log_group_size); + let group = (group_id % xgroups, group_id / xgroups); + let origin = (group.0 << log_group_size, group.1 << log_group_size); + let size = ( + (self.size.0 - origin.0).min(1 << log_group_size), + (self.size.1 - origin.1).min(1 << log_group_size), + ); + self.as_rect_mut().into_rect(Rect { origin, size }).unwrap() + } + + pub fn as_rect(&self) -> ImageRect<'_, T> { + ImageRect { + rect: Rect { + origin: (0, 0), + size: self.size, + }, + image: self, + } + } + + pub fn as_rect_mut(&mut self) -> ImageRectMut<'_, T> { + ImageRectMut { + rect: Rect { + origin: (0, 0), + size: self.size, + }, + image: self, + } + } + + pub fn try_clone(&self) -> Result<Self> { + self.as_rect().to_image() + } +} + +impl<'a, T: ImageDataType> ImageRect<'a, T> { + pub fn rect(self, rect: Rect) -> Result<ImageRect<'a, T>> { + rect.is_within(self.rect.size)?; + Ok(ImageRect { + rect: Rect { + origin: ( + rect.origin.0 + self.rect.origin.0, + rect.origin.1 + self.rect.origin.1, + ), + size: rect.size, + }, + image: self.image, + }) + } + + pub fn size(&self) -> (usize, usize) { + self.rect.size + } + + pub fn row(&self, row: usize) -> &'a [T] { + debug_assert!(row < self.rect.size.1); + let start = (row + self.rect.origin.1) * self.image.size.0 + self.rect.origin.0; + trace!( + "{self:?} img size {:?} rect size {:?} row {row} start {}", + self.image.size, self.rect.size, start + ); + &self.image.data[start..start + self.rect.size.0] + } + + pub fn to_image(&self) -> Result<Image<T>> { + let total_size = self.rect.size.0 * self.rect.size.1; + let mut data = vec![]; + data.try_reserve_exact(total_size)?; + + for row_idx in 0..self.rect.size.1 { + data.extend_from_slice(self.row(row_idx)); + } + + Ok(Image { + size: self.rect.size, + data, + }) + } + + pub fn iter(&self) -> impl Iterator<Item = T> + '_ { + (0..self.rect.size.1).flat_map(|x| self.row(x).iter().cloned()) + } + + #[cfg(test)] + pub fn check_equal(&self, other: ImageRect<T>) { + assert_eq!(self.rect.size, other.rect.size); + let mismatch_info = |x: usize, y: usize| -> String { + let mut msg = format!( + "mismatch at position {x}x{y}, values {:?} and {:?}", + self.row(y)[x], + other.row(y)[x] + ); + if self.rect.origin != (0, 0) { + msg = format!( + "{}; position in ground truth {}x{}", + msg, + x + self.rect.origin.0, + y + self.rect.origin.1 + ); + } + if other.rect.origin != (0, 0) { + msg = format!( + "{}; position in checked img {}x{}", + msg, + x + other.rect.origin.0, + y + other.rect.origin.1 + ); + } + msg + }; + for y in 0..self.rect.size.1 { + for x in 0..self.rect.size.0 { + assert_eq!(self.row(y)[x], other.row(y)[x], "{}", mismatch_info(x, y)); + } + } + } +} + +impl<'a, T: ImageDataType> PartialEq<ImageRect<'a, T>> for ImageRect<'a, T> { + fn eq(&self, other: &ImageRect<'a, T>) -> bool { + self.iter().zip(other.iter()).all(|(x, y)| x == y) + } +} + +impl<T: ImageDataType + Eq> Eq for ImageRect<'_, T> {} + +impl<'a, T: ImageDataType> ImageRectMut<'a, T> { + pub fn rect(&'a mut self, rect: Rect) -> Result<ImageRectMut<'a, T>> { + rect.is_within(self.rect.size)?; + Ok(ImageRectMut { + rect: Rect { + origin: ( + rect.origin.0 + self.rect.origin.0, + rect.origin.1 + self.rect.origin.1, + ), + size: rect.size, + }, + image: self.image, + }) + } + + pub fn into_rect(self, rect: Rect) -> Result<ImageRectMut<'a, T>> { + rect.is_within(self.rect.size)?; + Ok(ImageRectMut { + rect: Rect { + origin: ( + rect.origin.0 + self.rect.origin.0, + rect.origin.1 + self.rect.origin.1, + ), + size: rect.size, + }, + image: self.image, + }) + } + + pub fn size(&self) -> (usize, usize) { + self.rect.size + } + + #[instrument(skip_all)] + pub fn copy_from(&mut self, other: ImageRect<'_, T>) -> Result<()> { + if other.rect.size != self.rect.size { + return Err(Error::CopyOfDifferentSize( + other.rect.size.0, + other.rect.size.1, + self.rect.size.0, + self.rect.size.1, + )); + } + + for i in 0..self.rect.size.1 { + trace!("copying row {i} of {}", self.rect.size.1); + self.row(i).copy_from_slice(other.row(i)); + } + + Ok(()) + } + + fn row_offset(&self, row: usize) -> usize { + debug_assert!(row < self.rect.size.1); + (row + self.rect.origin.1) * self.image.size.0 + self.rect.origin.0 + } + + pub fn row(&mut self, row: usize) -> &mut [T] { + debug_assert!(row < self.rect.size.1); + let start = self.row_offset(row); + trace!( + "{self:?} img size {:?} row {row} start {}", + self.image.size, start + ); + &mut self.image.data[start..start + self.rect.size.0] + } + + pub fn as_rect(&'a self) -> ImageRect<'a, T> { + ImageRect { + rect: self.rect, + image: self.image, + } + } + + /// Applies `f` to all the pixels in this rect. As side information, `f` is passed the + /// coordinates of the pixel in the full image. + pub fn apply<F>(&mut self, mut f: F) + where + F: for<'b> FnMut((usize, usize), &'b mut T), + { + let origin = self.rect.origin; + (0..self.rect.size.1).for_each(|x| { + self.row(x) + .iter_mut() + .enumerate() + .for_each(|(y, v)| f((origin.0 + x, origin.1 + y), v)) + }); + } +} + +#[cfg(test)] +mod test { + use arbtest::arbitrary::Arbitrary; + + use crate::error::Result; + + use super::{Image, ImageDataType, Rect}; + + #[test] + fn huge_image() { + assert!(Image::<u8>::new((1 << 28, 1 << 28)).is_err()); + } + + #[test] + fn rect_basic() -> Result<()> { + let mut image = Image::<u8>::new((32, 42))?; + assert_eq!( + image + .as_rect_mut() + .rect(Rect { + origin: (31, 40), + size: (1, 1) + })? + .size(), + (1, 1) + ); + assert_eq!( + image + .as_rect_mut() + .rect(Rect { + origin: (0, 0), + size: (1, 1) + })? + .size(), + (1, 1) + ); + assert!( + image + .as_rect_mut() + .rect(Rect { + origin: (30, 30), + size: (3, 3) + }) + .is_err() + ); + image + .as_rect_mut() + .rect(Rect { + origin: (30, 30), + size: (1, 1), + })? + .row(0)[0] = 1; + assert_eq!(image.as_rect_mut().row(30)[30], 1); + Ok(()) + } + + fn f64_conversions<T: ImageDataType + Eq + for<'a> Arbitrary<'a>>() { + arbtest::arbtest(|u| { + let t = T::arbitrary(u)?; + assert_eq!(t, T::from_f64(t.to_f64())); + Ok(()) + }); + } + + #[test] + fn u8_f64_conv() { + f64_conversions::<u8>(); + } + + #[test] + fn u16_f64_conv() { + f64_conversions::<u16>(); + } + + #[test] + fn u32_f64_conv() { + f64_conversions::<u32>(); + } + + #[test] + fn i8_f64_conv() { + f64_conversions::<i8>(); + } + + #[test] + fn i16_f64_conv() { + f64_conversions::<i16>(); + } + + #[test] + fn i32_f64_conv() { + f64_conversions::<i32>(); + } + + #[test] + fn f32_f64_conv() { + arbtest::arbtest(|u| { + let t = f32::arbitrary(u)?; + if !t.is_nan() { + assert_eq!(t, f32::from_f64(t.to_f64())); + } + Ok(()) + }); + } +} diff --git a/third_party/rust/jxl/src/lib.rs b/third_party/rust/jxl/src/lib.rs @@ -0,0 +1,29 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +#![deny(unsafe_code)] +pub mod api; +pub mod bit_reader; +pub mod color; +pub mod container; +pub mod entropy_coding; +pub mod error; +pub mod features; +pub mod frame; +pub mod headers; +pub mod icc; +pub mod image; +pub mod render; +mod simd; +pub mod util; +pub mod var_dct; + +// TODO: Move these to a more appropriate location. +const GROUP_DIM: usize = 256; +const BLOCK_DIM: usize = 8; +const BLOCK_SIZE: usize = BLOCK_DIM * BLOCK_DIM; +const SIGMA_PADDING: usize = 2; +#[allow(clippy::excessive_precision)] +const MIN_SIGMA: f32 = -3.90524291751269967465540850526868; diff --git a/third_party/rust/jxl/src/render/internal.rs b/third_party/rust/jxl/src/render/internal.rs @@ -0,0 +1,62 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +use crate::image::{DataTypeTag, ImageDataType}; + +use super::{RenderPipelineExtendStage, RenderPipelineInOutStage, RenderPipelineInPlaceStage}; + +#[derive(Copy, Clone, Debug, PartialEq, Eq)] +pub enum RenderPipelineStageType { + InPlace, + InOut, + Extend, +} + +pub trait RenderPipelineStageInfo: super::simple_pipeline::RenderPipelineRunStage { + const TYPE: RenderPipelineStageType; + const BORDER: (u8, u8); + const SHIFT: (u8, u8); + const INPUT_TYPE: DataTypeTag; + const OUTPUT_TYPE: Option<DataTypeTag>; + type RowType<'a>; +} + +impl<T: ImageDataType> RenderPipelineStageInfo for RenderPipelineInPlaceStage<T> { + const TYPE: RenderPipelineStageType = RenderPipelineStageType::InPlace; + const BORDER: (u8, u8) = (0, 0); + const SHIFT: (u8, u8) = (0, 0); + const INPUT_TYPE: DataTypeTag = T::DATA_TYPE_ID; + const OUTPUT_TYPE: Option<DataTypeTag> = None; + type RowType<'a> = &'a mut [T]; +} + +pub type InOutChannel<'a, InputT, OutputT> = (&'a [&'a [InputT]], &'a mut [&'a mut [OutputT]]); + +impl< + InputT: ImageDataType, + OutputT: ImageDataType, + const BORDER_X: u8, + const BORDER_Y: u8, + const SHIFT_X: u8, + const SHIFT_Y: u8, +> RenderPipelineStageInfo + for RenderPipelineInOutStage<InputT, OutputT, BORDER_X, BORDER_Y, SHIFT_X, SHIFT_Y> +{ + const TYPE: RenderPipelineStageType = RenderPipelineStageType::InOut; + const BORDER: (u8, u8) = (BORDER_X, BORDER_Y); + const SHIFT: (u8, u8) = (SHIFT_X, SHIFT_Y); + const INPUT_TYPE: DataTypeTag = InputT::DATA_TYPE_ID; + const OUTPUT_TYPE: Option<DataTypeTag> = Some(OutputT::DATA_TYPE_ID); + type RowType<'a> = InOutChannel<'a, InputT, OutputT>; +} + +impl<T: ImageDataType> RenderPipelineStageInfo for RenderPipelineExtendStage<T> { + const TYPE: RenderPipelineStageType = RenderPipelineStageType::Extend; + const BORDER: (u8, u8) = (0, 0); + const SHIFT: (u8, u8) = (0, 0); + const INPUT_TYPE: DataTypeTag = T::DATA_TYPE_ID; + const OUTPUT_TYPE: Option<DataTypeTag> = None; + type RowType<'a> = &'a mut [T]; +} diff --git a/third_party/rust/jxl/src/render/mod.rs b/third_party/rust/jxl/src/render/mod.rs @@ -0,0 +1,151 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +use internal::RenderPipelineStageInfo; +use std::{any::Any, marker::PhantomData}; + +use crate::{ + api::JxlOutputBuffer, + error::Result, + image::{Image, ImageDataType}, +}; + +mod internal; +mod simple_pipeline; +pub mod stages; +#[cfg(test)] +mod test; + +pub use simple_pipeline::{ + SimpleRenderPipeline, SimpleRenderPipelineBuilder, + save::{SaveStage, SaveStageType}, +}; + +/// Modifies channels in-place. +/// +/// For these stages, `process_row_chunk` receives read-write access to each row in each channel. +pub struct RenderPipelineInPlaceStage<T: ImageDataType> { + _phantom: PhantomData<T>, +} + +/// Modifies data and writes it to a new buffer, of possibly different type. +/// +/// BORDER_X and BORDER_Y represent the amount of padding required on the input side. +/// SHIFT_X and SHIFT_Y represent the base 2 log of the number of rows/columns produced +/// for each row/column of input. +/// +/// For each channel: +/// - the input slice contains 1 + BORDER_Y * 2 slices, each of length +/// xsize + BORDER_X * 2, i.e. covering one input row and up to BORDER pixels of +/// padding on either side. +/// - the output slice contains 1 << SHIFT_Y slices, each of length xsize << SHIFT_X, the +/// corresponding output pixels. +/// +/// `process_row_chunk` is passed a pair of q(input, output)` slices. +pub struct RenderPipelineInOutStage< + InputT: ImageDataType, + OutputT: ImageDataType, + const BORDER_X: u8, + const BORDER_Y: u8, + const SHIFT_X: u8, + const SHIFT_Y: u8, +> { + _phantom: PhantomData<(InputT, OutputT)>, +} + +/// Does not directly modify the current image pixels, but extends the current image with +/// additional data. +/// +/// `uses_channel` must always return true, and stages of this type should override +/// `new_size` and `original_data_origin`. +/// `process_row_chunk` will be called with the *new* image coordinates, and will only be called +/// on row chunks outside of the original image data. +/// After stages of this type, no stage can have a non-0 SHIFT_X or SHIFT_Y. +pub struct RenderPipelineExtendStage<T: ImageDataType> { + _phantom: PhantomData<T>, +} + +// TODO(veluca): figure out how to modify the interface for concurrent usage. +pub trait RenderPipelineStage: Any + std::fmt::Display { + type Type: RenderPipelineStageInfo; + + /// Which channels are actually used by this stage. + /// Must always return `true` if `Self::Type` is `RenderPipelineExtendStage`. + fn uses_channel(&self, c: usize) -> bool; + + /// Process one chunk of row. The semantics of this function are detailed in the + /// documentation of the various types of stages. + fn process_row_chunk( + &self, + position: (usize, usize), + xsize: usize, + // one for each channel + row: &mut [<Self::Type as RenderPipelineStageInfo>::RowType<'_>], + state: Option<&mut dyn Any>, + ); + + /// Returns the new size of the image after this stage. Should be implemented by + /// `RenderPipelineExtendStage` stages. + fn new_size(&self, current_size: (usize, usize)) -> (usize, usize) { + current_size + } + + /// Returns the origin of the original image data in the output of this stage. + /// Should be implemented by `RenderPipelineExtendStage` stages. + fn original_data_origin(&self) -> (isize, isize) { + (0, 0) + } + + /// Initializes thread local state for the stage. Returns `Ok(None)` if no state is needed. + /// + /// This method returns `Box<dyn Any>` for dyn compatibility. `process_row_chunk` should + /// downcast the state to the desired type. + fn init_local_state(&self) -> Result<Option<Box<dyn Any>>> { + Ok(None) + } +} + +pub trait RenderPipelineBuilder: Sized { + type RenderPipeline: RenderPipeline; + fn new( + num_channels: usize, + size: (usize, usize), + downsampling_shift: usize, + log_group_size: usize, + ) -> Self; + fn add_stage<Stage: RenderPipelineStage>(self, stage: Stage) -> Result<Self>; + fn add_save_stage<T: ImageDataType>(self, stage: SaveStage<T>) -> Result<Self>; + fn build(self) -> Result<Self::RenderPipeline>; +} + +pub trait RenderPipeline { + type Builder: RenderPipelineBuilder<RenderPipeline = Self>; + + /// Obtains a buffer suitable for storing the input at channel `channel` of group `group_id`. + /// This *might* be a buffer that was used to store that channel for that group in a previous + /// pass, a new buffer, or a re-used buffer from i.e. previously decoded frames. + fn get_buffer_for_group<T: ImageDataType>( + &mut self, + channel: usize, + group_id: usize, + ) -> Result<Image<T>>; + + /// Gives back the buffer for a channel and group to the render pipeline, marking that + /// `num_passes` additional passes (wrt. the previous call to this method for the same channel + /// and group, or 0 if no previous call happend) were rendered into the input buffer. + fn set_buffer_for_group<T: ImageDataType>( + &mut self, + channel: usize, + group_id: usize, + num_passes: usize, + buf: Image<T>, + ); + + /// Renders new data that is available after the last call to `render`. + fn do_render(&mut self, buffers: &mut [Option<JxlOutputBuffer>]) -> Result<()>; + + fn into_stages(self) -> Vec<Box<dyn Any>>; + fn num_groups(&self) -> usize; +} diff --git a/third_party/rust/jxl/src/render/simple_pipeline/mod.rs b/third_party/rust/jxl/src/render/simple_pipeline/mod.rs @@ -0,0 +1,770 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +use std::any::Any; + +use save::SaveStage; + +use crate::{ + BLOCK_DIM, + error::{Error, Result}, + image::{DataTypeTag, Image, ImageDataType}, + render::internal::RenderPipelineStageInfo, + simd::round_up_size_to_two_cache_lines, + util::{ShiftRightCeil, tracing_wrappers::*}, +}; + +use super::{ + RenderPipeline, RenderPipelineBuilder, RenderPipelineExtendStage, RenderPipelineInOutStage, + RenderPipelineInPlaceStage, RenderPipelineStage, internal::RenderPipelineStageType, +}; + +// TODO(veluca): this should eventually become non-pub. +pub mod save; + +#[derive(Clone, Debug)] +struct ChannelInfo { + ty: Option<DataTypeTag>, + downsample: (u8, u8), +} + +pub struct SimpleRenderPipelineBuilder { + pipeline: SimpleRenderPipeline, + can_shift: bool, +} + +impl SimpleRenderPipelineBuilder { + #[instrument(level = "debug")] + pub(super) fn new_with_chunk_size( + num_channels: usize, + size: (usize, usize), + downsampling_shift: usize, + mut log_group_size: usize, + chunk_size: usize, + ) -> Self { + info!("creating simple pipeline"); + assert!(chunk_size <= u16::MAX as usize); + assert_ne!(chunk_size, 0); + // The number of pixels that a group encompasses in the final, upsampled image along one dimension is effectively multiplied by the upsampling factor. + log_group_size += downsampling_shift; + SimpleRenderPipelineBuilder { + pipeline: SimpleRenderPipeline { + channel_info: vec![vec![ + ChannelInfo { + ty: None, + downsample: (0, 0) + }; + num_channels + ]], + input_size: size, + log_group_size, + xgroups: size.0.shrc(log_group_size), + stages: vec![], + group_chan_ready_passes: vec![ + vec![0; num_channels]; + size.0.shrc(log_group_size) + * size.1.shrc(log_group_size) + ], + completed_passes: 0, + input_buffers: vec![], + chunk_size, + }, + can_shift: true, + } + } + + fn add_dyn_stage( + mut self, + stage: Box<dyn RunStage>, + input_type: DataTypeTag, + output_type: Option<DataTypeTag>, + shift: (u8, u8), + is_extend: bool, + ) -> Result<Self> { + let current_info = self.pipeline.channel_info.last().unwrap().clone(); + debug!( + last_stage_channel_info = ?current_info, + can_shift = self.can_shift, + "adding stage '{stage}'", + ); + let mut after_info = vec![]; + for (c, info) in current_info.iter().enumerate() { + if !stage.uses_channel(c) { + after_info.push(ChannelInfo { + ty: info.ty, + downsample: (0, 0), + }); + } else { + if let Some(ty) = info.ty + && ty != input_type + { + return Err(Error::PipelineChannelTypeMismatch( + stage.to_string(), + c, + input_type, + ty, + )); + } + after_info.push(ChannelInfo { + ty: Some(output_type.unwrap_or(input_type)), + downsample: shift, + }); + } + } + if !self.can_shift && shift != (0, 0) { + return Err(Error::PipelineShiftAfterExpand(stage.to_string())); + } + if is_extend { + self.can_shift = false; + } + debug!( + new_channel_info = ?after_info, + can_shift = self.can_shift, + "added stage '{stage}'", + ); + self.pipeline.channel_info.push(after_info); + self.pipeline.stages.push(stage); + Ok(self) + } +} + +impl RenderPipelineBuilder for SimpleRenderPipelineBuilder { + type RenderPipeline = SimpleRenderPipeline; + + fn new( + num_channels: usize, + size: (usize, usize), + downsampling_shift: usize, + log_group_size: usize, + ) -> Self { + // The 256 is an even number of cache lines, which makes round_up_size_to_two_cache_lines not round it up. + // This is fine since it's also an even number of SIMD lanes - and when using borders in stages we round up + // chunk_size + border * 2, and again get space for full SIMD lane loading. + // If the chunk size is less than an even number of SIMD lanes, then rounding up chunk_size + border * 2 might + // end up with an uneven number of SIMD lanes to read the full chunk. + // Example: chunk_size=3833, border=1, round_up_size_to_two_cache_lines(chunk_size + border * 2)=3840, + // reading the last border pixel will then read a lane outside the buffer. + Self::new_with_chunk_size(num_channels, size, downsampling_shift, log_group_size, 256) + } + + #[instrument(skip_all, err)] + fn add_stage<Stage: RenderPipelineStage>(self, stage: Stage) -> Result<Self> { + self.add_dyn_stage( + Box::new(stage), + Stage::Type::INPUT_TYPE, + Stage::Type::OUTPUT_TYPE, + Stage::Type::SHIFT, + Stage::Type::TYPE == RenderPipelineStageType::Extend, + ) + } + + #[instrument(skip_all, err)] + fn add_save_stage<T: ImageDataType>(self, stage: SaveStage<T>) -> Result<Self> { + self.add_dyn_stage(Box::new(stage), T::DATA_TYPE_ID, None, (0, 0), false) + } + + #[instrument(skip_all, err)] + fn build(mut self) -> Result<Self::RenderPipeline> { + let channel_info = &mut self.pipeline.channel_info; + let num_channels = channel_info[0].len(); + let mut cur_downsamples = vec![(0u8, 0u8); num_channels]; + for (s, stage) in self.pipeline.stages.iter().enumerate().rev() { + let [current_info, next_info, ..] = &mut channel_info[s..] else { + unreachable!() + }; + for chan in 0..num_channels { + let cur_chan = &mut current_info[chan]; + let next_chan = &mut next_info[chan]; + if stage.uses_channel(chan) { + assert_eq!(Some(stage.output_type()), next_chan.ty); + } + if cur_chan.ty.is_none() { + cur_chan.ty = if stage.uses_channel(chan) { + Some(stage.input_type()) + } else { + next_chan.ty + } + } + // Arithmetic overflows here should be very uncommon, so custom error variants + // are probably unwarranted. + let cur_downsample = &mut cur_downsamples[chan]; + let next_downsample = &mut next_chan.downsample; + let next_total_downsample = *cur_downsample; + cur_downsample.0 = cur_downsample + .0 + .checked_add(next_downsample.0) + .ok_or(Error::ArithmeticOverflow)?; + cur_downsample.1 = cur_downsample + .1 + .checked_add(next_downsample.1) + .ok_or(Error::ArithmeticOverflow)?; + *next_downsample = next_total_downsample; + } + } + for (chan, cur_downsample) in cur_downsamples.iter().enumerate() { + channel_info[0][chan].downsample = *cur_downsample; + } + #[cfg(feature = "tracing")] + { + for (s, (current_info, stage)) in channel_info + .iter() + .zip(self.pipeline.stages.iter()) + .enumerate() + { + debug!("final channel info before stage {s} '{stage}': {current_info:?}"); + } + debug!( + "final channel info after all stages {:?}", + channel_info.last().unwrap() + ); + } + + // Ensure all channels have been used, so that we know the types of all buffers at all + // stages. + for (c, chinfo) in channel_info.iter().flat_map(|x| x.iter().enumerate()) { + if chinfo.ty.is_none() { + return Err(Error::PipelineChannelUnused(c)); + } + } + + let input_buffers: Result<_> = channel_info[0] + .iter() + .map(|x| { + let xsize = self.pipeline.input_size.0.shrc(x.downsample.0); + let ysize = self.pipeline.input_size.1.shrc(x.downsample.1); + Image::new((xsize, ysize)) + }) + .collect(); + self.pipeline.input_buffers = input_buffers?; + + Ok(self.pipeline) + } +} + +/// A RenderPipeline that waits for all input of a pass to be ready before doing any rendering, and +/// prioritizes simplicity over memory usage and computational efficiency. +/// Eventually meant to be used only for verification purposes. +pub struct SimpleRenderPipeline { + channel_info: Vec<Vec<ChannelInfo>>, + input_size: (usize, usize), + log_group_size: usize, + xgroups: usize, + stages: Vec<Box<dyn RunStage>>, + group_chan_ready_passes: Vec<Vec<usize>>, + completed_passes: usize, + input_buffers: Vec<Image<f64>>, + chunk_size: usize, +} + +fn clone_images<T: ImageDataType>(images: &[Image<T>]) -> Result<Vec<Image<T>>> { + images.iter().map(|x| x.as_rect().to_image()).collect() +} + +impl SimpleRenderPipeline { + #[instrument(skip_all, err)] + fn do_render(&mut self) -> Result<()> { + let ready_passes = self + .group_chan_ready_passes + .iter() + .flat_map(|x| x.iter()) + .copied() + .min() + .unwrap(); + if ready_passes <= self.completed_passes { + debug!( + "no more ready passes ({} completed, {ready_passes} ready)", + self.completed_passes + ); + return Ok(()); + } + debug!( + "new ready passes ({} completed, {ready_passes} ready)", + self.completed_passes + ); + self.completed_passes = ready_passes; + + let mut current_buffers = clone_images(&self.input_buffers)?; + + let mut current_size = self.input_size; + + for (i, stage) in self.stages.iter().enumerate() { + debug!("running stage {i}: {stage}"); + let mut output_buffers = clone_images(¤t_buffers)?; + // Replace buffers of different sizes. + if stage.shift() != (0, 0) || stage.new_size(current_size) != current_size { + current_size = stage.new_size(current_size); + for (c, info) in self.channel_info[i + 1].iter().enumerate() { + if stage.uses_channel(c) { + let xsize = current_size.0.shrc(info.downsample.0); + let ysize = current_size.1.shrc(info.downsample.1); + debug!("reallocating channel {c} to new size {xsize}x{ysize}"); + output_buffers[c] = Image::new((xsize, ysize))?; + } + } + } + let input_buf: Vec<_> = current_buffers + .iter() + .enumerate() + .filter(|x| stage.uses_channel(x.0)) + .map(|x| x.1) + .collect(); + let mut output_buf: Vec<_> = output_buffers + .iter_mut() + .enumerate() + .filter(|x| stage.uses_channel(x.0)) + .map(|x| x.1) + .collect(); + let mut state = stage.init_local_state()?; + stage.run_stage_on( + self.chunk_size, + &input_buf, + &mut output_buf, + state.as_deref_mut(), + ); + current_buffers = output_buffers; + } + + Ok(()) + } + + fn group_position(&self, group_id: usize) -> (usize, usize) { + (group_id % self.xgroups, group_id / self.xgroups) + } + + fn group_offset(&self, group_id: usize) -> (usize, usize) { + let group = self.group_position(group_id); + ( + group.0 << self.log_group_size, + group.1 << self.log_group_size, + ) + } + + fn group_size(&self, group_id: usize) -> (usize, usize) { + let goffset = self.group_offset(group_id); + ( + self.input_size + .0 + .min(goffset.0 + (1 << self.log_group_size)) + - goffset.0, + self.input_size + .1 + .min(goffset.1 + (1 << self.log_group_size)) + - goffset.1, + ) + } + + fn group_size_for_channel( + &self, + channel: usize, + group_id: usize, + requested_data_type: DataTypeTag, + ) -> (usize, usize) { + let goffset = self.group_offset(group_id); + let ChannelInfo { downsample, ty } = self.channel_info[0][channel]; + if ty.unwrap() != requested_data_type { + panic!( + "Invalid pipeline usage: incorrect channel type, requested {:?}, but pipeline wants {ty:?}", + requested_data_type + ); + } + assert_eq!(goffset.0 % (1 << downsample.0), 0); + assert_eq!(goffset.1 % (1 << downsample.1), 0); + let group_size = self.group_size(group_id); + ( + group_size.0.shrc(downsample.0), + group_size.1.shrc(downsample.1), + ) + } +} + +impl RenderPipeline for SimpleRenderPipeline { + type Builder = SimpleRenderPipelineBuilder; + + #[instrument(skip_all, err)] + fn get_buffer_for_group<T: ImageDataType>( + &mut self, + channel: usize, + group_id: usize, + ) -> Result<Image<T>> { + let sz = self.group_size_for_channel(channel, group_id, T::DATA_TYPE_ID); + Image::<T>::new(sz) + } + + fn set_buffer_for_group<T: ImageDataType>( + &mut self, + channel: usize, + group_id: usize, + num_passes: usize, + buf: Image<T>, + ) { + debug!( + "filling data for group {}, channel {}, using type {:?}", + group_id, + channel, + T::DATA_TYPE_ID, + ); + let sz = self.group_size_for_channel(channel, group_id, T::DATA_TYPE_ID); + let goffset = self.group_offset(group_id); + let ChannelInfo { ty, downsample } = self.channel_info[0][channel]; + let off = (goffset.0 >> downsample.0, goffset.1 >> downsample.1); + debug!(?sz, input_buffers_sz=?self.input_buffers[channel].size(), offset=?off, ?downsample, ?goffset); + let bsz = buf.size(); + assert!(sz.0 <= bsz.0); + assert!(sz.1 <= bsz.1); + assert!(sz.0 + BLOCK_DIM > bsz.0); + assert!(sz.1 + BLOCK_DIM > bsz.1); + let ty = ty.unwrap(); + assert_eq!(ty, T::DATA_TYPE_ID); + for y in 0..sz.1 { + for x in 0..sz.0 { + self.input_buffers[channel].as_rect_mut().row(y + off.1)[x + off.0] = + buf.as_rect().row(y)[x].to_f64(); + } + } + self.group_chan_ready_passes[group_id][channel] += num_passes; + } + + // TODO(veluca): actually use the passed-in buffers. + fn do_render(&mut self, _: &mut [Option<crate::api::JxlOutputBuffer>]) -> Result<()> { + self.do_render() + } + + fn into_stages(self) -> Vec<Box<dyn std::any::Any>> { + self.stages.into_iter().map(|x| x.as_any()).collect() + } + + fn num_groups(&self) -> usize { + self.xgroups * self.input_size.1.shrc(self.log_group_size) + } +} + +pub trait RenderPipelineRunStage { + fn run_stage_on<S: RenderPipelineStage<Type = Self>>( + stage: &S, + chunk_size: usize, + input_buffers: &[&Image<f64>], + output_buffers: &mut [&mut Image<f64>], + state: Option<&mut dyn Any>, + ); +} + +impl<T: ImageDataType> RenderPipelineRunStage for RenderPipelineInPlaceStage<T> { + #[instrument(skip_all)] + fn run_stage_on<S: RenderPipelineStage<Type = Self>>( + stage: &S, + chunk_size: usize, + input_buffers: &[&Image<f64>], + output_buffers: &mut [&mut Image<f64>], + mut state: Option<&mut dyn Any>, + ) { + debug!("running inplace stage '{stage}' in simple pipeline"); + let numc = input_buffers.len(); + if numc == 0 { + return; + } + assert_eq!(output_buffers.len(), numc); + let size = input_buffers[0].size(); + for b in input_buffers.iter() { + assert_eq!(size, b.size()); + } + for b in output_buffers.iter() { + assert_eq!(size, b.size()); + } + let mut buffer = + vec![vec![T::default(); round_up_size_to_two_cache_lines::<T>(chunk_size)]; numc]; + for y in 0..size.1 { + for x in (0..size.0).step_by(chunk_size) { + let xsize = size.0.min(x + chunk_size) - x; + debug!("position: {x}x{y} xsize: {xsize}"); + for c in 0..numc { + let in_rect = input_buffers[c].as_rect(); + let in_row = in_rect.row(y); + for ix in 0..xsize { + buffer[c][ix] = T::from_f64(in_row[x + ix]); + } + } + let mut row: Vec<_> = buffer.iter_mut().map(|x| x as &mut [T]).collect(); + stage.process_row_chunk((x, y), xsize, &mut row, state.as_deref_mut()); + for c in 0..numc { + let mut out_rect = output_buffers[c].as_rect_mut(); + let out_row = out_rect.row(y); + for ix in 0..xsize { + out_row[x + ix] = buffer[c][ix].to_f64(); + } + } + } + } + } +} + +impl< + InputT: ImageDataType, + OutputT: ImageDataType, + const BORDER_X: u8, + const BORDER_Y: u8, + const SHIFT_X: u8, + const SHIFT_Y: u8, +> RenderPipelineRunStage + for RenderPipelineInOutStage<InputT, OutputT, BORDER_X, BORDER_Y, SHIFT_X, SHIFT_Y> +{ + #[instrument(skip_all)] + fn run_stage_on<S: RenderPipelineStage<Type = Self>>( + stage: &S, + chunk_size: usize, + input_buffers: &[&Image<f64>], + output_buffers: &mut [&mut Image<f64>], + mut state: Option<&mut dyn Any>, + ) { + assert_ne!(chunk_size, 0); + debug!("running inout stage '{stage}' in simple pipeline"); + let numc = input_buffers.len(); + if numc == 0 { + return; + } + assert_eq!(output_buffers.len(), numc); + let input_size = input_buffers[0].size(); + let output_size = output_buffers[0].size(); + for c in 1..numc { + assert_eq!(input_size, input_buffers[c].size()); + assert_eq!(output_size, output_buffers[c].size()); + } + debug!( + "input_size = {input_size:?} output_size = {output_size:?} SHIFT_X = {SHIFT_X} \ + SHIFT_Y = {SHIFT_Y} BORDER_X = {BORDER_X} BORDER_Y = {BORDER_Y} numc = {numc}" + ); + assert_eq!(input_size.0, output_size.0.div_ceil(1 << SHIFT_X)); + assert_eq!(input_size.1, output_size.1.div_ceil(1 << SHIFT_Y)); + let mut buffer_in = vec![ + vec![ + vec![ + InputT::default(); + // Double rounding make sure that we always have enough buffer for reading a whole SIMD lane. + round_up_size_to_two_cache_lines::<OutputT>( + round_up_size_to_two_cache_lines::<OutputT>(chunk_size) + + BORDER_X as usize * 2 + ) + ]; + BORDER_Y as usize * 2 + 1 + ]; + numc + ]; + let mut buffer_out = vec![ + vec![ + vec![ + OutputT::default(); + round_up_size_to_two_cache_lines::<OutputT>(chunk_size) + << SHIFT_X + ]; + 1 << SHIFT_Y + ]; + numc + ]; + + let mirror = |mut v: i64, size: i64| { + while v < 0 || v >= size { + if v < 0 { + v = -v - 1; + } + if v >= size { + v = size + (size - v) - 1; + } + } + v as usize + }; + for y in 0..input_size.1 { + for x in (0..input_size.0).step_by(chunk_size) { + let border_x = BORDER_X as i64; + let border_y = BORDER_Y as i64; + let xsize = input_size.0.min(x + chunk_size) - x; + let xs = xsize as i64; + debug!("position: {x}x{y} xsize: {xsize}"); + for c in 0..numc { + let in_rect = input_buffers[c].as_rect(); + for iy in -border_y..=border_y { + let imgy = mirror(y as i64 + iy, input_size.1 as i64); + let in_row = in_rect.row(imgy); + let buf_in_row = &mut buffer_in[c][(iy + border_y) as usize]; + for ix in (-border_x..0).chain(xs..xs + border_x) { + let imgx = mirror(x as i64 + ix, input_size.0 as i64); + buf_in_row[(ix + border_x) as usize] = InputT::from_f64(in_row[imgx]); + } + for ix in 0..xsize { + buf_in_row[ix + border_x as usize] = InputT::from_f64(in_row[x + ix]); + } + } + } + let buffer_in_ref: Vec<Vec<_>> = buffer_in + .iter() + .map(|x| x.iter().map(|x| x as &[_]).collect()) + .collect(); + let mut buffer_out_ref: Vec<Vec<_>> = buffer_out + .iter_mut() + .map(|x| x.iter_mut().map(|x| x as &mut [_]).collect::<Vec<_>>()) + .collect(); + let mut row: Vec<_> = buffer_in_ref + .iter() + .zip(buffer_out_ref.iter_mut()) + .map(|(itin, itout)| (itin as &[_], itout as &mut [_])) + .collect(); + stage.process_row_chunk((x, y), xsize, &mut row, state.as_deref_mut()); + let stripe_xsize = (xsize << SHIFT_X).min(output_size.0 - (x << SHIFT_X)); + let stripe_ysize = (1usize << SHIFT_Y).min(output_size.1 - (y << SHIFT_Y)); + for c in 0..numc { + let mut out_rect = output_buffers[c].as_rect_mut(); + for iy in 0..stripe_ysize { + let out_row = out_rect.row((y << SHIFT_Y) + iy); + for ix in 0..stripe_xsize { + out_row[(x << SHIFT_X) + ix] = buffer_out[c][iy][ix].to_f64(); + } + } + } + } + } + } +} + +impl<T: ImageDataType> RenderPipelineRunStage for RenderPipelineExtendStage<T> { + #[instrument(skip_all)] + fn run_stage_on<S: RenderPipelineStage<Type = Self>>( + stage: &S, + chunk_size: usize, + input_buffers: &[&Image<f64>], + output_buffers: &mut [&mut Image<f64>], + mut state: Option<&mut dyn Any>, + ) { + debug!("running extend stage '{stage}' in simple pipeline"); + let numc = input_buffers.len(); + if numc == 0 { + return; + } + assert_eq!(output_buffers.len(), numc); + let input_size = input_buffers[0].size(); + let output_size = output_buffers[0].size(); + for c in 1..numc { + assert_eq!(input_size, input_buffers[c].size()); + assert_eq!(output_size, output_buffers[c].size()); + } + assert_eq!(output_size, stage.new_size(input_size)); + let origin = stage.original_data_origin(); + assert!(origin.0 <= output_size.0 as isize); + assert!(origin.1 <= output_size.1 as isize); + assert!(origin.0 + input_size.0 as isize >= 0); + assert!(origin.1 + input_size.1 as isize >= 0); + let origin = stage.original_data_origin(); + debug!("input_size = {input_size:?} output_size = {output_size:?} origin = {origin:?}"); + // Compute the input rectangle + let x0 = origin.0.max(0) as usize; + let y0 = origin.1.max(0) as usize; + let x1 = (origin.0 + input_size.0 as isize).min(output_size.0 as isize) as usize; + let y1 = (origin.1 + input_size.1 as isize).min(output_size.1 as isize) as usize; + debug!("x0 = {x0} x1 = {x1} y0 = {y0} y1 = {y1}"); + let in_x0 = (x0 as isize - origin.0) as usize; + let in_x1 = (x1 as isize - origin.0) as usize; + let in_y0 = (y0 as isize - origin.1) as usize; + let in_y1 = (y1 as isize - origin.1) as usize; + debug!("in_x0 = {in_x0} in_x1 = {in_x1} in_y0 = {in_y0} in_y1 = {in_y1}"); + // First, copy the data in the middle. + for c in 0..numc { + for in_y in in_y0..in_y1 { + debug!("copy row: {in_y}"); + let in_row = input_buffers[c].as_rect().row(in_y); + let y = (in_y as isize + origin.1) as usize; + output_buffers[c].as_rect_mut().row(y)[x0..x1] + .copy_from_slice(&in_row[in_x0..in_x1]); + } + } + // Fill in rows above and below the original data. + let mut buffer = + vec![vec![T::default(); round_up_size_to_two_cache_lines::<T>(chunk_size)]; numc]; + for y in (0..y0).chain(y1..output_size.1) { + for x in (0..output_size.0).step_by(chunk_size) { + let xsize = output_size.0.min(x + chunk_size) - x; + debug!("position above/below: ({x},{y}) xsize: {xsize}"); + let mut row: Vec<_> = buffer.iter_mut().map(|x| x as &mut [T]).collect(); + stage.process_row_chunk((x, y), xsize, &mut row, state.as_deref_mut()); + for c in 0..numc { + for ix in 0..xsize { + output_buffers[c].as_rect_mut().row(y)[x + ix] = buffer[c][ix].to_f64(); + } + } + } + } + // Fill in left and right of the original data. + for y in y0..y1 { + for (x, xsize) in (0..x0) + .step_by(chunk_size) + .map(|x| (x, x0.min(x + chunk_size) - x)) + .chain( + (x1..output_size.0) + .step_by(chunk_size) + .map(|x| (x, output_size.0.min(x + chunk_size) - x)), + ) + { + let mut row: Vec<_> = buffer.iter_mut().map(|x| x as &mut [T]).collect(); + debug!("position on the side: ({x},{y}) xsize: {xsize}"); + stage.process_row_chunk((x, y), xsize, &mut row, state.as_deref_mut()); + for c in 0..numc { + for ix in 0..xsize { + output_buffers[c].as_rect_mut().row(y)[x + ix] = buffer[c][ix].to_f64(); + } + } + } + } + } +} + +trait RunStage: Any + std::fmt::Display { + fn run_stage_on( + &self, + chunk_size: usize, + input_buffers: &[&Image<f64>], + output_buffers: &mut [&mut Image<f64>], + state: Option<&mut dyn Any>, + ); + fn init_local_state(&self) -> Result<Option<Box<dyn Any>>>; + fn shift(&self) -> (u8, u8); + fn new_size(&self, size: (usize, usize)) -> (usize, usize); + fn uses_channel(&self, c: usize) -> bool; + fn as_any(self: Box<Self>) -> Box<dyn Any>; + fn input_type(&self) -> DataTypeTag; + fn output_type(&self) -> DataTypeTag; +} + +impl<T: RenderPipelineStage> RunStage for T { + fn run_stage_on( + &self, + chunk_size: usize, + input_buffers: &[&Image<f64>], + output_buffers: &mut [&mut Image<f64>], + state: Option<&mut dyn Any>, + ) { + T::Type::run_stage_on(self, chunk_size, input_buffers, output_buffers, state) + } + + fn init_local_state(&self) -> Result<Option<Box<dyn Any>>> { + T::init_local_state(self) + } + + fn shift(&self) -> (u8, u8) { + T::Type::SHIFT + } + + fn new_size(&self, size: (usize, usize)) -> (usize, usize) { + self.new_size(size) + } + + fn uses_channel(&self, c: usize) -> bool { + self.uses_channel(c) + } + fn as_any(self: Box<Self>) -> Box<dyn Any> { + self + } + fn input_type(&self) -> DataTypeTag { + T::Type::INPUT_TYPE + } + fn output_type(&self) -> DataTypeTag { + T::Type::OUTPUT_TYPE.unwrap_or(T::Type::INPUT_TYPE) + } +} diff --git a/third_party/rust/jxl/src/render/simple_pipeline/save.rs b/third_party/rust/jxl/src/render/simple_pipeline/save.rs @@ -0,0 +1,376 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +use std::{any::Any, sync::Mutex}; + +use crate::{ + error::Result, + headers::Orientation, + image::{DataTypeTag, Image, ImageDataType, Rect}, + simd::round_up_size_to_two_cache_lines, + util::tracing_wrappers::debug, +}; + +use super::RunStage; + +pub enum SaveStageType { + Output, + Reference, + Lf, +} + +pub struct SaveStage<T: ImageDataType> { + pub stage_type: SaveStageType, + buf: Mutex<Image<T>>, + channel: usize, + // TODO(szabadka): Have a fixed scale per data-type and make the datatype conversions do + // the scaling. + scale: T, + orientation: Orientation, +} + +#[allow(unused)] +impl<T: ImageDataType> SaveStage<T> { + pub(crate) fn new( + stage_type: SaveStageType, + channel: usize, + size: (usize, usize), + scale: T, + orientation: Orientation, + ) -> Result<SaveStage<T>> { + let buf_size = if orientation.is_transposing() { + (size.1, size.0) + } else { + size + }; + Ok(SaveStage { + stage_type, + channel, + buf: Mutex::new(Image::new(buf_size)?), + scale, + orientation, + }) + } + + pub(crate) fn new_with_buffer( + stage_type: SaveStageType, + channel: usize, + img: Image<T>, + scale: T, + orientation: Orientation, + ) -> SaveStage<T> { + SaveStage { + stage_type, + channel, + buf: Mutex::new(img), + scale, + orientation, + } + } + + pub(crate) fn buffer(&self) -> impl std::ops::Deref<Target = Image<T>> { + self.buf.lock().unwrap() + } + + pub(crate) fn into_buffer(self) -> Image<T> { + self.buf.into_inner().unwrap() + } +} + +impl<T: ImageDataType> std::fmt::Display for SaveStage<T> { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!( + f, + "save channel {} (type {:?}) scale {:?}", + self.channel, + T::DATA_TYPE_ID, + self.scale, + ) + } +} + +impl<T: ImageDataType> SaveStage<T> { + fn process_row_chunk( + &self, + position: (usize, usize), + mut xsize: usize, + row: &mut [&[T]], + _state: Option<&mut dyn std::any::Any>, + ) { + let input = row[0]; + let mut buf = self.buf.lock().unwrap(); + let mut outbuf_rect = buf.as_rect_mut(); + let (out_w, out_h) = outbuf_rect.size(); + + // Establish source dimensions based on the orientation. + let (w_src, h_src) = if self.orientation.is_transposing() { + (out_h, out_w) // Swapped + } else { + (out_w, out_h) + }; + + // Perform boundary checks against source dimensions. + if position.1 >= h_src { + return; + } + xsize = xsize.min(w_src - position.0); + if xsize == 0 { + return; + } + + match self.orientation { + // non-transposing cases + Orientation::Identity => { + let mut out_row = outbuf_rect + .rect(Rect { + origin: position, + size: (xsize, 1), + }) + .unwrap(); + for (out_pixel, &in_pixel) in out_row.row(0).iter_mut().zip(input.iter()) { + *out_pixel = in_pixel * self.scale; + } + } + Orientation::FlipHorizontal => { + let mut out_row = outbuf_rect + .rect(Rect { + origin: (0, position.1), + size: (out_w, 1), + }) + .unwrap(); + for (i, &in_pixel) in input.iter().enumerate().take(xsize) { + out_row.row(0)[out_w - 1 - position.0 - i] = in_pixel * self.scale; + } + } + Orientation::FlipVertical => { + let y_out = out_h - 1 - position.1; + let mut out_row = outbuf_rect + .rect(Rect { + origin: (position.0, y_out), + size: (xsize, 1), + }) + .unwrap(); + for (out_pixel, &in_pixel) in out_row.row(0).iter_mut().zip(input.iter()) { + *out_pixel = in_pixel * self.scale; + } + } + Orientation::Rotate180 => { + let y_out = out_h - 1 - position.1; + let mut out_row = outbuf_rect + .rect(Rect { + origin: (0, y_out), + size: (out_w, 1), + }) + .unwrap(); + for (i, &in_pixel) in input.iter().enumerate().take(xsize) { + out_row.row(0)[out_w - 1 - position.0 - i] = in_pixel * self.scale; + } + } + + // transposing cases + Orientation::Transpose + | Orientation::Rotate90 + | Orientation::AntiTranspose + | Orientation::Rotate270 => { + let y_src = position.1; + for (i, &in_pixel) in input.iter().enumerate().take(xsize) { + let x_src = position.0 + i; + + let (x_dest, y_dest) = match self.orientation { + Orientation::Transpose => (y_src, x_src), + Orientation::Rotate90 => (y_src, w_src - 1 - x_src), + Orientation::AntiTranspose => (h_src - 1 - y_src, w_src - 1 - x_src), + Orientation::Rotate270 => (h_src - 1 - y_src, x_src), + _ => unreachable!(), + }; + + // Final check to ensure we don't write out of the destination buffer bounds. + if x_dest < out_w && y_dest < out_h { + outbuf_rect.row(y_dest)[x_dest] = in_pixel * self.scale; + } + } + } + } + } +} + +impl<T: ImageDataType> RunStage for SaveStage<T> { + fn run_stage_on( + &self, + chunk_size: usize, + input_buffers: &[&Image<f64>], + _output_buffers: &mut [&mut Image<f64>], + mut state: Option<&mut dyn Any>, + ) { + debug!("running save stage '{self}' in simple pipeline"); + let numc = input_buffers.len(); + if numc == 0 { + return; + } + let size = input_buffers[0].size(); + for b in input_buffers.iter() { + assert_eq!(size, b.size()); + } + let mut buffer = + vec![vec![T::default(); round_up_size_to_two_cache_lines::<T>(chunk_size)]; numc]; + for y in 0..size.1 { + for x in (0..size.0).step_by(chunk_size) { + let xsize = size.0.min(x + chunk_size) - x; + debug!("position: {x}x{y} xsize: {xsize}"); + for c in 0..numc { + let in_rect = input_buffers[c].as_rect(); + let in_row = in_rect.row(y); + for ix in 0..xsize { + buffer[c][ix] = T::from_f64(in_row[x + ix]); + } + } + let mut row: Vec<_> = buffer.iter().map(|x| x as &[T]).collect(); + self.process_row_chunk((x, y), xsize, &mut row, state.as_deref_mut()); + } + } + } + + fn init_local_state(&self) -> Result<Option<Box<dyn Any>>> { + Ok(None) + } + fn shift(&self) -> (u8, u8) { + (0, 0) + } + fn new_size(&self, size: (usize, usize)) -> (usize, usize) { + size + } + fn uses_channel(&self, c: usize) -> bool { + self.channel == c + } + fn as_any(self: Box<Self>) -> Box<dyn Any> { + self + } + fn input_type(&self) -> DataTypeTag { + T::DATA_TYPE_ID + } + fn output_type(&self) -> DataTypeTag { + T::DATA_TYPE_ID + } +} + +#[cfg(test)] +mod test { + use super::*; + use rand::SeedableRng; + use rand_xorshift::XorShiftRng; + use test_log::test; + + #[test] + fn save_stage() -> Result<()> { + let save_stage = SaveStage::<u8>::new( + SaveStageType::Output, + 0, + (128, 128), + 1, + Orientation::Identity, + )?; + let mut rng = XorShiftRng::seed_from_u64(0); + let src = Image::<u8>::new_random((128, 128), &mut rng)?; + + for i in 0..128 { + save_stage.process_row_chunk((0, i), 128, &mut [src.as_rect().row(i)], None); + } + + src.as_rect().check_equal(save_stage.buffer().as_rect()); + + Ok(()) + } + + macro_rules! test_orientation { + ($test_name:ident, $orientation:expr, $transform:expr) => { + #[test] + fn $test_name() -> Result<()> { + // Source dimensions + let (w, h) = (32, 16); + let mut rng = XorShiftRng::seed_from_u64(0); + let src = Image::<u8>::new_random((w, h), &mut rng)?; + let orientation = $orientation; + + // SaveStage will create its buffer with the correct (possibly swapped) dimensions. + let save_stage = + SaveStage::<u8>::new(SaveStageType::Output, 0, (w, h), 1, orientation)?; + + for y in 0..h { + save_stage.process_row_chunk((0, y), w, &mut [src.as_rect().row(y)], None); + } + + let (out_w, out_h) = save_stage.buffer().size(); + + let mut expected = Image::<u8>::new((out_w, out_h))?; + + // The transform is a closure: |x_dest, y_dest, w_src, h_src| -> (x_src, y_src) + let transform = $transform; + + // Iterate over the DESTINATION image pixels. + for y_dest in 0..out_h { + for x_dest in 0..out_w { + // For each destination pixel, find its corresponding source pixel. + let (src_x, src_y) = transform(x_dest, y_dest, w, h); + expected.as_rect_mut().row(y_dest)[x_dest] = + src.as_rect().row(src_y)[src_x]; + } + } + + expected + .as_rect() + .check_equal(save_stage.buffer().as_rect()); + + Ok(()) + } + }; + } + test_orientation!(orientation_identity, Orientation::Identity, |x, y, _, _| ( + x, y + )); + + test_orientation!( + orientation_flip_horizontal, + Orientation::FlipHorizontal, + |x, y, w, _| (w - 1 - x, y) + ); + + test_orientation!( + orientation_flip_vertical, + Orientation::FlipVertical, + |x, y, _, h| (x, h - 1 - y) + ); + + test_orientation!( + orientation_rotate_180, + Orientation::Rotate180, + |x, y, w, h| (w - 1 - x, h - 1 - y) + ); + + // transposing orientations + + test_orientation!( + orientation_transpose, + Orientation::Transpose, + |x_dest, y_dest, _, _| (y_dest, x_dest) + ); + + test_orientation!( + orientation_rotate_90, + Orientation::Rotate90, + |x_dest, y_dest, w_src, _| (w_src - 1 - y_dest, x_dest) + ); + + test_orientation!( + orientation_anti_transpose, + Orientation::AntiTranspose, + |x_dest, y_dest, w_src, h_src| (w_src - 1 - y_dest, h_src - 1 - x_dest) + ); + + test_orientation!( + orientation_rotate_270, + Orientation::Rotate270, + |x_dest, y_dest, _, h_src| (y_dest, h_src - 1 - x_dest) + ); +} diff --git a/third_party/rust/jxl/src/render/stages/blending.rs b/third_party/rust/jxl/src/render/stages/blending.rs @@ -0,0 +1,178 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +use crate::{ + error::Result, + features::{ + blending::perform_blending, + patches::{PatchBlendMode, PatchBlending}, + }, + frame::ReferenceFrame, + headers::{FileHeader, extra_channels::ExtraChannelInfo, frame_header::*}, + render::{RenderPipelineInPlaceStage, RenderPipelineStage}, + util::slice, +}; + +pub struct BlendingStage { + pub frame_origin: (isize, isize), + pub image_size: (isize, isize), + pub blending_info: BlendingInfo, + pub ec_blending_info: Vec<BlendingInfo>, + pub extra_channels: Vec<ExtraChannelInfo>, + pub reference_frames: Vec<Option<ReferenceFrame>>, + pub zeros: Vec<f32>, +} + +impl From<&BlendingInfo> for PatchBlending { + fn from(info: &BlendingInfo) -> Self { + let mode = match info.mode { + BlendingMode::Replace => PatchBlendMode::None, + BlendingMode::Add => PatchBlendMode::Add, + BlendingMode::Mul => PatchBlendMode::Mul, + BlendingMode::Blend => PatchBlendMode::BlendBelow, + BlendingMode::AlphaWeightedAdd => PatchBlendMode::AlphaWeightedAddBelow, + }; + PatchBlending { + mode, + alpha_channel: info.alpha_channel as usize, + clamp: info.clamp, + } + } +} + +impl BlendingStage { + pub fn new( + frame_header: &FrameHeader, + file_header: &FileHeader, + reference_frames: &[Option<ReferenceFrame>], + ) -> Result<BlendingStage> { + Ok(BlendingStage { + frame_origin: (frame_header.x0 as isize, frame_header.y0 as isize), + image_size: ( + file_header.size.xsize() as isize, + file_header.size.ysize() as isize, + ), + blending_info: frame_header.blending_info.clone(), + ec_blending_info: frame_header.ec_blending_info.clone(), + extra_channels: file_header.image_metadata.extra_channel_info.clone(), + reference_frames: reference_frames.to_owned(), + zeros: vec![0f32; file_header.size.xsize() as usize], + }) + } +} + +impl std::fmt::Display for BlendingStage { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "blending") + } +} + +impl RenderPipelineStage for BlendingStage { + type Type = RenderPipelineInPlaceStage<f32>; + + fn uses_channel(&self, c: usize) -> bool { + c < 3 + self.extra_channels.len() + } + + fn process_row_chunk( + &self, + position: (usize, usize), + xsize: usize, + row: &mut [&mut [f32]], + _state: Option<&mut dyn std::any::Any>, + ) { + let num_ec = self.extra_channels.len(); + let fg_y0 = self.frame_origin.1 + position.1 as isize; + let mut fg_x0 = self.frame_origin.0 + position.0 as isize; + let mut fg_x1 = fg_x0 + xsize as isize; + let mut bg_x0: isize = 0; + let mut bg_x1: isize = xsize as isize; + + if fg_x1 <= 0 || fg_x0 >= self.image_size.0 || fg_y0 < 0 || fg_y0 >= self.image_size.1 { + return; + } + + if fg_x0 < 0 { + bg_x0 -= fg_x0; + fg_x0 = 0; + } + if fg_x1 > self.image_size.0 { + bg_x1 = bg_x0 + self.image_size.0 - fg_x0; + fg_x1 = self.image_size.0; + } + + let fg_x0: usize = fg_x0 as usize; + let fg_x1: usize = fg_x1 as usize; + let bg_x0: usize = bg_x0 as usize; + let bg_x1: usize = bg_x1 as usize; + let fg_y0: usize = fg_y0 as usize; + + // TODO(szabadka): Allocate a buffer for this when building the stage instead of when + // executing it. + let mut out = row + .iter_mut() + .map(|s| &mut s[..xsize]) + .collect::<Vec<&mut [f32]>>(); + + let mut fg = vec![self.zeros.as_slice(); 3 + num_ec]; + + for (c, fg_ptr) in fg.iter_mut().enumerate().take(3) { + if self.reference_frames[self.blending_info.source as usize].is_some() { + *fg_ptr = &(self.reference_frames[self.blending_info.source as usize] + .as_ref() + .unwrap() + .frame[c] + .as_rect() + .row(fg_y0)[fg_x0..fg_x1]); + } + } + for i in 0..num_ec { + if self.reference_frames[self.ec_blending_info[i].source as usize].is_some() { + fg[3 + i] = &(self.reference_frames[self.ec_blending_info[i].source as usize] + .as_ref() + .unwrap() + .frame[3 + i] + .as_rect() + .row(fg_y0)[fg_x0..fg_x1]); + } + } + + let blending_info = PatchBlending::from(&self.blending_info); + let ec_blending_info: Vec<PatchBlending> = self + .ec_blending_info + .iter() + .map(PatchBlending::from) + .collect(); + + perform_blending( + &mut slice!(&mut out, .., bg_x0..bg_x1), + &fg, + &blending_info, + &ec_blending_info, + &self.extra_channels, + ); + } +} + +#[cfg(test)] +mod test { + use test_log::test; + + use super::*; + use crate::error::Result; + use crate::util::test::read_headers_and_toc; + + #[test] + fn blending_consistency() -> Result<()> { + let (file_header, frame_header, _) = + read_headers_and_toc(include_bytes!("../../../resources/test/basic.jxl")).unwrap(); + let reference_frames: Vec<Option<ReferenceFrame>> = vec![None, None, None, None]; + crate::render::test::test_stage_consistency::<_, f32, f32>( + BlendingStage::new(&frame_header, &file_header, &reference_frames)?, + (500, 500), + 4, + ) + } +} diff --git a/third_party/rust/jxl/src/render/stages/chroma_upsample.rs b/third_party/rust/jxl/src/render/stages/chroma_upsample.rs @@ -0,0 +1,153 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +use crate::render::{RenderPipelineInOutStage, RenderPipelineStage}; + +pub struct HorizontalChromaUpsample { + channel: usize, +} + +impl HorizontalChromaUpsample { + pub fn new(channel: usize) -> HorizontalChromaUpsample { + HorizontalChromaUpsample { channel } + } +} + +impl std::fmt::Display for HorizontalChromaUpsample { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!( + f, + "chroma upsample of channel {}, horizontally", + self.channel + ) + } +} + +impl RenderPipelineStage for HorizontalChromaUpsample { + type Type = RenderPipelineInOutStage<f32, f32, 1, 0, 1, 0>; + + fn uses_channel(&self, c: usize) -> bool { + c == self.channel + } + + fn process_row_chunk( + &self, + _position: (usize, usize), + xsize: usize, + row: &mut [(&[&[f32]], &mut [&mut [f32]])], + _state: Option<&mut dyn std::any::Any>, + ) { + let (input, output) = &mut row[0]; + for i in 0..xsize { + let scaled_cur = input[0][i + 1] * 0.75; + let prev = input[0][i]; + let next = input[0][i + 2]; + let left = 0.25 * prev + scaled_cur; + let right = 0.25 * next + scaled_cur; + output[0][2 * i] = left; + output[0][2 * i + 1] = right; + } + } +} + +pub struct VerticalChromaUpsample { + channel: usize, +} + +impl VerticalChromaUpsample { + pub fn new(channel: usize) -> VerticalChromaUpsample { + VerticalChromaUpsample { channel } + } +} + +impl std::fmt::Display for VerticalChromaUpsample { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "chroma upsample of channel {}, vertically", self.channel) + } +} + +impl RenderPipelineStage for VerticalChromaUpsample { + type Type = RenderPipelineInOutStage<f32, f32, 0, 1, 0, 1>; + + fn uses_channel(&self, c: usize) -> bool { + c == self.channel + } + + fn process_row_chunk( + &self, + _position: (usize, usize), + xsize: usize, + row: &mut [(&[&[f32]], &mut [&mut [f32]])], + _state: Option<&mut dyn std::any::Any>, + ) { + let (input, output) = &mut row[0]; + for i in 0..xsize { + let scaled_cur = input[1][i] * 0.75; + let prev = input[0][i]; + let next = input[2][i]; + let up = 0.25 * prev + scaled_cur; + let down = 0.25 * next + scaled_cur; + output[0][i] = up; + output[1][i] = down; + } + } +} + +#[cfg(test)] +mod test { + use super::*; + use crate::{error::Result, image::Image, render::test::make_and_run_simple_pipeline}; + use test_log::test; + + #[test] + fn hchr_consistency() -> Result<()> { + crate::render::test::test_stage_consistency::<_, f32, f32>( + HorizontalChromaUpsample::new(0), + (500, 500), + 1, + ) + } + + #[test] + fn test_hchr() -> Result<()> { + let mut input = Image::new((3, 1))?; + input + .as_rect_mut() + .row(0) + .copy_from_slice(&[1.0f32, 2.0, 4.0]); + let stage = HorizontalChromaUpsample::new(0); + let output: Vec<Image<f32>> = + make_and_run_simple_pipeline(stage, &[input], (6, 1), 0, 256)?.1; + assert_eq!(output[0].as_rect().row(0), [1.0, 1.25, 1.75, 2.5, 3.5, 4.0]); + Ok(()) + } + + #[test] + fn vchr_consistency() -> Result<()> { + crate::render::test::test_stage_consistency::<_, f32, f32>( + VerticalChromaUpsample::new(0), + (500, 500), + 1, + ) + } + + #[test] + fn test_vchr() -> Result<()> { + let mut input = Image::new((1, 3))?; + input.as_rect_mut().row(0)[0] = 1.0f32; + input.as_rect_mut().row(1)[0] = 2.0f32; + input.as_rect_mut().row(2)[0] = 4.0f32; + let stage = VerticalChromaUpsample::new(0); + let output: Vec<Image<f32>> = + make_and_run_simple_pipeline(stage, &[input], (1, 6), 0, 256)?.1; + assert_eq!(output[0].as_rect().row(0)[0], 1.0); + assert_eq!(output[0].as_rect().row(1)[0], 1.25); + assert_eq!(output[0].as_rect().row(2)[0], 1.75); + assert_eq!(output[0].as_rect().row(3)[0], 2.5); + assert_eq!(output[0].as_rect().row(4)[0], 3.5); + assert_eq!(output[0].as_rect().row(5)[0], 4.0); + Ok(()) + } +} diff --git a/third_party/rust/jxl/src/render/stages/convert.rs b/third_party/rust/jxl/src/render/stages/convert.rs @@ -0,0 +1,222 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +use crate::{ + frame::quantizer::LfQuantFactors, + headers::bit_depth::BitDepth, + render::{RenderPipelineInOutStage, RenderPipelineStage}, +}; + +pub struct ConvertU8F32Stage { + channel: usize, +} + +impl ConvertU8F32Stage { + pub fn new(channel: usize) -> ConvertU8F32Stage { + ConvertU8F32Stage { channel } + } +} + +impl std::fmt::Display for ConvertU8F32Stage { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "convert U8 data to F32 in channel {}", self.channel) + } +} + +impl RenderPipelineStage for ConvertU8F32Stage { + type Type = RenderPipelineInOutStage<u8, f32, 0, 0, 0, 0>; + + fn uses_channel(&self, c: usize) -> bool { + c == self.channel + } + + fn process_row_chunk( + &self, + _position: (usize, usize), + xsize: usize, + row: &mut [(&[&[u8]], &mut [&mut [f32]])], + _state: Option<&mut dyn std::any::Any>, + ) { + let (input, output) = &mut row[0]; + for i in 0..xsize { + output[0][i] = input[0][i] as f32 * (1.0 / 255.0); + } + } +} + +pub struct ConvertModularXYBToF32Stage { + first_channel: usize, + scale: [f32; 3], +} + +impl ConvertModularXYBToF32Stage { + pub fn new(first_channel: usize, lf_quant: &LfQuantFactors) -> ConvertModularXYBToF32Stage { + ConvertModularXYBToF32Stage { + first_channel, + scale: lf_quant.quant_factors, + } + } +} + +impl std::fmt::Display for ConvertModularXYBToF32Stage { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!( + f, + "convert modular xyb data to F32 in channels {}..{} with scales {:?}", + self.first_channel, + self.first_channel + 2, + self.scale + ) + } +} + +impl RenderPipelineStage for ConvertModularXYBToF32Stage { + type Type = RenderPipelineInOutStage<i32, f32, 0, 0, 0, 0>; + + fn uses_channel(&self, c: usize) -> bool { + (self.first_channel..self.first_channel + 3).contains(&c) + } + + fn process_row_chunk( + &self, + _position: (usize, usize), + xsize: usize, + row: &mut [(&[&[i32]], &mut [&mut [f32]])], + _state: Option<&mut dyn std::any::Any>, + ) { + let [scale_x, scale_y, scale_b] = self.scale; + let [ + (input_y, output_x), + (input_x, output_y), + (input_b, output_b), + ] = row + else { + panic!( + "incorrect number of channels; expected 3, found {}", + row.len() + ); + }; + for i in 0..xsize { + output_x[0][i] = input_x[0][i] as f32 * scale_x; + output_y[0][i] = input_y[0][i] as f32 * scale_y; + output_b[0][i] = (input_b[0][i] + input_y[0][i]) as f32 * scale_b; + } + } +} + +pub struct ConvertModularToF32Stage { + channel: usize, + bit_depth: BitDepth, +} + +impl ConvertModularToF32Stage { + pub fn new(channel: usize, bit_depth: BitDepth) -> ConvertModularToF32Stage { + ConvertModularToF32Stage { channel, bit_depth } + } +} + +impl std::fmt::Display for ConvertModularToF32Stage { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!( + f, + "convert modular data to F32 in channel {} with bit depth {:?}", + self.channel, self.bit_depth + ) + } +} + +// Converts custom [bits]-bit float (with [exp_bits] exponent bits) stored as +// int back to binary32 float. +// TODO(sboukortt): SIMD +fn int_to_float(input: &[i32], output: &mut [f32], bit_depth: &BitDepth) { + assert_eq!(input.len(), output.len()); + let bits = bit_depth.bits_per_sample(); + let exp_bits = bit_depth.exponent_bits_per_sample(); + if bits == 32 { + assert!(exp_bits == 8); + for (&in_val, out_val) in input.iter().zip(output) { + *out_val = f32::from_bits(in_val as u32); + } + return; + } + let exp_bias = (1 << (exp_bits - 1)) - 1; + let sign_shift = bits - 1; + let mant_bits = bits - exp_bits - 1; + let mant_shift = 23 - mant_bits; + for (&in_val, out_val) in input.iter().zip(output) { + let mut f = in_val as u32; + let signbit = (f >> sign_shift) != 0; + f &= (1 << sign_shift) - 1; + if f == 0 { + *out_val = if signbit { -0.0 } else { 0.0 }; + continue; + } + let mut exp = (f >> mant_bits) as i32; + let mut mantissa = f & ((1 << mant_bits) - 1); + mantissa <<= mant_shift; + // Try to normalize only if there is space for maneuver. + if exp == 0 && exp_bits < 8 { + // subnormal number + while (mantissa & 0x800000) == 0 { + mantissa <<= 1; + exp -= 1; + } + exp += 1; + // remove leading 1 because it is implicit now + mantissa &= 0x7fffff; + } + exp -= exp_bias; + // broke up the arbitrary float into its parts, now reassemble into + // binary32 + exp += 127; + assert!(exp >= 0); + f = if signbit { 0x80000000 } else { 0 }; + f |= (exp as u32) << 23; + f |= mantissa; + *out_val = f32::from_bits(f); + } +} + +impl RenderPipelineStage for ConvertModularToF32Stage { + type Type = RenderPipelineInOutStage<i32, f32, 0, 0, 0, 0>; + + fn uses_channel(&self, c: usize) -> bool { + c == self.channel + } + + fn process_row_chunk( + &self, + _position: (usize, usize), + xsize: usize, + row: &mut [(&[&[i32]], &mut [&mut [f32]])], + _state: Option<&mut dyn std::any::Any>, + ) { + let (input, output) = &mut row[0]; + if self.bit_depth.floating_point_sample() { + int_to_float(&input[0][..xsize], &mut output[0][..xsize], &self.bit_depth); + } else { + let scale = 1.0 / ((1u64 << self.bit_depth.bits_per_sample()) - 1) as f32; + for i in 0..xsize { + output[0][i] = input[0][i] as f32 * scale; + } + } + } +} + +#[cfg(test)] +mod test { + use super::*; + use crate::error::Result; + use test_log::test; + + #[test] + fn u8_consistency() -> Result<()> { + crate::render::test::test_stage_consistency::<_, u8, f32>( + ConvertU8F32Stage::new(0), + (500, 500), + 1, + ) + } +} diff --git a/third_party/rust/jxl/src/render/stages/epf.rs b/third_party/rust/jxl/src/render/stages/epf.rs @@ -0,0 +1,555 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +use std::sync::Arc; + +use crate::{ + BLOCK_DIM, MIN_SIGMA, SIGMA_PADDING, + image::Image, + render::{RenderPipelineInOutStage, RenderPipelineStage}, + simd::{F32SimdVec, simd_function}, +}; + +/// 5x5 plus-shaped kernel with 5 SADs per pixel (3x3 plus-shaped). So this makes this filter a 7x7 filter. +pub struct Epf0Stage { + /// Multiplier for sigma in pass 0 + sigma_scale: f32, + /// (inverse) multiplier for sigma on borders + border_sad_mul: f32, + channel_scale: [f32; 3], + sigma: Arc<Image<f32>>, +} + +impl std::fmt::Display for Epf0Stage { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!( + f, + "EPF stage 0 with sigma scale: {}, border_sad_mul: {}", + self.sigma_scale, self.border_sad_mul + ) + } +} + +impl Epf0Stage { + pub fn new( + sigma_scale: f32, + border_sad_mul: f32, + channel_scale: [f32; 3], + sigma: Arc<Image<f32>>, + ) -> Self { + Self { + sigma, + sigma_scale, + channel_scale, + border_sad_mul, + } + } +} + +type EPFRow<'a> = [(&'a [&'a [f32]], &'a mut [&'a mut [f32]])]; + +simd_function!( + epf0_process_row_chunk_dispatch, + d: D, + fn epf0_process_row_chunk_simd( + stage: &Epf0Stage, + pos: (usize, usize), + xsize: usize, + row: &mut EPFRow, +) { + let (xpos, ypos) = pos; + assert!(row.len() == 3, "Expected 3 channels, got {}", row.len()); + + let row_sigma = stage.sigma.as_rect().row(ypos / BLOCK_DIM + SIGMA_PADDING); + + let sm = stage.sigma_scale * 1.65; + let bsm = sm * stage.border_sad_mul; + + let sad_mul_block = if ypos % BLOCK_DIM == 0 || ypos % BLOCK_DIM == BLOCK_DIM - 1 { + [bsm; 8] // border + } else { + [bsm, sm, sm, sm, sm, sm, sm, bsm] // center + }; + + let mut sigma_storage = vec![0.0; D::F32Vec::LEN]; + let mut sad_mul_storage = vec![0.0; D::F32Vec::LEN]; + + for x in (0..xsize).step_by(D::F32Vec::LEN) { + for (k, value) in sigma_storage.iter_mut().enumerate() { + let x = x + k; + *value = row_sigma[(x + xpos + SIGMA_PADDING * BLOCK_DIM) / BLOCK_DIM]; + } + + if sigma_storage.iter().all(|&sigma| sigma < MIN_SIGMA) { + for (input_c, output_c) in row.iter_mut() { + D::F32Vec::load(d, &input_c[3][3 + x..]).store(&mut output_c[0][x..]); + } + continue; + } + + for (k, value) in sad_mul_storage.iter_mut().enumerate() { + let x = x + k; + *value = sad_mul_block[(x + xpos) % BLOCK_DIM]; + } + + let sigma = D::F32Vec::load(d, &sigma_storage); + let sad_mul = D::F32Vec::load(d, &sad_mul_storage); + + // Compute SADs + let mut sads = [D::F32Vec::splat(d, 0.0); 12]; + for ((input_c, _), scale) in row.iter_mut().zip(stage.channel_scale) { + let scale = D::F32Vec::splat(d, scale); + + let p30 = D::F32Vec::load(d, &input_c[0][3 + x..]); + let p21 = D::F32Vec::load(d, &input_c[1][2 + x..]); + let p31 = D::F32Vec::load(d, &input_c[1][3 + x..]); + let p41 = D::F32Vec::load(d, &input_c[1][4 + x..]); + let p12 = D::F32Vec::load(d, &input_c[2][1 + x..]); + let p22 = D::F32Vec::load(d, &input_c[2][2 + x..]); + let p32 = D::F32Vec::load(d, &input_c[2][3 + x..]); + let p42 = D::F32Vec::load(d, &input_c[2][4 + x..]); + let p52 = D::F32Vec::load(d, &input_c[2][5 + x..]); + let p03 = D::F32Vec::load(d, &input_c[3][x..]); + let p13 = D::F32Vec::load(d, &input_c[3][1 + x..]); + let p23 = D::F32Vec::load(d, &input_c[3][2 + x..]); + let p33 = D::F32Vec::load(d, &input_c[3][3 + x..]); + let p43 = D::F32Vec::load(d, &input_c[3][4 + x..]); + let p53 = D::F32Vec::load(d, &input_c[3][5 + x..]); + let p63 = D::F32Vec::load(d, &input_c[3][6 + x..]); + let p14 = D::F32Vec::load(d, &input_c[4][1 + x..]); + let p24 = D::F32Vec::load(d, &input_c[4][2 + x..]); + let p34 = D::F32Vec::load(d, &input_c[4][3 + x..]); + let p44 = D::F32Vec::load(d, &input_c[4][4 + x..]); + let p54 = D::F32Vec::load(d, &input_c[4][5 + x..]); + let p25 = D::F32Vec::load(d, &input_c[5][2 + x..]); + let p35 = D::F32Vec::load(d, &input_c[5][3 + x..]); + let p45 = D::F32Vec::load(d, &input_c[5][4 + x..]); + let p36 = D::F32Vec::load(d, &input_c[6][3 + x..]); + let d32_30 = (p32 - p30).abs(); + let d32_21 = (p32 - p21).abs(); + let d32_31 = (p32 - p31).abs(); + let d32_41 = (p32 - p41).abs(); + let d32_12 = (p32 - p12).abs(); + let d32_22 = (p32 - p22).abs(); + let d32_42 = (p32 - p42).abs(); + let d32_52 = (p32 - p52).abs(); + let d32_23 = (p32 - p23).abs(); + let d32_34 = (p32 - p34).abs(); + let d32_43 = (p32 - p43).abs(); + let d32_33 = (p32 - p33).abs(); + let d23_21 = (p23 - p21).abs(); + let d23_12 = (p23 - p12).abs(); + let d23_22 = (p23 - p22).abs(); + let d23_03 = (p23 - p03).abs(); + let d23_13 = (p23 - p13).abs(); + let d23_33 = (p23 - p33).abs(); + let d23_43 = (p23 - p43).abs(); + let d23_14 = (p23 - p14).abs(); + let d23_24 = (p23 - p24).abs(); + let d23_34 = (p23 - p34).abs(); + let d23_25 = (p23 - p25).abs(); + let d33_31 = (p33 - p31).abs(); + let d33_22 = (p33 - p22).abs(); + let d33_42 = (p33 - p42).abs(); + let d33_13 = (p33 - p13).abs(); + let d33_43 = (p33 - p43).abs(); + let d33_53 = (p33 - p53).abs(); + let d33_24 = (p33 - p24).abs(); + let d33_34 = (p33 - p34).abs(); + let d33_44 = (p33 - p44).abs(); + let d33_35 = (p33 - p35).abs(); + let d43_41 = (p43 - p41).abs(); + let d43_42 = (p43 - p42).abs(); + let d43_52 = (p43 - p52).abs(); + let d43_53 = (p43 - p53).abs(); + let d43_63 = (p43 - p63).abs(); + let d43_34 = (p43 - p34).abs(); + let d43_44 = (p43 - p44).abs(); + let d43_54 = (p43 - p54).abs(); + let d43_45 = (p43 - p45).abs(); + let d34_14 = (p34 - p14).abs(); + let d34_24 = (p34 - p24).abs(); + let d34_44 = (p34 - p44).abs(); + let d34_54 = (p34 - p54).abs(); + let d34_25 = (p34 - p25).abs(); + let d34_35 = (p34 - p35).abs(); + let d34_45 = (p34 - p45).abs(); + let d34_36 = (p34 - p36).abs(); + sads[0] = scale.mul_add(d32_30 + d23_21 + d33_31 + d43_41 + d32_34, sads[0]); + sads[1] = scale.mul_add(d32_21 + d23_12 + d33_22 + d32_43 + d23_34, sads[1]); + sads[2] = scale.mul_add(d32_31 + d23_22 + d32_33 + d43_42 + d33_34, sads[2]); + sads[3] = scale.mul_add(d32_41 + d32_23 + d33_42 + d43_52 + d43_34, sads[3]); + sads[4] = scale.mul_add(d32_12 + d23_03 + d33_13 + d23_43 + d34_14, sads[4]); + sads[5] = scale.mul_add(d32_22 + d23_13 + d23_33 + d33_43 + d34_24, sads[5]); + sads[6] = scale.mul_add(d32_42 + d23_33 + d33_43 + d43_53 + d34_44, sads[6]); + sads[7] = scale.mul_add(d32_52 + d23_43 + d33_53 + d43_63 + d34_54, sads[7]); + sads[8] = scale.mul_add(d32_23 + d23_14 + d33_24 + d43_34 + d34_25, sads[8]); + sads[9] = scale.mul_add(d32_33 + d23_24 + d33_34 + d43_44 + d34_35, sads[9]); + sads[10] = scale.mul_add(d32_43 + d23_34 + d33_44 + d43_54 + d34_45, sads[10]); + sads[11] = scale.mul_add(d32_34 + d23_25 + d33_35 + d43_45 + d34_36, sads[11]); + } + // Compute output based on SADs + let inv_sigma = sigma * sad_mul; + let mut w = D::F32Vec::splat(d, 1.0); + for sad in sads.iter_mut() { + *sad = sad + .mul_add(inv_sigma, D::F32Vec::splat(d, 1.0)) + .max(D::F32Vec::splat(d, 0.0)); + w += *sad; + } + let inv_w = D::F32Vec::splat(d, 1.0) / w; + for (input_c, output_c) in row.iter_mut() { + let mut out = D::F32Vec::load(d, &input_c[3][3 + x..]); + for (row_idx, col_idx, sad_idx) in [ + (5, 3+x, 11), + (4, 4+x, 10), + (4, 3+x, 9), + (4, 2+x, 8), + (3, 5+x, 7), + (3, 4+x, 6), + (3, 2+x, 5), + (3, 1+x, 4), + (2, 4+x, 3), + (2, 3+x, 2), + (2, 2+x, 1), + (1, 3+x, 0), + ] { + out = D::F32Vec::load(d, &input_c[row_idx][col_idx..]).mul_add(sads[sad_idx], out); + } + (out * inv_w).store(&mut output_c[0][x..]); + } + } +}); + +impl RenderPipelineStage for Epf0Stage { + type Type = RenderPipelineInOutStage<f32, f32, 3, 3, 0, 0>; + + fn uses_channel(&self, c: usize) -> bool { + c < 3 + } + + fn process_row_chunk( + &self, + (xpos, ypos): (usize, usize), + xsize: usize, + row: &mut EPFRow, + _state: Option<&mut dyn std::any::Any>, + ) { + epf0_process_row_chunk_dispatch(self, (xpos, ypos), xsize, row); + } +} + +/// 3x3 plus-shaped kernel with 5 SADs per pixel (3x3 plus-shaped). So this makes this filter a 5x5 filter. +pub struct Epf1Stage { + /// Multiplier for sigma in pass 1 + sigma_scale: f32, + /// (inverse) multiplier for sigma on borders + border_sad_mul: f32, + channel_scale: [f32; 3], + sigma: Arc<Image<f32>>, +} + +impl std::fmt::Display for Epf1Stage { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!( + f, + "EPF stage 1 with sigma scale: {}, border_sad_mul: {}", + self.sigma_scale, self.border_sad_mul + ) + } +} + +impl Epf1Stage { + pub fn new( + sigma_scale: f32, + border_sad_mul: f32, + channel_scale: [f32; 3], + sigma: Arc<Image<f32>>, + ) -> Self { + Self { + sigma, + sigma_scale, + channel_scale, + border_sad_mul, + } + } +} + +simd_function!( +epf1_process_row_chunk_dispatch, +d: D, +fn epf1_process_row_chunk( + stage: &Epf1Stage, + pos: (usize, usize), + xsize: usize, + row: &mut EPFRow, +) { + let (xpos, ypos) = pos; + assert!(row.len() == 3, "Expected 3 channels, got {}", row.len()); + + let row_sigma = stage.sigma.as_rect().row(ypos / BLOCK_DIM + SIGMA_PADDING); + + let sm = stage.sigma_scale * 1.65; + let bsm = sm * stage.border_sad_mul; + + let sad_mul_block = if ypos % BLOCK_DIM == 0 || ypos % BLOCK_DIM == BLOCK_DIM - 1 { + [bsm; 8] // border + } else { + [bsm, sm, sm, sm, sm, sm, sm, bsm] // center + }; + + let mut sigma_storage = vec![0.0; D::F32Vec::LEN]; + let mut sad_mul_storage = vec![0.0; D::F32Vec::LEN]; + + for x in (0..xsize).step_by(D::F32Vec::LEN) { + for (k, value) in sigma_storage.iter_mut().enumerate() { + let x = x + k; + *value = row_sigma[(x + xpos + SIGMA_PADDING * BLOCK_DIM) / BLOCK_DIM]; + } + + if sigma_storage.iter().all(|&sigma| sigma < MIN_SIGMA) { + for (input_c, output_c) in row.iter_mut() { + D::F32Vec::load(d, &input_c[2][2 + x..]).store(&mut output_c[0][x..]); + } + continue; + } + + for (k, value) in sad_mul_storage.iter_mut().enumerate() { + let x = x + k; + *value = sad_mul_block[(x + xpos) % BLOCK_DIM]; + } + + let sigma = D::F32Vec::load(d, &sigma_storage); + let sad_mul = D::F32Vec::load(d, &sad_mul_storage); + + // Compute SADs + let mut sads = [D::F32Vec::splat(d, 0.0); 4]; + for ((input_c, _), scale) in row.iter_mut().zip(stage.channel_scale) { + let scale = D::F32Vec::splat(d, scale); + let p20 = D::F32Vec::load(d, &input_c[0][2 + x..]); + let p11 = D::F32Vec::load(d, &input_c[1][1 + x..]); + let p21 = D::F32Vec::load(d, &input_c[1][2 + x..]); + let p31 = D::F32Vec::load(d, &input_c[1][3 + x..]); + let p02 = D::F32Vec::load(d, &input_c[2][x..]); + let p12 = D::F32Vec::load(d, &input_c[2][1 + x..]); + let p22 = D::F32Vec::load(d, &input_c[2][2 + x..]); + let p32 = D::F32Vec::load(d, &input_c[2][3 + x..]); + let p42 = D::F32Vec::load(d, &input_c[2][4 + x..]); + let p13 = D::F32Vec::load(d, &input_c[3][1 + x..]); + let p23 = D::F32Vec::load(d, &input_c[3][2 + x..]); + let p33 = D::F32Vec::load(d, &input_c[3][3 + x..]); + let p24 = D::F32Vec::load(d, &input_c[4][2 + x..]); + let d20_21 = (p20 - p21).abs(); + let d11_21 = (p11 - p21).abs(); + let d22_21 = (p22 - p21).abs(); + let d31_21 = (p31 - p21).abs(); + let d02_12 = (p02 - p12).abs(); + let d11_12 = (p11 - p12).abs(); + let d12_22 = (p22 - p12).abs(); + let d31_32 = (p31 - p32).abs(); + let d22_32 = (p22 - p32).abs(); + let d42_32 = (p42 - p32).abs(); + let d13_12 = (p13 - p12).abs(); + let d22_23 = (p22 - p23).abs(); + let d13_23 = (p13 - p23).abs(); + let d33_23 = (p33 - p23).abs(); + let d33_32 = (p33 - p32).abs(); + let d24_23 = (p24 - p23).abs(); + sads[0] = (d20_21 + d11_12 + d22_21 + d31_32 + d22_23).mul_add(scale, sads[0]); + sads[1] = (d11_21 + d02_12 + d12_22 + d22_32 + d13_23).mul_add(scale, sads[1]); + sads[2] = (d31_21 + d12_22 + d22_32 + d42_32 + d33_23).mul_add(scale, sads[2]); + sads[3] = (d22_21 + d13_12 + d22_23 + d33_32 + d24_23).mul_add(scale, sads[3]); + } + + // Compute output based on SADs + let inv_sigma = sigma * sad_mul; + let mut w = D::F32Vec::splat(d, 1.0); + for sad in sads.iter_mut() { + *sad = sad + .mul_add(inv_sigma, D::F32Vec::splat(d, 1.0)) + .max(D::F32Vec::splat(d, 0.0)); + w += *sad; + } + let inv_w = D::F32Vec::splat(d, 1.0) / w; + for (input_c, output_c) in row.iter_mut() { + let mut out = D::F32Vec::load(d, &input_c[2][2 + x..]); + for (row_idx, col_idx, sad_idx) in [ + (3, 2+x, 3), + (2, 3+x, 2), + (2, 1+x, 1), + (1, 2+x, 0), + ] { + out = D::F32Vec::load(d, &input_c[row_idx][col_idx..]).mul_add(sads[sad_idx], out); + } + (out * inv_w).store(&mut output_c[0][x..]); + } + } +}); + +impl RenderPipelineStage for Epf1Stage { + type Type = RenderPipelineInOutStage<f32, f32, 2, 2, 0, 0>; + + fn uses_channel(&self, c: usize) -> bool { + c < 3 + } + + fn process_row_chunk( + &self, + (xpos, ypos): (usize, usize), + xsize: usize, + row: &mut EPFRow, + _state: Option<&mut dyn std::any::Any>, + ) { + epf1_process_row_chunk_dispatch(self, (xpos, ypos), xsize, row); + } +} + +/// 3x3 plus-shaped kernel with 1 SAD per pixel. So this makes this filter a 3x3 filter. +pub struct Epf2Stage { + /// Multiplier for sigma in pass 2 + sigma_scale: f32, + /// (inverse) multiplier for sigma on borders + border_sad_mul: f32, + channel_scale: [f32; 3], + sigma: Arc<Image<f32>>, +} + +impl std::fmt::Display for Epf2Stage { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!( + f, + "EPF stage 2 with sigma scale: {}, border_sad_mul: {}", + self.sigma_scale, self.border_sad_mul + ) + } +} + +impl Epf2Stage { + pub fn new( + sigma_scale: f32, + border_sad_mul: f32, + channel_scale: [f32; 3], + sigma: Arc<Image<f32>>, + ) -> Self { + Self { + sigma, + sigma_scale, + channel_scale, + border_sad_mul, + } + } +} + +simd_function!( +epf2_process_row_chunk_dispatch, +d: D, +fn epf2_process_row_chunk( + stage: &Epf2Stage, + pos: (usize, usize), + xsize: usize, + row: &mut EPFRow, +) { + let (xpos, ypos) = pos; + let [ + (input_x, output_x), + (input_y, output_y), + (input_b, output_b), + ] = row + else { + panic!("Expected 3 channels, got {}", row.len()); + }; + + let row_sigma = stage.sigma.as_rect().row(ypos / BLOCK_DIM + SIGMA_PADDING); + + let sm = stage.sigma_scale * 1.65; + let bsm = sm * stage.border_sad_mul; + + let sad_mul_block = if ypos % BLOCK_DIM == 0 || ypos % BLOCK_DIM == BLOCK_DIM - 1 { + [bsm; 8] // border + } else { + [bsm, sm, sm, sm, sm, sm, sm, bsm] // center + }; + + let mut sigma_storage = vec![0.0; D::F32Vec::LEN]; + let mut sad_mul_storage = vec![0.0; D::F32Vec::LEN]; + + for x in (0..xsize).step_by(D::F32Vec::LEN) { + for (k, value) in sigma_storage.iter_mut().enumerate() { + let x = x + k; + *value = row_sigma[(x + xpos + SIGMA_PADDING * BLOCK_DIM) / BLOCK_DIM]; + } + + if sigma_storage.iter().all(|&sigma| sigma < MIN_SIGMA) { + D::F32Vec::load(d, &input_x[1][1 + x..]).store(&mut output_x[0][x..]); + D::F32Vec::load(d, &input_y[1][1 + x..]).store(&mut output_y[0][x..]); + D::F32Vec::load(d, &input_b[1][1 + x..]).store(&mut output_b[0][x..]); + continue; + } + + for (k, value) in sad_mul_storage.iter_mut().enumerate() { + let x = x + k; + *value = sad_mul_block[(x + xpos) % BLOCK_DIM]; + } + + let sigma = D::F32Vec::load(d, &sigma_storage); + let vsm = D::F32Vec::load(d, &sad_mul_storage); + let inv_sigma = sigma * vsm; + + let x_cc = D::F32Vec::load(d, &input_x[1][1 + x..]); + let y_cc = D::F32Vec::load(d, &input_y[1][1 + x..]); + let b_cc = D::F32Vec::load(d, &input_b[1][1 + x..]); + + let mut w_acc = D::F32Vec::splat(d, 1.0); + let mut x_acc = x_cc; + let mut y_acc = y_cc; + let mut b_acc = b_cc; + + for (y_off, x_off) in [(0, 1), (1, 0), (1, 2), (2, 1)] { + let (cx, cy, cb) = ( + D::F32Vec::load(d, &input_x[y_off as usize][x_off + x..]), + D::F32Vec::load(d, &input_y[y_off as usize][x_off + x..]), + D::F32Vec::load(d, &input_b[y_off as usize][x_off + x..]), + ); + let sad = (cx - x_cc).abs().mul_add( + D::F32Vec::splat(d, stage.channel_scale[0]), + (cy - y_cc).abs().mul_add( + D::F32Vec::splat(d, stage.channel_scale[1]), + (cb - b_cc).abs() * D::F32Vec::splat(d, stage.channel_scale[2]), + ), + ); + let weight = sad + .mul_add(inv_sigma, D::F32Vec::splat(d, 1.0)) + .max(D::F32Vec::splat(d, 0.0)); + w_acc += weight; + x_acc = weight.mul_add(cx, x_acc); + y_acc = weight.mul_add(cy, y_acc); + b_acc = weight.mul_add(cb, b_acc); + } + + let inv_w = D::F32Vec::splat(d, 1.0) / w_acc; + + (x_acc * inv_w).store(&mut output_x[0][x..]); + (y_acc * inv_w).store(&mut output_y[0][x..]); + (b_acc * inv_w).store(&mut output_b[0][x..]); + } +}); + +impl RenderPipelineStage for Epf2Stage { + type Type = RenderPipelineInOutStage<f32, f32, 1, 1, 0, 0>; + + fn uses_channel(&self, c: usize) -> bool { + c < 3 + } + + fn process_row_chunk( + &self, + (xpos, ypos): (usize, usize), + xsize: usize, + row: &mut EPFRow, + _state: Option<&mut dyn std::any::Any>, + ) { + epf2_process_row_chunk_dispatch(self, (xpos, ypos), xsize, row); + } +} diff --git a/third_party/rust/jxl/src/render/stages/extend.rs b/third_party/rust/jxl/src/render/stages/extend.rs @@ -0,0 +1,125 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +use crate::{ + error::Result, + frame::ReferenceFrame, + headers::{FileHeader, extra_channels::ExtraChannelInfo, frame_header::*}, + render::{RenderPipelineExtendStage, RenderPipelineStage}, +}; + +pub struct ExtendToImageDimensionsStage { + pub frame_origin: (isize, isize), + pub image_size: (isize, isize), + pub blending_info: BlendingInfo, + pub ec_blending_info: Vec<BlendingInfo>, + pub extra_channels: Vec<ExtraChannelInfo>, + pub reference_frames: Vec<Option<ReferenceFrame>>, + pub zeros: Vec<f32>, +} + +impl ExtendToImageDimensionsStage { + pub fn new( + frame_header: &FrameHeader, + file_header: &FileHeader, + reference_frames: &[Option<ReferenceFrame>], + ) -> Result<ExtendToImageDimensionsStage> { + Ok(ExtendToImageDimensionsStage { + frame_origin: (frame_header.x0 as isize, frame_header.y0 as isize), + image_size: ( + file_header.size.xsize() as isize, + file_header.size.ysize() as isize, + ), + blending_info: frame_header.blending_info.clone(), + ec_blending_info: frame_header.ec_blending_info.clone(), + extra_channels: file_header.image_metadata.extra_channel_info.clone(), + reference_frames: reference_frames.to_owned(), + zeros: vec![0f32; file_header.size.xsize() as usize], + }) + } +} + +impl std::fmt::Display for ExtendToImageDimensionsStage { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "extend-to-image-dims") + } +} + +impl RenderPipelineStage for ExtendToImageDimensionsStage { + type Type = RenderPipelineExtendStage<f32>; + + fn uses_channel(&self, _c: usize) -> bool { + true + } + + fn new_size(&self, _current_size: (usize, usize)) -> (usize, usize) { + (self.image_size.0 as usize, self.image_size.1 as usize) + } + + fn original_data_origin(&self) -> (isize, isize) { + self.frame_origin + } + + fn process_row_chunk( + &self, + position: (usize, usize), + xsize: usize, + row: &mut [&mut [f32]], + _state: Option<&mut dyn std::any::Any>, + ) { + let num_ec = self.extra_channels.len(); + let num_c = 3 + num_ec; + let x0 = position.0; + let x1 = x0 + xsize; + let y0 = position.1; + + let mut bg = vec![self.zeros.as_slice(); num_c]; + for (c, bg_ptr) in bg.iter_mut().enumerate().take(3) { + if self.reference_frames[self.blending_info.source as usize].is_some() { + *bg_ptr = self.reference_frames[self.blending_info.source as usize] + .as_ref() + .unwrap() + .frame[c] + .as_rect() + .row(y0); + } + } + for i in 0..num_ec { + if self.reference_frames[self.ec_blending_info[i].source as usize].is_some() { + bg[3 + i] = self.reference_frames[self.ec_blending_info[i].source as usize] + .as_ref() + .unwrap() + .frame[3 + i] + .as_rect() + .row(y0); + } + } + + for c in 0..num_c { + row[c][0..xsize].copy_from_slice(&bg[c][x0..x1]); + } + } +} + +#[cfg(test)] +mod test { + use test_log::test; + + use super::*; + use crate::error::Result; + use crate::util::test::read_headers_and_toc; + + #[test] + fn extend_consistency() -> Result<()> { + let (file_header, frame_header, _) = + read_headers_and_toc(include_bytes!("../../../resources/test/basic.jxl")).unwrap(); + let reference_frames: Vec<Option<ReferenceFrame>> = vec![None, None, None, None]; + crate::render::test::test_stage_consistency::<_, f32, f32>( + ExtendToImageDimensionsStage::new(&frame_header, &file_header, &reference_frames)?, + (500, 500), + 4, + ) + } +} diff --git a/third_party/rust/jxl/src/render/stages/from_linear.rs b/third_party/rust/jxl/src/render/stages/from_linear.rs @@ -0,0 +1,261 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +use crate::color::tf; +use crate::headers::color_encoding::CustomTransferFunction; +use crate::render::{RenderPipelineInPlaceStage, RenderPipelineStage}; + +/// Apply transfer function to display-referred linear color samples. +#[derive(Debug)] +pub struct FromLinearStage { + first_channel: usize, + tf: TransferFunction, +} + +impl FromLinearStage { + pub fn new(first_channel: usize, tf: TransferFunction) -> Self { + Self { first_channel, tf } + } + + pub fn sdr(first_channel: usize, tf: CustomTransferFunction) -> Self { + let tf = TransferFunction::try_from(tf).expect("transfer function is not an SDR one"); + Self::new(first_channel, tf) + } + + pub fn pq(first_channel: usize, intensity_target: f32) -> Self { + let tf = TransferFunction::Pq { intensity_target }; + Self::new(first_channel, tf) + } + + pub fn hlg(first_channel: usize, intensity_target: f32, luminance_rgb: [f32; 3]) -> Self { + let tf = TransferFunction::Hlg { + intensity_target, + luminance_rgb, + }; + Self::new(first_channel, tf) + } +} + +impl std::fmt::Display for FromLinearStage { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + let channel = self.first_channel; + write!( + f, + "Apply transfer function {:?} to channel [{},{},{}]", + self.tf, + channel, + channel + 1, + channel + 2 + ) + } +} + +impl RenderPipelineStage for FromLinearStage { + type Type = RenderPipelineInPlaceStage<f32>; + + fn uses_channel(&self, c: usize) -> bool { + (self.first_channel..self.first_channel + 3).contains(&c) + } + + fn process_row_chunk( + &self, + _position: (usize, usize), + xsize: usize, + row: &mut [&mut [f32]], + _state: Option<&mut dyn std::any::Any>, + ) { + let [row_r, row_g, row_b] = row else { + panic!( + "incorrect number of channels; expected 3, found {}", + row.len() + ); + }; + + match self.tf { + TransferFunction::Bt709 => { + for row in row { + tf::linear_to_bt709(&mut row[..xsize]); + } + } + TransferFunction::Srgb => { + for row in row { + tf::linear_to_srgb_fast(&mut row[..xsize]); + } + } + TransferFunction::Pq { intensity_target } => { + for row in row { + tf::linear_to_pq(intensity_target, &mut row[..xsize]); + } + } + TransferFunction::Hlg { + intensity_target, + luminance_rgb, + } => { + let rows = [ + &mut row_r[..xsize], + &mut row_g[..xsize], + &mut row_b[..xsize], + ]; + tf::hlg_display_to_scene(intensity_target, luminance_rgb, rows); + + tf::scene_to_hlg(&mut row_r[..xsize]); + tf::scene_to_hlg(&mut row_g[..xsize]); + tf::scene_to_hlg(&mut row_b[..xsize]); + } + TransferFunction::Gamma(g) => { + for row in row { + for v in &mut row[..xsize] { + *v = crate::util::fast_powf(*v, g); + } + } + } + } + } +} + +#[derive(Clone, Debug)] +pub enum TransferFunction { + Bt709, + Srgb, + Pq { + intensity_target: f32, + }, + Hlg { + intensity_target: f32, + luminance_rgb: [f32; 3], + }, + /// Inverse gamma in range `(0, 1]` + Gamma(f32), +} + +impl TryFrom<CustomTransferFunction> for TransferFunction { + type Error = (); + + fn try_from(ctf: CustomTransferFunction) -> Result<Self, ()> { + use crate::headers::color_encoding::TransferFunction; + + if ctf.have_gamma { + Ok(Self::Gamma(ctf.gamma())) + } else { + match ctf.transfer_function { + TransferFunction::BT709 => Ok(Self::Bt709), + TransferFunction::Unknown => Err(()), + TransferFunction::Linear => Ok(Self::Gamma(1.0)), + TransferFunction::SRGB => Ok(Self::Srgb), + TransferFunction::PQ => Err(()), + TransferFunction::DCI => Ok(Self::Gamma(2.6_f32.recip())), + TransferFunction::HLG => Err(()), + } + } + } +} + +#[cfg(test)] +mod test { + use test_log::test; + + use super::*; + use crate::error::Result; + use crate::image::Image; + use crate::render::test::make_and_run_simple_pipeline; + use crate::util::test::assert_all_almost_abs_eq; + + const LUMINANCE_BT2020: [f32; 3] = [0.2627, 0.678, 0.0593]; + + #[test] + fn consistency_hlg() -> Result<()> { + crate::render::test::test_stage_consistency::<_, f32, f32>( + FromLinearStage::hlg(0, 1000f32, LUMINANCE_BT2020), + (500, 500), + 3, + ) + } + + #[test] + fn consistency_pq() -> Result<()> { + crate::render::test::test_stage_consistency::<_, f32, f32>( + FromLinearStage::pq(0, 10000f32), + (500, 500), + 3, + ) + } + + #[test] + fn consistency_srgb() -> Result<()> { + crate::render::test::test_stage_consistency::<_, f32, f32>( + FromLinearStage::new(0, TransferFunction::Srgb), + (500, 500), + 3, + ) + } + + #[test] + fn consistency_bt709() -> Result<()> { + crate::render::test::test_stage_consistency::<_, f32, f32>( + FromLinearStage::new(0, TransferFunction::Bt709), + (500, 500), + 3, + ) + } + + #[test] + fn consistency_gamma22() -> Result<()> { + crate::render::test::test_stage_consistency::<_, f32, f32>( + FromLinearStage::new(0, TransferFunction::Gamma(0.4545455)), + (500, 500), + 3, + ) + } + + #[test] + fn sdr_white_hlg() -> Result<()> { + let intensity_target = 1000f32; + let input_r = Image::new_constant((1, 1), 0.203)?; + let input_g = Image::new_constant((1, 1), 0.203)?; + let input_b = Image::new_constant((1, 1), 0.203)?; + + // 75% HLG + let stage = FromLinearStage::hlg(0, intensity_target, LUMINANCE_BT2020); + let output = make_and_run_simple_pipeline::<_, f32, f32>( + stage, + &[input_r, input_g, input_b], + (1, 1), + 0, + 256, + )? + .1; + + assert_all_almost_abs_eq(output[0].as_rect().row(0), &[0.75], 1e-3); + assert_all_almost_abs_eq(output[1].as_rect().row(0), &[0.75], 1e-3); + assert_all_almost_abs_eq(output[2].as_rect().row(0), &[0.75], 1e-3); + + Ok(()) + } + + #[test] + fn sdr_white_pq() -> Result<()> { + let intensity_target = 1000f32; + let input_r = Image::new_constant((1, 1), 0.203)?; + let input_g = Image::new_constant((1, 1), 0.203)?; + let input_b = Image::new_constant((1, 1), 0.203)?; + + // 58% PQ + let stage = FromLinearStage::pq(0, intensity_target); + let output = make_and_run_simple_pipeline::<_, f32, f32>( + stage, + &[input_r, input_g, input_b], + (1, 1), + 0, + 256, + )? + .1; + + assert_all_almost_abs_eq(output[0].as_rect().row(0), &[0.58], 1e-3); + assert_all_almost_abs_eq(output[1].as_rect().row(0), &[0.58], 1e-3); + assert_all_almost_abs_eq(output[2].as_rect().row(0), &[0.58], 1e-3); + + Ok(()) + } +} diff --git a/third_party/rust/jxl/src/render/stages/gaborish.rs b/third_party/rust/jxl/src/render/stages/gaborish.rs @@ -0,0 +1,142 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +use crate::render::internal::InOutChannel; +use crate::render::{RenderPipelineInOutStage, RenderPipelineStage}; +use crate::simd::{F32SimdVec, simd_function}; + +/// Apply Gabor-like filter to a channel. +#[derive(Debug)] +pub struct GaborishStage { + channel: usize, + kernel_top_bottom: [f32; 3], + kernel_center: [f32; 3], +} + +impl GaborishStage { + pub fn new(channel: usize, weight1: f32, weight2: f32) -> Self { + let weight_total = 1.0 + weight1 * 4.0 + weight2 * 4.0; + let kernel_top_bottom = [weight2, weight1, weight2].map(|x| x / weight_total); + let kernel_center = [weight1, 1.0, weight1].map(|x| x / weight_total); + Self { + channel, + kernel_top_bottom, + kernel_center, + } + } +} + +impl std::fmt::Display for GaborishStage { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "Gaborish filter for channel {}", self.channel) + } +} + +simd_function!( + gaborish_process_dispatch, + d: D, + fn gaborish_process( + stage: &GaborishStage, + xsize: usize, + row: &mut [InOutChannel<f32, f32>], + ) { + let (rows_in, ref mut rows_out) = row[0]; + let row_out = &mut rows_out[0]; + for idx in (0..xsize).step_by(D::F32Vec::LEN) { + let mut sum = D::F32Vec::splat(d, 0f32); + let row_and_kernel = std::iter::zip( + rows_in, + [stage.kernel_top_bottom, stage.kernel_center, stage.kernel_top_bottom], + ); + for (row_in, kernel) in row_and_kernel { + for (dx, weight) in kernel.iter().enumerate() { + sum = D::F32Vec::load(d, &row_in[idx + dx..]).mul_add(D::F32Vec::splat(d, *weight), sum); + } + } + sum.store(&mut row_out[idx..]); + } + } +); + +impl RenderPipelineStage for GaborishStage { + type Type = RenderPipelineInOutStage<f32, f32, 1, 1, 0, 0>; + + fn uses_channel(&self, c: usize) -> bool { + c == self.channel + } + + fn process_row_chunk( + &self, + _position: (usize, usize), + xsize: usize, + row: &mut [InOutChannel<f32, f32>], + _state: Option<&mut dyn std::any::Any>, + ) { + gaborish_process_dispatch(self, xsize, row); + } +} + +#[cfg(test)] +mod test { + use test_log::test; + + use super::*; + use crate::error::Result; + use crate::image::Image; + use crate::render::test::{create_in_out_rows, make_and_run_simple_pipeline}; + use crate::simd::{ScalarDescriptor, SimdDescriptor, test_all_instruction_sets}; + use crate::util::test::assert_all_almost_abs_eq; + + #[test] + fn consistency() -> Result<()> { + crate::render::test::test_stage_consistency::<_, f32, f32>( + GaborishStage::new(0, 0.115169525, 0.061248592), + (500, 500), + 1, + ) + } + + #[test] + fn checkerboard() -> Result<()> { + let mut image = Image::new((2, 2))?; + image.as_rect_mut().row(0).copy_from_slice(&[0.0, 1.0]); + image.as_rect_mut().row(1).copy_from_slice(&[1.0, 0.0]); + + let stage = GaborishStage::new(0, 0.115169525, 0.061248592); + let (_, output) = + make_and_run_simple_pipeline::<_, f32, f32>(stage, &[image], (2, 2), 0, 256)?; + let output = output[0].as_rect(); + + assert_all_almost_abs_eq(output.row(0), &[0.20686048, 0.7931395], 1e-6); + assert_all_almost_abs_eq(output.row(1), &[0.7931395, 0.20686048], 1e-6); + + Ok(()) + } + + fn gaborish_process_scalar_equivalent<D: SimdDescriptor>(d: D) { + arbtest::arbtest(|u| { + let stage = GaborishStage::new(0, 0.115169525, 0.061248592); + create_in_out_rows!(u, 1, 1, rows, xsize); + gaborish_process(d, &stage, xsize, rows); + let simd_result: Vec<Vec<f32>> = rows[0] + .1 + .iter() + .map(|inner_slice| inner_slice.to_vec()) + .collect(); + gaborish_process(ScalarDescriptor::new().unwrap(), &stage, xsize, rows); + let scalar_result: Vec<Vec<f32>> = rows[0] + .1 + .iter() + .map(|inner_slice| inner_slice.to_vec()) + .collect(); + for (index, scalar_row) in scalar_result.iter().take(xsize).enumerate() { + assert_all_almost_abs_eq(scalar_row, &simd_result[index], 1e-8); + } + Ok(()) + }); + } + + test_all_instruction_sets!(gaborish_process_scalar_equivalent); +} diff --git a/third_party/rust/jxl/src/render/stages/mod.rs b/third_party/rust/jxl/src/render/stages/mod.rs @@ -0,0 +1,37 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +mod blending; +mod chroma_upsample; +mod convert; +mod epf; +mod extend; +mod from_linear; +mod gaborish; +mod nearest_neighbor; +mod noise; +mod patches; +mod splines; +mod spot; +mod to_linear; +mod upsample; +mod xyb; +mod ycbcr; + +pub use blending::*; +pub use chroma_upsample::*; +pub use convert::*; +pub use epf::*; +pub use extend::*; +pub use from_linear::*; +pub use gaborish::*; +pub use noise::*; +pub use patches::*; +pub use splines::*; +pub use spot::*; +pub use to_linear::{ToLinearStage, TransferFunction as ToLinearTransferFunction}; +pub use upsample::*; +pub use xyb::*; +pub use ycbcr::*; diff --git a/third_party/rust/jxl/src/render/stages/nearest_neighbor.rs b/third_party/rust/jxl/src/render/stages/nearest_neighbor.rs @@ -0,0 +1,94 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +use crate::render::{RenderPipelineInOutStage, RenderPipelineStage}; +pub struct NearestNeighbourUpsample { + channel: usize, +} + +impl NearestNeighbourUpsample { + // TODO(veluca): this probably should be #[cfg(test)] + #[allow(dead_code)] + pub fn new(channel: usize) -> NearestNeighbourUpsample { + NearestNeighbourUpsample { channel } + } +} + +impl std::fmt::Display for NearestNeighbourUpsample { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!( + f, + "2x2 nearest neighbour upsample of channel {}", + self.channel + ) + } +} + +impl RenderPipelineStage for NearestNeighbourUpsample { + type Type = RenderPipelineInOutStage<u8, u8, 0, 0, 1, 1>; + + fn uses_channel(&self, c: usize) -> bool { + c == self.channel + } + + fn process_row_chunk( + &self, + _position: (usize, usize), + xsize: usize, + row: &mut [(&[&[u8]], &mut [&mut [u8]])], + _state: Option<&mut dyn std::any::Any>, + ) { + let (input, output) = &mut row[0]; + for i in 0..xsize { + output[0][i * 2] = input[0][i]; + output[0][i * 2 + 1] = input[0][i]; + output[1][i * 2] = input[0][i]; + output[1][i * 2 + 1] = input[0][i]; + } + } +} + +#[cfg(test)] +mod test { + use super::*; + use crate::{error::Result, image::Image, render::test::make_and_run_simple_pipeline}; + use rand::SeedableRng; + use test_log::test; + + #[test] + fn nn_consistency() -> Result<()> { + crate::render::test::test_stage_consistency::<_, u8, u8>( + NearestNeighbourUpsample::new(0), + (500, 500), + 1, + ) + } + + #[test] + fn test_nn() -> Result<()> { + let image_size = (500, 400); + let input_size = (image_size.0 / 2, image_size.1 / 2); + let mut rng = rand_xorshift::XorShiftRng::seed_from_u64(0); + let input = vec![Image::<u8>::new_random(input_size, &mut rng)?]; + let stage = NearestNeighbourUpsample::new(0); + let output: Vec<Image<u8>> = + make_and_run_simple_pipeline(stage, &input, image_size, 0, 256)?.1; + assert_eq!(image_size, output[0].size()); + for y in 0..image_size.1 { + for x in 0..image_size.0 { + let ix = x / 2; + let iy = y / 2; + let i = input[0].as_rect().row(iy)[ix]; + let o = output[0].as_rect().row(y)[x]; + assert_eq!( + i, o, + "mismatch at output position {x}x{y}: {i} vs output {o}" + ); + } + } + + Ok(()) + } +} diff --git a/third_party/rust/jxl/src/render/stages/noise.rs b/third_party/rust/jxl/src/render/stages/noise.rs @@ -0,0 +1,300 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +use crate::{ + features::noise::Noise, + frame::color_correlation_map::ColorCorrelationParams, + render::{RenderPipelineInOutStage, RenderPipelineInPlaceStage, RenderPipelineStage}, +}; + +pub struct ConvolveNoiseStage { + channel: usize, +} + +impl ConvolveNoiseStage { + pub fn new(channel: usize) -> ConvolveNoiseStage { + ConvolveNoiseStage { channel } + } +} + +impl std::fmt::Display for ConvolveNoiseStage { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "convolve noise for channel {}", self.channel,) + } +} + +impl RenderPipelineStage for ConvolveNoiseStage { + type Type = RenderPipelineInOutStage<f32, f32, 2, 2, 0, 0>; + + fn uses_channel(&self, c: usize) -> bool { + c == self.channel + } + + fn process_row_chunk( + &self, + _position: (usize, usize), + xsize: usize, + row: &mut [(&[&[f32]], &mut [&mut [f32]])], + _state: Option<&mut dyn std::any::Any>, + ) { + let (input, output) = &mut row[0]; + for x in 0..xsize { + let mut others = 0.0; + for i in 0..5 { + let offset = (x as i32 + i) as usize; + others += input[0][offset]; + others += input[1][offset]; + others += input[3][offset]; + others += input[4][offset]; + } + others += input[2][x]; + others += input[2][x + 1]; + others += input[2][x + 3]; + others += input[2][x + 4]; + output[0][x] = others * 0.16 + input[2][x + 2] * -3.84; + } + } +} + +pub struct AddNoiseStage { + noise: Noise, + first_channel: usize, + color_correlation: ColorCorrelationParams, +} + +impl AddNoiseStage { + #[allow(dead_code)] + pub fn new( + noise: Noise, + color_correlation: ColorCorrelationParams, + first_channel: usize, + ) -> AddNoiseStage { + assert!(first_channel > 2); + AddNoiseStage { + noise, + first_channel, + color_correlation, + } + } +} + +impl std::fmt::Display for AddNoiseStage { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!( + f, + "add noise for channels [{},{},{}]", + self.first_channel, + self.first_channel + 1, + self.first_channel + 2 + ) + } +} + +impl RenderPipelineStage for AddNoiseStage { + type Type = RenderPipelineInPlaceStage<f32>; + + fn uses_channel(&self, c: usize) -> bool { + c < 3 || (c >= self.first_channel && c < self.first_channel + 3) + } + + fn process_row_chunk( + &self, + _position: (usize, usize), + xsize: usize, + row: &mut [&mut [f32]], + _state: Option<&mut dyn std::any::Any>, + ) { + let norm_const = 0.22; + let ytox = self.color_correlation.y_to_x_lf(); + let ytob = self.color_correlation.y_to_b_lf(); + for x in 0..xsize { + let row_rnd_r = row[3][x]; + let row_rnd_g = row[4][x]; + let row_rnd_c = row[5][x]; + let vx = row[0][x]; + let vy = row[1][x]; + let in_g = vy - vx; + let in_r = vy + vx; + let noise_strength_g = self.noise.strength(in_g * 0.5); + let noise_strength_r = self.noise.strength(in_r * 0.5); + let addit_rnd_noise_red = row_rnd_r * norm_const; + let addit_rnd_noise_green = row_rnd_g * norm_const; + let addit_rnd_noise_correlated = row_rnd_c * norm_const; + let k_rg_corr = 0.9921875; + let k_rgn_corr = 0.0078125; + let red_noise = noise_strength_r + * (k_rgn_corr * addit_rnd_noise_red + k_rg_corr * addit_rnd_noise_correlated); + let green_noise = noise_strength_g + * (k_rgn_corr * addit_rnd_noise_green + k_rg_corr * addit_rnd_noise_correlated); + let rg_noise = red_noise + green_noise; + row[0][x] += ytox * rg_noise + red_noise - green_noise; + row[1][x] += rg_noise; + row[2][x] += ytob * rg_noise; + } + } +} + +#[cfg(test)] +mod test { + use crate::{ + error::Result, + features::noise::Noise, + frame::color_correlation_map::ColorCorrelationParams, + image::Image, + render::{ + stages::noise::{AddNoiseStage, ConvolveNoiseStage}, + test::make_and_run_simple_pipeline, + }, + util::test::assert_almost_abs_eq, + }; + use test_log::test; + + // TODO(firsching): Add more relevant ConvolveNoise tests as per discussions in https://github.com/libjxl/jxl-rs/pull/60. + + #[test] + fn convolve_noise_process_row_chunk() -> Result<()> { + let input: Image<f32> = Image::new_range((2, 2), 0.0, 1.0)?; + let stage = ConvolveNoiseStage::new(0); + let output: Vec<Image<f32>> = + make_and_run_simple_pipeline(stage, &[input], (2, 2), 0, 256)?.1; + let rect = output[0].as_rect(); + assert_almost_abs_eq(rect.row(0)[0], 7.2, 1e-6); + assert_almost_abs_eq(rect.row(0)[1], 2.4, 1e-6); + assert_almost_abs_eq(rect.row(1)[0], -2.4, 1e-6); + assert_almost_abs_eq(rect.row(1)[1], -7.2, 1e-6); + Ok(()) + } + + #[test] + fn convolve_noise_consistency() -> Result<()> { + crate::render::test::test_stage_consistency::<_, f32, f32>( + ConvolveNoiseStage::new(0), + (500, 500), + 1, + ) + } + + // TODO(firsching): Add more relevant AddNoise tests as per discussions in https://github.com/libjxl/jxl-rs/pull/60. + + #[test] + fn add_noise_process_row_chunk() -> Result<()> { + let xsize = 8; + let ysize = 8; + let input_c0: Image<f32> = Image::new_range((xsize, ysize), 0.1, 0.1)?; + let input_c1: Image<f32> = Image::new_range((xsize, ysize), 0.1, 0.1)?; + let input_c2: Image<f32> = Image::new_range((xsize, ysize), 0.1, 0.1)?; + let input_c3: Image<f32> = Image::new_range((xsize, ysize), 0.1, 0.1)?; + let input_c4: Image<f32> = Image::new_range((xsize, ysize), 0.1, 0.1)?; + let input_c5: Image<f32> = Image::new_range((xsize, ysize), 0.1, 0.1)?; + let stage = AddNoiseStage::new( + Noise { + lut: [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0], + }, + ColorCorrelationParams::default(), + 3, + ); + let output: Vec<Image<f32>> = make_and_run_simple_pipeline( + stage, + &[input_c0, input_c1, input_c2, input_c3, input_c4, input_c5], + (xsize, ysize), + 0, + 256, + )? + .1; + // Golden data generated by libjxl. + let want_out = [ + [ + [ + 0.100000, 0.200000, 0.300000, 0.400000, 0.500000, 0.600000, 0.700000, 0.800000, + ], + [0.900000, 1.000000, 1.1, 1.2, 1.3, 1.4, 1.5, 1.6], + [1.7, 1.8, 1.9, 2.000000, 2.1, 2.2, 2.3, 2.4], + [ + 2.5, 2.6, 2.7, 2.799999, 2.899999, 2.999999, 3.099999, 3.199999, + ], + [ + 3.299999, 3.399999, 3.499999, 3.599999, 3.699999, 3.799999, 3.899998, 3.999998, + ], + [ + 4.099998, 4.199998, 4.299998, 4.399998, 4.499998, 4.599998, 4.699998, 4.799998, + ], + [ + 4.899998, 4.999998, 5.099998, 5.199997, 5.299997, 5.399997, 5.499997, 5.599997, + ], + [ + 5.699997, 5.799997, 5.899997, 5.999997, 6.099997, 6.199996, 6.299996, 6.399996, + ], + ], + [ + [ + 0.144000, 0.288000, 0.432000, 0.576000, 0.720000, 0.864000, 1.008, 1.152, + ], + [1.296, 1.44, 1.584, 1.728, 1.872, 2.016, 2.16, 2.304], + [2.448, 2.592, 2.736001, 2.88, 3.024, 3.168, 3.312, 3.456], + [ + 3.6, 3.743999, 3.888, 4.031999, 4.175999, 4.319999, 4.463999, 4.607999, + ], + [ + 4.751998, 4.895998, 5.039998, 5.183998, 5.327998, 5.471998, 5.615998, 5.759997, + ], + [ + 5.903998, 6.047997, 6.191998, 6.335998, 6.479997, 6.623997, 6.767997, 6.911997, + ], + [ + 7.055997, 7.199996, 7.343997, 7.487996, 7.631996, 7.775996, 7.919996, 8.063995, + ], + [ + 8.207995, 8.351995, 8.495996, 8.639996, 8.783995, 8.927995, 9.071995, 9.215995, + ], + ], + [ + [ + 0.144000, 0.288000, 0.432000, 0.576000, 0.720000, 0.864000, 1.008, 1.152, + ], + [1.296, 1.44, 1.584, 1.728, 1.872, 2.016, 2.16, 2.304], + [2.448, 2.592, 2.736001, 2.88, 3.024, 3.168, 3.312, 3.456], + [ + 3.6, 3.743999, 3.888, 4.031999, 4.175999, 4.319999, 4.463999, 4.607999, + ], + [ + 4.751998, 4.895998, 5.039998, 5.183998, 5.327998, 5.471998, 5.615998, 5.759997, + ], + [ + 5.903998, 6.047997, 6.191998, 6.335998, 6.479997, 6.623997, 6.767997, 6.911997, + ], + [ + 7.055997, 7.199996, 7.343997, 7.487996, 7.631996, 7.775996, 7.919996, 8.063995, + ], + [ + 8.207995, 8.351995, 8.495996, 8.639996, 8.783995, 8.927995, 9.071995, 9.215995, + ], + ], + ]; + for c in 0..3 { + let rect = output[c].as_rect(); + for y in 0..rect.size().1 { + for x in 0..rect.size().0 { + assert_almost_abs_eq(rect.row(y)[x], want_out[c][y][x], 1e-5); + } + } + } + Ok(()) + } + + #[test] + fn add_noise_consistency() -> Result<()> { + crate::render::test::test_stage_consistency::<_, f32, f32>( + AddNoiseStage::new( + Noise { + lut: [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + }, + ColorCorrelationParams::default(), + 3, + ), + (500, 500), + 6, + ) + } +} diff --git a/third_party/rust/jxl/src/render/stages/patches.rs b/third_party/rust/jxl/src/render/stages/patches.rs @@ -0,0 +1,63 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +use std::{any::Any, sync::Arc}; + +use crate::{ + features::patches::PatchesDictionary, + frame::ReferenceFrame, + headers::extra_channels::ExtraChannelInfo, + render::{RenderPipelineInPlaceStage, RenderPipelineStage}, + util::NewWithCapacity as _, +}; + +pub struct PatchesStage { + pub patches: PatchesDictionary, + pub extra_channels: Vec<ExtraChannelInfo>, + pub decoder_state: Arc<Vec<Option<ReferenceFrame>>>, +} + +impl std::fmt::Display for PatchesStage { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "patches") + } +} + +impl RenderPipelineStage for PatchesStage { + type Type = RenderPipelineInPlaceStage<f32>; + + fn uses_channel(&self, c: usize) -> bool { + c < 3 + self.extra_channels.len() + } + + fn process_row_chunk( + &self, + position: (usize, usize), + xsize: usize, + row: &mut [&mut [f32]], + state: Option<&mut dyn Any>, + ) { + let state: &mut Vec<usize> = state.unwrap().downcast_mut().unwrap(); + self.patches.add_one_row( + row, + position, + xsize, + &self.extra_channels, + &self.decoder_state, + state, + ); + } + + fn init_local_state(&self) -> crate::error::Result<Option<Box<dyn Any>>> { + let patches_for_row_result = Vec::<usize>::new_with_capacity(self.patches.positions.len())?; + Ok(Some(Box::new(patches_for_row_result) as Box<dyn Any>)) + } +} + +#[cfg(test)] +mod test { + #[test] + fn process_row_chunk() {} +} diff --git a/third_party/rust/jxl/src/render/stages/splines.rs b/third_party/rust/jxl/src/render/stages/splines.rs @@ -0,0 +1,164 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +use crate::{ + features::spline::Splines, + frame::color_correlation_map::ColorCorrelationParams, + render::{RenderPipelineInPlaceStage, RenderPipelineStage}, +}; + +pub struct SplinesStage { + splines: Splines, +} + +impl SplinesStage { + pub fn new( + mut splines: Splines, + frame_size: (usize, usize), + color_correlation_params: &ColorCorrelationParams, + ) -> Self { + splines + .initialize_draw_cache( + frame_size.0 as u64, + frame_size.1 as u64, + color_correlation_params, + ) + .unwrap(); + SplinesStage { splines } + } +} + +impl std::fmt::Display for SplinesStage { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "splines") + } +} + +impl RenderPipelineStage for SplinesStage { + type Type = RenderPipelineInPlaceStage<f32>; + + fn uses_channel(&self, c: usize) -> bool { + c < 3 + } + + fn process_row_chunk( + &self, + position: (usize, usize), + xsize: usize, + row: &mut [&mut [f32]], + _state: Option<&mut dyn std::any::Any>, + ) { + self.splines.draw_segments(row, position, xsize); + } +} + +#[cfg(test)] +mod test { + use crate::features::spline::{Point, QuantizedSpline, Splines}; + use crate::frame::color_correlation_map::ColorCorrelationParams; + use crate::render::test::make_and_run_simple_pipeline; + use crate::util::test::{self, assert_all_almost_abs_eq, read_pfm}; + use crate::{error::Result, image::Image, render::stages::splines::SplinesStage}; + use test_log::test; + + #[test] + fn splines_process_row_chunk() -> Result<(), test::Error> { + let want_image = read_pfm(include_bytes!("../../../resources/test/splines.pfm"))?; + let target_images = [ + Image::<f32>::new_constant((320, 320), 0.0)?, + Image::<f32>::new_constant((320, 320), 0.0)?, + Image::<f32>::new_constant((320, 320), 0.0)?, + ]; + let size = target_images[0].size(); + let mut splines = Splines::create( + 0, + vec![QuantizedSpline { + control_points: vec![ + (109, 105), + (-130, -261), + (-66, 193), + (227, -52), + (-170, 290), + ], + color_dct: [ + [ + 168, 119, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 0, 0, + ], + [ + 9, 0, 7, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 0, + ], + [ + -10, 7, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 0, 0, + ], + ], + sigma_dct: [ + 4, 0, 0, 0, 0, 0, 0, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, + ], + }], + vec![Point { x: 9.0, y: 54.0 }], + ); + splines.initialize_draw_cache( + size.0 as u64, + size.1 as u64, + &ColorCorrelationParams::default(), + )?; + let stage = SplinesStage { splines }; + let output: Vec<Image<f32>> = + make_and_run_simple_pipeline(stage, &target_images, size, 0, 256)?.1; + for c in 0..3 { + for row in 0..size.1 { + assert_all_almost_abs_eq( + output[c].as_rect().row(row), + want_image[c].as_rect().row(row), + 1e-3, + ); + } + } + Ok(()) + } + + #[test] + fn splines_consistency() -> Result<()> { + let mut splines = Splines::create( + 0, + vec![QuantizedSpline { + control_points: vec![ + (109, 105), + (-130, -261), + (-66, 193), + (227, -52), + (-170, 290), + ], + color_dct: [ + [ + 168, 119, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 0, 0, + ], + [ + 9, 0, 7, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 0, + ], + [ + -10, 7, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 0, 0, + ], + ], + sigma_dct: [ + 4, 0, 0, 0, 0, 0, 0, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, + ], + }], + vec![Point { x: 9.0, y: 54.0 }], + ); + splines.initialize_draw_cache(500, 500, &ColorCorrelationParams::default())?; + let stage = SplinesStage { splines }; + + crate::render::test::test_stage_consistency::<_, f32, f32>(stage, (500, 500), 6) + } +} diff --git a/third_party/rust/jxl/src/render/stages/spot.rs b/third_party/rust/jxl/src/render/stages/spot.rs @@ -0,0 +1,128 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +use crate::render::{RenderPipelineInPlaceStage, RenderPipelineStage}; + +/// Render spot color +pub struct SpotColorStage { + /// Spot color channel index + spot_c: usize, + /// Spot color in linear RGBA + spot_color: [f32; 4], +} + +impl std::fmt::Display for SpotColorStage { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "spot color stage for channel {}", self.spot_c) + } +} + +impl SpotColorStage { + #[allow(unused, reason = "remove once we actually use this")] + pub fn new(spot_c_offset: usize, spot_color: [f32; 4]) -> Self { + Self { + spot_c: 3 + spot_c_offset, + spot_color, + } + } +} + +impl RenderPipelineStage for SpotColorStage { + type Type = RenderPipelineInPlaceStage<f32>; + + fn uses_channel(&self, c: usize) -> bool { + c < 3 || c == self.spot_c + } + + // `row` should only contain color channels and the spot channel. + fn process_row_chunk( + &self, + _position: (usize, usize), + xsize: usize, + row: &mut [&mut [f32]], + _state: Option<&mut dyn std::any::Any>, + ) { + let [row_r, row_g, row_b, row_s] = row else { + panic!( + "incorrect number of channels; expected 4, found {}", + row.len() + ); + }; + + let scale = self.spot_color[3]; + assert!( + xsize <= row_r.len() + && xsize <= row_g.len() + && xsize <= row_b.len() + && xsize <= row_s.len() + ); + for idx in 0..xsize { + let mix = scale * row_s[idx]; + row_r[idx] = mix * self.spot_color[0] + (1.0 - mix) * row_r[idx]; + row_g[idx] = mix * self.spot_color[1] + (1.0 - mix) * row_g[idx]; + row_b[idx] = mix * self.spot_color[2] + (1.0 - mix) * row_b[idx]; + } + } +} + +#[cfg(test)] +mod test { + use test_log::test; + + use super::*; + use crate::error::Result; + use crate::image::Image; + use crate::render::test::make_and_run_simple_pipeline; + use crate::util::test::assert_all_almost_abs_eq; + + #[test] + fn consistency() -> Result<()> { + crate::render::test::test_stage_consistency::<_, f32, f32>( + SpotColorStage::new(0, [0.0; 4]), + (500, 500), + 4, + ) + } + + #[test] + fn srgb_primaries() -> Result<()> { + let mut input_r = Image::new((3, 1))?; + let mut input_g = Image::new((3, 1))?; + let mut input_b = Image::new((3, 1))?; + let mut input_s = Image::new((3, 1))?; + input_r + .as_rect_mut() + .row(0) + .copy_from_slice(&[1.0, 0.0, 0.0]); + input_g + .as_rect_mut() + .row(0) + .copy_from_slice(&[0.0, 1.0, 0.0]); + input_b + .as_rect_mut() + .row(0) + .copy_from_slice(&[0.0, 0.0, 1.0]); + input_s + .as_rect_mut() + .row(0) + .copy_from_slice(&[1.0, 1.0, 1.0]); + + let stage = SpotColorStage::new(0, [0.5; 4]); + let output = make_and_run_simple_pipeline::<_, f32, f32>( + stage, + &[input_r, input_g, input_b, input_s], + (3, 1), + 0, + 256, + )? + .1; + + assert_all_almost_abs_eq(output[0].as_rect().row(0), &[0.75, 0.25, 0.25], 1e-6); + assert_all_almost_abs_eq(output[1].as_rect().row(0), &[0.25, 0.75, 0.25], 1e-6); + assert_all_almost_abs_eq(output[2].as_rect().row(0), &[0.25, 0.25, 0.75], 1e-6); + + Ok(()) + } +} diff --git a/third_party/rust/jxl/src/render/stages/to_linear.rs b/third_party/rust/jxl/src/render/stages/to_linear.rs @@ -0,0 +1,232 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +use crate::color::tf; +use crate::headers::color_encoding::CustomTransferFunction; +use crate::render::stages::from_linear; +use crate::render::{RenderPipelineInPlaceStage, RenderPipelineStage}; + +/// Convert encoded non-linear color samples to display-referred linear color samples. +#[derive(Debug)] +pub struct ToLinearStage { + first_channel: usize, + tf: TransferFunction, +} + +impl ToLinearStage { + pub fn new(first_channel: usize, tf: TransferFunction) -> Self { + Self { first_channel, tf } + } + + #[allow(unused, reason = "tirr-c: remove once we use this!")] + pub fn sdr(first_channel: usize, tf: CustomTransferFunction) -> Self { + let tf = TransferFunction::try_from(tf).expect("transfer function is not an SDR one"); + Self::new(first_channel, tf) + } + + #[allow(unused, reason = "tirr-c: remove once we use this!")] + pub fn pq(first_channel: usize, intensity_target: f32) -> Self { + let tf = TransferFunction::Pq { intensity_target }; + Self::new(first_channel, tf) + } + + #[allow(unused, reason = "tirr-c: remove once we use this!")] + pub fn hlg(first_channel: usize, intensity_target: f32, luminance_rgb: [f32; 3]) -> Self { + let tf = TransferFunction::Hlg { + intensity_target, + luminance_rgb, + }; + Self::new(first_channel, tf) + } +} + +impl std::fmt::Display for ToLinearStage { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + let channel = self.first_channel; + write!( + f, + "Convert transfer function {:?} to display-referred linear TF for channel [{},{},{}]", + self.tf, + channel, + channel + 1, + channel + 2 + ) + } +} + +impl RenderPipelineStage for ToLinearStage { + type Type = RenderPipelineInPlaceStage<f32>; + + fn uses_channel(&self, c: usize) -> bool { + (self.first_channel..self.first_channel + 3).contains(&c) + } + + fn process_row_chunk( + &self, + _position: (usize, usize), + xsize: usize, + row: &mut [&mut [f32]], + _state: Option<&mut dyn std::any::Any>, + ) { + let [row_r, row_g, row_b] = row else { + panic!( + "incorrect number of channels; expected 3, found {}", + row.len() + ); + }; + + match self.tf { + TransferFunction::Bt709 => { + for row in row { + tf::bt709_to_linear(&mut row[..xsize]); + } + } + TransferFunction::Srgb => { + for row in row { + tf::srgb_to_linear(&mut row[..xsize]); + } + } + TransferFunction::Pq { intensity_target } => { + for row in row { + tf::pq_to_linear(intensity_target, &mut row[..xsize]); + } + } + TransferFunction::Hlg { + intensity_target, + luminance_rgb, + } => { + tf::hlg_to_scene(&mut row_r[..xsize]); + tf::hlg_to_scene(&mut row_g[..xsize]); + tf::hlg_to_scene(&mut row_b[..xsize]); + + let rows = [ + &mut row_r[..xsize], + &mut row_g[..xsize], + &mut row_b[..xsize], + ]; + tf::hlg_scene_to_display(intensity_target, luminance_rgb, rows); + } + TransferFunction::Gamma(g) => { + for row in row { + for v in &mut row[..xsize] { + *v = crate::util::fast_powf(*v, g); + } + } + } + } + } +} + +pub type TransferFunction = from_linear::TransferFunction; + +#[cfg(test)] +mod test { + use test_log::test; + + use super::*; + use crate::error::Result; + use crate::image::Image; + use crate::render::test::make_and_run_simple_pipeline; + use crate::util::test::assert_all_almost_abs_eq; + + const LUMINANCE_BT2020: [f32; 3] = [0.2627, 0.678, 0.0593]; + + #[test] + fn consistency_hlg() -> Result<()> { + crate::render::test::test_stage_consistency::<_, f32, f32>( + ToLinearStage::hlg(0, 1000f32, LUMINANCE_BT2020), + (500, 500), + 3, + ) + } + + #[test] + fn consistency_pq() -> Result<()> { + crate::render::test::test_stage_consistency::<_, f32, f32>( + ToLinearStage::pq(0, 10000f32), + (500, 500), + 3, + ) + } + + #[test] + fn consistency_srgb() -> Result<()> { + crate::render::test::test_stage_consistency::<_, f32, f32>( + ToLinearStage::new(0, TransferFunction::Srgb), + (500, 500), + 3, + ) + } + + #[test] + fn consistency_bt709() -> Result<()> { + crate::render::test::test_stage_consistency::<_, f32, f32>( + ToLinearStage::new(0, TransferFunction::Bt709), + (500, 500), + 3, + ) + } + + #[test] + fn consistency_gamma22() -> Result<()> { + crate::render::test::test_stage_consistency::<_, f32, f32>( + ToLinearStage::new(0, TransferFunction::Gamma(0.4545455)), + (500, 500), + 3, + ) + } + + #[test] + fn sdr_white_hlg() -> Result<()> { + let intensity_target = 1000f32; + // Reversed version of FromLinear test + let input_r = Image::new_constant((1, 1), 0.75)?; + let input_g = Image::new_constant((1, 1), 0.75)?; + let input_b = Image::new_constant((1, 1), 0.75)?; + + // 75% HLG + let stage = ToLinearStage::hlg(0, intensity_target, LUMINANCE_BT2020); + let output = make_and_run_simple_pipeline::<_, f32, f32>( + stage, + &[input_r, input_g, input_b], + (1, 1), + 0, + 256, + )? + .1; + + assert_all_almost_abs_eq(output[0].as_rect().row(0), &[0.203], 1e-3); + assert_all_almost_abs_eq(output[1].as_rect().row(0), &[0.203], 1e-3); + assert_all_almost_abs_eq(output[2].as_rect().row(0), &[0.203], 1e-3); + + Ok(()) + } + + #[test] + fn sdr_white_pq() -> Result<()> { + let intensity_target = 1000f32; + // Reversed version of FromLinear test + let input_r = Image::new_constant((1, 1), 0.5807)?; + let input_g = Image::new_constant((1, 1), 0.5807)?; + let input_b = Image::new_constant((1, 1), 0.5807)?; + + // 58% PQ + let stage = ToLinearStage::pq(0, intensity_target); + let output = make_and_run_simple_pipeline::<_, f32, f32>( + stage, + &[input_r, input_g, input_b], + (1, 1), + 0, + 256, + )? + .1; + + assert_all_almost_abs_eq(output[0].as_rect().row(0), &[0.203], 1e-3); + assert_all_almost_abs_eq(output[1].as_rect().row(0), &[0.203], 1e-3); + assert_all_almost_abs_eq(output[2].as_rect().row(0), &[0.203], 1e-3); + + Ok(()) + } +} diff --git a/third_party/rust/jxl/src/render/stages/upsample.rs b/third_party/rust/jxl/src/render/stages/upsample.rs @@ -0,0 +1,500 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +use crate::{ + headers::CustomTransformData, + render::{RenderPipelineInOutStage, RenderPipelineStage}, +}; + +pub struct Upsample<const N: usize, const SHIFT: u8> { + kernel: [[[[f32; 5]; 5]; N]; N], + channel: usize, +} + +impl<const N: usize, const SHIFT: u8> Upsample<N, SHIFT> { + pub fn new(ups_factors: &CustomTransformData, channel: usize) -> Self { + const { assert!(SHIFT >= 1 && SHIFT <= 3) } + const { assert!(1 << SHIFT == N) } + + let weights: &[f32] = match N { + 2 => &ups_factors.weights2, + 4 => &ups_factors.weights4, + 8 => &ups_factors.weights8, + _ => unreachable!(), + }; + + let mut kernel = [[[[0.0; 5]; 5]; N]; N]; + let n = N / 2; + for i in 0..5 * n { + for j in 0..5 * n { + let y = i.min(j); + let x = i.max(j); + let y = y as isize; + let x = x as isize; + let n = n as isize; + let index = (5 * n * y - y * (y - 1) / 2 + x - y) as usize; + // Filling in the top left corner from the weights + kernel[j / 5][i / 5][j % 5][i % 5] = weights[index]; + // Mirroring to get the rest of the kernel. + kernel[(2 * n as usize - 1) - j / 5][i / 5][4 - (j % 5)][i % 5] = weights[index]; + kernel[j / 5][(2 * n as usize - 1) - i / 5][j % 5][4 - (i % 5)] = weights[index]; + kernel[(2 * n as usize - 1) - j / 5][(2 * n as usize - 1) - i / 5][4 - (j % 5)] + [4 - (i % 5)] = weights[index]; + } + } + + Self { kernel, channel } + } +} + +impl<const N: usize, const SHIFT: u8> std::fmt::Display for Upsample<N, SHIFT> { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{N}x{N} upsampling of channel {}", self.channel) + } +} + +impl<const N: usize, const SHIFT: u8> RenderPipelineStage for Upsample<N, SHIFT> { + type Type = RenderPipelineInOutStage<f32, f32, 2, 2, SHIFT, SHIFT>; + + fn uses_channel(&self, c: usize) -> bool { + c == self.channel + } + /// Processes a chunk of a row, applying NxN upsampling using a 5x5 kernel. + /// Each input value expands into a NxN region in the output, based on neighboring inputs. + fn process_row_chunk( + &self, + _position: (usize, usize), + xsize: usize, + row: &mut [(&[&[f32]], &mut [&mut [f32]])], + _state: Option<&mut dyn std::any::Any>, + ) { + let (input, output) = &mut row[0]; + + for x in 0..xsize { + // Upsample this input value into a NxN region in the output + let mut minval = input[0][x]; + let mut maxval = minval; + for di in 0..N { + for dj in 0..N { + // Iterate over the input rows and columns + let mut output_val = 0.0; + for i in 0..5 { + for j in 0..5 { + let input_value = input[i][j + x]; + output_val += input_value * self.kernel[di][dj][i % 5][j % 5]; + minval = input_value.min(minval); + maxval = input_value.max(maxval); + } + } + output[di][dj + N * x] = output_val.clamp(minval, maxval); + } + } + } + } +} + +pub type Upsample2x = Upsample<2, 1>; +pub type Upsample4x = Upsample<4, 2>; +pub type Upsample8x = Upsample<8, 3>; + +#[cfg(test)] +mod test { + + use super::*; + use crate::{ + error::Result, headers::CustomTransformDataNonserialized, image::Image, + render::test::make_and_run_simple_pipeline, util::test::assert_almost_abs_eq, + }; + use test_log::test; + + fn ups_factors() -> CustomTransformData { + CustomTransformData::default(&CustomTransformDataNonserialized { xyb_encoded: true }) + } + + #[test] + fn upsample2x_consistency() -> Result<()> { + crate::render::test::test_stage_consistency::<_, f32, f32>( + Upsample2x::new(&ups_factors(), 0), + (500, 500), + 1, + ) + } + + #[test] + fn upsample4x_consistency() -> Result<()> { + crate::render::test::test_stage_consistency::<_, f32, f32>( + Upsample4x::new(&ups_factors(), 0), + (500, 500), + 1, + ) + } + + #[test] + fn upsample8x_consistency() -> Result<()> { + crate::render::test::test_stage_consistency::<_, f32, f32>( + Upsample8x::new(&ups_factors(), 0), + (504, 504), + 1, + ) + } + + #[test] + fn upsample2x_constant() -> Result<()> { + let image_size = (238, 412); + let input_size = (image_size.0 / 2, image_size.1 / 2); + let val = 0.777f32; + let input = Image::new_constant(input_size, val)?; + let stage = Upsample2x::new(&ups_factors(), 0); + let output: Vec<Image<f32>> = + make_and_run_simple_pipeline(stage, &[input], image_size, 0, 123)?.1; + for x in 0..image_size.0 { + for y in 0..image_size.1 { + assert_almost_abs_eq(output[0].as_rect().row(y)[x], val, 0.0000001); + } + } + Ok(()) + } + + #[test] + fn upsample4x_constant() -> Result<()> { + let image_size = (240, 412); + let input_size = (image_size.0 / 4, image_size.1 / 4); + let val = 0.777f32; + let input = Image::new_constant(input_size, val)?; + let stage = Upsample4x::new(&ups_factors(), 0); + let output: Vec<Image<f32>> = + make_and_run_simple_pipeline(stage, &[input], image_size, 0, 123)?.1; + for x in 0..image_size.0 { + for y in 0..image_size.1 { + assert_almost_abs_eq(output[0].as_rect().row(y)[x], val, 0.00001); + } + } + Ok(()) + } + + #[test] + fn upsample8x_constant() -> Result<()> { + let image_size = (240, 416); + let input_size = (image_size.0 / 8, image_size.1 / 8); + let val = 0.777f32; + let input = Image::new_constant(input_size, val)?; + let stage = Upsample8x::new(&ups_factors(), 0); + let output: Vec<Image<f32>> = + make_and_run_simple_pipeline(stage, &[input], image_size, 0, 123)?.1; + for x in 0..image_size.0 { + for y in 0..image_size.1 { + assert_almost_abs_eq(output[0].as_rect().row(y)[x], val, 0.00001); + } + } + Ok(()) + } + + #[test] + fn test_upsample2() -> Result<()> { + let eps = 0.0000001; + let mut input = Image::new((7, 7))?; + // Put a single "1.0" in the middle of the image. + input.as_rect_mut().row(3)[3] = 1.0f32; + let ups_factors = ups_factors(); + let stage = Upsample2x::new(&ups_factors, 0); + let output: Vec<Image<f32>> = + make_and_run_simple_pipeline(stage, &[input], (14, 14), 0, 77)?.1; + assert_eq!(output[0].as_rect().size(), (14, 14)); + // Check we have a border with zeros + for i in 0..14 { + for j in 0..2 { + assert_almost_abs_eq(output[0].as_rect().row(j)[i], 0.0, eps); + assert_almost_abs_eq(output[0].as_rect().row(i)[j], 0.0, eps); + assert_almost_abs_eq(output[0].as_rect().row(13 - j)[i], 0.0, eps); + assert_almost_abs_eq(output[0].as_rect().row(i)[13 - j], 0.0, eps); + } + } + // Define the mapping for the symmetric top-left kernel + let index_map = [ + [0, 1, 2, 3, 4], + [1, 5, 6, 7, 8], + [2, 6, 9, 10, 11], + [3, 7, 10, 12, 13], + [4, 8, 11, 13, 14], + ]; + + // Validate weights from the kernel + let kernel_size = 5; + let kernel_offset = 2; + let weights = &ups_factors.weights2; + for di in 0..2 { + for dj in 0..2 { + for i in 0..kernel_size { + for j in 0..kernel_size { + let output_value = output[0].as_rect().row(kernel_offset + di + 2 * i) + [kernel_offset + dj + 2 * j]; + let mapped_i = if di == 0 { kernel_size - 1 - i } else { i }; + let mapped_j = if dj == 0 { kernel_size - 1 - j } else { j }; + let weight_index = index_map[mapped_i][mapped_j]; + assert_almost_abs_eq( + output_value, + weights[weight_index].clamp(0.0, 1.0), + eps, + ); + } + } + } + } + + Ok(()) + } + + #[test] + fn test_upsample4() -> Result<()> { + let eps = 0.0000001; + let mut input = Image::new((7, 7))?; + // Put a single "1.0" in the middle of the image. + input.as_rect_mut().row(3)[3] = 1.0f32; + let ups_factors = ups_factors(); + let stage = Upsample4x::new(&ups_factors, 0); + let output: Vec<Image<f32>> = + make_and_run_simple_pipeline(stage, &[input], (28, 28), 0, 1024)?.1; + + assert_eq!(output[0].as_rect().size(), (28, 28)); + + // Check we have a border with zeros + for i in 0..28 { + for j in 0..4 { + assert_almost_abs_eq(output[0].as_rect().row(j)[i], 0.0, eps); + assert_almost_abs_eq(output[0].as_rect().row(i)[j], 0.0, eps); + assert_almost_abs_eq(output[0].as_rect().row(27 - j)[i], 0.0, eps); + assert_almost_abs_eq(output[0].as_rect().row(i)[27 - j], 0.0, eps); + } + } + + // Define the mapping for the symmetric top-left kernel + let index_map = [ + [0, 1, 2, 3, 4, 5, 6, 7, 8, 9], + [1, 10, 11, 12, 13, 14, 15, 16, 17, 18], + [2, 11, 19, 20, 21, 22, 23, 24, 25, 26], + [3, 12, 20, 27, 28, 29, 30, 31, 32, 33], + [4, 13, 21, 28, 34, 35, 36, 37, 38, 39], + [5, 14, 22, 29, 35, 40, 41, 42, 43, 44], + [6, 15, 23, 30, 36, 41, 45, 46, 47, 48], + [7, 16, 24, 31, 37, 42, 46, 49, 50, 51], + [8, 17, 25, 32, 38, 43, 47, 50, 52, 53], + [9, 18, 26, 33, 39, 44, 48, 51, 53, 54], + ]; + + // Validate weights from the kernel + let kernel_size = 5; + let kernel_offset = 4; + let weights = &ups_factors.weights4; + let row_size = output[0].as_rect().size().0; + let column_size = row_size; + for di in 0..4 { + for dj in 0..4 { + for ki in 0..kernel_size { + for kj in 0..kernel_size { + let i = kernel_size * di + ki; + let j = kernel_size * dj + kj; + let offset_i = kernel_offset + i; + let offset_j = kernel_offset + j; + // Testing symmetry + let output_value = output[0].as_rect().row(offset_i)[offset_j]; + let output_value_mirrored_right = + output[0].as_rect().row(row_size - offset_i - 1)[offset_j]; + let output_value_mirrored_down = output[0] + .as_rect() + .row(row_size - offset_i - 1)[column_size - offset_j - 1]; + let output_value_mirrored_down_right = output[0] + .as_rect() + .row(row_size - offset_i - 1)[column_size - offset_j - 1]; + + assert_almost_abs_eq(output_value, output_value_mirrored_right, eps); + assert_almost_abs_eq(output_value, output_value_mirrored_down, eps); + assert_almost_abs_eq(output_value, output_value_mirrored_down_right, eps); + + // Testing if we get the expected weights, appropriately mapped. + let mapped_i = if (i % 4) < 2 { + 4 - (i / 4) + (i % 2) * 5 + } else { + i / 4 + (1 - (i % 2)) * 5 + }; + let mapped_j = if (j % 4) < 2 { + 4 - (j / 4) + (j % 2) * 5 + } else { + j / 4 + (1 - (j % 2)) * 5 + }; + let weight_index = index_map[mapped_i][mapped_j]; + assert_almost_abs_eq( + output_value, + weights[weight_index].clamp(0.0, 1.0), + eps, + ); + } + } + } + } + + Ok(()) + } + + #[test] + fn test_upsample8() -> Result<()> { + let eps = 0.0000001; + let mut input = Image::new((7, 7))?; + // Put a single "1.0" in the middle of the image. + input.as_rect_mut().row(3)[3] = 1.0f32; + let ups_factors = ups_factors(); + let stage = Upsample8x::new(&ups_factors, 0); + let output: Vec<Image<f32>> = + make_and_run_simple_pipeline(stage, &[input], (56, 56), 0, 1024)?.1; + + assert_eq!(output[0].as_rect().size(), (56, 56)); + + // Check we have a border with zeros + for i in 0..56 { + for j in 0..8 { + assert_almost_abs_eq(output[0].as_rect().row(j)[i], 0.0, eps); + assert_almost_abs_eq(output[0].as_rect().row(i)[j], 0.0, eps); + assert_almost_abs_eq(output[0].as_rect().row(55 - j)[i], 0.0, eps); + assert_almost_abs_eq(output[0].as_rect().row(i)[55 - j], 0.0, eps); + } + } + + // Define the mapping for the symmetric top-left kernel + let index_map = [ + [ + 0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d, + 0x0e, 0x0f, 0x10, 0x11, 0x12, 0x13, + ], + [ + 0x01, 0x14, 0x15, 0x16, 0x17, 0x18, 0x19, 0x1a, 0x1b, 0x1c, 0x1d, 0x1e, 0x1f, 0x20, + 0x21, 0x22, 0x23, 0x24, 0x25, 0x26, + ], + [ + 0x02, 0x15, 0x27, 0x28, 0x29, 0x2a, 0x2b, 0x2c, 0x2d, 0x2e, 0x2f, 0x30, 0x31, 0x32, + 0x33, 0x34, 0x35, 0x36, 0x37, 0x38, + ], + [ + 0x03, 0x16, 0x28, 0x39, 0x3a, 0x3b, 0x3c, 0x3d, 0x3e, 0x3f, 0x40, 0x41, 0x42, 0x43, + 0x44, 0x45, 0x46, 0x47, 0x48, 0x49, + ], + [ + 0x04, 0x17, 0x29, 0x3a, 0x4a, 0x4b, 0x4c, 0x4d, 0x4e, 0x4f, 0x50, 0x51, 0x52, 0x53, + 0x54, 0x55, 0x56, 0x57, 0x58, 0x59, + ], + [ + 0x05, 0x18, 0x2a, 0x3b, 0x4b, 0x5a, 0x5b, 0x5c, 0x5d, 0x5e, 0x5f, 0x60, 0x61, 0x62, + 0x63, 0x64, 0x65, 0x66, 0x67, 0x68, + ], + [ + 0x06, 0x19, 0x2b, 0x3c, 0x4c, 0x5b, 0x69, 0x6a, 0x6b, 0x6c, 0x6d, 0x6e, 0x6f, 0x70, + 0x71, 0x72, 0x73, 0x74, 0x75, 0x76, + ], + [ + 0x07, 0x1a, 0x2c, 0x3d, 0x4d, 0x5c, 0x6a, 0x77, 0x78, 0x79, 0x7a, 0x7b, 0x7c, 0x7d, + 0x7e, 0x7f, 0x80, 0x81, 0x82, 0x83, + ], + [ + 0x08, 0x1b, 0x2d, 0x3e, 0x4e, 0x5d, 0x6b, 0x78, 0x84, 0x85, 0x86, 0x87, 0x88, 0x89, + 0x8a, 0x8b, 0x8c, 0x8d, 0x8e, 0x8f, + ], + [ + 0x09, 0x1c, 0x2e, 0x3f, 0x4f, 0x5e, 0x6c, 0x79, 0x85, 0x90, 0x91, 0x92, 0x93, 0x94, + 0x95, 0x96, 0x97, 0x98, 0x99, 0x9a, + ], + [ + 0x0a, 0x1d, 0x2f, 0x40, 0x50, 0x5f, 0x6d, 0x7a, 0x86, 0x91, 0x9b, 0x9c, 0x9d, 0x9e, + 0x9f, 0xa0, 0xa1, 0xa2, 0xa3, 0xa4, + ], + [ + 0x0b, 0x1e, 0x30, 0x41, 0x51, 0x60, 0x6e, 0x7b, 0x87, 0x92, 0x9c, 0xa5, 0xa6, 0xa7, + 0xa8, 0xa9, 0xaa, 0xab, 0xac, 0xad, + ], + [ + 0x0c, 0x1f, 0x31, 0x42, 0x52, 0x61, 0x6f, 0x7c, 0x88, 0x93, 0x9d, 0xa6, 0xae, 0xaf, + 0xb0, 0xb1, 0xb2, 0xb3, 0xb4, 0xb5, + ], + [ + 0x0d, 0x20, 0x32, 0x43, 0x53, 0x62, 0x70, 0x7d, 0x89, 0x94, 0x9e, 0xa7, 0xaf, 0xb6, + 0xb7, 0xb8, 0xb9, 0xba, 0xbb, 0xbc, + ], + [ + 0x0e, 0x21, 0x33, 0x44, 0x54, 0x63, 0x71, 0x7e, 0x8a, 0x95, 0x9f, 0xa8, 0xb0, 0xb7, + 0xbd, 0xbe, 0xbf, 0xc0, 0xc1, 0xc2, + ], + [ + 0x0f, 0x22, 0x34, 0x45, 0x55, 0x64, 0x72, 0x7f, 0x8b, 0x96, 0xa0, 0xa9, 0xb1, 0xb8, + 0xbe, 0xc3, 0xc4, 0xc5, 0xc6, 0xc7, + ], + [ + 0x10, 0x23, 0x35, 0x46, 0x56, 0x65, 0x73, 0x80, 0x8c, 0x97, 0xa1, 0xaa, 0xb2, 0xb9, + 0xbf, 0xc4, 0xc8, 0xc9, 0xca, 0xcb, + ], + [ + 0x11, 0x24, 0x36, 0x47, 0x57, 0x66, 0x74, 0x81, 0x8d, 0x98, 0xa2, 0xab, 0xb3, 0xba, + 0xc0, 0xc5, 0xc9, 0xcc, 0xcd, 0xce, + ], + [ + 0x12, 0x25, 0x37, 0x48, 0x58, 0x67, 0x75, 0x82, 0x8e, 0x99, 0xa3, 0xac, 0xb4, 0xbb, + 0xc1, 0xc6, 0xca, 0xcd, 0xcf, 0xd0, + ], + [ + 0x13, 0x26, 0x38, 0x49, 0x59, 0x68, 0x76, 0x83, 0x8f, 0x9a, 0xa4, 0xad, 0xb5, 0xbc, + 0xc2, 0xc7, 0xcb, 0xce, 0xd0, 0xd1, + ], + ]; + + // Validate weights from the kernel + let kernel_size = 5; + let kernel_offset = 8; + let weights = &ups_factors.weights8; + let row_size = output[0].as_rect().size().0; + let column_size = row_size; + for di in 0..8 { + for dj in 0..8 { + for ki in 0..kernel_size { + for kj in 0..kernel_size { + let i = kernel_size * di + ki; + let j = kernel_size * dj + kj; + let offset_i = kernel_offset + i; + let offset_j = kernel_offset + j; + // Testing symmetry + let output_value = output[0].as_rect().row(offset_i)[offset_j]; + let output_value_mirrored_right = + output[0].as_rect().row(row_size - offset_i - 1)[offset_j]; + let output_value_mirrored_down = output[0] + .as_rect() + .row(row_size - offset_i - 1)[column_size - offset_j - 1]; + let output_value_mirrored_down_right = output[0] + .as_rect() + .row(row_size - offset_i - 1)[column_size - offset_j - 1]; + + assert_almost_abs_eq(output_value, output_value_mirrored_right, eps); + assert_almost_abs_eq(output_value, output_value_mirrored_down, eps); + assert_almost_abs_eq(output_value, output_value_mirrored_down_right, eps); + + // Testing if we get the expected weights, appropriately mapped. + let mapped_i = if (i % 8) < 4 { + 4 - (i / 8) + (i % 4) * 5 + } else { + i / 8 + (3 - (i % 4)) * 5 + }; + let mapped_j = if (j % 8) < 4 { + 4 - (j / 8) + (j % 4) * 5 + } else { + j / 8 + (3 - (j % 4)) * 5 + }; + let weight_index = index_map[mapped_i][mapped_j]; + assert_almost_abs_eq( + output_value, + weights[weight_index].clamp(0.0, 1.0), + eps, + ); + } + } + } + } + + Ok(()) + } +} diff --git a/third_party/rust/jxl/src/render/stages/xyb.rs b/third_party/rust/jxl/src/render/stages/xyb.rs @@ -0,0 +1,369 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +use crate::api::{ + JxlColorEncoding, JxlPrimaries, JxlTransferFunction, JxlWhitePoint, adapt_to_xyz_d50, + primaries_to_xyz, primaries_to_xyz_d50, +}; +use crate::error::Result; +use crate::headers::{FileHeader, OpsinInverseMatrix}; +use crate::render::stages::from_linear; +use crate::render::{RenderPipelineInPlaceStage, RenderPipelineStage}; +use crate::simd::{F32SimdVec, simd_function}; +use crate::util::{Matrix3x3, inv_3x3_matrix, mul_3x3_matrix}; + +const SRGB_LUMINANCES: [f32; 3] = [0.2126, 0.7152, 0.0722]; + +#[derive(Clone)] +pub struct OutputColorInfo { + // Luminance of each primary. + pub luminances: [f32; 3], + pub intensity_target: f32, + pub opsin: OpsinInverseMatrix, + pub tf: from_linear::TransferFunction, +} + +#[cfg(test)] +impl Default for OutputColorInfo { + fn default() -> Self { + use crate::headers::encodings::Empty; + Self { + luminances: SRGB_LUMINANCES, + intensity_target: 255.0, + opsin: OpsinInverseMatrix::default(&Empty {}), + tf: from_linear::TransferFunction::Srgb, + } + } +} + +impl OutputColorInfo { + fn opsin_matrix_to_matrix3x3(matrix: [f32; 9]) -> Matrix3x3<f64> { + [ + [matrix[0] as f64, matrix[1] as f64, matrix[2] as f64], + [matrix[3] as f64, matrix[4] as f64, matrix[5] as f64], + [matrix[6] as f64, matrix[7] as f64, matrix[8] as f64], + ] + } + + fn matrix3x3_to_opsin_matrix(matrix: Matrix3x3<f64>) -> [f32; 9] { + [ + matrix[0][0] as f32, + matrix[0][1] as f32, + matrix[0][2] as f32, + matrix[1][0] as f32, + matrix[1][1] as f32, + matrix[1][2] as f32, + matrix[2][0] as f32, + matrix[2][1] as f32, + matrix[2][2] as f32, + ] + } + + pub fn from_header(header: &FileHeader) -> Result<Self> { + let srgb_output = OutputColorInfo { + luminances: SRGB_LUMINANCES, + intensity_target: header.image_metadata.tone_mapping.intensity_target, + opsin: header.transform_data.opsin_inverse_matrix.clone(), + tf: from_linear::TransferFunction::Srgb, + }; + if header.image_metadata.color_encoding.want_icc { + return Ok(srgb_output); + } + + let tf; + let mut inverse_matrix = Self::opsin_matrix_to_matrix3x3( + header.transform_data.opsin_inverse_matrix.inverse_matrix, + ); + let mut luminances = SRGB_LUMINANCES; + let desired_colorspace = + JxlColorEncoding::from_internal(&header.image_metadata.color_encoding)?; + match &desired_colorspace { + JxlColorEncoding::XYB { .. } => { + return Ok(srgb_output); + } + JxlColorEncoding::RgbColorSpace { + white_point, + primaries, + transfer_function, + .. + } => { + tf = transfer_function; + if *primaries != JxlPrimaries::SRGB || *white_point != JxlWhitePoint::D65 { + let [r, g, b] = JxlPrimaries::SRGB.to_xy_coords(); + let w = JxlWhitePoint::D65.to_xy_coords(); + let srgb_to_xyzd50 = + primaries_to_xyz_d50(r.0, r.1, g.0, g.1, b.0, b.1, w.0, w.1)?; + let [r, g, b] = primaries.to_xy_coords(); + let w = white_point.to_xy_coords(); + let original_to_xyz = primaries_to_xyz(r.0, r.1, g.0, g.1, b.0, b.1, w.0, w.1)?; + luminances = original_to_xyz[1].map(|lum| lum as f32); + let adapt_to_d50 = adapt_to_xyz_d50(w.0, w.1)?; + let original_to_xyzd50 = mul_3x3_matrix(&adapt_to_d50, &original_to_xyz); + let xyzd50_to_original = inv_3x3_matrix(&original_to_xyzd50)?; + let srgb_to_original = mul_3x3_matrix(&xyzd50_to_original, &srgb_to_xyzd50); + inverse_matrix = mul_3x3_matrix(&srgb_to_original, &inverse_matrix); + } + } + + JxlColorEncoding::GrayscaleColorSpace { + transfer_function, .. + } => { + tf = transfer_function; + let f64_luminances = luminances.map(|lum| lum as f64); + let srgb_to_luminance: Matrix3x3<f64> = + [f64_luminances, f64_luminances, f64_luminances]; + inverse_matrix = mul_3x3_matrix(&srgb_to_luminance, &inverse_matrix); + } + } + + let mut opsin = header.transform_data.opsin_inverse_matrix.clone(); + opsin.inverse_matrix = Self::matrix3x3_to_opsin_matrix(inverse_matrix); + let intensity_target = header.image_metadata.tone_mapping.intensity_target; + let from_linear_tf = match tf { + JxlTransferFunction::PQ => from_linear::TransferFunction::Pq { intensity_target }, + JxlTransferFunction::HLG => from_linear::TransferFunction::Hlg { + intensity_target, + luminance_rgb: luminances, + }, + JxlTransferFunction::BT709 => from_linear::TransferFunction::Bt709, + JxlTransferFunction::Linear => from_linear::TransferFunction::Gamma(1.0), + JxlTransferFunction::SRGB => from_linear::TransferFunction::Srgb, + JxlTransferFunction::DCI => from_linear::TransferFunction::Gamma(2.6_f32.recip()), + JxlTransferFunction::Gamma(g) => from_linear::TransferFunction::Gamma(*g), + }; + Ok(OutputColorInfo { + luminances, + intensity_target, + opsin, + tf: from_linear_tf, + }) + } +} + +/// Convert XYB to linear RGB with appropriate primaries, where 1.0 corresponds to `intensity_target` nits. +pub struct XybStage { + first_channel: usize, + output_color_info: OutputColorInfo, +} + +impl XybStage { + pub fn new(first_channel: usize, output_color_info: OutputColorInfo) -> Self { + Self { + first_channel, + output_color_info, + } + } +} + +impl std::fmt::Display for XybStage { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + let channel = self.first_channel; + write!( + f, + "XYB to linear for channel [{},{},{}]", + channel, + channel + 1, + channel + 2 + ) + } +} + +simd_function!( + xyb_process_dispatch, + d: D, + fn xyb_process( + opsin: &OpsinInverseMatrix, + intensity_target: f32, + xsize: usize, + row_x: &mut [f32], + row_y: &mut [f32], + row_b: &mut [f32], + ) { + let OpsinInverseMatrix { + inverse_matrix: mat, + opsin_biases: bias, + .. + } = opsin; + // TODO(veluca): consider computing the cbrt in advance. + let bias_cbrt = bias.map(|x| D::F32Vec::splat(d, x.cbrt())); + let intensity_scale = 255.0 / intensity_target; + let scaled_bias = bias.map(|x| D::F32Vec::splat(d, x * intensity_scale)); + let mat = mat.map(|x| D::F32Vec::splat(d, x)); + let intensity_scale = D::F32Vec::splat(d, intensity_scale); + + for idx in (0..xsize).step_by(D::F32Vec::LEN) { + let x = D::F32Vec::load(d, &row_x[idx..]); + let y = D::F32Vec::load(d, &row_y[idx..]); + let b = D::F32Vec::load(d, &row_b[idx..]); + + // Mix and apply bias + let l = y + x - bias_cbrt[0]; + let m = y - x - bias_cbrt[1]; + let s = b - bias_cbrt[2]; + + // Apply biased inverse gamma and scale (1.0 corresponds to `intensity_target` nits) + let l2 = l * l; + let m2 = m * m; + let s2 = s * s; + let scaled_l = l * intensity_scale; + let scaled_m = m * intensity_scale; + let scaled_s = s * intensity_scale; + let l = l2.mul_add(scaled_l, scaled_bias[0]); + let m = m2.mul_add(scaled_m, scaled_bias[1]); + let s = s2.mul_add(scaled_s, scaled_bias[2]); + + // Apply opsin inverse matrix (linear LMS to linear sRGB) + let r = mat[0].mul_add(l, mat[1].mul_add(m, mat[2] * s)); + let g = mat[3].mul_add(l, mat[4].mul_add(m, mat[5] * s)); + let b = mat[6].mul_add(l, mat[7].mul_add(m, mat[8] * s)); + r.store(&mut row_x[idx..]); + g.store(&mut row_y[idx..]); + b.store(&mut row_b[idx..]); + } + } +); + +impl RenderPipelineStage for XybStage { + type Type = RenderPipelineInPlaceStage<f32>; + + fn uses_channel(&self, c: usize) -> bool { + (self.first_channel..self.first_channel + 3).contains(&c) + } + + fn process_row_chunk( + &self, + _position: (usize, usize), + xsize: usize, + row: &mut [&mut [f32]], + _state: Option<&mut dyn std::any::Any>, + ) { + let [row_x, row_y, row_b] = row else { + panic!( + "incorrect number of channels; expected 3, found {}", + row.len() + ); + }; + + xyb_process_dispatch( + &self.output_color_info.opsin, + self.output_color_info.intensity_target, + xsize, + row_x, + row_y, + row_b, + ); + } +} + +#[cfg(test)] +mod test { + use test_log::test; + + use super::*; + use crate::error::Result; + use crate::headers::encodings::Empty; + use crate::image::Image; + use crate::render::test::make_and_run_simple_pipeline; + use crate::simd::{ + ScalarDescriptor, SimdDescriptor, round_up_size_to_two_cache_lines, + test_all_instruction_sets, + }; + use crate::util::test::assert_all_almost_abs_eq; + + #[test] + fn consistency() -> Result<()> { + crate::render::test::test_stage_consistency::<_, f32, f32>( + XybStage::new(0, OutputColorInfo::default()), + (500, 500), + 3, + ) + } + + #[test] + fn srgb_primaries() -> Result<()> { + let mut input_x = Image::new((3, 1))?; + let mut input_y = Image::new((3, 1))?; + let mut input_b = Image::new((3, 1))?; + input_x + .as_rect_mut() + .row(0) + .copy_from_slice(&[0.028100073, -0.015386105, 0.0]); + input_y + .as_rect_mut() + .row(0) + .copy_from_slice(&[0.4881882, 0.71478134, 0.2781282]); + input_b + .as_rect_mut() + .row(0) + .copy_from_slice(&[0.471659, 0.43707693, 0.66613984]); + + let stage = XybStage::new(0, OutputColorInfo::default()); + let output = make_and_run_simple_pipeline::<_, f32, f32>( + stage, + &[input_x, input_y, input_b], + (3, 1), + 0, + 256, + )? + .1; + + assert_all_almost_abs_eq(output[0].as_rect().row(0), &[1.0, 0.0, 0.0], 1e-6); + assert_all_almost_abs_eq(output[1].as_rect().row(0), &[0.0, 1.0, 0.0], 1e-6); + assert_all_almost_abs_eq(output[2].as_rect().row(0), &[0.0, 0.0, 1.0], 1e-6); + + Ok(()) + } + + fn xyb_process_scalar_equivalent<D: SimdDescriptor>(d: D) { + let opsin = OpsinInverseMatrix::default(&Empty {}); + arbtest::arbtest(|u| { + let xsize = u.arbitrary_len::<usize>()?; + let intensity_target = u.arbitrary::<u8>()? as f32 * 2.0 + 1.0; + let mut row_x = vec![0.0; round_up_size_to_two_cache_lines::<f32>(xsize)]; + let mut row_y = vec![0.0; round_up_size_to_two_cache_lines::<f32>(xsize)]; + let mut row_b = vec![0.0; round_up_size_to_two_cache_lines::<f32>(xsize)]; + + for i in 0..xsize { + row_x[i] = u.arbitrary::<i16>()? as f32 * (1.0 / i16::MAX as f32); + row_y[i] = u.arbitrary::<i16>()? as f32 * (1.0 / i16::MAX as f32); + row_b[i] = u.arbitrary::<i16>()? as f32 * (1.0 / i16::MAX as f32); + } + + let mut scalar_x = row_x.clone(); + let mut scalar_y = row_y.clone(); + let mut scalar_b = row_b.clone(); + + xyb_process( + d, + &opsin, + intensity_target, + xsize, + &mut row_x, + &mut row_y, + &mut row_b, + ); + + xyb_process( + ScalarDescriptor::new().unwrap(), + &opsin, + intensity_target, + xsize, + &mut scalar_x, + &mut scalar_y, + &mut scalar_b, + ); + + for i in 0..xsize { + assert!((row_x[i] - scalar_x[i]).abs() < 1e-8); + assert!((row_y[i] - scalar_y[i]).abs() < 1e-8); + assert!((row_b[i] - scalar_b[i]).abs() < 1e-8); + } + + Ok(()) + }); + } + + test_all_instruction_sets!(xyb_process_scalar_equivalent); +} diff --git a/third_party/rust/jxl/src/render/stages/ycbcr.rs b/third_party/rust/jxl/src/render/stages/ycbcr.rs @@ -0,0 +1,125 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +use crate::render::{RenderPipelineInPlaceStage, RenderPipelineStage}; + +/// Convert YCbCr to RGB +pub struct YcbcrToRgbStage { + first_channel: usize, +} + +impl YcbcrToRgbStage { + pub fn new(first_channel: usize) -> Self { + Self { first_channel } + } +} + +impl std::fmt::Display for YcbcrToRgbStage { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + let channel = self.first_channel; + write!( + f, + "YCbCr to RGB for channel [{},{},{}]", + channel, + channel + 1, + channel + 2 + ) + } +} + +impl RenderPipelineStage for YcbcrToRgbStage { + type Type = RenderPipelineInPlaceStage<f32>; + + fn uses_channel(&self, c: usize) -> bool { + (self.first_channel..self.first_channel + 3).contains(&c) + } + + fn process_row_chunk( + &self, + _position: (usize, usize), + xsize: usize, + row: &mut [&mut [f32]], + _state: Option<&mut dyn std::any::Any>, + ) { + // pixels are stored in `Cb Y Cr` order to mimic XYB colorspace + let [row_cb, row_y, row_cr] = row else { + panic!( + "incorrect number of channels; expected 3, found {}", + row.len() + ); + }; + + assert!(xsize <= row_cb.len() && xsize <= row_y.len() && xsize <= row_cr.len()); + for idx in 0..xsize { + let y = row_y[idx] + 128.0 / 255.0; // shift Y from [-0.5, 0.5] to [0, 1], matching JPEG spec + let cb = row_cb[idx]; + let cr = row_cr[idx]; + + // Full-range BT.601 as defined by JFIF Clause 7: + // https://www.itu.int/rec/T-REC-T.871-201105-I/en + row_cb[idx] = cr.mul_add(1.402, y); + row_y[idx] = cr.mul_add( + -0.299 * 1.402 / 0.587, + cb.mul_add(-0.114 * 1.772 / 0.587, y), + ); + row_cr[idx] = cb.mul_add(1.772, y); + } + } +} + +#[cfg(test)] +mod test { + use test_log::test; + + use super::*; + use crate::error::Result; + use crate::image::Image; + use crate::render::test::make_and_run_simple_pipeline; + use crate::util::test::assert_all_almost_abs_eq; + + #[test] + fn consistency() -> Result<()> { + crate::render::test::test_stage_consistency::<_, f32, f32>( + YcbcrToRgbStage::new(0), + (500, 500), + 3, + ) + } + + #[test] + fn srgb_primaries() -> Result<()> { + let mut input_y = Image::new((3, 1))?; + let mut input_cb = Image::new((3, 1))?; + let mut input_cr = Image::new((3, 1))?; + input_y + .as_rect_mut() + .row(0) + .copy_from_slice(&[-0.20296079, 0.08503921, -0.3879608]); + input_cb + .as_rect_mut() + .row(0) + .copy_from_slice(&[-0.16873589, -0.3312641, 0.5]); + input_cr + .as_rect_mut() + .row(0) + .copy_from_slice(&[0.5, -0.41868758, -0.08131241]); + + let stage = YcbcrToRgbStage::new(0); + let output = make_and_run_simple_pipeline::<_, f32, f32>( + stage, + &[input_cb, input_y, input_cr], + (3, 1), + 0, + 256, + )? + .1; + + assert_all_almost_abs_eq(output[0].as_rect().row(0), &[1.0, 0.0, 0.0], 1e-6); + assert_all_almost_abs_eq(output[1].as_rect().row(0), &[0.0, 1.0, 0.0], 1e-6); + assert_all_almost_abs_eq(output[2].as_rect().row(0), &[0.0, 0.0, 1.0], 1e-6); + + Ok(()) + } +} diff --git a/third_party/rust/jxl/src/render/test.rs b/third_party/rust/jxl/src/render/test.rs @@ -0,0 +1,183 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +use crate::{ + error::Result, + headers::Orientation, + image::{Image, ImageDataType}, + util::{ShiftRightCeil, tracing_wrappers::instrument}, +}; +use rand::SeedableRng; + +use super::{ + RenderPipeline, RenderPipelineBuilder, RenderPipelineStage, SaveStage, SaveStageType, + internal::RenderPipelineStageInfo, simple_pipeline::SimpleRenderPipelineBuilder, +}; + +pub(super) fn make_and_run_simple_pipeline< + S: RenderPipelineStage, + InputT: ImageDataType, + OutputT: ImageDataType + std::ops::Mul<Output = OutputT>, +>( + stage: S, + input_images: &[Image<InputT>], + image_size: (usize, usize), + downsampling_shift: usize, + chunk_size: usize, +) -> Result<(S, Vec<Image<OutputT>>)> { + let final_size = stage.new_size(image_size); + const LOG_GROUP_SIZE: usize = 8; + let all_channels = (0..input_images.len()).collect::<Vec<_>>(); + let uses_channel: Vec<_> = all_channels + .iter() + .map(|x| stage.uses_channel(*x)) + .collect(); + let mut pipeline = SimpleRenderPipelineBuilder::new_with_chunk_size( + input_images.len(), + image_size, + downsampling_shift, + LOG_GROUP_SIZE, + chunk_size, + ) + .add_stage(stage)?; + for i in 0..input_images.len() { + pipeline = pipeline.add_save_stage(SaveStage::<OutputT>::new( + SaveStageType::Output, + i, + final_size, + OutputT::from_f64(1.0), + Orientation::Identity, + )?)?; + } + let mut pipeline = pipeline.build()?; + + for g in 0..pipeline.num_groups() { + for &c in all_channels.iter() { + let log_group_size = if uses_channel[c] { + ( + LOG_GROUP_SIZE - S::Type::SHIFT.0 as usize, + LOG_GROUP_SIZE - S::Type::SHIFT.1 as usize, + ) + } else { + (LOG_GROUP_SIZE, LOG_GROUP_SIZE) + }; + pipeline.set_buffer_for_group( + c, + g, + 1, + input_images[c].group_rect(g, log_group_size).to_image()?, + ); + } + } + + // TODO(veluca): pass actual output buffers. + pipeline.do_render(&mut [])?; + + let mut stages = pipeline.into_stages().into_iter(); + let stage = stages + .next() + .unwrap() + .downcast::<S>() + .expect("first stage is always the tested stage"); + + let outputs = stages + .map(|s| { + s.downcast::<SaveStage<OutputT>>() + .expect("all later stages are always SaveStage") + .into_buffer() + }) + .collect(); + + Ok((*stage, outputs)) +} + +#[instrument(skip(stage), err)] +pub(super) fn test_stage_consistency< + S: RenderPipelineStage, + InputT: ImageDataType, + OutputT: ImageDataType + std::ops::Mul<Output = OutputT>, +>( + stage: S, + image_size: (usize, usize), + num_image_channels: usize, +) -> Result<()> { + let mut rng = rand_xorshift::XorShiftRng::seed_from_u64(0); + let images: Result<Vec<_>> = (0..num_image_channels) + .map(|c| { + let size = if stage.uses_channel(c) { + ( + image_size.0.shrc(S::Type::SHIFT.0), + image_size.1.shrc(S::Type::SHIFT.1), + ) + } else { + image_size + }; + Image::new_random(size, &mut rng) + }) + .collect(); + let images = images?; + + let (stage, base_output) = + make_and_run_simple_pipeline::<_, InputT, OutputT>(stage, &images, image_size, 0, 256)?; + + let mut stage = Some(stage); + + arbtest::arbtest(move |p| { + let chunk_size = p.arbitrary::<u16>()?.saturating_add(1) as usize; + let (s, output) = make_and_run_simple_pipeline::<_, InputT, OutputT>( + stage.take().unwrap(), + &images, + image_size, + 0, + chunk_size, + ) + .unwrap_or_else(|_| panic!("error running pipeline with chunk size {chunk_size}")); + stage = Some(s); + + for (o, bo) in output.iter().zip(base_output.iter()) { + bo.as_rect().check_equal(o.as_rect()); + } + + Ok(()) + }); + Ok(()) +} + +macro_rules! create_in_out_rows { + ($u:expr, $border_x:expr, $border_y:expr, $rows:ident, $xsize:ident) => { + use crate::simd::round_up_size_to_two_cache_lines; + let $xsize: usize = 1 + $u.arbitrary::<usize>()? % 4095; + let mut row_vecs = vec![( + vec![ + vec![ + 0f32; + round_up_size_to_two_cache_lines::<f32>( + round_up_size_to_two_cache_lines::<f32>($xsize) + $border_x * 2 + ) + ]; + 1 + $border_y * 2 + ], + vec![vec![0f32; round_up_size_to_two_cache_lines::<f32>($xsize)]], + )]; + + let mut row_vecs_refs: Vec<(Vec<&[f32]>, Vec<&mut [f32]>)> = row_vecs + .iter_mut() + .map(|(left, right)| { + ( + left.iter().map(|v| v.as_slice()).collect(), + right.iter_mut().map(|v| v.as_mut_slice()).collect(), + ) + }) + .collect(); + + let mut outer: Vec<(&[&[f32]], &mut [&mut [f32]])> = row_vecs_refs + .iter_mut() + .map(|(left, right)| (left.as_slice(), right.as_mut_slice())) + .collect(); + + let $rows = &mut outer[..]; + }; +} +pub(crate) use create_in_out_rows; diff --git a/third_party/rust/jxl/src/simd/mod.rs b/third_party/rust/jxl/src/simd/mod.rs @@ -0,0 +1,221 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +use std::{ + fmt::Debug, + ops::{Add, AddAssign, Div, DivAssign, Mul, MulAssign, Sub, SubAssign}, +}; + +#[cfg(target_arch = "x86_64")] +mod x86_64; + +mod scalar; + +#[cfg(target_arch = "x86_64")] +pub(crate) use x86_64::{avx::AvxDescriptor, avx512::Avx512Descriptor, simd_function}; + +#[cfg(not(target_arch = "x86_64"))] +pub(crate) use scalar::simd_function; + +#[cfg(all(test, target_arch = "x86_64"))] +pub(crate) use x86_64::test_all_instruction_sets; + +#[cfg(all(test, not(target_arch = "x86_64")))] +pub(crate) use scalar::test_all_instruction_sets; + +pub(crate) use scalar::ScalarDescriptor; + +const CACHE_LINE_BYTE_SIZE: usize = 64; + +pub fn round_up_size_to_two_cache_lines<T>(size: usize) -> usize { + let elements_per_cache_line = CACHE_LINE_BYTE_SIZE / std::mem::size_of::<T>() * 2; + size.div_ceil(elements_per_cache_line) * elements_per_cache_line +} + +pub trait SimdDescriptor: Sized + Copy + Debug + Send + Sync { + type F32Vec: F32SimdVec<Descriptor = Self>; + + fn new() -> Option<Self>; +} + +pub trait F32SimdVec: + Sized + + Copy + + Debug + + Send + + Sync + + Add<Self, Output = Self> + + Mul<Self, Output = Self> + + Sub<Self, Output = Self> + + Div<Self, Output = Self> + + AddAssign<Self> + + MulAssign<Self> + + SubAssign<Self> + + DivAssign<Self> +{ + type Descriptor: SimdDescriptor; + + const LEN: usize; + + /// Converts v to an array of v. + fn splat(d: Self::Descriptor, v: f32) -> Self; + + fn mul_add(self, mul: Self, add: Self) -> Self; + + // Requires `mem.len() >= Self::LEN` or it will panic. + fn load(d: Self::Descriptor, mem: &[f32]) -> Self; + + // Requires `mem.len() >= Self::LEN` or it will panic. + fn store(&self, mem: &mut [f32]); + + fn abs(self) -> Self; + + fn max(self, other: Self) -> Self; +} + +#[cfg(test)] +mod test { + use arbtest::arbitrary::Unstructured; + + use crate::{ + simd::{F32SimdVec, ScalarDescriptor, SimdDescriptor, test_all_instruction_sets}, + util::test::assert_all_almost_rel_eq, + }; + + enum Distribution { + Floats, + NonZeroFloats, + } + + fn arb_vec<D: SimdDescriptor>(_: D, u: &mut Unstructured, dist: Distribution) -> Vec<f32> { + let mut res = vec![0.0; D::F32Vec::LEN]; + for v in res.iter_mut() { + match dist { + Distribution::Floats => { + *v = u.arbitrary::<i32>().unwrap() as f32 + / (1.0 + u.arbitrary::<u32>().unwrap() as f32) + } + Distribution::NonZeroFloats => { + let sign = if u.arbitrary::<bool>().unwrap() { + 1.0 + } else { + -1.0 + }; + *v = sign * (1.0 + u.arbitrary::<u32>().unwrap() as f32) + / (1.0 + u.arbitrary::<u32>().unwrap() as f32); + } + } + } + res + } + + macro_rules! test_instruction { + ($name:ident, |$a:ident: $a_dist:ident| $block:expr) => { + fn $name<D: SimdDescriptor>(d: D) { + fn compute<D: SimdDescriptor>(d: D, a: &[f32]) -> Vec<f32> { + let closure = |$a: D::F32Vec| $block; + let mut res = vec![0f32; a.len()]; + for idx in (0..a.len()).step_by(D::F32Vec::LEN) { + closure(D::F32Vec::load(d, &a[idx..])).store(&mut res[idx..]); + } + res + } + arbtest::arbtest(|u| { + let a = arb_vec(d, u, Distribution::$a_dist); + let scalar_res = compute(ScalarDescriptor::new().unwrap(), &a); + let simd_res = compute(d, &a); + assert_all_almost_rel_eq(&scalar_res, &simd_res, 1e-8); + Ok(()) + }) + .size_min(64); + } + test_all_instruction_sets!($name); + }; + ($name:ident, |$a:ident: $a_dist:ident, $b:ident: $b_dist:ident| $block:expr) => { + fn $name<D: SimdDescriptor>(d: D) { + fn compute<D: SimdDescriptor>(d: D, a: &[f32], b: &[f32]) -> Vec<f32> { + let closure = |$a: D::F32Vec, $b: D::F32Vec| $block; + let mut res = vec![0f32; a.len()]; + for idx in (0..a.len()).step_by(D::F32Vec::LEN) { + closure(D::F32Vec::load(d, &a[idx..]), D::F32Vec::load(d, &b[idx..])) + .store(&mut res[idx..]); + } + res + } + arbtest::arbtest(|u| { + let a = arb_vec(d, u, Distribution::$a_dist); + let b = arb_vec(d, u, Distribution::$b_dist); + let scalar_res = compute(ScalarDescriptor::new().unwrap(), &a, &b); + let simd_res = compute(d, &a, &b); + assert_all_almost_rel_eq(&scalar_res, &simd_res, 1e-8); + Ok(()) + }) + .size_min(128); + } + test_all_instruction_sets!($name); + }; + ($name:ident, |$a:ident: $a_dist:ident, $b:ident: $b_dist:ident, $c:ident: $c_dist:ident| $block:expr) => { + fn $name<D: SimdDescriptor>(d: D) { + fn compute<D: SimdDescriptor>(d: D, a: &[f32], b: &[f32], c: &[f32]) -> Vec<f32> { + let closure = |$a: D::F32Vec, $b: D::F32Vec, $c: D::F32Vec| $block; + let mut res = vec![0f32; a.len()]; + for idx in (0..a.len()).step_by(D::F32Vec::LEN) { + closure( + D::F32Vec::load(d, &a[idx..]), + D::F32Vec::load(d, &b[idx..]), + D::F32Vec::load(d, &c[idx..]), + ) + .store(&mut res[idx..]); + } + res + } + arbtest::arbtest(|u| { + let a = arb_vec(d, u, Distribution::$a_dist); + let b = arb_vec(d, u, Distribution::$b_dist); + let c = arb_vec(d, u, Distribution::$c_dist); + let scalar_res = compute(ScalarDescriptor::new().unwrap(), &a, &b, &c); + let simd_res = compute(d, &a, &b, &c); + assert_all_almost_rel_eq(&scalar_res, &simd_res, 1e-8); + Ok(()) + }) + .size_min(172); + } + test_all_instruction_sets!($name); + }; + } + + test_instruction!(add, |a: Floats, b: Floats| { a + b }); + test_instruction!(mul, |a: Floats, b: Floats| { a * b }); + test_instruction!(sub, |a: Floats, b: Floats| { a - b }); + test_instruction!(div, |a: Floats, b: NonZeroFloats| { a / b }); + + test_instruction!(add_assign, |a: Floats, b: Floats| { + let mut res = a; + res += b; + res + }); + test_instruction!(mul_assign, |a: Floats, b: Floats| { + let mut res = a; + res *= b; + res + }); + test_instruction!(sub_assign, |a: Floats, b: Floats| { + let mut res = a; + res -= b; + res + }); + test_instruction!(div_assign, |a: Floats, b: NonZeroFloats| { + let mut res = a; + res /= b; + res + }); + + test_instruction!(mul_add, |a: Floats, b: Floats, c: Floats| { + a.mul_add(b, c) + }); + + test_instruction!(abs, |a: Floats| { a.abs() }); + test_instruction!(max, |a: Floats, b: Floats| { a.max(b) }); +} diff --git a/third_party/rust/jxl/src/simd/scalar.rs b/third_party/rust/jxl/src/simd/scalar.rs @@ -0,0 +1,139 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +use std::ops::{Add, AddAssign, Div, DivAssign, Mul, MulAssign, Sub, SubAssign}; + +use super::{F32SimdVec, SimdDescriptor}; + +#[derive(Clone, Copy, Debug)] +pub struct ScalarDescriptor; + +impl SimdDescriptor for ScalarDescriptor { + type F32Vec = F32VecScalar; + fn new() -> Option<Self> { + Some(Self) + } +} + +#[derive(Clone, Copy, Debug)] +pub struct F32VecScalar(f32); + +impl F32SimdVec for F32VecScalar { + type Descriptor = ScalarDescriptor; + + const LEN: usize = 1; + + fn load(_d: Self::Descriptor, mem: &[f32]) -> Self { + Self(mem[0]) + } + + fn store(&self, mem: &mut [f32]) { + mem[0] = self.0; + } + + fn mul_add(self, mul: Self, add: Self) -> Self { + Self(self.0.mul_add(mul.0, add.0)) + } + + fn splat(_d: Self::Descriptor, v: f32) -> Self { + Self(v) + } + + fn abs(self) -> Self { + Self(self.0.abs()) + } + + fn max(self, other: Self) -> Self { + Self(self.0.max(other.0)) + } +} + +impl Add<F32VecScalar> for F32VecScalar { + type Output = F32VecScalar; + fn add(self, rhs: F32VecScalar) -> Self::Output { + Self(self.0 + rhs.0) + } +} + +impl Sub<F32VecScalar> for F32VecScalar { + type Output = F32VecScalar; + fn sub(self, rhs: F32VecScalar) -> Self::Output { + Self(self.0 - rhs.0) + } +} + +impl Mul<F32VecScalar> for F32VecScalar { + type Output = F32VecScalar; + fn mul(self, rhs: F32VecScalar) -> Self::Output { + Self(self.0 * rhs.0) + } +} + +impl Div<F32VecScalar> for F32VecScalar { + type Output = F32VecScalar; + fn div(self, rhs: F32VecScalar) -> Self::Output { + Self(self.0 / rhs.0) + } +} + +impl AddAssign<F32VecScalar> for F32VecScalar { + fn add_assign(&mut self, rhs: F32VecScalar) { + self.0 += rhs.0; + } +} + +impl SubAssign<F32VecScalar> for F32VecScalar { + fn sub_assign(&mut self, rhs: F32VecScalar) { + self.0 -= rhs.0; + } +} + +impl MulAssign<F32VecScalar> for F32VecScalar { + fn mul_assign(&mut self, rhs: F32VecScalar) { + self.0 *= rhs.0; + } +} + +impl DivAssign<F32VecScalar> for F32VecScalar { + fn div_assign(&mut self, rhs: F32VecScalar) { + self.0 /= rhs.0; + } +} + +#[allow(unused_macros)] +macro_rules! simd_function { + ( + $dname:ident, + $descr:ident: $descr_ty:ident, + $pub:vis fn $name:ident($($arg:ident: $ty:ty),* $(,)?) $(-> $ret:ty )? $body: block + ) => { + $pub fn $name<$descr_ty: crate::simd::SimdDescriptor>($descr: $descr_ty, $($arg: $ty),*) $(-> $ret)? $body + $pub fn $dname($($arg: $ty),*) $(-> $ret)? { + use crate::simd::SimdDescriptor; + $name(crate::simd::ScalarDescriptor::new().unwrap(), $($arg),*) + } + }; +} + +#[allow(unused_imports)] +pub(crate) use simd_function; + +#[allow(unused_macros)] +macro_rules! test_all_instruction_sets { + ( + $name:ident + ) => { + paste::paste! { + #[test] + fn [<$name _scalar>]() { + use crate::simd::SimdDescriptor; + $name(crate::simd::ScalarDescriptor::new().unwrap()) + } + } + }; +} + +#[allow(unused_imports)] +pub(crate) use test_all_instruction_sets; diff --git a/third_party/rust/jxl/src/simd/x86_64/avx.rs b/third_party/rust/jxl/src/simd/x86_64/avx.rs @@ -0,0 +1,149 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +use std::{ + arch::x86_64::{ + __m256, _mm256_add_ps, _mm256_andnot_si256, _mm256_castps_si256, _mm256_castsi256_ps, + _mm256_div_ps, _mm256_fmadd_ps, _mm256_loadu_ps, _mm256_max_ps, _mm256_mul_ps, + _mm256_set1_epi32, _mm256_set1_ps, _mm256_storeu_ps, _mm256_sub_ps, + }, + ops::{Add, AddAssign, Div, DivAssign, Mul, MulAssign, Sub, SubAssign}, +}; + +use super::super::{F32SimdVec, SimdDescriptor}; + +// Safety invariant: this type is only ever constructed if avx2 and fma are available. +#[derive(Clone, Copy, Debug)] +pub struct AvxDescriptor; + +impl SimdDescriptor for AvxDescriptor { + type F32Vec = F32VecAvx; + fn new() -> Option<Self> { + if is_x86_feature_detected!("avx2") && is_x86_feature_detected!("fma") { + // SAFETY: we just checked avx2 and fma. + Some(Self) + } else { + None + } + } +} + +// TODO(veluca): retire this macro once we have #[unsafe(target_feature)]. +macro_rules! fn_avx { + ( + $this:ident: $self_ty:ty, + fn $name:ident($($arg:ident: $ty:ty),* $(,)?) $(-> $ret:ty )? $body: block) => { + #[inline(always)] + fn $name(self: $self_ty, $($arg: $ty),*) $(-> $ret)? { + #[target_feature(enable = "fma,avx2")] + #[inline] + fn inner($this: $self_ty, $($arg: $ty),*) $(-> $ret)? { + $body + } + // SAFETY: `self.1` is constructed iff avx2 and fma are available. + unsafe { inner(self, $($arg),*) } + } + }; +} + +#[derive(Clone, Copy, Debug)] +#[repr(transparent)] +pub struct F32VecAvx(__m256, AvxDescriptor); + +impl F32SimdVec for F32VecAvx { + type Descriptor = AvxDescriptor; + + const LEN: usize = 8; + + #[inline(always)] + fn load(d: Self::Descriptor, mem: &[f32]) -> Self { + assert!(mem.len() >= Self::LEN); + // SAFETY: we just checked that `mem` has enough space. Moreover, we know avx is available + // from the safety invariant on `d`. + Self(unsafe { _mm256_loadu_ps(mem.as_ptr()) }, d) + } + + #[inline(always)] + fn store(&self, mem: &mut [f32]) { + assert!(mem.len() >= Self::LEN); + // SAFETY: we just checked that `mem` has enough space. Moreover, we know avx is available + // from the safety invariant on `self.1`. + unsafe { _mm256_storeu_ps(mem.as_mut_ptr(), self.0) } + } + + fn_avx!(this: F32VecAvx, fn mul_add(mul: F32VecAvx, add: F32VecAvx) -> F32VecAvx { + F32VecAvx(_mm256_fmadd_ps(this.0, mul.0, add.0), this.1) + }); + + fn splat(d: Self::Descriptor, v: f32) -> Self { + // SAFETY: We know avx is available from the safety invariant on `d`. + unsafe { Self(_mm256_set1_ps(v), d) } + } + + fn_avx!(this: F32VecAvx, fn abs() -> F32VecAvx { + F32VecAvx( + _mm256_castsi256_ps(_mm256_andnot_si256( + _mm256_set1_epi32(0b10000000000000000000000000000000u32 as i32), + _mm256_castps_si256(this.0), + )), + this.1) + }); + + fn_avx!(this: F32VecAvx, fn max(other: F32VecAvx) -> F32VecAvx { + F32VecAvx(_mm256_max_ps(this.0, other.0), this.1) + }); +} + +impl Add<F32VecAvx> for F32VecAvx { + type Output = F32VecAvx; + fn_avx!(this: F32VecAvx, fn add(rhs: F32VecAvx) -> F32VecAvx { + F32VecAvx(_mm256_add_ps(this.0, rhs.0), this.1) + }); +} + +impl Sub<F32VecAvx> for F32VecAvx { + type Output = F32VecAvx; + fn_avx!(this: F32VecAvx, fn sub(rhs: F32VecAvx) -> F32VecAvx { + F32VecAvx(_mm256_sub_ps(this.0, rhs.0), this.1) + }); +} + +impl Mul<F32VecAvx> for F32VecAvx { + type Output = F32VecAvx; + fn_avx!(this: F32VecAvx, fn mul(rhs: F32VecAvx) -> F32VecAvx { + F32VecAvx(_mm256_mul_ps(this.0, rhs.0), this.1) + }); +} + +impl Div<F32VecAvx> for F32VecAvx { + type Output = F32VecAvx; + fn_avx!(this: F32VecAvx, fn div(rhs: F32VecAvx) -> F32VecAvx { + F32VecAvx(_mm256_div_ps(this.0, rhs.0), this.1) + }); +} + +impl AddAssign<F32VecAvx> for F32VecAvx { + fn_avx!(this: &mut F32VecAvx, fn add_assign(rhs: F32VecAvx) { + this.0 = _mm256_add_ps(this.0, rhs.0) + }); +} + +impl SubAssign<F32VecAvx> for F32VecAvx { + fn_avx!(this: &mut F32VecAvx, fn sub_assign(rhs: F32VecAvx) { + this.0 = _mm256_sub_ps(this.0, rhs.0) + }); +} + +impl MulAssign<F32VecAvx> for F32VecAvx { + fn_avx!(this: &mut F32VecAvx, fn mul_assign(rhs: F32VecAvx) { + this.0 = _mm256_mul_ps(this.0, rhs.0) + }); +} + +impl DivAssign<F32VecAvx> for F32VecAvx { + fn_avx!(this: &mut F32VecAvx, fn div_assign(rhs: F32VecAvx) { + this.0 = _mm256_div_ps(this.0, rhs.0) + }); +} diff --git a/third_party/rust/jxl/src/simd/x86_64/avx512.rs b/third_party/rust/jxl/src/simd/x86_64/avx512.rs @@ -0,0 +1,149 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +use std::{ + arch::x86_64::{ + __m512, _mm512_add_ps, _mm512_andnot_si512, _mm512_castps_si512, _mm512_castsi512_ps, + _mm512_div_ps, _mm512_fmadd_ps, _mm512_loadu_ps, _mm512_max_ps, _mm512_mul_ps, + _mm512_set1_epi32, _mm512_set1_ps, _mm512_storeu_ps, _mm512_sub_ps, + }, + ops::{Add, AddAssign, Div, DivAssign, Mul, MulAssign, Sub, SubAssign}, +}; + +use super::super::{F32SimdVec, SimdDescriptor}; + +// Safety invariant: this type is only ever constructed if avx512f is available. +#[derive(Clone, Copy, Debug)] +pub struct Avx512Descriptor; + +impl SimdDescriptor for Avx512Descriptor { + type F32Vec = F32VecAvx512; + fn new() -> Option<Self> { + if is_x86_feature_detected!("avx512f") { + // SAFETY: we just checked avx512f. + Some(Self) + } else { + None + } + } +} + +// TODO(veluca): retire this macro once we have #[unsafe(target_feature)]. +macro_rules! fn_avx { + ( + $this:ident: $self_ty:ty, + fn $name:ident($($arg:ident: $ty:ty),* $(,)?) $(-> $ret:ty )? $body: block) => { + #[inline(always)] + fn $name(self: $self_ty, $($arg: $ty),*) $(-> $ret)? { + #[target_feature(enable = "avx512f")] + #[inline] + fn inner($this: $self_ty, $($arg: $ty),*) $(-> $ret)? { + $body + } + // SAFETY: `self.1` is constructed iff avx512f is available. + unsafe { inner(self, $($arg),*) } + } + }; +} + +#[derive(Clone, Copy, Debug)] +#[repr(transparent)] +pub struct F32VecAvx512(__m512, Avx512Descriptor); + +impl F32SimdVec for F32VecAvx512 { + type Descriptor = Avx512Descriptor; + + const LEN: usize = 16; + + #[inline(always)] + fn load(d: Self::Descriptor, mem: &[f32]) -> Self { + assert!(mem.len() >= Self::LEN); + // SAFETY: we just checked that `mem` has enough space. Moreover, we know avx512f is available + // from the safety invariant on `d`. + Self(unsafe { _mm512_loadu_ps(mem.as_ptr()) }, d) + } + + #[inline(always)] + fn store(&self, mem: &mut [f32]) { + assert!(mem.len() >= Self::LEN); + // SAFETY: we just checked that `mem` has enough space. Moreover, we know avx512f is available + // from the safety invariant on `self.1`. + unsafe { _mm512_storeu_ps(mem.as_mut_ptr(), self.0) } + } + + fn_avx!(this: F32VecAvx512, fn mul_add(mul: F32VecAvx512, add: F32VecAvx512) -> F32VecAvx512 { + F32VecAvx512(_mm512_fmadd_ps(this.0, mul.0, add.0), this.1) + }); + + fn splat(d: Self::Descriptor, v: f32) -> Self { + // SAFETY: We know avx512f is available from the safety invariant on `d`. + unsafe { Self(_mm512_set1_ps(v), d) } + } + + fn_avx!(this: F32VecAvx512, fn abs() -> F32VecAvx512 { + F32VecAvx512( + _mm512_castsi512_ps(_mm512_andnot_si512( + _mm512_set1_epi32(0b10000000000000000000000000000000u32 as i32), + _mm512_castps_si512(this.0), + )), + this.1) + }); + + fn_avx!(this: F32VecAvx512, fn max(other: F32VecAvx512) -> F32VecAvx512 { + F32VecAvx512(_mm512_max_ps(this.0, other.0), this.1) + }); +} + +impl Add<F32VecAvx512> for F32VecAvx512 { + type Output = F32VecAvx512; + fn_avx!(this: F32VecAvx512, fn add(rhs: F32VecAvx512) -> F32VecAvx512 { + F32VecAvx512(_mm512_add_ps(this.0, rhs.0), this.1) + }); +} + +impl Sub<F32VecAvx512> for F32VecAvx512 { + type Output = F32VecAvx512; + fn_avx!(this: F32VecAvx512, fn sub(rhs: F32VecAvx512) -> F32VecAvx512 { + F32VecAvx512(_mm512_sub_ps(this.0, rhs.0), this.1) + }); +} + +impl Mul<F32VecAvx512> for F32VecAvx512 { + type Output = F32VecAvx512; + fn_avx!(this: F32VecAvx512, fn mul(rhs: F32VecAvx512) -> F32VecAvx512 { + F32VecAvx512(_mm512_mul_ps(this.0, rhs.0), this.1) + }); +} + +impl Div<F32VecAvx512> for F32VecAvx512 { + type Output = F32VecAvx512; + fn_avx!(this: F32VecAvx512, fn div(rhs: F32VecAvx512) -> F32VecAvx512 { + F32VecAvx512(_mm512_div_ps(this.0, rhs.0), this.1) + }); +} + +impl AddAssign<F32VecAvx512> for F32VecAvx512 { + fn_avx!(this: &mut F32VecAvx512, fn add_assign(rhs: F32VecAvx512) { + this.0 = _mm512_add_ps(this.0, rhs.0) + }); +} + +impl SubAssign<F32VecAvx512> for F32VecAvx512 { + fn_avx!(this: &mut F32VecAvx512, fn sub_assign(rhs: F32VecAvx512) { + this.0 = _mm512_sub_ps(this.0, rhs.0) + }); +} + +impl MulAssign<F32VecAvx512> for F32VecAvx512 { + fn_avx!(this: &mut F32VecAvx512, fn mul_assign(rhs: F32VecAvx512) { + this.0 = _mm512_mul_ps(this.0, rhs.0) + }); +} + +impl DivAssign<F32VecAvx512> for F32VecAvx512 { + fn_avx!(this: &mut F32VecAvx512, fn div_assign(rhs: F32VecAvx512) { + this.0 = _mm512_div_ps(this.0, rhs.0) + }); +} diff --git a/third_party/rust/jxl/src/simd/x86_64/mod.rs b/third_party/rust/jxl/src/simd/x86_64/mod.rs @@ -0,0 +1,73 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +#![allow(unsafe_code)] + +pub(super) mod avx; +pub(super) mod avx512; + +macro_rules! simd_function { + ( + $dname:ident, + $descr:ident: $descr_ty:ident, + $pub:vis fn $name:ident($($arg:ident: $ty:ty),* $(,)?) $(-> $ret:ty )? $body: block + ) => { + #[inline(always)] + $pub fn $name<$descr_ty: crate::simd::SimdDescriptor>($descr: $descr_ty, $($arg: $ty),*) $(-> $ret)? $body + #[allow(unsafe_code)] + $pub fn $dname($($arg: $ty),*) $(-> $ret)? { + use crate::simd::SimdDescriptor; + if let Some(d) = crate::simd::Avx512Descriptor::new() { + #[target_feature(enable = "avx512f")] + fn inner(d: crate::simd::Avx512Descriptor, $($arg: $ty),*) $(-> $ret)? { + $name(d, $($arg),*) + } + // SAFETY: we just checked for avx512f. + return unsafe { inner(d, $($arg),*) }; + } + if let Some(d) = crate::simd::AvxDescriptor::new() { + #[target_feature(enable = "avx2,fma")] + fn inner(d: crate::simd::AvxDescriptor, $($arg: $ty),*) $(-> $ret)? { + $name(d, $($arg),*) + } + // SAFETY: we just checked for avx2 and fma. + return unsafe { inner(d, $($arg),*) }; + } + $name(crate::simd::ScalarDescriptor::new().unwrap(), $($arg),*) + } + }; +} + +pub(crate) use simd_function; + +#[allow(unused_macros)] +macro_rules! test_all_instruction_sets { + ( + $name:ident + ) => { + paste::paste! { + #[test] + fn [<$name _scalar>]() { + use crate::simd::SimdDescriptor; + $name(crate::simd::ScalarDescriptor::new().unwrap()) + } + #[test] + fn [<$name _avx>]() { + use crate::simd::SimdDescriptor; + let Some(d) = crate::simd::AvxDescriptor::new() else { return; }; + $name(d) + } + #[test] + fn [<$name _avx512>]() { + use crate::simd::SimdDescriptor; + let Some(d) = crate::simd::Avx512Descriptor::new() else { return; }; + $name(d) + } + } + }; +} + +#[allow(unused_imports)] +pub(crate) use test_all_instruction_sets; diff --git a/third_party/rust/jxl/src/util.rs b/third_party/rust/jxl/src/util.rs @@ -0,0 +1,30 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +#[cfg(test)] +pub mod test; + +mod bits; +mod concat_slice; +mod fast_math; +mod linalg; +mod log2; +pub mod ndarray; +mod rational_poly; +mod shift_right_ceil; +pub mod tracing_wrappers; +mod vec_helpers; +mod xorshift128plus; + +pub use bits::*; +pub use concat_slice::*; +pub use fast_math::*; +pub use linalg::*; +pub use log2::*; +pub(crate) use ndarray::*; +pub use rational_poly::*; +pub use shift_right_ceil::*; +pub use vec_helpers::*; +pub use xorshift128plus::*; diff --git a/third_party/rust/jxl/src/util/bits.rs b/third_party/rust/jxl/src/util/bits.rs @@ -0,0 +1,23 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +pub fn value_of_lowest_1_bit(t: u32) -> u32 { + t & t.wrapping_neg() +} +#[cfg(test)] +mod test { + use super::*; + + #[test] + fn test_value_of_lowest_1_bit() { + assert_eq!(value_of_lowest_1_bit(0b0001), 1); + assert_eq!(value_of_lowest_1_bit(0b1111), 1); + assert_eq!(value_of_lowest_1_bit(0b0010), 2); + assert_eq!(value_of_lowest_1_bit(0b0100), 4); + assert_eq!(value_of_lowest_1_bit(0b1010), 2); + assert_eq!(value_of_lowest_1_bit(0b1000_0000), 128); + assert_eq!(value_of_lowest_1_bit(0), 0); + } +} diff --git a/third_party/rust/jxl/src/util/concat_slice.rs b/third_party/rust/jxl/src/util/concat_slice.rs @@ -0,0 +1,123 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +use crate::error::Error; + +pub struct ConcatSlice<'first, 'second> { + slices: (&'first [u8], &'second [u8]), + ptr: usize, +} + +impl<'first, 'second> ConcatSlice<'first, 'second> { + pub fn new(slice0: &'first [u8], slice1: &'second [u8]) -> Self { + Self { + slices: (slice0, slice1), + ptr: 0, + } + } + + #[allow(clippy::len_without_is_empty)] + pub fn len(&self) -> usize { + self.slices.0.len() + self.slices.1.len() + } + + pub fn remaining_slices(&self) -> (&'first [u8], &'second [u8]) { + let (slice0, slice1) = self.slices; + let total_len = self.len(); + let ptr = self.ptr; + if ptr >= total_len { + (&[], &[]) + } else if let Some(second_slice_ptr) = ptr.checked_sub(slice0.len()) { + (&[], &slice1[second_slice_ptr..]) + } else { + (&slice0[ptr..], slice1) + } + } + + pub fn advance(&mut self, bytes: usize) { + self.ptr += bytes; + } + + pub fn peek<'out>(&self, out_buf: &'out mut [u8]) -> &'out mut [u8] { + let (slice0, slice1) = self.remaining_slices(); + let total_len = slice0.len() + slice1.len(); + + let out_bytes = out_buf.len().min(total_len); + let out_buf = &mut out_buf[..out_bytes]; + + if out_bytes <= slice0.len() { + out_buf.copy_from_slice(&slice0[..out_bytes]); + } else { + let (out_first, out_second) = out_buf.split_at_mut(slice0.len()); + out_first.copy_from_slice(slice0); + out_second.copy_from_slice(&slice1[..out_second.len()]); + } + + out_buf + } + + pub fn fill_vec(&mut self, max_bytes: Option<usize>, v: &mut Vec<u8>) -> Result<usize, Error> { + let (slice0, slice1) = self.remaining_slices(); + let total_len = slice0.len() + slice1.len(); + + let out_bytes = max_bytes.unwrap_or(usize::MAX).min(total_len); + v.try_reserve(out_bytes)?; + + if out_bytes <= slice0.len() { + v.extend_from_slice(&slice0[..out_bytes]); + } else { + let second_slice_len = out_bytes - slice0.len(); + v.extend_from_slice(slice0); + v.extend_from_slice(&slice1[..second_slice_len]); + } + + self.advance(out_bytes); + Ok(out_bytes) + } +} + +#[cfg(test)] +mod test { + use super::*; + + #[test] + fn peek_advance() { + let mut reader = ConcatSlice::new(&[0, 1, 2, 3], &[4, 5, 6, 7]); + let mut buf = [0u8; 8]; + + let actual = reader.peek(&mut buf[..1]); + assert_eq!(actual, &[0]); + reader.advance(actual.len()); + + let actual = reader.peek(&mut buf[..2]); + assert_eq!(actual, &[1, 2]); + reader.advance(actual.len()); + + let actual = reader.peek(&mut buf[..3]); + assert_eq!(actual, &[3, 4, 5]); + reader.advance(actual.len()); + + let actual = reader.peek(&mut buf); + assert_eq!(actual, &[6, 7]); + reader.advance(actual.len()); + + let actual = reader.peek(&mut buf); + assert!(actual.is_empty()); + } + + #[test] + fn fill_vec() { + let mut reader = ConcatSlice::new(&[0, 1, 2, 3], &[4, 5, 6, 7]); + let mut v = Vec::new(); + + let count = reader.fill_vec(Some(3), &mut v).unwrap(); + assert_eq!(count, 3); + assert_eq!(&v, &[0, 1, 2]); + + let count = reader.fill_vec(None, &mut v).unwrap(); + assert_eq!(count, 5); + assert_eq!(&v, &[0, 1, 2, 3, 4, 5, 6, 7]); + } +} diff --git a/third_party/rust/jxl/src/util/fast_math.rs b/third_party/rust/jxl/src/util/fast_math.rs @@ -0,0 +1,191 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +#![allow(clippy::excessive_precision)] + +use std::f32::consts::{PI, SQRT_2}; + +use super::eval_rational_poly; + +const POW2F_NUMER_COEFFS: [f32; 3] = [1.01749063e1, 4.88687798e1, 9.85506591e1]; +const POW2F_DENOM_COEFFS: [f32; 4] = [2.10242958e-1, -2.22328856e-2, -1.94414990e1, 9.85506633e1]; + +#[inline] +pub fn fast_cos(x: f32) -> f32 { + // Step 1: range reduction to [0, 2pi) + let pi2 = PI * 2.0; + let pi2_inv = 0.5 / PI; + let npi2 = (x * pi2_inv).floor() * pi2; + let xmodpi2 = x - npi2; + // Step 2: range reduction to [0, pi] + let x_pi = xmodpi2.min(pi2 - xmodpi2); + // Step 3: range reduction to [0, pi/2] + let above_pihalf = x_pi >= PI / 2.0; + let x_pihalf = if above_pihalf { PI - x_pi } else { x_pi }; + // Step 4: Taylor-like approximation, scaled by 2**0.75 to make angle + // duplication steps faster, on x/4. + let xs = x_pihalf * 0.25; + let x2 = xs * xs; + let x4 = x2 * x2; + let cosx_prescaling = x4 * 0.06960438 + (x2 * -0.84087373 + 1.68179268); + // Step 5: angle duplication. + let cosx_scale1 = cosx_prescaling * cosx_prescaling - SQRT_2; + let cosx_scale2 = cosx_scale1 * cosx_scale1 - 1.0; + // Step 6: change sign if needed. + if above_pihalf { + -cosx_scale2 + } else { + cosx_scale2 + } +} + +#[inline] +pub fn fast_erff(x: f32) -> f32 { + // Formula from + // https://en.wikipedia.org/wiki/Error_function#Numerical_approximations + // but constants have been recomputed. + let absx = x.abs(); + // Compute 1 - 1 / ((((x * a + b) * x + c) * x + d) * x + 1)**4 + let denom1 = absx * 7.77394369e-02 + 2.05260015e-04; + let denom2 = denom1 * absx + 2.32120216e-01; + let denom3 = denom2 * absx + 2.77820801e-01; + let denom4 = denom3 * absx + 1.0; + let denom5 = denom4 * denom4; + let inv_denom5 = 1.0 / denom5; + let result = -inv_denom5 * inv_denom5 + 1.0; + result.copysign(x) +} + +#[inline] +pub fn fast_pow2f(x: f32) -> f32 { + let x_floor = x.floor(); + let exp = f32::from_bits(((x_floor as i32 + 127) as u32) << 23); + let frac = x - x_floor; + + let num = frac + POW2F_NUMER_COEFFS[0]; + let num = num * frac + POW2F_NUMER_COEFFS[1]; + let num = num * frac + POW2F_NUMER_COEFFS[2]; + let num = num * exp; + + let den = POW2F_DENOM_COEFFS[0] * frac + POW2F_DENOM_COEFFS[1]; + let den = den * frac + POW2F_DENOM_COEFFS[2]; + let den = den * frac + POW2F_DENOM_COEFFS[3]; + + num / den +} + +const LOG2F_P: [f32; 3] = [ + -1.8503833400518310e-6, + 1.4287160470083755, + 7.4245873327820566e-1, +]; +const LOG2F_Q: [f32; 3] = [ + 9.9032814277590719e-1, + 1.0096718572241148, + 1.7409343003366853e-1, +]; + +#[inline] +pub fn fast_log2f(x: f32) -> f32 { + let x_bits = x.to_bits() as i32; + let exp_bits = x_bits.wrapping_sub(0x3f2aaaab); + let exp_shifted = exp_bits >> 23; + let mantissa = f32::from_bits((x_bits.wrapping_sub(exp_shifted << 23)) as u32); + let exp_val = exp_shifted as f32; + + let x = mantissa - 1.0; + eval_rational_poly(x, LOG2F_P, LOG2F_Q) + exp_val +} + +// Max relative error: ~3e-5 +#[inline] +pub fn fast_powf(base: f32, exp: f32) -> f32 { + fast_pow2f(fast_log2f(base) * exp) +} + +pub fn floor_log2_nonzero<T: num_traits::Unsigned + num_traits::PrimInt>(x: T) -> u32 { + (size_of::<T>() * 8 - 1) as u32 ^ x.leading_zeros() +} + +#[cfg(test)] +mod test { + use test_log::test; + + use crate::util::test::assert_almost_abs_eq; + + use super::*; + + #[test] + fn test_fast_erff() { + // Golden data copied from https://en.wikipedia.org/wiki/Error_function#Table_of_values. + let golden = [ + (0.0, 0.0), + (0.02, 0.022564575), + (0.04, 0.045111106), + (0.06, 0.067621594), + (0.08, 0.090078126), + (0.1, 0.112462916), + (0.2, 0.222702589), + (0.3, 0.328626759), + (0.4, 0.428392355), + (0.5, 0.520499878), + (0.6, 0.603856091), + (0.7, 0.677801194), + (0.8, 0.742100965), + (0.9, 0.796908212), + (1.0, 0.842700793), + (1.1, 0.880205070), + (1.2, 0.910313978), + (1.3, 0.934007945), + (1.4, 0.952285120), + (1.5, 0.966105146), + (1.6, 0.976348383), + (1.7, 0.983790459), + (1.8, 0.989090502), + (1.9, 0.992790429), + (2.0, 0.995322265), + (2.1, 0.997020533), + (2.2, 0.998137154), + (2.3, 0.998856823), + (2.4, 0.999311486), + (2.5, 0.999593048), + (3.0, 0.999977910), + (3.5, 0.999999257), + ]; + for (x, erf_x) in golden { + assert_almost_abs_eq(fast_erff(x), erf_x, 6e-4); + assert_almost_abs_eq(fast_erff(-x), -erf_x, 6e-4); + } + } + + #[test] + fn test_fast_cos() { + for i in 0..100 { + let x = i as f32 / 100.0 * (5.0 * PI) - (2.5 * PI); + assert_almost_abs_eq(fast_cos(x), x.cos(), 1e-4); + } + } + + #[test] + fn fast_powf_arb() { + arbtest::arbtest(|u| { + // (0.0, 128.0] + let base = u.int_in_range(1..=1 << 24)? as f32 / (1 << 17) as f32; + // [-4.0, 4.0] + let exp = u.int_in_range(-(1i32 << 22)..=1 << 22)? as f32 / (1 << 20) as f32; + + let expected = base.powf(exp); + let actual = fast_powf(base, exp); + let abs_error = (actual - expected).abs(); + let rel_error = abs_error / expected; + assert!( + rel_error < 3e-5, + "base: {base}, exp: {exp}, rel_error: {rel_error}, expected: {expected}, \ + actual: {actual}", + ); + Ok(()) + }); + } +} diff --git a/third_party/rust/jxl/src/util/linalg.rs b/third_party/rust/jxl/src/util/linalg.rs @@ -0,0 +1,140 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +use crate::error::Error; + +pub type Matrix3x3<T> = [[T; 3]; 3]; +pub type Vector3<T> = [T; 3]; + +pub fn matmul3_vec(m: [f32; 9], v: [f32; 3]) -> [f32; 3] { + [ + v[0] * m[0] + v[1] * m[1] + v[2] * m[2], + v[0] * m[3] + v[1] * m[4] + v[2] * m[5], + v[0] * m[6] + v[1] * m[7] + v[2] * m[8], + ] +} + +pub fn mul_3x3_vector(matrix: &Matrix3x3<f64>, vector: &Vector3<f64>) -> Vector3<f64> { + std::array::from_fn(|i| { + matrix[i] + .iter() + .zip(vector.iter()) + .map(|(&matrix_element, &vector_element)| matrix_element * vector_element) + .sum() + }) +} + +pub fn mul_3x3_matrix(mat1: &Matrix3x3<f64>, mat2: &Matrix3x3<f64>) -> Matrix3x3<f64> { + std::array::from_fn(|i| std::array::from_fn(|j| (0..3).map(|k| mat1[i][k] * mat2[k][j]).sum())) +} + +fn det2x2(a: f64, b: f64, c: f64, d: f64) -> f64 { + a * d - b * c +} + +fn calculate_cofactor(m: &Matrix3x3<f64>, r: usize, c: usize) -> f64 { + // Determine the actual row and column indices for the 2x2 submatrix + // by excluding the current row 'r' and column 'c'. + // Ensure they are taken in ascending order to form the submatrix consistently. + let mut sub_rows = [0; 2]; + let mut sub_cols = [0; 2]; + + let mut current_idx = 0; + for i in 0..3 { + if i != r { + sub_rows[current_idx] = i; + current_idx += 1; + } + } + + current_idx = 0; + for i in 0..3 { + if i != c { + sub_cols[current_idx] = i; + current_idx += 1; + } + } + + let minor_val = det2x2( + m[sub_rows[0]][sub_cols[0]], + m[sub_rows[0]][sub_cols[1]], + m[sub_rows[1]][sub_cols[0]], + m[sub_rows[1]][sub_cols[1]], + ); + + // Apply the checkerboard pattern sign for the cofactor + if (r + c) % 2 == 0 { + minor_val + } else { + -minor_val + } +} + +/// Calculates the inverse of a 3x3 matrix. +pub fn inv_3x3_matrix(m: &Matrix3x3<f64>) -> Result<Matrix3x3<f64>, Error> { + let cofactor_matrix: [[f64; 3]; 3] = std::array::from_fn(|r_idx| { + std::array::from_fn(|c_idx| calculate_cofactor(m, r_idx, c_idx)) + }); + + let det = m[0] + .iter() + .zip(cofactor_matrix[0].iter()) + .map(|(&m_element, &cof_element)| m_element * cof_element) + .sum::<f64>(); + + // Check for numerical singularity. + const EPSILON: f64 = 1e-12; + if det.abs() < EPSILON { + return Err(Error::MatrixInversionFailed(det.abs())); + } + + let inv_det = 1.0 / det; + + let adjugate_matrix: [[f64; 3]; 3] = + std::array::from_fn(|r_idx| std::array::from_fn(|c_idx| cofactor_matrix[c_idx][r_idx])); + + // Inverse matrix = (1/det) * Adjugate matrix. + Ok(std::array::from_fn(|r_idx| { + std::array::from_fn(|c_idx| adjugate_matrix[r_idx][c_idx] * inv_det) + })) +} + +#[cfg(test)] +mod test { + use super::*; + + fn assert_matrix_eq(a: &Matrix3x3<f64>, b: &Matrix3x3<f64>, epsilon: f64) { + for r in 0..3 { + for c in 0..3 { + assert!( + (a[r][c] - b[r][c]).abs() < epsilon, + "Matrices differ at [{}][{}]: expected {}, got {}. Diff: {}", + r, + c, + b[r][c], + a[r][c], + (a[r][c] - b[r][c]).abs() + ); + } + } + } + + #[test] + fn test_3x3_inverse() { + // Random matrix (https://xkcd.com/221/) + let m: Matrix3x3<f64> = [[1.0f64, -3.0, -2.0], [2.0, 2.0, 1.0], [2.0, 1.0, 1.0]]; + + let expected_inv: Matrix3x3<f64> = [[0.2, 0.2, 0.2], [0., 1., -1.], [-0.4, -1.4, 1.6]]; + + match inv_3x3_matrix(&m) { + Ok(inv_m) => { + assert_matrix_eq(&inv_m, &expected_inv, 1e-12); + } + Err(e) => { + panic!("Matrix inversion failed unexpectedly: {e:?}"); + } + } + } +} diff --git a/third_party/rust/jxl/src/util/log2.rs b/third_party/rust/jxl/src/util/log2.rs @@ -0,0 +1,85 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +pub trait FloorLog2 { + fn floor_log2(&self) -> Self; +} + +pub trait CeilLog2 { + fn ceil_log2(&self) -> Self; +} + +impl FloorLog2 for u32 { + fn floor_log2(&self) -> Self { + debug_assert_ne!(*self, 0); + 0u32.leading_zeros() - self.leading_zeros() - 1 + } +} + +impl FloorLog2 for u64 { + fn floor_log2(&self) -> Self { + debug_assert_ne!(*self, 0); + (0u64.leading_zeros() - self.leading_zeros() - 1) as u64 + } +} + +impl FloorLog2 for usize { + fn floor_log2(&self) -> Self { + debug_assert_ne!(*self, 0); + (0usize.leading_zeros() - self.leading_zeros() - 1) as usize + } +} + +impl<T> CeilLog2 for T +where + T: FloorLog2, + T: std::ops::Add<Output = Self>, + T: std::ops::Sub<Output = Self>, + T: std::ops::BitAnd<Output = Self>, + T: std::cmp::PartialEq, + T: From<u8>, + T: Copy, +{ + fn ceil_log2(&self) -> Self { + if (*self & (*self - 1.into())) != 0.into() { + self.floor_log2() + 1.into() + } else { + self.floor_log2() + } + } +} + +#[cfg(test)] +mod test { + use super::*; + #[test] + fn test_floor() { + assert_eq!(0, 1u32.floor_log2()); + assert_eq!(1, 2u32.floor_log2()); + assert_eq!(1, 3u32.floor_log2()); + assert_eq!(2, 4u32.floor_log2()); + } + #[test] + fn test_ceil() { + assert_eq!(0, 1u32.ceil_log2()); + assert_eq!(1, 2u32.ceil_log2()); + assert_eq!(2, 3u32.ceil_log2()); + assert_eq!(2, 4u32.ceil_log2()); + } + #[test] + fn test_floor_us() { + assert_eq!(0, 1usize.floor_log2()); + assert_eq!(1, 2usize.floor_log2()); + assert_eq!(1, 3usize.floor_log2()); + assert_eq!(2, 4usize.floor_log2()); + } + #[test] + fn test_ceil_us() { + assert_eq!(0, 1usize.ceil_log2()); + assert_eq!(1, 2usize.ceil_log2()); + assert_eq!(2, 3usize.ceil_log2()); + assert_eq!(2, 4usize.ceil_log2()); + } +} diff --git a/third_party/rust/jxl/src/util/ndarray.rs b/third_party/rust/jxl/src/util/ndarray.rs @@ -0,0 +1,24 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +macro_rules! slice { + (&$data:expr, $range:expr) => { + $data[$range] + }; + (&mut $data:expr, $range:expr) => { + $data[$range] + }; + (&mut $data:expr, $range:expr $(, $rest_ranges:expr)+) => { + $data[$range].iter_mut().map(|inner_data| { + &mut slice!(&inner_data, $($rest_ranges),+) + }).collect::<Vec<_>>() + }; + (&$data:expr, $range:expr $(, $rest_ranges:expr)+) => { + $data[$range].iter().map(|inner_data| { + &slice!(&inner_data, $($rest_ranges),+) + }).collect::<Vec<_>>() + }; +} +pub(crate) use slice; diff --git a/third_party/rust/jxl/src/util/rational_poly.rs b/third_party/rust/jxl/src/util/rational_poly.rs @@ -0,0 +1,15 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +/// Computes `(p0 + p1 x + p2 x^2 + ...) / (q0 + q1 x + q2 x^2 + ...)`. +/// +/// # Panics +/// Panics if either `P` or `Q` is zero. +#[inline] +pub fn eval_rational_poly<const P: usize, const Q: usize>(x: f32, p: [f32; P], q: [f32; Q]) -> f32 { + let yp = p.into_iter().rev().reduce(|yp, p| yp * x + p).unwrap(); + let yq = q.into_iter().rev().reduce(|yq, q| yq * x + q).unwrap(); + yp / yq +} diff --git a/third_party/rust/jxl/src/util/shift_right_ceil.rs b/third_party/rust/jxl/src/util/shift_right_ceil.rs @@ -0,0 +1,41 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +use std::ops::{Add, Shl, Shr, Sub}; + +pub trait ShiftRightCeil: Copy { + fn shrc<T: Copy>(self, rhs: T) -> Self + where + Self: Shr<T, Output = Self> + Shl<T, Output = Self>; +} + +impl<S: Copy + Add<Self, Output = Self> + Sub<Self, Output = Self> + From<u8>> ShiftRightCeil + for S +{ + fn shrc<T: Copy>(self, rhs: T) -> Self + where + Self: Shr<T, Output = Self> + Shl<T, Output = Self>, + { + (self + (Self::from(1u8) << rhs) - Self::from(1u8)) >> rhs + } +} + +#[cfg(test)] +mod test { + use crate::util::ShiftRightCeil; + + #[test] + fn test_shrc() { + assert_eq!(1u8, 1u8.shrc(1u8)); + assert_eq!(1u8, 2u8.shrc(1u8)); + assert_eq!(2u8, 9u8.shrc(3u8)); + assert_eq!(1u32, 1u32.shrc(1u32)); + assert_eq!(1u32, 2u32.shrc(1u32)); + assert_eq!(2u32, 9u32.shrc(3u32)); + assert_eq!(1u32, 1u32.shrc(1u8)); + assert_eq!(1u32, 2u32.shrc(1u8)); + assert_eq!(2u32, 9u32.shrc(3u8)); + } +} diff --git a/third_party/rust/jxl/src/util/test.rs b/third_party/rust/jxl/src/util/test.rs @@ -0,0 +1,347 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +use std::{ + fmt::Debug, + io::{BufRead, BufReader, Cursor, Read, Write}, + num::{ParseFloatError, ParseIntError}, +}; + +use crate::{ + bit_reader::BitReader, + container::ContainerParser, + error::Error as JXLError, + headers::{FileHeader, JxlHeader, encodings::*, frame_header::TocNonserialized}, + image::Image, +}; + +use num_traits::AsPrimitive; +use thiserror::Error; + +#[derive(Error, Debug)] +pub enum Error { + #[error("Invalid PFM: {0}")] + InvalidPFM(String), +} + +impl From<ParseFloatError> for Error { + fn from(value: ParseFloatError) -> Self { + Error::InvalidPFM(value.to_string()) + } +} + +impl From<ParseIntError> for Error { + fn from(value: ParseIntError) -> Self { + Error::InvalidPFM(value.to_string()) + } +} + +impl From<std::io::Error> for Error { + fn from(value: std::io::Error) -> Self { + Error::InvalidPFM(value.to_string()) + } +} + +impl From<JXLError> for Error { + fn from(value: JXLError) -> Self { + Error::InvalidPFM(value.to_string()) + } +} + +fn rel_error_gt<T: AsPrimitive<f64>>(left: T, right: T, max_rel_error: T) -> bool { + let left_f64: f64 = left.as_(); + let right_f64: f64 = right.as_(); + let error = (left_f64 - right_f64).abs(); + matches!( + (2.0 * error / (left_f64.abs() + right_f64.abs() + 1e-16)) + .partial_cmp(&max_rel_error.as_()), + Some(std::cmp::Ordering::Greater) | None + ) +} + +fn abs_error_gt<T: AsPrimitive<f64>>(left: T, right: T, max_abs_error: T) -> bool { + let left_f64: f64 = left.as_(); + let right_f64: f64 = right.as_(); + matches!( + (left_f64 - right_f64) + .abs() + .partial_cmp(&max_abs_error.as_()), + Some(std::cmp::Ordering::Greater) | None + ) +} + +pub fn assert_almost_eq<T: AsPrimitive<f64> + Debug + Copy>( + left: T, + right: T, + max_abs_error: T, + max_rel_error: T, +) { + if abs_error_gt(left, right, max_abs_error) || rel_error_gt(left, right, max_rel_error) { + panic!( + "assertion failed: `(left ≈ right)`\n left: `{left:?}`,\n right: `{right:?}`,\n max_abs_error: `{max_abs_error:?}`,\n max_rel_error: `{max_rel_error:?}`" + ); + } +} + +pub fn assert_almost_rel_eq<T: AsPrimitive<f64> + Debug + Copy>( + left: T, + right: T, + max_rel_error: T, +) { + if rel_error_gt(left, right, max_rel_error) { + panic!( + "assertion failed: `(left ≈ right)`\n left: `{left:?}`,\n right: `{right:?}`,\n max_rel_error: `{max_rel_error:?}`" + ); + } +} + +pub fn assert_almost_abs_eq<T: AsPrimitive<f64> + Debug + Copy>( + left: T, + right: T, + max_abs_error: T, +) { + if abs_error_gt(left, right, max_abs_error) { + panic!( + "assertion failed: `(left ≈ right)`\n left: `{left:?}`,\n right: `{right:?}`,\n max_abs_error: `{max_abs_error:?}`" + ); + } +} + +fn assert_same_len<T: AsPrimitive<f64> + Debug + Copy>(left: &[T], right: &[T]) { + if left.as_ref().len() != right.as_ref().len() { + panic!( + "assertion failed: `(left ≈ right)`\n left.len(): `{}`,\n right.len(): `{}`", + left.as_ref().len(), + right.as_ref().len() + ); + } +} + +pub fn assert_all_almost_eq<T: AsPrimitive<f64> + Debug + Copy, V: AsRef<[T]> + Debug>( + left: V, + right: V, + max_abs_error: T, + max_rel_error: T, +) { + assert_same_len(left.as_ref(), right.as_ref()); + for (idx, (left_val, right_val)) in left + .as_ref() + .iter() + .copied() + .zip(right.as_ref().iter().copied()) + .enumerate() + { + if abs_error_gt(left_val, right_val, max_abs_error) + || rel_error_gt(left_val, right_val, max_rel_error) + { + panic!( + "assertion failed: `(left ≈ right)`\n left: `{left:?}`,\n right: `{right:?}`,\n max_abs_error: `{max_abs_error:?}`,\n max_rel_error: `{max_rel_error:?}`,\n left[{idx}]: `{left_val:?}`,\n right[{idx}]: `{right_val:?}`", + ); + } + } +} + +pub fn assert_all_almost_rel_eq<T: AsPrimitive<f64> + Debug + Copy, V: AsRef<[T]> + Debug>( + left: V, + right: V, + max_rel_error: T, +) { + assert_same_len(left.as_ref(), right.as_ref()); + for (idx, (left_val, right_val)) in left + .as_ref() + .iter() + .copied() + .zip(right.as_ref().iter().copied()) + .enumerate() + { + if rel_error_gt(left_val, right_val, max_rel_error) { + panic!( + "assertion failed: `(left ≈ right)`\n left: `{left:?}`,\n right: `{right:?}`,\n max_rel_error: `{max_rel_error:?}`,\n left[{idx}]: `{left_val:?}`,\n right[{idx}]: `{right_val:?}`", + ); + } + } +} + +pub fn assert_all_almost_abs_eq<T: AsPrimitive<f64> + Debug + Copy, V: AsRef<[T]> + Debug>( + left: V, + right: V, + max_abs_error: T, +) { + assert_same_len(left.as_ref(), right.as_ref()); + for (idx, (left_val, right_val)) in left + .as_ref() + .iter() + .copied() + .zip(right.as_ref().iter().copied()) + .enumerate() + { + if abs_error_gt(left_val, right_val, max_abs_error) { + panic!( + "assertion failed: `(left ≈ right)`\n left: `{left:?}`,\n right: `{right:?}`,\n max_abs_error: `{max_abs_error:?}`,\n left[{idx}]: `{left_val:?}`,\n right[{idx}]: `{right_val:?}`", + ); + } + } +} + +pub fn read_headers_and_toc(image: &[u8]) -> Result<(FileHeader, FrameHeader, Toc), JXLError> { + let codestream = ContainerParser::collect_codestream(image).unwrap(); + let mut br = BitReader::new(&codestream); + let file_header = FileHeader::read(&mut br)?; + + let frame_header = + FrameHeader::read_unconditional(&(), &mut br, &file_header.frame_header_nonserialized())?; + let num_toc_entries = frame_header.num_toc_entries(); + let toc = Toc::read_unconditional( + &(), + &mut br, + &TocNonserialized { + num_entries: num_toc_entries as u32, + }, + )?; + Ok((file_header, frame_header, toc)) +} + +pub fn write_pfm(image: Vec<Image<f32>>, mut buf: impl Write) -> Result<(), Error> { + if image.len() == 1 { + buf.write_all(b"Pf\n")?; + } else if image.len() == 3 { + buf.write_all(b"PF\n")?; + } else { + return Err(Error::InvalidPFM(format!( + "invalid number of channels: {}", + image.len() + ))); + } + let size = image[0].size(); + for c in image.iter().skip(1) { + assert_eq!(size, c.size()); + } + buf.write_fmt(format_args!("{} {}\n", size.0, size.1))?; + buf.write_all(b"1.0\n")?; + let mut b: [u8; 4]; + for row in 0..size.1 { + for col in 0..size.0 { + for c in image.iter() { + b = c.as_rect().row(size.1 - row - 1)[col].to_be_bytes(); + buf.write_all(&b)?; + } + } + } + buf.flush()?; + Ok(()) +} + +pub fn read_pfm(b: &[u8]) -> Result<Vec<Image<f32>>, Error> { + let mut bf = BufReader::new(Cursor::new(b)); + let mut line = String::new(); + bf.read_line(&mut line)?; + let channels = match line.trim() { + "Pf" => 1, + "PF" => 3, + &_ => return Err(Error::InvalidPFM(format!("invalid PFM type header {line}"))), + }; + line.clear(); + bf.read_line(&mut line)?; + let mut dims = line.split_whitespace(); + let xres = if let Some(xres_str) = dims.next() { + xres_str.trim().parse()? + } else { + return Err(Error::InvalidPFM(format!( + "invalid PFM resolution header {line}", + ))); + }; + let yres = if let Some(yres_str) = dims.next() { + yres_str.trim().parse()? + } else { + return Err(Error::InvalidPFM(format!( + "invalid PFM resolution header {line}", + ))); + }; + line.clear(); + bf.read_line(&mut line)?; + let endianness: f32 = line.trim().parse()?; + + let mut res = Vec::<Image<f32>>::new(); + for _ in 0..channels { + let img = Image::new((xres, yres))?; + res.push(img); + } + + let mut buf = [0u8; 4]; + for row in 0..yres { + for col in 0..xres { + for chan in res.iter_mut() { + bf.read_exact(&mut buf)?; + chan.as_rect_mut().row(yres - row - 1)[col] = if endianness < 0.0 { + f32::from_le_bytes(buf) + } else { + f32::from_be_bytes(buf) + } + } + } + } + + Ok(res) +} + +use crate::headers::frame_header::{FrameHeader, Toc}; + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_with_floats() { + assert_almost_abs_eq(1.0000001f64, 1.0000002, 0.000001); + assert_almost_abs_eq(1.0, 1.1, 0.2); + } + + #[test] + fn test_with_integers() { + assert_almost_abs_eq(100, 101, 2); + assert_almost_abs_eq(777u32, 770, 7); + assert_almost_abs_eq(500i64, 498, 3); + } + + #[test] + #[should_panic] + fn test_panic_float() { + assert_almost_abs_eq(1.0, 1.2, 0.1); + } + #[test] + #[should_panic] + fn test_panic_integer() { + assert_almost_abs_eq(100, 105, 2); + } + + #[test] + #[should_panic] + fn test_nan_comparison() { + assert_almost_abs_eq(f64::NAN, f64::NAN, 0.1); + } + + #[test] + #[should_panic] + fn test_nan_tolerance() { + assert_almost_abs_eq(1.0, 1.0, f64::NAN); + } + + #[test] + fn test_infinity_tolerance() { + assert_almost_abs_eq(1.0, 1.0, f64::INFINITY); + } + + #[test] + #[should_panic] + fn test_nan_comparison_with_infinity_tolerance() { + assert_almost_abs_eq(f32::NAN, f32::NAN, f32::INFINITY); + } + + #[test] + #[should_panic] + fn test_infinity_comparison_with_infinity_tolerance() { + assert_almost_abs_eq(f32::INFINITY, f32::INFINITY, f32::INFINITY); + } +} diff --git a/third_party/rust/jxl/src/util/tracing_wrappers.rs b/third_party/rust/jxl/src/util/tracing_wrappers.rs @@ -0,0 +1,26 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +#![allow(unused_imports)] + +#[cfg(feature = "tracing")] +mod private { + pub use tracing::{debug, error, info, instrument, trace, warn}; +} + +#[cfg(not(feature = "tracing"))] +mod private { + macro_rules! fake_log { + ($($_: tt)*) => {}; + } + pub(crate) use fake_log as debug; + pub(crate) use fake_log as error; + pub(crate) use fake_log as info; + pub(crate) use fake_log as trace; + pub(crate) use fake_log as warn; + pub use jxl_macros::noop as instrument; +} + +pub use private::*; diff --git a/third_party/rust/jxl/src/util/vec_helpers.rs b/third_party/rust/jxl/src/util/vec_helpers.rs @@ -0,0 +1,33 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +// TODO(firsching): as soon as "Vec::try_with_capacity" is available from the +// standard library use this instead of the functions here. +pub trait NewWithCapacity { + type Output; + type Error; + fn new_with_capacity(capacity: usize) -> Result<Self::Output, Self::Error>; +} + +impl<T> NewWithCapacity for Vec<T> { + type Output = Vec<T>; + type Error = std::collections::TryReserveError; + + fn new_with_capacity(capacity: usize) -> Result<Self::Output, Self::Error> { + let mut vec = Vec::new(); + vec.try_reserve(capacity)?; + Ok(vec) + } +} + +impl NewWithCapacity for String { + type Output = String; + type Error = std::collections::TryReserveError; + fn new_with_capacity(capacity: usize) -> Result<Self::Output, Self::Error> { + let mut s = String::new(); + s.try_reserve(capacity)?; + Ok(s) + } +} diff --git a/third_party/rust/jxl/src/util/xorshift128plus.rs b/third_party/rust/jxl/src/util/xorshift128plus.rs @@ -0,0 +1,733 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +// Adapted from https://github.com/vpxyz/xorshift/blob/master/xorshift128plus/ +// (MIT-license) + +pub struct Xorshift128Plus { + s0: [u64; Self::N], + s1: [u64; Self::N], +} + +impl Xorshift128Plus { + pub const N: usize = 8; + + pub fn new_with_seed(seed: u64) -> Self { + let mut s0 = [0; Self::N]; + let mut s1 = [0; Self::N]; + + s0[0] = Self::split_mix_64(seed + 0x9E3779B97F4A7C15); + s1[0] = Self::split_mix_64(s0[0]); + + for i in 1..Self::N { + s0[i] = Self::split_mix_64(s1[i - 1]); + s1[i] = Self::split_mix_64(s0[i]); + } + + Self { s0, s1 } + } + + pub fn new_with_seeds(seed1: u32, seed2: u32, seed3: u32, seed4: u32) -> Self { + let mut s0 = [0; Self::N]; + let mut s1 = [0; Self::N]; + + s0[0] = Self::split_mix_64( + (((seed1 as u64) << 32) + seed2 as u64).wrapping_add(0x9E3779B97F4A7C15), + ); + s1[0] = Self::split_mix_64( + (((seed3 as u64) << 32) + seed4 as u64).wrapping_add(0x9E3779B97F4A7C15), + ); + for i in 1..Self::N { + s0[i] = Self::split_mix_64(s0[i - 1]); + s1[i] = Self::split_mix_64(s1[i - 1]); + } + + Self { s0, s1 } + } + + pub fn fill(&mut self, random_bits: &mut [u64; Self::N]) { + for ((s0, s1), random_bits) in self + .s0 + .iter_mut() + .zip(self.s1.iter_mut()) + .zip(random_bits.iter_mut()) + { + let mut new_s1 = *s0; + *s0 = *s1; + let bits = new_s1.wrapping_add(*s0); // b, c + new_s1 ^= new_s1 << 23; + *random_bits = bits; + new_s1 ^= *s0 ^ (new_s1 >> 18) ^ (*s0 >> 5); + *s1 = new_s1; + } + } + + fn split_mix_64(mut z: u64) -> u64 { + z = (z ^ (z >> 30)).wrapping_mul(0xBF58476D1CE4E5B9); + z = (z ^ (z >> 27)).wrapping_mul(0x94D049BB133111EB); + z ^ (z >> 31) + } +} + +#[cfg(test)] +mod test { + use crate::util::xorshift128plus::Xorshift128Plus; + + #[test] + fn xorshift128plus_golden() { + const NUM_VECTORS: usize = 64; + const EXPECTED: [[u64; Xorshift128Plus::N]; NUM_VECTORS] = [ + [ + 0x6E901576D477CBB1, + 0xE9E53789195DA2A2, + 0xB681F6DDA5E0AE99, + 0x8EFD18CE21FD6896, + 0xA898A80DF75CF532, + 0x50CEB2C9E2DE7E32, + 0x3CA7C2FEB25C0DD0, + 0xA4D0866B80B4D836, + ], + [ + 0x8CD6A1E6233D3A26, + 0x3D4603ADE98B112D, + 0xDC427AF674019E36, + 0xE28B4D230705AC53, + 0x7297E9BBA88783DD, + 0x34D3D23CFCD9B41A, + 0x5A223615ADBE96B8, + 0xE5EB529027CFBD01, + ], + [ + 0xC1894CF00DFAC6A2, + 0x18EDF8AE9085E404, + 0x8E936625296B4CCD, + 0x31971EF3A14A899B, + 0xBE87535FCE0BF26A, + 0x576F7A752BC6649F, + 0xA44CBADCE0C6B937, + 0x3DBA819BB17A353A, + ], + [ + 0x27CE38DFCC1C5EB6, + 0x920BEB5606340256, + 0x3986CBC40C9AFC2C, + 0xE22BCB3EEB1E191E, + 0x6E1FCDD3602A8FBA, + 0x052CB044E5415A29, + 0x46266646EFB9ECD7, + 0x8F44914618D29335, + ], + [ + 0xDD30AEDF72A362C5, + 0xBC1D824E16BB98F4, + 0x9EA6009C2AA3D2F1, + 0xF65C0FBBE17AF081, + 0x22424D06A8738991, + 0x8A62763F2B7611D2, + 0x2F3E89F722637939, + 0x84D338BEF50AFD50, + ], + [ + 0x00F46494898E2B0B, + 0x81239DC4FB8E8003, + 0x414AD93EC5773FE7, + 0x791473C450E4110F, + 0x87F127BF68C959AC, + 0x6429282D695EF67B, + 0x661082E11546CBA8, + 0x5815D53FA5436BFD, + ], + [ + 0xB3DEADAB9BE6E0F9, + 0xAA1B7B8F7CED0202, + 0x4C5ED437699D279E, + 0xA4471727F1CB39D3, + 0xE439DA193F802F70, + 0xF89401BB04FA6493, + 0x3B08045A4FE898BA, + 0x32137BFE98227950, + ], + [ + 0xFBAE4A092897FEF3, + 0x0639F6CE56E71C8E, + 0xF0AD6465C07F0C1E, + 0xFF8E28563361DCE5, + 0xC2013DB7F86BC6B9, + 0x8EFCC0503330102F, + 0x3F6B767EA5C4DA40, + 0xB9864B950B2232E1, + ], + [ + 0x76EB58DE8E5EC22A, + 0x9BBBF49A18B32F4F, + 0xC8405F02B2B2FAB9, + 0xC3E122A5F146BC34, + 0xC90BB046660F5765, + 0xB933981310DBECCF, + 0x5A2A7BFC9126FD1C, + 0x8BB388C94DF87901, + ], + [ + 0x753EB89AD63EF3C3, + 0xF24AAF40C89D65AD, + 0x23F68931C1A6AA6D, + 0xF47E79BF702C6DD0, + 0xA3AD113244EE7EAE, + 0xD42CBEA28F793DC3, + 0xD896FCF1820F497C, + 0x042B86D2818948C1, + ], + [ + 0x8F2A4FC5A4265763, + 0xEC499E6F95EAA10C, + 0xE3786D4ECCD0DEB5, + 0xC725C53D3AC4CC43, + 0x065A4ACBBF83610E, + 0x35C61C9FEF167129, + 0x7B720AEAA7D70048, + 0x14206B841377D039, + ], + [ + 0xAD27D78BF96055F6, + 0x5F43B20FF47ADCD4, + 0xE184C2401E2BF71E, + 0x30B263D78990045D, + 0xC22F00EBFF9BA201, + 0xAE7F86522B53A562, + 0x2853312BC039F0A4, + 0x868D619E6549C3C8, + ], + [ + 0xFD5493D8AE9A8371, + 0x773D5E224DF61B3B, + 0x5377C54FBB1A8280, + 0xCAD4DE3B8265CAFA, + 0xCDF3F19C91EBD5F6, + 0xC8EA0F182D73BD78, + 0x220502D593433FF1, + 0xB81205E612DC31B1, + ], + [ + 0x8F32A39EAEDA4C70, + 0x1D4B0914AA4DAC7F, + 0x56EF1570F3A8B405, + 0x29812CB17404A592, + 0x97A2AAF69CAE90F2, + 0x12BF5E02778BBFE5, + 0x9D4B55AD42A05FD2, + 0x06C2BAB5E6086620, + ], + [ + 0x8DB4B9648302B253, + 0xD756AD9E3AEA12C7, + 0x68709B7F11D4B188, + 0x7CC299DDCD707A4B, + 0x97B860C370A7661D, + 0xCECD314FC20E64F5, + 0x55F412CDFB4C7EC3, + 0x55EE97591193B525, + ], + [ + 0xCF70F3ACA96E6254, + 0x022FEDECA2E09F46, + 0x686823DB60AE1ECF, + 0xFD36190D3739830E, + 0x74E1C09027F68120, + 0xB5883A835C093842, + 0x93E1EFB927E9E4E3, + 0xB2721E249D7E5EBE, + ], + [ + 0x69B6E21C44188CB8, + 0x5D6CFB853655A7AA, + 0x3E001A0B425A66DC, + 0x8C57451103A5138F, + 0x7BF8B4BE18EAB402, + 0x494102EB8761A365, + 0xB33796A9F6A81F0E, + 0x10005AB3BCCFD960, + ], + [ + 0xB2CF25740AE965DC, + 0x6F7C1DF7EF53D670, + 0x648DD6087AC2251E, + 0x040955D9851D487D, + 0xBD550FC7E21A7F66, + 0x57408F484DEB3AB5, + 0x481E24C150B506C1, + 0x72C0C3EAF91A40D6, + ], + [ + 0x1997A481858A5D39, + 0x539718F4BEF50DC1, + 0x2EC4DC4787E7E368, + 0xFF1CE78879419845, + 0xE219A93DD6F6DD30, + 0x85328618D02FEC1A, + 0xC86E02D969181B20, + 0xEBEC8CD8BBA34E6E, + ], + [ + 0x28B55088A16CE947, + 0xDD25AC11E6350195, + 0xBD1F176694257B1C, + 0x09459CCF9FCC9402, + 0xF8047341E386C4E4, + 0x7E8E9A9AD984C6C0, + 0xA4661E95062AA092, + 0x70A9947005ED1152, + ], + [ + 0x4C01CF75DBE98CCD, + 0x0BA076CDFC7373B9, + 0x6C5E7A004B57FB59, + 0x336B82297FD3BC56, + 0x7990C0BE74E8D60F, + 0xF0275CC00EC5C8C8, + 0x6CF29E682DFAD2E9, + 0xFA4361524BD95D72, + ], + [ + 0x631D2A19FF62F018, + 0x41C43863B985B3FA, + 0xE052B2267038EFD9, + 0xE2A535FAC575F430, + 0xE004EEA90B1FF5B8, + 0x42DFE2CA692A1F26, + 0x90FB0BFC9A189ECC, + 0x4484102BD3536BD0, + ], + [ + 0xD027134E9ACCA5A5, + 0xBBAB4F966D476A9B, + 0x713794A96E03D693, + 0x9F6335E6B94CD44A, + 0xC5090C80E7471617, + 0x6D9C1B0C87B58E33, + 0x1969CE82E31185A5, + 0x2099B97E87754EBE, + ], + [ + 0x60EBAF4ED934350F, + 0xC26FBF0BA5E6ECFF, + 0x9E54150F0312EC57, + 0x0973B48364ED0041, + 0x800A523241426CFC, + 0x03AB5EC055F75989, + 0x8CF315935DEEB40A, + 0x83D3FC0190BD1409, + ], + [ + 0x26D35394CF720A51, + 0xCE9EAA15243CBAFE, + 0xE2B45FBAF21B29E0, + 0xDB92E98EDE73F9E0, + 0x79B16F5101C26387, + 0x1AC15959DE88C86F, + 0x387633AEC6D6A580, + 0xA6FC05807BFC5EB8, + ], + [ + 0x2D26C8E47C6BADA9, + 0x820E6EC832D52D73, + 0xB8432C3E0ED0EE5B, + 0x0F84B3C4063AAA87, + 0xF393E4366854F651, + 0x749E1B4D2366A567, + 0x805EACA43480D004, + 0x244EBF3AA54400A5, + ], + [ + 0xBFDC3763AA79F75A, + 0x9E3A74CC751F41DB, + 0xF401302A149DBC55, + 0x6B25F7973D7BF7BC, + 0x13371D34FDBC3DAE, + 0xC5E1998C8F484DCD, + 0x7031B8AE5C364464, + 0x3847F0C4F3DA2C25, + ], + [ + 0x24C6387D2C0F1225, + 0x77CCE960255C67A4, + 0x21A0947E497B10EB, + 0xBB5DB73A825A9D7E, + 0x26294A41999E553D, + 0x3953E0089F87D925, + 0x3DAE6E5D4E5EAAFE, + 0x74B545460341A7AA, + ], + [ + 0x710E5EB08A7DB820, + 0x7E43C4E77CAEA025, + 0xD4C91529C8B060C1, + 0x09AE26D8A7B0CA29, + 0xAB9F356BB360A772, + 0xB68834A25F19F6E9, + 0x79B8D9894C5734E2, + 0xC6847E7C8FFD265F, + ], + [ + 0x10C4BCB06A5111E6, + 0x57CB50955B6A2516, + 0xEF53C87798B6995F, + 0xAB38E15BBD8D0197, + 0xA51C6106EFF73C93, + 0x83D7F0E2270A7134, + 0x0923FD330397FCE5, + 0xF9DE54EDFE58FB45, + ], + [ + 0x07D44833ACCD1A94, + 0xAAD3C9E945E2F9F3, + 0xABF4C879B876AA37, + 0xF29C69A21B301619, + 0x2DDCE959111C788B, + 0x7CEDB48F8AC1729B, + 0x93F3BA9A02B659BE, + 0xF20A87FF17933CBE, + ], + [ + 0x8E96EBE93180CFE6, + 0x94CAA12873937079, + 0x05F613D9380D4189, + 0xBCAB40C1DC79F38A, + 0x0AD8907B7C61D19E, + 0x88534E189D103910, + 0x2DB2FAABA160AB8F, + 0xA070E7506B06F15C, + ], + [ + 0x6FB1FCDAFFEF87A9, + 0xE735CF25337A090D, + 0x172C6EDCEFEF1825, + 0x76957EA49EF0542D, + 0x819BF4CD250F7C49, + 0xD6FF23E4AD00C4D4, + 0xE79673C1EC358FF0, + 0xAC9C048144337938, + ], + [ + 0x4C5387FF258B3AF4, + 0xEDB68FAEC2CB1AA3, + 0x02A624E67B4E1DA4, + 0x5C44797A38E08AF2, + 0x36546A70E9411B4B, + 0x47C17B24D2FD9675, + 0x101957AAA020CA26, + 0x47A1619D4779F122, + ], + [ + 0xF84B8BCDC92D9A3C, + 0x951D7D2C74B3066B, + 0x7AC287C06EDDD9B2, + 0x4C38FC476608D38F, + 0x224D793B19CB4BCD, + 0x835A255899BF1A41, + 0x4AD250E9F62DB4AB, + 0xD9B44F4B58781096, + ], + [ + 0xABBAF99A8EB5C6B8, + 0xFB568E900D3A9F56, + 0x11EDF63D23C5DF11, + 0xA9C3011D3FA7C5A8, + 0xAEDD3CF11AFFF725, + 0xABCA472B5F1EDD6B, + 0x0600B6BB5D879804, + 0xDB4DE007F22191A0, + ], + [ + 0xD76CC9EFF0CE9392, + 0xF5E0A772B59BA49A, + 0x7D1AE1ED0C1261B5, + 0x79224A33B5EA4F4A, + 0x6DD825D80C40EA60, + 0x47FC8E747E51C953, + 0x695C05F72888BF98, + 0x1A012428440B9015, + ], + [ + 0xD754DD61F9B772BF, + 0xC4A2FCF4C0F9D4EB, + 0x461167CDF67A24A2, + 0x434748490EBCB9D4, + 0x274DD9CDCA5781DE, + 0x36BAC63BA9A85209, + 0x30324DAFDA36B70F, + 0x337570DB4FE6DAB3, + ], + [ + 0xF46CBDD57C551546, + 0x8E02507E676DA3E3, + 0xD826245A8C15406D, + 0xDFB38A5B71113B72, + 0x5EA38454C95B16B5, + 0x28C054FB87ABF3E1, + 0xAA2724C0BA1A8096, + 0xECA83EC980304F2F, + ], + [ + 0x6AA76EC294EB3303, + 0x42D4CDB2A8032E3B, + 0x7999EDF75DCD8735, + 0xB422BFFE696CCDCC, + 0x8F721461FD7CCDFE, + 0x148E1A5814FDE253, + 0x4DC941F4375EF8FF, + 0x27B2A9E0EB5B49CF, + ], + [ + 0xCEA592EF9343EBE1, + 0xF7D38B5FA7698903, + 0x6CCBF352203FEAB6, + 0x830F3095FCCDA9C5, + 0xDBEEF4B81B81C8F4, + 0x6D7EB9BCEECA5CF9, + 0xC58ABB0FBE436C69, + 0xE4B97E6DB2041A4B, + ], + [ + 0x7E40FC772978AF14, + 0xCDDA4BBAE28354A1, + 0xE4F993B832C32613, + 0xD3608093C68A4B35, + 0x9A3B60E01BEE3699, + 0x03BEF248F3288713, + 0x70B9294318F3E9B4, + 0x8D2ABB913B8610DE, + ], + [ + 0x37F209128E7D8B2C, + 0x81D2AB375BD874BC, + 0xA716A1B7373F7408, + 0x0CEE97BEC4706540, + 0xA40C5FD9CDBC1512, + 0x73CAF6C8918409E7, + 0x45E11BCEDF0BBAA1, + 0x612C612BFF6E6605, + ], + [ + 0xF8ECB14A12D0F649, + 0xDA683CD7C01BA1AC, + 0xA2203F7510E124C1, + 0x7F83E52E162F3C78, + 0x77D2BB73456ACADB, + 0x37FC34FC840BBA6F, + 0x3076BC7D4C6EBC1F, + 0x4F514123632B5FA9, + ], + [ + 0x44D789DED935E884, + 0xF8291591E09FEC9F, + 0xD9CED2CF32A2E4B7, + 0x95F70E1EB604904A, + 0xDE438FE43C14F6AB, + 0x4C8D23E4FAFCF8D8, + 0xC716910A3067EB86, + 0x3D6B7915315095D3, + ], + [ + 0x3170FDBADAB92095, + 0x8F1963933FC5650B, + 0x72F94F00ABECFEAB, + 0x6E3AE826C6AAB4CE, + 0xA677A2BF31068258, + 0x9660CDC4F363AF10, + 0xD81A15A152379EF1, + 0x5D7D285E1080A3F9, + ], + [ + 0xDAD5DDFF9A2249B3, + 0x6F9721D926103FAE, + 0x1418CBB83FFA349A, + 0xE71A30AD48C012B2, + 0xBE76376C63751132, + 0x3496467ACA713AE6, + 0x8D7EC01369F991A3, + 0xD8C73A88B96B154E, + ], + [ + 0x8B5D9C74AEB4833A, + 0xF914FB3F867B912F, + 0xB894EA034936B1DC, + 0x8A16D21BE51C4F5B, + 0x31FF048ED582D98E, + 0xB95AB2F4DC65B820, + 0x04082B9170561AF7, + 0xA215610A5DC836FA, + ], + [ + 0xB2ADE592C092FAAC, + 0x7A1E683BCBF13294, + 0xC7A4DBF86858C096, + 0x3A49940F97BFF316, + 0xCAE5C06B82C46703, + 0xC7F413A0F951E2BD, + 0x6665E7BB10EB5916, + 0x86F84A5A94EDE319, + ], + [ + 0x4EA199D8FAA79CA3, + 0xDFA26E5BF1981704, + 0x0F5E081D37FA4E01, + 0x9CB632F89CD675CD, + 0x4A09DB89D48C0304, + 0x88142742EA3C7672, + 0xAC4F149E6D2E9BDB, + 0x6D9E1C23F8B1C6C6, + ], + [ + 0xD58BE47B92DEC0E9, + 0x8E57573645E34328, + 0x4CC094CCB5FB5126, + 0x5F1D66AF6FB40E3C, + 0x2BA15509132D3B00, + 0x0D6545646120E567, + 0x3CF680C45C223666, + 0x96B28E32930179DA, + ], + [ + 0x5900C45853AC7990, + 0x61881E3E8B7FF169, + 0x4DE5F835DF2230FF, + 0x4427A9E7932F73FF, + 0x9B641BAD379A8C8D, + 0xDF271E5BF98F4E5C, + 0xDFDA16DB830FF5EE, + 0x371C7E7CFB89C0E9, + ], + [ + 0x4410A8576247A250, + 0x6AD2DA12B45AC0D9, + 0x18DFC72AAC85EECC, + 0x06FC8BB2A0EF25C8, + 0xEB287619C85E6118, + 0x19553ECA67F25A2C, + 0x3B9557F1DCEC5BAA, + 0x7BAD9E8B710D1079, + ], + [ + 0x34F365D66BD22B28, + 0xE6E124B9F10F835D, + 0x0573C38ABF2B24DC, + 0xD32E6AF10A0125AE, + 0x383590ACEA979519, + 0x8376ED7A39E28205, + 0xF0B7F184DCBDA435, + 0x062A203390E31794, + ], + [ + 0xA2AFFD7E41918760, + 0x7F90FC1BD0819C86, + 0x5033C08E5A969533, + 0x2707AF5C6D039590, + 0x57BBD5980F17DF9C, + 0xD3FE6E61D763268A, + 0x9E0A0AE40F335A3B, + 0x43CF4EB0A99613C5, + ], + [ + 0xD4D2A397CE1A7C2E, + 0x3DF7CE7CC3212DAD, + 0x0880F0D5D356C75A, + 0xA8AFC44DD03B1346, + 0x79263B46C13A29E0, + 0x11071B3C0ED58E7A, + 0xED46DC9F538406BF, + 0x2C94974F2B94843D, + ], + [ + 0xE246E13C39AB5D5E, + 0xAC1018489D955B20, + 0x8601B558771852B8, + 0x110BD4C06DB40173, + 0x738FC8A18CCA0EBB, + 0x6673E09BE0EA76E5, + 0x024BC7A0C7527877, + 0x45E6B4652E2EC34E, + ], + [ + 0xD1ED26A1A375CDC8, + 0xAABC4E896A617CB8, + 0x0A9C9E8E57D753C6, + 0xA3774A75FEB4C30E, + 0x30B816C01C93E49E, + 0xF405BABC06D2408C, + 0xCC0CE6B4CE788ABC, + 0x75E7922D0447956C, + ], + [ + 0xD07C1676A698BC95, + 0x5F9AEA4840E2D860, + 0xD5FC10D58BDF6F02, + 0xF190A2AD4BC2EEA7, + 0x0C24D11F51726931, + 0xDB646899A16B6512, + 0x7BC10670047B1DD8, + 0x2413A5ABCD45F092, + ], + [ + 0x4E66892190CFD923, + 0xF10162440365EC8E, + 0x158ACA5A6A2280AE, + 0x0D60ED11C0224166, + 0x7CD2E9A71B9D7488, + 0x450D7289706AB2A3, + 0x88FAE34EC9A0D7DC, + 0x96FF9103575A97DA, + ], + [ + 0x77990FAC6046C446, + 0xB174B5FB30C76676, + 0xE352CE3EB56CF82A, + 0xC6039B6873A9A082, + 0xE3F80F3AE333148A, + 0xB853BA24BA3539B9, + 0xE8863E52ECCB0C74, + 0x309B4CC1092CC245, + ], + [ + 0xBC2B70BEE8388D9F, + 0xE48D92AE22216DCE, + 0xF15F3BF3E2C15D8F, + 0x1DD964D4812D8B24, + 0xD56AF02FB4665E4C, + 0x98002200595BD9A3, + 0x049246D50BB8FA12, + 0x1B542DF485B579B9, + ], + [ + 0x2347409ADFA8E497, + 0x36015C2211D62498, + 0xE9F141F32EB82690, + 0x1F839912D0449FB9, + 0x4E4DCFFF2D02D97C, + 0xF8A03AB4C0F625C9, + 0x0605F575795DAC5C, + 0x4746C9BEA0DDA6B1, + ], + [ + 0xCA5BB519ECE7481B, + 0xFD496155E55CA945, + 0xF753B9DBB1515F81, + 0x50549E8BAC0F70E7, + 0x8614FB0271E21C60, + 0x60C72947EB0F0070, + 0xA6511C10AEE742B6, + 0x48FB48F2CACCB43E, + ], + ]; + + let mut rng = Xorshift128Plus::new_with_seed(12345); + for expected_row in EXPECTED.iter() { + let mut row = [0; Xorshift128Plus::N]; + rng.fill(&mut row); + for (&actual, &expected) in row.iter().zip(expected_row) { + assert_eq!(actual, expected); + } + } + } +} diff --git a/third_party/rust/jxl/src/var_dct.rs b/third_party/rust/jxl/src/var_dct.rs @@ -0,0 +1,9 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +pub mod dct; +pub mod dct_scales; +mod dct_slow; +pub mod transform; diff --git a/third_party/rust/jxl/src/var_dct/dct.rs b/third_party/rust/jxl/src/var_dct/dct.rs @@ -0,0 +1,830 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +use std::f64::consts::SQRT_2; + +use super::dct_scales::WcMultipliers; + +struct CoeffBundle<const N: usize, const SZ: usize>; + +pub struct DCT1DImpl<const SIZE: usize>; +pub struct IDCT1DImpl<const SIZE: usize>; + +pub trait DCT1D { + fn do_dct<const COLUMNS: usize>(data: &mut [[f32; COLUMNS]]); +} +pub trait IDCT1D { + fn do_idct<const COLUMNS: usize>(data: &mut [[f32; COLUMNS]]); +} + +impl DCT1D for DCT1DImpl<1> { + fn do_dct<const COLUMNS: usize>(_data: &mut [[f32; COLUMNS]]) { + // Do nothing + } +} +impl IDCT1D for IDCT1DImpl<1> { + fn do_idct<const COLUMNS: usize>(_data: &mut [[f32; COLUMNS]]) { + // Do nothing + } +} + +impl DCT1D for DCT1DImpl<2> { + fn do_dct<const COLUMNS: usize>(data: &mut [[f32; COLUMNS]]) { + for i in 0..COLUMNS { + let temp0 = data[0][i]; + let temp1 = data[1][i]; + data[0][i] = temp0 + temp1; + data[1][i] = temp0 - temp1; + } + } +} + +impl IDCT1D for IDCT1DImpl<2> { + fn do_idct<const COLUMNS: usize>(data: &mut [[f32; COLUMNS]]) { + for i in 0..COLUMNS { + let temp0 = data[0][i]; + let temp1 = data[1][i]; + data[0][i] = temp0 + temp1; + data[1][i] = temp0 - temp1; + } + } +} + +macro_rules! define_dct_1d { + ($n:literal, $nhalf: literal) => { + // Helper functions for CoeffBundle operating on $nhalf rows + impl<const SZ: usize> CoeffBundle<$nhalf, SZ> { + /// Adds a_in1[i] and a_in2[$nhalf - 1 - i], storing in a_out[i]. + fn add_reverse(a_in1: &[[f32; SZ]], a_in2: &[[f32; SZ]], a_out: &mut [[f32; SZ]]) { + const N_HALF_CONST: usize = $nhalf; + for i in 0..N_HALF_CONST { + for j in 0..SZ { + a_out[i][j] = a_in1[i][j] + a_in2[N_HALF_CONST - 1 - i][j]; + } + } + } + + /// Subtracts a_in2[$nhalf - 1 - i] from a_in1[i], storing in a_out[i]. + fn sub_reverse(a_in1: &[[f32; SZ]], a_in2: &[[f32; SZ]], a_out: &mut [[f32; SZ]]) { + const N_HALF_CONST: usize = $nhalf; + for i in 0..N_HALF_CONST { + for j in 0..SZ { + a_out[i][j] = a_in1[i][j] - a_in2[N_HALF_CONST - 1 - i][j]; + } + } + } + + /// Applies the B transform (forward DCT step). + /// Operates on a slice of $nhalf rows. + fn b(coeff: &mut [[f32; SZ]]) { + const N_HALF_CONST: usize = $nhalf; + for j in 0..SZ { + coeff[0][j] = coeff[0][j] * (SQRT_2 as f32) + coeff[1][j]; + } + // empty in the case N_HALF_CONST == 2 + #[allow(clippy::reversed_empty_ranges)] + for i in 1..(N_HALF_CONST - 1) { + for j in 0..SZ { + coeff[i][j] += coeff[i + 1][j]; + } + } + } + } + + // Helper functions for CoeffBundle operating on $n rows + impl<const SZ: usize> CoeffBundle<$n, SZ> { + /// Multiplies the second half of `coeff` by WcMultipliers. + fn multiply(coeff: &mut [[f32; SZ]]) { + const N_CONST: usize = $n; + const N_HALF_CONST: usize = $nhalf; + for i in 0..N_HALF_CONST { + let mul_val = WcMultipliers::<N_CONST>::K_MULTIPLIERS[i]; + for j in 0..SZ { + coeff[N_HALF_CONST + i][j] *= mul_val; + } + } + } + + /// De-interleaves `a_in` into `a_out`. + /// Even indexed rows of `a_out` get first half of `a_in`. + /// Odd indexed rows of `a_out` get second half of `a_in`. + fn inverse_even_odd(a_in: &[[f32; SZ]], a_out: &mut [[f32; SZ]]) { + const N_HALF_CONST: usize = $nhalf; + for i in 0..N_HALF_CONST { + for j in 0..SZ { + a_out[2 * i][j] = a_in[i][j]; + } + } + for i in 0..N_HALF_CONST { + for j in 0..SZ { + a_out[2 * i + 1][j] = a_in[N_HALF_CONST + i][j]; + } + } + } + } + + impl DCT1D for DCT1DImpl<$n> { + fn do_dct<const COLUMNS: usize>(data: &mut [[f32; COLUMNS]]) { + const { assert!($nhalf * 2 == $n, "N/2 * 2 must be N") } + assert!( + data.len() == $n, + "Input data must have $n rows for DCT1DImpl<$n>" + ); + + let mut tmp_buffer = [[0.0f32; COLUMNS]; $n]; + + // 1. AddReverse + // + // Inputs: first N/2 rows of data, second N/2 rows of data + // Output: first N/2 rows of tmp_buffer + CoeffBundle::<$nhalf, COLUMNS>::add_reverse( + &data[0..$nhalf], + &data[$nhalf..$n], + &mut tmp_buffer[0..$nhalf], + ); + + // 2. First Recursive Call (do_dct) + // first half + DCT1DImpl::<$nhalf>::do_dct::<COLUMNS>(&mut tmp_buffer[0..$nhalf]); + + // 3. SubReverse + // Inputs: first N/2 rows of data, second N/2 rows of data + // Output: second N/2 rows of tmp_buffer + CoeffBundle::<$nhalf, COLUMNS>::sub_reverse( + &data[0..$nhalf], + &data[$nhalf..$n], + &mut tmp_buffer[$nhalf..$n], + ); + + // 4. Multiply(tmp); + // Operates on the entire tmp_buffer. + CoeffBundle::<$n, COLUMNS>::multiply(&mut tmp_buffer); + + // 5. Second Recursive Call (do_dct) + // second half. + DCT1DImpl::<$nhalf>::do_dct::<COLUMNS>(&mut tmp_buffer[$nhalf..$n]); + + // 6. B + // Operates on the second N/2 rows of tmp_buffer. + CoeffBundle::<$nhalf, COLUMNS>::b(&mut tmp_buffer[$nhalf..$n]); + + // 7. InverseEvenOdd + CoeffBundle::<$n, COLUMNS>::inverse_even_odd(&tmp_buffer, data); + } + } + }; +} +define_dct_1d!(4, 2); +define_dct_1d!(8, 4); +define_dct_1d!(16, 8); +define_dct_1d!(32, 16); +define_dct_1d!(64, 32); +define_dct_1d!(128, 64); +define_dct_1d!(256, 128); + +macro_rules! define_idct_1d { + ($n:literal, $nhalf: literal) => { + impl<const SZ: usize> CoeffBundle<$nhalf, SZ> { + fn b_transpose(coeff: &mut [[f32; SZ]]) { + for i in (1..$nhalf).rev() { + for j in 0..SZ { + coeff[i][j] += coeff[i - 1][j]; + } + } + for j in 0..SZ { + coeff[0][j] *= SQRT_2 as f32; + } + } + } + + impl<const SZ: usize> CoeffBundle<$n, SZ> { + fn forward_even_odd(a_in: &[[f32; SZ]], a_out: &mut [[f32; SZ]]) { + for i in 0..($nhalf) { + for j in 0..SZ { + a_out[i][j] = a_in[2 * i][j]; + } + } + for i in ($nhalf)..$n { + for j in 0..SZ { + a_out[i][j] = a_in[2 * (i - $nhalf) + 1][j]; + } + } + } + fn multiply_and_add(coeff: &[[f32; SZ]], out: &mut [[f32; SZ]]) { + for i in 0..($nhalf) { + for j in 0..SZ { + let mul = WcMultipliers::<$n>::K_MULTIPLIERS[i]; + let in1 = coeff[i][j]; + let in2 = coeff[$nhalf + i][j]; + out[i][j] = mul * in2 + in1; + out[($n - i - 1)][j] = -mul * in2 + in1; + } + } + } + } + + impl IDCT1D for IDCT1DImpl<$n> { + fn do_idct<const COLUMNS: usize>(data: &mut [[f32; COLUMNS]]) { + const { assert!($nhalf * 2 == $n, "N/2 * 2 must be N") } + + // We assume `data` is arranged as a nxCOLUMNS matrix. + + let mut tmp = [[0.0f32; COLUMNS]; $n]; + + // 1. ForwardEvenOdd + CoeffBundle::<$n, COLUMNS>::forward_even_odd(data, &mut tmp); + // 2. First Recursive Call (IDCT1DImpl::do_idct) + // first half + IDCT1DImpl::<$nhalf>::do_idct::<COLUMNS>(&mut tmp[0..$nhalf]); + // 3. BTranspose. + // only the second half + CoeffBundle::<$nhalf, COLUMNS>::b_transpose(&mut tmp[$nhalf..$n]); + // 4. Second Recursive Call (IDCT1DImpl::do_idct) + // second half + IDCT1DImpl::<$nhalf>::do_idct::<COLUMNS>(&mut tmp[$nhalf..$n]); + // 5. MultiplyAndAdd. + CoeffBundle::<$n, COLUMNS>::multiply_and_add(&tmp, data); + } + } + }; +} +define_idct_1d!(4, 2); +define_idct_1d!(8, 4); +define_idct_1d!(16, 8); +define_idct_1d!(32, 16); +define_idct_1d!(64, 32); +define_idct_1d!(128, 64); +define_idct_1d!(256, 128); + +fn transpose<const ROWS: usize, const COLS: usize>(input: &[f32], output: &mut [f32]) { + assert_eq!(input.len(), ROWS * COLS); + assert_eq!(output.len(), ROWS * COLS); + + for r in 0..ROWS { + for c in 0..COLS { + let input_idx = r * COLS + c; + let output_idx = c * ROWS + r; + output[output_idx] = input[input_idx]; + } + } +} + +pub fn dct2d<const ROWS: usize, const COLS: usize>(data: &mut [f32]) +where + DCT1DImpl<ROWS>: DCT1D, + DCT1DImpl<COLS>: DCT1D, +{ + assert_eq!(data.len(), ROWS * COLS, "Data length mismatch"); + + // Copy data from flat slice `data` into a temporary Vec of arrays (rows). + let mut temp_rows: Vec<[f32; COLS]> = vec![[0.0f32; COLS]; ROWS]; + for (r, column) in temp_rows.iter_mut().enumerate() { + let start = r * COLS; + let end = start + COLS; + column.copy_from_slice(&data[start..end]); + } + + DCT1DImpl::<ROWS>::do_dct::<COLS>(&mut temp_rows); + + for (r, column) in temp_rows.iter().enumerate() { + let start = r * COLS; + let end = start + COLS; + data[start..end].copy_from_slice(column); + } + + // Create a temporary flat buffer for the transposed data. + let mut transposed_data = vec![0.0f32; ROWS * COLS]; + transpose::<ROWS, COLS>(data, &mut transposed_data); + + // Copy data from flat `transposed_data` into a temporary Vec of arrays. + let mut temp_cols: Vec<[f32; ROWS]> = vec![[0.0f32; ROWS]; COLS]; + for (c, row) in temp_cols.iter_mut().enumerate() { + let start = c * ROWS; + let end = start + ROWS; + row.copy_from_slice(&transposed_data[start..end]); + } + + // Perform DCT on the temporary structure (treating original columns as rows). + DCT1DImpl::<COLS>::do_dct::<ROWS>(&mut temp_cols); + + // Copy results back from the temporary structure into the flat `transposed_data`. + for (c, row) in temp_cols.iter().enumerate() { + let start = c * ROWS; + let end = start + ROWS; + transposed_data[start..end].copy_from_slice(row); + } + transpose::<COLS, ROWS>(&transposed_data, data); +} + +pub fn idct2d<const ROWS: usize, const COLS: usize>(data: &mut [f32]) +where + IDCT1DImpl<ROWS>: IDCT1D, + IDCT1DImpl<COLS>: IDCT1D, +{ + assert_eq!(data.len(), ROWS * COLS, "Data length mismatch"); + + // Create a temporary flat buffer for the transposed data. + let mut transposed_data = vec![0.0f32; ROWS * COLS]; + if ROWS < COLS { + transpose::<ROWS, COLS>(data, &mut transposed_data); + } else { + transposed_data.copy_from_slice(data); + } + + // Copy data from flat `transposed_data` into a temporary Vec of arrays. + let mut temp_cols: Vec<[f32; ROWS]> = vec![[0.0f32; ROWS]; COLS]; + for (c, row) in temp_cols.iter_mut().enumerate() { + let start = c * ROWS; + let end = start + ROWS; + row.copy_from_slice(&transposed_data[start..end]); + } + + // Perform IDCT on the temporary structure (treating original columns as rows). + IDCT1DImpl::<COLS>::do_idct::<ROWS>(&mut temp_cols); + + // Copy results back from the temporary structure into the flat `transposed_data`. + for (c, row) in temp_cols.iter().enumerate() { + let start = c * ROWS; + let end = start + ROWS; + transposed_data[start..end].copy_from_slice(row); + } + + transpose::<COLS, ROWS>(&transposed_data, data); + + // Copy data from flat slice `data` into a temporary Vec of arrays (rows). + let mut temp_rows: Vec<[f32; COLS]> = vec![[0.0f32; COLS]; ROWS]; + for (r, column) in temp_rows.iter_mut().enumerate() { + let start = r * COLS; + let end = start + COLS; + column.copy_from_slice(&data[start..end]); + } + IDCT1DImpl::<ROWS>::do_idct::<COLS>(&mut temp_rows); + + for (r, column) in temp_rows.iter().enumerate() { + let start = r * COLS; + let end = start + COLS; + data[start..end].copy_from_slice(column); + } +} + +pub fn compute_scaled_dct<const ROWS: usize, const COLS: usize>( + mut from: [[f32; COLS]; ROWS], + to: &mut [f32], +) where + DCT1DImpl<ROWS>: DCT1D, + DCT1DImpl<COLS>: DCT1D, +{ + DCT1DImpl::<ROWS>::do_dct::<COLS>(&mut from); + let mut transposed_dct_buffer = [[0.0; ROWS]; COLS]; + #[allow(clippy::needless_range_loop)] + for y in 0..ROWS { + for x in 0..COLS { + transposed_dct_buffer[x][y] = from[y][x]; + } + } + DCT1DImpl::<COLS>::do_dct::<ROWS>(&mut transposed_dct_buffer); + if ROWS < COLS { + for y in 0..ROWS { + for x in 0..COLS { + to[y * COLS + x] = transposed_dct_buffer[x][y] / (ROWS * COLS) as f32; + } + } + } else { + for y in 0..COLS { + for x in 0..ROWS { + to[y * ROWS + x] = transposed_dct_buffer[y][x] / (ROWS * COLS) as f32; + } + } + } +} + +#[cfg(test)] +mod tests { + use crate::{ + util::test::{assert_all_almost_abs_eq, assert_almost_abs_eq}, + var_dct::{ + dct::{DCT1D, DCT1DImpl, IDCT1D, IDCT1DImpl, compute_scaled_dct, dct2d, idct2d}, + dct_slow::{dct1d, idct1d}, + }, + }; + use test_log::test; + macro_rules! test_dct1d_eq_slow_n { + ($test_name:ident, $n_val:expr, $tolerance:expr) => { + #[test] + fn $test_name() { + const N: usize = $n_val; + const M: usize = 1; + const NM: usize = N * M; + + // Generate input data for the reference dct1d. + // Results in vec![vec![1.0], vec![2.0], ..., vec![N.0]] + let input_matrix_for_ref: Vec<Vec<f64>> = + std::array::from_fn::<f64, NM, _>(|i| (i + 1) as f64) + .chunks(M) + .map(|row_slice| row_slice.to_vec()) + .collect(); + + let output_matrix_slow: Vec<Vec<f64>> = dct1d(&input_matrix_for_ref); + + // DCT1DImpl expects data in [[f32; M]; N] format. + let mut input_arr_2d = [[0.0f32; M]; N]; + for r_idx in 0..N { + for c_idx in 0..M { + input_arr_2d[r_idx][c_idx] = input_matrix_for_ref[r_idx][c_idx] as f32; + } + } + + let mut output = input_arr_2d; + DCT1DImpl::<N>::do_dct::<M>(&mut output); + + for i in 0..N { + assert_almost_abs_eq(output[i][0], output_matrix_slow[i][0] as f32, $tolerance); + } + } + }; + } + + macro_rules! test_idct1d_eq_slow_n { + ($test_name:ident, $n_val:expr, $tolerance:expr) => { + #[test] + fn $test_name() { + const N: usize = $n_val; + const M: usize = 1; + const NM: usize = N * M; + + // Generate input data for the reference idct1d. + // Results in vec![vec![1.0], vec![2.0], ..., vec![N.0]] + let input_matrix_for_ref: Vec<Vec<f64>> = + std::array::from_fn::<f64, NM, _>(|i| (i + 1) as f64) + .chunks(M) + .map(|row_slice| row_slice.to_vec()) + .collect(); + + let output_matrix_slow: Vec<Vec<f64>> = idct1d(&input_matrix_for_ref); + + // IDCT1DImpl expects input coefficient data in [[f32; M]; N] format. + let mut input_arr_2d = [[0.0f32; M]; N]; + for r_idx in 0..N { + for c_idx in 0..M { + input_arr_2d[r_idx][c_idx] = input_matrix_for_ref[r_idx][c_idx] as f32; + } + } + + let mut output = input_arr_2d; + IDCT1DImpl::<N>::do_idct::<M>(&mut output); + + for i in 0..N { + assert_almost_abs_eq(output[i][0], output_matrix_slow[i][0] as f32, $tolerance); + } + } + }; + } + + test_dct1d_eq_slow_n!(test_dct1d_1x1_eq_slow, 1, 1e-6); + test_idct1d_eq_slow_n!(test_idct1d_1x1_eq_slow, 1, 1e-6); + test_dct1d_eq_slow_n!(test_dct1d_2x1_eq_slow, 2, 1e-6); + test_idct1d_eq_slow_n!(test_idct1d_2x1_eq_slow, 2, 1e-6); + test_dct1d_eq_slow_n!(test_dct1d_4x1_eq_slow, 4, 1e-6); + test_idct1d_eq_slow_n!(test_idct1d_4x1_eq_slow, 4, 1e-6); + test_dct1d_eq_slow_n!(test_dct1d_8x1_eq_slow, 8, 1e-5); + test_idct1d_eq_slow_n!(test_idct1d_8x1_eq_slow, 8, 1e-5); + test_dct1d_eq_slow_n!(test_dct1d_16x1_eq_slow, 16, 1e-4); + test_idct1d_eq_slow_n!(test_idct1d_16x1_eq_slow, 16, 1e-4); + test_dct1d_eq_slow_n!(test_dct1d_32x1_eq_slow, 32, 1e-3); + test_idct1d_eq_slow_n!(test_idct1d_32x1_eq_slow, 32, 1e-3); + test_dct1d_eq_slow_n!(test_dct1d_64x1_eq_slow, 64, 1e-2); + test_idct1d_eq_slow_n!(test_idct1d_64x1_eq_slow, 64, 1e-2); + test_dct1d_eq_slow_n!(test_dct1d_128x1_eq_slow, 128, 1e-2); + test_idct1d_eq_slow_n!(test_idct1d_128x1_eq_slow, 128, 1e-2); + test_dct1d_eq_slow_n!(test_dct1d_256x1_eq_slow, 256, 1e-1); + test_idct1d_eq_slow_n!(test_idct1d_256x1_eq_slow, 256, 1e-1); + + #[test] + fn test_idct1d_8x3_eq_slow() { + const N: usize = 8; + const M: usize = 3; + const NM: usize = N * M; // 24 + + // Initialize an N x M matrix with data from 1.0 to 24.0 + let input_coeffs_matrix_for_ref: Vec<Vec<f64>> = + std::array::from_fn::<f64, NM, _>(|i| (i + 1) as f64) + .chunks(M) + .map(|row_slice| row_slice.to_vec()) + .collect(); + + let output_matrix_slow: Vec<Vec<f64>> = idct1d(&input_coeffs_matrix_for_ref); + + // Prepare input for the implementation under test (IDCT1DImpl) + // IDCT1DImpl expects data in [[f32; M]; N] format. + let mut input_coeffs_for_fast_impl = [[0.0f32; M]; N]; + for r in 0..N { + for c in 0..M { + // Use the same source coefficient values as the reference IDCT + input_coeffs_for_fast_impl[r][c] = input_coeffs_matrix_for_ref[r][c] as f32; + } + } + + // This will be modified in-place by IDCT1DImpl + let mut output_fast_impl = input_coeffs_for_fast_impl; + + // Call the implementation under test (operates on 2D data) + IDCT1DImpl::<N>::do_idct::<M>(&mut output_fast_impl); + + // Compare results element-wise + for r_idx in 0..N { + for c_idx in 0..M { + assert_almost_abs_eq( + output_fast_impl[r_idx][c_idx], + output_matrix_slow[r_idx][c_idx] as f32, + 1e-5, + ); + } + } + } + + #[test] + fn test_dct1d_8x3_eq_slow() { + const N: usize = 8; + const M: usize = 3; + const NM: usize = N * M; // 24 + + // Initialize a 3 x 8 marix with data from 1.0 to 24.0 + let input_matrix_for_ref: Vec<Vec<f64>> = + std::array::from_fn::<f64, NM, _>(|i| (i + 1) as f64) + .chunks(M) + .map(|row_slice| row_slice.to_vec()) + .collect(); + + let output_matrix_slow: Vec<Vec<f64>> = dct1d(&input_matrix_for_ref); + + // Prepare input for the implementation under test (DCT1DImpl) + // DCT1DImpl expects data in [[f32; M]; N] format. + let mut input_for_fast_impl = [[0.0f32; M]; N]; + for r in 0..N { + for c in 0..M { + // Use the same source values as the reference DCT + input_for_fast_impl[r][c] = input_matrix_for_ref[r][c] as f32; + } + } + + // This will be modified in-place by DCT1DImpl + let mut output_fast_impl = input_for_fast_impl; + + // Call the implementation under test (operates on 2D data) + DCT1DImpl::<N>::do_dct::<M>(&mut output_fast_impl); + + // Compare results element-wise + for r_freq_idx in 0..N { + for c_col_idx in 0..M { + assert_almost_abs_eq( + output_fast_impl[r_freq_idx][c_col_idx], + output_matrix_slow[r_freq_idx][c_col_idx] as f32, + 1e-5, + ); + } + } + } + + // TODO(firsching): possibly change these tests to test against slow + // (i)dct method (after adding 2d-variant there) + macro_rules! test_idct2d_exists_n_m { + ($test_name:ident, $n_val:expr, $m_val:expr) => { + #[test] + fn $test_name() { + const N: usize = $n_val; + const M: usize = $m_val; + let mut data = [0.0f32; M * N]; + idct2d::<N, M>(&mut data); + } + }; + } + macro_rules! test_dct2d_exists_n_m { + ($test_name:ident, $n_val:expr, $m_val:expr) => { + #[test] + fn $test_name() { + const N: usize = $n_val; + const M: usize = $m_val; + let mut data = [0.0f32; M * N]; + dct2d::<N, M>(&mut data); + } + }; + } + test_dct2d_exists_n_m!(test_dct2d_exists_1_1, 1, 1); + test_idct2d_exists_n_m!(test_idct2d_exists_1_1, 1, 1); + test_dct2d_exists_n_m!(test_dct2d_exists_1_2, 1, 2); + test_idct2d_exists_n_m!(test_idct2d_exists_1_2, 1, 2); + test_dct2d_exists_n_m!(test_dct2d_exists_1_4, 1, 4); + test_idct2d_exists_n_m!(test_idct2d_exists_1_4, 1, 4); + test_dct2d_exists_n_m!(test_dct2d_exists_1_8, 1, 8); + test_idct2d_exists_n_m!(test_idct2d_exists_1_8, 1, 8); + test_dct2d_exists_n_m!(test_dct2d_exists_1_16, 1, 16); + test_idct2d_exists_n_m!(test_idct2d_exists_1_16, 1, 16); + test_dct2d_exists_n_m!(test_dct2d_exists_1_32, 1, 32); + test_idct2d_exists_n_m!(test_idct2d_exists_1_32, 1, 32); + test_dct2d_exists_n_m!(test_dct2d_exists_1_64, 1, 64); + test_idct2d_exists_n_m!(test_idct2d_exists_1_64, 1, 64); + test_dct2d_exists_n_m!(test_dct2d_exists_1_128, 1, 128); + test_idct2d_exists_n_m!(test_idct2d_exists_1_128, 1, 128); + test_dct2d_exists_n_m!(test_dct2d_exists_1_256, 1, 256); + test_idct2d_exists_n_m!(test_idct2d_exists_1_256, 1, 256); + test_dct2d_exists_n_m!(test_dct2d_exists_2_1, 2, 1); + test_idct2d_exists_n_m!(test_idct2d_exists_2_1, 2, 1); + test_dct2d_exists_n_m!(test_dct2d_exists_2_2, 2, 2); + test_idct2d_exists_n_m!(test_idct2d_exists_2_2, 2, 2); + test_dct2d_exists_n_m!(test_dct2d_exists_2_4, 2, 4); + test_idct2d_exists_n_m!(test_idct2d_exists_2_4, 2, 4); + test_dct2d_exists_n_m!(test_dct2d_exists_2_8, 2, 8); + test_idct2d_exists_n_m!(test_idct2d_exists_2_8, 2, 8); + test_dct2d_exists_n_m!(test_dct2d_exists_2_16, 2, 16); + test_idct2d_exists_n_m!(test_idct2d_exists_2_16, 2, 16); + test_dct2d_exists_n_m!(test_dct2d_exists_2_32, 2, 32); + test_idct2d_exists_n_m!(test_idct2d_exists_2_32, 2, 32); + test_dct2d_exists_n_m!(test_dct2d_exists_2_64, 2, 64); + test_idct2d_exists_n_m!(test_idct2d_exists_2_64, 2, 64); + test_dct2d_exists_n_m!(test_dct2d_exists_2_128, 2, 128); + test_idct2d_exists_n_m!(test_idct2d_exists_2_128, 2, 128); + test_dct2d_exists_n_m!(test_dct2d_exists_2_256, 2, 256); + test_idct2d_exists_n_m!(test_idct2d_exists_2_256, 2, 256); + test_dct2d_exists_n_m!(test_dct2d_exists_4_1, 4, 1); + test_idct2d_exists_n_m!(test_idct2d_exists_4_1, 4, 1); + test_dct2d_exists_n_m!(test_dct2d_exists_4_2, 4, 2); + test_idct2d_exists_n_m!(test_idct2d_exists_4_2, 4, 2); + test_dct2d_exists_n_m!(test_dct2d_exists_4_4, 4, 4); + test_idct2d_exists_n_m!(test_idct2d_exists_4_4, 4, 4); + test_dct2d_exists_n_m!(test_dct2d_exists_4_8, 4, 8); + test_idct2d_exists_n_m!(test_idct2d_exists_4_8, 4, 8); + test_dct2d_exists_n_m!(test_dct2d_exists_4_16, 4, 16); + test_idct2d_exists_n_m!(test_idct2d_exists_4_16, 4, 16); + test_dct2d_exists_n_m!(test_dct2d_exists_4_32, 4, 32); + test_idct2d_exists_n_m!(test_idct2d_exists_4_32, 4, 32); + test_dct2d_exists_n_m!(test_dct2d_exists_4_64, 4, 64); + test_idct2d_exists_n_m!(test_idct2d_exists_4_64, 4, 64); + test_dct2d_exists_n_m!(test_dct2d_exists_4_128, 4, 128); + test_idct2d_exists_n_m!(test_idct2d_exists_4_128, 4, 128); + test_dct2d_exists_n_m!(test_dct2d_exists_4_256, 4, 256); + test_idct2d_exists_n_m!(test_idct2d_exists_4_256, 4, 256); + test_dct2d_exists_n_m!(test_dct2d_exists_8_1, 8, 1); + test_idct2d_exists_n_m!(test_idct2d_exists_8_1, 8, 1); + test_dct2d_exists_n_m!(test_dct2d_exists_8_2, 8, 2); + test_idct2d_exists_n_m!(test_idct2d_exists_8_2, 8, 2); + test_dct2d_exists_n_m!(test_dct2d_exists_8_4, 8, 4); + test_idct2d_exists_n_m!(test_idct2d_exists_8_4, 8, 4); + test_dct2d_exists_n_m!(test_dct2d_exists_8_8, 8, 8); + test_idct2d_exists_n_m!(test_idct2d_exists_8_8, 8, 8); + test_dct2d_exists_n_m!(test_dct2d_exists_8_16, 8, 16); + test_idct2d_exists_n_m!(test_idct2d_exists_8_16, 8, 16); + test_dct2d_exists_n_m!(test_dct2d_exists_8_32, 8, 32); + test_idct2d_exists_n_m!(test_idct2d_exists_8_32, 8, 32); + test_dct2d_exists_n_m!(test_dct2d_exists_8_64, 8, 64); + test_idct2d_exists_n_m!(test_idct2d_exists_8_64, 8, 64); + test_dct2d_exists_n_m!(test_dct2d_exists_8_128, 8, 128); + test_idct2d_exists_n_m!(test_idct2d_exists_8_128, 8, 128); + test_dct2d_exists_n_m!(test_dct2d_exists_8_256, 8, 256); + test_idct2d_exists_n_m!(test_idct2d_exists_8_256, 8, 256); + test_dct2d_exists_n_m!(test_dct2d_exists_16_1, 16, 1); + test_idct2d_exists_n_m!(test_idct2d_exists_16_1, 16, 1); + test_dct2d_exists_n_m!(test_dct2d_exists_16_2, 16, 2); + test_idct2d_exists_n_m!(test_idct2d_exists_16_2, 16, 2); + test_dct2d_exists_n_m!(test_dct2d_exists_16_4, 16, 4); + test_idct2d_exists_n_m!(test_idct2d_exists_16_4, 16, 4); + test_dct2d_exists_n_m!(test_dct2d_exists_16_8, 16, 8); + test_idct2d_exists_n_m!(test_idct2d_exists_16_8, 16, 8); + test_dct2d_exists_n_m!(test_dct2d_exists_16_16, 16, 16); + test_idct2d_exists_n_m!(test_idct2d_exists_16_16, 16, 16); + test_dct2d_exists_n_m!(test_dct2d_exists_16_32, 16, 32); + test_idct2d_exists_n_m!(test_idct2d_exists_16_32, 16, 32); + test_dct2d_exists_n_m!(test_dct2d_exists_16_64, 16, 64); + test_idct2d_exists_n_m!(test_idct2d_exists_16_64, 16, 64); + test_dct2d_exists_n_m!(test_dct2d_exists_16_128, 16, 128); + test_idct2d_exists_n_m!(test_idct2d_exists_16_128, 16, 128); + test_dct2d_exists_n_m!(test_dct2d_exists_16_256, 16, 256); + test_idct2d_exists_n_m!(test_idct2d_exists_16_256, 16, 256); + test_dct2d_exists_n_m!(test_dct2d_exists_32_1, 32, 1); + test_idct2d_exists_n_m!(test_idct2d_exists_32_1, 32, 1); + test_dct2d_exists_n_m!(test_dct2d_exists_32_2, 32, 2); + test_idct2d_exists_n_m!(test_idct2d_exists_32_2, 32, 2); + test_dct2d_exists_n_m!(test_dct2d_exists_32_4, 32, 4); + test_idct2d_exists_n_m!(test_idct2d_exists_32_4, 32, 4); + test_dct2d_exists_n_m!(test_dct2d_exists_32_8, 32, 8); + test_idct2d_exists_n_m!(test_idct2d_exists_32_8, 32, 8); + test_dct2d_exists_n_m!(test_dct2d_exists_32_16, 32, 16); + test_idct2d_exists_n_m!(test_idct2d_exists_32_16, 32, 16); + test_dct2d_exists_n_m!(test_dct2d_exists_32_32, 32, 32); + test_idct2d_exists_n_m!(test_idct2d_exists_32_32, 32, 32); + test_dct2d_exists_n_m!(test_dct2d_exists_32_64, 32, 64); + test_idct2d_exists_n_m!(test_idct2d_exists_32_64, 32, 64); + test_dct2d_exists_n_m!(test_dct2d_exists_32_128, 32, 128); + test_idct2d_exists_n_m!(test_idct2d_exists_32_128, 32, 128); + test_dct2d_exists_n_m!(test_dct2d_exists_32_256, 32, 256); + test_idct2d_exists_n_m!(test_idct2d_exists_32_256, 32, 256); + test_dct2d_exists_n_m!(test_dct2d_exists_64_1, 64, 1); + test_idct2d_exists_n_m!(test_idct2d_exists_64_1, 64, 1); + test_dct2d_exists_n_m!(test_dct2d_exists_64_2, 64, 2); + test_idct2d_exists_n_m!(test_idct2d_exists_64_2, 64, 2); + test_dct2d_exists_n_m!(test_dct2d_exists_64_4, 64, 4); + test_idct2d_exists_n_m!(test_idct2d_exists_64_4, 64, 4); + test_dct2d_exists_n_m!(test_dct2d_exists_64_8, 64, 8); + test_idct2d_exists_n_m!(test_idct2d_exists_64_8, 64, 8); + test_dct2d_exists_n_m!(test_dct2d_exists_64_16, 64, 16); + test_idct2d_exists_n_m!(test_idct2d_exists_64_16, 64, 16); + test_dct2d_exists_n_m!(test_dct2d_exists_64_32, 64, 32); + test_idct2d_exists_n_m!(test_idct2d_exists_64_32, 64, 32); + test_dct2d_exists_n_m!(test_dct2d_exists_64_64, 64, 64); + test_idct2d_exists_n_m!(test_idct2d_exists_64_64, 64, 64); + test_dct2d_exists_n_m!(test_dct2d_exists_64_128, 64, 128); + test_idct2d_exists_n_m!(test_idct2d_exists_64_128, 64, 128); + test_dct2d_exists_n_m!(test_dct2d_exists_64_256, 64, 256); + test_idct2d_exists_n_m!(test_idct2d_exists_64_256, 64, 256); + test_dct2d_exists_n_m!(test_dct2d_exists_128_1, 128, 1); + test_idct2d_exists_n_m!(test_idct2d_exists_128_1, 128, 1); + test_dct2d_exists_n_m!(test_dct2d_exists_128_2, 128, 2); + test_idct2d_exists_n_m!(test_idct2d_exists_128_2, 128, 2); + test_dct2d_exists_n_m!(test_dct2d_exists_128_4, 128, 4); + test_idct2d_exists_n_m!(test_idct2d_exists_128_4, 128, 4); + test_dct2d_exists_n_m!(test_dct2d_exists_128_8, 128, 8); + test_idct2d_exists_n_m!(test_idct2d_exists_128_8, 128, 8); + test_dct2d_exists_n_m!(test_dct2d_exists_128_16, 128, 16); + test_idct2d_exists_n_m!(test_idct2d_exists_128_16, 128, 16); + test_dct2d_exists_n_m!(test_dct2d_exists_128_32, 128, 32); + test_idct2d_exists_n_m!(test_idct2d_exists_128_32, 128, 32); + test_dct2d_exists_n_m!(test_dct2d_exists_128_64, 128, 64); + test_idct2d_exists_n_m!(test_idct2d_exists_128_64, 128, 64); + test_dct2d_exists_n_m!(test_dct2d_exists_128_128, 128, 128); + test_idct2d_exists_n_m!(test_idct2d_exists_128_128, 128, 128); + test_dct2d_exists_n_m!(test_dct2d_exists_128_256, 128, 256); + test_idct2d_exists_n_m!(test_idct2d_exists_128_256, 128, 256); + test_dct2d_exists_n_m!(test_dct2d_exists_256_1, 256, 1); + test_idct2d_exists_n_m!(test_idct2d_exists_256_1, 256, 1); + test_dct2d_exists_n_m!(test_dct2d_exists_256_2, 256, 2); + test_idct2d_exists_n_m!(test_idct2d_exists_256_2, 256, 2); + test_dct2d_exists_n_m!(test_dct2d_exists_256_4, 256, 4); + test_idct2d_exists_n_m!(test_idct2d_exists_256_4, 256, 4); + test_dct2d_exists_n_m!(test_dct2d_exists_256_8, 256, 8); + test_idct2d_exists_n_m!(test_idct2d_exists_256_8, 256, 8); + test_dct2d_exists_n_m!(test_dct2d_exists_256_16, 256, 16); + test_idct2d_exists_n_m!(test_idct2d_exists_256_16, 256, 16); + test_dct2d_exists_n_m!(test_dct2d_exists_256_32, 256, 32); + test_idct2d_exists_n_m!(test_idct2d_exists_256_32, 256, 32); + test_dct2d_exists_n_m!(test_dct2d_exists_256_64, 256, 64); + test_idct2d_exists_n_m!(test_idct2d_exists_256_64, 256, 64); + test_dct2d_exists_n_m!(test_dct2d_exists_256_128, 256, 128); + test_idct2d_exists_n_m!(test_idct2d_exists_256_128, 256, 128); + test_dct2d_exists_n_m!(test_dct2d_exists_256_256, 256, 256); + test_idct2d_exists_n_m!(test_idct2d_exists_256_256, 256, 256); + + #[test] + fn test_compute_scaled_dct_wide() { + let input = [ + [86.0, 239.0, 213.0, 36.0, 34.0, 142.0, 248.0, 87.0], + [128.0, 122.0, 131.0, 72.0, 156.0, 112.0, 248.0, 55.0], + [120.0, 31.0, 246.0, 177.0, 119.0, 154.0, 176.0, 248.0], + [21.0, 151.0, 107.0, 101.0, 202.0, 71.0, 246.0, 48.0], + ]; + + let mut output = [0.0; 4 * 8]; + + compute_scaled_dct::<4, 8>(input, &mut output); + + assert_all_almost_abs_eq( + output, + [ + 135.219, -13.1026, 0.573698, -6.19682, -29.5938, 11.5028, -13.3955, 21.9205, + 1.4572, 11.3448, 16.3991, 2.50104, -20.549, 0.363681, 3.94596, -4.05406, -8.21875, + 6.57931, 0.601308, 1.51804, -20.5312, -9.29264, -19.6983, -0.850355, 12.4189, + -5.0881, 5.82096, -20.1997, 3.87769, 2.80762, 24.6634, -8.93341, + ], + 1e-3, + ); + } + + #[test] + fn test_compute_scaled_dct_tall() { + let input = [ + [86.0, 239.0, 213.0, 36.0], + [34.0, 142.0, 248.0, 87.0], + [128.0, 122.0, 131.0, 72.0], + [156.0, 112.0, 248.0, 55.0], + [120.0, 31.0, 246.0, 177.0], + [119.0, 154.0, 176.0, 248.0], + [21.0, 151.0, 107.0, 101.0], + [202.0, 71.0, 246.0, 48.0], + ]; + + let mut output = [0.0; 8 * 4]; + + compute_scaled_dct::<8, 4>(input, &mut output); + + assert_all_almost_abs_eq( + output, + [ + 135.219, -0.899633, -4.54363, 9.7776, 7.65625, -7.7203, 10.5073, -11.9921, + -8.31418, 5.39457, 11.3896, -17.5006, 11.6535, 12.6257, 9.27026, -0.767252, + -29.5938, -19.9538, -17.5214, -0.467021, -3.28125, -7.67861, 11.3504, 5.01615, + 24.9226, -4.19572, -7.10474, -16.7029, 24.2961, -16.8923, -3.32708, -4.09777, + ], + 1e-3, + ); + } +} diff --git a/third_party/rust/jxl/src/var_dct/dct_scales.rs b/third_party/rust/jxl/src/var_dct/dct_scales.rs @@ -0,0 +1,396 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. +#![allow(clippy::excessive_precision)] + +pub trait HasDctResampleScales<const N: usize> { + const SCALES: [f32; N]; +} + +pub struct DctResampleScales<const FROM: usize, const TO: usize>; + +impl HasDctResampleScales<1> for DctResampleScales<1, 8> { + const SCALES: [f32; 1] = [1.000000000000000000]; +} + +impl HasDctResampleScales<2> for DctResampleScales<2, 16> { + const SCALES: [f32; 2] = [1.000000000000000000, 1.108937353592731823]; +} + +impl HasDctResampleScales<4> for DctResampleScales<4, 32> { + const SCALES: [f32; 4] = [ + 1.000000000000000000, + 1.025760096781116015, + 1.108937353592731823, + 1.270559368765487251, + ]; +} + +impl HasDctResampleScales<8> for DctResampleScales<8, 64> { + const SCALES: [f32; 8] = [ + 1.0000000000000000, + 1.0063534990068217, + 1.0257600967811158, + 1.0593017296817173, + 1.1089373535927318, + 1.1777765381970435, + 1.2705593687654873, + 1.3944898413647777, + ]; +} + +impl HasDctResampleScales<16> for DctResampleScales<16, 128> { + const SCALES: [f32; 16] = [ + 1.0, + 1.0015830492062623, + 1.0063534990068217, + 1.0143759095928793, + 1.0257600967811158, + 1.0406645869480142, + 1.0593017296817173, + 1.0819447744633812, + 1.1089373535927318, + 1.1407059950032632, + 1.1777765381970435, + 1.2207956782315876, + 1.2705593687654873, + 1.3280505578213306, + 1.3944898413647777, + 1.4714043176061107, + ]; +} + +impl HasDctResampleScales<32> for DctResampleScales<32, 256> { + const SCALES: [f32; 32] = [ + 1.0, + 1.0003954307206069, + 1.0015830492062623, + 1.0035668445360069, + 1.0063534990068217, + 1.009952439375063, + 1.0143759095928793, + 1.0196390660647288, + 1.0257600967811158, + 1.0327603660498115, + 1.0406645869480142, + 1.049501024072585, + 1.0593017296817173, + 1.0701028169146336, + 1.0819447744633812, + 1.0948728278734026, + 1.1089373535927318, + 1.124194353004584, + 1.1407059950032632, + 1.158541237256391, + 1.1777765381970435, + 1.1984966740820495, + 1.2207956782315876, + 1.244777922949508, + 1.2705593687654873, + 1.2982690107339132, + 1.3280505578213306, + 1.3600643892400104, + 1.3944898413647777, + 1.4315278911623237, + 1.4714043176061107, + 1.5143734423314616, + ]; +} + +pub fn dct_total_resample_scale<const SIZE: usize, const DCT_SIZE: usize>(x: usize) -> f32 +where + DctResampleScales<SIZE, DCT_SIZE>: HasDctResampleScales<SIZE>, +{ + DctResampleScales::<SIZE, DCT_SIZE>::SCALES[x] +} + +pub struct WcMultipliers<const N: usize>; + +// Constants for DCT implementation. Generated by the following snippet: +// for i in range(N // 2): +// print(1.0 / (2 * math.cos((i + 0.5) * math.pi / N)), end=", ") +impl WcMultipliers<4> { + pub const K_MULTIPLIERS: [f32; 2] = [0.541196100146197, 1.3065629648763764]; +} +impl WcMultipliers<8> { + pub const K_MULTIPLIERS: [f32; 4] = [ + 0.5097955791041592, + 0.6013448869350453, + 0.8999762231364156, + 2.5629154477415055, + ]; +} + +impl WcMultipliers<16> { + pub const K_MULTIPLIERS: [f32; 8] = [ + 0.5024192861881557, + 0.5224986149396889, + 0.5669440348163577, + 0.6468217833599901, + 0.7881546234512502, + 1.060677685990347, + 1.7224470982383342, + 5.101148618689155, + ]; +} + +impl WcMultipliers<32> { + pub const K_MULTIPLIERS: [f32; 16] = [ + 0.5006029982351963, + 0.5054709598975436, + 0.5154473099226246, + 0.5310425910897841, + 0.5531038960344445, + 0.5829349682061339, + 0.6225041230356648, + 0.6748083414550057, + 0.7445362710022986, + 0.8393496454155268, + 0.9725682378619608, + 1.1694399334328847, + 1.4841646163141662, + 2.057781009953411, + 3.407608418468719, + 10.190008123548033, + ]; +} + +impl WcMultipliers<64> { + pub const K_MULTIPLIERS: [f32; 32] = [ + 0.500150636020651, + 0.5013584524464084, + 0.5037887256810443, + 0.5074711720725553, + 0.5124514794082247, + 0.5187927131053328, + 0.52657731515427, + 0.535909816907992, + 0.5469204379855088, + 0.5597698129470802, + 0.57465518403266, + 0.5918185358574165, + 0.6115573478825099, + 0.6342389366884031, + 0.6603198078137061, + 0.6903721282002123, + 0.7251205223771985, + 0.7654941649730891, + 0.8127020908144905, + 0.8683447152233481, + 0.9345835970364075, + 1.0144082649970547, + 1.1120716205797176, + 1.233832737976571, + 1.3892939586328277, + 1.5939722833856311, + 1.8746759800084078, + 2.282050068005162, + 2.924628428158216, + 4.084611078129248, + 6.796750711673633, + 20.373878167231453, + ]; +} + +impl WcMultipliers<128> { + pub const K_MULTIPLIERS: [f32; 64] = [ + 0.5000376519155477, + 0.5003390374428216, + 0.5009427176380873, + 0.5018505174842379, + 0.5030651913013697, + 0.5045904432216454, + 0.5064309549285542, + 0.5085924210498143, + 0.5110815927066812, + 0.5139063298475396, + 0.5170756631334912, + 0.5205998663018917, + 0.524490540114724, + 0.5287607092074876, + 0.5334249333971333, + 0.538499435291984, + 0.5440022463817783, + 0.549953374183236, + 0.5563749934898856, + 0.5632916653417023, + 0.5707305880121454, + 0.5787218851348208, + 0.5872989370937893, + 0.5964987630244563, + 0.606362462272146, + 0.6169357260050706, + 0.6282694319707711, + 0.6404203382416639, + 0.6534518953751283, + 0.6674352009263413, + 0.6824501259764195, + 0.6985866506472291, + 0.7159464549705746, + 0.7346448236478627, + 0.7548129391165311, + 0.776600658233963, + 0.8001798956216941, + 0.8257487738627852, + 0.8535367510066064, + 0.8838110045596234, + 0.9168844461846523, + 0.9531258743921193, + 0.9929729612675466, + 1.036949040910389, + 1.0856850642580145, + 1.1399486751015042, + 1.2006832557294167, + 1.2690611716991191, + 1.346557628206286, + 1.4350550884414341, + 1.5369941008524954, + 1.6555965242641195, + 1.7952052190778898, + 1.961817848571166, + 2.163957818751979, + 2.4141600002500763, + 2.7316450287739396, + 3.147462191781909, + 3.7152427383269746, + 4.5362909369693565, + 5.827688377844654, + 8.153848602466814, + 13.58429025728446, + 40.744688103351834, + ]; +} + +impl WcMultipliers<256> { + pub const K_MULTIPLIERS: [f32; 128] = [ + 0.5000094125358878, + 0.500084723455784, + 0.5002354020255269, + 0.5004615618093246, + 0.5007633734146156, + 0.5011410648064231, + 0.5015949217281668, + 0.502125288230386, + 0.5027325673091954, + 0.5034172216566842, + 0.5041797745258774, + 0.5050208107132756, + 0.5059409776624396, + 0.5069409866925212, + 0.5080216143561264, + 0.509183703931388, + 0.5104281670536573, + 0.5117559854927805, + 0.5131682130825206, + 0.5146659778093218, + 0.516250484068288, + 0.5179230150949777, + 0.5196849355823947, + 0.5215376944933958, + 0.5234828280796439, + 0.52552196311921, + 0.5276568203859896, + 0.5298892183652453, + 0.5322210772308335, + 0.5346544231010253, + 0.537191392591309, + 0.5398342376841637, + 0.5425853309375497, + 0.545447171055775, + 0.5484223888484947, + 0.551513753605893, + 0.554724179920619, + 0.5580567349898085, + 0.5615146464335654, + 0.5651013106696203, + 0.5688203018875696, + 0.5726753816701664, + 0.5766705093136241, + 0.5808098529038624, + 0.5850978012111273, + 0.58953897647151, + 0.5941382481306648, + 0.5989007476325463, + 0.6038318843443582, + 0.6089373627182432, + 0.614223200800649, + 0.6196957502119484, + 0.6253617177319102, + 0.6312281886412079, + 0.6373026519855411, + 0.6435930279473415, + 0.6501076975307724, + 0.6568555347890955, + 0.6638459418498757, + 0.6710888870233562, + 0.6785949463131795, + 0.6863753486870501, + 0.6944420255086364, + 0.7028076645818034, + 0.7114857693151208, + 0.7204907235796304, + 0.7298378629074134, + 0.7395435527641373, + 0.749625274727372, + 0.7601017215162176, + 0.7709929019493761, + 0.7823202570613161, + 0.7941067887834509, + 0.8063772028037925, + 0.8191580674598145, + 0.83247799080191, + 0.8463678182968619, + 0.860860854031955, + 0.8759931087426972, + 0.8918035785352535, + 0.9083345588266809, + 0.9256319988042384, + 0.9437459026371479, + 0.962730784794803, + 0.9826461881778968, + 1.0035572754078206, + 1.0255355056139732, + 1.048659411496106, + 1.0730154944316674, + 1.0986992590905857, + 1.1258164135986009, + 1.1544842669978943, + 1.184833362908442, + 1.217009397314603, + 1.2511754798461228, + 1.287514812536712, + 1.326233878832723, + 1.3675662599582539, + 1.411777227500661, + 1.459169302866857, + 1.5100890297227016, + 1.5649352798258847, + 1.6241695131835794, + 1.6883285509131505, + 1.7580406092704062, + 1.8340456094306077, + 1.9172211551275689, + 2.0086161135167564, + 2.1094945286246385, + 2.22139377701127, + 2.346202662531156, + 2.486267909203593, + 2.644541877144861, + 2.824791402350551, + 3.0318994541759925, + 3.2723115884254845, + 3.5547153325075804, + 3.891107790700307, + 4.298537526449054, + 4.802076008665048, + 5.440166215091329, + 6.274908408039339, + 7.413566756422303, + 9.058751453879703, + 11.644627325175037, + 16.300023088031555, + 27.163977662448232, + 81.48784219222516, + ]; +} diff --git a/third_party/rust/jxl/src/var_dct/dct_slow.rs b/third_party/rust/jxl/src/var_dct/dct_slow.rs @@ -0,0 +1,309 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +#![cfg(test)] + +use std::f64::consts::FRAC_1_SQRT_2; +use std::f64::consts::PI; +use std::f64::consts::SQRT_2; + +#[inline(always)] +fn alpha(u: usize) -> f64 { + if u == 0 { FRAC_1_SQRT_2 } else { 1.0 } +} + +pub fn dct1d(input_matrix: &[Vec<f64>]) -> Vec<Vec<f64>> { + let num_rows = input_matrix.len(); + + if num_rows == 0 { + return Vec::new(); + } + + let num_cols = input_matrix[0].len(); + + let mut output_matrix = vec![vec![0.0f64; num_cols]; num_rows]; + + let scale: f64 = SQRT_2; + + // Precompute the DCT matrix (size: n_rows x n_rows) + let mut dct_coeff_matrix = vec![vec![0.0f64; num_rows]; num_rows]; + for (u_freq, row) in dct_coeff_matrix.iter_mut().enumerate() { + let alpha_u_val = alpha(u_freq); + for (y_spatial, coeff) in row.iter_mut().enumerate() { + *coeff = alpha_u_val + * ((y_spatial as f64 + 0.5) * u_freq as f64 * PI / num_rows as f64).cos() + * scale; + } + } + + // Perform the DCT calculation column by column + for x_col_idx in 0..num_cols { + for u_freq_idx in 0..num_rows { + let mut sum = 0.0; + for (y_spatial_idx, col) in input_matrix.iter().enumerate() { + // This access `input_matrix[y_spatial_idx][x_col_idx]` assumes the input_matrix + // is rectangular. If not, it might panic here. + sum += dct_coeff_matrix[u_freq_idx][y_spatial_idx] * col[x_col_idx]; + } + output_matrix[u_freq_idx][x_col_idx] = sum; + } + } + + output_matrix +} + +pub fn idct1d(input_matrix: &[Vec<f64>]) -> Vec<Vec<f64>> { + let num_rows = input_matrix.len(); + + if num_rows == 0 { + return Vec::new(); + } + + let num_cols = input_matrix[0].len(); + + let mut output_matrix = vec![vec![0.0f64; num_cols]; num_rows]; + + let scale: f64 = SQRT_2; + + // Precompute the DCT matrix (size: num_rows x num_rows) + let mut dct_coeff_matrix = vec![vec![0.0f64; num_rows]; num_rows]; + for (u_freq, row) in dct_coeff_matrix.iter_mut().enumerate() { + let alpha_u_val = alpha(u_freq); + for (y_def_idx, coeff) in row.iter_mut().enumerate() { + *coeff = alpha_u_val + * ((y_def_idx as f64 + 0.5) * u_freq as f64 * PI / num_rows as f64).cos() + * scale; + } + } + + // Perform the IDCT calculation column by column + for x_col_idx in 0..num_cols { + for (y_row_idx, row) in output_matrix.iter_mut().enumerate() { + let mut sum = 0.0; + for (u_freq_idx, col) in input_matrix.iter().enumerate() { + // This access input_coeffs_matrix[u_freq_idx][x_col_idx] assumes input_coeffs_matrix + // is rectangular. If not, it might panic here. + sum += dct_coeff_matrix[u_freq_idx][y_row_idx] * col[x_col_idx]; + } + row[x_col_idx] = sum; + } + } + + output_matrix +} + +#[cfg(test)] +mod tests { + use test_log::test; + + use crate::util::test::assert_all_almost_abs_eq; + + use super::*; + + #[test] + fn test_slow_dct1d() { + const N_ROWS: usize = 8; + const M_COLS: usize = 1; + + let flat_input_data: [f64; N_ROWS] = [0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0]; + + // Prepare input_matrix for dct1d + // It expects Vec<Vec<f64>> structured as input_matrix[row_idx][col_idx] + // For N_ROWS=8, M_COLS=1, this means 8 rows, each containing a Vec with 1 element. + let input_matrix: Vec<Vec<f64>> = + flat_input_data.iter().map(|&value| vec![value]).collect(); + + // Call the refactored dct1d function which returns a new matrix + let output_matrix: Vec<Vec<f64>> = dct1d(&input_matrix); + + // Extract the first (and only) column from output_matrix for comparison + let mut result_column: Vec<f64> = Vec::with_capacity(N_ROWS); + if M_COLS > 0 { + for row in output_matrix.iter() { + result_column.push(row[0]); + } + } + + let expected = [ + 2.80000000e+01, + -1.82216412e+01, + -1.38622135e-15, + -1.90481783e+00, + 0.00000000e+00, + -5.68239222e-01, + -1.29520973e-15, + -1.43407825e-01, + ]; + // Ensure assert_all_almost_abs_eq can compare Vec<f64> (or slice) with [f64; N] + assert_all_almost_abs_eq(result_column.as_slice(), expected.as_slice(), 1e-7); + } + + #[test] + fn test_slow_dct1d_same_on_columns() { + const N_ROWS: usize = 8; + const M_COLS: usize = 5; + + // Prepare input_matrix for dct1d + // It expects Vec<Vec<f64>> structured as input_matrix[row_idx][col_idx]. + // Each column of the input should be [0.0, 1.0, ..., N_ROWS-1.0]. + let input_matrix: Vec<Vec<f64>> = (0..N_ROWS).map(|r| vec![r as f64; M_COLS]).collect(); + + // Call the refactored dct1d function which returns a new matrix + let output_matrix: Vec<Vec<f64>> = dct1d(&input_matrix); + + // Expected output for a single column [0.0 .. N_ROWS-1.0] + let single_column_dct_expected = [ + 2.80000000e+01, + -1.82216412e+01, + -1.38622135e-15, + -1.90481783e+00, + 0.00000000e+00, + -5.68239222e-01, + -1.29520973e-15, + -1.43407825e-01, + ]; + + for r_freq_idx in 0..N_ROWS { + let actual_row_slice: &[f64] = output_matrix[r_freq_idx].as_slice(); + let expected_row_values: Vec<f64> = + vec![single_column_dct_expected[r_freq_idx]; M_COLS]; + assert_all_almost_abs_eq(actual_row_slice, expected_row_values.as_slice(), 1e-7); + } + } + + #[test] + fn test_slow_idct1d() { + const N: usize = 8; + const M: usize = 1; + + let flat_input_data: [f64; N] = [0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0]; + + let input_coeffs_matrix_p1: Vec<Vec<f64>> = + flat_input_data.iter().map(|&value| vec![value]).collect(); + // Prepare input_matrix for dct1d + // It expects Vec<Vec<f64>> structured as input_matrix[row_idx][col_idx]. + // Each column of the input should be [0.0, 1.0, ..., N_ROWS-1.0].k + let output_matrix: Vec<Vec<f64>> = idct1d(&input_coeffs_matrix_p1); + + let mut result_column: Vec<f64> = Vec::with_capacity(N); + if M > 0 { + for row_vec in output_matrix.iter() { + result_column.push(row_vec[0]); + } + } + + let expected = [ + 20.63473963, + -22.84387206, + 8.99218712, + -7.77138893, + 4.05078387, + -3.47821595, + 1.32990088, + -0.91413457, + ]; + assert_all_almost_abs_eq(result_column.as_slice(), expected.as_slice(), 1e-7); + } + + #[test] + fn test_slow_idct1d_same_on_columns() { + const N_ROWS: usize = 8; + const M_COLS: usize = 5; + + // Prepare input_matrix for idct1d + // It expects Vec<Vec<f64>> structured as input_matrix[row_idx][col_idx]. + // Each column of the input should be [0.0, 1.0, ..., N_ROWS-1.0]. + let input_matrix: Vec<Vec<f64>> = (0..N_ROWS).map(|r| vec![r as f64; M_COLS]).collect(); + + // Call the refactored idct1d function which returns a new matrix + let output_matrix: Vec<Vec<f64>> = idct1d(&input_matrix); + + // Expected spatial output for a single input coefficient column [0.0 .. N_FREQUENCIES-1.0] + // This is taken from the single-column test `test_slow_idct1d` + let single_column_idct_expected = [ + 20.63473963, + -22.84387206, + 8.99218712, + -7.77138893, + 4.05078387, + -3.47821595, + 1.32990088, + -0.91413457, + ]; + + // Verify each row of output_spatial_matrix. + // The row output_spatial_matrix[r_spatial_idx] should consist of M_COLS elements, + // all equal to single_column_idct_expected[r_spatial_idx]. + for r_spatial_idx in 0..N_ROWS { + // Iterate over spatial output rows + let actual_row_slice: &[f64] = output_matrix[r_spatial_idx].as_slice(); + let expected_row_values: Vec<f64> = + vec![single_column_idct_expected[r_spatial_idx]; M_COLS]; + assert_all_almost_abs_eq(actual_row_slice, expected_row_values.as_slice(), 1e-7); + } + } + + #[test] + fn test_dct_idct_scaling() { + const N_ROWS: usize = 7; + const M_COLS: usize = 13; + let input_matrix: Vec<Vec<f64>> = (0..N_ROWS) + .map(|r_idx| { + (0..M_COLS) + // some arbitrary pattern + .map(|c_idx| (r_idx + c_idx) as f64 * 7.7) + .collect::<Vec<f64>>() + }) + .collect::<Vec<Vec<f64>>>(); + + let dct_output = dct1d(&input_matrix); + let idct_output = idct1d(&dct_output); + + // Verify that idct1d(dct1d(input)) == N_ROWS * input + for r_idx in 0..N_ROWS { + let expected_current_row_scaled: Vec<f64> = input_matrix[r_idx] + .iter() + .map(|&val| val * (N_ROWS as f64)) + .collect(); + + assert_all_almost_abs_eq( + idct_output[r_idx].as_slice(), + expected_current_row_scaled.as_slice(), + 1e-7, + ); + } + } + + #[test] + fn test_idct_dct_scaling() { + const N_ROWS: usize = 17; + const M_COLS: usize = 11; + let input_matrix: Vec<Vec<f64>> = (0..N_ROWS) + .map(|r_idx| { + (0..M_COLS) + // some arbitrary pattern + .map(|c_idx| (r_idx + c_idx) as f64 * 12.34) + .collect::<Vec<f64>>() + }) + .collect::<Vec<Vec<f64>>>(); + + let idct_output = idct1d(&input_matrix); + let dct_output = dct1d(&idct_output); + + // Verify that dct1d(idct1d(input)) == N_ROWS * input + for r_idx in 0..N_ROWS { + let expected_current_row_scaled: Vec<f64> = input_matrix[r_idx] + .iter() + .map(|&val| val * (N_ROWS as f64)) + .collect(); + + assert_all_almost_abs_eq( + dct_output[r_idx].as_slice(), + expected_current_row_scaled.as_slice(), + 1e-7, + ); + } + } +} diff --git a/third_party/rust/jxl/src/var_dct/transform.rs b/third_party/rust/jxl/src/var_dct/transform.rs @@ -0,0 +1,573 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +use crate::{ + BLOCK_DIM, + error::{Error, Result}, + frame::transform_map::*, + var_dct::dct::*, +}; + +fn idct2_top_block(s: usize, block_in: &[f32], block_out: &mut [f32]) { + let num_2x2 = s / 2; + for y in 0..num_2x2 { + for x in 0..num_2x2 { + let c00 = block_in[y * BLOCK_DIM + x]; + let c01 = block_in[y * BLOCK_DIM + num_2x2 + x]; + let c10 = block_in[(y + num_2x2) * BLOCK_DIM + x]; + let c11 = block_in[(y + num_2x2) * BLOCK_DIM + num_2x2 + x]; + let r00 = c00 + c01 + c10 + c11; + let r01 = c00 + c01 - c10 - c11; + let r10 = c00 - c01 + c10 - c11; + let r11 = c00 - c01 - c10 + c11; + block_out[y * 2 * BLOCK_DIM + x * 2] = r00; + block_out[y * 2 * BLOCK_DIM + x * 2 + 1] = r01; + block_out[(y * 2 + 1) * BLOCK_DIM + x * 2] = r10; + block_out[(y * 2 + 1) * BLOCK_DIM + x * 2 + 1] = r11; + } + } +} + +#[allow(clippy::excessive_precision)] +#[allow(clippy::approx_constant)] +fn avfidct4x4(coeffs: &[f32], pixels: &mut [f32]) { + let afv4x4basis: Vec<f32> = vec![ + 0.25, + 0.25, + 0.25, + 0.25, + 0.25, + 0.25, + 0.25, + 0.25, + 0.25, + 0.25, + 0.25, + 0.25, + 0.25, + 0.25, + 0.25, + 0.25, + 0.876902929799142, + 0.2206518106944235, + -0.10140050393753763, + -0.1014005039375375, + 0.2206518106944236, + -0.10140050393753777, + -0.10140050393753772, + -0.10140050393753763, + -0.10140050393753758, + -0.10140050393753769, + -0.1014005039375375, + -0.10140050393753768, + -0.10140050393753768, + -0.10140050393753759, + -0.10140050393753763, + -0.10140050393753741, + 0.0, + 0.0, + 0.40670075830260755, + 0.44444816619734445, + 0.0, + 0.0, + 0.19574399372042936, + 0.2929100136981264, + -0.40670075830260716, + -0.19574399372042872, + 0.0, + 0.11379074460448091, + -0.44444816619734384, + -0.29291001369812636, + -0.1137907446044814, + 0.0, + 0.0, + 0.0, + -0.21255748058288748, + 0.3085497062849767, + 0.0, + 0.4706702258572536, + -0.1621205195722993, + 0.0, + -0.21255748058287047, + -0.16212051957228327, + -0.47067022585725277, + -0.1464291867126764, + 0.3085497062849487, + 0.0, + -0.14642918671266536, + 0.4251149611657548, + 0.0, + -0.7071067811865474, + 0.0, + 0.0, + 0.7071067811865476, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + -0.4105377591765233, + 0.6235485373547691, + -0.06435071657946274, + -0.06435071657946266, + 0.6235485373547694, + -0.06435071657946284, + -0.0643507165794628, + -0.06435071657946274, + -0.06435071657946272, + -0.06435071657946279, + -0.06435071657946266, + -0.06435071657946277, + -0.06435071657946277, + -0.06435071657946273, + -0.06435071657946274, + -0.0643507165794626, + 0.0, + 0.0, + -0.4517556589999482, + 0.15854503551840063, + 0.0, + -0.04038515160822202, + 0.0074182263792423875, + 0.39351034269210167, + -0.45175565899994635, + 0.007418226379244351, + 0.1107416575309343, + 0.08298163094882051, + 0.15854503551839705, + 0.3935103426921022, + 0.0829816309488214, + -0.45175565899994796, + 0.0, + 0.0, + -0.304684750724869, + 0.5112616136591823, + 0.0, + 0.0, + -0.290480129728998, + -0.06578701549142804, + 0.304684750724884, + 0.2904801297290076, + 0.0, + -0.23889773523344604, + -0.5112616136592012, + 0.06578701549142545, + 0.23889773523345467, + 0.0, + 0.0, + 0.0, + 0.3017929516615495, + 0.25792362796341184, + 0.0, + 0.16272340142866204, + 0.09520022653475037, + 0.0, + 0.3017929516615503, + 0.09520022653475055, + -0.16272340142866173, + -0.35312385449816297, + 0.25792362796341295, + 0.0, + -0.3531238544981624, + -0.6035859033230976, + 0.0, + 0.0, + 0.40824829046386274, + 0.0, + 0.0, + 0.0, + 0.0, + -0.4082482904638628, + -0.4082482904638635, + 0.0, + 0.0, + -0.40824829046386296, + 0.0, + 0.4082482904638634, + 0.408248290463863, + 0.0, + 0.0, + 0.0, + 0.1747866975480809, + 0.0812611176717539, + 0.0, + 0.0, + -0.3675398009862027, + -0.307882213957909, + -0.17478669754808135, + 0.3675398009862011, + 0.0, + 0.4826689115059883, + -0.08126111767175039, + 0.30788221395790305, + -0.48266891150598584, + 0.0, + 0.0, + 0.0, + -0.21105601049335784, + 0.18567180916109802, + 0.0, + 0.0, + 0.49215859013738733, + -0.38525013709251915, + 0.21105601049335806, + -0.49215859013738905, + 0.0, + 0.17419412659916217, + -0.18567180916109904, + 0.3852501370925211, + -0.1741941265991621, + 0.0, + 0.0, + 0.0, + -0.14266084808807264, + -0.3416446842253372, + 0.0, + 0.7367497537172237, + 0.24627107722075148, + -0.08574019035519306, + -0.14266084808807344, + 0.24627107722075137, + 0.14883399227113567, + -0.04768680350229251, + -0.3416446842253373, + -0.08574019035519267, + -0.047686803502292804, + -0.14266084808807242, + 0.0, + 0.0, + -0.13813540350758585, + 0.3302282550303788, + 0.0, + 0.08755115000587084, + -0.07946706605909573, + -0.4613374887461511, + -0.13813540350758294, + -0.07946706605910261, + 0.49724647109535086, + 0.12538059448563663, + 0.3302282550303805, + -0.4613374887461554, + 0.12538059448564315, + -0.13813540350758452, + 0.0, + 0.0, + -0.17437602599651067, + 0.0702790691196284, + 0.0, + -0.2921026642334881, + 0.3623817333531167, + 0.0, + -0.1743760259965108, + 0.36238173335311646, + 0.29210266423348785, + -0.4326608024727445, + 0.07027906911962818, + 0.0, + -0.4326608024727457, + 0.34875205199302267, + 0.0, + 0.0, + 0.11354987314994337, + -0.07417504595810355, + 0.0, + 0.19402893032594343, + -0.435190496523228, + 0.21918684838857466, + 0.11354987314994257, + -0.4351904965232251, + 0.5550443808910661, + -0.25468277124066463, + -0.07417504595810233, + 0.2191868483885728, + -0.25468277124066413, + 0.1135498731499429, + ]; + for i in 0..16 { + let mut pixel: f32 = 0.0; + for j in 0..16 { + pixel += coeffs[j] * afv4x4basis[j * 16 + i]; + } + pixels[i] = pixel; + } +} + +fn afv_transform_to_pixels(afv_kind: usize, coefficients: &[f32], pixels: &mut [f32]) { + let afv_x = afv_kind & 1; + let afv_y = afv_kind / 2; + let block00 = coefficients[0]; + let block01 = coefficients[1]; + let block10 = coefficients[8]; + let dcs: [f32; 3] = [ + (block00 + block10 + block01) * 4.0, + block00 + block10 - block01, + block00 - block10, + ]; + // IAFV: (even, even) positions. + let mut coeff: Vec<f32> = vec![0.0; 4 * 4]; + for iy in 0..4 { + for ix in 0..4 { + coeff[iy * 4 + ix] = if ix == 0 && iy == 0 { + dcs[0] + } else { + coefficients[iy * 2 * 8 + ix * 2] + }; + } + } + let mut block: Vec<f32> = vec![0.0; 4 * 8]; + avfidct4x4(&coeff, &mut block); + for iy in 0..4 { + let block_y = if afv_y == 1 { 3 - iy } else { iy }; + for ix in 0..4 { + let block_x = if afv_x == 1 { 3 - ix } else { ix }; + pixels[(iy + afv_y * 4) * 8 + afv_x * 4 + ix] = block[block_y * 4 + block_x]; + } + } + // IDCT4x4 in (odd, even) positions. + for iy in 0..4 { + for ix in 0..4 { + block[iy * 4 + ix] = if ix == 0 && iy == 0 { + dcs[1] + } else { + coefficients[iy * 2 * 8 + ix * 2 + 1] + }; + } + } + idct2d::<4, 4>(&mut block[0..16]); + for iy in 0..4 { + for ix in 0..4 { + pixels[(iy + afv_y * 4) * 8 + (1 - afv_x) * 4 + ix] = block[iy * 4 + ix]; + } + } + // IDCT4x8. + for iy in 0..4 { + for ix in 0..8 { + block[iy * 8 + ix] = if ix == 0 && iy == 0 { + dcs[2] + } else { + coefficients[(1 + iy * 2) * 8 + ix] + }; + } + } + idct2d::<4, 8>(&mut block); + for iy in 0..4 { + for ix in 0..8 { + pixels[(iy + (1 - afv_y) * 4) * 8 + ix] = block[iy * 8 + ix]; + } + } +} + +pub fn transform_to_pixels( + transform_type: HfTransformType, + transform_buffer: &mut [f32], +) -> Result<(), Error> { + match transform_type { + HfTransformType::DCT => { + idct2d::<8, 8>(&mut transform_buffer[0..64]); + } + HfTransformType::DCT16X16 => { + idct2d::<16, 16>(&mut transform_buffer[0..256]); + } + HfTransformType::DCT32X32 => { + idct2d::<32, 32>(&mut transform_buffer[0..1024]); + } + HfTransformType::DCT16X8 => { + idct2d::<16, 8>(&mut transform_buffer[0..128]); + } + HfTransformType::DCT8X16 => { + idct2d::<8, 16>(&mut transform_buffer[0..128]); + } + HfTransformType::DCT32X8 => { + idct2d::<32, 8>(&mut transform_buffer[0..256]); + } + HfTransformType::DCT8X32 => { + idct2d::<8, 32>(&mut transform_buffer[0..256]); + } + HfTransformType::DCT32X16 => { + idct2d::<32, 16>(&mut transform_buffer[0..512]); + } + HfTransformType::DCT16X32 => { + idct2d::<16, 32>(&mut transform_buffer[0..512]); + } + HfTransformType::DCT64X64 => { + idct2d::<64, 64>(&mut transform_buffer[0..4096]); + } + HfTransformType::DCT64X32 => { + idct2d::<64, 32>(&mut transform_buffer[0..2048]); + } + HfTransformType::DCT32X64 => { + idct2d::<32, 64>(&mut transform_buffer[0..2048]); + } + HfTransformType::DCT128X128 => { + idct2d::<128, 128>(&mut transform_buffer[0..16384]); + } + HfTransformType::DCT128X64 => { + idct2d::<128, 64>(&mut transform_buffer[0..8192]); + } + HfTransformType::DCT64X128 => { + idct2d::<64, 128>(&mut transform_buffer[0..8192]); + } + HfTransformType::DCT256X256 => { + idct2d::<256, 256>(&mut transform_buffer[0..65536]); + } + HfTransformType::DCT256X128 => { + idct2d::<256, 128>(&mut transform_buffer[0..32768]); + } + HfTransformType::DCT128X256 => { + idct2d::<128, 256>(&mut transform_buffer[0..32768]); + } + HfTransformType::AFV0 => { + let block: Vec<f32> = transform_buffer[0..64].to_vec(); + afv_transform_to_pixels(0, &block, &mut transform_buffer[0..64]); + } + HfTransformType::AFV1 => { + let block: Vec<f32> = transform_buffer[0..64].to_vec(); + afv_transform_to_pixels(1, &block, &mut transform_buffer[0..64]); + } + HfTransformType::AFV2 => { + let block: Vec<f32> = transform_buffer[0..64].to_vec(); + afv_transform_to_pixels(2, &block, &mut transform_buffer[0..64]); + } + HfTransformType::AFV3 => { + let block: Vec<f32> = transform_buffer[0..64].to_vec(); + afv_transform_to_pixels(3, &block, &mut transform_buffer[0..64]); + } + HfTransformType::IDENTITY => { + let coefficients: Vec<f32> = transform_buffer[0..64].to_vec(); + let block00 = coefficients[0]; + let block01 = coefficients[1]; + let block10 = coefficients[8]; + let block11 = coefficients[9]; + let dcs: [f32; 4] = [ + block00 + block01 + block10 + block11, + block00 + block01 - block10 - block11, + block00 - block01 + block10 - block11, + block00 - block01 - block10 + block11, + ]; + for y in 0..2 { + for x in 0..2 { + let block_dc = dcs[y * 2 + x]; + let mut residual_sum = 0.0; + for iy in 0..4 { + for ix in 0..4 { + if ix == 0 && iy == 0 { + continue; + } + residual_sum += coefficients[(y + iy * 2) * 8 + x + ix * 2]; + } + } + transform_buffer[(4 * y + 1) * 8 + 4 * x + 1] = + block_dc - residual_sum * (1.0 / 16.0); + for iy in 0..4 { + for ix in 0..4 { + if ix == 1 && iy == 1 { + continue; + } + transform_buffer[(y * 4 + iy) * 8 + x * 4 + ix] = coefficients + [(y + iy * 2) * 8 + x + ix * 2] + + transform_buffer[(4 * y + 1) * 8 + 4 * x + 1]; + } + } + transform_buffer[y * 4 * 8 + x * 4] = coefficients[(y + 2) * 8 + x + 2] + + transform_buffer[(4 * y + 1) * 8 + 4 * x + 1]; + } + } + } + HfTransformType::DCT2X2 => { + let mut tmp: Vec<f32> = transform_buffer[0..64].to_vec(); + idct2_top_block(2, &tmp, &mut transform_buffer[0..64]); + idct2_top_block(4, &transform_buffer[0..64], &mut tmp); + idct2_top_block(8, &tmp, &mut transform_buffer[0..64]); + } + HfTransformType::DCT4X4 => { + let coefficients: Vec<f32> = transform_buffer[0..64].to_vec(); + let block00 = coefficients[0]; + let block01 = coefficients[1]; + let block10 = coefficients[8]; + let block11 = coefficients[9]; + let dcs: [f32; 4] = [ + block00 + block01 + block10 + block11, + block00 + block01 - block10 - block11, + block00 - block01 + block10 - block11, + block00 - block01 - block10 + block11, + ]; + for y in 0..2 { + for x in 0..2 { + let mut block: Vec<f32> = vec![0.0; 4 * 4]; + block[0] = dcs[y * 2 + x]; + for iy in 0..4 { + for ix in 0..4 { + if ix == 0 && iy == 0 { + continue; + } + block[iy * 4 + ix] = coefficients[(y + iy * 2) * 8 + x + ix * 2]; + } + } + idct2d::<4, 4>(&mut block); + for iy in 0..4 { + for ix in 0..4 { + transform_buffer[(y * 4 + iy) * 8 + x * 4 + ix] = block[iy * 4 + ix]; + } + } + } + } + } + HfTransformType::DCT8X4 => { + let coefficients: Vec<f32> = transform_buffer[0..64].to_vec(); + let block0 = coefficients[0]; + let block1 = coefficients[8]; + let dcs: [f32; 2] = [block0 + block1, block0 - block1]; + for x in 0..2 { + let mut block: Vec<f32> = vec![0.0; 8 * 4]; + for iy in 0..4 { + for ix in 0..8 { + block[iy * 8 + ix] = if ix == 0 && iy == 0 { + dcs[x] + } else { + coefficients[(x + iy * 2) * 8 + ix] + } + } + } + idct2d::<8, 4>(&mut block); + for iy in 0..8 { + for ix in 0..4 { + transform_buffer[iy * 8 + x * 4 + ix] = block[iy * 4 + ix]; + } + } + } + } + HfTransformType::DCT4X8 => { + let coefficients: Vec<f32> = transform_buffer[0..64].to_vec(); + let block0 = coefficients[0]; + let block1 = coefficients[8]; + let dcs: [f32; 2] = [block0 + block1, block0 - block1]; + for y in 0..2 { + let mut block: Vec<f32> = vec![0.0; 4 * 8]; + for iy in 0..4 { + for ix in 0..8 { + block[iy * 8 + ix] = if ix == 0 && iy == 0 { + dcs[y] + } else { + coefficients[(y + iy * 2) * 8 + ix] + } + } + } + idct2d::<4, 8>(&mut block); + for iy in 0..4 { + for ix in 0..8 { + transform_buffer[(y * 4 + iy) * 8 + ix] = block[iy * 8 + ix]; + } + } + } + } + }; + Ok(()) +} diff --git a/third_party/rust/jxl_macros/.cargo-checksum.json b/third_party/rust/jxl_macros/.cargo-checksum.json @@ -0,0 +1 @@ +{"files":{"Cargo.lock":"44fcf448601ff9cec27d61d146769830c2b315e03fe51c256eb1d011d3675e8e","Cargo.toml":"92593888805c19646e81cbccefaff73b0ab6d271ec16423358c4ed5e0712228c","README.md":"3bb0a5cac8cebc1b90a75a94f981ba895cc7afff28ce57b4d88697ba9ee0a92a","src/lib.rs":"17581ecec32de19ae2d4b6a62ed80070ec934735ea7e5368632a48940070f161"},"package":"5169e84285ee08cee04f115f2658cafba62e7ff54e8eb91a3842129ca12b003f"} +\ No newline at end of file diff --git a/third_party/rust/jxl_macros/Cargo.lock b/third_party/rust/jxl_macros/Cargo.lock @@ -0,0 +1,70 @@ +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. +version = 4 + +[[package]] +name = "jxl_macros" +version = "0.1.1" +dependencies = [ + "proc-macro-error2", + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "proc-macro-error-attr2" +version = "2.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "96de42df36bb9bba5542fe9f1a054b8cc87e172759a1868aa05c1f3acc89dfc5" +dependencies = [ + "proc-macro2", + "quote", +] + +[[package]] +name = "proc-macro-error2" +version = "2.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "11ec05c52be0a07b08061f7dd003e7d7092e0472bc731b4af7bb1ef876109802" +dependencies = [ + "proc-macro-error-attr2", + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "proc-macro2" +version = "1.0.92" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "37d3544b3f2748c54e147655edb5025752e2303145b5aefb3c3ea2c78b973bb0" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "quote" +version = "1.0.37" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b5b9d34b8991d19d98081b46eacdd8eb58c6f2b201139f7c5f643cc155a633af" +dependencies = [ + "proc-macro2", +] + +[[package]] +name = "syn" +version = "2.0.90" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "919d3b74a5dd0ccd15aeb8f93e7006bd9e14c295087c9896a110f490752bcf31" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "unicode-ident" +version = "1.0.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e91b56cd4cadaeb79bbf1a5645f6b4f8dc5bde8834ad5894a8db35fda9efa1fe" diff --git a/third_party/rust/jxl_macros/Cargo.toml b/third_party/rust/jxl_macros/Cargo.toml @@ -0,0 +1,63 @@ +# THIS FILE IS AUTOMATICALLY GENERATED BY CARGO +# +# When uploading crates to the registry Cargo will automatically +# "normalize" Cargo.toml files for maximal compatibility +# with all versions of Cargo and also rewrite `path` dependencies +# to registry (e.g., crates.io) dependencies. +# +# If you are reading this file be aware that the original Cargo.toml +# will likely look very different (and much more reasonable). +# See Cargo.toml.orig for the original contents. + +[package] +edition = "2024" +name = "jxl_macros" +version = "0.1.1" +authors = ["Luca Versari <veluca93@gmail.com>"] +build = false +autolib = false +autobins = false +autoexamples = false +autotests = false +autobenches = false +description = "High performance Rust implementation of a JPEG XL decoder - supporting macros" +readme = "README.md" +keywords = [ + "jpeg-xl", + "decoder", +] +categories = ["multimedia::images"] +license = "BSD-3-Clause" +repository = "https://github.com/libjxl/jxl-rs" +resolver = "2" + +[features] +test = [] + +[lib] +name = "jxl_macros" +path = "src/lib.rs" +proc-macro = true + +[dependencies.proc-macro-error2] +version = "2.0.1" + +[dependencies.proc-macro2] +version = "1.0" + +[dependencies.quote] +version = "1.0" + +[dependencies.syn] +version = "2.0.90" +features = [ + "extra-traits", + "full", +] + +[lints.clippy] +missing_safety_doc = "deny" +undocumented_unsafe_blocks = "deny" + +[lints.rust] +unsafe_op_in_unsafe_fn = "deny" diff --git a/third_party/rust/jxl_macros/README.md b/third_party/rust/jxl_macros/README.md @@ -0,0 +1,7 @@ +# JPEG XL in Rust + +This is a work-in-progress, and currently incomplete reimplementation of a JPEG XL decoder in Rust, which aims to be conforming, safe and fast. + +Refer to [the main libjxl repository](https://github.com/libjxl/libjxl) for +more information, including contributing instructions. + diff --git a/third_party/rust/jxl_macros/src/lib.rs b/third_party/rust/jxl_macros/src/lib.rs @@ -0,0 +1,753 @@ +// Copyright (c) the JPEG XL Project Authors. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +extern crate proc_macro; + +use proc_macro::TokenStream; +use proc_macro_error2::{abort, proc_macro_error}; +use proc_macro2::TokenStream as TokenStream2; +use quote::quote; +use syn::{DeriveInput, Meta, parse_macro_input}; + +fn get_bits(expr_call: &syn::ExprCall) -> syn::Expr { + if let syn::Expr::Path(ep) = &*expr_call.func { + if !ep.path.is_ident("Bits") { + abort!( + expr_call, + "Unexpected function name in coder: {}", + ep.path.get_ident().unwrap() + ); + } + if expr_call.args.len() != 1 { + abort!( + expr_call, + "Unexpected number of arguments for Bits() in coder: {}", + expr_call.args.len() + ); + } + return expr_call.args[0].clone(); + } + abort!(expr_call, "Unexpected function call in coder"); +} + +fn parse_single_coder(input: &syn::Expr, extra_lit: Option<&syn::ExprLit>) -> TokenStream2 { + match &input { + syn::Expr::Lit(lit) => match extra_lit { + None => quote! {U32::Val(#lit)}, + Some(elit) => quote! {U32::Val(#lit + #elit)}, + }, + syn::Expr::Call(expr_call) => { + let bits = get_bits(expr_call); + match extra_lit { + None => quote! {U32::Bits(#bits)}, + Some(elit) => quote! {U32::BitsOffset{n: #bits, off: #elit}}, + } + } + syn::Expr::Binary(syn::ExprBinary { + attrs: _, + left, + op: syn::BinOp::Add(_), + right, + }) => { + let (left, right) = if let syn::Expr::Lit(_) = **left { + (right, left) + } else { + (left, right) + }; + match (&**left, &**right) { + (syn::Expr::Call(expr_call), syn::Expr::Lit(lit)) => { + let bits = get_bits(expr_call); + match extra_lit { + None => quote! {U32::BitsOffset{n: #bits, off: #lit}}, + Some(elit) => quote! {U32::BitsOffset{n: #bits, off: #lit + #elit}}, + } + } + _ => abort!( + input, + "Unexpected expression in coder, must be Bits(a) + b, Bits(a), or b" + ), + } + } + _ => abort!( + input, + "Unexpected expression in coder, must be Bits(a) + b, Bits(a), or b" + ), + } +} + +fn parse_coder(input: &syn::Expr) -> TokenStream2 { + let parse_u2s = |expr_call: &syn::ExprCall, lit: Option<&syn::ExprLit>| { + if let syn::Expr::Path(ep) = &*expr_call.func { + if !ep.path.is_ident("u2S") { + let coder = parse_single_coder(input, None); + return quote! {U32Coder::Direct(#coder)}; + } + if expr_call.args.len() != 4 { + abort!( + input, + "Unexpected number of arguments for U32() in coder: {}", + expr_call.args.len() + ); + } + let args = vec![ + parse_single_coder(&expr_call.args[0], lit), + parse_single_coder(&expr_call.args[1], lit), + parse_single_coder(&expr_call.args[2], lit), + parse_single_coder(&expr_call.args[3], lit), + ]; + return quote! {U32Coder::Select(#(#args),*)}; + } + abort!(input, "Unexpected function call in coder"); + }; + + match &input { + syn::Expr::Call(expr_call) => parse_u2s(expr_call, None), + syn::Expr::Binary(syn::ExprBinary { + attrs: _, + left, + op: syn::BinOp::Add(_), + right, + }) => { + let (left, right) = if let syn::Expr::Lit(_) = **left { + (right, left) + } else { + (left, right) + }; + match (&**left, &**right) { + (syn::Expr::Call(expr_call), syn::Expr::Lit(lit)) => { + parse_u2s(expr_call, Some(lit)) + } + _ => abort!( + input, + "Unexpected expression in coder, must be (u2S|Bits)(a) + b, (u2S|Bits)(a), or b" + ), + } + } + _ => parse_single_coder(input, None), + } +} + +fn parse_size_coder(mut input: syn::Expr) -> TokenStream2 { + match input { + syn::Expr::Call(syn::ExprCall { + ref func, + ref mut args, + .. + }) => { + if args.len() != 1 { + abort!(input, "Expected 1 argument in sized_coder inner call"); + } + + match &**func { + syn::Expr::Path(expr_path) if expr_path.path.is_ident("implicit") => { + let arg = args.first().unwrap().clone(); + parse_coder(&arg) + } + syn::Expr::Path(expr_path) if expr_path.path.is_ident("explicit") => { + quote! { U32Coder::Direct(U32::Val(#args)) } + } + _ => abort!( + input, + "Unexpected expression in size_coder, must be 'implicit()' or 'explicit()'" + ), + } + } + _ => abort!( + input, + "Unexpected expression in size_coder, must be 'implicit()' or 'explicit()'" + ), + } +} + +fn prettify_condition(cond: &syn::Expr) -> String { + (quote! {#cond}) + .to_string() + .replace(" . ", ".") + .replace("! ", "!") + .replace(" :: ", "::") +} + +#[derive(Debug)] +struct Condition { + expr: Option<syn::Expr>, + has_all_default: bool, + pretty: String, +} + +impl Condition { + fn get_expr(&self, all_default_field: &Option<syn::Ident>) -> Option<TokenStream2> { + if self.has_all_default { + let all_default = all_default_field.as_ref().unwrap(); + match &self.expr { + Some(expr) => Some(quote! { !#all_default && (#expr) }), + None => Some(quote! { !#all_default }), + } + } else { + self.expr.as_ref().map(|expr| quote! {#expr}) + } + } + fn get_pretty(&self, all_default_field: &Option<syn::Ident>) -> String { + if self.has_all_default { + let all_default = all_default_field.as_ref().unwrap(); + let all_default = "!".to_owned() + "e! {#all_default}.to_string(); + match &self.expr { + Some(_) => all_default + " && (" + &self.pretty + ")", + None => all_default, + } + } else { + self.pretty.clone() + } + } +} + +#[derive(Debug, Clone)] +struct U32 { + coder: TokenStream2, +} + +#[derive(Debug)] +#[allow(clippy::large_enum_variant)] +enum Coder { + WithoutConfig, + U32(U32), + Select(Condition, U32, U32), + Vector(U32, Box<Coder>), +} + +impl Coder { + fn ty(&self) -> TokenStream2 { + match self { + Coder::WithoutConfig => quote! {()}, + Coder::U32(..) => quote! {U32Coder}, + Coder::Select(..) => quote! {SelectCoder<U32Coder>}, + Coder::Vector(_, value_coder) => { + let value_coder_ty = value_coder.ty(); + quote! {VectorCoder<#value_coder_ty>} + } + } + } + + fn config(&self, all_default_field: &Option<syn::Ident>) -> TokenStream2 { + match self { + Coder::WithoutConfig => quote! { () }, + Coder::U32(U32 { coder }) => quote! { #coder }, + Coder::Select(condition, U32 { coder: coder_true }, U32 { coder: coder_false }) => { + let cnd = condition.get_expr(all_default_field).unwrap(); + quote! { + SelectCoder{use_true: #cnd, coder_true: #coder_true, coder_false: #coder_false} + } + } + Coder::Vector(U32 { coder }, value_coder) => { + let value_coder = value_coder.config(all_default_field); + quote! {VectorCoder{size_coder: #coder, value_coder: #value_coder}} + } + } + } +} + +#[derive(Debug)] +enum FieldKind { + Unconditional(Coder), + Conditional(Condition, Coder), + Defaulted(Condition, Coder), +} + +#[derive(Debug)] +struct Field { + name: proc_macro2::Ident, + kind: FieldKind, + ty: syn::Type, + default: Option<TokenStream2>, + default_element: Option<TokenStream2>, + nonserialized_inits: Vec<TokenStream2>, +} + +impl Field { + fn parse(f: &syn::Field, num: usize, all_default_field: &mut Option<syn::Ident>) -> Field { + let mut condition = None; + let mut default = None; + let mut coder = None; + + let mut select_coder = None; + let mut coder_true = None; + let mut coder_false = None; + + let mut is_all_default = false; + + let mut size_coder = None; + + let mut nonserialized = vec![]; + + let mut default_element = None; + + // Parse attributes. + for a in &f.attrs { + match a.path().get_ident().map(syn::Ident::to_string).as_deref() { + Some("coder") => { + if coder.is_some() { + abort!(f, "Repeated coder"); + } + let coder_ast = a.parse_args::<syn::Expr>().unwrap(); + coder = Some(Coder::U32(U32 { + coder: parse_coder(&coder_ast), + })); + } + Some("default") => { + if default.is_some() { + abort!(f, "Repeated default"); + } + let default_expr = a.parse_args::<syn::Expr>().unwrap(); + default = Some(quote! {#default_expr}); + } + Some("default_element") => { + if default_element.is_some() { + abort!(f, "Repeated default_element") + } + let default_element_expr = a.parse_args::<syn::Expr>().unwrap(); + default_element = Some(quote! { #default_element_expr }) + } + Some("condition") => { + if condition.is_some() { + abort!(f, "Repeated condition"); + } + let condition_ast = a.parse_args::<syn::Expr>().unwrap(); + let pretty_cond = prettify_condition(&condition_ast); + condition = Some(Condition { + expr: Some(condition_ast), + has_all_default: all_default_field.is_some(), + pretty: pretty_cond, + }); + } + Some("all_default") => { + if num != 0 { + abort!(f, "all_default is not the first field"); + } + if default.is_some() { + abort!(f, "all_default has an implicit default"); + } + is_all_default = true; + default = Some(quote! { true }); + } + Some("select_coder") => { + if select_coder.is_some() { + abort!(f, "Repeated select_coder"); + } + let condition_ast = a.parse_args::<syn::Expr>().unwrap(); + let pretty_cond = prettify_condition(&condition_ast); + select_coder = Some(Condition { + expr: Some(condition_ast), + has_all_default: false, + pretty: pretty_cond, + }); + } + Some("coder_false") => { + if coder_false.is_some() { + abort!(f, "Repeated coder_false"); + } + let coder_ast = a.parse_args::<syn::Expr>().unwrap(); + coder_false = Some(U32 { + coder: parse_coder(&coder_ast), + }); + } + Some("coder_true") => { + if coder_true.is_some() { + abort!(f, "Repeated coder_true"); + } + let coder_ast = a.parse_args::<syn::Expr>().unwrap(); + coder_true = Some(U32 { + coder: parse_coder(&coder_ast), + }); + } + Some("size_coder") => { + if size_coder.is_some() { + abort!(f, "Repeated size_coder"); + } + let coder_ast = a.parse_args::<syn::Expr>().unwrap(); + size_coder = Some(U32 { + coder: parse_size_coder(coder_ast), + }); + } + Some("nonserialized") => { + let Meta::List(ns) = &a.meta else { + abort!(a, "Invalid attribute"); + }; + let stream = &ns.tokens; + nonserialized.push(quote! {#stream}); + } + _ => {} + } + } + + if default.is_some() && default_element.is_some() { + abort!(f, "default is incompatible with default_element"); + } + + if let Some(select_coder) = select_coder { + if coder_true.is_none() || coder_false.is_none() { + abort!( + f, + "Invalid field, select_coder is set but coder_true or coder_false are not" + ) + } + if coder.is_some() { + abort!(f, "Invalid field, select_coder and coder are both present") + } + coder = Some(Coder::Select( + select_coder, + coder_true.unwrap(), + coder_false.unwrap(), + )) + } + + let condition = if condition.is_some() || all_default_field.is_none() { + condition + } else { + Some(Condition { + expr: None, + has_all_default: true, + pretty: String::new(), + }) + }; + + // Assume nested field if no coder. + let mut coder = coder.unwrap_or_else(|| Coder::WithoutConfig); + + if let Some(c) = size_coder { + if default.is_none() { + default = Some(quote! { Vec::new() }); + } + + coder = Coder::Vector(c, Box::new(coder)) + } + + let ident = f.ident.as_ref().unwrap(); + + let kind = match (condition, default.is_some()) { + (None, _) => FieldKind::Unconditional(coder), + (Some(cond), false) => FieldKind::Conditional(cond, coder), + (Some(cond), true) => FieldKind::Defaulted(cond, coder), + }; + if is_all_default { + *all_default_field = Some(f.ident.as_ref().unwrap().clone()); + } + Field { + name: ident.clone(), + kind, + ty: f.ty.clone(), + default, + default_element, + nonserialized_inits: nonserialized, + } + } + + // Produces reading code (possibly with tracing). + fn read_fun(&self, all_default_field: &Option<syn::Ident>) -> TokenStream2 { + let ident = &self.name; + let ty = &self.ty; + let nonserialized_inits = &self.nonserialized_inits; + match &self.kind { + FieldKind::Unconditional(coder) => { + let cfg_ty = coder.ty(); + let cfg = coder.config(all_default_field); + let trc = quote! { + crate::util::tracing_wrappers::trace!("Setting {} to {:?}. total_bits_read: {}, peek: {:08b}", stringify!(#ident), #ident, br.total_bits_read(), br.peek(8)); + }; + quote! { + let #ident = { + let cfg = #cfg; + type NS = <#ty as UnconditionalCoder<#cfg_ty>>::Nonserialized; + let nonserialized = NS { #(#nonserialized_inits),* }; + <#ty>::read_unconditional(&cfg, br, &nonserialized)? + }; + #trc + } + } + FieldKind::Conditional(condition, coder) => { + let cfg_ty = coder.ty(); + let cfg = coder.config(all_default_field); + let cnd = condition.get_expr(all_default_field).unwrap(); + let pretty_cnd = condition.get_pretty(all_default_field); + let trc = quote! { + crate::util::tracing_wrappers::trace!("{} is {}, setting {} to {:?}. total_bits_read: {}, peek {:08b}", #pretty_cnd, #cnd, stringify!(#ident), #ident, br.total_bits_read(), br.peek(8)); + }; + quote! { + let #ident = { + let cond = #cnd; + let cfg = #cfg; + type NS = <#ty as ConditionalCoder<#cfg_ty>>::Nonserialized; + let nonserialized = NS { #(#nonserialized_inits),* }; + <#ty>::read_conditional(&cfg, cond, br, &nonserialized)? + }; + #trc + } + } + FieldKind::Defaulted(condition, coder) => { + let cfg_ty = coder.ty(); + let cfg = coder.config(all_default_field); + let cnd = condition.get_expr(all_default_field).unwrap(); + let pretty_cnd = condition.get_pretty(all_default_field); + let default = &self.default; + let trc = quote! { + crate::util::tracing_wrappers::trace!("{} is {}, setting {} to {:?}. total_bits_read: {}, peek {:08b}", #pretty_cnd, #cnd, stringify!(#ident), #ident, br.total_bits_read(), br.peek(8)); + }; + + let (read_fn, default) = if let Some(def) = &self.default_element { + (quote! { read_defaulted_element }, Some(def)) + } else { + (quote! { read_defaulted }, default.as_ref()) + }; + + quote! { + let #ident = { + let cond = #cnd; + let cfg = #cfg; + type NS = <#ty as DefaultedCoder<#cfg_ty>>::Nonserialized; + let field_nonserialized = NS { #(#nonserialized_inits),* }; + let default = #default; + <#ty>::#read_fn(&cfg, cond, default, br, &field_nonserialized)? + }; + #trc + } + } + } + } + + // Produces default code. + fn default_code(&self) -> TokenStream2 { + let ident = &self.name; + let ty = &self.ty; + let nonserialized_inits = &self.nonserialized_inits; + let default = &self.default; + match &self.kind { + FieldKind::Defaulted(_, coder) => { + let cfg_ty = coder.ty(); + let default = &self.default; + + quote! { + let #ident = { + type NS = <#ty as DefaultedCoder<#cfg_ty>>::Nonserialized; + let field_nonserialized = NS { #(#nonserialized_inits),* }; + #default + }; + } + } + _ => quote! { let #ident = #default; }, + } + } +} + +fn derive_struct(input: &DeriveInput) -> TokenStream2 { + let name = &input.ident; + + let validate = input.attrs.iter().any(|a| a.path().is_ident("validate")); + let nonserialized: Vec<_> = input + .attrs + .iter() + .filter_map(|a| { + if a.path().is_ident("nonserialized") { + Some(a.parse_args::<syn::Expr>().unwrap()) + } else { + None + } + }) + .collect(); + if nonserialized.len() > 1 { + abort!(input, "repeated nonserialized"); + } + let nonserialized = if nonserialized.is_empty() { + quote! {Empty} + } else { + let v = &nonserialized[0]; + quote! {#v} + }; + + let data = if let syn::Data::Struct(struct_data) = &input.data { + struct_data + } else { + abort!(input, "derive_struct didn't get a struct"); + }; + + let fields = if let syn::Fields::Named(syn::FieldsNamed { + brace_token: _, + named, + }) = &data.fields + { + named + } else { + abort!(data.fields, "only named fields are supported (for now?)"); + }; + + let mut all_default_field = None; + + let fields: Vec<_> = fields + .iter() + .enumerate() + .map(|(n, f)| Field::parse(f, n, &mut all_default_field)) + .collect(); + let fields_read = fields.iter().map(|x| x.read_fun(&all_default_field)); + let fields_names = fields.iter().map(|x| &x.name); + + let impl_default = if fields.iter().all(|x| x.default.is_some()) { + let field_init = fields.iter().map(Field::default_code); + let struct_init = fields.iter().map(|f| { + let ident = &f.name; + quote! { #ident } + }); + quote! { + impl #name { + pub fn default(nonserialized: &#nonserialized) -> #name { + #(#field_init)* + #name { + #(#struct_init),* + } + } + } + + } + } else { + quote! {} + }; + + let impl_validate = if validate { + quote! { return_value.check(nonserialized)?; } + } else { + quote! {} + }; + + let align = match input.attrs.iter().any(|a| a.path().is_ident("aligned")) { + true => quote! { br.jump_to_byte_boundary()?; }, + false => quote! {}, + }; + + quote! { + #impl_default + impl crate::headers::encodings::UnconditionalCoder<()> for #name { + type Nonserialized = #nonserialized; + fn read_unconditional(_: &(), br: &mut BitReader, nonserialized: &Self::Nonserialized) -> Result<#name, Error> { + use crate::headers::encodings::UnconditionalCoder; + use crate::headers::encodings::ConditionalCoder; + use crate::headers::encodings::DefaultedCoder; + use crate::headers::encodings::DefaultedElementCoder; + #align + #(#fields_read)* + let return_value = #name { + #(#fields_names),* + }; + #impl_validate + Ok(return_value) + } + } + } +} + +fn derive_enum(input: &DeriveInput) -> TokenStream2 { + let name = &input.ident; + quote! { + impl crate::headers::encodings::UnconditionalCoder<U32Coder> for #name { + type Nonserialized = Empty; + fn read_unconditional(config: &U32Coder, br: &mut BitReader, _: &Empty) -> Result<#name, Error> { + use num_traits::FromPrimitive; + let u = u32::read_unconditional(config, br, &Empty{})?; + if let Some(e) = #name::from_u32(u) { + Ok(e) + } else { + Err(Error::InvalidEnum(u, stringify!(#name).to_string())) + } + } + } + impl crate::headers::encodings::UnconditionalCoder<()> for #name { + type Nonserialized = Empty; + fn read_unconditional(config: &(), br: &mut BitReader, nonserialized: &Empty) -> Result<#name, Error> { + #name::read_unconditional( + &U32Coder::Select( + U32::Val(0), U32::Val(1), + U32::BitsOffset{n: 4, off: 2}, + U32::BitsOffset{n: 6, off: 18}), br, nonserialized) + } + } + } +} + +#[proc_macro_error] +#[proc_macro_derive( + UnconditionalCoder, + attributes( + coder, + condition, + default, + default_element, + all_default, + select_coder, + coder_true, + coder_false, + validate, + size_coder, + nonserialized, + aligned, + ) +)] +pub fn derive_jxl_headers(input: TokenStream) -> TokenStream { + let input = parse_macro_input!(input as DeriveInput); + + match &input.data { + syn::Data::Struct(_) => derive_struct(&input).into(), + syn::Data::Enum(_) => derive_enum(&input).into(), + _ => abort!(input, "Only implemented for struct"), + } +} + +#[proc_macro_attribute] +pub fn noop(_attr: TokenStream, item: TokenStream) -> TokenStream { + item +} + +#[cfg(feature = "test")] +#[proc_macro] +pub fn for_each_test_file(input: TokenStream) -> TokenStream { + use std::{fs, path::Path}; + use syn::Ident; + + let fn_name = parse_macro_input!(input as Ident); + let root_test_dir = Path::new(env!("CARGO_MANIFEST_DIR")) + .join("..") + .join("jxl") + .join("resources") + .join("test"); + let conformance_test_dir = root_test_dir.join("conformance_test_images"); + + let mut tests = vec![]; + + for test_dir in [root_test_dir, conformance_test_dir] { + for entry in fs::read_dir(&test_dir).unwrap() { + let entry = entry.unwrap(); + let path = entry.path(); + if path.extension().is_some_and(|ext| ext == "jxl") { + let pathname = path.to_string_lossy(); + let relative_path = path + .strip_prefix(&test_dir) + .unwrap() + .to_string_lossy() + .replace('/', "_slash_"); + let test_name = format!( + "{}_{}", + fn_name, + relative_path.strip_suffix(".jxl").unwrap() + ); + let test_name = Ident::new(&test_name, fn_name.span()); + tests.push(quote! { + #[test] + fn #test_name() { + #fn_name(&Path::new(#pathname)).unwrap() + } + }); + } + } + } + + quote! { + #(#tests)* + } + .into() +} diff --git a/third_party/rust/proc-macro-error-attr2/.cargo-checksum.json b/third_party/rust/proc-macro-error-attr2/.cargo-checksum.json @@ -0,0 +1 @@ +{"files":{"Cargo.toml":"8584e30fb2d9e38659e27d3b7b7e9e69c093b8f84de0bc1d7ff96735de8be22d","LICENSE-APACHE":"6fd0f3522047150ca7c1939f02bc4a15662a4741a89bc03ae784eefa18caa299","LICENSE-MIT":"544b3aed1fd723d0cadea567affdcfe0431e43e18d997a718f9d67256b814fde","src/lib.rs":"0a7709be2204e227cedb390121b7802f0d1dd533cd7f2745aee0656f4f788e0a","src/parse.rs":"ea729b99efa701153d3d5a248e1a6a82de7eb7a839e9e069c1f8fdc25d319c36","src/settings.rs":"1a7d1b941869712c78e1758abdc5df8720828265f318033d1820f57c2cf74094"},"package":"96de42df36bb9bba5542fe9f1a054b8cc87e172759a1868aa05c1f3acc89dfc5"} +\ No newline at end of file diff --git a/third_party/rust/proc-macro-error-attr2/Cargo.toml b/third_party/rust/proc-macro-error-attr2/Cargo.toml @@ -0,0 +1,44 @@ +# THIS FILE IS AUTOMATICALLY GENERATED BY CARGO +# +# When uploading crates to the registry Cargo will automatically +# "normalize" Cargo.toml files for maximal compatibility +# with all versions of Cargo and also rewrite `path` dependencies +# to registry (e.g., crates.io) dependencies. +# +# If you are reading this file be aware that the original Cargo.toml +# will likely look very different (and much more reasonable). +# See Cargo.toml.orig for the original contents. + +[package] +edition = "2021" +rust-version = "1.61" +name = "proc-macro-error-attr2" +version = "2.0.0" +authors = [ + "CreepySkeleton <creepy-skeleton@yandex.ru>", + "GnomedDev <david2005thomas@gmail.com>", +] +build = false +autobins = false +autoexamples = false +autotests = false +autobenches = false +description = "Attribute macro for the proc-macro-error2 crate" +readme = false +license = "MIT OR Apache-2.0" +repository = "https://github.com/GnomedDev/proc-macro-error-2" + +[lib] +name = "proc_macro_error_attr2" +path = "src/lib.rs" +proc-macro = true + +[dependencies.proc-macro2] +version = "1" + +[dependencies.quote] +version = "1" + +[lints.clippy.pedantic] +level = "warn" +priority = -1 diff --git a/third_party/rust/proc-macro-error-attr2/LICENSE-APACHE b/third_party/rust/proc-macro-error-attr2/LICENSE-APACHE @@ -0,0 +1,201 @@ + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + +TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + +1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + +2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + +3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + +4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + +5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + +6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + +7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + +8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + +9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + +END OF TERMS AND CONDITIONS + +APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + +Copyright 2019-2020 CreepySkeleton <creepy-skeleton@yandex.ru> + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. diff --git a/third_party/rust/proc-macro-error-attr2/LICENSE-MIT b/third_party/rust/proc-macro-error-attr2/LICENSE-MIT @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2019-2020 CreepySkeleton + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/third_party/rust/proc-macro-error-attr2/src/lib.rs b/third_party/rust/proc-macro-error-attr2/src/lib.rs @@ -0,0 +1,111 @@ +//! This is `#[proc_macro_error]` attribute to be used with +//! [`proc-macro-error`](https://docs.rs/proc-macro-error2/). There you go. + +use crate::parse::parse_input; +use crate::parse::Attribute; +use proc_macro::TokenStream; +use proc_macro2::{Literal, Span, TokenStream as TokenStream2, TokenTree}; +use quote::{quote, quote_spanned}; + +use crate::settings::{ + parse_settings, + Setting::{AllowNotMacro, AssertUnwindSafe, ProcMacroHack}, + Settings, +}; + +mod parse; +mod settings; + +type Result<T> = std::result::Result<T, Error>; + +struct Error { + span: Span, + message: String, +} + +impl Error { + fn new(span: Span, message: String) -> Self { + Error { span, message } + } + + fn into_compile_error(self) -> TokenStream2 { + let mut message = Literal::string(&self.message); + message.set_span(self.span); + quote_spanned!(self.span=> compile_error!{#message}) + } +} + +#[proc_macro_attribute] +pub fn proc_macro_error(attr: TokenStream, input: TokenStream) -> TokenStream { + match impl_proc_macro_error(attr.into(), input.clone().into()) { + Ok(ts) => ts, + Err(e) => { + let error = e.into_compile_error(); + let input = TokenStream2::from(input); + + quote!(#input #error).into() + } + } +} + +fn impl_proc_macro_error(attr: TokenStream2, input: TokenStream2) -> Result<TokenStream> { + let (attrs, signature, body) = parse_input(input)?; + let mut settings = parse_settings(attr)?; + + let is_proc_macro = is_proc_macro(&attrs); + if is_proc_macro { + settings.set(AssertUnwindSafe); + } + + if detect_proc_macro_hack(&attrs) { + settings.set(ProcMacroHack); + } + + if settings.is_set(ProcMacroHack) { + settings.set(AllowNotMacro); + } + + if !(settings.is_set(AllowNotMacro) || is_proc_macro) { + return Err(Error::new( + Span::call_site(), + "#[proc_macro_error] attribute can be used only with procedural macros\n\n \ + = hint: if you are really sure that #[proc_macro_error] should be applied \ + to this exact function, use #[proc_macro_error(allow_not_macro)]\n" + .into(), + )); + } + + let body = gen_body(&body, &settings); + + let res = quote! { + #(#attrs)* + #(#signature)* + { #body } + }; + Ok(res.into()) +} + +fn gen_body(block: &TokenTree, settings: &Settings) -> proc_macro2::TokenStream { + let is_proc_macro_hack = settings.is_set(ProcMacroHack); + let closure = if settings.is_set(AssertUnwindSafe) { + quote!(::std::panic::AssertUnwindSafe(|| #block )) + } else { + quote!(|| #block) + }; + + quote!( ::proc_macro_error2::entry_point(#closure, #is_proc_macro_hack) ) +} + +fn detect_proc_macro_hack(attrs: &[Attribute]) -> bool { + attrs + .iter() + .any(|attr| attr.path_is_ident("proc_macro_hack")) +} + +fn is_proc_macro(attrs: &[Attribute]) -> bool { + attrs.iter().any(|attr| { + attr.path_is_ident("proc_macro") + || attr.path_is_ident("proc_macro_derive") + || attr.path_is_ident("proc_macro_attribute") + }) +} diff --git a/third_party/rust/proc-macro-error-attr2/src/parse.rs b/third_party/rust/proc-macro-error-attr2/src/parse.rs @@ -0,0 +1,89 @@ +use crate::{Error, Result}; +use proc_macro2::{Delimiter, Ident, Span, TokenStream, TokenTree}; +use quote::ToTokens; +use std::iter::Peekable; + +pub(crate) fn parse_input( + input: TokenStream, +) -> Result<(Vec<Attribute>, Vec<TokenTree>, TokenTree)> { + let mut input = input.into_iter().peekable(); + let mut attrs = Vec::new(); + + while let Some(attr) = parse_next_attr(&mut input)? { + attrs.push(attr); + } + + let sig = parse_signature(&mut input); + let body = input.next().ok_or_else(|| { + Error::new( + Span::call_site(), + "`#[proc_macro_error]` can be applied only to functions".to_string(), + ) + })?; + + Ok((attrs, sig, body)) +} + +fn parse_next_attr( + input: &mut Peekable<impl Iterator<Item = TokenTree>>, +) -> Result<Option<Attribute>> { + let shebang = match input.peek() { + Some(TokenTree::Punct(ref punct)) if punct.as_char() == '#' => input.next().unwrap(), + _ => return Ok(None), + }; + + let group = match input.peek() { + Some(TokenTree::Group(ref group)) if group.delimiter() == Delimiter::Bracket => { + let res = group.clone(); + input.next(); + res + } + other => { + let span = other.map_or(Span::call_site(), TokenTree::span); + return Err(Error::new(span, "expected `[`".to_string())); + } + }; + + let path = match group.stream().into_iter().next() { + Some(TokenTree::Ident(ident)) => Some(ident), + _ => None, + }; + + Ok(Some(Attribute { + shebang, + group: TokenTree::Group(group), + path, + })) +} + +fn parse_signature(input: &mut Peekable<impl Iterator<Item = TokenTree>>) -> Vec<TokenTree> { + let mut sig = Vec::new(); + loop { + match input.peek() { + Some(TokenTree::Group(ref group)) if group.delimiter() == Delimiter::Brace => { + return sig; + } + None => return sig, + _ => sig.push(input.next().unwrap()), + } + } +} + +pub(crate) struct Attribute { + pub(crate) shebang: TokenTree, + pub(crate) group: TokenTree, + pub(crate) path: Option<Ident>, +} + +impl Attribute { + pub(crate) fn path_is_ident(&self, ident: &str) -> bool { + self.path.as_ref().map_or(false, |p| *p == ident) + } +} + +impl ToTokens for Attribute { + fn to_tokens(&self, ts: &mut TokenStream) { + self.shebang.to_tokens(ts); + self.group.to_tokens(ts); + } +} diff --git a/third_party/rust/proc-macro-error-attr2/src/settings.rs b/third_party/rust/proc-macro-error-attr2/src/settings.rs @@ -0,0 +1,72 @@ +use crate::{Error, Result}; +use proc_macro2::{Ident, Span, TokenStream, TokenTree}; + +macro_rules! decl_settings { + ($($val:expr => $variant:ident),+ $(,)*) => { + #[derive(PartialEq, Clone, Copy)] + pub(crate) enum Setting { + $($variant),* + } + + fn ident_to_setting(ident: Ident) -> Result<Setting> { + match &*ident.to_string() { + $($val => Ok(Setting::$variant),)* + _ => { + let possible_vals = [$($val),*] + .iter() + .map(|v| format!("`{}`", v)) + .collect::<Vec<_>>() + .join(", "); + + Err(Error::new( + ident.span(), + format!("unknown setting `{}`, expected one of {}", ident, possible_vals))) + } + } + } + }; +} + +decl_settings! { + "assert_unwind_safe" => AssertUnwindSafe, + "allow_not_macro" => AllowNotMacro, + "proc_macro_hack" => ProcMacroHack, +} + +pub(crate) fn parse_settings(input: TokenStream) -> Result<Settings> { + let mut input = input.into_iter(); + let mut res = Settings(Vec::new()); + loop { + match input.next() { + Some(TokenTree::Ident(ident)) => { + res.0.push(ident_to_setting(ident)?); + } + None => return Ok(res), + other => { + let span = other.map_or(Span::call_site(), |tt| tt.span()); + return Err(Error::new(span, "expected identifier".to_string())); + } + } + + match input.next() { + Some(TokenTree::Punct(ref punct)) if punct.as_char() == ',' => {} + None => return Ok(res), + other => { + let span = other.map_or(Span::call_site(), |tt| tt.span()); + return Err(Error::new(span, "expected `,`".to_string())); + } + } + } +} + +pub(crate) struct Settings(Vec<Setting>); + +impl Settings { + pub(crate) fn is_set(&self, setting: Setting) -> bool { + self.0.iter().any(|s| *s == setting) + } + + pub(crate) fn set(&mut self, setting: Setting) { + self.0.push(setting); + } +} diff --git a/third_party/rust/proc-macro-error2/.cargo-checksum.json b/third_party/rust/proc-macro-error2/.cargo-checksum.json @@ -0,0 +1 @@ +{"files":{"CHANGELOG.md":"be4c548354e0713a7ce362a6d668d5304c940e02c9e7c95c0f58763691088e8d","Cargo.toml":"9419aea8db386d6d06fabe4b2b808b853e233e84df0484d598b81a1836d64ae8","LICENSE-APACHE":"4665f973ccb9393807a7fb1264add6b3513d19abc6b357e4cb52c6fe59cc6a3b","LICENSE-MIT":"544b3aed1fd723d0cadea567affdcfe0431e43e18d997a718f9d67256b814fde","README.md":"186fa2911fda58f18b9c012f844e6110cc497791643dd01854988b2bf4f36b45","src/diagnostic.rs":"638f2c7113f71694d4249e1c61b814d60e65cabca771eb197b531a9e33fd5926","src/dummy.rs":"ae20aeaa2262067b281492db68cbc8669ae65bfcbac91f6d471d0f8f89047148","src/imp/delegate.rs":"ba738cf74ae42fbe058bf3251a412c6131f332d74eed0c1644877a35565cf8dd","src/imp/fallback.rs":"125aa1d68c96c471f39af2c441891a7fc61ca8366635bccb6eea94027d9aa084","src/lib.rs":"eeada6be6ecd1172526ffd680372e9a1d5667a68820a85d91b8719f39ab879d4","src/macros.rs":"3be6feccd343cd9dc4bf03780f3107909bf70e02c6c7c72912e4b160dc6a68fc","src/sealed.rs":"dcf569c4c7ce1d372ff51b0fa73fa67f995bdca8e520cb225cde789c71276439","tests/macro-errors.rs":"ff16f408bd53576995522af1bd6919ef5791d571c03dabd19aed44d7f08e7058","tests/ok.rs":"8379b6066c0b8bd939fee7d92bf6306f19f8300b0b03c201cc78a45f4a58f831","tests/runtime-errors.rs":"fdab67272feaa97fb853e5f015325cd411dca59e4d2f6fe00d0aeaf8546e3760","tests/ui/abort.rs":"25629a3319c8dec464b8bd37d243379348a4c974b22c0df01b55bca499b66ecf","tests/ui/abort.stderr":"41dfc53c80681a3fedc334663e45a3d5237212458800fa4ff5914e24b1e15080","tests/ui/append_dummy.rs":"42942587c0720f3cfca2de7a2a1a38b1c14b35bcfd737c0f95a445884332df95","tests/ui/append_dummy.stderr":"9ea462a1fcc5cc10571c86920c24cc8f223e0fe60d33be55c02005d57faeb6ee","tests/ui/children_messages.rs":"aac44ffcfc73726d9e347735a8a7dea32c7cc0decb2360e3cbdbdcfe46d15f09","tests/ui/children_messages.stderr":"0b09c7eeefa62c31ecd07670980fd8345c243dac89cb1402415ff7f800798bb8","tests/ui/dummy.rs":"e7134b23dc19ef9a58d4b253483e45bcfa882c5c373ce1811cba5acdc5e6287c","tests/ui/dummy.stderr":"3d9394b1a4ce864276ee24be2e3c828f79f76f137bfe0ed9dca6a9b6b7e6ef2f","tests/ui/emit.rs":"ae9dfcc4dbc3e4571db627b8940ae5f448db87c1bf7bc6927c85ec9b9c58b53d","tests/ui/emit.stderr":"03be9ee199a6819d30ed340f64c558fe05214739060ada29c09e48cf55ee0d3a","tests/ui/explicit_span_range.rs":"64e7a187856475d1ad501124355db46ee73cfc3250a0da9b5112338100334b0d","tests/ui/explicit_span_range.stderr":"d8a1288274101ee06baad91280d023732acac0f2a5843c65ea107720698549ad","tests/ui/misuse.rs":"e504820bf8c52854073a6ecfd3486c2759deee7514079ec80ebee5a614a7ccea","tests/ui/misuse.stderr":"d6ec6d638f0f107e8faa97741cd5354583a1ffa4bee07eccc8a4069dd838f9f3","tests/ui/multiple_tokens.rs":"c212bcd27d8551a42bbf2e83a254bc13adb02fbb09cb7df82777b1b426f07451","tests/ui/multiple_tokens.stderr":"c7026c7931c42e204f2c67623cc2755137720182bfc5e4b79e61ae0d7d647f72","tests/ui/not_proc_macro.rs":"23034d95ba47cda5a349b672c3ce5c93fe5ed0d51b4cdb3bf17f9cdad434de32","tests/ui/not_proc_macro.stderr":"12dbf434f8f6f2c4517cba2b07e1f28e0d87f4e87f27527c0c4341373fd6563c","tests/ui/option_ext.rs":"38a7377e4ff1858a5b95116e285c744daf6a010e1deb8fdd9d53532ac745ee70","tests/ui/option_ext.stderr":"986b9d212ad26e7e4854991fa2107fb0d52274b8a92ead4aadab2bb41ac2281b","tests/ui/result_ext.rs":"84f6dfe83fa619d2eac4a6df1f285b0e36755d15ecff75445b1d0dc96ebd2b28","tests/ui/result_ext.stderr":"006105f86d2533309191dbce61f6f388e5d03c94474f1bcd1481db314e4a66ec","tests/ui/to_tokens_span.rs":"4b89ddd4a727422da558c02887ffd712da9534d230593eb7b8e367a0f62eb2cf","tests/ui/to_tokens_span.stderr":"41cc86b25d20bf568b8f14c8572f991e4cf2a0f7c4e35978d35191974a0fa2b4","tests/ui/unknown_setting.rs":"ac88d5775207ed57b921a3139586e055ddcf5f683227fad5723a2a8ddfb7392e","tests/ui/unknown_setting.stderr":"746ef8d8cf2a60acf84b4ba81ae034682aaa8b5f5275803a34fce63378a0f0df","tests/ui/unrelated_panic.rs":"444e74b52373e260ec150de951b27b72f93b0acc3763e064cefbf320189d4bf1","tests/ui/unrelated_panic.stderr":"a21b02875cfc58271c2d8cbdde40a77f41f82c4b829a985f42f579c6e85c4524"},"package":"11ec05c52be0a07b08061f7dd003e7d7092e0472bc731b4af7bb1ef876109802"} +\ No newline at end of file diff --git a/third_party/rust/proc-macro-error2/CHANGELOG.md b/third_party/rust/proc-macro-error2/CHANGELOG.md @@ -0,0 +1,180 @@ +# v2.0.1 (2024-09-06) + +* Fixed a span location issue due to mistake in refactoring (#2) + +# v2.0.0 (2024-09-05) + +No changes, simply releasing pre-release as full release. + +# v2.0.0-pre.1 (2024-09-01) + +* __Crate has been renamed to `proc-macro-error2`, due to the old maintainer's inactivity.__ + +* `syn` has been upgraded to `2` +* MSRV has been bumped to `1.61` +* Warnings have been fixed, including `clippy::pedantic` lints +* CI has been converted to GitHub actions, and testing infrastructure significantly simplified. +* Automatic nightly detection has been removed, use the `nightly` feature for improved diagnostics at the cost of stablity. + +# v1.0.4 (2020-7-31) + +* `SpanRange` facility is now public. +* Docs have been improved. +* Introduced the `syn-error` feature so you can opt-out from the `syn` dependency. + +# v1.0.3 (2020-6-26) + +* Corrected a few typos. +* Fixed the `emit_call_site_warning` macro. + +# v1.0.2 (2020-4-9) + +* An obsolete note was removed from documentation. + +# v1.0.1 (2020-4-9) + +* `proc-macro-hack` is now well tested and supported. Not sure about `proc-macro-nested`, + please fill a request if you need it. +* Fixed `emit_call_site_error`. +* Documentation improvements. + +# v1.0.0 (2020-3-25) + +I believe the API can be considered stable because it's been a few months without +breaking changes, and I also don't think this crate will receive much further evolution. +It's perfect, admit it. + +Hence, meet the new, stable release! + +### Improvements + +* Supported nested `#[proc_macro_error]` attributes. Well, you aren't supposed to do that, + but I caught myself doing it by accident on one occasion and the behavior was... surprising. + Better to handle this smooth. + +# v0.4.12 (2020-3-23) + +* Error message on macros' misuse is now a bit more understandable. + +# v0.4.11 (2020-3-02) + +* `build.rs` no longer fails when `rustc` date could not be determined, + (thanks to [`Fabian Möller`](https://gitlab.com/CreepySkeleton/proc-macro-error/issues/8) + for noticing and to [`Igor Gnatenko`](https://gitlab.com/CreepySkeleton/proc-macro-error/-/merge_requests/25) + for fixing). + +# v0.4.10 (2020-2-29) + +* `proc-macro-error` doesn't depend on syn\[full\] anymore, the compilation + is \~30secs faster. + +# v0.4.9 (2020-2-13) + +* New function: `append_dummy`. + +# v0.4.8 (2020-2-01) + +* Support for children messages + +# v0.4.7 (2020-1-31) + +* Now any type that implements `quote::ToTokens` can be used instead of spans. + This allows for high quality error messages. + +# v0.4.6 (2020-1-31) + +* `From<syn::Error>` implementation doesn't lose span info anymore, see + [#6](https://gitlab.com/CreepySkeleton/proc-macro-error/issues/6). + +# v0.4.5 (2020-1-20) +Just a small intermediate release. + +* Fix some bugs. +* Populate license files into subfolders. + +# v0.4.4 (2019-11-13) +* Fix `abort_if_dirty` + warnings bug +* Allow trailing commas in macros + +# v0.4.2 (2019-11-7) +* FINALLY fixed `__pme__suggestions not found` bug + +# v0.4.1 (2019-11-7) YANKED +* Fixed `__pme__suggestions not found` bug +* Documentation improvements, links checked + +# v0.4.0 (2019-11-6) YANKED + +## New features +* "help" messages that can have their own span on nightly, they + inherit parent span on stable. + ```rust + let cond_help = if condition { Some("some help message") else { None } }; + abort!( + span, // parent span + "something's wrong, {} wrongs in total", 10; // main message + help = "here's a help for you, {}", "take it"; // unconditional help message + help =? cond_help; // conditional help message, must be Option + note = note_span => "don't forget the note, {}", "would you?" // notes can have their own span but it's effective only on nightly + ) + ``` +* Warnings via `emit_warning` and `emit_warning_call_site`. Nightly only, they're ignored on stable. +* Now `proc-macro-error` delegates to `proc_macro::Diagnostic` on nightly. + +## Breaking changes +* `MacroError` is now replaced by `Diagnostic`. Its API resembles `proc_macro::Diagnostic`. +* `Diagnostic` does not implement `From<&str/String>` so `Result<T, &str/String>::abort_or_exit()` + won't work anymore (nobody used it anyway). +* `macro_error!` macro is replaced with `diagnostic!`. + +## Improvements +* Now `proc-macro-error` renders notes exactly just like rustc does. +* We don't parse a body of a function annotated with `#[proc_macro_error]` anymore, + only looking at the signature. This should somewhat decrease expansion time for large functions. + +# v0.3.3 (2019-10-16) +* Now you can use any word instead of "help", undocumented. + +# v0.3.2 (2019-10-16) +* Introduced support for "help" messages, undocumented. + +# v0.3.0 (2019-10-8) + +## The crate has been completely rewritten from scratch! + +## Changes (most are breaking): +* Renamed macros: + * `span_error` => `abort` + * `call_site_error` => `abort_call_site` +* `filter_macro_errors` was replaced by `#[proc_macro_error]` attribute. +* `set_dummy` now takes `TokenStream` instead of `Option<TokenStream>` +* Support for multiple errors via `emit_error` and `emit_call_site_error` +* New `macro_error` macro for building errors in format=like style. +* `MacroError` API had been reconsidered. It also now implements `quote::ToTokens`. + +# v0.2.6 (2019-09-02) +* Introduce support for dummy implementations via `dummy::set_dummy` +* `multi::*` is now deprecated, will be completely rewritten in v0.3 + +# v0.2.0 (2019-08-15) + +## Breaking changes +* `trigger_error` replaced with `MacroError::trigger` and `filter_macro_error_panics` + is hidden from docs. + This is not quite a breaking change since users weren't supposed to use these functions directly anyway. +* All dependencies are updated to `v1.*`. + +## New features +* Ability to stack multiple errors via `multi::MultiMacroErrors` and emit them at once. + +## Improvements +* Now `MacroError` implements `std::fmt::Display` instead of `std::string::ToString`. +* `MacroError::span` inherent method. +* `From<MacroError> for proc_macro/proc_macro2::TokenStream` implementations. +* `AsRef/AsMut<String> for MacroError` implementations. + +# v0.1.x (2019-07-XX) + +## New features +* An easy way to report errors inside within a proc-macro via `span_error`, + `call_site_error` and `filter_macro_errors`. diff --git a/third_party/rust/proc-macro-error2/Cargo.toml b/third_party/rust/proc-macro-error2/Cargo.toml @@ -0,0 +1,91 @@ +# THIS FILE IS AUTOMATICALLY GENERATED BY CARGO +# +# When uploading crates to the registry Cargo will automatically +# "normalize" Cargo.toml files for maximal compatibility +# with all versions of Cargo and also rewrite `path` dependencies +# to registry (e.g., crates.io) dependencies. +# +# If you are reading this file be aware that the original Cargo.toml +# will likely look very different (and much more reasonable). +# See Cargo.toml.orig for the original contents. + +[package] +edition = "2021" +rust-version = "1.61" +name = "proc-macro-error2" +version = "2.0.1" +authors = [ + "CreepySkeleton <creepy-skeleton@yandex.ru>", + "GnomedDev <david2005thomas@gmail.com>", +] +build = false +autobins = false +autoexamples = false +autotests = false +autobenches = false +description = "Almost drop-in replacement to panics in proc-macros" +readme = "README.md" +keywords = [ + "proc-macro", + "error", + "errors", +] +categories = ["development-tools::procedural-macro-helpers"] +license = "MIT OR Apache-2.0" +repository = "https://github.com/GnomedDev/proc-macro-error-2" + +[lib] +name = "proc_macro_error2" +path = "src/lib.rs" + +[[test]] +name = "macro-errors" +path = "tests/macro-errors.rs" + +[[test]] +name = "ok" +path = "tests/ok.rs" + +[[test]] +name = "runtime-errors" +path = "tests/runtime-errors.rs" + +[dependencies.proc-macro-error-attr2] +version = "=2.0.0" + +[dependencies.proc-macro2] +version = "1" + +[dependencies.quote] +version = "1" + +[dependencies.syn] +version = "2" +optional = true +default-features = false + +[dev-dependencies.syn] +version = "2" +features = ["full"] + +[dev-dependencies.trybuild] +version = "1.0.99" +features = ["diff"] + +[features] +default = ["syn-error"] +nightly = [] +syn-error = ["dep:syn"] + +[lints.clippy.module_name_repetitions] +level = "allow" +priority = 0 + +[lints.clippy.pedantic] +level = "warn" +priority = -1 + +[lints.rust.unexpected_cfgs] +level = "warn" +priority = 0 +check-cfg = ["cfg(run_ui_tests)"] diff --git a/third_party/rust/proc-macro-error2/LICENSE-APACHE b/third_party/rust/proc-macro-error2/LICENSE-APACHE @@ -0,0 +1,201 @@ + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + +TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + +1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + +2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + +3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + +4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + +5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + +6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + +7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + +8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + +9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + +END OF TERMS AND CONDITIONS + +APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + +Copyright 2019-2020 CreepySkeleton <creepy-skeleton@yandex.ru> + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. diff --git a/third_party/rust/proc-macro-error2/LICENSE-MIT b/third_party/rust/proc-macro-error2/LICENSE-MIT @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2019-2020 CreepySkeleton + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/third_party/rust/proc-macro-error2/README.md b/third_party/rust/proc-macro-error2/README.md @@ -0,0 +1,250 @@ +# Makes error reporting in procedural macros nice and easy + +[](https://docs.rs/proc-macro-error2) +[](https://github.com/rust-secure-code/safety-dance/) + +This crate aims to make error reporting in proc-macros simple and easy to use. +Migrate from `panic!`-based errors for as little effort as possible! + +Also, you can explicitly [append a dummy token stream][crate::dummy] to your errors. + +To achieve this, this crate serves as a tiny shim around `proc_macro::Diagnostic` and +`compile_error!`. It detects the most preferable way to emit errors based on compiler's version. +When the underlying diagnostic type is finally stabilized, this crate will be simply +delegating to it, requiring no changes in your code! + +So you can just use this crate and have *both* some of `proc_macro::Diagnostic` functionality +available on stable ahead of time and your error-reporting code future-proof. + +```toml +[dependencies] +proc-macro-error2 = "2.0" +``` + +*Supports rustc 1.61 and up* + +[Documentation and guide][guide] + +## Quick example + +Code: + +```rust +#[proc_macro] +#[proc_macro_error] +pub fn make_fn(input: TokenStream) -> TokenStream { + let mut input = TokenStream2::from(input).into_iter(); + let name = input.next().unwrap(); + if let Some(second) = input.next() { + abort! { second, + "I don't like this part!"; + note = "I see what you did there..."; + help = "I need only one part, you know?"; + } + } + + quote!( fn #name() {} ).into() +} +``` + +This is how the error is rendered in a terminal: + +<p align="center"> +<img src="https://user-images.githubusercontent.com/50968528/78830016-d3b46a80-79d6-11ea-9de2-972e8d7904ef.png" width="600"> +</p> + +And this is what your users will see in their IDE: + +<p align="center"> +<img src="https://user-images.githubusercontent.com/50968528/78830547-a9af7800-79d7-11ea-822e-59e29bda335c.png" width="600"> +</p> + +## Examples + +### Panic-like usage + +```rust +use proc_macro_error2::{ + proc_macro_error, + abort, + abort_call_site, + ResultExt, + OptionExt, +}; +use proc_macro::TokenStream; +use syn::{DeriveInput, parse_macro_input}; +use quote::quote; + +// This is your main entry point +#[proc_macro] +// This attribute *MUST* be placed on top of the #[proc_macro] function +#[proc_macro_error] +pub fn make_answer(input: TokenStream) -> TokenStream { + let input = parse_macro_input!(input as DeriveInput); + + if let Err(err) = some_logic(&input) { + // we've got a span to blame, let's use it + // This immediately aborts the proc-macro and shows the error + // + // You can use `proc_macro::Span`, `proc_macro2::Span`, and + // anything that implements `quote::ToTokens` (almost every type from + // `syn` and `proc_macro2`) + abort!(err, "You made an error, go fix it: {}", err.msg); + } + + // `Result` has some handy shortcuts if your error type implements + // `Into<Diagnostic>`. `Option` has one unconditionally. + more_logic(&input).expect_or_abort("What a careless user, behave!"); + + if !more_logic_for_logic_god(&input) { + // We don't have an exact location this time, + // so just highlight the proc-macro invocation itself + abort_call_site!( + "Bad, bad user! Now go stand in the corner and think about what you did!"); + } + + // Now all the processing is done, return `proc_macro::TokenStream` + quote!(/* stuff */).into() +} +``` + +### `proc_macro::Diagnostic`-like usage + +```rust +use proc_macro_error2::*; +use proc_macro::TokenStream; +use syn::{spanned::Spanned, DeriveInput, ItemStruct, Fields, Attribute , parse_macro_input}; +use quote::quote; + +fn process_attrs(attrs: &[Attribute]) -> Vec<Attribute> { + attrs + .iter() + .filter_map(|attr| match process_attr(attr) { + Ok(res) => Some(res), + Err(msg) => { + emit_error!(attr, "Invalid attribute: {}", msg); + None + } + }) + .collect() +} + +fn process_fields(_attrs: &Fields) -> Vec<TokenStream> { + // processing fields in pretty much the same way as attributes + unimplemented!() +} + +#[proc_macro] +#[proc_macro_error] +pub fn make_answer(input: TokenStream) -> TokenStream { + let input = parse_macro_input!(input as ItemStruct); + let attrs = process_attrs(&input.attrs); + + // abort right now if some errors were encountered + // at the attributes processing stage + abort_if_dirty(); + + let fields = process_fields(&input.fields); + + // no need to think about emitted errors + // #[proc_macro_error] will handle them for you + // + // just return a TokenStream as you normally would + quote!(/* stuff */).into() +} +``` + +## Real world examples + +* [`structopt-derive`](https://github.com/TeXitoi/structopt/tree/master/structopt-derive) + (abort-like usage) +* [`auto-impl`](https://github.com/auto-impl-rs/auto_impl/) (emit-like usage) + +## Limitations + +- Warnings are emitted only on nightly, they are ignored on stable. +- "help" suggestions can't have their own span info on stable, + (essentially inheriting the parent span). +- If your macro happens to trigger a panic, no errors will be displayed. This is not a + technical limitation but rather intentional design. `panic` is not for error reporting. + +## MSRV policy + +The MSRV is currently `1.61`, and this is considered a breaking change to increase. + +However, if an existing dependency requires a higher MSRV without a semver breaking update, this may be raised. + +## Motivation + +Error handling in proc-macros sucks. There's not much of a choice today: +you either "bubble up" the error up to the top-level of the macro and convert it to +a [`compile_error!`][compl_err] invocation or just use a good old panic. Both these ways suck: + +- Former sucks because it's quite redundant to unroll a proper error handling + just for critical errors that will crash the macro anyway; so people mostly + choose not to bother with it at all and use panic. Simple `.expect` is too tempting. + + Also, if you do decide to implement this `Result`-based architecture in your macro + you're going to have to rewrite it entirely once [`proc_macro::Diagnostic`][] is finally + stable. Not cool. + +- Later sucks because there's no way to carry out the span info via `panic!`. + `rustc` will highlight the invocation itself but not some specific token inside it. + + Furthermore, panics aren't for error-reporting at all; panics are for bug-detecting + (like unwrapping on `None` or out-of-range indexing) or for early development stages + when you need a prototype ASAP so error handling can wait. Mixing these usages only + messes things up. + +- There is [`proc_macro::Diagnostic`][] which is awesome but it has been experimental + for more than a year and is unlikely to be stabilized any time soon. + + This crate's API is intentionally designed to be compatible with `proc_macro::Diagnostic` + and delegates to it whenever possible. Once `Diagnostics` is stable this crate + will **always** delegate to it, no code changes will be required on user side. + +That said, we need a solution, but this solution must meet these conditions: + +- It must be better than `panic!`. The main point: it must offer a way to carry the span information + over to user. +- It must take as little effort as possible to migrate from `panic!`. Ideally, a new + macro with similar semantics plus ability to carry out span info. +- It must maintain compatibility with [`proc_macro::Diagnostic`][] . +- **It must be usable on stable**. + +This crate aims to provide such a mechanism. All you have to do is annotate your top-level +`#[proc_macro]` function with `#[proc_macro_error]` attribute and change panics to +[`abort!`]/[`abort_call_site!`] where appropriate, see [the Guide][guide]. + +## Disclaimer +Please note that **this crate is not intended to be used in any way other +than error reporting in procedural macros**, use `Result` and `?` (possibly along with one of the +many helpers out there) for anything else. + +<br> + +#### License + +<sup> +Licensed under either of <a href="LICENSE-APACHE">Apache License, Version +2.0</a> or <a href="LICENSE-MIT">MIT license</a> at your option. +</sup> + +<br> + +<sub> +Unless you explicitly state otherwise, any contribution intentionally submitted +for inclusion in this crate by you, as defined in the Apache-2.0 license, shall +be dual licensed as above, without any additional terms or conditions. +</sub> + + +[compl_err]: https://doc.rust-lang.org/std/macro.compile_error.html +[`proc_macro::Diagnostic`]: https://doc.rust-lang.org/proc_macro/struct.Diagnostic.html + +[crate::dummy]: https://docs.rs/proc-macro-error2/1/proc_macro_error/dummy/index.html +[crate::multi]: https://docs.rs/proc-macro-error2/1/proc_macro_error/multi/index.html + +[`abort_call_site!`]: https://docs.rs/proc-macro-error2/1/proc_macro_error/macro.abort_call_site.html +[`abort!`]: https://docs.rs/proc-macro-error2/1/proc_macro_error/macro.abort.html +[guide]: https://docs.rs/proc-macro-error2 diff --git a/third_party/rust/proc-macro-error2/src/diagnostic.rs b/third_party/rust/proc-macro-error2/src/diagnostic.rs @@ -0,0 +1,360 @@ +use crate::{abort_now, check_correctness, sealed::Sealed, SpanRange}; +use proc_macro2::Span; +use proc_macro2::TokenStream; + +use quote::{quote_spanned, ToTokens}; + +/// Represents a diagnostic level +/// +/// # Warnings +/// +/// Warnings are ignored on stable/beta +#[derive(Debug, PartialEq)] +#[non_exhaustive] +pub enum Level { + Error, + Warning, +} + +/// Represents a single diagnostic message +#[derive(Debug)] +#[must_use = "A diagnostic does nothing unless emitted"] +pub struct Diagnostic { + pub(crate) level: Level, + pub(crate) span_range: SpanRange, + pub(crate) msg: String, + pub(crate) suggestions: Vec<(SuggestionKind, String, Option<SpanRange>)>, + pub(crate) children: Vec<(SpanRange, String)>, +} + +/// A collection of methods that do not exist in `proc_macro::Diagnostic` +/// but still useful to have around. +/// +/// This trait is sealed and cannot be implemented outside of `proc_macro_error`. +pub trait DiagnosticExt: Sealed { + /// Create a new diagnostic message that points to the `span_range`. + /// + /// This function is the same as `Diagnostic::spanned` but produces considerably + /// better error messages for multi-token spans on stable. + fn spanned_range(span_range: SpanRange, level: Level, message: String) -> Self; + + /// Add another error message to self such that it will be emitted right after + /// the main message. + /// + /// This function is the same as `Diagnostic::span_error` but produces considerably + /// better error messages for multi-token spans on stable. + #[must_use] + fn span_range_error(self, span_range: SpanRange, msg: String) -> Self; + + /// Attach a "help" note to your main message, the note will have it's own span on nightly. + /// + /// This function is the same as `Diagnostic::span_help` but produces considerably + /// better error messages for multi-token spans on stable. + /// + /// # Span + /// + /// The span is ignored on stable, the note effectively inherits its parent's (main message) span + #[must_use] + fn span_range_help(self, span_range: SpanRange, msg: String) -> Self; + + /// Attach a note to your main message, the note will have it's own span on nightly. + /// + /// This function is the same as `Diagnostic::span_note` but produces considerably + /// better error messages for multi-token spans on stable. + /// + /// # Span + /// + /// The span is ignored on stable, the note effectively inherits its parent's (main message) span + #[must_use] + fn span_range_note(self, span_range: SpanRange, msg: String) -> Self; +} + +impl DiagnosticExt for Diagnostic { + fn spanned_range(span_range: SpanRange, level: Level, message: String) -> Self { + Diagnostic { + level, + span_range, + msg: message, + suggestions: vec![], + children: vec![], + } + } + + fn span_range_error(mut self, span_range: SpanRange, msg: String) -> Self { + self.children.push((span_range, msg)); + self + } + + fn span_range_help(mut self, span_range: SpanRange, msg: String) -> Self { + self.suggestions + .push((SuggestionKind::Help, msg, Some(span_range))); + self + } + + fn span_range_note(mut self, span_range: SpanRange, msg: String) -> Self { + self.suggestions + .push((SuggestionKind::Note, msg, Some(span_range))); + self + } +} + +impl Diagnostic { + /// Create a new diagnostic message that points to `Span::call_site()` + pub fn new(level: Level, message: String) -> Self { + Diagnostic::spanned(Span::call_site(), level, message) + } + + /// Create a new diagnostic message that points to the `span` + pub fn spanned(span: Span, level: Level, message: String) -> Self { + Diagnostic::spanned_range( + SpanRange { + first: span, + last: span, + }, + level, + message, + ) + } + + /// Add another error message to self such that it will be emitted right after + /// the main message. + pub fn span_error(self, span: Span, msg: String) -> Self { + self.span_range_error( + SpanRange { + first: span, + last: span, + }, + msg, + ) + } + + /// Attach a "help" note to your main message, the note will have it's own span on nightly. + /// + /// # Span + /// + /// The span is ignored on stable, the note effectively inherits its parent's (main message) span + pub fn span_help(self, span: Span, msg: String) -> Self { + self.span_range_help( + SpanRange { + first: span, + last: span, + }, + msg, + ) + } + + /// Attach a "help" note to your main message. + pub fn help(mut self, msg: String) -> Self { + self.suggestions.push((SuggestionKind::Help, msg, None)); + self + } + + /// Attach a note to your main message, the note will have it's own span on nightly. + /// + /// # Span + /// + /// The span is ignored on stable, the note effectively inherits its parent's (main message) span + pub fn span_note(self, span: Span, msg: String) -> Self { + self.span_range_note( + SpanRange { + first: span, + last: span, + }, + msg, + ) + } + + /// Attach a note to your main message + pub fn note(mut self, msg: String) -> Self { + self.suggestions.push((SuggestionKind::Note, msg, None)); + self + } + + /// The message of main warning/error (no notes attached) + #[must_use] + pub fn message(&self) -> &str { + &self.msg + } + + /// Abort the proc-macro's execution and display the diagnostic. + /// + /// # Warnings + /// + /// Warnings are not emitted on stable and beta, but this function will abort anyway. + pub fn abort(self) -> ! { + self.emit(); + abort_now() + } + + /// Display the diagnostic while not aborting macro execution. + /// + /// # Warnings + /// + /// Warnings are ignored on stable/beta + pub fn emit(self) { + check_correctness(); + crate::imp::emit_diagnostic(self); + } +} + +/// **NOT PUBLIC API! NOTHING TO SEE HERE!!!** +#[doc(hidden)] +impl Diagnostic { + pub fn span_suggestion(self, span: Span, suggestion: &str, msg: String) -> Self { + match suggestion { + "help" | "hint" => self.span_help(span, msg), + _ => self.span_note(span, msg), + } + } + + pub fn suggestion(self, suggestion: &str, msg: String) -> Self { + match suggestion { + "help" | "hint" => self.help(msg), + _ => self.note(msg), + } + } +} + +impl ToTokens for Diagnostic { + fn to_tokens(&self, ts: &mut TokenStream) { + use std::borrow::Cow; + + fn ensure_lf(buf: &mut String, s: &str) { + if s.ends_with('\n') { + buf.push_str(s); + } else { + buf.push_str(s); + buf.push('\n'); + } + } + + fn diag_to_tokens( + span_range: SpanRange, + level: &Level, + msg: &str, + suggestions: &[(SuggestionKind, String, Option<SpanRange>)], + ) -> TokenStream { + if *level == Level::Warning { + return TokenStream::new(); + } + + let message = if suggestions.is_empty() { + Cow::Borrowed(msg) + } else { + let mut message = String::new(); + ensure_lf(&mut message, msg); + message.push('\n'); + + for (kind, note, _span) in suggestions { + message.push_str(" = "); + message.push_str(kind.name()); + message.push_str(": "); + ensure_lf(&mut message, note); + } + message.push('\n'); + + Cow::Owned(message) + }; + + let mut msg = proc_macro2::Literal::string(&message); + msg.set_span(span_range.last); + let group = quote_spanned!(span_range.last=> { #msg } ); + quote_spanned!(span_range.first=> compile_error!#group) + } + + ts.extend(diag_to_tokens( + self.span_range, + &self.level, + &self.msg, + &self.suggestions, + )); + ts.extend( + self.children + .iter() + .map(|(span_range, msg)| diag_to_tokens(*span_range, &Level::Error, msg, &[])), + ); + } +} + +#[derive(Debug)] +pub(crate) enum SuggestionKind { + Help, + Note, +} + +impl SuggestionKind { + fn name(&self) -> &'static str { + match self { + SuggestionKind::Note => "note", + SuggestionKind::Help => "help", + } + } +} + +#[cfg(feature = "syn-error")] +impl From<syn::Error> for Diagnostic { + fn from(err: syn::Error) -> Self { + use proc_macro2::{Delimiter, TokenTree}; + + fn gut_error(ts: &mut impl Iterator<Item = TokenTree>) -> Option<(SpanRange, String)> { + let start_span = ts.next()?.span(); + ts.next().expect(":1"); + ts.next().expect("core"); + ts.next().expect(":2"); + ts.next().expect(":3"); + ts.next().expect("compile_error"); + ts.next().expect("!"); + + let lit = match ts.next().unwrap() { + TokenTree::Group(group) => { + // Currently `syn` builds `compile_error!` invocations + // exclusively in `ident{"..."}` (braced) form which is not + // followed by `;` (semicolon). + // + // But if it changes to `ident("...");` (parenthesized) + // or `ident["..."];` (bracketed) form, + // we will need to skip the `;` as well. + // Highly unlikely, but better safe than sorry. + + if group.delimiter() == Delimiter::Parenthesis + || group.delimiter() == Delimiter::Bracket + { + ts.next().unwrap(); // ; + } + + match group.stream().into_iter().next().unwrap() { + TokenTree::Literal(lit) => lit, + _ => unreachable!(""), + } + } + _ => unreachable!(""), + }; + + let last = lit.span(); + let mut msg = lit.to_string(); + + // "abc" => abc + msg.pop(); + msg.remove(0); + + Some(( + SpanRange { + first: start_span, + last, + }, + msg, + )) + } + + let mut ts = err.to_compile_error().into_iter(); + + let (span_range, msg) = gut_error(&mut ts).unwrap(); + let mut res = Diagnostic::spanned_range(span_range, Level::Error, msg); + + while let Some((span_range, msg)) = gut_error(&mut ts) { + res = res.span_range_error(span_range, msg); + } + + res + } +} diff --git a/third_party/rust/proc-macro-error2/src/dummy.rs b/third_party/rust/proc-macro-error2/src/dummy.rs @@ -0,0 +1,151 @@ +//! Facility to emit dummy implementations (or whatever) in case +//! an error happen. +//! +//! `compile_error!` does not abort a compilation right away. This means +//! `rustc` doesn't just show you the error and abort, it carries on the +//! compilation process looking for other errors to report. +//! +//! Let's consider an example: +//! +//! ```rust,ignore +//! use proc_macro::TokenStream; +//! use proc_macro_error2::*; +//! +//! trait MyTrait { +//! fn do_thing(); +//! } +//! +//! // this proc macro is supposed to generate MyTrait impl +//! #[proc_macro_derive(MyTrait)] +//! #[proc_macro_error] +//! fn example(input: TokenStream) -> TokenStream { +//! // somewhere deep inside +//! abort!(span, "something's wrong"); +//! +//! // this implementation will be generated if no error happened +//! quote! { +//! impl MyTrait for #name { +//! fn do_thing() {/* whatever */} +//! } +//! } +//! } +//! +//! // ================ +//! // in main.rs +//! +//! // this derive triggers an error +//! #[derive(MyTrait)] // first BOOM! +//! struct Foo; +//! +//! fn main() { +//! Foo::do_thing(); // second BOOM! +//! } +//! ``` +//! +//! The problem is: the generated token stream contains only `compile_error!` +//! invocation, the impl was not generated. That means user will see two compilation +//! errors: +//! +//! ```text +//! error: something's wrong +//! --> $DIR/probe.rs:9:10 +//! | +//! 9 |#[proc_macro_derive(MyTrait)] +//! | ^^^^^^^ +//! +//! error[E0599]: no function or associated item named `do_thing` found for type `Foo` in the current scope +//! --> src\main.rs:3:10 +//! | +//! 1 | struct Foo; +//! | ----------- function or associated item `do_thing` not found for this +//! 2 | fn main() { +//! 3 | Foo::do_thing(); // second BOOM! +//! | ^^^^^^^^ function or associated item not found in `Foo` +//! ``` +//! +//! But the second error is meaningless! We definitely need to fix this. +//! +//! Most used approach in cases like this is "dummy implementation" - +//! omit `impl MyTrait for #name` and fill functions bodies with `unimplemented!()`. +//! +//! This is how you do it: +//! +//! ```rust,ignore +//! use proc_macro::TokenStream; +//! use proc_macro_error2::*; +//! +//! trait MyTrait { +//! fn do_thing(); +//! } +//! +//! // this proc macro is supposed to generate MyTrait impl +//! #[proc_macro_derive(MyTrait)] +//! #[proc_macro_error] +//! fn example(input: TokenStream) -> TokenStream { +//! // first of all - we set a dummy impl which will be appended to +//! // `compile_error!` invocations in case a trigger does happen +//! set_dummy(quote! { +//! impl MyTrait for #name { +//! fn do_thing() { unimplemented!() } +//! } +//! }); +//! +//! // somewhere deep inside +//! abort!(span, "something's wrong"); +//! +//! // this implementation will be generated if no error happened +//! quote! { +//! impl MyTrait for #name { +//! fn do_thing() {/* whatever */} +//! } +//! } +//! } +//! +//! // ================ +//! // in main.rs +//! +//! // this derive triggers an error +//! #[derive(MyTrait)] // first BOOM! +//! struct Foo; +//! +//! fn main() { +//! Foo::do_thing(); // no more errors! +//! } +//! ``` + +use proc_macro2::TokenStream; +use std::cell::RefCell; + +use crate::check_correctness; + +thread_local! { + static DUMMY_IMPL: RefCell<Option<TokenStream>> = const { RefCell::new(None) }; +} + +/// Sets dummy token stream which will be appended to `compile_error!(msg);...` +/// invocations in case you'll emit any errors. +/// +/// See [guide](../index.html#guide). +#[allow(clippy::must_use_candidate)] // Mutates thread local state +pub fn set_dummy(dummy: TokenStream) -> Option<TokenStream> { + check_correctness(); + DUMMY_IMPL.with(|old_dummy| old_dummy.replace(Some(dummy))) +} + +/// Same as [`set_dummy`] but, instead of resetting, appends tokens to the +/// existing dummy (if any). Behaves as `set_dummy` if no dummy is present. +pub fn append_dummy(dummy: TokenStream) { + check_correctness(); + DUMMY_IMPL.with(|old_dummy| { + let mut cell = old_dummy.borrow_mut(); + if let Some(ts) = cell.as_mut() { + ts.extend(dummy); + } else { + *cell = Some(dummy); + } + }); +} + +pub(crate) fn cleanup() -> Option<TokenStream> { + DUMMY_IMPL.with(|old_dummy| old_dummy.replace(None)) +} diff --git a/third_party/rust/proc-macro-error2/src/imp/delegate.rs b/third_party/rust/proc-macro-error2/src/imp/delegate.rs @@ -0,0 +1,68 @@ +//! This implementation uses [`proc_macro::Diagnostic`], nightly only. + +use std::cell::Cell; + +use proc_macro::{Diagnostic as PDiag, Level as PLevel}; + +use crate::{ + abort_now, check_correctness, + diagnostic::{Diagnostic, Level, SuggestionKind}, +}; + +pub fn abort_if_dirty() { + check_correctness(); + if IS_DIRTY.with(|c| c.get()) { + abort_now() + } +} + +pub(crate) fn cleanup() -> Vec<Diagnostic> { + IS_DIRTY.with(|c| c.set(false)); + vec![] +} + +pub(crate) fn emit_diagnostic(diag: Diagnostic) { + let Diagnostic { + level, + span_range, + msg, + suggestions, + children, + } = diag; + + let span = span_range.collapse().unwrap(); + + let level = match level { + Level::Warning => PLevel::Warning, + Level::Error => { + IS_DIRTY.with(|c| c.set(true)); + PLevel::Error + } + }; + + let mut res = PDiag::spanned(span, level, msg); + + for (kind, msg, span) in suggestions { + res = match (kind, span) { + (SuggestionKind::Note, Some(span_range)) => { + res.span_note(span_range.collapse().unwrap(), msg) + } + (SuggestionKind::Help, Some(span_range)) => { + res.span_help(span_range.collapse().unwrap(), msg) + } + (SuggestionKind::Note, None) => res.note(msg), + (SuggestionKind::Help, None) => res.help(msg), + } + } + + for (span_range, msg) in children { + let span = span_range.collapse().unwrap(); + res = res.span_error(span, msg); + } + + res.emit() +} + +thread_local! { + static IS_DIRTY: Cell<bool> = Cell::new(false); +} diff --git a/third_party/rust/proc-macro-error2/src/imp/fallback.rs b/third_party/rust/proc-macro-error2/src/imp/fallback.rs @@ -0,0 +1,30 @@ +//! This implementation uses self-written stable facilities. + +use crate::{ + abort_now, check_correctness, + diagnostic::{Diagnostic, Level}, +}; +use std::cell::RefCell; + +pub fn abort_if_dirty() { + check_correctness(); + ERR_STORAGE.with(|storage| { + if !storage.borrow().is_empty() { + abort_now() + } + }); +} + +pub(crate) fn cleanup() -> Vec<Diagnostic> { + ERR_STORAGE.with(|storage| storage.replace(Vec::new())) +} + +pub(crate) fn emit_diagnostic(diag: Diagnostic) { + if diag.level == Level::Error { + ERR_STORAGE.with(|storage| storage.borrow_mut().push(diag)); + } +} + +thread_local! { + static ERR_STORAGE: RefCell<Vec<Diagnostic>> = const { RefCell::new(Vec::new()) }; +} diff --git a/third_party/rust/proc-macro-error2/src/lib.rs b/third_party/rust/proc-macro-error2/src/lib.rs @@ -0,0 +1,565 @@ +//! # proc-macro-error2 +//! +//! This crate aims to make error reporting in proc-macros simple and easy to use. +//! Migrate from `panic!`-based errors for as little effort as possible! +//! +//! (Also, you can explicitly [append a dummy token stream](dummy/index.html) to your errors). +//! +//! To achieve his, this crate serves as a tiny shim around `proc_macro::Diagnostic` and +//! `compile_error!`. It detects the best way of emitting available based on compiler's version. +//! When the underlying diagnostic type is finally stabilized, this crate will simply be +//! delegating to it requiring no changes in your code! +//! +//! So you can just use this crate and have *both* some of `proc_macro::Diagnostic` functionality +//! available on stable ahead of time *and* your error-reporting code future-proof. +//! +//! ## Cargo features +//! +//! This crate provides *enabled by default* `syn-error` feature that gates +//! `impl From<syn::Error> for Diagnostic` conversion. If you don't use `syn` and want +//! to cut off some of compilation time, you can disable it via +//! +//! ```toml +//! [dependencies] +//! proc-macro-error2 = { version = "2.0.0", default-features = false } +//! ``` +//! +//! ***Please note that disabling this feature makes sense only if you don't depend on `syn` +//! directly or indirectly, and you very likely do.** +//! +//! ## Real world examples +//! +//! * [`structopt-derive`](https://github.com/TeXitoi/structopt/tree/master/structopt-derive) +//! (abort-like usage) +//! * [`auto-impl`](https://github.com/auto-impl-rs/auto_impl/) (emit-like usage) +//! +//! ## Limitations +//! +//! - Warnings are emitted only on nightly, they are ignored on stable. +//! - "help" suggestions can't have their own span info on stable, +//! (essentially inheriting the parent span). +//! - If a panic occurs somewhere in your macro no errors will be displayed. This is not a +//! technical limitation but rather intentional design. `panic` is not for error reporting. +//! +//! ### `#[proc_macro_error]` attribute +//! +//! **This attribute MUST be present on the top level of your macro** (the function +//! annotated with any of `#[proc_macro]`, `#[proc_macro_derive]`, `#[proc_macro_attribute]`). +//! +//! This attribute performs the setup and cleanup necessary to make things work. +//! +//! In most cases you'll need the simple `#[proc_macro_error]` form without any +//! additional settings. Feel free to [skip the "Syntax" section](#macros). +//! +//! #### Syntax +//! +//! `#[proc_macro_error]` or `#[proc_macro_error(settings...)]`, where `settings...` +//! is a comma-separated list of: +//! +//! - `proc_macro_hack`: +//! +//! In order to correctly cooperate with `#[proc_macro_hack]`, `#[proc_macro_error]` +//! attribute must be placed *before* (above) it, like this: +//! +//! ```no_run +//! # use proc_macro2::TokenStream; +//! # const IGNORE: &str = " +//! #[proc_macro_error] +//! #[proc_macro_hack] +//! #[proc_macro] +//! # "; +//! fn my_macro(input: TokenStream) -> TokenStream { +//! unimplemented!() +//! } +//! ``` +//! +//! If, for some reason, you can't place it like that you can use +//! `#[proc_macro_error(proc_macro_hack)]` instead. +//! +//! # Note +//! +//! If `proc-macro-hack` was detected (by any means) `allow_not_macro` +//! and `assert_unwind_safe` will be applied automatically. +//! +//! - `allow_not_macro`: +//! +//! By default, the attribute checks that it's applied to a proc-macro. +//! If none of `#[proc_macro]`, `#[proc_macro_derive]` nor `#[proc_macro_attribute]` are +//! present it will panic. It's the intention - this crate is supposed to be used only with +//! proc-macros. +//! +//! This setting is made to bypass the check, useful in certain circumstances. +//! +//! Pay attention: the function this attribute is applied to must return +//! `proc_macro::TokenStream`. +//! +//! This setting is implied if `proc-macro-hack` was detected. +//! +//! - `assert_unwind_safe`: +//! +//! By default, your code must be [unwind safe]. If your code is not unwind safe, +//! but you believe it's correct, you can use this setting to bypass the check. +//! You would need this for code that uses `lazy_static` or `thread_local` with +//! `Cell/RefCell` inside (and the like). +//! +//! This setting is implied if `#[proc_macro_error]` is applied to a function +//! marked as `#[proc_macro]`, `#[proc_macro_derive]` or `#[proc_macro_attribute]`. +//! +//! This setting is also implied if `proc-macro-hack` was detected. +//! +//! ## Macros +//! +//! Most of the time you want to use the macros. Syntax is described in the next section below. +//! +//! You'll need to decide how you want to emit errors: +//! +//! * Emit the error and abort. Very much panic-like usage. Served by [`abort!`] and +//! [`abort_call_site!`]. +//! * Emit the error but do not abort right away, looking for other errors to report. +//! Served by [`emit_error!`] and [`emit_call_site_error!`]. +//! +//! You **can** mix these usages. +//! +//! `abort` and `emit_error` take a "source span" as the first argument. This source +//! will be used to highlight the place the error originates from. It must be one of: +//! +//! * *Something* that implements [`ToTokens`] (most types in `syn` and `proc-macro2` do). +//! This source is the preferable one since it doesn't lose span information on multi-token +//! spans, see [this issue](https://gitlab.com/CreepySkeleton/proc-macro-error/-/issues/6) +//! for details. +//! * [`proc_macro::Span`] +//! * [`proc-macro2::Span`] +//! +//! The rest is your message in format-like style. +//! +//! See [the next section](#syntax-1) for detailed syntax. +//! +//! - [`abort!`]: +//! +//! Very much panic-like usage - abort right away and show the error. +//! Expands to [`!`] (never type). +//! +//! - [`abort_call_site!`]: +//! +//! Shortcut for `abort!(Span::call_site(), ...)`. Expands to [`!`] (never type). +//! +//! - [`emit_error!`]: +//! +//! [`proc_macro::Diagnostic`]-like usage - emit the error but keep going, +//! looking for other errors to report. +//! The compilation will fail nonetheless. Expands to [`()`] (unit type). +//! +//! - [`emit_call_site_error!`]: +//! +//! Shortcut for `emit_error!(Span::call_site(), ...)`. Expands to [`()`] (unit type). +//! +//! - [`emit_warning!`]: +//! +//! Like `emit_error!` but emit a warning instead of error. The compilation won't fail +//! because of warnings. +//! Expands to [`()`] (unit type). +//! +//! **Beware**: warnings are nightly only, they are completely ignored on stable. +//! +//! - [`emit_call_site_warning!`]: +//! +//! Shortcut for `emit_warning!(Span::call_site(), ...)`. Expands to [`()`] (unit type). +//! +//! - [`diagnostic`]: +//! +//! Build an instance of `Diagnostic` in format-like style. +//! +//! #### Syntax +//! +//! All the macros have pretty much the same syntax: +//! +//! 1. ```ignore +//! abort!(single_expr) +//! ``` +//! Shortcut for `Diagnostic::from(expr).abort()`. +//! +//! 2. ```ignore +//! abort!(span, message) +//! ``` +//! The first argument is an expression the span info should be taken from. +//! +//! The second argument is the error message, it must implement [`ToString`]. +//! +//! 3. ```ignore +//! abort!(span, format_literal, format_args...) +//! ``` +//! +//! This form is pretty much the same as 2, except `format!(format_literal, format_args...)` +//! will be used to for the message instead of [`ToString`]. +//! +//! That's it. `abort!`, `emit_warning`, `emit_error` share this exact syntax. +//! +//! `abort_call_site!`, `emit_call_site_warning`, `emit_call_site_error` lack 1 form +//! and do not take span in 2'th and 3'th forms. Those are essentially shortcuts for +//! `macro!(Span::call_site(), args...)`. +//! +//! `diagnostic!` requires a [`Level`] instance between `span` and second argument +//! (1'th form is the same). +//! +//! > **Important!** +//! > +//! > If you have some type from `proc_macro` or `syn` to point to, do not call `.span()` +//! > on it but rather use it directly: +//! > ```no_run +//! > # use proc_macro_error2::abort; +//! > # let input = proc_macro2::TokenStream::new(); +//! > let ty: syn::Type = syn::parse2(input).unwrap(); +//! > abort!(ty, "BOOM"); +//! > // ^^ <-- avoid .span() +//! > ``` +//! > +//! > `.span()` calls work too, but you may experience regressions in message quality. +//! +//! #### Note attachments +//! +//! 3. Every macro can have "note" attachments (only 2 and 3 form). +//! ```ignore +//! let opt_help = if have_some_info { Some("did you mean `this`?") } else { None }; +//! +//! abort!( +//! span, message; // <--- attachments start with `;` (semicolon) +//! +//! help = "format {} {}", "arg1", "arg2"; // <--- every attachment ends with `;`, +//! // maybe except the last one +//! +//! note = "to_string"; // <--- one arg uses `.to_string()` instead of `format!()` +//! +//! yay = "I see what {} did here", "you"; // <--- "help =" and "hint =" are mapped +//! // to Diagnostic::help, +//! // anything else is Diagnostic::note +//! +//! wow = note_span => "custom span"; // <--- attachments can have their own span +//! // it takes effect only on nightly though +//! +//! hint =? opt_help; // <-- "optional" attachment, get displayed only if `Some` +//! // must be single `Option` expression +//! +//! note =? note_span => opt_help // <-- optional attachments can have custom spans too +//! ); +//! ``` +//! + +//! ### Diagnostic type +//! +//! [`Diagnostic`] type is intentionally designed to be API compatible with [`proc_macro::Diagnostic`]. +//! Not all API is implemented, only the part that can be reasonably implemented on stable. +//! +//! +//! [`abort!`]: macro.abort.html +//! [`abort_call_site!`]: macro.abort_call_site.html +//! [`emit_warning!`]: macro.emit_warning.html +//! [`emit_error!`]: macro.emit_error.html +//! [`emit_call_site_warning!`]: macro.emit_call_site_error.html +//! [`emit_call_site_error!`]: macro.emit_call_site_warning.html +//! [`diagnostic!`]: macro.diagnostic.html +//! [`Diagnostic`]: struct.Diagnostic.html +//! +//! [`proc_macro::Span`]: https://doc.rust-lang.org/proc_macro/struct.Span.html +//! [`proc_macro::Diagnostic`]: https://doc.rust-lang.org/proc_macro/struct.Diagnostic.html +//! +//! [unwind safe]: https://doc.rust-lang.org/std/panic/trait.UnwindSafe.html#what-is-unwind-safety +//! [`!`]: https://doc.rust-lang.org/std/primitive.never.html +//! [`()`]: https://doc.rust-lang.org/std/primitive.unit.html +//! [`ToString`]: https://doc.rust-lang.org/std/string/trait.ToString.html +//! +//! [`proc-macro2::Span`]: https://docs.rs/proc-macro2/1.0.10/proc_macro2/struct.Span.html +//! [`ToTokens`]: https://docs.rs/quote/1.0.3/quote/trait.ToTokens.html +//! + +#![cfg_attr(feature = "nightly", feature(proc_macro_diagnostic))] +#![forbid(unsafe_code)] + +extern crate proc_macro; + +pub use crate::{ + diagnostic::{Diagnostic, DiagnosticExt, Level}, + dummy::{append_dummy, set_dummy}, +}; +pub use proc_macro_error_attr2::proc_macro_error; + +use proc_macro2::Span; +use quote::{quote, ToTokens}; + +use std::cell::Cell; +use std::panic::{catch_unwind, resume_unwind, UnwindSafe}; + +pub mod dummy; + +mod diagnostic; +mod macros; +mod sealed; + +#[cfg(not(feature = "nightly"))] +#[path = "imp/fallback.rs"] +mod imp; + +#[cfg(feature = "nightly")] +#[path = "imp/delegate.rs"] +mod imp; + +#[derive(Debug, Clone, Copy)] +#[must_use = "A SpanRange does nothing unless used"] +pub struct SpanRange { + pub first: Span, + pub last: Span, +} + +impl SpanRange { + /// Create a range with the `first` and `last` spans being the same. + pub fn single_span(span: Span) -> Self { + SpanRange { + first: span, + last: span, + } + } + + /// Create a `SpanRange` resolving at call site. + pub fn call_site() -> Self { + SpanRange::single_span(Span::call_site()) + } + + /// Construct span range from a `TokenStream`. This method always preserves all the + /// range. + /// + /// ### Note + /// + /// If the stream is empty, the result is `SpanRange::call_site()`. If the stream + /// consists of only one `TokenTree`, the result is `SpanRange::single_span(tt.span())` + /// that doesn't lose anything. + pub fn from_tokens(ts: &dyn ToTokens) -> Self { + let mut spans = ts.to_token_stream().into_iter().map(|tt| tt.span()); + let first = spans.next().unwrap_or_else(Span::call_site); + let last = spans.last().unwrap_or(first); + + SpanRange { first, last } + } + + /// Join two span ranges. The resulting range will start at `self.first` and end at + /// `other.last`. + pub fn join_range(self, other: SpanRange) -> Self { + SpanRange { + first: self.first, + last: other.last, + } + } + + /// Collapse the range into single span, preserving as much information as possible. + #[must_use] + pub fn collapse(self) -> Span { + self.first.join(self.last).unwrap_or(self.first) + } +} + +/// This traits expands `Result<T, Into<Diagnostic>>` with some handy shortcuts. +pub trait ResultExt { + type Ok; + + /// Behaves like `Result::unwrap`: if self is `Ok` yield the contained value, + /// otherwise abort macro execution via `abort!`. + fn unwrap_or_abort(self) -> Self::Ok; + + /// Behaves like `Result::expect`: if self is `Ok` yield the contained value, + /// otherwise abort macro execution via `abort!`. + /// If it aborts then resulting error message will be preceded with `message`. + fn expect_or_abort(self, msg: &str) -> Self::Ok; +} + +/// This traits expands `Option` with some handy shortcuts. +pub trait OptionExt { + type Some; + + /// Behaves like `Option::expect`: if self is `Some` yield the contained value, + /// otherwise abort macro execution via `abort_call_site!`. + /// If it aborts the `message` will be used for [`compile_error!`][compl_err] invocation. + /// + /// [compl_err]: https://doc.rust-lang.org/std/macro.compile_error.html + fn expect_or_abort(self, msg: &str) -> Self::Some; +} + +/// Abort macro execution and display all the emitted errors, if any. +/// +/// Does nothing if no errors were emitted (warnings do not count). +pub fn abort_if_dirty() { + imp::abort_if_dirty(); +} + +impl<T, E: Into<Diagnostic>> ResultExt for Result<T, E> { + type Ok = T; + + fn unwrap_or_abort(self) -> T { + match self { + Ok(res) => res, + Err(e) => e.into().abort(), + } + } + + fn expect_or_abort(self, message: &str) -> T { + match self { + Ok(res) => res, + Err(e) => { + let mut e = e.into(); + e.msg = format!("{}: {}", message, e.msg); + e.abort() + } + } + } +} + +impl<T> OptionExt for Option<T> { + type Some = T; + + fn expect_or_abort(self, message: &str) -> T { + match self { + Some(res) => res, + None => abort_call_site!(message), + } + } +} + +/// This is the entry point for a proc-macro. +/// +/// **NOT PUBLIC API, SUBJECT TO CHANGE WITHOUT ANY NOTICE** +#[doc(hidden)] +pub fn entry_point<F>(f: F, proc_macro_hack: bool) -> proc_macro::TokenStream +where + F: FnOnce() -> proc_macro::TokenStream + UnwindSafe, +{ + ENTERED_ENTRY_POINT.with(|flag| flag.set(flag.get() + 1)); + let caught = catch_unwind(f); + let dummy = dummy::cleanup(); + let err_storage = imp::cleanup(); + ENTERED_ENTRY_POINT.with(|flag| flag.set(flag.get() - 1)); + + let gen_error = || { + if proc_macro_hack { + quote! {{ + macro_rules! proc_macro_call { + () => ( unimplemented!() ) + } + + #(#err_storage)* + #dummy + + unimplemented!() + }} + } else { + quote!( #(#err_storage)* #dummy ) + } + }; + + match caught { + Ok(ts) => { + if err_storage.is_empty() { + ts + } else { + gen_error().into() + } + } + + Err(boxed) => match boxed.downcast::<AbortNow>() { + Ok(_) => gen_error().into(), + Err(boxed) => resume_unwind(boxed), + }, + } +} + +fn abort_now() -> ! { + check_correctness(); + std::panic::panic_any(AbortNow) +} + +thread_local! { + static ENTERED_ENTRY_POINT: Cell<usize> = const { Cell::new(0) }; +} + +struct AbortNow; + +fn check_correctness() { + assert!( + ENTERED_ENTRY_POINT.with(Cell::get) != 0, + "proc-macro-error2 API cannot be used outside of `entry_point` invocation, \ + perhaps you forgot to annotate your #[proc_macro] function with `#[proc_macro_error]" + ); +} + +/// **ALL THE STUFF INSIDE IS NOT PUBLIC API!!!** +#[doc(hidden)] +pub mod __export { + // reexports for use in macros + pub use proc_macro; + pub use proc_macro2; + + use proc_macro2::Span; + use quote::ToTokens; + + use crate::SpanRange; + + // inspired by + // https://github.com/dtolnay/case-studies/blob/master/autoref-specialization/README.md#simple-application + + pub trait SpanAsSpanRange { + #[allow(non_snake_case)] + fn FIRST_ARG_MUST_EITHER_BE_Span_OR_IMPLEMENT_ToTokens_OR_BE_SpanRange(&self) -> SpanRange; + } + + pub trait Span2AsSpanRange { + #[allow(non_snake_case)] + fn FIRST_ARG_MUST_EITHER_BE_Span_OR_IMPLEMENT_ToTokens_OR_BE_SpanRange(&self) -> SpanRange; + } + + pub trait ToTokensAsSpanRange { + #[allow(non_snake_case)] + fn FIRST_ARG_MUST_EITHER_BE_Span_OR_IMPLEMENT_ToTokens_OR_BE_SpanRange(&self) -> SpanRange; + } + + pub trait SpanRangeAsSpanRange { + #[allow(non_snake_case)] + fn FIRST_ARG_MUST_EITHER_BE_Span_OR_IMPLEMENT_ToTokens_OR_BE_SpanRange(&self) -> SpanRange; + } + + impl<T: ToTokens> ToTokensAsSpanRange for &T { + fn FIRST_ARG_MUST_EITHER_BE_Span_OR_IMPLEMENT_ToTokens_OR_BE_SpanRange(&self) -> SpanRange { + let mut ts = self.to_token_stream().into_iter(); + let first = match ts.next() { + Some(t) => t.span(), + None => Span::call_site(), + }; + + let last = match ts.last() { + Some(t) => t.span(), + None => first, + }; + + SpanRange { first, last } + } + } + + impl Span2AsSpanRange for Span { + fn FIRST_ARG_MUST_EITHER_BE_Span_OR_IMPLEMENT_ToTokens_OR_BE_SpanRange(&self) -> SpanRange { + SpanRange { + first: *self, + last: *self, + } + } + } + + impl SpanAsSpanRange for proc_macro::Span { + fn FIRST_ARG_MUST_EITHER_BE_Span_OR_IMPLEMENT_ToTokens_OR_BE_SpanRange(&self) -> SpanRange { + SpanRange { + first: (*self).into(), + last: (*self).into(), + } + } + } + + impl SpanRangeAsSpanRange for SpanRange { + fn FIRST_ARG_MUST_EITHER_BE_Span_OR_IMPLEMENT_ToTokens_OR_BE_SpanRange(&self) -> SpanRange { + *self + } + } +} diff --git a/third_party/rust/proc-macro-error2/src/macros.rs b/third_party/rust/proc-macro-error2/src/macros.rs @@ -0,0 +1,288 @@ +// FIXME: this can be greatly simplified via $()? +// as soon as MRSV hits 1.32 + +/// Build [`Diagnostic`](struct.Diagnostic.html) instance from provided arguments. +/// +/// # Syntax +/// +/// See [the guide](index.html#guide). +/// +#[macro_export] +macro_rules! diagnostic { + // from alias + ($err:expr) => { $crate::Diagnostic::from($err) }; + + // span, message, help + ($span:expr, $level:expr, $fmt:expr, $($args:expr),+ ; $($rest:tt)+) => {{ + #[allow(unused_imports)] + use $crate::__export::{ + ToTokensAsSpanRange, + Span2AsSpanRange, + SpanAsSpanRange, + SpanRangeAsSpanRange + }; + use $crate::DiagnosticExt; + let span_range = (&$span).FIRST_ARG_MUST_EITHER_BE_Span_OR_IMPLEMENT_ToTokens_OR_BE_SpanRange(); + + let diag = $crate::Diagnostic::spanned_range( + span_range, + $level, + format!($fmt, $($args),*) + ); + $crate::__pme__suggestions!(diag $($rest)*); + diag + }}; + + ($span:expr, $level:expr, $msg:expr ; $($rest:tt)+) => {{ + #[allow(unused_imports)] + use $crate::__export::{ + ToTokensAsSpanRange, + Span2AsSpanRange, + SpanAsSpanRange, + SpanRangeAsSpanRange + }; + use $crate::DiagnosticExt; + let span_range = (&$span).FIRST_ARG_MUST_EITHER_BE_Span_OR_IMPLEMENT_ToTokens_OR_BE_SpanRange(); + + let diag = $crate::Diagnostic::spanned_range(span_range, $level, $msg.to_string()); + $crate::__pme__suggestions!(diag $($rest)*); + diag + }}; + + // span, message, no help + ($span:expr, $level:expr, $fmt:expr, $($args:expr),+) => {{ + #[allow(unused_imports)] + use $crate::__export::{ + ToTokensAsSpanRange, + Span2AsSpanRange, + SpanAsSpanRange, + SpanRangeAsSpanRange + }; + use $crate::DiagnosticExt; + let span_range = (&$span).FIRST_ARG_MUST_EITHER_BE_Span_OR_IMPLEMENT_ToTokens_OR_BE_SpanRange(); + + $crate::Diagnostic::spanned_range( + span_range, + $level, + format!($fmt, $($args),*) + ) + }}; + + ($span:expr, $level:expr, $msg:expr) => {{ + #[allow(unused_imports)] + use $crate::__export::{ + ToTokensAsSpanRange, + Span2AsSpanRange, + SpanAsSpanRange, + SpanRangeAsSpanRange + }; + use $crate::DiagnosticExt; + let span_range = (&$span).FIRST_ARG_MUST_EITHER_BE_Span_OR_IMPLEMENT_ToTokens_OR_BE_SpanRange(); + + $crate::Diagnostic::spanned_range(span_range, $level, $msg.to_string()) + }}; + + + // trailing commas + + ($span:expr, $level:expr, $fmt:expr, $($args:expr),+, ; $($rest:tt)+) => { + $crate::diagnostic!($span, $level, $fmt, $($args),* ; $($rest)*) + }; + ($span:expr, $level:expr, $msg:expr, ; $($rest:tt)+) => { + $crate::diagnostic!($span, $level, $msg ; $($rest)*) + }; + ($span:expr, $level:expr, $fmt:expr, $($args:expr),+,) => { + $crate::diagnostic!($span, $level, $fmt, $($args),*) + }; + ($span:expr, $level:expr, $msg:expr,) => { + $crate::diagnostic!($span, $level, $msg) + }; + // ($err:expr,) => { $crate::diagnostic!($err) }; +} + +/// Abort proc-macro execution right now and display the error. +/// +/// # Syntax +/// +/// See [the guide](index.html#guide). +#[macro_export] +macro_rules! abort { + ($err:expr) => { + $crate::diagnostic!($err).abort() + }; + + ($span:expr, $($tts:tt)*) => { + $crate::diagnostic!($span, $crate::Level::Error, $($tts)*).abort() + }; +} + +/// Shortcut for `abort!(Span::call_site(), msg...)`. This macro +/// is still preferable over plain panic, panics are not for error reporting. +/// +/// # Syntax +/// +/// See [the guide](index.html#guide). +/// +#[macro_export] +macro_rules! abort_call_site { + ($($tts:tt)*) => { + $crate::abort!($crate::__export::proc_macro2::Span::call_site(), $($tts)*) + }; +} + +/// Emit an error while not aborting the proc-macro right away. +/// +/// # Syntax +/// +/// See [the guide](index.html#guide). +/// +#[macro_export] +macro_rules! emit_error { + ($err:expr) => { + $crate::diagnostic!($err).emit() + }; + + ($span:expr, $($tts:tt)*) => {{ + let level = $crate::Level::Error; + $crate::diagnostic!($span, level, $($tts)*).emit() + }}; +} + +/// Shortcut for `emit_error!(Span::call_site(), ...)`. This macro +/// is still preferable over plain panic, panics are not for error reporting.. +/// +/// # Syntax +/// +/// See [the guide](index.html#guide). +/// +#[macro_export] +macro_rules! emit_call_site_error { + ($($tts:tt)*) => { + $crate::emit_error!($crate::__export::proc_macro2::Span::call_site(), $($tts)*) + }; +} + +/// Emit a warning. Warnings are not errors and compilation won't fail because of them. +/// +/// **Does nothing on stable** +/// +/// # Syntax +/// +/// See [the guide](index.html#guide). +/// +#[macro_export] +macro_rules! emit_warning { + ($span:expr, $($tts:tt)*) => { + $crate::diagnostic!($span, $crate::Level::Warning, $($tts)*).emit() + }; +} + +/// Shortcut for `emit_warning!(Span::call_site(), ...)`. +/// +/// **Does nothing on stable** +/// +/// # Syntax +/// +/// See [the guide](index.html#guide). +/// +#[macro_export] +macro_rules! emit_call_site_warning { + ($($tts:tt)*) => {{ + $crate::emit_warning!($crate::__export::proc_macro2::Span::call_site(), $($tts)*) + }}; +} + +#[doc(hidden)] +#[macro_export] +macro_rules! __pme__suggestions { + ($var:ident) => (); + + ($var:ident $help:ident =? $msg:expr) => { + let $var = if let Some(msg) = $msg { + $var.suggestion(stringify!($help), msg.to_string()) + } else { + $var + }; + }; + ($var:ident $help:ident =? $span:expr => $msg:expr) => { + let $var = if let Some(msg) = $msg { + $var.span_suggestion($span.into(), stringify!($help), msg.to_string()) + } else { + $var + }; + }; + + ($var:ident $help:ident =? $msg:expr ; $($rest:tt)*) => { + $crate::__pme__suggestions!($var $help =? $msg); + $crate::__pme__suggestions!($var $($rest)*); + }; + ($var:ident $help:ident =? $span:expr => $msg:expr ; $($rest:tt)*) => { + $crate::__pme__suggestions!($var $help =? $span => $msg); + $crate::__pme__suggestions!($var $($rest)*); + }; + + + ($var:ident $help:ident = $msg:expr) => { + let $var = $var.suggestion(stringify!($help), $msg.to_string()); + }; + ($var:ident $help:ident = $fmt:expr, $($args:expr),+) => { + let $var = $var.suggestion( + stringify!($help), + format!($fmt, $($args),*) + ); + }; + ($var:ident $help:ident = $span:expr => $msg:expr) => { + let $var = $var.span_suggestion($span.into(), stringify!($help), $msg.to_string()); + }; + ($var:ident $help:ident = $span:expr => $fmt:expr, $($args:expr),+) => { + let $var = $var.span_suggestion( + $span.into(), + stringify!($help), + format!($fmt, $($args),*) + ); + }; + + ($var:ident $help:ident = $msg:expr ; $($rest:tt)*) => { + $crate::__pme__suggestions!($var $help = $msg); + $crate::__pme__suggestions!($var $($rest)*); + }; + ($var:ident $help:ident = $fmt:expr, $($args:expr),+ ; $($rest:tt)*) => { + $crate::__pme__suggestions!($var $help = $fmt, $($args),*); + $crate::__pme__suggestions!($var $($rest)*); + }; + ($var:ident $help:ident = $span:expr => $msg:expr ; $($rest:tt)*) => { + $crate::__pme__suggestions!($var $help = $span => $msg); + $crate::__pme__suggestions!($var $($rest)*); + }; + ($var:ident $help:ident = $span:expr => $fmt:expr, $($args:expr),+ ; $($rest:tt)*) => { + $crate::__pme__suggestions!($var $help = $span => $fmt, $($args),*); + $crate::__pme__suggestions!($var $($rest)*); + }; + + // trailing commas + + ($var:ident $help:ident = $msg:expr,) => { + $crate::__pme__suggestions!($var $help = $msg) + }; + ($var:ident $help:ident = $fmt:expr, $($args:expr),+,) => { + $crate::__pme__suggestions!($var $help = $fmt, $($args)*) + }; + ($var:ident $help:ident = $span:expr => $msg:expr,) => { + $crate::__pme__suggestions!($var $help = $span => $msg) + }; + ($var:ident $help:ident = $span:expr => $fmt:expr, $($args:expr),*,) => { + $crate::__pme__suggestions!($var $help = $span => $fmt, $($args)*) + }; + ($var:ident $help:ident = $msg:expr, ; $($rest:tt)*) => { + $crate::__pme__suggestions!($var $help = $msg; $($rest)*) + }; + ($var:ident $help:ident = $fmt:expr, $($args:expr),+, ; $($rest:tt)*) => { + $crate::__pme__suggestions!($var $help = $fmt, $($args),*; $($rest)*) + }; + ($var:ident $help:ident = $span:expr => $msg:expr, ; $($rest:tt)*) => { + $crate::__pme__suggestions!($var $help = $span => $msg; $($rest)*) + }; + ($var:ident $help:ident = $span:expr => $fmt:expr, $($args:expr),+, ; $($rest:tt)*) => { + $crate::__pme__suggestions!($var $help = $span => $fmt, $($args),*; $($rest)*) + }; +} diff --git a/third_party/rust/proc-macro-error2/src/sealed.rs b/third_party/rust/proc-macro-error2/src/sealed.rs @@ -0,0 +1,3 @@ +pub trait Sealed {} + +impl Sealed for crate::Diagnostic {} diff --git a/third_party/rust/proc-macro-error2/tests/macro-errors.rs b/third_party/rust/proc-macro-error2/tests/macro-errors.rs @@ -0,0 +1,6 @@ +#[test] +#[cfg(run_ui_tests)] +fn ui() { + let t = trybuild::TestCases::new(); + t.compile_fail("tests/ui/*.rs"); +} diff --git a/third_party/rust/proc-macro-error2/tests/ok.rs b/third_party/rust/proc-macro-error2/tests/ok.rs @@ -0,0 +1,8 @@ +use test_crate::*; + +ok!(it_works); + +#[test] +fn check_it_works() { + it_works(); +} diff --git a/third_party/rust/proc-macro-error2/tests/runtime-errors.rs b/third_party/rust/proc-macro-error2/tests/runtime-errors.rs @@ -0,0 +1,13 @@ +use proc_macro_error2::*; + +#[test] +#[should_panic = "proc-macro-error2 API cannot be used outside of"] +fn missing_attr_emit() { + emit_call_site_error!("You won't see me"); +} + +#[test] +#[should_panic = "proc-macro-error2 API cannot be used outside of"] +fn missing_attr_abort() { + abort_call_site!("You won't see me"); +} diff --git a/third_party/rust/proc-macro-error2/tests/ui/abort.rs b/third_party/rust/proc-macro-error2/tests/ui/abort.rs @@ -0,0 +1,10 @@ +use test_crate::*; + +abort_from!(one, two); +abort_to_string!(one, two); +abort_format!(one, two); +direct_abort!(one, two); +abort_notes!(one, two); +abort_call_site_test!(one, two); + +fn main() {} diff --git a/third_party/rust/proc-macro-error2/tests/ui/abort.stderr b/third_party/rust/proc-macro-error2/tests/ui/abort.stderr @@ -0,0 +1,48 @@ +error: abort!(span, from) test + --> tests/ui/abort.rs:3:13 + | +3 | abort_from!(one, two); + | ^^^ + +error: abort!(span, single_expr) test + --> tests/ui/abort.rs:4:18 + | +4 | abort_to_string!(one, two); + | ^^^ + +error: abort!(span, expr1, expr2) test + --> tests/ui/abort.rs:5:15 + | +5 | abort_format!(one, two); + | ^^^ + +error: Diagnostic::abort() test + --> tests/ui/abort.rs:6:15 + | +6 | direct_abort!(one, two); + | ^^^ + +error: This is an error + + = note: simple note + = help: simple help + = help: simple hint + = note: simple yay + = note: format note + = note: Some note + = note: spanned simple note + = note: spanned format note + = note: Some note + + --> tests/ui/abort.rs:7:14 + | +7 | abort_notes!(one, two); + | ^^^ + +error: abort_call_site! test + --> tests/ui/abort.rs:8:1 + | +8 | abort_call_site_test!(one, two); + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + | + = note: this error originates in the macro `abort_call_site_test` (in Nightly builds, run with -Z macro-backtrace for more info) diff --git a/third_party/rust/proc-macro-error2/tests/ui/append_dummy.rs b/third_party/rust/proc-macro-error2/tests/ui/append_dummy.rs @@ -0,0 +1,12 @@ +use test_crate::*; + +enum NeedDefault { + A, + B, +} + +append_dummy!(need_default); + +fn main() { + let _ = NeedDefault::default(); +} diff --git a/third_party/rust/proc-macro-error2/tests/ui/append_dummy.stderr b/third_party/rust/proc-macro-error2/tests/ui/append_dummy.stderr @@ -0,0 +1,5 @@ +error: append_dummy test + --> tests/ui/append_dummy.rs:8:15 + | +8 | append_dummy!(need_default); + | ^^^^^^^^^^^^ diff --git a/third_party/rust/proc-macro-error2/tests/ui/children_messages.rs b/third_party/rust/proc-macro-error2/tests/ui/children_messages.rs @@ -0,0 +1,5 @@ +use test_crate::*; + +children_messages!(one, two, three, four); + +fn main() {} diff --git a/third_party/rust/proc-macro-error2/tests/ui/children_messages.stderr b/third_party/rust/proc-macro-error2/tests/ui/children_messages.stderr @@ -0,0 +1,23 @@ +error: main macro message + --> tests/ui/children_messages.rs:3:20 + | +3 | children_messages!(one, two, three, four); + | ^^^ + +error: child message + --> tests/ui/children_messages.rs:3:25 + | +3 | children_messages!(one, two, three, four); + | ^^^ + +error: main syn::Error + --> tests/ui/children_messages.rs:3:30 + | +3 | children_messages!(one, two, three, four); + | ^^^^^ + +error: child syn::Error + --> tests/ui/children_messages.rs:3:37 + | +3 | children_messages!(one, two, three, four); + | ^^^^ diff --git a/third_party/rust/proc-macro-error2/tests/ui/dummy.rs b/third_party/rust/proc-macro-error2/tests/ui/dummy.rs @@ -0,0 +1,12 @@ +use test_crate::*; + +enum NeedDefault { + A, + B, +} + +dummy!(need_default); + +fn main() { + let _ = NeedDefault::default(); +} diff --git a/third_party/rust/proc-macro-error2/tests/ui/dummy.stderr b/third_party/rust/proc-macro-error2/tests/ui/dummy.stderr @@ -0,0 +1,5 @@ +error: set_dummy test + --> tests/ui/dummy.rs:8:8 + | +8 | dummy!(need_default); + | ^^^^^^^^^^^^ diff --git a/third_party/rust/proc-macro-error2/tests/ui/emit.rs b/third_party/rust/proc-macro-error2/tests/ui/emit.rs @@ -0,0 +1,6 @@ +use test_crate::*; + +emit!(one, two, three, four, five); +emit_notes!(one, two); + +fn main() {} diff --git a/third_party/rust/proc-macro-error2/tests/ui/emit.stderr b/third_party/rust/proc-macro-error2/tests/ui/emit.stderr @@ -0,0 +1,48 @@ +error: emit!(span, from) test + --> tests/ui/emit.rs:3:7 + | +3 | emit!(one, two, three, four, five); + | ^^^ + +error: emit!(span, expr1, expr2) test + --> tests/ui/emit.rs:3:12 + | +3 | emit!(one, two, three, four, five); + | ^^^ + +error: emit!(span, single_expr) test + --> tests/ui/emit.rs:3:17 + | +3 | emit!(one, two, three, four, five); + | ^^^^^ + +error: Diagnostic::emit() test + --> tests/ui/emit.rs:3:24 + | +3 | emit!(one, two, three, four, five); + | ^^^^ + +error: emit_call_site_error!(expr) test + --> tests/ui/emit.rs:3:1 + | +3 | emit!(one, two, three, four, five); + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + | + = note: this error originates in the macro `emit` (in Nightly builds, run with -Z macro-backtrace for more info) + +error: This is an error + + = note: simple note + = help: simple help + = help: simple hint + = note: simple yay + = note: format note + = note: Some note + = note: spanned simple note + = note: spanned format note + = note: Some note + + --> tests/ui/emit.rs:4:13 + | +4 | emit_notes!(one, two); + | ^^^ diff --git a/third_party/rust/proc-macro-error2/tests/ui/explicit_span_range.rs b/third_party/rust/proc-macro-error2/tests/ui/explicit_span_range.rs @@ -0,0 +1,5 @@ +use test_crate::*; + +explicit_span_range!(one, two, three, four); + +fn main() {} diff --git a/third_party/rust/proc-macro-error2/tests/ui/explicit_span_range.stderr b/third_party/rust/proc-macro-error2/tests/ui/explicit_span_range.stderr @@ -0,0 +1,5 @@ +error: explicit SpanRange + --> tests/ui/explicit_span_range.rs:3:22 + | +3 | explicit_span_range!(one, two, three, four); + | ^^^^^^^^^^^^^^^ diff --git a/third_party/rust/proc-macro-error2/tests/ui/misuse.rs b/third_party/rust/proc-macro-error2/tests/ui/misuse.rs @@ -0,0 +1,10 @@ +use proc_macro_error2::abort; + +struct Foo; + +#[allow(unused)] +fn foo() { + abort!(Foo, "BOOM"); +} + +fn main() {} diff --git a/third_party/rust/proc-macro-error2/tests/ui/misuse.stderr b/third_party/rust/proc-macro-error2/tests/ui/misuse.stderr @@ -0,0 +1,24 @@ +error[E0599]: the method `FIRST_ARG_MUST_EITHER_BE_Span_OR_IMPLEMENT_ToTokens_OR_BE_SpanRange` exists for reference `&Foo`, but its trait bounds were not satisfied + --> tests/ui/misuse.rs:7:5 + | +3 | struct Foo; + | ---------- doesn't satisfy `Foo: quote::to_tokens::ToTokens` +... +7 | abort!(Foo, "BOOM"); + | ^^^^^^^^^^^^^^^^^^^ method cannot be called on `&Foo` due to unsatisfied trait bounds + | + = note: the following trait bounds were not satisfied: + `Foo: quote::to_tokens::ToTokens` + which is required by `&Foo: ToTokensAsSpanRange` +note: the trait `quote::to_tokens::ToTokens` must be implemented + --> $CARGO/quote-1.0.37/src/to_tokens.rs + | + | pub trait ToTokens { + | ^^^^^^^^^^^^^^^^^^ + = help: items from traits can only be used if the trait is implemented and in scope + = note: the following traits define an item `FIRST_ARG_MUST_EITHER_BE_Span_OR_IMPLEMENT_ToTokens_OR_BE_SpanRange`, perhaps you need to implement one of them: + candidate #1: `Span2AsSpanRange` + candidate #2: `SpanAsSpanRange` + candidate #3: `SpanRangeAsSpanRange` + candidate #4: `ToTokensAsSpanRange` + = note: this error originates in the macro `$crate::diagnostic` which comes from the expansion of the macro `abort` (in Nightly builds, run with -Z macro-backtrace for more info) diff --git a/third_party/rust/proc-macro-error2/tests/ui/multiple_tokens.rs b/third_party/rust/proc-macro-error2/tests/ui/multiple_tokens.rs @@ -0,0 +1,4 @@ +#[test_crate::multiple_tokens] +type T = (); + +fn main() {} diff --git a/third_party/rust/proc-macro-error2/tests/ui/multiple_tokens.stderr b/third_party/rust/proc-macro-error2/tests/ui/multiple_tokens.stderr @@ -0,0 +1,5 @@ +error: ... + --> tests/ui/multiple_tokens.rs:2:1 + | +2 | type T = (); + | ^^^^^^^^^^^^ diff --git a/third_party/rust/proc-macro-error2/tests/ui/not_proc_macro.rs b/third_party/rust/proc-macro-error2/tests/ui/not_proc_macro.rs @@ -0,0 +1,4 @@ +use proc_macro_error2::proc_macro_error; + +#[proc_macro_error] +fn main() {} diff --git a/third_party/rust/proc-macro-error2/tests/ui/not_proc_macro.stderr b/third_party/rust/proc-macro-error2/tests/ui/not_proc_macro.stderr @@ -0,0 +1,9 @@ +error: #[proc_macro_error] attribute can be used only with procedural macros + + = hint: if you are really sure that #[proc_macro_error] should be applied to this exact function, use #[proc_macro_error(allow_not_macro)] + --> tests/ui/not_proc_macro.rs:3:1 + | +3 | #[proc_macro_error] + | ^^^^^^^^^^^^^^^^^^^ + | + = note: this error originates in the attribute macro `proc_macro_error` (in Nightly builds, run with -Z macro-backtrace for more info) diff --git a/third_party/rust/proc-macro-error2/tests/ui/option_ext.rs b/third_party/rust/proc-macro-error2/tests/ui/option_ext.rs @@ -0,0 +1,5 @@ +use test_crate::*; + +option_ext!(one, two); + +fn main() {} diff --git a/third_party/rust/proc-macro-error2/tests/ui/option_ext.stderr b/third_party/rust/proc-macro-error2/tests/ui/option_ext.stderr @@ -0,0 +1,7 @@ +error: Option::expect_or_abort() test + --> tests/ui/option_ext.rs:3:1 + | +3 | option_ext!(one, two); + | ^^^^^^^^^^^^^^^^^^^^^ + | + = note: this error originates in the macro `option_ext` (in Nightly builds, run with -Z macro-backtrace for more info) diff --git a/third_party/rust/proc-macro-error2/tests/ui/result_ext.rs b/third_party/rust/proc-macro-error2/tests/ui/result_ext.rs @@ -0,0 +1,6 @@ +use test_crate::*; + +result_unwrap_or_abort!(one, two); +result_expect_or_abort!(one, two); + +fn main() {} diff --git a/third_party/rust/proc-macro-error2/tests/ui/result_ext.stderr b/third_party/rust/proc-macro-error2/tests/ui/result_ext.stderr @@ -0,0 +1,11 @@ +error: Result::unwrap_or_abort() test + --> tests/ui/result_ext.rs:3:25 + | +3 | result_unwrap_or_abort!(one, two); + | ^^^ + +error: BOOM: Result::expect_or_abort() test + --> tests/ui/result_ext.rs:4:25 + | +4 | result_expect_or_abort!(one, two); + | ^^^ diff --git a/third_party/rust/proc-macro-error2/tests/ui/to_tokens_span.rs b/third_party/rust/proc-macro-error2/tests/ui/to_tokens_span.rs @@ -0,0 +1,5 @@ +use test_crate::*; + +to_tokens_span!(std::option::Option); + +fn main() {} diff --git a/third_party/rust/proc-macro-error2/tests/ui/to_tokens_span.stderr b/third_party/rust/proc-macro-error2/tests/ui/to_tokens_span.stderr @@ -0,0 +1,11 @@ +error: whole type + --> tests/ui/to_tokens_span.rs:3:17 + | +3 | to_tokens_span!(std::option::Option); + | ^^^^^^^^^^^^^^^^^^^ + +error: explicit .span() + --> tests/ui/to_tokens_span.rs:3:17 + | +3 | to_tokens_span!(std::option::Option); + | ^^^ diff --git a/third_party/rust/proc-macro-error2/tests/ui/unknown_setting.rs b/third_party/rust/proc-macro-error2/tests/ui/unknown_setting.rs @@ -0,0 +1,4 @@ +use proc_macro_error2::proc_macro_error; + +#[proc_macro_error(allow_not_macro, assert_unwind_safe, trololo)] +fn main() {} diff --git a/third_party/rust/proc-macro-error2/tests/ui/unknown_setting.stderr b/third_party/rust/proc-macro-error2/tests/ui/unknown_setting.stderr @@ -0,0 +1,5 @@ +error: unknown setting `trololo`, expected one of `assert_unwind_safe`, `allow_not_macro`, `proc_macro_hack` + --> tests/ui/unknown_setting.rs:3:57 + | +3 | #[proc_macro_error(allow_not_macro, assert_unwind_safe, trololo)] + | ^^^^^^^ diff --git a/third_party/rust/proc-macro-error2/tests/ui/unrelated_panic.rs b/third_party/rust/proc-macro-error2/tests/ui/unrelated_panic.rs @@ -0,0 +1,5 @@ +use test_crate::*; + +unrelated_panic!(); + +fn main() {} diff --git a/third_party/rust/proc-macro-error2/tests/ui/unrelated_panic.stderr b/third_party/rust/proc-macro-error2/tests/ui/unrelated_panic.stderr @@ -0,0 +1,7 @@ +error: proc macro panicked + --> tests/ui/unrelated_panic.rs:3:1 + | +3 | unrelated_panic!(); + | ^^^^^^^^^^^^^^^^^^ + | + = help: message: unrelated panic test diff --git a/toolkit/content/license.html b/toolkit/content/license.html @@ -2286,6 +2286,8 @@ product. <li><code>third_party/rust/uritemplate-next/</code></li> #ifdef MOZ_JXL <li><code>third_party/jpeg-xl/</code></li> + <li><code>third_party/rust/jxl</code></li> + <li><code>third_party/rust/jxl_macros</code></li> #endif <li><code>third_party/xsimd/</code></li> <li><code>third_party/zstd/</code></li> @@ -4160,8 +4162,14 @@ product. <li><code>devtools/client/netmonitor/src/components/messages/parsers/stomp/byte.js</code>, <code>devtools/client/netmonitor/src/components/messages/parsers/stomp/frame.js</code> and <code>devtools/client/netmonitor/src/components/messages/parsers/stomp/parser.js</code></li> + <li><code>third_party/rust/proc-macro2</code></li> <li><code>third_party/rust/synstructure</code></li> <li><code>third_party/rust/void</code></li> +#ifdef MOZ_JXL + <li><code>third_party/rust/array-init</code></li> + <li><code>third_party/rust/proc-macro-error-attr2</code> and + <code>third_party/rust/proc-macro-error2</code></li> +#endif <li><code>js/src/zydis</code> (unless otherwise specified)</li> <li><code>js/src/vm/Float16.h</code>(the code contained in the half namespace)</li> <li><code>js/src/xsum</code></li> diff --git a/toolkit/library/rust/shared/Cargo.toml b/toolkit/library/rust/shared/Cargo.toml @@ -13,6 +13,7 @@ mozglue-static = { path = "../../../../mozglue/static/rust" } geckoservo = { path = "../../../../servo/ports/geckolib" } kvstore = { path = "../../../components/kvstore" } mp4parse_capi = { git = "https://github.com/mozilla/mp4parse-rust", rev = "f955be5d2a04a631c0f1777d6f35370ea1a99e2d", features = ["missing-pixi-permitted"] } +image_jxl = { path = "../../../../image/rust/jxl", optional = true } nserror = { path = "../../../../xpcom/rust/nserror" } nsstring = { path = "../../../../xpcom/rust/nsstring" } netwerk_helper = { path = "../../../../netwerk/base/rust-helper" }