Skip to content

Commit

Permalink
Copy ReadExt/WriteExt from bitcoin
Browse files Browse the repository at this point in the history
Otherwise we miss them when upgrading to 0.32
  • Loading branch information
RCasatta committed Aug 19, 2024
1 parent 4aad862 commit c08382c
Show file tree
Hide file tree
Showing 4 changed files with 240 additions and 84 deletions.
73 changes: 11 additions & 62 deletions src/encode.rs
Original file line number Diff line number Diff line change
Expand Up @@ -27,9 +27,6 @@ use crate::transaction::{Transaction, TxIn, TxOut};

pub use bitcoin::{self, consensus::encode::MAX_VEC_SIZE};

// Use the ReadExt/WriteExt traits as is from upstream
pub use bitcoin::consensus::encode::{ReadExt, WriteExt};

use crate::taproot::TapLeafHash;

/// Encoding error
Expand Down Expand Up @@ -207,7 +204,7 @@ impl Decodable for sha256::Midstate {
}
}

pub(crate) fn consensus_encode_with_size<S: io::Write>(
pub(crate) fn consensus_encode_with_size<S: crate::WriteExt>(
data: &[u8],
mut s: S,
) -> Result<usize, Error> {
Expand Down Expand Up @@ -245,64 +242,16 @@ impl Decodable for crate::locktime::Time {
}
}

// TODO reuse bitcoin's `WriteExt::emit_varint`, `ReadExt::read_varint` when available

/// A variable sized integer.
pub struct VarInt(pub u64);
impl Encodable for VarInt {
fn consensus_encode<W: io::Write>(&self, mut e: W) -> Result<usize, Error> {
match self.0 {
i @ 0..=0xFC => {
e.emit_u8(i as u8)?;
Ok(1)
}
i @ 0xFD..=0xFFFF => {
e.emit_u8(0xFD)?;
e.emit_u16(i as u16)?;
Ok(3)
}
i @ 0x10000..=0xFFFFFFFF => {
e.emit_u8(0xFE)?;
e.emit_u32(i as u32)?;
Ok(5)
}
i => {
e.emit_u8(0xFF)?;
e.emit_u64(i)?;
Ok(9)
}
}
fn consensus_encode<W: crate::WriteExt>(&self, mut e: W) -> Result<usize, Error> {
Ok(e.emit_varint(self.0)?)
}
}
impl Decodable for VarInt {
fn consensus_decode<D: io::Read>(mut d: D) -> Result<Self, Error> {
match d.read_u8()? {
0xFF => {
let x = d.read_u64()?;
if x < 0x100000000 {
Err(Error::NonMinimalVarInt)
} else {
Ok(VarInt(x))
}
}
0xFE => {
let x = d.read_u32()?;
if x < 0x10000 {
Err(Error::NonMinimalVarInt)
} else {
Ok(VarInt(x as u64))
}
}
0xFD => {
let x = d.read_u16()?;
if x < 0xFD {
Err(Error::NonMinimalVarInt)
} else {
Ok(VarInt(x as u64))
}
}
n => Ok(VarInt(n as u64)),
}
fn consensus_decode<D: crate::ReadExt>(mut d: D) -> Result<Self, Error> {
Ok(VarInt(d.read_varint()?))
}
}
impl VarInt {
Expand All @@ -321,14 +270,14 @@ impl VarInt {
macro_rules! impl_int {
($ty:ident, $meth_dec:ident, $meth_enc:ident) => {
impl Encodable for $ty {
fn consensus_encode<W: io::Write>(&self, mut w: W) -> Result<usize, Error> {
fn consensus_encode<W: crate::WriteExt>(&self, mut w: W) -> Result<usize, Error> {
w.$meth_enc(*self)?;
Ok(mem::size_of::<$ty>())
}
}
impl Decodable for $ty {
fn consensus_decode<R: io::Read>(mut r: R) -> Result<Self, Error> {
Ok(ReadExt::$meth_dec(&mut r)?)
fn consensus_decode<R: crate::ReadExt>(mut r: R) -> Result<Self, Error> {
crate::ReadExt::$meth_dec(&mut r)
}
}
};
Expand Down Expand Up @@ -411,7 +360,7 @@ macro_rules! impl_array {
( $size:literal ) => {
impl Encodable for [u8; $size] {
#[inline]
fn consensus_encode<W: WriteExt>(
fn consensus_encode<W: crate::WriteExt>(
&self,
mut w: W,
) -> core::result::Result<usize, Error> {
Expand All @@ -422,7 +371,7 @@ macro_rules! impl_array {

impl Decodable for [u8; $size] {
#[inline]
fn consensus_decode<R: ReadExt>(mut r: R) -> core::result::Result<Self, Error> {
fn consensus_decode<R: crate::ReadExt>(mut r: R) -> core::result::Result<Self, Error> {
let mut ret = [0; $size];
r.read_slice(&mut ret)?;
Ok(ret)
Expand Down Expand Up @@ -452,7 +401,7 @@ impl Encodable for Vec<u8> {
}
}
impl Decodable for Vec<u8> {
fn consensus_decode<D: io::Read>(mut d: D) -> Result<Self, Error> {
fn consensus_decode<D: crate::ReadExt>(mut d: D) -> Result<Self, Error> {
let s = VarInt::consensus_decode(&mut d)?.0 as usize;
let mut v = vec![0; s];
d.read_slice(&mut v)?;
Expand Down
194 changes: 194 additions & 0 deletions src/ext.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,194 @@
use std::io;

use crate::encode;

/// Extensions of `Write` to encode data as per Bitcoin consensus.
pub trait WriteExt: io::Write {
/// Outputs a 64-bit unsigned integer.
fn emit_u64(&mut self, v: u64) -> Result<(), io::Error>;
/// Outputs a 32-bit unsigned integer.
fn emit_u32(&mut self, v: u32) -> Result<(), io::Error>;
/// Outputs a 16-bit unsigned integer.
fn emit_u16(&mut self, v: u16) -> Result<(), io::Error>;
/// Outputs an 8-bit unsigned integer.
fn emit_u8(&mut self, v: u8) -> Result<(), io::Error>;

/// Outputs a 64-bit signed integer.
fn emit_i64(&mut self, v: i64) -> Result<(), io::Error>;
/// Outputs a 32-bit signed integer.
fn emit_i32(&mut self, v: i32) -> Result<(), io::Error>;
/// Outputs a 16-bit signed integer.
fn emit_i16(&mut self, v: i16) -> Result<(), io::Error>;
/// Outputs an 8-bit signed integer.
fn emit_i8(&mut self, v: i8) -> Result<(), io::Error>;

/// Outputs a variable sized integer.
fn emit_varint(&mut self, v: u64) -> Result<usize, io::Error>;

/// Outputs a boolean.
fn emit_bool(&mut self, v: bool) -> Result<(), io::Error>;

/// Outputs a byte slice.
fn emit_slice(&mut self, v: &[u8]) -> Result<usize, io::Error>;
}

/// Extensions of `Read` to decode data as per Bitcoin consensus.
pub trait ReadExt: io::Read {
/// Reads a 64-bit unsigned integer.
fn read_u64(&mut self) -> Result<u64, encode::Error>;
/// Reads a 32-bit unsigned integer.
fn read_u32(&mut self) -> Result<u32, encode::Error>;
/// Reads a 16-bit unsigned integer.
fn read_u16(&mut self) -> Result<u16, encode::Error>;
/// Reads an 8-bit unsigned integer.
fn read_u8(&mut self) -> Result<u8, encode::Error>;

/// Reads a 64-bit signed integer.
fn read_i64(&mut self) -> Result<i64, encode::Error>;
/// Reads a 32-bit signed integer.
fn read_i32(&mut self) -> Result<i32, encode::Error>;
/// Reads a 16-bit signed integer.
fn read_i16(&mut self) -> Result<i16, encode::Error>;
/// Reads an 8-bit signed integer.
fn read_i8(&mut self) -> Result<i8, encode::Error>;

/// Reads a variable sized integer.
fn read_varint(&mut self) -> Result<u64, encode::Error>;

/// Reads a boolean.
fn read_bool(&mut self) -> Result<bool, encode::Error>;

/// Reads a byte slice.
fn read_slice(&mut self, slice: &mut [u8]) -> Result<(), encode::Error>;
}

macro_rules! encoder_fn {
($name:ident, $val_type:ty) => {
#[inline]
fn $name(&mut self, v: $val_type) -> core::result::Result<(), io::Error> {
self.write_all(&v.to_le_bytes())
}
};
}

macro_rules! decoder_fn {
($name:ident, $val_type:ty, $byte_len: expr) => {
#[inline]
fn $name(&mut self) -> core::result::Result<$val_type, encode::Error> {
let mut val = [0; $byte_len];
self.read_exact(&mut val[..]).map_err(encode::Error::Io)?;
Ok(<$val_type>::from_le_bytes(val))
}
};
}

impl<W: io::Write + ?Sized> WriteExt for W {
encoder_fn!(emit_u64, u64);
encoder_fn!(emit_u32, u32);
encoder_fn!(emit_u16, u16);
encoder_fn!(emit_i64, i64);
encoder_fn!(emit_i32, i32);
encoder_fn!(emit_i16, i16);

#[inline]
fn emit_i8(&mut self, v: i8) -> Result<(), io::Error> {
self.write_all(&[v as u8])
}
#[inline]
fn emit_u8(&mut self, v: u8) -> Result<(), io::Error> {
self.write_all(&[v])
}
#[inline]
fn emit_bool(&mut self, v: bool) -> Result<(), io::Error> {
self.write_all(&[v as u8])
}
#[inline]
fn emit_slice(&mut self, v: &[u8]) -> Result<usize, io::Error> {
self.write_all(v)?;
Ok(v.len())
}
#[inline]
fn emit_varint(&mut self, v: u64) -> Result<usize, io::Error> {
match v {
i @ 0..=0xFC => {
self.emit_u8(i as u8)?;
Ok(1)
}
i @ 0xFD..=0xFFFF => {
self.emit_u8(0xFD)?;
self.emit_u16(i as u16)?;
Ok(3)
}
i @ 0x10000..=0xFFFFFFFF => {
self.emit_u8(0xFE)?;
self.emit_u32(i as u32)?;
Ok(5)
}
i => {
self.emit_u8(0xFF)?;
self.emit_u64(i)?;
Ok(9)
}
}
}
}

impl<R: io::Read + ?Sized> ReadExt for R {
decoder_fn!(read_u64, u64, 8);
decoder_fn!(read_u32, u32, 4);
decoder_fn!(read_u16, u16, 2);
decoder_fn!(read_i64, i64, 8);
decoder_fn!(read_i32, i32, 4);
decoder_fn!(read_i16, i16, 2);

#[inline]
fn read_u8(&mut self) -> Result<u8, encode::Error> {
let mut slice = [0u8; 1];
self.read_exact(&mut slice)?;
Ok(slice[0])
}
#[inline]
fn read_i8(&mut self) -> Result<i8, encode::Error> {
let mut slice = [0u8; 1];
self.read_exact(&mut slice)?;
Ok(slice[0] as i8)
}
#[inline]
fn read_bool(&mut self) -> Result<bool, encode::Error> {
ReadExt::read_i8(self).map(|bit| bit != 0)
}
#[inline]
fn read_slice(&mut self, slice: &mut [u8]) -> Result<(), encode::Error> {
self.read_exact(slice).map_err(encode::Error::Io)
}
#[inline]
fn read_varint(&mut self) -> Result<u64, encode::Error> {
match self.read_u8()? {
0xFF => {
let x = self.read_u64()?;
if x < 0x100000000 {
Err(encode::Error::NonMinimalVarInt)
} else {
Ok(x)
}
}
0xFE => {
let x = self.read_u32()?;
if x < 0x10000 {
Err(encode::Error::NonMinimalVarInt)
} else {
Ok(x as u64)
}
}
0xFD => {
let x = self.read_u16()?;
if x < 0xFD {
Err(encode::Error::NonMinimalVarInt)
} else {
Ok(x as u64)
}
}
n => Ok(n as u64),
}
}
}
Loading

0 comments on commit c08382c

Please sign in to comment.