Merge pull request RustPython#3725 from youknowone/opt-nth

Optimize comprehension item add operations

Author: youknowone

bytecode/src/lib.rs

@@ -393,13 +393,6 @@ pub enum Instruction {
     GetAIter,
     GetANext,
     EndAsyncFor,
-
-    /// Reverse order evaluation in MapAdd
-    /// required to support named expressions of Python 3.8 in dict comprehension
-    /// today (including Py3.9) only required in dict comprehension.
-    MapAddRev {
-        i: u32,
-    },
 }
 
 use self::Instruction::*;
@@ -1024,7 +1017,7 @@ impl Instruction {
             }
             BuildSlice { step } => -2 - (*step as i32) + 1,
             ListAppend { .. } | SetAdd { .. } => -1,
-            MapAdd { .. } | MapAddRev { .. } => -2,
+            MapAdd { .. } => -2,
             PrintExpr => -1,
             LoadBuildClass => 1,
             UnpackSequence { size } => -1 + *size as i32,
@@ -1187,7 +1180,7 @@ impl Instruction {
             BuildSlice { step } => w!(BuildSlice, step),
             ListAppend { i } => w!(ListAppend, i),
             SetAdd { i } => w!(SetAdd, i),
-            MapAddRev { i } => w!(MapAddRev, i),
+            MapAdd { i } => w!(MapAdd, i),
             PrintExpr => w!(PrintExpr),
             LoadBuildClass => w!(LoadBuildClass),
             UnpackSequence { size } => w!(UnpackSequence, size),
@@ -1199,7 +1192,6 @@ impl Instruction {
             GetAIter => w!(GetAIter),
             GetANext => w!(GetANext),
             EndAsyncFor => w!(EndAsyncFor),
-            MapAdd { i } => w!(MapAdd, i),
         }
     }
 }

compiler/src/compile.rs

@@ -2095,7 +2095,7 @@ impl Compiler {
                     &|compiler| {
                         compiler.compile_comprehension_element(elt)?;
                         compiler.emit(Instruction::ListAppend {
-                            i: (1 + generators.len()) as u32,
+                            i: generators.len() as u32,
                         });
                         Ok(())
                     },
@@ -2112,7 +2112,7 @@ impl Compiler {
                     &|compiler| {
                         compiler.compile_comprehension_element(elt)?;
                         compiler.emit(Instruction::SetAdd {
-                            i: (1 + generators.len()) as u32,
+                            i: generators.len() as u32,
                         });
                         Ok(())
                     },
@@ -2136,8 +2136,8 @@ impl Compiler {
                         compiler.compile_expression(key)?;
                         compiler.compile_expression(value)?;
 
-                        compiler.emit(Instruction::MapAddRev {
-                            i: (1 + generators.len()) as u32,
+                        compiler.emit(Instruction::MapAdd {
+                            i: generators.len() as u32,
                         });
 
                         Ok(())

vm/src/frame.rs

@@ -4,7 +4,7 @@ use crate::{
         asyncgenerator::PyAsyncGenWrappedValue,
         function::{PyCell, PyCellRef, PyFunction},
         tuple::{PyTuple, PyTupleTyped},
-        PyBaseExceptionRef, PyCode, PyCoroutine, PyDict, PyDictRef, PyGenerator, PyListRef, PySet,
+        PyBaseExceptionRef, PyCode, PyCoroutine, PyDict, PyDictRef, PyGenerator, PyList, PySet,
         PySlice, PyStr, PyStrRef, PyTraceback, PyTypeRef,
     },
     bytecode,
@@ -668,45 +668,33 @@ impl ExecutingFrame<'_> {
             } => self.execute_build_map(vm, *size, *unpack, *for_call),
             bytecode::Instruction::BuildSlice { step } => self.execute_build_slice(vm, *step),
             bytecode::Instruction::ListAppend { i } => {
+                let item = self.pop_value();
                 let obj = self.nth_value(*i);
-                let list: PyListRef = unsafe {
+                let list: &Py<PyList> = unsafe {
                     // SAFETY: trust compiler
-                    obj.downcast_unchecked()
+                    obj.downcast_unchecked_ref()
                 };
-                let item = self.pop_value();
                 list.append(item);
                 Ok(None)
             }
             bytecode::Instruction::SetAdd { i } => {
+                let item = self.pop_value();
                 let obj = self.nth_value(*i);
-                let set: PyRef<PySet> = unsafe {
+                let set: &Py<PySet> = unsafe {
                     // SAFETY: trust compiler
-                    obj.downcast_unchecked()
+                    obj.downcast_unchecked_ref()
                 };
-                let item = self.pop_value();
                 set.add(item, vm)?;
                 Ok(None)
             }
             bytecode::Instruction::MapAdd { i } => {
-                let obj = self.nth_value(*i + 1);
-                let dict: PyDictRef = unsafe {
-                    // SAFETY: trust compiler
-                    obj.downcast_unchecked()
-                };
-                let key = self.pop_value();
                 let value = self.pop_value();
-                dict.set_item(&*key, value, vm)?;
-                Ok(None)
-            }
-            bytecode::Instruction::MapAddRev { i } => {
-                // change order of evalutio of key and value to support Py3.8 Named expressions in dict comprehension
-                let obj = self.nth_value(*i + 1);
-                let dict: PyDictRef = unsafe {
+                let key = self.pop_value();
+                let obj = self.nth_value(*i);
+                let dict: &Py<PyDict> = unsafe {
                     // SAFETY: trust compiler
-                    obj.downcast_unchecked()
+                    obj.downcast_unchecked_ref()
                 };
-                let value = self.pop_value();
-                let key = self.pop_value();
                 dict.set_item(&*key, value, vm)?;
                 Ok(None)
             }
@@ -1811,8 +1799,9 @@ impl ExecutingFrame<'_> {
     }
 
     #[inline]
-    fn nth_value(&self, depth: u32) -> PyObjectRef {
-        self.state.stack[self.state.stack.len() - depth as usize - 1].clone()
+    fn nth_value(&self, depth: u32) -> &PyObject {
+        let stack = &self.state.stack;
+        &stack[stack.len() - depth as usize - 1]
     }
 
     #[cold]