use crate::{MattenError, Tensor};
pub(crate) const MAX_NDIM: usize = 8;
pub(crate) const MAX_ELEMENTS: usize = 1 << 20;
#[cfg(feature = "json")]
pub(crate) const MAX_JSON_ELEMENTS: usize = 1 << 24;
#[cfg(all(feature = "dynamic", feature = "json"))]
pub(crate) const MAX_DYNAMIC_ELEMENTS: usize = 1 << 24;
pub(crate) const MAX_SLICE_STR_BYTES: usize = 512;
pub(crate) const MAX_PARSE_BYTES: usize = 128 * 1024 * 1024;
#[derive(Debug, Clone, PartialEq)]
pub struct MattenLimits {
pub max_dimensions: usize,
pub max_elements: usize,
pub max_parse_bytes: usize,
}
impl Default for MattenLimits {
fn default() -> Self {
Self {
max_dimensions: MAX_NDIM,
max_elements: MAX_ELEMENTS,
max_parse_bytes: MAX_PARSE_BYTES,
}
}
}
impl MattenLimits {
pub fn strict() -> Self {
Self {
max_dimensions: 4,
max_elements: 1024,
max_parse_bytes: 64 * 1024,
}
}
pub(crate) fn check_elements(
&self,
requested: usize,
operation: &'static str,
) -> Result<(), MattenError> {
if requested > self.max_elements {
Err(MattenError::Allocation {
requested_elements: requested,
message: format!(
"{operation} requested {requested} elements, exceeding the \
limit of {} (MattenLimits::max_elements); use smaller shapes \
or increase the limit",
self.max_elements
),
})
} else {
Ok(())
}
}
pub(crate) fn check_shape(
&self,
shape: &[usize],
operation: &'static str,
) -> Result<usize, MattenError> {
if shape.len() > self.max_dimensions {
return Err(MattenError::Shape {
operation,
message: format!(
"rank {} exceeds the maximum supported rank of {} (shape {shape:?})",
shape.len(),
self.max_dimensions
),
});
}
let len = crate::shape::checked_shape_len(shape, operation)?;
self.check_elements(len, operation)?;
Ok(len)
}
}
impl Tensor {
pub fn try_zeros(shape: &[usize]) -> Result<Tensor, MattenError> {
Tensor::try_zeros_with_limits(shape, &MattenLimits::default())
}
pub fn try_zeros_with_limits(
shape: &[usize],
limits: &MattenLimits,
) -> Result<Tensor, MattenError> {
let len = limits.check_shape(shape, "try_zeros")?;
Ok(Tensor {
data: vec![0.0f64; len],
shape: shape.to_vec(),
#[cfg(feature = "dynamic")]
dynamic: None,
})
}
pub fn try_ones(shape: &[usize]) -> Result<Tensor, MattenError> {
Tensor::try_ones_with_limits(shape, &MattenLimits::default())
}
pub fn try_ones_with_limits(
shape: &[usize],
limits: &MattenLimits,
) -> Result<Tensor, MattenError> {
let len = limits.check_shape(shape, "try_ones")?;
Ok(Tensor {
data: vec![1.0f64; len],
shape: shape.to_vec(),
#[cfg(feature = "dynamic")]
dynamic: None,
})
}
pub fn try_full(shape: &[usize], value: f64) -> Result<Tensor, MattenError> {
Tensor::try_full_with_limits(shape, value, &MattenLimits::default())
}
pub fn try_full_with_limits(
shape: &[usize],
value: f64,
limits: &MattenLimits,
) -> Result<Tensor, MattenError> {
let len = limits.check_shape(shape, "try_full")?;
Ok(Tensor {
data: vec![value; len],
shape: shape.to_vec(),
#[cfg(feature = "dynamic")]
dynamic: None,
})
}
}