Use &mut when collecting parameters

Summary:
`&mut [Option<Value>]` is supposed to provide better optimization than `&[Cell<Option<Value>>]` because compiler knows `&mut` pointer is unique.

The only non-obvious part is `eval.current_frame.locals_mut()` call in this function:

https://www.internalfb.com/code/fbsource/[D63300287-V1]/fbcode/buck2/starlark-rust/starlark/src/eval/compiler/def.rs?lines=738-759

This is safe because:
- `current_frame` was just allocated in `alloca_frame`, and not stored anywhere except `eval`, so it is unique
- `collect_inline` does not have access to `eval`, so it cannot reach `eval.current_frame`

Reviewed By: blackm00n, JakobDegen

Differential Revision: D63300287

fbshipit-source-id: 26a4347e8a2aa799c7602c8777bc3c62c0358032

starlark/src/__derive_refs/parse_args.rs

@@ -15,8 +15,6 @@
  * limitations under the License.
  */
 
-use std::cell::Cell;
-
 use crate::eval::Arguments;
 use crate::eval::ParametersSpec;
 use crate::values::FrozenValue;
@@ -33,7 +31,7 @@ pub fn parse_signature<'v, const N: usize>(
     parser: &ParametersSpec<FrozenValue>,
     args: &Arguments<'v, '_>,
     heap: &'v Heap,
-) -> crate::Result<[Cell<Option<Value<'v>>>; N]> {
+) -> crate::Result<[Option<Value<'v>>; N]> {
     parser.collect_into(args, heap)
 }
 

starlark/src/eval/bc/frame.rs

@@ -17,7 +17,6 @@
 
 //! Local variables and stack, in single allocation.
 
-use std::cell::Cell;
 use std::mem;
 use std::mem::MaybeUninit;
 use std::ptr;
@@ -96,7 +95,7 @@ impl<'v> BcFramePtr<'v> {
     }
 
     #[inline(always)]
-    fn frame_mut(&mut self) -> &mut BcFrame<'v> {
+    fn frame_mut<'a>(self) -> &'a mut BcFrame<'v> {
         debug_assert!(self.is_inititalized());
         unsafe {
             let frame = (self.slots_ptr as *mut u8).sub(BcFrame::offset_of_slots()) as *mut BcFrame;
@@ -120,7 +119,7 @@ impl<'v> BcFramePtr<'v> {
     }
 
     #[inline(always)]
-    pub(crate) fn set_slot(mut self, slot: LocalSlotIdCapturedOrNot, value: Value<'v>) {
+    pub(crate) fn set_slot(self, slot: LocalSlotIdCapturedOrNot, value: Value<'v>) {
         self.frame_mut().set_slot(slot, value)
     }
 
@@ -130,7 +129,7 @@ impl<'v> BcFramePtr<'v> {
     }
 
     #[inline(always)]
-    pub(crate) fn set_bc_slot(mut self, slot: BcSlotOut, value: Value<'v>) {
+    pub(crate) fn set_bc_slot(self, slot: BcSlotOut, value: Value<'v>) {
         self.frame_mut().set_bc_slot(slot, value)
     }
 
@@ -145,7 +144,7 @@ impl<'v> BcFramePtr<'v> {
     }
 
     #[inline(always)]
-    pub(crate) fn set_iter_index(mut self, loop_depth: LoopDepth, index: usize) {
+    pub(crate) fn set_iter_index(self, loop_depth: LoopDepth, index: usize) {
         self.frame_mut().set_iter_index(loop_depth, index)
     }
 
@@ -154,8 +153,8 @@ impl<'v> BcFramePtr<'v> {
     }
 
     #[inline(always)]
-    pub(crate) fn locals(&self) -> &[Cell<Option<Value<'v>>>] {
-        self.frame().locals()
+    pub(crate) unsafe fn locals_mut<'a>(self) -> &'a mut [Option<Value<'v>>] {
+        self.frame_mut().locals_mut()
     }
 }
 
@@ -176,16 +175,6 @@ impl<'v> BcFrame<'v> {
         }
     }
 
-    #[inline(always)]
-    fn locals(&self) -> &[Cell<Option<Value<'v>>>] {
-        unsafe {
-            slice::from_raw_parts(
-                self.slots.as_ptr() as *const Cell<Option<Value>>,
-                self.local_count as usize,
-            )
-        }
-    }
-
     #[inline(always)]
     fn locals_mut(&mut self) -> &mut [Option<Value<'v>>] {
         unsafe { slice::from_raw_parts_mut(self.slots.as_mut_ptr(), self.local_count as usize) }
@@ -361,7 +350,7 @@ pub(crate) fn alloca_frame<'v, 'a, 'e, R>(
         local_count,
         max_stack_size,
         loop_depth,
-        |eval, mut frame| {
+        |eval, frame| {
             // TODO(nga): no need to fill the slots for parameters.
             frame.frame_mut().init();
             let old_frame = mem::replace(&mut eval.current_frame, frame);

starlark/src/eval/compiler/call.rs

@@ -17,8 +17,6 @@
 
 //! Compile function calls.
 
-use std::cell::Cell;
-
 use starlark_derive::VisitSpanMut;
 use starlark_syntax::slice_vec_ext::VecExt;
 
@@ -178,15 +176,15 @@ impl CallCompiled {
         };
 
         args.all_values_generic(expr_to_value, |arguments| {
-            let slots = vec![Cell::new(None); fun.parameters.len()];
+            let mut slots = vec![None; fun.parameters.len()];
             fun.parameters
-                .collect(arguments.frozen_to_v(), &slots, ctx.heap())
+                .collect(arguments.frozen_to_v(), &mut slots, ctx.heap())
                 .ok()?;
 
             let slots = slots
                 .into_try_map(|value| {
                     // Value must be set, but better ignore optimization here than panic.
-                    let value = value.get().ok_or(())?;
+                    let value = value.ok_or(())?;
                     // Everything should be frozen here, but if not,
                     // it is safer to abandon optimization.
                     value.unpack_frozen().ok_or(())

starlark/src/eval/compiler/def.rs

@@ -751,7 +751,11 @@ where
             bc.max_stack_size,
             bc.max_loop_depth,
             |eval| {
-                let slots = eval.current_frame.locals();
+                // SAFETY: `slots` is unique: `alloca_frame` just allocated the frame,
+                //   so there are no references to the frame except `eval.current_frame`.
+                //   We use `slots` only in `collect_inline`,
+                //   which does not have access to `eval` thus cannot access the frame indirectly.
+                let slots = unsafe { eval.current_frame.locals_mut() };
                 self.parameters.collect_inline(args, slots, eval.heap())?;
                 self.invoke_raw(me, eval)
             },

starlark/src/eval/runtime/params/parser.rs

@@ -15,8 +15,6 @@
  * limitations under the License.
  */
 
-use std::cell::Cell;
-
 use starlark_syntax::other_error;
 
 use crate::values::UnpackValue;
@@ -26,19 +24,19 @@ use crate::values::Value;
 /// [`ParametersSpec`](crate::eval::ParametersSpec).
 ///
 /// This is created with [`ParametersSpec::parser`](crate::eval::ParametersSpec::parser).
-pub struct ParametersParser<'v, 'a>(std::slice::Iter<'a, Cell<Option<Value<'v>>>>);
+pub struct ParametersParser<'v, 'a>(std::slice::Iter<'a, Option<Value<'v>>>);
 
 impl<'v, 'a> ParametersParser<'v, 'a> {
     /// Create a parameter parser, which stored parameters into provided slots reference.
-    pub(crate) fn new(slots: &'a [Cell<Option<Value<'v>>>]) -> Self {
+    pub(crate) fn new(slots: &'a [Option<Value<'v>>]) -> Self {
         Self(slots.iter())
     }
 
     fn get_next(&mut self, name: &str) -> anyhow::Result<Option<Value<'v>>> {
         let Some(v) = self.0.next() else {
             return Err(other_error!("Requested parameter `{name}`, which is after the number of parameters in provided signature").into_anyhow());
         };
-        Ok(v.get())
+        Ok(*v)
     }
 
     /// Obtain the next parameter, corresponding to

starlark/src/eval/runtime/params/spec.rs

@@ -15,7 +15,6 @@
  * limitations under the License.
  */
 
-use std::cell::Cell;
 use std::cmp;
 use std::collections::HashMap;
 use std::fmt;
@@ -505,7 +504,7 @@ impl<'v> ParametersSpec<Value<'v>> {
     fn collect_impl(
         &self,
         args: &Arguments<'v, '_>,
-        slots: &[Cell<Option<Value<'v>>>],
+        slots: &mut [Option<Value<'v>>],
         heap: &'v Heap,
     ) -> crate::Result<()> {
         self.collect_inline(&args.0, slots, heap)
@@ -519,9 +518,9 @@ impl<'v> ParametersSpec<Value<'v>> {
         &self,
         args: &Arguments<'v, '_>,
         heap: &'v Heap,
-    ) -> crate::Result<[Cell<Option<Value<'v>>>; N]> {
-        let slots = [(); N].map(|_| Cell::new(None));
-        self.collect(args, &slots, heap)?;
+    ) -> crate::Result<[Option<Value<'v>>; N]> {
+        let mut slots = [(); N].map(|_| None);
+        self.collect(args, &mut slots, heap)?;
         Ok(slots)
     }
 
@@ -531,7 +530,7 @@ impl<'v> ParametersSpec<Value<'v>> {
     fn collect_inline_impl<'a, A: ArgumentsImpl<'v, 'a>>(
         &self,
         args: &A,
-        slots: &[Cell<Option<Value<'v>>>],
+        slots: &mut [Option<Value<'v>>],
         heap: &'v Heap,
     ) -> crate::Result<()>
     where
@@ -546,8 +545,8 @@ impl<'v> ParametersSpec<Value<'v>> {
             && args.args().is_none()
             && args.kwargs().is_none()
         {
-            for (v, s) in args.pos().iter().zip(slots.iter()) {
-                s.set(Some(*v));
+            for (v, s) in args.pos().iter().zip(slots.iter_mut()) {
+                *s = Some(*v);
             }
 
             return Ok(());
@@ -559,7 +558,7 @@ impl<'v> ParametersSpec<Value<'v>> {
     fn collect_slow<'a, A: ArgumentsImpl<'v, 'a>>(
         &self,
         args: &A,
-        slots: &[Cell<Option<Value<'v>>>],
+        slots: &mut [Option<Value<'v>>],
         heap: &'v Heap,
     ) -> crate::Result<()>
     where
@@ -606,14 +605,14 @@ impl<'v> ParametersSpec<Value<'v>> {
         // First deal with positional parameters
         if args.pos().len() <= (self.indices.num_positional as usize) {
             // fast path for when we don't need to bounce down to filling in args
-            for (v, s) in args.pos().iter().zip(slots.iter()) {
-                s.set(Some(*v));
+            for (v, s) in args.pos().iter().zip(slots.iter_mut()) {
+                *s = Some(*v);
             }
             next_position = args.pos().len();
         } else {
             for v in args.pos() {
                 if next_position < (self.indices.num_positional as usize) {
-                    slots[next_position].set(Some(*v));
+                    slots[next_position] = Some(*v);
                     next_position += 1;
                 } else {
                     star_args.push(*v);
@@ -635,7 +634,7 @@ impl<'v> ParametersSpec<Value<'v>> {
                         kwargs.insert(Hashed::new_unchecked(name.small_hash(), *name_value), *v);
                     }
                     Some(i) => {
-                        slots[i].set(Some(*v));
+                        slots[i] = Some(*v);
                         lowest_name = cmp::min(lowest_name, i);
                     }
                 }
@@ -649,7 +648,7 @@ impl<'v> ParametersSpec<Value<'v>> {
                 .map_err(|_| FunctionError::ArgsArrayIsNotIterable)?
             {
                 if next_position < (self.indices.num_positional as usize) {
-                    slots[next_position].set(Some(v));
+                    slots[next_position] = Some(v);
                     next_position += 1;
                 } else {
                     star_args.push(v);
@@ -679,9 +678,9 @@ impl<'v> ParametersSpec<Value<'v>> {
                                 {
                                     None => kwargs.insert(Hashed::new_unchecked(k.hash(), s), v),
                                     Some(i) => {
-                                        let this_slot = &slots[*i as usize];
-                                        let repeat = this_slot.get().is_some();
-                                        this_slot.set(Some(v));
+                                        let this_slot = &mut slots[*i as usize];
+                                        let repeat = this_slot.is_some();
+                                        *this_slot = Some(v);
                                         repeat
                                     }
                                 };
@@ -706,11 +705,11 @@ impl<'v> ParametersSpec<Value<'v>> {
         for index in next_position..kinds.len() {
             // The number of locals must be at least the number of parameters, see `collect`
             // which reserves `max(_, kinds.len())`.
-            let slot = unsafe { slots.get_unchecked(index) };
+            let slot = unsafe { slots.get_unchecked_mut(index) };
             let def = unsafe { kinds.get_unchecked(index) };
 
             // We know that up to next_position got filled positionally, so we don't need to check those
-            if slot.get().is_some() {
+            if slot.is_some() {
                 continue;
             }
             match def {
@@ -722,7 +721,7 @@ impl<'v> ParametersSpec<Value<'v>> {
                     .into());
                 }
                 ParameterKind::Defaulted(x) => {
-                    slot.set(Some(x.to_value()));
+                    *slot = Some(x.to_value());
                 }
                 _ => {}
             }
@@ -732,7 +731,7 @@ impl<'v> ParametersSpec<Value<'v>> {
         // Note that we deliberately give warnings about missing parameters _before_ giving warnings
         // about unexpected extra parameters, so if a user misspells an argument they get a better error.
         if let Some(args_pos) = self.indices.args {
-            slots[args_pos as usize].set(Some(heap.alloc_tuple(&star_args)));
+            slots[args_pos as usize] = Some(heap.alloc_tuple(&star_args));
         } else if unlikely(!star_args.is_empty()) {
             return Err(FunctionError::ExtraPositionalArg {
                 count: star_args.len(),
@@ -742,7 +741,7 @@ impl<'v> ParametersSpec<Value<'v>> {
         }
 
         if let Some(kwargs_pos) = self.indices.kwargs {
-            slots[kwargs_pos as usize].set(Some(kwargs.alloc(heap)));
+            slots[kwargs_pos as usize] = Some(kwargs.alloc(heap));
         } else if let Some(kwargs) = kwargs.kwargs {
             return Err(FunctionError::ExtraNamedArg {
                 names: kwargs.keys().map(|x| x.as_str().to_owned()).collect(),
@@ -862,7 +861,7 @@ impl<'v> ParametersSpec<Value<'v>> {
     {
         eval.alloca_init(
             self.len(),
-            || Cell::new(None),
+            || None,
             |slots, eval| {
                 self.collect_inline(&args.0, slots, eval.heap())?;
                 let mut parser = ParametersParser::new(slots);
@@ -888,7 +887,7 @@ impl<'v, V: ValueLike<'v>> ParametersSpec<V> {
         &self,
         args: &Arguments<'v, '_>,
         heap: &'v Heap,
-    ) -> crate::Result<[Cell<Option<Value<'v>>>; N]> {
+    ) -> crate::Result<[Option<Value<'v>>; N]> {
         self.as_value().collect_into_impl(args, heap)
     }
 
@@ -898,7 +897,7 @@ impl<'v, V: ValueLike<'v>> ParametersSpec<V> {
     pub fn collect(
         &self,
         args: &Arguments<'v, '_>,
-        slots: &[Cell<Option<Value<'v>>>],
+        slots: &mut [Option<Value<'v>>],
         heap: &'v Heap,
     ) -> crate::Result<()> {
         self.as_value().collect_impl(args, slots, heap)
@@ -940,7 +939,7 @@ impl<'v, V: ValueLike<'v>> ParametersSpec<V> {
     pub(crate) fn collect_inline<'a, A: ArgumentsImpl<'v, 'a>>(
         &self,
         args: &A,
-        slots: &[Cell<Option<Value<'v>>>],
+        slots: &mut [Option<Value<'v>>],
         heap: &'v Heap,
     ) -> crate::Result<()>
     where

starlark_derive/src/module/render/fun.rs

@@ -443,7 +443,7 @@ fn render_binding_arg(arg: &StarArg) -> BindingArg {
 
     let source = match arg.source {
         StarArgSource::This => quote! { __this },
-        StarArgSource::Argument(i) => quote! { __args[#i].get() },
+        StarArgSource::Argument(i) => quote! { __args[#i] },
         StarArgSource::Required(i) => quote! {  Some(__required[#i])},
         StarArgSource::Optional(i) => quote! { __optional[#i] },
         ref s => unreachable!("unknown source: {:?}", s),