From 54651c97c5d382c4b9960e0eb05bd0902d4efd55 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jon=C3=A1=C5=A1=20Fiala?= <jonas.fiala@inf.ethz.ch>
Date: Mon, 13 Nov 2023 19:16:19 +0100
Subject: [PATCH] Make top crates iterator

---
 mir-state-analysis/tests/top_crates.rs | 73 +++++++++++++++++---------
 1 file changed, 49 insertions(+), 24 deletions(-)

diff --git a/mir-state-analysis/tests/top_crates.rs b/mir-state-analysis/tests/top_crates.rs
index 84dee437ac0..65e4c1fdb5a 100644
--- a/mir-state-analysis/tests/top_crates.rs
+++ b/mir-state-analysis/tests/top_crates.rs
@@ -17,8 +17,8 @@ fn get(url: &str) -> reqwest::Result<reqwest::blocking::Response> {
 
 pub fn top_crates_range(range: std::ops::Range<usize>) {
     std::fs::create_dir_all("tmp").unwrap();
-    let top_crates = top_crates_by_download_count(range.end - 1);
-    for (i, krate) in top_crates.into_iter().enumerate().skip(range.start) {
+    let top_crates = CratesIter::top(range);
+    for (i, krate) in top_crates {
         let version = krate.version.unwrap_or(krate.newest_version);
         println!("Starting: {i} ({})", krate.name);
         run_on_crate(&krate.name, &version);
@@ -92,27 +92,52 @@ struct CratesList {
     crates: Vec<Crate>,
 }
 
-/// Create a list of top ``count`` crates.
-fn top_crates_by_download_count(mut count: usize) -> Vec<Crate> {
-    const PAGE_SIZE: usize = 100;
-    let page_count = count / PAGE_SIZE + 2;
-    let mut sources = Vec::new();
-    for page in 1..page_count {
-        let url = format!(
-            "https://crates.io/api/v1/crates?page={page}&per_page={PAGE_SIZE}&sort=downloads"
-        );
-        let resp = get(&url).expect("Could not fetch top crates");
-        assert!(
-            resp.status().is_success(),
-            "Response status: {}",
-            resp.status()
-        );
-        let page_crates: CratesList = match serde_json::from_reader(resp) {
-            Ok(page_crates) => page_crates,
-            Err(e) => panic!("Invalid JSON {e}"),
-        };
-        sources.extend(page_crates.crates.into_iter().take(count));
-        count -= std::cmp::min(PAGE_SIZE, count);
+const PAGE_SIZE: usize = 100;
+struct CratesIter {
+    curr_idx: usize,
+    curr_page: usize,
+    crates: Vec<Crate>,
+}
+
+impl CratesIter {
+    pub fn new(start: usize) -> Self {
+        Self {
+            curr_idx: start,
+            curr_page: start / PAGE_SIZE + 1,
+            crates: Vec::new(),
+        }
+    }
+    pub fn top(range: std::ops::Range<usize>) -> impl Iterator<Item = (usize, Crate)> {
+        Self::new(range.start).take(range.len())
+    }
+}
+
+impl Iterator for CratesIter {
+    type Item = (usize, Crate);
+    fn next(&mut self) -> Option<Self::Item> {
+        if self.crates.is_empty() {
+            let url = format!(
+                "https://crates.io/api/v1/crates?page={}&per_page={PAGE_SIZE}&sort=downloads",
+                self.curr_page,
+            );
+            let resp = get(&url).expect("Could not fetch top crates");
+            assert!(
+                resp.status().is_success(),
+                "Response status: {}",
+                resp.status()
+            );
+            let page_crates: CratesList = match serde_json::from_reader(resp) {
+                Ok(page_crates) => page_crates,
+                Err(e) => panic!("Invalid JSON {e}"),
+            };
+            assert_eq!(page_crates.crates.len(), PAGE_SIZE);
+            self.crates = page_crates.crates;
+            self.crates.reverse();
+            self.crates
+                .truncate(self.crates.len() - self.curr_idx % PAGE_SIZE);
+            self.curr_page += 1;
+        }
+        self.curr_idx += 1;
+        Some((self.curr_idx - 1, self.crates.pop().unwrap()))
     }
-    sources
 }