use crate::bitswap_pb;
use crate::block::Block;
use crate::error::BitswapError;
use crate::prefix::Prefix;
use core::convert::TryFrom;
use hash_hasher::{HashedMap, HashedSet};
use libipld::Cid;
use prost::Message as ProstMessage;
use std::mem;
pub type Priority = i32;
#[derive(Debug, Default)]
pub struct Ledger {
sent_want_list: HashedMap<Cid, Priority>,
pub(crate) received_want_list: HashedMap<Cid, Priority>,
message: Message,
}
impl Ledger {
pub fn new() -> Self {
Self::default()
}
pub fn add_block(&mut self, block: Block) {
self.message.add_block(block);
}
pub fn want_block(&mut self, cid: &Cid, priority: Priority) {
self.message.want_block(cid, priority);
}
pub fn cancel_block(&mut self, cid: &Cid) {
self.message.cancel_block(cid);
}
pub fn wantlist(&self) -> Vec<(Cid, Priority)> {
self.received_want_list
.iter()
.map(|(cid, prio)| (*cid, *prio))
.collect()
}
pub fn send(&mut self) -> Option<Message> {
if self.message.is_empty() {
return None;
}
for cid in self.message.cancel() {
self.sent_want_list.remove(cid);
}
for (cid, priority) in self.message.want() {
self.sent_want_list.insert(*cid, *priority);
}
Some(mem::take(&mut self.message))
}
}
#[derive(Clone, PartialEq, Eq, Default)]
pub struct Message {
want: HashedMap<Cid, Priority>,
cancel: HashedSet<Cid>,
full: bool,
pub(crate) blocks: Vec<Block>,
}
impl Message {
pub fn is_empty(&self) -> bool {
self.want.is_empty() && self.cancel.is_empty() && self.blocks.is_empty()
}
pub fn blocks(&self) -> &[Block] {
&self.blocks
}
pub fn want(&self) -> &HashedMap<Cid, Priority> {
&self.want
}
pub fn cancel(&self) -> &HashedSet<Cid> {
&self.cancel
}
pub fn add_block(&mut self, block: Block) {
self.blocks.push(block);
}
pub fn remove_block(&mut self, cid: &Cid) {
self.blocks.retain(|block| block.cid() != cid);
}
pub fn want_block(&mut self, cid: &Cid, priority: Priority) {
self.want.insert(cid.to_owned(), priority);
}
pub fn cancel_block(&mut self, cid: &Cid) {
self.cancel.insert(cid.to_owned());
}
#[allow(unused)]
pub fn remove_want_block(&mut self, cid: &Cid) {
self.want.remove(cid);
}
}
impl From<&Message> for Vec<u8> {
fn from(val: &Message) -> Self {
let mut proto = bitswap_pb::Message::default();
let mut wantlist = bitswap_pb::message::Wantlist::default();
for (cid, priority) in val.want() {
let entry = bitswap_pb::message::wantlist::Entry {
block: cid.to_bytes(),
priority: *priority,
..Default::default()
};
wantlist.entries.push(entry);
}
for cid in val.cancel() {
let entry = bitswap_pb::message::wantlist::Entry {
block: cid.to_bytes(),
cancel: true,
..Default::default()
};
wantlist.entries.push(entry);
}
for block in val.blocks() {
let payload = bitswap_pb::message::Block {
prefix: Prefix::from(block.cid()).to_bytes(),
data: block.data().to_vec(),
};
proto.payload.push(payload);
}
if !wantlist.entries.is_empty() {
proto.wantlist = Some(wantlist);
}
let mut res = Vec::with_capacity(proto.encoded_len());
proto
.encode(&mut res)
.expect("there is no situation in which the protobuf message can be invalid");
res
}
}
impl Message {
pub fn to_bytes(&self) -> Vec<u8> {
self.into()
}
pub fn from_bytes(bytes: &[u8]) -> Result<Self, BitswapError> {
Self::try_from(bytes)
}
}
impl From<()> for Message {
fn from(_: ()) -> Self {
Default::default()
}
}
impl TryFrom<&[u8]> for Message {
type Error = BitswapError;
fn try_from(bytes: &[u8]) -> Result<Self, Self::Error> {
let proto: bitswap_pb::Message = bitswap_pb::Message::decode(bytes)?;
let mut message = Message::default();
for entry in proto.wantlist.unwrap_or_default().entries {
let cid = Cid::try_from(entry.block)?;
if entry.cancel {
message.cancel_block(&cid);
} else {
message.want_block(&cid, entry.priority);
}
}
for payload in proto.payload {
let prefix = Prefix::new(&payload.prefix)?;
let cid = prefix.to_cid(&payload.data)?;
let block = Block::new(cid, payload.data)?;
message.add_block(block);
}
Ok(message)
}
}
impl std::fmt::Debug for Message {
fn fmt(&self, fmt: &mut std::fmt::Formatter) -> Result<(), std::fmt::Error> {
let mut first = true;
for (cid, priority) in self.want() {
if first {
first = false;
} else {
write!(fmt, ", ")?;
}
write!(fmt, "want: {cid} {priority}")?;
}
for cid in self.cancel() {
if first {
first = false;
} else {
write!(fmt, ", ")?;
}
write!(fmt, "cancel: {cid}")?;
}
for block in self.blocks() {
if first {
first = false;
} else {
write!(fmt, ", ")?;
}
write!(fmt, "block: {}", block.cid())?;
}
if first {
write!(fmt, "(empty message)")?;
}
Ok(())
}
}