Add a method to update_only_files to chunk_tracker

crates/librqbit/src/chunk_tracker.rs

@@ -1,5 +1,11 @@
+use std::collections::HashSet;
+
 use anyhow::Context;
-use librqbit_core::lengths::{ChunkInfo, Lengths, ValidPieceIndex};
+use buffers::ByteBufOwned;
+use librqbit_core::{
+    lengths::{ChunkInfo, Lengths, ValidPieceIndex},
+    torrent_metainfo::{TorrentMetaV1Info, TorrentMetaV1Owned},
+};
 use peer_binary_protocol::Piece;
 use tracing::{debug, trace};
 
@@ -29,6 +35,12 @@ pub struct ChunkTracker {
     total_selected_bytes: u64,
 }
 
+#[derive(Debug, PartialEq, Eq, Clone, Copy)]
+pub struct HaveNeeded {
+    pub have_bytes: u64,
+    pub needed_bytes: u64,
+}
+
 // Comput the have-status of chunks.
 //
 // Save as "have_pieces", but there's one bit per chunk (not per piece).
@@ -227,15 +239,96 @@ impl ChunkTracker {
         }
         Some(ChunkMarkingResult::NotCompleted)
     }
+
+    // NOTE: this doesn't validate new_only_files.
+    // E.g. if there are indices there that don't make
+    // sense, they will be ignored.
+    pub fn update_only_files(
+        &mut self,
+        file_lengths_iterator: impl IntoIterator<Item = u64>,
+        // TODO: maybe make this a BF
+        new_only_files: &HashSet<usize>,
+    ) -> anyhow::Result<HaveNeeded> {
+        let mut piece_it = self.lengths.iter_piece_infos();
+        let mut current_piece = piece_it
+            .next()
+            .context("bug: iter_piece_infos() returned empty iterator")?;
+        let mut current_piece_needed = false;
+        let mut current_piece_remaining = current_piece.len;
+        let mut have_bytes = 0u64;
+        let mut needed_bytes = 0u64;
+
+        for (idx, len) in file_lengths_iterator.into_iter().enumerate() {
+            let file_required = new_only_files.contains(&idx);
+
+            let mut remaining_file_len = len;
+
+            while remaining_file_len > 0 {
+                current_piece_needed |= len > 0 && file_required;
+                let shift = std::cmp::min(current_piece_remaining as u64, remaining_file_len);
+                assert!(shift > 0);
+                remaining_file_len -= shift;
+                current_piece_remaining -= shift as u32;
+
+                dbg!(
+                    idx,
+                    shift,
+                    remaining_file_len,
+                    current_piece_remaining,
+                    current_piece_needed,
+                    file_required,
+                    current_piece
+                );
+
+                if current_piece_remaining == 0 {
+                    let current_piece_have = self.have[current_piece.piece_index.get() as usize];
+                    if current_piece_have {
+                        have_bytes += current_piece.len as u64;
+                    }
+                    if current_piece_needed {
+                        needed_bytes += current_piece.len as u64;
+                    }
+                    match (current_piece_needed, current_piece_have) {
+                        (true, true) => {}
+                        (true, false) => {
+                            dbg!(self.mark_piece_broken_if_not_have(current_piece.piece_index))
+                        }
+                        (false, true) => {}
+                        (false, false) => {
+                            // don't need the piece, and don't have it - cancel downloading it
+                            dbg!(self
+                                .needed_pieces
+                                .set(current_piece.piece_index.get() as usize, false));
+                        }
+                    }
+
+                    if current_piece.piece_index != self.lengths.last_piece_id() {
+                        current_piece = piece_it.next().context(
+                            "bug: iter_piece_infos() pieces ended earlier than expected",
+                        )?;
+                        current_piece_needed = false;
+                        current_piece_remaining = current_piece.len;
+                    }
+                }
+            }
+        }
+
+        Ok(HaveNeeded {
+            have_bytes,
+            needed_bytes,
+        })
+    }
 }
 
 #[cfg(test)]
 mod tests {
+    use std::collections::HashSet;
+
     use librqbit_core::{constants::CHUNK_SIZE, lengths::Lengths};
 
-    use crate::type_aliases::BF;
+    use crate::{chunk_tracker::HaveNeeded, type_aliases::BF};
 
-    use super::compute_chunk_have_status;
+    use super::{compute_chunk_have_status, ChunkTracker};
 
     #[test]
     fn test_compute_chunk_status() {
@@ -334,4 +427,124 @@ mod tests {
             }
         }
     }
+
+    #[test]
+    fn test_update_only_files() {
+        let piece_len = CHUNK_SIZE * 2 + 1;
+        let total_len = piece_len as u64 * 2 + 1;
+        let l = Lengths::new(total_len, piece_len).unwrap();
+        assert_eq!(l.total_pieces(), 3);
+        assert_eq!(l.total_chunks(), 7);
+
+        let all_files = [
+            piece_len as u64, // piece 0 and boundary
+            1,                // piece 1
+            0,                // piece 1 (or none)
+            piece_len as u64, // piece 1 and 2
+        ];
+
+        let bf_len = l.piece_bitfield_bytes();
+        let initial_have = BF::from_boxed_slice(vec![0u8; bf_len].into_boxed_slice());
+        let initial_needed = BF::from_boxed_slice(vec![u8::MAX; bf_len].into_boxed_slice());
+
+        // Initially, we need all files and all pieces.
+        let mut ct = ChunkTracker::new(
+            initial_needed.clone(),
+            initial_have.clone(),
+            l,
+            l.total_length(),
+        )
+        .unwrap();
+
+        // Select all file, no changes.
+        assert_eq!(
+            ct.update_only_files(all_files.into_iter(), &HashSet::from_iter([0, 1, 2, 3]))
+                .unwrap(),
+            HaveNeeded {
+                have_bytes: 0,
+                needed_bytes: total_len
+            }
+        );
+        assert_eq!(ct.have, initial_have);
+        assert_eq!(ct.needed_pieces, initial_needed);
+
+        // Select only the first file.
+        println!("Select only the first file.");
+        assert_eq!(
+            ct.update_only_files(all_files, &HashSet::from_iter([0]))
+                .unwrap(),
+            HaveNeeded {
+                have_bytes: 0,
+                needed_bytes: all_files[0],
+            }
+        );
+        assert_eq!(ct.needed_pieces[0], true);
+        assert_eq!(ct.needed_pieces[1], false);
+        assert_eq!(ct.needed_pieces[2], false);
+
+        // Select only the second file.
+        assert_eq!(
+            ct.update_only_files(all_files, &HashSet::from_iter([1]))
+                .unwrap(),
+            HaveNeeded {
+                have_bytes: 0,
+                needed_bytes: piece_len as u64,
+            }
+        );
+        assert_eq!(ct.needed_pieces[0], false);
+        assert_eq!(ct.needed_pieces[1], true);
+        assert_eq!(ct.needed_pieces[2], false);
+
+        // Select only the third file (zero sized one!).
+        assert_eq!(
+            ct.update_only_files(all_files, &HashSet::from_iter([2]))
+                .unwrap(),
+            HaveNeeded {
+                have_bytes: 0,
+                needed_bytes: 0,
+            }
+        );
+        assert_eq!(ct.needed_pieces[0], false);
+        assert_eq!(ct.needed_pieces[1], false);
+        assert_eq!(ct.needed_pieces[2], false);
+
+        // Select only the fourth file.
+        assert_eq!(
+            ct.update_only_files(all_files, &HashSet::from_iter([3]))
+                .unwrap(),
+            HaveNeeded {
+                have_bytes: 0,
+                needed_bytes: (piece_len + 1) as u64,
+            }
+        );
+        assert_eq!(ct.needed_pieces[0], false);
+        assert_eq!(ct.needed_pieces[1], true);
+        assert_eq!(ct.needed_pieces[2], true);
+
+        // Select first and last file
+        assert_eq!(
+            ct.update_only_files(all_files.clone(), &HashSet::from_iter([0, 3]))
+                .unwrap(),
+            HaveNeeded {
+                have_bytes: 0,
+                needed_bytes: all_files[0] + all_files[3] + 1,
+            }
+        );
+        assert_eq!(ct.needed_pieces[0], true);
+        assert_eq!(ct.needed_pieces[1], true);
+        assert_eq!(ct.needed_pieces[2], true);
+
+        // Select all files
+        assert_eq!(
+            ct.update_only_files(all_files.clone(), &HashSet::from_iter([0, 1, 2, 3]))
+                .unwrap(),
+            HaveNeeded {
+                have_bytes: 0,
+                needed_bytes: total_len,
+            }
+        );
+        assert_eq!(ct.needed_pieces[0], true);
+        assert_eq!(ct.needed_pieces[1], true);
+        assert_eq!(ct.needed_pieces[2], true);
+    }
 }