use num_traits::{Signed, ToPrimitive};
use tycho_types::cell::{LoadMode, RefsIter};
use tycho_types::error::Error;
use tycho_types::prelude::*;
use tycho_vm_proc::vm_module;
use crate::GasConsumer;
use crate::error::VmResult;
use crate::saferc::SafeRc;
use crate::stack::StackValueType;
use crate::state::VmState;
pub struct SizeOps;
#[vm_module]
impl SizeOps {
#[op(code = "f940", fmt = "CDATASIZEQ", args(is_slice = false, q = true))]
#[op(code = "f941", fmt = "CDATASIZE", args(is_slice = false, q = false))]
#[op(code = "f942", fmt = "SDATASIZEQ", args(is_slice = true, q = true))]
#[op(code = "f943", fmt = "SDATASIZE", args(is_slice = true, q = false))]
pub fn exec_compute_data_size(st: &mut VmState, is_slice: bool, q: bool) -> VmResult<i32> {
let stack = SafeRc::make_mut(&mut st.stack);
let bound = ok!(stack.pop_int_or_nan());
let value = if is_slice {
ok!(stack.pop_cs()).into_dyn_value()
} else {
ok!(stack.pop_cell()).into_dyn_value()
};
let bound = match bound {
Some(bound) if !bound.is_negative() => SafeRc::unwrap_or_clone(bound),
_ => vm_bail!(IntegerOverflow),
};
let limit = bound.to_u64().unwrap_or(u64::MAX);
let mut storage = StorageStatExt::with_limit(&st.gas, limit);
let ok = if is_slice {
let Some(slice) = value.as_cell_slice() else {
vm_bail!(InvalidType {
expected: StackValueType::Slice as _,
actual: value.raw_ty()
})
};
let cs = slice.apply();
ok!(storage.add_slice(&cs))
} else {
let Some(cell) = value.as_cell() else {
vm_bail!(InvalidType {
expected: StackValueType::Cell as _,
actual: value.raw_ty()
})
};
ok!(storage.add_cell(cell.as_ref()))
};
if ok {
ok!(stack.push_int(storage.stats.cell_count));
ok!(stack.push_int(storage.stats.bit_count));
ok!(stack.push_int(storage.stats.ref_count));
} else if !q {
vm_bail!(CellError(Error::CellOverflow));
}
if q {
ok!(stack.push_bool(ok))
}
Ok(0)
}
}
struct StorageStatExt<'a, 'c: 'a, 'l> {
gas: &'c GasConsumer<'l>,
visited: ahash::HashSet<&'a HashBytes>,
stack: Vec<RefsIter<'a>>,
stats: CellTreeStatsExt,
limit: u64,
}
impl<'a, 'c: 'a, 'l> StorageStatExt<'a, 'c, 'l> {
pub fn with_limit(gas: &'c GasConsumer<'l>, limit: u64) -> Self {
Self {
gas,
visited: Default::default(),
stack: Vec::new(),
stats: Default::default(),
limit,
}
}
fn add_cell(&mut self, mut cell: &'a DynCell) -> VmResult<bool> {
if !self.visited.insert(cell.repr_hash()) {
return Ok(true);
}
if self.stats.cell_count >= self.limit {
return Ok(false);
}
cell = self.gas.load_dyn_cell(cell, LoadMode::UseGas)?;
let refs = cell.references();
self.stats.bit_count += cell.bit_len() as u64;
self.stats.ref_count += refs.len() as u64;
self.stats.cell_count += 1;
self.stack.clear();
self.stack.push(refs);
self.reduce_stack()
}
fn add_slice(&mut self, slice: &CellSlice<'a>) -> VmResult<bool> {
let refs = slice.references();
self.stats.bit_count += slice.size_bits() as u64;
self.stats.ref_count += refs.len() as u64;
self.stack.clear();
self.stack.push(refs);
self.reduce_stack()
}
fn reduce_stack(&mut self) -> VmResult<bool> {
'outer: while let Some(item) = self.stack.last_mut() {
for mut cell in item.by_ref() {
if !self.visited.insert(cell.repr_hash()) {
continue;
}
if self.stats.cell_count >= self.limit {
return Ok(false);
}
cell = self.gas.load_dyn_cell(cell, LoadMode::UseGas)?;
let next = cell.references();
let ref_count = next.len();
self.stats.bit_count += cell.bit_len() as u64;
self.stats.ref_count += ref_count as u64;
self.stats.cell_count += 1;
if ref_count > 0 {
self.stack.push(next);
continue 'outer;
}
}
self.stack.pop();
}
Ok(true)
}
}
#[derive(Default)]
struct CellTreeStatsExt {
bit_count: u64,
ref_count: u64,
cell_count: u64,
}
#[cfg(test)]
mod tests {
use rand::Rng;
use tracing_test::traced_test;
use tycho_types::prelude::*;
use crate::OwnedCellSlice;
#[test]
#[traced_test]
fn data_size() {
let empty_cell = Cell::default();
assert_run_vm!("CDATASIZE", [cell empty_cell.clone(), int 10] => [int 1, int 0, int 0]);
assert_run_vm!("CDATASIZE", [cell empty_cell.clone(), int 0] => [int 0], exit_code: 8);
assert_run_vm!("CDATASIZEQ", [cell empty_cell.clone(), int 10] => [int 1, int 0, int 0, int -1]);
assert_run_vm!("CDATASIZEQ", [cell empty_cell.clone(), int 0] => [int 0]);
assert_run_vm!("CDATASIZEQ", [cell empty_cell.clone(), nan] => [int 0], exit_code: 4);
let empty_slice = OwnedCellSlice::new_allow_exotic(empty_cell);
assert_run_vm!("SDATASIZE", [slice empty_slice.clone(), int 10] => [int 0, int 0, int 0]);
assert_run_vm!("SDATASIZE", [slice empty_slice.clone(), int 0] => [int 0, int 0, int 0]);
assert_run_vm!("SDATASIZEQ", [slice empty_slice.clone(), int 10] => [int 0, int 0, int 0, int -1]);
assert_run_vm!("SDATASIZEQ", [slice empty_slice, int 0] => [int 0, int 0, int 0, int -1]);
let plain_cell = CellBuilder::build_from((123u8, 123123123u32, false)).unwrap();
assert_run_vm!("CDATASIZE", [cell plain_cell.clone(), int 1] => [int 1, int 8 + 32 + 1, int 0]);
assert_run_vm!("CDATASIZEQ", [cell plain_cell.clone(), int 1] => [int 1, int 8 + 32 + 1, int 0, int -1]);
let plain_slice = OwnedCellSlice::new_allow_exotic(plain_cell);
assert_run_vm!("SDATASIZE", [slice plain_slice.clone(), int 1] => [int 0, int 8 + 32 + 1, int 0]);
assert_run_vm!("SDATASIZEQ", [slice plain_slice.clone(), int 1] => [int 0, int 8 + 32 + 1, int 0, int -1]);
let one_ref_cell =
CellBuilder::build_from((123u8, 123123123u32, false, Cell::default())).unwrap();
assert_run_vm!("CDATASIZE", [cell one_ref_cell.clone(), int 2] => [int 2, int 8 + 32 + 1, int 1]);
assert_run_vm!("CDATASIZEQ", [cell one_ref_cell.clone(), int 2] => [int 2, int 8 + 32 + 1, int 1, int -1]);
assert_run_vm!("CDATASIZE", [cell one_ref_cell.clone(), int 1] => [int 0], exit_code: 8);
assert_run_vm!("CDATASIZEQ", [cell one_ref_cell.clone(), int 1] => [int 0]);
let one_ref_slice = OwnedCellSlice::new_allow_exotic(one_ref_cell);
assert_run_vm!("SDATASIZE", [slice one_ref_slice.clone(), int 1] => [int 1, int 8 + 32 + 1, int 1]);
assert_run_vm!("SDATASIZEQ", [slice one_ref_slice.clone(), int 1] => [int 1, int 8 + 32 + 1, int 1, int -1]);
assert_run_vm!("SDATASIZE", [slice one_ref_slice.clone(), int 0] => [int 0], exit_code: 8);
assert_run_vm!("SDATASIZEQ", [slice one_ref_slice.clone(), int 0] => [int 0]);
let mut deep_cell = CellBuilder::build_from(HashBytes::ZERO).unwrap();
for _ in 0..100 {
deep_cell = CellBuilder::build_from((
deep_cell.clone(),
deep_cell.clone(),
deep_cell.clone(),
deep_cell.clone(),
))
.unwrap();
}
assert_run_vm!("CDATASIZE", [cell deep_cell.clone(), int 101] => [int 101, int 256, int 100 * 4]);
assert_run_vm!("CDATASIZEQ", [cell deep_cell.clone(), int 100] => [int 0]);
let deep_slice = OwnedCellSlice::new_allow_exotic(deep_cell);
assert_run_vm!("SDATASIZE", [slice deep_slice.clone(), int 100] => [int 100, int 256, int 100 * 4]);
assert_run_vm!("SDATASIZEQ", [slice deep_slice.clone(), int 99] => [int 0]);
fn make_huge_cell(rng: &mut impl Rng, depth: u8) -> Cell {
if depth == 0 {
CellBuilder::build_from(rng.random::<HashBytes>()).unwrap()
} else {
CellBuilder::build_from((
rng.random::<HashBytes>(),
make_huge_cell(rng, depth - 1),
make_huge_cell(rng, depth - 1),
))
.unwrap()
}
}
let huge_cell = make_huge_cell(&mut rand::rng(), 10);
assert_run_vm!("CDATASIZE", [cell huge_cell.clone(), int 2048] => [int 2047, int 524032, int 2046]);
assert_run_vm!("CDATASIZEQ", [cell huge_cell.clone(), int 100] => [int 0]);
assert_run_vm!("CDATASIZE", gas: 10000, [cell huge_cell.clone(), int 2048] => [int 10026], exit_code: -14);
assert_run_vm!("CDATASIZEQ", gas: 10000, [cell huge_cell.clone(), int 2048] => [int 10026], exit_code: -14);
let huge_slice = OwnedCellSlice::new_allow_exotic(huge_cell);
assert_run_vm!("SDATASIZE", [slice huge_slice.clone(), int 2048] => [int 2046, int 524032, int 2046]);
assert_run_vm!("SDATASIZEQ", [slice huge_slice.clone(), int 100] => [int 0]);
assert_run_vm!("SDATASIZE", gas: 10000, [slice huge_slice.clone(), int 2048] => [int 10026], exit_code: -14);
assert_run_vm!("SDATASIZEQ", gas: 10000, [slice huge_slice.clone(), int 2048] => [int 10026], exit_code: -14);
}
}