Move nxt function from pyobject to objiter file. Also fix formatting issues.

Author: windelbouwman

vm/src/builtins.rs

@@ -4,9 +4,9 @@ use std::collections::HashMap;
 use std::io::{self, Write};
 
 use super::compile;
+use super::obj::objiter;
 use super::obj::objstr;
 use super::obj::objtype;
-use super::obj::objiter;
 use super::objbool;
 use super::pyobject::{
     AttributeProtocol, DictProtocol, IdProtocol, PyContext, PyFuncArgs, PyObject, PyObjectKind,
@@ -223,11 +223,7 @@ fn builtin_issubclass(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
 }
 
 fn builtin_iter(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
-    arg_check!(
-        vm,
-        args,
-        required = [(iter_target, None)]
-    );
+    arg_check!(vm, args, required = [(iter_target, None)]);
     objiter::get_iter(vm, iter_target)
 }
 
@@ -409,10 +405,7 @@ pub fn make_module(ctx: &PyContext) -> PyObjectRef {
         String::from("issubclass"),
         ctx.new_rustfunc(builtin_issubclass),
     );
-    dict.insert(
-        String::from("iter"),
-        ctx.new_rustfunc(builtin_iter),
-    );
+    dict.insert(String::from("iter"), ctx.new_rustfunc(builtin_iter));
     dict.insert(String::from("len"), ctx.new_rustfunc(builtin_len));
     dict.insert(String::from("list"), ctx.list_type());
     dict.insert(String::from("locals"), ctx.new_rustfunc(builtin_locals));

vm/src/obj/objiter.rs

@@ -3,18 +3,23 @@
  */
 
 use super::super::pyobject::{
-    AttributeProtocol, PyContext, PyFuncArgs, PyObject, PyObjectKind, PyObjectRef,
-    PyResult, TypeProtocol,
+    AttributeProtocol, PyContext, PyFuncArgs, PyObject, PyObjectKind, PyObjectRef, PyResult,
+    TypeProtocol,
 };
 use super::super::vm::VirtualMachine;
 use super::objstr;
 use super::objtype; // Required for arg_check! to use isinstance
 
+/*
+ * This helper function is called at multiple places. First, it is called
+ * in the vm when a for loop is entered. Next, it is used when the builtin
+ * function 'iter' is called.
+ */
 pub fn get_iter(vm: &mut VirtualMachine, iter_target: &PyObjectRef) -> PyResult {
     // Check what we are going to iterate over:
     let iterated_obj = if objtype::isinstance(iter_target, vm.ctx.iter_type()) {
         // If object is already an iterator, return that one.
-        return Ok(iter_target.clone())
+        return Ok(iter_target.clone());
     } else if objtype::isinstance(iter_target, vm.ctx.list_type()) {
         iter_target.clone()
     } else {
@@ -37,46 +42,45 @@ pub fn get_iter(vm: &mut VirtualMachine, iter_target: &PyObjectRef) -> PyResult
 
 // Sequence iterator:
 fn iter_new(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
-    arg_check!(
-        vm,
-        args,
-        required = [(iter_target, None)]
-    );
+    arg_check!(vm, args, required = [(iter_target, None)]);
 
     get_iter(vm, iter_target)
 }
 
 fn iter_iter(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
-    arg_check!(
-        vm,
-        args,
-        required = [(iter, Some(vm.ctx.iter_type()))]
-    );
+    arg_check!(vm, args, required = [(iter, Some(vm.ctx.iter_type()))]);
     // Return self:
     Ok(iter.clone())
 }
 
 fn iter_next(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
-    arg_check!(
-        vm,
-        args,
-        required = [(iter, Some(vm.ctx.iter_type()))]
-    );
+    arg_check!(vm, args, required = [(iter, Some(vm.ctx.iter_type()))]);
 
-    let next_obj: Option<PyObjectRef> = {
-        // We require a mutable pyobject here to update the iterator:
-        let mut iterator: &mut PyObject = &mut iter.borrow_mut();
-        iterator.nxt()
-    };
-
-    // Return next item, or StopIteration
-    match next_obj {
-        Some(value) => Ok(value),
-        None => {
-            let stop_iteration_type = vm.ctx.exceptions.stop_iteration.clone();
-            let stop_iteration = vm.new_exception(stop_iteration_type, "End of iterator".to_string());
-            Err(stop_iteration)
+    if let PyObjectKind::Iterator {
+        ref mut position,
+        iterated_obj: ref iterated_obj_ref,
+    } = iter.borrow_mut().kind
+    {
+        let iterated_obj = &*iterated_obj_ref.borrow_mut();
+        match iterated_obj.kind {
+            PyObjectKind::List { ref elements } => {
+                if *position < elements.len() {
+                    let obj_ref = elements[*position].clone();
+                    *position += 1;
+                    Ok(obj_ref)
+                } else {
+                    let stop_iteration_type = vm.ctx.exceptions.stop_iteration.clone();
+                    let stop_iteration =
+                        vm.new_exception(stop_iteration_type, "End of iterator".to_string());
+                    Err(stop_iteration)
+                }
+            }
+            _ => {
+                panic!("NOT IMPL");
+            }
         }
+    } else {
+        panic!("NOT IMPL");
     }
 }
 
@@ -86,4 +90,3 @@ pub fn init(context: &PyContext) {
     iter_type.set_attr("__iter__", context.new_rustfunc(iter_iter));
     iter_type.set_attr("__next__", context.new_rustfunc(iter_next));
 }
-

vm/src/pyobject.rs

@@ -758,35 +758,6 @@ impl PyObject {
         }
     }
 
-    // Implement iterator protocol:
-    pub fn nxt(&mut self) -> Option<PyObjectRef> {
-        match self.kind {
-            PyObjectKind::Iterator {
-                ref mut position,
-                iterated_obj: ref iterated_obj_ref,
-            } => {
-                let iterated_obj = &*iterated_obj_ref.borrow_mut();
-                match iterated_obj.kind {
-                    PyObjectKind::List { ref elements } => {
-                        if *position < elements.len() {
-                            let obj_ref = elements[*position].clone();
-                            *position += 1;
-                            Some(obj_ref)
-                        } else {
-                            None
-                        }
-                    }
-                    _ => {
-                        panic!("NOT IMPL");
-                    }
-                }
-            }
-            _ => {
-                panic!("NOT IMPL");
-            }
-        }
-    }
-
     // Move this object into a reference object, transferring ownership.
     pub fn into_ref(self) -> PyObjectRef {
         Rc::new(RefCell::new(self))

vm/src/vm.rs

@@ -832,7 +832,7 @@ impl VirtualMachine {
                     Ok(iter_obj) => {
                         self.push_value(iter_obj);
                         None
-                    },
+                    }
                     Err(err) => Some(Err(err)),
                 }
             }
@@ -851,12 +851,16 @@ impl VirtualMachine {
                     }
                     Err(next_error) => {
                         // Check if we have stopiteration, or something else:
-                        if objtype::isinstance(&next_error, self.ctx.exceptions.stop_iteration.clone()) {
+                        if objtype::isinstance(
+                            &next_error,
+                            self.ctx.exceptions.stop_iteration.clone(),
+                        ) {
                             // Pop iterator from stack:
                             self.pop_value();
 
                             // End of for loop
-                            let end_label = if let Block::Loop { start: _, end } = self.last_block() {
+                            let end_label = if let Block::Loop { start: _, end } = self.last_block()
+                            {
                                 *end
                             } else {
                                 panic!("Wrong block type")