From 8941c8cbc6a54195b437256f75a415cf54bc6d3e Mon Sep 17 00:00:00 2001
From: Palmer Cox <palmer@fiscalnote.com>
Date: Fri, 12 Nov 2021 15:53:36 -0500
Subject: [PATCH] Create public APIs to check the validity of the final octet /
 quintet

One use case for this information is for error correction - if we know
that the trailing octet or quintet can't possibly be valid, we may be
able to do better error correction.
---
 src/decode_impl.rs | 14 +++++++++++++-
 src/encode_impl.rs | 13 ++++++++++++-
 src/lib.rs         |  8 ++++++--
 3 files changed, 31 insertions(+), 4 deletions(-)

diff --git a/src/decode_impl.rs b/src/decode_impl.rs
index a2539c7..98f515c 100644
--- a/src/decode_impl.rs
+++ b/src/decode_impl.rs
@@ -1,5 +1,8 @@
 use crate::error::{zbase32_error, ZBase32Error, ZBase32ErrorInfo};
-use crate::stateful_decoder::{HaveOctets, NeedQuintets, NextOctetResult, ProvideQuintetResult};
+use crate::stateful_decoder::{
+    quintet_has_valid_trailing_bits, HaveOctets, NeedQuintets, NextOctetResult,
+    ProvideQuintetResult,
+};
 use crate::tables::{CHARACTER_MIN_VALUE, CHARACTER_TO_QUINTET};
 use crate::util::{required_octets_buffer_len, required_quintets_buffer_len};
 use core::iter::Peekable;
@@ -128,6 +131,15 @@ fn calc_last_quintet_bits(bits: u64) -> Option<u8> {
     }
 }
 
+/// Determine if the last quintet is, given the number of bits to decode
+pub fn is_last_quintet_valid(bits: u64, quintet: u8) -> bool {
+    if let Some(last_quintet_bits) = calc_last_quintet_bits(bits) {
+        quintet <= 31 && quintet_has_valid_trailing_bits(last_quintet_bits, quintet)
+    } else {
+        false
+    }
+}
+
 /// Convert a buffer of quintet integer values (ie, integers of the
 /// range 0-31) into a buffer of octet (byte) values.
 ///
diff --git a/src/encode_impl.rs b/src/encode_impl.rs
index c484879..a703b37 100644
--- a/src/encode_impl.rs
+++ b/src/encode_impl.rs
@@ -1,5 +1,7 @@
 use crate::error::{zbase32_error, ZBase32Error, ZBase32ErrorInfo};
-use crate::stateful_encoder::{HaveQuintets, NeedOctets, NextQuintetResult, ProvideOctetResult};
+use crate::stateful_encoder::{
+    octet_has_valid_trailing_bits, HaveQuintets, NeedOctets, NextQuintetResult, ProvideOctetResult,
+};
 use crate::tables::QUINTET_TO_CHARACTER;
 use crate::util::{required_octets_buffer_len, required_quintets_buffer_len};
 use core::iter::Peekable;
@@ -123,6 +125,15 @@ fn calc_last_octet_bits(bits: u64) -> Option<u8> {
     }
 }
 
+/// Determine if the last octet is valid, given the number of bits to encode
+pub fn is_last_octet_valid(bits: u64, octet: u8) -> bool {
+    if let Some(last_octet_bits) = calc_last_octet_bits(bits) {
+        octet_has_valid_trailing_bits(last_octet_bits, octet)
+    } else {
+        false
+    }
+}
+
 /// Convert a buffer of octets (bytes) into a buffer of quintet values (ie, integers between 0-31).
 ///
 /// The length of `in_octets` must match the value returned by
diff --git a/src/lib.rs b/src/lib.rs
index 4e18924..aaa162c 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -134,12 +134,16 @@ pub use encode_impl::encode;
 
 pub mod low_level_decode {
     //! Low-level decoding functionality
-    pub use crate::decode_impl::{character_to_quintet, decode_slices, quintets_to_octets};
+    pub use crate::decode_impl::{
+        character_to_quintet, decode_slices, is_last_quintet_valid, quintets_to_octets,
+    };
     pub use crate::util::required_octets_buffer_len;
 }
 
 pub mod low_level_encode {
     //! Low-level encoding functionality
-    pub use crate::encode_impl::{encode_slices, octets_to_quintets, quintet_to_character};
+    pub use crate::encode_impl::{
+        encode_slices, is_last_octet_valid, octets_to_quintets, quintet_to_character,
+    };
     pub use crate::util::required_quintets_buffer_len;
 }