feat(array): implement TrustedLen for ArrayIterator and `ArrayImp…

…lIterator` (#2309)

* impl

* add test

src/common/src/array/iterator.rs

@@ -12,6 +12,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use std::iter::TrustedLen;
+
 use super::Array;
 use crate::array::ArrayImpl;
 use crate::types::DatumRef;
@@ -27,6 +29,8 @@ impl<'a, A: Array> ArrayIterator<'a, A> {
     }
 }
 
+unsafe impl<'a, A: Array> TrustedLen for ArrayIterator<'a, A> {}
+
 impl<'a, A: Array> Iterator for ArrayIterator<'a, A> {
     type Item = Option<A::RefItem<'a>>;
 
@@ -39,6 +43,11 @@ impl<'a, A: Array> Iterator for ArrayIterator<'a, A> {
             Some(item)
         }
     }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        let size = self.data.len() - self.pos;
+        (size, Some(size))
+    }
 }
 
 pub struct ArrayImplIterator<'a> {
@@ -52,6 +61,8 @@ impl<'a> ArrayImplIterator<'a> {
     }
 }
 
+unsafe impl<'a> TrustedLen for ArrayImplIterator<'a> {}
+
 impl<'a> Iterator for ArrayImplIterator<'a> {
     type Item = DatumRef<'a>;
 
@@ -64,4 +75,52 @@ impl<'a> Iterator for ArrayImplIterator<'a> {
             Some(item)
         }
     }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        let size = self.data.len() - self.pos;
+        (size, Some(size))
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use paste::paste;
+
+    use super::*;
+    use crate::array::ArrayBuilder;
+    use crate::for_all_variants;
+
+    macro_rules! test_trusted_len {
+        ([], $( { $variant_name:ident, $suffix_name:ident, $array:ty, $builder:ty } ),*) => {
+            $(
+                paste! {
+                    #[test]
+                    fn [<test_trusted_len_for_ $suffix_name _array>]() {
+                        use crate::array::$builder;
+                        let mut builder = $builder::new(3).unwrap();
+                        for _ in 0..3 {
+                            builder.append_null().unwrap();
+                        }
+                        let array = builder.finish().unwrap();
+                        let mut iter = array.iter();
+
+                        assert_eq!(iter.size_hint(), (3, Some(3))); iter.next();
+                        assert_eq!(iter.size_hint(), (2, Some(2))); iter.next();
+                        assert_eq!(iter.size_hint(), (1, Some(1))); iter.next();
+                        assert_eq!(iter.size_hint(), (0, Some(0)));
+
+                        let array_impl = ArrayImpl::from(array);
+                        let mut iter = array_impl.iter();
+
+                        assert_eq!(iter.size_hint(), (3, Some(3))); iter.next();
+                        assert_eq!(iter.size_hint(), (2, Some(2))); iter.next();
+                        assert_eq!(iter.size_hint(), (1, Some(1))); iter.next();
+                        assert_eq!(iter.size_hint(), (0, Some(0)));
+                    }
+                }
+            )*
+        };
+    }
+
+    for_all_variants! {test_trusted_len}
 }

src/common/src/lib.rs

@@ -33,6 +33,7 @@
 #![feature(fn_traits)]
 #![feature(type_alias_impl_trait)]
 #![feature(test)]
+#![feature(trusted_len)]
 
 #[macro_use]
 pub mod error;