LibGfx/JPEG2000: Implement tag trees

A tag tree is a data structure used for deserialiing 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,61 @@ 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;
+        auto make_read_bit = [&]() {
+            return [&]() -> bool { return bits[index++]; };
+        };
+        EXPECT_EQ(1u, MUST(tree.read_value(0, 0, make_read_bit())));
+        EXPECT_EQ(index, 5u);
+        EXPECT_EQ(3u, MUST(tree.read_value(1, 0, make_read_bit())));
+        EXPECT_EQ(index, 8u);
+        EXPECT_EQ(2u, MUST(tree.read_value(2, 0, make_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;
+        auto make_read_bit = [&]() {
+            return [&]() -> bool {
+                return bits[index++];
+            };
+        };
+        u32 next_layer = 1;
+        EXPECT_EQ(0u, MUST(tree.read_value(0, 0, make_read_bit(), next_layer)));
+        EXPECT_EQ(index, 3u);
+        EXPECT_EQ(0u, MUST(tree.read_value(1, 0, make_read_bit(), next_layer)));
+        EXPECT_EQ(index, 4u);
+        EXPECT_EQ(1u, MUST(tree.read_value(2, 0, make_read_bit(), next_layer)));
+        EXPECT_EQ(index, 5u);
+        EXPECT_EQ(1u, MUST(tree.read_value(0, 1, make_read_bit(), next_layer)));
+        EXPECT_EQ(index, 6u);
+        EXPECT_EQ(1u, MUST(tree.read_value(1, 1, make_read_bit(), next_layer)));
+        EXPECT_EQ(index, 7u);
+        EXPECT_EQ(1u, MUST(tree.read_value(2, 1, make_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,113 @@ 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<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>()> 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, move(read_bit), value, stop_at);
+            }
+
+            bool bit = TRY(read_bit());
+            if (!bit)
+                value++;
+            else
+                state = Final;
+        }
+    }
+};
+
+TagTree::TagTree() = default;
+TagTree::TagTree(TagTree const&) = default;
+TagTree::~TagTree() = default;
+
+ErrorOr<TagTree> TagTree::create(u32 x_count, u32 y_count)
+{
+    // In leaf-to-root order.
+    Vector<u32> nodes_in_level;
+    Vector<u32> width_in_level;
+    Vector<u32> height_in_level;
+    size_t number_of_nodes = 0;
+    do {
+        nodes_in_level.append(x_count * y_count);
+        width_in_level.append(x_count);
+        height_in_level.append(y_count);
+        number_of_nodes += nodes_in_level.last();
+        x_count = max(ceil_div(x_count, 2), 1);
+        y_count = max(ceil_div(y_count, 2), 1);
+    } while (width_in_level.last() > 1 || height_in_level.last() > 1);
+
+    // In root-to-leaf order.
+    Vector<TagTreeNode> nodes;
+    TRY(nodes.try_resize(number_of_nodes));
+    u32 level_offset = 0, node_offset = 0;
+    for (ssize_t level = nodes_in_level.size() - 1; level >= 0; --level) {
+        level_offset += nodes_in_level[level];
+        for (u32 y = 0; y < height_in_level[level]; ++y) {
+            for (u32 x = 0; x < width_in_level[level]; ++x, ++node_offset) {
+                nodes[node_offset].level = level;
+                if (level == 0)
+                    continue;
+                for (u32 y_offset = 0; y_offset < 2; ++y_offset) {
+                    for (u32 x_offset = 0; x_offset < 2; ++x_offset) {
+                        u32 child_x = x * 2 + x_offset;
+                        u32 child_y = y * 2 + y_offset;
+                        if (child_x >= width_in_level[level - 1] || child_y >= height_in_level[level - 1])
+                            continue;
+                        u32 child_index = (y_offset << 1) | x_offset;
+                        nodes[node_offset].children[child_index] = &nodes[level_offset + child_y * width_in_level[level - 1] + child_x];
+                    }
+                }
+            }
+        }
+    }
+
+    TagTree tree;
+    tree.nodes = move(nodes);
+    tree.root = &tree.nodes[0];
+    return tree;
+}
+
+ErrorOr<u32> TagTree::read_value(u32 x, u32 y, Function<ErrorOr<bool>()> read_bit, Optional<u32> stop_at) const
+{
+    return root->read_value(x, y, move(read_bit), 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 const&);
+    ~TagTree();
+
+    static ErrorOr<TagTree> create(u32 x_count, u32 y_count);
+
+    ErrorOr<u32> read_value(u32 x, u32 y, Function<ErrorOr<bool>()> read_bit, Optional<u32> stop_at = {}) const;
+
+private:
+    Vector<TagTreeNode> nodes;
+    TagTreeNode* root { nullptr };
+};
+
+}
+
 struct JPEG2000LoadingContext;
 
 class JPEG2000ImageDecoderPlugin : public ImageDecoderPlugin {