1 files changed, 18 insertions, 16 deletions
diff --git a/src/jpg/segment.zig b/src/jpg/segment.zig
index ec1346a..9b55780 100644
--- a/src/jpg/segment.zig
+++ b/src/jpg/segment.zig
@@ -1,7 +1,7 @@
 const std = @import("std");
 const jpg = @import("root.zig");
 
-pub const Segment = union(jpg.Marker.Kind) {
+pub const Segment = union(enum) {
     temporary: void,
     reset: void,
     start_of_image: void,
@@ -13,6 +13,7 @@ pub const Segment = union(jpg.Marker.Kind) {
     define_number_of_lines: void,
     define_hierarchy_progression: void,
     extend_huffman_table: void,
+    jfif: jpg.Jfif,
     application_segment: void,
     comment: void,
     start_of_scan: void,
@@ -20,6 +21,8 @@ pub const Segment = union(jpg.Marker.Kind) {
     const Self = @This();
 
     pub fn read(gpa: std.mem.Allocator, reader: *std.Io.Reader) !Self {
+        _ = gpa;
+
         if (try reader.peekByte() != 0xff) {
             return error.InvalidMarkerHeader;
         }
@@ -34,26 +37,22 @@ pub const Segment = union(jpg.Marker.Kind) {
             .reset => .{ .reset = void{} },
             .start_of_image => .{ .start_of_image = void{} },
             .end_of_image => .{ .end_of_image = void{} },
-            .start_of_frame => data: {
-                const buffer = try read_data_with_length(gpa, reader);
-                defer gpa.free(buffer);
-                break :data .{ .start_of_frame = void{} };
-            },
-            .application_segment => data: {
-                const buffer = try read_data_with_length(gpa, reader);
-                defer gpa.free(buffer);
-                break :data .{ .application_segment = void{} };
+            .application_segment => |n| switch (n) {
+                0 => jfif: {
+                    // NOTE: we can toss the length of this segment
+                    //       since the JFIF segment has a predefined size
+                    reader.toss(2);
+                    break :jfif .{
+                        .jfif = try jpg.Jfif.read(reader)
+                    };
+                },
+                else => unreachable,
             },
             else => {
                 std.debug.panic("unimplemented tag '{s}'\n", .{@tagName(marker)});
             },
         };
     }
-
-    fn read_data_with_length(gpa: std.mem.Allocator, reader: *std.Io.Reader) ![]u8 {
-        const length = try reader.takeInt(u16, .big);
-        return try reader.readAlloc(gpa, @intCast(length));
-    }
 };
 
 
@@ -65,5 +64,8 @@ test "file" {
     try std.testing.expectEqual(.start_of_image, std.meta.activeTag(segment));
 
     segment = try .read(std.testing.allocator, &reader);
-    try std.testing.expectEqual(.application_segment, std.meta.activeTag(segment));
+    try std.testing.expectEqual(.jfif, std.meta.activeTag(segment));
+
+    segment = try .read(std.testing.allocator, &reader);
+    try std.testing.expectEqual(.jfif, std.meta.activeTag(segment));
 }