lib/png/src/encode/model.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2
3 const alpha_vector_bytes = 64;
4 const AlphaVector = @Vector(alpha_vector_bytes, u8);
5 const alpha_rgb_fill = alphaRgbFill();
6 const opaque_vector: AlphaVector = @splat(0xff);
7 const rgba_vector_bytes = 16;
8 const rgb_vector_bytes = 12;
9 const RgbaVector = @Vector(rgba_vector_bytes, u8);
10 const RgbVector = @Vector(rgb_vector_bytes, u8);
11 const rgba_to_rgb_shuffle: @Vector(rgb_vector_bytes, i32) = .{
12 0, 1, 2, 4, 5, 6, 8, 9, 10, 12, 13, 14,
13 };
14
15 pub const Compression = enum {
16 compressed,
17 stored,
18 };
19
20 pub const Options = struct {
21 compression: Compression = .compressed,
22 };
23
24 pub const ImageView = struct {
25 rgba8: []const u8,
26 width: i32,
27 height: i32,
28 options: Options = .{},
29 };
30
31 pub const Color = enum(u8) {
32 rgb = 2,
33 rgba = 6,
34
35 pub fn bytesPerPixel(self: Color) usize {
36 return switch (self) {
37 .rgb => 3,
38 .rgba => 4,
39 };
40 }
41 };
42
43 pub const EncodeError = error{
44 CapacityOverflow,
45 InvalidDimensions,
46 PixelSizeMismatch,
47 OutputCapacityExceeded,
48 };
49
50 pub const Exhaustion = error{
51 EncodeInputMismatch,
52 EncodeStorageInUse,
53 ImagePixelCapacityExceeded,
54 };
55
56 pub fn pixelCount(width: i32, height: i32) EncodeError!usize {
57 if (width <= 0 or height <= 0) return error.InvalidDimensions;
58 return std.math.mul(usize, @intCast(width), @intCast(height)) catch {
59 return error.CapacityOverflow;
60 };
61 }
62
63 pub fn pixelByteCount(width: i32, height: i32) EncodeError!usize {
64 return std.math.mul(usize, try pixelCount(width, height), 4) catch {
65 return error.CapacityOverflow;
66 };
67 }
68
69 pub fn colorForPixels(pixels: []const u8) Color {
70 var index: usize = 0;
71 while (pixels.len - index >= alpha_vector_bytes) : (index += alpha_vector_bytes) {
72 const block: AlphaVector = @bitCast(pixels[index..][0..alpha_vector_bytes].*);
73 if (!@reduce(.And, (block | alpha_rgb_fill) == opaque_vector)) return .rgba;
74 }
75 index += 3;
76 while (index < pixels.len) : (index += 4) {
77 if (pixels[index] != 0xff) return .rgba;
78 }
79 return .rgb;
80 }
81
82 fn alphaRgbFill() AlphaVector {
83 var bytes: [alpha_vector_bytes]u8 = @splat(0xff);
84 var index: usize = 3;
85 while (index < bytes.len) : (index += 4) bytes[index] = 0;
86 return @bitCast(bytes);
87 }
88
89 pub fn imageHash(image: ImageView) u64 {
90 var hash = std.hash.Wyhash.init(0);
91 hash.update(std.mem.asBytes(&image.width));
92 hash.update(std.mem.asBytes(&image.height));
93 const compression: u8 = @backingInt(image.options.compression);
94 hash.update(std.mem.asBytes(&compression));
95 hash.update(image.rgba8);
96 return hash.final();
97 }
98
99 pub const RawSource = struct {
100 pixels: []const u8,
101 width: usize,
102 height: usize,
103 color: Color,
104 row: usize = 0,
105 row_offset: usize = 0,
106
107 pub fn init(image: ImageView, color: Color) RawSource {
108 std.debug.assert(image.width > 0);
109 std.debug.assert(image.height > 0);
110 return .{
111 .pixels = image.rgba8,
112 .width = @intCast(image.width),
113 .height = @intCast(image.height),
114 .color = color,
115 };
116 }
117
118 pub fn read(self: *RawSource, output: []u8) usize {
119 var written: usize = 0;
120 while (written < output.len and self.row < self.height) {
121 if (self.row_offset == 0) {
122 output[written] = 0;
123 written += 1;
124 self.row_offset = 1;
125 continue;
126 }
127 written += switch (self.color) {
128 .rgba => self.readRgba(output[written..]),
129 .rgb => self.readRgb(output[written..]),
130 };
131 if (self.row_offset == self.rowBytes() + 1) {
132 self.row += 1;
133 self.row_offset = 0;
134 }
135 }
136 return written;
137 }
138
139 pub fn remaining(self: RawSource) usize {
140 if (self.row == self.height) return 0;
141 const stride = self.rowBytes() + 1;
142 return (self.height - self.row - 1) * stride + stride - self.row_offset;
143 }
144
145 fn readRgba(self: *RawSource, output: []u8) usize {
146 const data_offset = self.row_offset - 1;
147 const count = @min(output.len, self.rowBytes() - data_offset);
148 const source_offset = self.row * self.width * 4 + data_offset;
149 @memcpy(output[0..count], self.pixels[source_offset..][0..count]);
150 self.row_offset += count;
151 return count;
152 }
153
154 fn readRgb(self: *RawSource, output: []u8) usize {
155 const data_offset = self.row_offset - 1;
156 const count = @min(output.len, self.rowBytes() - data_offset);
157 const source_row = self.row * self.width * 4;
158 var output_index: usize = 0;
159 while (output_index < count and (data_offset + output_index) % 3 != 0) : (output_index += 1) {
160 const offset = data_offset + output_index;
161 const pixel = offset / 3;
162 const channel = offset % 3;
163 output[output_index] = self.pixels[source_row + pixel * 4 + channel];
164 }
165 while (count - output_index >= rgb_vector_bytes) : (output_index += rgb_vector_bytes) {
166 const pixel = (data_offset + output_index) / 3;
167 const source_offset = source_row + pixel * 4;
168 const rgba: RgbaVector = @bitCast(self.pixels[source_offset..][0..rgba_vector_bytes].*);
169 const rgb: RgbVector = @shuffle(u8, rgba, undefined, rgba_to_rgb_shuffle);
170 output[output_index..][0..rgb_vector_bytes].* = @bitCast(rgb);
171 }
172 for (output[output_index..count], data_offset + output_index..) |*byte, offset| {
173 const pixel = offset / 3;
174 const channel = offset % 3;
175 byte.* = self.pixels[source_row + pixel * 4 + channel];
176 }
177 self.row_offset += count;
178 return count;
179 }
180
181 fn rowBytes(self: RawSource) usize {
182 return self.width * self.color.bytesPerPixel();
183 }
184 };
185
186 test "raw source streams RGB and RGBA scanlines without ownership" {
187 const opaque_pixels = [_]u8{ 1, 2, 3, 255, 4, 5, 6, 255 };
188 var rgb = RawSource.init(.{ .rgba8 = &opaque_pixels, .width = 2, .height = 1 }, .rgb);
189 var rgb_bytes: [7]u8 = undefined;
190 try std.testing.expectEqual(rgb_bytes.len, rgb.read(&rgb_bytes));
191 try std.testing.expectEqualSlices(u8, &.{ 0, 1, 2, 3, 4, 5, 6 }, &rgb_bytes);
192 try std.testing.expectEqual(@as(usize, 0), rgb.remaining());
193
194 const alpha = [_]u8{ 1, 2, 3, 7, 4, 5, 6, 8 };
195 var rgba = RawSource.init(.{ .rgba8 = &alpha, .width = 2, .height = 1 }, .rgba);
196 var rgba_bytes: [9]u8 = undefined;
197 try std.testing.expectEqual(rgba_bytes.len, rgba.read(&rgba_bytes));
198 try std.testing.expectEqualSlices(u8, &.{ 0, 1, 2, 3, 7, 4, 5, 6, 8 }, &rgba_bytes);
199 try std.testing.expectEqual(@as(usize, 0), rgba.remaining());
200 }
201
202 test "color planning matches scalar alpha scans across every vector tail and alignment" {
203 const maximum_bytes = 257;
204 var storage: [maximum_bytes + 3]u8 = undefined;
205 for (0..maximum_bytes + 1) |len| {
206 for (0..4) |offset| {
207 const pixels = storage[offset..][0..len];
208 for (pixels, 0..) |*byte, index| {
209 byte.* = if (index % 4 == 3) 0xff else @truncate(index *% 37 +% offset *% 11);
210 }
211 for (0..4) |pattern| {
212 const alpha_count = len / 4;
213 const changed_alpha: ?usize = if (pattern == 0 or alpha_count == 0)
214 null
215 else switch (pattern) {
216 1 => 0,
217 2 => alpha_count / 2,
218 3 => alpha_count - 1,
219 else => unreachable,
220 };
221 if (changed_alpha) |alpha| pixels[alpha * 4 + 3] = 0x7f;
222
223 var expected: Color = .rgb;
224 var alpha_index: usize = 3;
225 while (alpha_index < pixels.len) : (alpha_index += 4) {
226 if (pixels[alpha_index] != 0xff) {
227 expected = .rgba;
228 break;
229 }
230 }
231 try std.testing.expectEqual(expected, colorForPixels(pixels));
232 if (changed_alpha) |alpha| pixels[alpha * 4 + 3] = 0xff;
233 }
234 }
235 }
236 }
237
238 test "RGB streaming matches scalar projection across fragments and alignments" {
239 const width = 65;
240 const height = 3;
241 const pixel_bytes = width * height * 4;
242 const stream_bytes = height * (width * 3 + 1);
243 var source_storage: [pixel_bytes + 3]u8 = undefined;
244 var expected: [stream_bytes]u8 = undefined;
245 var actual: [stream_bytes]u8 = undefined;
246 var scratch_storage: [80 + 3]u8 = undefined;
247
248 for (0..4) |source_offset| {
249 const source = source_storage[source_offset..][0..pixel_bytes];
250 for (source, 0..) |*byte, index| byte.* = @truncate(index *% 43 +% source_offset *% 17);
251 var expected_index: usize = 0;
252 for (0..height) |row| {
253 expected[expected_index] = 0;
254 expected_index += 1;
255 for (0..width) |pixel| {
256 const input = (row * width + pixel) * 4;
257 @memcpy(expected[expected_index..][0..3], source[input..][0..3]);
258 expected_index += 3;
259 }
260 }
261 try std.testing.expectEqual(expected.len, expected_index);
262
263 for (1..81) |fragment_bytes| {
264 for (0..4) |destination_offset| {
265 var raw = RawSource.init(.{
266 .rgba8 = source,
267 .width = width,
268 .height = height,
269 }, .rgb);
270 var actual_index: usize = 0;
271 while (raw.remaining() != 0) {
272 const scratch = scratch_storage[destination_offset..][0..fragment_bytes];
273 const written = raw.read(scratch);
274 try std.testing.expect(written != 0);
275 @memcpy(actual[actual_index..][0..written], scratch[0..written]);
276 actual_index += written;
277 }
278 try std.testing.expectEqual(actual.len, actual_index);
279 try std.testing.expectEqualSlices(u8, &expected, &actual);
280 }
281 }
282 }
283 }