From 4a27c1369a8b2c16e598760c345080f561647db3 Mon Sep 17 00:00:00 2001
From: Sebastian Geisler <sebastian@blockstream.io>
Date: Thu, 8 Nov 2018 19:13:55 -0800
Subject: [PATCH] avoid heap use when encoding base58

---
 src/util/base58.rs | 72 ++++++++++++++++++++++++++++++++--------------
 1 file changed, 51 insertions(+), 21 deletions(-)

diff --git a/src/util/base58.rs b/src/util/base58.rs
index 28384a09..a355b42e 100644
--- a/src/util/base58.rs
+++ b/src/util/base58.rs
@@ -14,7 +14,7 @@
 
 //! Base58 encoder and decoder
 
-use std::{error, fmt, str};
+use std::{error, fmt, str, slice, iter};
 
 use byteorder::{ByteOrder, LittleEndian};
 use util::hash::Sha256dHash;
@@ -65,6 +65,44 @@ impl error::Error for Error {
     }
 }
 
+/// Vector-like object that holds the first 100 elements on the stack. If more space is needed it
+/// will be allocated on the heap.
+struct SmallVec<T> {
+    len: usize,
+    stack: [T; 100],
+    heap: Vec<T>,
+}
+
+impl<T: Default + Copy> SmallVec<T> {
+    pub fn new() -> SmallVec<T> {
+        use std::default::Default;
+        SmallVec {
+            len: 0,
+            stack: [T::default(); 100],
+            heap: Vec::new(),
+        }
+    }
+
+    pub fn push(&mut self, val: T) {
+        if self.len < 100 {
+            self.stack[self.len] = val;
+            self.len += 1;
+        } else {
+            self.heap.push(val);
+        }
+    }
+
+    pub fn iter(&self) -> iter::Chain<slice::Iter<T>, slice::Iter<T>> {
+        // If len<100 then we just append an empty vec
+        self.stack[0..self.len].iter().chain(self.heap.iter())
+    }
+
+    pub fn iter_mut(&mut self) -> iter::Chain<slice::IterMut<T>, slice::IterMut<T>> {
+        // If len<100 then we just append an empty vec
+        self.stack[0..self.len].iter_mut().chain(self.heap.iter_mut())
+    }
+}
+
 static BASE58_CHARS: &'static [u8] = b"123456789ABCDEFGHJKLMNPQRSTUVWXYZabcdefghijkmnopqrstuvwxyz";
 
 static BASE58_DIGITS: [Option<u8>; 128] = [
@@ -134,14 +172,12 @@ pub fn from_check(data: &str) -> Result<Vec<u8>, Error> {
     Ok(ret)
 }
 
-fn encode_iter_utf8<I>(data: I) -> Vec<u8>
+fn format_iter<I, W>(writer: &mut W, data: I) -> Result<(), fmt::Error>
 where
     I: Iterator<Item = u8> + Clone,
+    W: fmt::Write
 {
-    let (len, _) = data.size_hint();
-
-    // 7/5 is just over log_58(256)
-    let mut ret = Vec::with_capacity(1 + len * 7 / 5);
+    let mut ret = SmallVec::new();
 
     let mut leading_zero_count = 0;
     let mut leading_zeroes = true;
@@ -169,29 +205,23 @@ where
     for _ in 0..leading_zero_count {
         ret.push(0);
     }
-    ret.reverse();
-    for ch in ret.iter_mut() {
-        *ch = BASE58_CHARS[*ch as usize];
+
+    for ch in ret.iter().rev() {
+        writer.write_char(BASE58_CHARS[*ch as usize] as char)?;
     }
-    ret
+
+    Ok(())
 }
 
 fn encode_iter<I>(data: I) -> String
 where
     I: Iterator<Item = u8> + Clone,
 {
-    let ret = encode_iter_utf8(data);
-    String::from_utf8(ret).unwrap()
+    let mut ret = String::new();
+    format_iter(&mut ret, data).expect("writing into string shouldn't fail");
+    ret
 }
 
-/// Directly encode a slice as base58 into a `Formatter`.
-fn encode_iter_to_fmt<I>(fmt: &mut fmt::Formatter, data: I) -> fmt::Result
-where
-    I: Iterator<Item = u8> + Clone,
-{
-    let ret = encode_iter_utf8(data);
-    fmt.write_str(str::from_utf8(&ret).unwrap())
-}
 
 /// Directly encode a slice as base58
 pub fn encode_slice(data: &[u8]) -> String {
@@ -216,7 +246,7 @@ pub fn check_encode_slice_to_fmt(fmt: &mut fmt::Formatter, data: &[u8]) -> fmt::
     let iter = data.iter()
         .cloned()
         .chain(checksum[0..4].iter().cloned());
-    encode_iter_to_fmt(fmt, iter)
+    format_iter(fmt, iter)
 }
 
 #[cfg(test)]