LibGfx/JPEG2000: Implement tag trees

A tag tree is a data structure used for deserializing JPEG2000
packet headers.

We don't use them for anything yet, except from tests.

The implementation feels a bit awkward to me, but we can always polish
it later.

The spec thankfully includes two concrete examples. The code is
correct enough to pass those -- I added them as test.

Tests/LibGfx/TestImageDecoder.cpp

@@ -606,6 +606,59 @@ TEST_CASE(test_jpeg2000_gray)
     EXPECT_EQ(icc_bytes->size(), 912u);
 }
 
+TEST_CASE(test_jpeg2000_tag_tree)
+{
+    {
+        // The example from the NOTE at the end of B.10.2 Tag trees:
+        auto tree = TRY_OR_FAIL(Gfx::JPEG2000::TagTree::create(6, 3));
+        auto bits = to_array<u8>({
+            0, 1, 1, 1, 1, // q3(0, 0)
+            0, 0, 1,       // q3(1, 0)
+            1, 0, 1,       // q3(2, 0)
+        });
+        size_t index = 0;
+        Function<ErrorOr<bool>()> read_bit = [&]() -> bool {
+            return bits[index++];
+        };
+        EXPECT_EQ(1u, MUST(tree.read_value(0, 0, read_bit)));
+        EXPECT_EQ(index, 5u);
+        EXPECT_EQ(3u, MUST(tree.read_value(1, 0, read_bit)));
+        EXPECT_EQ(index, 8u);
+        EXPECT_EQ(2u, MUST(tree.read_value(2, 0, read_bit)));
+        EXPECT_EQ(index, 11u);
+    }
+
+    {
+        // The inclusion tag tree bits from Table B.5 – Example packet header bit stream.
+        auto tree = TRY_OR_FAIL(Gfx::JPEG2000::TagTree::create(3, 2));
+        auto bits = to_array<u8>({
+            1, 1, 1, // Code-block 0, 0 included for the first time (partial inclusion tag tree)
+            1,       // Code-block 1, 0 included for the first time (partial inclusion tag tree)
+            0,       // Code-block 2, 0 not yet included (partial tag tree)
+            0,       // Code-block 0, 1 not yet included
+            0,       // Code-block 1, 2 not yet included
+            // Code-block 2, 1 not yet included (no data needed, already conveyed by partial tag tree for code-block 2, 0)
+        });
+        size_t index = 0;
+        Function<ErrorOr<bool>()> read_bit = [&]() -> bool {
+            return bits[index++];
+        };
+        u32 next_layer = 1;
+        EXPECT_EQ(0u, MUST(tree.read_value(0, 0, read_bit, next_layer)));
+        EXPECT_EQ(index, 3u);
+        EXPECT_EQ(0u, MUST(tree.read_value(1, 0, read_bit, next_layer)));
+        EXPECT_EQ(index, 4u);
+        EXPECT_EQ(1u, MUST(tree.read_value(2, 0, read_bit, next_layer)));
+        EXPECT_EQ(index, 5u);
+        EXPECT_EQ(1u, MUST(tree.read_value(0, 1, read_bit, next_layer)));
+        EXPECT_EQ(index, 6u);
+        EXPECT_EQ(1u, MUST(tree.read_value(1, 1, read_bit, next_layer)));
+        EXPECT_EQ(index, 7u);
+        EXPECT_EQ(1u, MUST(tree.read_value(2, 1, read_bit, next_layer)));
+        EXPECT_EQ(index, 7u); // Didn't change!
+    }
+}
+
 TEST_CASE(test_pam_rgb)
 {
     auto file = TRY_OR_FAIL(Core::MappedFile::map(TEST_INPUT("pnm/2x1.pam"sv)));

Userland/Libraries/LibGfx/ImageFormats/JPEG2000Loader.cpp

@@ -804,6 +804,104 @@ static ErrorOr<void> decode_jpeg2000_header(JPEG2000LoadingContext& context, Rea
     return {};
 }
 
+namespace JPEG2000 {
+
+// Tag trees are used to store the code-block inclusion bits and the zero bit-plane information.
+// B.10.2 Tag trees
+// "At every node of this tree the minimum integer of the (up to four) nodes below it is recorded. [...]
+//  Level 0 is the lowest level of the tag tree; it contains the top node. [...]
+//  Each node has a [...] current value, [...] initialized to zero. A 0 bit in the tag tree means that the minimum
+//  (or the value in the case of the highest level) is larger than the current value and a 1 bit means that the minimum
+//  (or the value in the case of the highest level) is equal to the current value.
+//  For each contiguous 0 bit in the tag tree the current value is incremented by one.
+//  Nodes at higher levels cannot be coded until lower level node values are fixed (i.e, a 1 bit is coded). [...]
+//  Only the information needed for the current code-block is stored at the current point in the packet header."
+// The example in Figure B.13 / Table B.5 is useful to understand what exactly "only the information needed" means.
+struct TagTreeNode {
+    u32 value { 0 };
+    enum State {
+        Pending,
+        Final,
+    };
+    State state { Pending };
+    Array<OwnPtr<TagTreeNode>, 4> children {};
+    u32 level { 0 }; // 0 for leaf nodes, 1 for the next level, etc.
+
+    bool is_leaf() const { return level == 0; }
+
+    ErrorOr<u32> read_value(u32 x, u32 y, Function<ErrorOr<bool>()> const& read_bit, u32 start_value, Optional<u32> stop_at = {})
+    {
+        value = max(value, start_value);
+        while (true) {
+            if (stop_at.has_value() && value == stop_at.value())
+                return value;
+
+            if (state == Final) {
+                if (is_leaf())
+                    return value;
+                u32 x_index = (x >> (level - 1)) & 1;
+                u32 y_index = (y >> (level - 1)) & 1;
+                return children[y_index * 2 + x_index]->read_value(x, y, read_bit, value, stop_at);
+            }
+
+            bool bit = TRY(read_bit());
+            if (!bit)
+                value++;
+            else
+                state = Final;
+        }
+    }
+
+    static ErrorOr<NonnullOwnPtr<TagTreeNode>> create(u32 x_count, u32 y_count, u32 level)
+    {
+        VERIFY(x_count > 0);
+        VERIFY(y_count > 0);
+
+        auto node = TRY(try_make<TagTreeNode>());
+        node->level = level;
+        if (node->is_leaf()) {
+            VERIFY(x_count == 1);
+            VERIFY(y_count == 1);
+            return node;
+        }
+
+        VERIFY(x_count > 1 || y_count > 1);
+        u32 top_left_x_child_count = min(x_count, 1u << (max(level, 1) - 1));
+        u32 top_left_y_child_count = min(y_count, 1u << (max(level, 1) - 1));
+        for (u32 y = 0; y < 2; ++y) {
+            for (u32 x = 0; x < 2; ++x) {
+                u32 child_x_count = x == 1 ? x_count - top_left_x_child_count : top_left_x_child_count;
+                u32 child_y_count = y == 1 ? y_count - top_left_y_child_count : top_left_y_child_count;
+                if (child_x_count == 0 || child_y_count == 0)
+                    continue;
+                node->children[y * 2 + x] = TRY(create(child_x_count, child_y_count, level - 1));
+            }
+        }
+        return node;
+    }
+};
+
+TagTree::TagTree(NonnullOwnPtr<TagTreeNode> root)
+    : m_root(move(root))
+{
+}
+
+TagTree::TagTree(TagTree&&) = default;
+TagTree::~TagTree() = default;
+
+ErrorOr<TagTree> TagTree::create(u32 x_count, u32 y_count)
+{
+    auto level = ceil(log2(max(x_count, y_count)));
+    return TagTree { TRY(TagTreeNode::create(x_count, y_count, level)) };
+}
+
+ErrorOr<u32> TagTree::read_value(u32 x, u32 y, Function<ErrorOr<bool>()> const& read_bit, Optional<u32> stop_at) const
+{
+    return m_root->read_value(x, y, read_bit, m_root->value, stop_at);
+}
+
+}
+
 bool JPEG2000ImageDecoderPlugin::sniff(ReadonlyBytes data)
 {
     return data.starts_with(jp2_id_string);

Userland/Libraries/LibGfx/ImageFormats/JPEG2000Loader.h

@@ -10,6 +10,26 @@
 
 namespace Gfx {
 
+namespace JPEG2000 {
+
+struct TagTreeNode;
+class TagTree {
+public:
+    TagTree(TagTree&&);
+    ~TagTree();
+
+    static ErrorOr<TagTree> create(u32 x_count, u32 y_count);
+
+    ErrorOr<u32> read_value(u32 x, u32 y, Function<ErrorOr<bool>()> const& read_bit, Optional<u32> stop_at = {}) const;
+
+private:
+    TagTree(NonnullOwnPtr<TagTreeNode>);
+
+    NonnullOwnPtr<TagTreeNode> m_root;
+};
+
+}
+
 struct JPEG2000LoadingContext;
 
 class JPEG2000ImageDecoderPlugin : public ImageDecoderPlugin {