From 0bf6424b2995a46e827a5c2988b7184467788f90 Mon Sep 17 00:00:00 2001
From: Joel Wurtz <jwurtz@jolicode.com>
Date: Tue, 17 Oct 2023 11:42:33 +0200
Subject: [PATCH] feat(iterator): add helper for zend_object_iterator

---
 src/types/iterator.rs | 106 ++++++++++++++++++++++++++++++++++++++++++
 src/types/mod.rs      |   1 +
 src/types/object.rs   |  11 +++++
 3 files changed, 118 insertions(+)
 create mode 100644 src/types/iterator.rs

diff --git a/src/types/iterator.rs b/src/types/iterator.rs
new file mode 100644
index 0000000..7703f0e
--- /dev/null
+++ b/src/types/iterator.rs
@@ -0,0 +1,106 @@
+use crate::convert::FromZvalMut;
+use crate::ffi::zend_object_iterator;
+use crate::flags::DataType;
+use crate::types::{ZendLong, Zval};
+use std::convert::TryInto;
+
+pub type ZendIterator = zend_object_iterator;
+
+pub struct Iter<'a> {
+    zi: &'a mut ZendIterator,
+}
+
+impl ZendIterator {
+    pub fn iter(&mut self) -> Iter {
+        self.index = 0;
+        self.rewind();
+
+        Iter { zi: self }
+    }
+
+    pub fn valid(&mut self) -> bool {
+        if let Some(valid) = unsafe { (*self.funcs).valid } {
+            unsafe { valid(&mut *self) != 0 }
+        } else {
+            true
+        }
+    }
+
+    pub fn rewind(&mut self) {
+        if let Some(rewind) = unsafe { (*self.funcs).rewind } {
+            unsafe {
+                rewind(&mut *self);
+            }
+        }
+    }
+
+    pub fn move_forward(&mut self) {
+        if let Some(move_forward) = unsafe { (*self.funcs).move_forward } {
+            unsafe {
+                move_forward(&mut *self);
+            }
+        }
+    }
+
+    pub fn get_current_data<'a>(&mut self) -> Option<&'a Zval> {
+        let get_current_data = unsafe { (*self.funcs).get_current_data }?;
+        let value = unsafe { &*get_current_data(&mut *self) };
+
+        Some(value)
+    }
+
+    pub fn get_current_key(&mut self) -> Option<Zval> {
+        let get_current_key = unsafe { (*self.funcs).get_current_key }?;
+        let mut key = Zval::new();
+        unsafe {
+            get_current_key(&mut *self, &mut key);
+        }
+
+        Some(key)
+    }
+}
+
+impl<'a> Iterator for Iter<'a> {
+    type Item = (u64, Option<String>, &'a Zval);
+
+    fn next(&mut self) -> Option<Self::Item> {
+        // Call next when index > 0, so next is really called at the start of each iteration, which allow to work better with generator iterator
+        if self.zi.index > 0 {
+            self.zi.move_forward();
+
+            if !self.zi.valid() {
+                return None;
+            }
+        }
+
+        self.zi.index += 1;
+
+        let key = self.zi.get_current_key();
+        let value = self.zi.get_current_data()?;
+        let real_index = self.zi.index - 1;
+
+        Some(match key {
+            Some(key) => match key.is_long() {
+                false => (real_index, key.try_into().ok(), value),
+                true => (
+                    key.long().unwrap_or(real_index as ZendLong) as u64,
+                    None,
+                    value,
+                ),
+            },
+            None => (real_index, None, value),
+        })
+    }
+}
+
+impl<'a> FromZvalMut<'a> for &'a mut ZendIterator {
+    const TYPE: DataType = DataType::Object(Some("Traversable"));
+
+    fn from_zval_mut(zval: &'a mut Zval) -> Option<Self> {
+        let zend_object = zval.object()?;
+        let ce = zend_object.get_class_entry_mut();
+        let iterator = unsafe { ce.get_iterator?(&mut *ce, &mut *zval, 0) };
+
+        unsafe { iterator.as_mut() }
+    }
+}
diff --git a/src/types/mod.rs b/src/types/mod.rs
index 888e056..84785eb 100644
--- a/src/types/mod.rs
+++ b/src/types/mod.rs
@@ -6,6 +6,7 @@
 mod array;
 mod callable;
 mod class_object;
+mod iterator;
 mod long;
 mod object;
 mod string;
diff --git a/src/types/object.rs b/src/types/object.rs
index ca8e526..0b83796 100644
--- a/src/types/object.rs
+++ b/src/types/object.rs
@@ -103,6 +103,17 @@ impl ZendObject {
         unsafe { self.ce.as_ref() }.expect("Could not retrieve class entry.")
     }
 
+    /// Returns the [`ClassEntry`] associated with this object.
+    ///
+    /// # Panics
+    ///
+    /// Panics if the class entry is invalid.
+    pub fn get_class_entry_mut(&self) -> &'static mut ClassEntry {
+        // SAFETY: it is OK to panic here since PHP would segfault anyway
+        // when encountering an object with no class entry.
+        unsafe { self.ce.as_mut() }.expect("Could not retrieve class entry.")
+    }
+
     /// Attempts to retrieve the class name of the object.
     pub fn get_class_name(&self) -> Result<String> {
         unsafe {