use std::sync::Arc;
use super::super::{MethodOutcome, arg1};
use crate::{
error::{EvalError, InterpreterError},
eval::{control_flow::iterate_value, literals::value_to_key},
pyset::SetBody,
state::estimate_value_size,
value::{ExceptionValue, SharedSet, Value},
};
fn set_value(body: SetBody) -> Value {
Value::Set(crate::value::shared_set(body))
}
fn arg_body(arg: &Value) -> Result<SetBody, EvalError> {
match arg {
Value::Set(s) => Ok(s.lock().clone()),
Value::Frozenset(f) => Ok((**f).clone()),
other => Ok(SetBody::from_items(iterate_value(other)?)),
}
}
fn apply_union(acc: &mut SetBody, arg: &Value) -> Result<(), EvalError> {
match arg {
Value::Set(s) => acc.merge_from(&s.lock()),
Value::Frozenset(f) => acc.merge_from(f),
other => {
for item in iterate_value(other)? {
acc.add_value(item);
}
}
}
Ok(())
}
fn intersect_arg(acc: &SetBody, arg: &Value) -> Result<SetBody, EvalError> {
match arg {
Value::Set(s) => Ok(acc.intersection_with(&s.lock())),
Value::Frozenset(f) => Ok(acc.intersection_with(f)),
other => {
let mut r = SetBody::empty();
for item in iterate_value(other)? {
if acc.contains(&item) {
r.add_value(item);
}
}
Ok(r)
}
}
}
fn apply_difference(acc: &mut SetBody, arg: &Value) -> Result<(), EvalError> {
match arg {
Value::Set(s) => acc.difference_from(&s.lock()),
Value::Frozenset(f) => acc.difference_from(f),
other => {
let items = iterate_value(other)?;
acc.difference_from(&SetBody::from_items(items));
}
}
Ok(())
}
fn body_bytes(body: &SetBody) -> usize {
body.iter_ordered().iter().map(estimate_value_size).sum()
}
fn delta_outcome(old: usize, new: usize) -> MethodOutcome {
if new >= old {
MethodOutcome::grew(Value::None, new - old)
} else {
MethodOutcome::shrank(Value::None, old - new)
}
}
pub(crate) fn dispatch_set_method(
shared: &SharedSet,
method: &str,
args: &[Value],
kwargs: &indexmap::IndexMap<String, Value>,
) -> Result<MethodOutcome, EvalError> {
crate::eval::functions::reject_kwargs(method, kwargs)?;
match method {
"copy" => Ok(MethodOutcome::pure(set_value(shared.lock().copied()))),
"union" => {
let mut acc = shared.lock().copied();
for arg in args {
apply_union(&mut acc, arg)?;
}
Ok(MethodOutcome::pure(set_value(acc)))
}
"intersection" => {
if args.is_empty() {
return Ok(MethodOutcome::pure(set_value(shared.lock().copied())));
}
let mut acc = shared.lock().clone();
for arg in args {
acc = intersect_arg(&acc, arg)?;
}
Ok(MethodOutcome::pure(set_value(acc)))
}
"difference" => {
let mut acc = shared.lock().copied();
for arg in args {
apply_difference(&mut acc, arg)?;
}
Ok(MethodOutcome::pure(set_value(acc)))
}
"symmetric_difference" => {
let other = arg_body(arg1(method, args)?)?;
let result = shared.lock().symmetric_difference_with(&other);
Ok(MethodOutcome::pure(set_value(result)))
}
"issubset" => {
let other = arg_body(arg1(method, args)?)?;
let result = shared.lock().iter_ordered().iter().all(|v| other.contains(v));
Ok(MethodOutcome::pure(Value::Bool(result)))
}
"issuperset" => {
let other = iterate_value(arg1(method, args)?)?;
let body = shared.lock();
let result = other.iter().all(|v| body.contains(v));
Ok(MethodOutcome::pure(Value::Bool(result)))
}
"isdisjoint" => {
let other = iterate_value(arg1(method, args)?)?;
let body = shared.lock();
let result = !other.iter().any(|v| body.contains(v));
Ok(MethodOutcome::pure(Value::Bool(result)))
}
"add" => {
let arg = arg1(method, args)?;
if !matches!(arg, Value::Instance(_)) {
value_to_key(arg)?;
}
let size = estimate_value_size(arg);
if shared.lock().add_value(arg.clone()) {
Ok(MethodOutcome::grew(Value::None, size))
} else {
Ok(MethodOutcome::pure(Value::None))
}
}
"remove" => {
let arg = arg1(method, args)?;
let freed = estimate_value_size(arg);
if shared.lock().discard_value(arg) {
Ok(MethodOutcome::shrank(Value::None, freed))
} else {
Err(EvalError::Exception(ExceptionValue::new("KeyError", arg.repr())))
}
}
"discard" => {
let arg = arg1(method, args)?;
let freed = estimate_value_size(arg);
if shared.lock().discard_value(arg) {
Ok(MethodOutcome::shrank(Value::None, freed))
} else {
Ok(MethodOutcome::pure(Value::None))
}
}
"pop" => match shared.lock().pop_first() {
Some(val) => {
let freed = estimate_value_size(&val);
Ok(MethodOutcome::shrank(val, freed))
}
None => {
Err(EvalError::Exception(ExceptionValue::new(
"KeyError",
"'pop from an empty set'",
)))
}
},
"clear" => {
let mut body = shared.lock();
let freed: usize = body.iter_ordered().iter().map(estimate_value_size).sum();
body.clear();
Ok(MethodOutcome::shrank(Value::None, freed))
}
"update" => {
let mut acc = shared.lock().clone();
let old = body_bytes(&acc);
for arg in args {
apply_union(&mut acc, arg)?;
}
let new = body_bytes(&acc);
*shared.lock() = acc;
Ok(delta_outcome(old, new))
}
"intersection_update" => {
let mut acc = shared.lock().clone();
let old = body_bytes(&acc);
for arg in args {
acc = intersect_arg(&acc, arg)?;
}
let new = body_bytes(&acc);
*shared.lock() = acc;
Ok(delta_outcome(old, new))
}
"difference_update" => {
let mut acc = shared.lock().clone();
let old = body_bytes(&acc);
for arg in args {
apply_difference(&mut acc, arg)?;
}
let new = body_bytes(&acc);
*shared.lock() = acc;
Ok(delta_outcome(old, new))
}
"symmetric_difference_update" => {
let other = arg_body(arg1(method, args)?)?;
let mut acc = shared.lock().clone();
let old = body_bytes(&acc);
for item in other.iter_ordered() {
if !acc.discard_value(&item) {
acc.add_value(item);
}
}
let new = body_bytes(&acc);
*shared.lock() = acc;
Ok(delta_outcome(old, new))
}
_ => Err(InterpreterError::AttributeError(format!(
"'set' object has no attribute '{method}'"
))
.into()),
}
}
pub(crate) fn dispatch_frozenset_method(
body: &SetBody,
method: &str,
args: &[Value],
kwargs: &indexmap::IndexMap<String, Value>,
) -> Result<MethodOutcome, EvalError> {
const FROZENSET_METHODS: &[&str] = &[
"copy",
"union",
"intersection",
"difference",
"symmetric_difference",
"issubset",
"issuperset",
"isdisjoint",
];
if !FROZENSET_METHODS.contains(&method) {
return Err(InterpreterError::AttributeError(format!(
"'frozenset' object has no attribute '{method}'"
))
.into());
}
let scratch = crate::value::shared_set(body.clone());
let outcome = dispatch_set_method(&scratch, method, args, kwargs)?;
let value = match outcome.value {
Value::Set(v) => Value::Frozenset(Arc::new(v.lock().clone())),
other => other,
};
Ok(MethodOutcome::pure(value))
}