From 8a5550a43ce2011d6893d87684ea2743a3ccc2a7 Mon Sep 17 00:00:00 2001
From: Andrew Gunnerson <accounts+github@chiller3.com>
Date: Fri, 29 Nov 2024 21:21:11 -0500
Subject: [PATCH] Add new traits for zero padding and unpadding strings

Signed-off-by: Andrew Gunnerson <accounts+github@chiller3.com>
---
 avbroot/src/format/padding.rs | 34 +++++++++++++++++++++++++++++++++-
 1 file changed, 33 insertions(+), 1 deletion(-)

diff --git a/avbroot/src/format/padding.rs b/avbroot/src/format/padding.rs
index 067846a..2c349f0 100644
--- a/avbroot/src/format/padding.rs
+++ b/avbroot/src/format/padding.rs
@@ -1,4 +1,4 @@
-// SPDX-FileCopyrightText: 2023 Andrew Gunnerson
+// SPDX-FileCopyrightText: 2023-2024 Andrew Gunnerson
 // SPDX-License-Identifier: GPL-3.0-only
 
 use std::io::{self, Read, Seek, Write};
@@ -45,3 +45,35 @@ pub fn write_zeros(mut writer: impl Write + Seek, page_size: u64) -> io::Result<
 
     Ok(padding)
 }
+
+pub trait ZeroPadding {
+    /// Trim trailing zeros. Intermediate zeros before the last non-zero byte
+    /// are kept.
+    fn trim_end_padding(&self) -> &[u8];
+
+    /// Return the slice as an array padded with zeros at the end.
+    fn to_padded_array<const N: usize>(&self) -> Option<[u8; N]>;
+}
+
+impl ZeroPadding for [u8] {
+    fn trim_end_padding(&self) -> &[u8] {
+        let first_ending_zero = self
+            .iter()
+            .rposition(|b| *b != 0)
+            .map(|pos| pos + 1)
+            .unwrap_or_default();
+
+        &self[..first_ending_zero]
+    }
+
+    fn to_padded_array<const N: usize>(&self) -> Option<[u8; N]> {
+        if self.len() > N {
+            return None;
+        }
+
+        let mut result = [0u8; N];
+        result[..self.len()].copy_from_slice(self);
+
+        Some(result)
+    }
+}