use indexmap::IndexMap;
use super::{super::MethodOutcome, dict};
use crate::{
error::{EvalError, InterpreterError},
eval::{control_flow::iterate_value, literals::value_to_key},
value::{Value, ValueKey, shared_list},
};
pub(crate) fn dispatch_counter_method(
map: &mut IndexMap<ValueKey, Value>,
method: &str,
args: &[Value],
kwargs: &IndexMap<String, Value>,
) -> Result<MethodOutcome, EvalError> {
match method {
"keys" | "values" | "items" | "get" | "copy" | "pop" | "setdefault" | "clear" => {
let outcome = dict::dispatch_dict_method(map, method, args, kwargs)?;
let value = if method == "copy" {
match outcome.value {
Value::Dict(d) => Value::Counter(d.lock().clone()),
other => other,
}
} else {
outcome.value
};
Ok(MethodOutcome { value, mem_delta: outcome.mem_delta })
}
"most_common" => {
crate::eval::functions::reject_kwargs(method, kwargs)?;
let mut entries: Vec<(ValueKey, i64)> = map
.iter()
.map(|(k, v)| {
let n = match v {
Value::Int(i) => *i,
Value::Bool(b) => i64::from(*b),
_ => 0,
};
(k.clone(), n)
})
.collect();
entries.sort_by_key(|entry| std::cmp::Reverse(entry.1));
let n = match args.first() {
Some(Value::Int(i)) => usize::try_from((*i).max(0)).unwrap_or(entries.len()),
Some(Value::Bool(b)) => usize::from(*b),
None | Some(Value::None) => entries.len(),
Some(_) => {
return Err(InterpreterError::TypeError(
"most_common(): n must be an integer or None".into(),
)
.into());
}
};
let result: Vec<Value> = entries
.into_iter()
.take(n)
.map(|(k, n)| Value::Tuple(vec![k.to_value(), Value::Int(n)]))
.collect();
Ok(MethodOutcome::pure(Value::List(shared_list(result))))
}
"elements" => {
crate::eval::functions::reject_kwargs(method, kwargs)?;
let mut out: Vec<Value> = Vec::new();
for (key, value) in map.iter() {
let n = match value {
Value::Int(i) => *i,
Value::Bool(b) => i64::from(*b),
_ => 0,
};
if n > 0 {
let key_val = key.to_value();
for _ in 0..n {
out.push(key_val.clone());
}
}
}
Ok(MethodOutcome::pure(Value::List(shared_list(out))))
}
"subtract" => {
let mut delta = counter_apply_in_place(map, args.first(), |cur, delta| cur - delta)?;
delta += counter_apply_kwargs(map, kwargs, |cur, delta| cur - delta);
Ok(counter_mem_outcome(delta))
}
"update" => {
let mut delta = counter_apply_in_place(map, args.first(), |cur, delta| cur + delta)?;
delta += counter_apply_kwargs(map, kwargs, |cur, delta| cur + delta);
Ok(counter_mem_outcome(delta))
}
"total" => {
crate::eval::functions::reject_kwargs(method, kwargs)?;
let total: i64 = map
.values()
.map(|v| match v {
Value::Int(i) => *i,
Value::Bool(b) => i64::from(*b),
_ => 0,
})
.sum();
Ok(MethodOutcome::pure(Value::Int(total)))
}
_ => Err(InterpreterError::AttributeError(format!(
"'Counter' object has no attribute '{method}'"
))
.into()),
}
}
fn counter_mem_outcome(mem_delta: isize) -> MethodOutcome {
MethodOutcome { value: Value::None, mem_delta }
}
fn counter_apply_kwargs(
map: &mut IndexMap<ValueKey, Value>,
kwargs: &IndexMap<String, Value>,
op: fn(i64, i64) -> i64,
) -> isize {
let mut mem_delta = 0;
for (name, value) in kwargs {
let amount = match value {
Value::Int(i) => *i,
Value::Bool(b) => i64::from(*b),
_ => continue,
};
let key = ValueKey::String(name.as_str().into());
let cur = match map.get(&key) {
Some(Value::Int(i)) => *i,
Some(Value::Bool(b)) => i64::from(*b),
_ => 0,
};
mem_delta += crate::eval::functions::methods::dict::insert_entry(
map,
key,
Value::Int(op(cur, amount)),
);
}
mem_delta
}
fn counter_apply_in_place(
map: &mut IndexMap<ValueKey, Value>,
other_arg: Option<&Value>,
op: fn(i64, i64) -> i64,
) -> Result<isize, EvalError> {
let Some(other) = other_arg else { return Ok(0) };
let insert_entry = crate::eval::functions::methods::dict::insert_entry;
let mut mem_delta = 0;
let other_map = match other {
Value::Dict(m) => Some(m.lock().clone()),
Value::Counter(m) => Some(m.clone()),
_ => None,
};
if let Some(other_map) = other_map {
for (k, v) in &other_map {
let amount = match v {
Value::Int(i) => *i,
Value::Bool(b) => i64::from(*b),
_ => 0,
};
let cur = match map.get(k) {
Some(Value::Int(i)) => *i,
Some(Value::Bool(b)) => i64::from(*b),
_ => 0,
};
mem_delta += insert_entry(map, k.clone(), Value::Int(op(cur, amount)));
}
return Ok(mem_delta);
}
for item in iterate_value(other)? {
let key = value_to_key(&item)?;
let cur = match map.get(&key) {
Some(Value::Int(i)) => *i,
Some(Value::Bool(b)) => i64::from(*b),
_ => 0,
};
mem_delta += insert_entry(map, key, Value::Int(op(cur, 1)));
}
Ok(mem_delta)
}