From ae2d3d544333b5601bfe1a2a5ac5ebe1be058b1d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Motiejus=20Jak=C5=A1tys?= <motiejus@jakstys.lt>
Date: Sat, 19 Feb 2022 21:23:33 +0200
Subject: [PATCH] generalize padding funcs

---
 src/padding.zig | 33 +++++++++++++++++----------------
 src/shell.zig   |  4 ++--
 src/user.zig    |  2 +-
 3 files changed, 20 insertions(+), 19 deletions(-)

diff --git a/src/padding.zig b/src/padding.zig
index cfd9a13..66ff91e 100644
--- a/src/padding.zig
+++ b/src/padding.zig
@@ -2,28 +2,29 @@ const std = @import("std");
 
 // rounds up a u12 to the nearest factor of 4 and returns the difference
 // (padding)
-pub fn roundUp4Padding(comptime T: type, n: T) T {
-    return roundUp4(T, n) - n;
+pub fn roundUpPadding(comptime T: type, comptime nbits: u8, n: T) T {
+    return roundUp(T, nbits, n) - n;
 }
 
 // rounds up a u12 to the nearest factor of 4.
-pub fn roundUp4(comptime T: type, n: T) T {
-    return ((n + 3) & ~@as(T, 3));
+pub fn roundUp(comptime T: type, comptime nbits: u8, n: T) T {
+    const factor = comptime (1 << nbits) - 1;
+    return ((n + factor) & ~@as(T, factor));
 }
 
 const testing = std.testing;
 
 test "padding" {
-    try testing.expectEqual(roundUp4Padding(u12, 0), 0);
-    try testing.expectEqual(roundUp4Padding(u12, 1), 3);
-    try testing.expectEqual(roundUp4Padding(u12, 2), 2);
-    try testing.expectEqual(roundUp4Padding(u12, 3), 1);
-    try testing.expectEqual(roundUp4Padding(u12, 4), 0);
-    try testing.expectEqual(roundUp4Padding(u12, 40), 0);
-    try testing.expectEqual(roundUp4Padding(u12, 41), 3);
-    try testing.expectEqual(roundUp4Padding(u12, 42), 2);
-    try testing.expectEqual(roundUp4Padding(u12, 43), 1);
-    try testing.expectEqual(roundUp4Padding(u12, 44), 0);
-    try testing.expectEqual(roundUp4Padding(u12, 4091), 1);
-    try testing.expectEqual(roundUp4Padding(u12, 4092), 0);
+    try testing.expectEqual(roundUpPadding(u12, 2, 0), 0);
+    try testing.expectEqual(roundUpPadding(u12, 2, 1), 3);
+    try testing.expectEqual(roundUpPadding(u12, 2, 2), 2);
+    try testing.expectEqual(roundUpPadding(u12, 2, 3), 1);
+    try testing.expectEqual(roundUpPadding(u12, 2, 4), 0);
+    try testing.expectEqual(roundUpPadding(u12, 2, 40), 0);
+    try testing.expectEqual(roundUpPadding(u12, 2, 41), 3);
+    try testing.expectEqual(roundUpPadding(u12, 2, 42), 2);
+    try testing.expectEqual(roundUpPadding(u12, 2, 43), 1);
+    try testing.expectEqual(roundUpPadding(u12, 2, 44), 0);
+    try testing.expectEqual(roundUpPadding(u12, 2, 4091), 1);
+    try testing.expectEqual(roundUpPadding(u12, 2, 4092), 0);
 }
diff --git a/src/shell.zig b/src/shell.zig
index 9dcacfb..a1dc1e0 100644
--- a/src/shell.zig
+++ b/src/shell.zig
@@ -71,7 +71,7 @@ pub const ShellWriter = struct {
                 });
 
                 fullOffset += len;
-                const padding = pad.roundUp4Padding(u12, fullOffset);
+                const padding = pad.roundUpPadding(u12, 2, fullOffset);
                 fullOffset += padding;
                 //const stderr = std.io.getStdErr().writer();
                 //stderr.print("\n", .{}) catch unreachable;
@@ -209,7 +209,7 @@ test "basic shellpopcon" {
     try testing.expectEqual(sections.getIndex(nobody), null);
     try testing.expectEqual(
         sections.sectionBlob().len,
-        pad.roundUp4(u12, bash.len) + pad.roundUp4(u12, zsh.len) + pad.roundUp4(u12, long.len),
+        pad.roundUp(u12, 2, bash.len) + pad.roundUp(u12, 2, zsh.len) + pad.roundUp(u12, 2, long.len),
     );
 
     const shellReader = ShellReader.init(
diff --git a/src/user.zig b/src/user.zig
index fd46454..4168e7c 100644
--- a/src/user.zig
+++ b/src/user.zig
@@ -74,7 +74,7 @@ pub const UserWriter = struct {
             try self.appendTo.appendSlice(user.shell);
         }
 
-        const padding = pad.roundUp4Padding(u64, self.appendTo.items.len);
+        const padding = pad.roundUpPadding(u64, 2, self.appendTo.items.len);
         try self.appendTo.appendNTimes(0, padding);
     }
 };