Skip to content

Commit

Permalink
Make ArgNames unique
Browse files Browse the repository at this point in the history
Summary:
`ArgNames` struct inside `Arguments` contains names of named arguments.

Make it unique, so we can do `insert_unique_unchecked` in more places.

Reviewed By: blackm00n

Differential Revision: D63615329

fbshipit-source-id: 8a3625c76d07dfd7311e133725cd61247ce341ec
  • Loading branch information
stepancheg authored and facebook-github-bot committed Oct 1, 2024
1 parent d0d384d commit 952a5d5
Show file tree
Hide file tree
Showing 6 changed files with 82 additions and 40 deletions.
2 changes: 1 addition & 1 deletion starlark-rust/starlark/src/eval.rs
Original file line number Diff line number Diff line change
Expand Up @@ -158,7 +158,7 @@ impl<'v, 'a, 'e> Evaluator<'v, 'a, 'e> {
let params = Arguments(ArgumentsFull {
pos: positional,
named: &named,
names: ArgNames::new(&names),
names: ArgNames::new_check_unique(&names)?,
args: None,
kwargs: None,
});
Expand Down
6 changes: 3 additions & 3 deletions starlark-rust/starlark/src/eval/bc/call.rs
Original file line number Diff line number Diff line change
Expand Up @@ -136,7 +136,7 @@ impl<S: ArgSymbol> BcCallArgs<S> for BcCallArgsFull<S> {
ArgumentsFull {
pos,
named,
names: ArgNames::new(coerce(&self.names)),
names: ArgNames::new_unique(coerce(&self.names)),
args,
kwargs,
}
Expand All @@ -150,7 +150,7 @@ impl<S: ArgSymbol> BcCallArgs<S> for BcCallArgsPos {
ArgumentsFull {
pos,
named: &[],
names: ArgNames::new(&[]),
names: ArgNames::new_unique(&[]),
args: None,
kwargs: None,
}
Expand All @@ -175,7 +175,7 @@ impl BcCallArgsForDef for BcCallArgsFull<ResolvedArgName> {
ArgumentsFull {
pos,
named,
names: ArgNames::new(coerce(&self.names)),
names: ArgNames::new_unique(coerce(&self.names)),
args,
kwargs,
}
Expand Down
2 changes: 1 addition & 1 deletion starlark-rust/starlark/src/eval/compiler/args.rs
Original file line number Diff line number Diff line change
Expand Up @@ -122,7 +122,7 @@ impl ArgsCompiledValue {
Some(handler(&Arguments(ArgumentsFull {
pos: &pos,
named: &named,
names: ArgNames::new(coerce(&self.names)),
names: ArgNames::new_unique(coerce(&self.names)),
args,
kwargs,
})))
Expand Down
67 changes: 35 additions & 32 deletions starlark-rust/starlark/src/eval/runtime/arguments.rs
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,8 @@ use std::marker::PhantomData;
use dupe::Clone_;
use dupe::Dupe_;
use either::Either;
use starlark_map::small_set::SmallSet;
use starlark_syntax::value_error;
use thiserror::Error;

use crate::cast::transmute;
Expand Down Expand Up @@ -120,26 +122,45 @@ unsafe impl Coerce<ResolvedArgName> for ResolvedArgName {}

#[derive(Debug, Clone_, Dupe_)]
pub(crate) struct ArgNames<'a, 'v, S: ArgSymbol> {
/// Names are not guaranteed to be unique here.
/// Names are guaranteed to be unique here.
names: &'a [(S, StringValue<'v>)],
}

impl<'a, 'v, S: ArgSymbol> Default for ArgNames<'a, 'v, S> {
fn default() -> Self {
ArgNames { names: &[] }
Self::new_unique(&[])
}
}

impl<'a, 'v, S: ArgSymbol> Copy for ArgNames<'a, 'v, S> {}

impl<'a, 'v, S: ArgSymbol> ArgNames<'a, 'v, S> {
/// Names are allowed to be not-unique.
/// Names must be unique.
/// String in `Symbol` must be equal to the `StringValue`,
/// it is caller responsibility to ensure that.
pub(crate) fn new(names: &'a [(S, StringValue<'v>)]) -> ArgNames<'a, 'v, S> {
///
/// When this invariant is violated, it is memory safe,
/// but behavior will be incorrect (errors in wrong places, missing errors, panics, etc.)
pub(crate) fn new_unique(names: &'a [(S, StringValue<'v>)]) -> ArgNames<'a, 'v, S> {
ArgNames { names }
}

pub(crate) fn new_check_unique(
names: &'a [(S, StringValue<'v>)],
) -> crate::Result<ArgNames<'a, 'v, S>> {
let mut set = SmallSet::with_capacity(names.len());
for (s, name) in names {
if !set.insert_hashed(Hashed::new_unchecked(s.small_hash(), name.as_str())) {
return Err(value_error!(
"Argument `{}` occurs more than once",
name.as_str()
));
}
}
Ok(Self::new_unique(names))
}

/// Unique names.
pub(crate) fn names(&self) -> &'a [(S, StringValue<'v>)] {
self.names
}
Expand Down Expand Up @@ -264,14 +285,10 @@ impl<'v, 'a> Arguments<'v, 'a> {
None => {
let mut result = SmallMap::with_capacity(self.0.names.names().len());
for (k, v) in self.0.names.names().iter().zip(self.0.named) {
let old =
result.insert_hashed(Hashed::new_unchecked(k.0.small_hash(), k.1), *v);
if unlikely(old.is_some()) {
return Err(FunctionError::RepeatedArg {
name: k.1.as_str().to_owned(),
}
.into());
}
result.insert_hashed_unique_unchecked(
Hashed::new_unchecked(k.0.small_hash(), k.1),
*v,
);
}
Ok(result)
}
Expand All @@ -286,14 +303,10 @@ impl<'v, 'a> Arguments<'v, 'a> {
let mut result =
SmallMap::with_capacity(self.0.names.names().len() + kwargs.len());
for (k, v) in self.0.names.names().iter().zip(self.0.named) {
let old =
result.insert_hashed(Hashed::new_unchecked(k.0.small_hash(), k.1), *v);
if unlikely(old.is_some()) {
return Err(FunctionError::RepeatedArg {
name: k.1.as_str().to_owned(),
}
.into());
}
result.insert_hashed_unique_unchecked(
Hashed::new_unchecked(k.0.small_hash(), k.1),
*v,
);
}
for (k, v) in kwargs.iter_hashed() {
let s = Arguments::unpack_kwargs_key_as_value(*k.key())?;
Expand Down Expand Up @@ -623,14 +636,13 @@ mod tests {
let named = [Value::new_none()];
p.0.named = &named;
let names = [(Symbol::new("test"), heap.alloc_str("test"))];
p.0.names = ArgNames::new(&names);
p.0.names = ArgNames::new_check_unique(&names).unwrap();
assert!(p.no_named_args().is_err());
assert_eq!(p.len().unwrap(), 1);
}

#[test]
fn test_names_map_repeated_name_in_arg_names() {
let named = vec![Value::testing_new_int(10), Value::new_bool(true)];
let names = vec![
(
Symbol::new("a"),
Expand All @@ -641,15 +653,6 @@ mod tests {
const_frozen_string!("a").to_string_value(),
),
];
let error = Arguments(ArgumentsFull {
pos: &[],
named: &named,
names: ArgNames::new(&names),
args: None,
kwargs: None,
})
.names_map()
.unwrap_err();
assert!(error.to_string().contains("occurs more than once"));
assert!(ArgNames::new_check_unique(&names).is_err());
}
}
19 changes: 18 additions & 1 deletion starlark-rust/starlark/src/eval/runtime/params/spec.rs
Original file line number Diff line number Diff line change
Expand Up @@ -593,6 +593,20 @@ impl<'v> ParametersSpec<Value<'v>> {
}
}

#[inline(always)]
fn insert_unique_unchecked(&mut self, key: Hashed<StringValue<'v>>, val: Value<'v>) {
match &mut self.kwargs {
None => {
let mut mp = SmallMap::with_capacity(12);
mp.insert_hashed_unique_unchecked(key, val);
self.kwargs = Some(mp);
}
Some(mp) => {
mp.insert_hashed_unique_unchecked(key, val);
}
}
}

fn alloc(self, heap: &'v Heap) -> Value<'v> {
let kwargs = match self.kwargs {
Some(kwargs) => Dict::new(coerce(kwargs)),
Expand Down Expand Up @@ -639,7 +653,10 @@ impl<'v> ParametersSpec<Value<'v>> {
// Safe to use new_unchecked because hash for the Value and str are the same
match name.get_index_from_param_spec(self) {
None => {
kwargs.insert(Hashed::new_unchecked(name.small_hash(), *name_value), *v);
kwargs.insert_unique_unchecked(
Hashed::new_unchecked(name.small_hash(), *name_value),
*v,
);
}
Some(i) => {
slots[i] = Some(*v);
Expand Down
26 changes: 24 additions & 2 deletions starlark-rust/starlark/src/stdlib/partial.rs
Original file line number Diff line number Diff line change
Expand Up @@ -19,11 +19,13 @@ use std::fmt;
use std::fmt::Display;

use allocative::Allocative;
use hashbrown::HashTable;
use starlark_derive::starlark_module;
use starlark_derive::starlark_value;
use starlark_derive::NoSerialize;
use starlark_syntax::slice_vec_ext::SliceExt;
use starlark_syntax::slice_vec_ext::VecExt;
use starlark_syntax::value_error;

use crate as starlark;
use crate::any::ProvidesStaticType;
Expand Down Expand Up @@ -60,7 +62,7 @@ pub fn partial(builder: &mut GlobalsBuilder) {
#[starlark(kwargs)] kwargs: DictRef<'v>,
) -> anyhow::Result<Partial<'v>> {
debug_assert!(Tuple::from_value(args).is_some());
let names = kwargs
let names: Vec<_> = kwargs
.keys()
.map(|x| {
let x = StringValue::new(x).unwrap();
Expand All @@ -72,11 +74,16 @@ pub fn partial(builder: &mut GlobalsBuilder) {
)
})
.collect();
let mut names_index = HashTable::with_capacity(names.len());
for (i, (k, _)) in names.iter().enumerate() {
names_index.insert_unique(k.hash(), i, |i| names[*i].0.hash());
}
Ok(Partial {
func,
pos: args,
named: kwargs.values().collect(),
names,
names_index,
})
}
}
Expand All @@ -89,6 +96,7 @@ struct PartialGen<V, S> {
pos: V,
named: Vec<V>,
names: Vec<(Symbol, S)>,
names_index: HashTable<usize>,
}

impl<'v, V: ValueLike<'v>, S> PartialGen<V, S> {
Expand Down Expand Up @@ -132,6 +140,7 @@ impl<'v> Freeze for Partial<'v> {
names: self
.names
.into_try_map(|(s, x)| anyhow::Ok((s, x.freeze(freezer)?)))?,
names_index: self.names_index,
})
}
}
Expand Down Expand Up @@ -159,13 +168,26 @@ where
let self_named = coerce(&self.named);
let self_names = coerce(&self.names);

for (symbol, _) in args.0.names.names() {
if self
.names_index
.find(symbol.hash(), |i| &self.names[*i].0 == symbol)
.is_some()
{
return Err(value_error!(
"partial() got multiple values for argument `{}`",
symbol.as_str(),
));
}
}

eval.alloca_concat(self_pos, args.0.pos, |pos, eval| {
eval.alloca_concat(self_named, args.0.named, |named, eval| {
eval.alloca_concat(self_names, args.0.names.names(), |names, eval| {
let params = Arguments(ArgumentsFull {
pos,
named,
names: ArgNames::new(names),
names: ArgNames::new_unique(names),
args: args.0.args,
kwargs: args.0.kwargs,
});
Expand Down

0 comments on commit 952a5d5

Please sign in to comment.