use std::iter;
use std::sync::Arc;
use arbitrary::Result;
use arbitrary::Unstructured;
use vortex_buffer::BufferString;
use vortex_buffer::ByteBuffer;
use vortex_dtype::DType;
use vortex_dtype::DecimalDType;
use vortex_dtype::NativeDecimalType;
use vortex_dtype::PType;
use vortex_dtype::half::f16;
use vortex_dtype::match_each_decimal_value_type;
use crate::DecimalValue;
use crate::InnerScalarValue;
use crate::PValue;
use crate::Scalar;
use crate::ScalarValue;
pub fn random_scalar(u: &mut Unstructured, dtype: &DType) -> Result<Scalar> {
Ok(Scalar::new(dtype.clone(), random_scalar_value(u, dtype)?))
}
fn random_scalar_value(u: &mut Unstructured, dtype: &DType) -> Result<ScalarValue> {
match dtype {
DType::Null => Ok(ScalarValue(InnerScalarValue::Null)),
DType::Bool(_) => Ok(ScalarValue(InnerScalarValue::Bool(u.arbitrary()?))),
DType::Primitive(p, _) => Ok(ScalarValue(InnerScalarValue::Primitive(random_pvalue(
u, p,
)?))),
DType::Decimal(decimal_type, _) => random_decimal(u, decimal_type),
DType::Utf8(_) => Ok(ScalarValue(InnerScalarValue::BufferString(Arc::new(
BufferString::from(u.arbitrary::<String>()?),
)))),
DType::Binary(_) => Ok(ScalarValue(InnerScalarValue::Buffer(Arc::new(
ByteBuffer::from(u.arbitrary::<Vec<u8>>()?),
)))),
DType::Struct(sdt, _) => Ok(ScalarValue(InnerScalarValue::List(
sdt.fields()
.map(|d| random_scalar_value(u, &d))
.collect::<Result<Vec<_>>>()?
.into(),
))),
DType::List(edt, _) => Ok(ScalarValue(InnerScalarValue::List(
iter::from_fn(|| {
u.arbitrary()
.unwrap_or(false)
.then(|| random_scalar_value(u, edt))
})
.collect::<Result<Vec<_>>>()?
.into(),
))),
DType::FixedSizeList(edt, size, _) => Ok(ScalarValue(InnerScalarValue::List(
(0..*size)
.map(|_| random_scalar_value(u, edt))
.collect::<Result<Vec<_>>>()?
.into(),
))),
DType::Extension(..) => {
unreachable!("Can't yet generate arbitrary scalars for ext dtype")
}
}
}
fn random_pvalue(u: &mut Unstructured, ptype: &PType) -> Result<PValue> {
Ok(match ptype {
PType::U8 => PValue::U8(u.arbitrary()?),
PType::U16 => PValue::U16(u.arbitrary()?),
PType::U32 => PValue::U32(u.arbitrary()?),
PType::U64 => PValue::U64(u.arbitrary()?),
PType::I8 => PValue::I8(u.arbitrary()?),
PType::I16 => PValue::I16(u.arbitrary()?),
PType::I32 => PValue::I32(u.arbitrary()?),
PType::I64 => PValue::I64(u.arbitrary()?),
PType::F16 => PValue::F16(f16::from_bits(u.arbitrary()?)),
PType::F32 => PValue::F32(u.arbitrary()?),
PType::F64 => PValue::F64(u.arbitrary()?),
})
}
pub fn random_decimal(u: &mut Unstructured, decimal_type: &DecimalDType) -> Result<ScalarValue> {
let precision = decimal_type.precision();
let value = match_each_decimal_value_type!(
DecimalType::smallest_decimal_value_type(decimal_type),
|D| {
DecimalValue::from(u.int_in_range(
D::MIN_BY_PRECISION[precision as usize]..=D::MAX_BY_PRECISION[precision as usize],
)?)
}
);
Ok(ScalarValue(InnerScalarValue::Decimal(value)))
}