From 03eed121fbf05cb75bcfb6a8cb57e788589eaf47 Mon Sep 17 00:00:00 2001
From: Kangzhi Shi <shikangzhi@gmail.com>
Date: Thu, 15 Oct 2020 15:42:11 +0200
Subject: [PATCH] manually loop for list

---
 Lib/test/test_bytes.py |  1 -
 vm/src/bytesinner.rs   |  6 ++----
 vm/src/iterator.rs     |  4 ++--
 vm/src/vm.rs           | 40 ++++++++++++++++++++++++++++++++--------
 4 files changed, 36 insertions(+), 15 deletions(-)

diff --git a/Lib/test/test_bytes.py b/Lib/test/test_bytes.py
index 51b378480..b28c2d5c7 100644
--- a/Lib/test/test_bytes.py
+++ b/Lib/test/test_bytes.py
@@ -113,7 +113,6 @@ class BaseBytesTest:
         b = self.type2test([1, 2, 3])
         self.assertEqual(b, b"\x01\x02\x03")
 
-    # @unittest.skip("TODO")
     def test_from_mutating_list(self):
         # Issue #34973: Crash in bytes constructor with mutating list.
         class X:
diff --git a/vm/src/bytesinner.rs b/vm/src/bytesinner.rs
index 4ac017d86..e7e0cdf33 100644
--- a/vm/src/bytesinner.rs
+++ b/vm/src/bytesinner.rs
@@ -3,6 +3,7 @@ use itertools::Itertools;
 use num_bigint::BigInt;
 use num_traits::ToPrimitive;
 
+use crate::anystr::{self, AnyStr, AnyStrContainer, AnyStrWrapper};
 use crate::builtins::bytearray::PyByteArray;
 use crate::builtins::bytes::{PyBytes, PyBytesRef};
 use crate::builtins::int::{self, PyInt, PyIntRef};
@@ -14,12 +15,9 @@ use crate::pyobject::{
     BorrowValue, Either, PyComparisonValue, PyIterable, PyObjectRef, PyResult, TryFromObject,
     TypeProtocol,
 };
+use crate::sliceable::PySliceableSequence;
 use crate::slots::PyComparisonOp;
 use crate::vm::VirtualMachine;
-use crate::{
-    anystr::{self, AnyStr, AnyStrContainer, AnyStrWrapper},
-    sliceable::PySliceableSequence,
-};
 use rustpython_common::hash;
 
 #[derive(Debug, Default, Clone)]
diff --git a/vm/src/iterator.rs b/vm/src/iterator.rs
index ef0fc573c..2dfb229d7 100644
--- a/vm/src/iterator.rs
+++ b/vm/src/iterator.rs
@@ -80,7 +80,7 @@ pub fn get_all<T: TryFromObject>(vm: &VirtualMachine, iter_obj: &PyObjectRef) ->
 
 pub fn try_map<F, R>(vm: &VirtualMachine, iter_obj: &PyObjectRef, mut f: F) -> PyResult<Vec<R>>
 where
-    F: FnMut(&PyObjectRef) -> PyResult<R>,
+    F: FnMut(PyObjectRef) -> PyResult<R>,
 {
     let cap = length_hint(vm, iter_obj.clone())?.unwrap_or(0);
     // TODO: fix extend to do this check (?), see test_extend in Lib/test/list_tests.py,
@@ -90,7 +90,7 @@ where
     }
     let mut results = Vec::with_capacity(cap);
     while let Some(element) = get_next_object(vm, iter_obj)? {
-        results.push(f(&element)?);
+        results.push(f(element)?);
     }
     results.shrink_to_fit();
     Ok(results)
diff --git a/vm/src/vm.rs b/vm/src/vm.rs
index 1e971e1b0..32c3ef871 100644
--- a/vm/src/vm.rs
+++ b/vm/src/vm.rs
@@ -997,22 +997,46 @@ impl VirtualMachine {
         T: TryFromObject,
         F: FnMut(T) -> PyResult<R>,
     {
-        self.map_iterable_object(obj, |elem| f(T::try_from_object(self, elem.clone())?))
+        self.map_iterable_object(obj, |elem| f(T::try_from_object(self, elem)?))
     }
 
-    pub fn map_iterable_object<F, R>(&self, obj: &PyObjectRef, f: F) -> PyResult<PyResult<Vec<R>>>
+    pub fn map_iterable_object<F, R>(
+        &self,
+        obj: &PyObjectRef,
+        mut f: F,
+    ) -> PyResult<PyResult<Vec<R>>>
     where
-        F: FnMut(&PyObjectRef) -> PyResult<R>,
+        F: FnMut(PyObjectRef) -> PyResult<R>,
     {
-        Ok(match_class!(match obj {
-            ref l @ PyList => l.borrow_value().iter().map(f).collect(),
-            ref t @ PyTuple => t.borrow_value().iter().map(f).collect(),
+        match_class!(match obj {
+            ref l @ PyList => {
+                let mut i: usize = 0;
+                let mut results = Vec::with_capacity(l.borrow_value().len());
+                loop {
+                    let elem = {
+                        let elements = &*l.borrow_value();
+                        if i >= elements.len() {
+                            results.shrink_to_fit();
+                            return Ok(Ok(results));
+                        } else {
+                            elements[i].clone()
+                        }
+                        // free the lock
+                    };
+                    match f(elem) {
+                        Ok(result) => results.push(result),
+                        Err(err) => return Ok(Err(err)),
+                    }
+                    i += 1;
+                }
+            }
+            ref t @ PyTuple => Ok(t.borrow_value().iter().cloned().map(f).collect()),
             // TODO: put internal iterable type
             obj => {
                 let iter = iterator::get_iter(self, obj)?;
-                iterator::try_map(self, &iter, f)
+                Ok(iterator::try_map(self, &iter, f))
             }
-        }))
+        })
     }
 
     // get_attribute should be used for full attribute access (usually from user code).