cas: implement split algo [part #2]

Summary: implement split algo

Reviewed By: muirdm

Differential Revision: D62855251

fbshipit-source-id: 2e24b94eb2fa681cfb46907d8abba0aaf8075724

Author: Liubov Dmitrieva

eden/scm/lib/cas-client/re-cas-common/src/lib.rs

@@ -88,3 +88,78 @@ impl $crate::CasClient for $struct {
 }
 };
 }
+
+pub fn split_up_to_max_bytes(
+    digests: &[CasDigest],
+    max_bytes: u64,
+) -> impl Iterator<Item = &[CasDigest]> {
+    struct Iter<'a> {
+        pub digests: &'a [CasDigest],
+        pub max_bytes: u64,
+    }
+
+    impl<'a> Iterator for Iter<'a> {
+        type Item = &'a [CasDigest];
+
+        fn next(&mut self) -> Option<&'a [CasDigest]> {
+            if self.digests.is_empty() {
+                return None;
+            }
+            let mut split_at = 0;
+            let mut sum = 0;
+            for (i, digest) in self.digests.iter().enumerate() {
+                sum += digest.size;
+                if sum > self.max_bytes {
+                    break;
+                }
+                split_at = i + 1;
+            }
+            if split_at == 0 {
+                // We didn't find a split point meaning that there is a single file above the threshold,
+                // so just return this first item.
+                split_at = 1;
+            }
+            let (batch, rest) = self.digests.split_at(split_at);
+            self.digests = rest;
+            Some(batch)
+        }
+    }
+    Iter { digests, max_bytes }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::str::FromStr;
+
+    use super::*;
+
+    fn split_and_format_results(digests: &[CasDigest], max_bytes: u64) -> String {
+        split_up_to_max_bytes(digests, max_bytes)
+            .map(|v| {
+                v.iter()
+                    .map(|d| d.size.to_string())
+                    .collect::<Vec<_>>()
+                    .join(",")
+            })
+            .collect::<Vec<_>>()
+            .join("|")
+    }
+
+    #[test]
+    fn test_split_up_to_max_bytes() {
+        let hash =
+            Blake3::from_str("2078b4229b5353de0268efc7f64b68f3c99fb8829e9c052117b4e1e090b2603a")
+                .unwrap();
+        let digests = vec![
+            CasDigest { hash, size: 200 },
+            CasDigest { hash, size: 200 },
+            CasDigest { hash, size: 400 },
+        ];
+
+        assert_eq!(split_and_format_results(&digests, 200), "200|200|400");
+        assert_eq!(split_and_format_results(&digests, 400), "200,200|400");
+        assert_eq!(split_and_format_results(&digests, 500), "200,200|400");
+        assert_eq!(split_and_format_results(&digests, 5000), "200,200,400");
+        assert_eq!(split_and_format_results(&digests, 10), "200|200|400");
+    }
+}