From be726e9aa3bbf154558d8e5ea5459e1b89df7aa1 Mon Sep 17 00:00:00 2001 From: Igor Katson Date: Mon, 10 Jul 2023 14:47:40 +0100 Subject: [PATCH] Just messing around with itertools Either --- Cargo.lock | 16 +++++ crates/librqbit_core/Cargo.toml | 3 +- crates/librqbit_core/src/torrent_metainfo.rs | 70 ++++++++------------ 3 files changed, 45 insertions(+), 44 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 9116077..e2eefc0 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -438,6 +438,12 @@ dependencies = [ "winapi", ] +[[package]] +name = "either" +version = "1.8.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7fcaabb2fef8c910e7f4c7ce9f67a1283a1715879a7c230ca9d6d1ae31f16d91" + [[package]] name = "encoding_rs" version = "0.8.32" @@ -860,6 +866,15 @@ dependencies = [ "windows-sys", ] +[[package]] +name = "itertools" +version = "0.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b1c173a5686ce8bfa551b3563d0c2170bf24ca44da99c7ca4bfdab5418c3fe57" +dependencies = [ + "either", +] + [[package]] name = "itoa" version = "1.0.8" @@ -935,6 +950,7 @@ dependencies = [ "buffers", "clone_to_owned", "hex 0.4.3", + "itertools", "log", "parking_lot", "serde", diff --git a/crates/librqbit_core/Cargo.toml b/crates/librqbit_core/Cargo.toml index 9835643..5ff5ee8 100644 --- a/crates/librqbit_core/Cargo.toml +++ b/crates/librqbit_core/Cargo.toml @@ -15,4 +15,5 @@ parking_lot = "0.12" serde = {version = "1", features=["derive"]} buffers = {path="../buffers"} bencode = {path = "../bencode", default-features=false} -clone_to_owned = {path="../clone_to_owned"} \ No newline at end of file +clone_to_owned = {path="../clone_to_owned"} +itertools = "0.11.0" diff --git a/crates/librqbit_core/src/torrent_metainfo.rs b/crates/librqbit_core/src/torrent_metainfo.rs index 83c8861..d712b2a 100644 --- a/crates/librqbit_core/src/torrent_metainfo.rs +++ b/crates/librqbit_core/src/torrent_metainfo.rs @@ -1,9 +1,10 @@ -use std::path::PathBuf; +use std::{iter::once, path::PathBuf}; use anyhow::Context; use bencode::BencodeDeserializer; use buffers::{ByteBuf, ByteString}; use clone_to_owned::CloneToOwned; +use itertools::Either; use serde::Deserialize; use crate::id20::Id20; @@ -46,7 +47,7 @@ pub struct TorrentMetaV1 { impl TorrentMetaV1 { pub fn iter_announce(&self) -> impl Iterator { - std::iter::once(&self.announce).chain(self.announce_list.iter().flatten()) + once(&self.announce).chain(self.announce_list.iter().flatten()) } } @@ -112,26 +113,13 @@ impl<'a, ByteBuf> FileIteratorName<'a, ByteBuf> { where ByteBuf: AsRef<[u8]>, { - let single_it = std::iter::once(match self { - FileIteratorName::Single(n) => Some(*n), - FileIteratorName::Tree(_) => None, - }); - let multi_it = match self { - FileIteratorName::Single(_) => &[], - FileIteratorName::Tree(t) => *t, - } - .iter() - .map(|p| Some(Some(p))); - - let it = single_it.chain(multi_it).flatten(); - - it.map(|part| { - let part = match part { - Some(part) => part, - None => return Ok("torrent-content"), - }; - let bit = std::str::from_utf8(part.as_ref()) - .context("cannot decode filename bit as UTF-8")?; + let it = match self { + FileIteratorName::Single(None) => return Either::Left(once(Ok("torrent-content"))), + FileIteratorName::Single(Some(name)) => Either::Left(once((*name).as_ref())), + FileIteratorName::Tree(t) => Either::Right(t.iter().map(|bb| bb.as_ref())), + }; + Either::Right(it.map(|part: &'a [u8]| { + let bit = std::str::from_utf8(part).context("cannot decode filename bit as UTF-8")?; if bit == ".." { anyhow::bail!("path traversal detected, \"..\" in filename bit {:?}", bit); } @@ -143,7 +131,7 @@ impl<'a, ByteBuf> FileIteratorName<'a, ByteBuf> { ); } Ok(bit) - }) + })) } } @@ -154,6 +142,7 @@ impl> TorrentMetaV1Info { let expected_hash = self.pieces.as_ref().get(start..end)?; Some(expected_hash) } + pub fn compare_hash(&self, piece: u32, hash: [u8; 20]) -> Option { let start = piece as usize * 20; let end = start + 20; @@ -161,36 +150,31 @@ impl> TorrentMetaV1Info { Some(expected_hash == hash) } - fn is_single_file(&self) -> anyhow::Result { + pub fn iter_filenames_and_lengths( + &self, + ) -> anyhow::Result, u64)>> { match (self.length, self.files.as_ref()) { + // Single-file + (Some(length), None) => Ok(Either::Left(once(( + FileIteratorName::Single(self.name.as_ref()), + length, + )))), + + // Multi-file (None, Some(files)) => { if files.is_empty() { anyhow::bail!("expected multi-file torrent to have at least one file") } - Ok(false) + Ok(Either::Right( + files + .iter() + .map(|f| (FileIteratorName::Tree(&f.path), f.length)), + )) } - (Some(_), None) => Ok(true), _ => anyhow::bail!("torrent can't be both in single and multi-file mode"), } } - pub fn iter_filenames_and_lengths( - &self, - ) -> anyhow::Result, u64)>> { - self.is_single_file()?; - - let single_it = std::iter::once(match (self.name.as_ref(), self.length) { - (n, Some(l)) => Some((FileIteratorName::Single(n), l)), - _ => None, - }); - let multi_it = self - .files - .as_deref() - .unwrap_or_default() - .iter() - .map(|f| Some((FileIteratorName::Tree(&f.path), f.length))); - Ok(single_it.chain(multi_it).flatten()) - } pub fn iter_file_lengths(&self) -> anyhow::Result + '_> { Ok(self.iter_filenames_and_lengths()?.map(|(_, l)| l)) }