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 BitWriter { 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 { + /// 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) -> Vec + where + F: FnOnce(&mut BitWriter<&mut Vec>) -> io::Result, + { + 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::>::unit(); + let _ = BitWriter::vec(|w| program.encode(w)); + } + + #[test] + fn empty_vec() { + let vec = BitWriter::vec(|_| Ok(0)); + assert!(vec.is_empty()); + } +}