From fc91f4e32788f03d6d303fdc681704c68a3065a9 Mon Sep 17 00:00:00 2001
From: Christian Lewe <ch_lewe@agdsn.me>
Date: Tue, 21 Nov 2023 22:05:15 +0100
Subject: [PATCH] Encode: Add BitWriter::vec

We often write bits to a byte vector.
BitWriter::vec generalizes this workflow by writing the result of any
bit operation into a byte vector, which it returns.
The caller doesn't have to worry about I/O errors that can never occur.
---
 src/bit_encoding/bitwriter.rs | 39 +++++++++++++++++++++++++++++++++++
 1 file changed, 39 insertions(+)

diff --git a/src/bit_encoding/bitwriter.rs b/src/bit_encoding/bitwriter.rs
index a6a38fc0..d5a60aba 100644
--- a/src/bit_encoding/bitwriter.rs
+++ b/src/bit_encoding/bitwriter.rs
@@ -109,3 +109,42 @@ impl<W: io::Write> BitWriter<W> {
         Ok(len)
     }
 }
+
+// Implement for io::Sink to make vec() an associated function
+// Implementing for a generic W would require callers to specify W, which is annoying
+impl BitWriter<io::Sink> {
+    /// Write the result of a bit operation into a byte vector and return the vector.
+    ///
+    /// I/O to a vector never fails.
+    pub fn vec<F>(f: F) -> Vec<u8>
+    where
+        F: FnOnce(&mut BitWriter<&mut Vec<u8>>) -> io::Result<usize>,
+    {
+        let mut bytes = Vec::new();
+        let mut bits = BitWriter::new(&mut bytes);
+        f(&mut bits).expect("I/O to vector never fails");
+        bits.flush_all().expect("I/O to vector never fails");
+        bytes
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::jet::Core;
+    use crate::node::CoreConstructible;
+    use crate::ConstructNode;
+    use std::sync::Arc;
+
+    #[test]
+    fn vec() {
+        let program = Arc::<ConstructNode<Core>>::unit();
+        let _ = BitWriter::vec(|w| program.encode(w));
+    }
+
+    #[test]
+    fn empty_vec() {
+        let vec = BitWriter::vec(|_| Ok(0));
+        assert!(vec.is_empty());
+    }
+}