Implement try_from_iter()

yutannihilation · Mar 30, 2024 · e428cba · e428cba
1 parent bee3c73
commit e428cba
Show file tree

Hide file tree

Showing 7 changed files with 148 additions and 0 deletions.
diff --git a/savvy-ffi/src/lib.rs b/savvy-ffi/src/lib.rs
@@ -66,6 +66,7 @@ extern "C" {
 // Allocation and attributes
 extern "C" {
     pub fn Rf_xlength(arg1: SEXP) -> R_xlen_t;
+    pub fn SETLENGTH(x: SEXP, v: R_xlen_t);
     pub fn Rf_allocVector(arg1: SEXPTYPE, arg2: R_xlen_t) -> SEXP;
     pub fn Rf_install(arg1: *const ::std::os::raw::c_char) -> SEXP;
     pub fn Rf_getAttrib(arg1: SEXP, arg2: SEXP) -> SEXP;

diff --git a/src/sexp/complex.rs b/src/sexp/complex.rs
@@ -121,6 +121,35 @@ impl OwnedComplexSexp {
             raw,
         })
     }
+
+    pub fn try_from_iter<I>(iter: I) -> crate::error::Result<Self>
+    where
+        I: Iterator<Item = Complex64>,
+    {
+        let iter = iter.into_iter();
+        match iter.size_hint() {
+            (_, Some(upper)) => {
+                let mut out = unsafe { Self::new_without_init(upper)? };
+                let mut actual_len = 0;
+                for (i, v) in iter.enumerate() {
+                    out.set_elt(i, v)?;
+                    actual_len = i;
+                }
+
+                if actual_len != upper {
+                    unsafe {
+                        savvy_ffi::SETLENGTH(out.inner, actual_len as _);
+                    }
+                }
+
+                Ok(out)
+            }
+            (_, None) => {
+                let v: Vec<I::Item> = iter.collect();
+                v.try_into()
+            }
+        }
+    }
 }
 
 impl Drop for OwnedComplexSexp {

diff --git a/src/sexp/integer.rs b/src/sexp/integer.rs
@@ -150,6 +150,35 @@ impl OwnedIntegerSexp {
             raw,
         })
     }
+
+    pub fn try_from_iter<I>(iter: I) -> crate::error::Result<Self>
+    where
+        I: Iterator<Item = i32>,
+    {
+        let iter = iter.into_iter();
+        match iter.size_hint() {
+            (_, Some(upper)) => {
+                let mut out = unsafe { Self::new_without_init(upper)? };
+                let mut actual_len = 0;
+                for (i, v) in iter.enumerate() {
+                    out.set_elt(i, v)?;
+                    actual_len = i;
+                }
+
+                if actual_len != upper {
+                    unsafe {
+                        savvy_ffi::SETLENGTH(out.inner, actual_len as _);
+                    }
+                }
+
+                Ok(out)
+            }
+            (_, None) => {
+                let v: Vec<I::Item> = iter.collect();
+                v.try_into()
+            }
+        }
+    }
 }
 
 impl Drop for OwnedIntegerSexp {

diff --git a/src/sexp/logical.rs b/src/sexp/logical.rs
@@ -145,6 +145,35 @@ impl OwnedLogicalSexp {
             raw,
         })
     }
+
+    pub fn try_from_iter<I>(iter: I) -> crate::error::Result<Self>
+    where
+        I: Iterator<Item = bool>,
+    {
+        let iter = iter.into_iter();
+        match iter.size_hint() {
+            (_, Some(upper)) => {
+                let mut out = unsafe { Self::new_without_init(upper)? };
+                let mut actual_len = 0;
+                for (i, v) in iter.enumerate() {
+                    out.set_elt(i, v)?;
+                    actual_len = i;
+                }
+
+                if actual_len != upper {
+                    unsafe {
+                        savvy_ffi::SETLENGTH(out.inner, actual_len as _);
+                    }
+                }
+
+                Ok(out)
+            }
+            (_, None) => {
+                let v: Vec<I::Item> = iter.collect();
+                v.try_into()
+            }
+        }
+    }
 }
 
 impl Drop for OwnedLogicalSexp {

diff --git a/src/sexp/real.rs b/src/sexp/real.rs
@@ -146,6 +146,35 @@ impl OwnedRealSexp {
             raw,
         })
     }
+
+    pub fn try_from_iter<I>(iter: I) -> crate::error::Result<Self>
+    where
+        I: Iterator<Item = f64>,
+    {
+        let iter = iter.into_iter();
+        match iter.size_hint() {
+            (_, Some(upper)) => {
+                let mut out = unsafe { Self::new_without_init(upper)? };
+                let mut actual_len = 0;
+                for (i, v) in iter.enumerate() {
+                    out.set_elt(i, v)?;
+                    actual_len = i;
+                }
+
+                if actual_len != upper {
+                    unsafe {
+                        savvy_ffi::SETLENGTH(out.inner, actual_len as _);
+                    }
+                }
+
+                Ok(out)
+            }
+            (_, None) => {
+                let v: Vec<I::Item> = iter.collect();
+                v.try_into()
+            }
+        }
+    }
 }
 
 impl Drop for OwnedRealSexp {

diff --git a/src/sexp/string.rs b/src/sexp/string.rs
@@ -105,6 +105,36 @@ impl OwnedStringSexp {
 
         Ok(Self { inner, token, len })
     }
+
+    pub fn try_from_iter<I, U>(iter: I) -> crate::error::Result<Self>
+    where
+        I: Iterator<Item = U>,
+        U: AsRef<str>,
+    {
+        let iter = iter.into_iter();
+        match iter.size_hint() {
+            (_, Some(upper)) => {
+                let mut out = Self::new(upper)?;
+                let mut actual_len = 0;
+                for (i, v) in iter.enumerate() {
+                    out.set_elt(i, v.as_ref())?;
+                    actual_len = i;
+                }
+
+                if actual_len != upper {
+                    unsafe {
+                        savvy_ffi::SETLENGTH(out.inner, actual_len as _);
+                    }
+                }
+
+                Ok(out)
+            }
+            (_, None) => {
+                let v: Vec<I::Item> = iter.collect();
+                v.try_into()
+            }
+        }
+    }
 }
 
 unsafe fn str_to_charsxp(v: &str) -> crate::error::Result<SEXP> {

diff --git a/xtask/src/main.rs b/xtask/src/main.rs
@@ -84,6 +84,7 @@ fn show() -> Result<(), DynError> {
         .allowlist_function("R_IsNA")
         // Allocation and attributes
         .allowlist_function("Rf_xlength")
+        .allowlist_function("SETLENGTH")
         .allowlist_function("Rf_allocVector")
         .allowlist_function("Rf_install")
         .allowlist_function("Rf_getAttrib")