From e5753e943bd929e0cc5b07bb22ae2e5002c6123b Mon Sep 17 00:00:00 2001
From: Andrii Dmytrenko <refresh.xss@gmail.com>
Date: Mon, 18 Mar 2024 18:23:15 +0000
Subject: [PATCH] Expose `distribute_indexes_ruby`

---
 src/distribute_indexes.rs | 12 ++++++++++++
 src/lib.rs                |  1 +
 src/ruby_ext.rs           | 15 +++++++++------
 3 files changed, 22 insertions(+), 6 deletions(-)

diff --git a/src/distribute_indexes.rs b/src/distribute_indexes.rs
index 2d6736c..51b0535 100644
--- a/src/distribute_indexes.rs
+++ b/src/distribute_indexes.rs
@@ -118,4 +118,16 @@ mod tests {
             vec![0, 3, 4, 1, 6, 5, 2, 7, 8, 10, 9, 12, 11, 13, 14]
         );
     }
+
+    #[test]
+    fn test_distribute2() {
+        let data = vec!["Picture", "Post", "Video", "Video", "Picture", "Post", "Picture", "Picture", "Video"];
+        let result = distribute_ids(&data, 3);
+        let result_with_labels = result
+            .iter()
+            .map(|idx| data[*idx])
+            .collect::<Vec<_>>();
+        assert_eq!(result_with_labels, vec!["Picture", "Post", "Video", "Picture", "Post", "Video", "Picture", "Video", "Picture"]);
+        assert_eq!(result, vec![0, 1, 2, 4, 5, 3, 6, 8, 7]);
+    }
 }
diff --git a/src/lib.rs b/src/lib.rs
index aa58c53..2e617b3 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -7,4 +7,5 @@ mod ruby_ext;
 
 pub use distribute_csv::distribute as distribute_csv;
 pub use distribute_indexes::distribute;
+pub use distribute_indexes::distribute_ids;
 pub use distributing_iterator::DistributingIterator;
diff --git a/src/ruby_ext.rs b/src/ruby_ext.rs
index ecd892a..1b1c09a 100644
--- a/src/ruby_ext.rs
+++ b/src/ruby_ext.rs
@@ -1,19 +1,22 @@
-use magnus::{error::Result, function, exception, Error};
+use magnus::{error::Result, exception, function, Error, Value};
 
-use crate::distribute_csv;
+use crate::{distribute_ids, distribute_csv};
 
 fn distribute_csv_ruby(data: String, field: String, spread: u64) -> Result<String> {
     match distribute_csv(&data, &field, spread) {
         Ok(result) => Ok(result),
-        Err(e) => {
-            Err(Error::new(exception::standard_error(), format!("{:?}", e)))
-        }
+        Err(e) => Err(Error::new(exception::standard_error(), format!("{:?}", e))),
     }
 }
 
+fn distribute_indexes_ruby(data: Vec<String>, spread: usize) -> Result<Vec<usize>> {
+    Ok(distribute_ids(&data, spread))
+}
+
 #[magnus::init]
 fn init(ruby: &magnus::Ruby) -> Result<()> {
     let module = ruby.define_module("DistributingIterator")?;
-    module.define_module_function("distribute", function!(distribute_csv_ruby, 3))?;
+    module.define_module_function("distribute_csv", function!(distribute_csv_ruby, 3))?;
+    module.define_module_function("distribute_indexes", function!(distribute_indexes_ruby, 2))?;
     Ok(())
 }