Add buffering during encode

Currently we write char at a time to both writers (`fmt::Write` and `std::io::Write`). We can improve performance by first collecting the ASCII bytes of the encoded bech32 string into a buffer on the stack then writing the buffer as a single call. This is purely an optimization.
rust-bitcoin · Oct 16, 2023 · c2316bd · c2316bd
1 parent 487655c
commit c2316bd
Show file tree

Hide file tree

Showing 2 changed files with 114 additions and 36 deletions.
diff --git a/src/lib.rs b/src/lib.rs
@@ -160,6 +160,9 @@ pub use {
     crate::primitives::{Bech32, Bech32m, NoChecksum},
 };
 
+/// A bech32 string is at most 90 characters long.
+pub const BECH32_MAX_LENGTH: usize = 90;
+
 /// Decodes a bech32 encoded string.
 ///
 /// If this function succeeds the input string was found to be well formed (hrp, separator, bech32
@@ -265,11 +268,20 @@ pub fn encode_lower_to_fmt<Ck: Checksum, W: fmt::Write>(
     hrp: Hrp,
     data: &[u8],
 ) -> Result<(), fmt::Error> {
+    let mut buf = [0u8; BECH32_MAX_LENGTH];
+    let mut pos = 0;
+
     let iter = data.iter().copied().bytes_to_fes();
     let chars = iter.with_checksum::<Ck>(&hrp).chars();
-    for c in chars {
-        fmt.write_char(c)?;
-    }
+
+    buf.iter_mut().zip(chars).for_each(|(b, c)| {
+        *b = c as u8;
+        pos += 1;
+    });
+
+    let s = core::str::from_utf8(&buf[..pos]).expect("we only write ASCII");
+    fmt.write_str(s)?;
+
     Ok(())
 }
 
@@ -283,11 +295,20 @@ pub fn encode_upper_to_fmt<Ck: Checksum, W: fmt::Write>(
     hrp: Hrp,
     data: &[u8],
 ) -> Result<(), fmt::Error> {
+    let mut buf = [0u8; BECH32_MAX_LENGTH];
+    let mut pos = 0;
+
     let iter = data.iter().copied().bytes_to_fes();
     let chars = iter.with_checksum::<Ck>(&hrp).chars();
-    for c in chars {
-        fmt.write_char(c.to_ascii_uppercase())?;
-    }
+
+    buf.iter_mut().zip(chars).for_each(|(b, c)| {
+        *b = c.to_ascii_uppercase() as u8;
+        pos += 1;
+    });
+
+    let s = core::str::from_utf8(&buf[..pos]).expect("we only write ASCII");
+    fmt.write_str(s)?;
+
     Ok(())
 }
 
@@ -316,11 +337,19 @@ pub fn encode_lower_to_writer<Ck: Checksum, W: std::io::Write>(
     hrp: Hrp,
     data: &[u8],
 ) -> Result<(), std::io::Error> {
+    let mut buf = [0u8; BECH32_MAX_LENGTH];
+    let mut pos = 0;
+
     let iter = data.iter().copied().bytes_to_fes();
     let chars = iter.with_checksum::<Ck>(&hrp).chars();
-    for c in chars {
-        w.write_all(&[c as u8])?;
-    }
+
+    buf.iter_mut().zip(chars).for_each(|(b, c)| {
+        *b = c as u8;
+        pos += 1;
+    });
+
+    w.write_all(&buf[..pos])?;
+
     Ok(())
 }
 
@@ -335,11 +364,19 @@ pub fn encode_upper_to_writer<Ck: Checksum, W: std::io::Write>(
     hrp: Hrp,
     data: &[u8],
 ) -> Result<(), std::io::Error> {
+    let mut buf = [0u8; BECH32_MAX_LENGTH];
+    let mut pos = 0;
+
     let iter = data.iter().copied().bytes_to_fes();
     let chars = iter.with_checksum::<Ck>(&hrp).chars();
-    for c in chars {
-        w.write_all(&[c.to_ascii_uppercase() as u8])?;
-    }
+
+    buf.iter_mut().zip(chars).for_each(|(b, c)| {
+        *b = c.to_ascii_uppercase() as u8;
+        pos += 1;
+    });
+
+    w.write_all(&buf[..pos])?;
+
     Ok(())
 }
 

diff --git a/src/segwit.rs b/src/segwit.rs
@@ -54,6 +54,7 @@ use crate::primitives::iter::{ByteIterExt, Fe32IterExt};
 use crate::primitives::segwit;
 use crate::primitives::segwit::{InvalidWitnessVersionError, WitnessLengthError};
 use crate::primitives::{Bech32, Bech32m};
+use crate::BECH32_MAX_LENGTH;
 
 #[rustfmt::skip]                // Keep public re-exports separate.
 #[doc(inline)]
@@ -144,19 +145,30 @@ pub fn encode_lower_to_fmt_unchecked<W: fmt::Write>(
     witness_version: Fe32,
     witness_program: &[u8],
 ) -> fmt::Result {
+    let mut buf = [0u8; BECH32_MAX_LENGTH];
+    let mut pos = 0;
+
     let iter = witness_program.iter().copied().bytes_to_fes();
     match witness_version {
         VERSION_0 => {
-            for c in iter.with_checksum::<Bech32>(hrp).with_witness_version(VERSION_0).chars() {
-                fmt.write_char(c)?;
-            }
+            let chars = iter.with_checksum::<Bech32>(hrp).with_witness_version(VERSION_0).chars();
+            buf.iter_mut().zip(chars).for_each(|(b, c)| {
+                *b = c as u8;
+                pos += 1;
+            });
         }
         version => {
-            for c in iter.with_checksum::<Bech32m>(hrp).with_witness_version(version).chars() {
-                fmt.write_char(c)?;
-            }
+            let chars = iter.with_checksum::<Bech32m>(hrp).with_witness_version(version).chars();
+            buf.iter_mut().zip(chars).for_each(|(b, c)| {
+                *b = c as u8;
+                pos += 1;
+            });
         }
     }
+
+    let s = core::str::from_utf8(&buf[..pos]).expect("we only write ASCII");
+    fmt.write_str(s)?;
+
     Ok(())
 }
 
@@ -173,20 +185,30 @@ pub fn encode_upper_to_fmt_unchecked<W: fmt::Write>(
     witness_version: Fe32,
     witness_program: &[u8],
 ) -> fmt::Result {
+    let mut buf = [0u8; BECH32_MAX_LENGTH];
+    let mut pos = 0;
+
     let iter = witness_program.iter().copied().bytes_to_fes();
     match witness_version {
         VERSION_0 => {
-            for c in iter.with_checksum::<Bech32>(hrp).with_witness_version(VERSION_0).chars() {
-                fmt.write_char(c.to_ascii_uppercase())?;
-            }
+            let chars = iter.with_checksum::<Bech32>(hrp).with_witness_version(VERSION_0).chars();
+            buf.iter_mut().zip(chars).for_each(|(b, c)| {
+                *b = c.to_ascii_uppercase() as u8;
+                pos += 1;
+            });
         }
         version => {
-            for c in iter.with_checksum::<Bech32m>(hrp).with_witness_version(version).chars() {
-                fmt.write_char(c.to_ascii_uppercase())?;
-            }
+            let chars = iter.with_checksum::<Bech32m>(hrp).with_witness_version(version).chars();
+            buf.iter_mut().zip(chars).for_each(|(b, c)| {
+                *b = c.to_ascii_uppercase() as u8;
+                pos += 1;
+            });
         }
     }
 
+    let s = core::str::from_utf8(&buf[..pos]).expect("we only write ASCII");
+    fmt.write_str(s)?;
+
     Ok(())
 }
 
@@ -217,19 +239,29 @@ pub fn encode_lower_to_writer_unchecked<W: std::io::Write>(
     witness_version: Fe32,
     witness_program: &[u8],
 ) -> std::io::Result<()> {
+    let mut buf = [0u8; BECH32_MAX_LENGTH];
+    let mut pos = 0;
+
     let iter = witness_program.iter().copied().bytes_to_fes();
     match witness_version {
         VERSION_0 => {
-            for c in iter.with_checksum::<Bech32>(hrp).with_witness_version(VERSION_0).chars() {
-                w.write_all(&[c as u8])?;
-            }
+            let chars = iter.with_checksum::<Bech32>(hrp).with_witness_version(VERSION_0).chars();
+            buf.iter_mut().zip(chars).for_each(|(b, c)| {
+                *b = c as u8;
+                pos += 1;
+            });
         }
         version => {
-            for c in iter.with_checksum::<Bech32m>(hrp).with_witness_version(version).chars() {
-                w.write_all(&[c as u8])?;
-            }
+            let chars = iter.with_checksum::<Bech32m>(hrp).with_witness_version(version).chars();
+            buf.iter_mut().zip(chars).for_each(|(b, c)| {
+                *b = c as u8;
+                pos += 1;
+            });
         }
     }
+
+    w.write_all(&buf[..pos])?;
+
     Ok(())
 }
 
@@ -247,20 +279,29 @@ pub fn encode_upper_to_writer_unchecked<W: std::io::Write>(
     witness_version: Fe32,
     witness_program: &[u8],
 ) -> std::io::Result<()> {
+    let mut buf = [0u8; BECH32_MAX_LENGTH];
+    let mut pos = 0;
+
     let iter = witness_program.iter().copied().bytes_to_fes();
     match witness_version {
         VERSION_0 => {
-            for c in iter.with_checksum::<Bech32>(hrp).with_witness_version(VERSION_0).chars() {
-                w.write_all(&[c.to_ascii_uppercase() as u8])?;
-            }
+            let chars = iter.with_checksum::<Bech32>(hrp).with_witness_version(VERSION_0).chars();
+            buf.iter_mut().zip(chars).for_each(|(b, c)| {
+                *b = c.to_ascii_uppercase() as u8;
+                pos += 1;
+            });
         }
         version => {
-            for c in iter.with_checksum::<Bech32m>(hrp).with_witness_version(version).chars() {
-                w.write_all(&[c.to_ascii_uppercase() as u8])?;
-            }
+            let chars = iter.with_checksum::<Bech32m>(hrp).with_witness_version(version).chars();
+            buf.iter_mut().zip(chars).for_each(|(b, c)| {
+                *b = c.to_ascii_uppercase() as u8;
+                pos += 1;
+            });
         }
     }
 
+    w.write_all(&buf[..pos])?;
+
     Ok(())
 }