5 files changed, 138 insertions, 12 deletions

diff --git a/Libraries/LibDraw/ImageLoader.cpp b/Libraries/LibDraw/ImageLoader.cpp
new file mode 100644
index 0000000000..2cfdc0ce11
--- /dev/null
+++ b/Libraries/LibDraw/ImageLoader.cpp
@@ -0,0 +1,11 @@
+#include <LibDraw/ImageLoader.h>
+#include <LibDraw/PNGLoader.h>
+
+ImageLoader::ImageLoader(const u8* data, size_t size)
+{
+    m_plugin = make<PNGImageLoaderPlugin>(data, size);
+}
+
+ImageLoader::~ImageLoader()
+{
+}
diff --git a/Libraries/LibDraw/ImageLoader.h b/Libraries/LibDraw/ImageLoader.h
new file mode 100644
index 0000000000..3b72b5cba7
--- /dev/null
+++ b/Libraries/LibDraw/ImageLoader.h
@@ -0,0 +1,35 @@
+#pragma once
+
+#include <AK/NonnullRefPtr.h>
+#include <AK/OwnPtr.h>
+#include <AK/RefCounted.h>
+#include <LibDraw/Size.h>
+
+class GraphicsBitmap;
+
+class ImageLoaderPlugin {
+public:
+    virtual ~ImageLoaderPlugin() {}
+
+    virtual Size size() = 0;
+    virtual RefPtr<GraphicsBitmap> bitmap() = 0;
+
+protected:
+    ImageLoaderPlugin() {}
+};
+
+class ImageLoader : public RefCounted<ImageLoader> {
+public:
+    static NonnullRefPtr<ImageLoader> create(const u8* data, size_t size) { return adopt(*new ImageLoader(data, size)); }
+    ~ImageLoader();
+
+    Size size() const { return m_plugin->size(); }
+    int width() const { return size().width(); }
+    int height() const { return size().height(); }
+    RefPtr<GraphicsBitmap> bitmap() const { return m_plugin->bitmap(); }
+
+private:
+    ImageLoader(const u8*, size_t);
+
+    mutable OwnPtr<ImageLoaderPlugin> m_plugin;
+};
diff --git a/Libraries/LibDraw/Makefile b/Libraries/LibDraw/Makefile
index 653c372a2a..e866614ec4 100644
--- a/Libraries/LibDraw/Makefile
+++ b/Libraries/LibDraw/Makefile
@@ -8,6 +8,7 @@ OBJS = \
     GraphicsBitmap.o \
     Painter.o \
     PNGLoader.o \
+    ImageLoader.o \
     Rect.o \
     StylePainter.o \
     Emoji.o
diff --git a/Libraries/LibDraw/PNGLoader.cpp b/Libraries/LibDraw/PNGLoader.cpp
index d848045a83..f86d37b438 100644
--- a/Libraries/LibDraw/PNGLoader.cpp
+++ b/Libraries/LibDraw/PNGLoader.cpp
@@ -61,6 +61,14 @@ struct [[gnu::packed]] Quad16
 };
 
 struct PNGLoadingContext {
+    enum class State {
+        NotDecoded,
+        ChunksDecoded,
+        BitmapDecoded,
+    };
+    State state { State::NotDecoded };
+    const u8* data { nullptr };
+    size_t data_size { 0 };
     int width { -1 };
     int height { -1 };
     u8 bit_depth { 0 };
@@ -382,23 +390,22 @@ template<bool has_alpha, u8 filter_type>
     }
 }
 
-static RefPtr<GraphicsBitmap> load_png_impl(const u8* data, int data_size)
+static bool decode_png_chunks(PNGLoadingContext& context)
 {
+    ASSERT(context.state == PNGLoadingContext::State::NotDecoded);
 #ifdef PNG_STOPWATCH_DEBUG
     Stopwatch sw("load_png_impl: total");
 #endif
-    const u8* data_ptr = data;
-    int data_remaining = data_size;
+    const u8* data_ptr = context.data;
+    int data_remaining = context.data_size;
 
     const u8 png_header[8] = { 0x89, 'P', 'N', 'G', 13, 10, 26, 10 };
-    if (memcmp(data, png_header, sizeof(png_header))) {
+    if (memcmp(context.data, png_header, sizeof(png_header))) {
         dbgprintf("Invalid PNG header\n");
-        return nullptr;
+        return false;
     }
 
-    PNGLoadingContext context;
-
-    context.compressed_data.ensure_capacity(data_size);
+    context.compressed_data.ensure_capacity(context.data_size);
 
     data_ptr += sizeof(png_header);
     data_remaining -= sizeof(png_header);
@@ -410,11 +417,18 @@ static RefPtr<GraphicsBitmap> load_png_impl(const u8* data, int data_size)
         Streamer streamer(data_ptr, data_remaining);
         while (!streamer.at_end()) {
             if (!process_chunk(streamer, context)) {
-                return nullptr;
+                return false;
             }
         }
     }
 
+    context.state = PNGLoadingContext::State::ChunksDecoded;
+    return true;
+}
+
+static bool decode_png_bitmap(PNGLoadingContext& context)
+{
+    ASSERT(context.state == PNGLoadingContext::State::ChunksDecoded);
     {
 #ifdef PNG_STOPWATCH_DEBUG
         Stopwatch sw("load_png_impl: uncompress");
@@ -423,7 +437,7 @@ static RefPtr<GraphicsBitmap> load_png_impl(const u8* data, int data_size)
         unsigned long destlen = context.decompression_buffer_size;
         int ret = puff(context.decompression_buffer, &destlen, context.compressed_data.data() + 2, &srclen);
         if (ret < 0)
-            return nullptr;
+            return false;
         context.compressed_data.clear();
     }
 
@@ -436,12 +450,12 @@ static RefPtr<GraphicsBitmap> load_png_impl(const u8* data, int data_size)
         for (int y = 0; y < context.height; ++y) {
             u8 filter;
             if (!streamer.read(filter))
-                return nullptr;
+                return false;
 
             context.scanlines.append({ filter });
             auto& scanline_buffer = context.scanlines.last().data;
             if (!streamer.wrap_bytes(scanline_buffer, context.width * context.bytes_per_pixel))
-                return nullptr;
+                return false;
         }
     }
 
@@ -458,6 +472,22 @@ static RefPtr<GraphicsBitmap> load_png_impl(const u8* data, int data_size)
     context.decompression_buffer = nullptr;
     context.decompression_buffer_size = 0;
 
+    context.state = PNGLoadingContext::State::BitmapDecoded;
+    return true;
+}
+
+static RefPtr<GraphicsBitmap> load_png_impl(const u8* data, int data_size)
+{
+    PNGLoadingContext context;
+    context.data = data;
+    context.data_size = data_size;
+
+    if (!decode_png_chunks(context))
+        return nullptr;
+
+    if (!decode_png_bitmap(context))
+        return nullptr;
+
     return context.bitmap;
 }
 
@@ -576,3 +606,37 @@ static bool process_chunk(Streamer& streamer, PNGLoadingContext& context)
         return process_tRNS(chunk_data, context);
     return true;
 }
+
+PNGImageLoaderPlugin::PNGImageLoaderPlugin(const u8* data, size_t size)
+{
+    m_context = make<PNGLoadingContext>();
+    m_context->data = data;
+    m_context->data_size = size;
+}
+
+PNGImageLoaderPlugin::~PNGImageLoaderPlugin()
+{
+}
+
+Size PNGImageLoaderPlugin::size()
+{
+    if (m_context->state == PNGLoadingContext::State::NotDecoded) {
+        bool success = decode_png_chunks(*m_context);
+        ASSERT(success);
+    }
+
+    return { m_context->width, m_context->height };
+}
+
+RefPtr<GraphicsBitmap> PNGImageLoaderPlugin::bitmap()
+{
+    if (m_context->state != PNGLoadingContext::State::BitmapDecoded) {
+        // NOTE: This forces the chunk decoding to happen.
+        size();
+        bool success = decode_png_bitmap(*m_context);
+        ASSERT(success);
+    }
+
+    ASSERT(m_context->bitmap);
+    return m_context->bitmap;
+}
diff --git a/Libraries/LibDraw/PNGLoader.h b/Libraries/LibDraw/PNGLoader.h
index 323579ffc6..baf33f10b6 100644
--- a/Libraries/LibDraw/PNGLoader.h
+++ b/Libraries/LibDraw/PNGLoader.h
@@ -1,6 +1,21 @@
 #pragma once
 
 #include <LibDraw/GraphicsBitmap.h>
+#include <LibDraw/ImageLoader.h>
 
 RefPtr<GraphicsBitmap> load_png(const StringView& path);
 RefPtr<GraphicsBitmap> load_png_from_memory(const u8*, size_t);
+
+struct PNGLoadingContext;
+
+class PNGImageLoaderPlugin final : public ImageLoaderPlugin {
+public:
+    virtual ~PNGImageLoaderPlugin() override;
+    PNGImageLoaderPlugin(const u8*, size_t);
+
+    virtual Size size() override;
+    virtual RefPtr<GraphicsBitmap> bitmap() override;
+
+private:
+    OwnPtr<PNGLoadingContext> m_context;
+};