terminal/kitty-gfx: centralize all image loading on LoadingImage

src/terminal/kitty/graphics_exec.zig

@@ -8,7 +8,7 @@ const command = @import("graphics_command.zig");
 const image = @import("graphics_image.zig");
 const Command = command.Command;
 const Response = command.Response;
-const ChunkedImage = image.ChunkedImage;
+const LoadingImage = image.LoadingImage;
 const Image = image.Image;
 
 const log = std.log.scoped(.kitty_gfx);
@@ -83,11 +83,11 @@ fn query(alloc: Allocator, cmd: *Command) Response {
     };
 
     // Attempt to load the image. If we cannot, then set an appropriate error.
-    var img = Image.load(alloc, cmd) catch |err| {
+    var loading = LoadingImage.init(alloc, cmd) catch |err| {
         encodeError(&result, err);
         return result;
     };
-    img.deinit(alloc);
+    loading.deinit(alloc);
 
     return result;
 }
@@ -201,55 +201,60 @@ fn loadAndAddImage(
     const storage = &terminal.screen.kitty_images;
 
     // Determine our image. This also handles chunking and early exit.
-    var img = if (storage.chunk) |chunk| img: {
+    var loading: LoadingImage = if (storage.loading) |loading| loading: {
         // Note: we do NOT want to call "cmd.toOwnedData" here because
         // we're _copying_ the data. We want the command data to be freed.
-        try chunk.addData(alloc, cmd.data);
+        try loading.addData(alloc, cmd.data);
 
         // If we have more then we're done
-        if (t.more_chunks) return chunk.image;
+        if (t.more_chunks) return loading.image;
 
-        // We have no more chunks. Complete and validate the image.
-        // At this point no matter what we want to clear out our chunked
-        // state. If we hit a validation error or something we don't want
-        // the chunked image hanging around in-memory.
+        // We have no more chunks. We're going to be completing the
+        // image so we want to destroy the pointer to the loading
+        // image and copy it out.
         defer {
-            chunk.destroy(alloc);
-            storage.chunk = null;
+            alloc.destroy(loading);
+            storage.loading = null;
         }
 
-        break :img try chunk.complete(alloc);
-    } else img: {
-        const img = try Image.load(alloc, cmd);
-        _ = cmd.toOwnedData();
-        break :img img;
-    };
-    errdefer img.deinit(alloc);
+        break :loading loading.*;
+    } else try LoadingImage.init(alloc, cmd);
+
+    // We only want to deinit on error. If we're chunking, then we don't
+    // want to deinit at all. If we're not chunking, then we'll deinit
+    // after we've copied the image out.
+    errdefer loading.deinit(alloc);
 
     // If the image has no ID, we assign one
-    if (img.id == 0) {
-        img.id = storage.next_id;
+    if (loading.image.id == 0) {
+        loading.image.id = storage.next_id;
         storage.next_id +%= 1;
     }
 
     // If this is chunked, this is the beginning of a new chunked transmission.
     // (We checked for an in-progress chunk above.)
     if (t.more_chunks) {
-        // We allocate the chunk on the heap because its rare and we
+        // We allocate the pointer on the heap because its rare and we
         // don't want to always pay the memory cost to keep it around.
-        const chunk_ptr = try alloc.create(ChunkedImage);
-        errdefer alloc.destroy(chunk_ptr);
-        chunk_ptr.* = try ChunkedImage.init(alloc, img);
-        storage.chunk = chunk_ptr;
-        return img;
+        const loading_ptr = try alloc.create(LoadingImage);
+        errdefer alloc.destroy(loading_ptr);
+        loading_ptr.* = loading;
+        storage.loading = loading_ptr;
+        return loading.image;
     }
 
     // Dump the image data before it is decompressed
     // img.debugDump() catch unreachable;
 
     // Validate and store our image
-    try img.complete(alloc);
+    var img = try loading.complete(alloc);
+    errdefer img.deinit(alloc);
     try storage.addImage(alloc, img);
+
+    // Ensure we deinit the loading state because we're done. The image
+    // won't be deinit because of "complete" above.
+    loading.deinit(alloc);
+
     return img;
 }
 

src/terminal/kitty/graphics_image.zig

@@ -14,9 +14,11 @@ const max_dimension = 10000;
 /// Maximum size in bytes, taken from Kitty.
 const max_size = 400 * 1024 * 1024; // 400MB
 
-/// A chunked image is an image that is in-progress and being constructed
-/// using chunks (the "m" parameter in the protocol).
-pub const ChunkedImage = struct {
+/// An image that is still being loaded. The image should be initialized
+/// using init on the first chunk and then addData for each subsequent
+/// chunk. Once all chunks have been added, complete should be called
+/// to finalize the image.
+pub const LoadingImage = struct {
     /// The in-progress image. The first chunk must have all the metadata
     /// so this comes from that initially.
     image: Image,
@@ -24,31 +26,66 @@ pub const ChunkedImage = struct {
     /// The data that is being built up.
     data: std.ArrayListUnmanaged(u8) = .{},
 
-    /// Initialize a chunked image from the first image part.
-    pub fn init(alloc: Allocator, image: Image) !ChunkedImage {
-        // Copy our initial set of data
-        var data = try std.ArrayListUnmanaged(u8).initCapacity(alloc, image.data.len * 2);
-        errdefer data.deinit(alloc);
-        try data.appendSlice(alloc, image.data);
+    /// Initialize a chunked immage from the first image transmission.
+    /// If this is a multi-chunk image, this should only be the FIRST
+    /// chunk.
+    pub fn init(alloc: Allocator, cmd: *command.Command) !LoadingImage {
+        // We must have data to load an image
+        if (cmd.data.len == 0) return error.InvalidData;
+
+        // Build our initial image from the properties sent via the control.
+        // These can be overwritten by the data loading process. For example,
+        // PNG loading sets the width/height from the data.
+        const t = cmd.transmission().?;
+        var result: LoadingImage = .{
+            .image = .{
+                .id = t.image_id,
+                .number = t.image_number,
+                .width = t.width,
+                .height = t.height,
+                .compression = t.compression,
+                .format = switch (t.format) {
+                    .rgb => .rgb,
+                    .rgba => .rgba,
+                    else => unreachable,
+                },
+            },
+        };
+
+        // Load the base64 encoded data from the transmission medium.
+        const raw_data = switch (t.medium) {
+            .direct => direct: {
+                const data = cmd.data;
+                _ = cmd.toOwnedData();
+                break :direct data;
+            },
+
+            else => {
+                std.log.warn("unimplemented medium={}", .{t.medium});
+                return error.UnsupportedMedium;
+            },
+        };
+        defer alloc.free(raw_data);
+
+        // Add the data
+        try result.addData(alloc, raw_data);
 
-        // Set data to empty so it doesn't get freed.
-        var result: ChunkedImage = .{ .image = image, .data = data };
-        result.image.data = "";
         return result;
     }
 
-    pub fn deinit(self: *ChunkedImage, alloc: Allocator) void {
+    pub fn deinit(self: *LoadingImage, alloc: Allocator) void {
         self.image.deinit(alloc);
         self.data.deinit(alloc);
     }
 
-    pub fn destroy(self: *ChunkedImage, alloc: Allocator) void {
+    pub fn destroy(self: *LoadingImage, alloc: Allocator) void {
         self.deinit(alloc);
         alloc.destroy(self);
     }
 
-    /// Adds a chunk of base64-encoded data to the image.
-    pub fn addData(self: *ChunkedImage, alloc: Allocator, data: []const u8) !void {
+    /// Adds a chunk of base64-encoded data to the image. Use this if the
+    /// image is coming in chunks (the "m" parameter in the protocol).
+    pub fn addData(self: *LoadingImage, alloc: Allocator, data: []const u8) !void {
         const Base64Decoder = std.base64.standard.Decoder;
 
         // Grow our array list by size capacity if it needs it
@@ -69,10 +106,12 @@ pub const ChunkedImage = struct {
     }
 
     /// Complete the chunked image, returning a completed image.
-    pub fn complete(self: *ChunkedImage, alloc: Allocator) !Image {
+    pub fn complete(self: *LoadingImage, alloc: Allocator) !Image {
         var result = self.image;
         result.data = try self.data.toOwnedSlice(alloc);
+        errdefer result.deinit(alloc);
         self.image = .{};
+        try result.complete(alloc);
         return result;
     }
 };
@@ -180,83 +219,6 @@ pub const Image = struct {
         if (actual_len != expected_len) return error.InvalidData;
     }
 
-    /// Load an image from a transmission. The data in the command will be
-    /// owned by the image if successful. Note that you still must deinit
-    /// the command, all the state change will be done internally.
-    ///
-    /// If the command represents a chunked image then this image will
-    /// be incomplete. The caller is expected to inspect the command
-    /// and determine if it is a chunked image.
-    pub fn load(alloc: Allocator, cmd: *command.Command) !Image {
-        const t = cmd.transmission().?;
-
-        // We must have data to load an image
-        if (cmd.data.len == 0) return error.InvalidData;
-
-        // Load the data
-        const raw_data = switch (t.medium) {
-            .direct => direct: {
-                const data = cmd.data;
-                _ = cmd.toOwnedData();
-                break :direct data;
-            },
-
-            else => {
-                std.log.warn("unimplemented medium={}", .{t.medium});
-                return error.UnsupportedMedium;
-            },
-        };
-
-        // We always free the raw data because it is base64 decoded below
-        defer alloc.free(raw_data);
-
-        // We base64 the data immediately
-        const decoded_data = base64Decode(alloc, raw_data) catch |err| {
-            log.warn("failed to calculate base64 decoded size: {}", .{err});
-            return error.InvalidData;
-        };
-
-        // If we loaded an image successfully then we take ownership
-        // of the command data and we need to make sure to clean up on error.
-        errdefer if (decoded_data.len > 0) alloc.free(decoded_data);
-
-        const img = switch (t.format) {
-            .rgb, .rgba => try loadPacked(t, decoded_data),
-            else => return error.UnsupportedFormat,
-        };
-
-        return img;
-    }
-
-    /// Read the temporary file data from a command. This will also DELETE
-    /// the temporary file if it is successful and the temporary file is
-    /// in a safe, well-known location.
-    fn readTemporaryFile(alloc: Allocator, path: []const u8) ![]const u8 {
-        _ = alloc;
-        _ = path;
-        return "";
-    }
-
-    /// Load a package image format, i.e. RGB or RGBA.
-    fn loadPacked(
-        t: command.Transmission,
-        data: []const u8,
-    ) !Image {
-        return Image{
-            .id = t.image_id,
-            .number = t.image_number,
-            .width = t.width,
-            .height = t.height,
-            .compression = t.compression,
-            .format = switch (t.format) {
-                .rgb => .rgb,
-                .rgba => .rgba,
-                else => unreachable,
-            },
-            .data = data,
-        };
-    }
-
     pub fn deinit(self: *Image, alloc: Allocator) void {
         if (self.data.len > 0) alloc.free(self.data);
     }
@@ -312,8 +274,8 @@ test "image load with invalid RGB data" {
         .data = try alloc.dupe(u8, "AAAA"),
     };
     defer cmd.deinit(alloc);
-    var img = try Image.load(alloc, &cmd);
-    defer img.deinit(alloc);
+    var loading = try LoadingImage.init(alloc, &cmd);
+    defer loading.deinit(alloc);
 }
 
 test "image load with image too wide" {
@@ -330,9 +292,9 @@ test "image load with image too wide" {
         .data = try alloc.dupe(u8, "AAAA"),
     };
     defer cmd.deinit(alloc);
-    var img = try Image.load(alloc, &cmd);
-    defer img.deinit(alloc);
-    try testing.expectError(error.DimensionsTooLarge, img.complete(alloc));
+    var loading = try LoadingImage.init(alloc, &cmd);
+    defer loading.deinit(alloc);
+    try testing.expectError(error.DimensionsTooLarge, loading.complete(alloc));
 }
 
 test "image load with image too tall" {
@@ -349,9 +311,9 @@ test "image load with image too tall" {
         .data = try alloc.dupe(u8, "AAAA"),
     };
     defer cmd.deinit(alloc);
-    var img = try Image.load(alloc, &cmd);
-    defer img.deinit(alloc);
-    try testing.expectError(error.DimensionsTooLarge, img.complete(alloc));
+    var loading = try LoadingImage.init(alloc, &cmd);
+    defer loading.deinit(alloc);
+    try testing.expectError(error.DimensionsTooLarge, loading.complete(alloc));
 }
 
 test "image load: rgb, zlib compressed, direct" {
@@ -373,9 +335,10 @@ test "image load: rgb, zlib compressed, direct" {
         ),
     };
     defer cmd.deinit(alloc);
-    var img = try Image.load(alloc, &cmd);
+    var loading = try LoadingImage.init(alloc, &cmd);
+    defer loading.deinit(alloc);
+    var img = try loading.complete(alloc);
     defer img.deinit(alloc);
-    try img.complete(alloc);
 
     // should be decompressed
     try testing.expect(img.compression == .none);
@@ -400,9 +363,10 @@ test "image load: rgb, not compressed, direct" {
         ),
     };
     defer cmd.deinit(alloc);
-    var img = try Image.load(alloc, &cmd);
+    var loading = try LoadingImage.init(alloc, &cmd);
+    defer loading.deinit(alloc);
+    var img = try loading.complete(alloc);
     defer img.deinit(alloc);
-    try img.complete(alloc);
 
     // should be decompressed
     try testing.expect(img.compression == .none);

src/terminal/kitty/graphics_storage.zig

@@ -5,7 +5,7 @@ const ArenaAllocator = std.heap.ArenaAllocator;
 
 const point = @import("../point.zig");
 const command = @import("graphics_command.zig");
-const ChunkedImage = @import("graphics_image.zig").ChunkedImage;
+const LoadingImage = @import("graphics_image.zig").LoadingImage;
 const Image = @import("graphics_image.zig").Image;
 const Command = command.Command;
 const ScreenPoint = point.ScreenPoint;
@@ -29,11 +29,11 @@ pub const ImageStorage = struct {
     /// The set of placements for loaded images.
     placements: PlacementMap = .{},
 
-    /// Non-null if there is a chunked image in progress.
-    chunk: ?*ChunkedImage = null,
+    /// Non-null if there is an in-progress loading image.
+    loading: ?*LoadingImage = null,
 
     pub fn deinit(self: *ImageStorage, alloc: Allocator) void {
-        if (self.chunk) |chunk| chunk.destroy(alloc);
+        if (self.loading) |loading| loading.destroy(alloc);
 
         var it = self.images.iterator();
         while (it.next()) |kv| kv.value_ptr.deinit(alloc);