Skip to content

Commit

Permalink
Rewrite load_cell_data
Browse files Browse the repository at this point in the history
  • Loading branch information
@mohanson
mohanson committed Mar 9, 2024
1 parent 6be39c2 commit d4b52f1
Show file tree
Hide file tree
Showing 4 changed files with 242 additions and 296 deletions.
201 changes: 88 additions & 113 deletions script/src/syscalls/load_cell_data.rs
View file
Original file line number Diff line number Diff line change
@@ -1,162 +1,137 @@
use crate::types::Indices;
use crate::v2_types::{DataPieceId, TxData};
use crate::{
cost_model::transferred_byte_cycles,
syscalls::{
utils::store_data, Source, SourceEntry, INDEX_OUT_OF_BOUND,
LOAD_CELL_DATA_AS_CODE_SYSCALL_NUMBER, LOAD_CELL_DATA_SYSCALL_NUMBER, SLICE_OUT_OF_BOUND,
SUCCESS,
INDEX_OUT_OF_BOUND, LOAD_CELL_DATA_AS_CODE_SYSCALL_NUMBER, LOAD_CELL_DATA_SYSCALL_NUMBER,
SLICE_OUT_OF_BOUND, SUCCESS,
},
};
use ckb_traits::CellDataProvider;
use ckb_types::core::cell::{CellMeta, ResolvedTransaction};
use ckb_traits::{CellDataProvider, ExtensionProvider, HeaderProvider};
use ckb_vm::{
memory::{Memory, FLAG_EXECUTABLE, FLAG_FREEZED},
registers::{A0, A1, A2, A3, A4, A5, A7},
snapshot2::{DataSource, Snapshot2Context},
Error as VMError, Register, SupportMachine, Syscalls,
};
use std::sync::Arc;
use std::sync::{Arc, Mutex};

pub struct LoadCellData<DL> {
data_loader: DL,
rtx: Arc<ResolvedTransaction>,
outputs: Arc<Vec<CellMeta>>,
group_inputs: Indices,
group_outputs: Indices,
pub struct LoadCellData<DL>
where
DL: CellDataProvider + HeaderProvider + ExtensionProvider + Send + Sync + Clone + 'static,
{
snapshot2_context: Arc<Mutex<Snapshot2Context<DataPieceId, TxData<DL>>>>,
}

impl<DL: CellDataProvider> LoadCellData<DL> {
impl<DL> LoadCellData<DL>
where
DL: CellDataProvider + HeaderProvider + ExtensionProvider + Send + Sync + Clone + 'static,
{
pub fn new(
data_loader: DL,
rtx: Arc<ResolvedTransaction>,
outputs: Arc<Vec<CellMeta>>,
group_inputs: Indices,
group_outputs: Indices,
snapshot2_context: Arc<Mutex<Snapshot2Context<DataPieceId, TxData<DL>>>>,
) -> LoadCellData<DL> {
LoadCellData {
data_loader,
rtx,
outputs,
group_inputs,
group_outputs,
}
}

#[inline]
fn resolved_inputs(&self) -> &Vec<CellMeta> {
&self.rtx.resolved_inputs
}

#[inline]
fn resolved_cell_deps(&self) -> &Vec<CellMeta> {
&self.rtx.resolved_cell_deps
LoadCellData { snapshot2_context }
}

fn fetch_cell(&self, source: Source, index: usize) -> Result<&CellMeta, u8> {
match source {
Source::Transaction(SourceEntry::Input) => {
self.resolved_inputs().get(index).ok_or(INDEX_OUT_OF_BOUND)
}
Source::Transaction(SourceEntry::Output) => {
self.outputs.get(index).ok_or(INDEX_OUT_OF_BOUND)
fn load_data<Mac: SupportMachine>(&self, machine: &mut Mac) -> Result<(), VMError> {
let index = machine.registers()[A3].to_u64();
let source = machine.registers()[A4].to_u64();
let data_piece_id = match DataPieceId::try_from((source, index, 0)) {
Ok(id) => id,
Err(_) => {
machine.set_register(A0, Mac::REG::from_u8(INDEX_OUT_OF_BOUND));
return Ok(());
}
Source::Transaction(SourceEntry::CellDep) => self
.resolved_cell_deps()
.get(index)
.ok_or(INDEX_OUT_OF_BOUND),
Source::Transaction(SourceEntry::HeaderDep) => Err(INDEX_OUT_OF_BOUND),
Source::Group(SourceEntry::Input) => self
.group_inputs
.get(index)
.ok_or(INDEX_OUT_OF_BOUND)
.and_then(|actual_index| {
self.resolved_inputs()
.get(*actual_index)
.ok_or(INDEX_OUT_OF_BOUND)
}),
Source::Group(SourceEntry::Output) => self
.group_outputs
.get(index)
.ok_or(INDEX_OUT_OF_BOUND)
.and_then(|actual_index| self.outputs.get(*actual_index).ok_or(INDEX_OUT_OF_BOUND)),
Source::Group(SourceEntry::CellDep) => Err(INDEX_OUT_OF_BOUND),
Source::Group(SourceEntry::HeaderDep) => Err(INDEX_OUT_OF_BOUND),
}
};
let addr = machine.registers()[A0].to_u64();
let size_addr = machine.registers()[A1].clone();
let size = machine.memory_mut().load64(&size_addr)?.to_u64();
let offset = machine.registers()[A2].to_u64();
let mut sc = self
.snapshot2_context
.lock()
.map_err(|e| VMError::Unexpected(e.to_string()))?;

let (wrote_size, full_size) =
match sc.store_bytes(machine, addr, &data_piece_id, offset, size) {
Ok(val) => val,
Err(VMError::External(m)) if m == "INDEX_OUT_OF_BOUND" => {
// This comes from TxData results in an out of bound error, to
// mimic current behavior, we would return INDEX_OUT_OF_BOUND error.
machine.set_register(A0, Mac::REG::from_u8(INDEX_OUT_OF_BOUND));
return Ok(());
}
Err(e) => return Err(e),
};
machine
.memory_mut()
.store64(&size_addr, &Mac::REG::from_u64(full_size))?;
machine.add_cycles_no_checking(transferred_byte_cycles(wrote_size))?;
machine.set_register(A0, Mac::REG::from_u8(SUCCESS));
Ok(())
}

fn load_data_as_code<Mac: SupportMachine>(&self, machine: &mut Mac) -> Result<(), VMError> {
let addr = machine.registers()[A0].to_u64();
let memory_size = machine.registers()[A1].to_u64();
let content_offset = machine.registers()[A2].to_u64();
let content_size = machine.registers()[A3].to_u64();

let index = machine.registers()[A4].to_u64();
let source = Source::parse_from_u64(machine.registers()[A5].to_u64())?;

let cell = self.fetch_cell(source, index as usize);
if let Err(err) = cell {
machine.set_register(A0, Mac::REG::from_u8(err));
return Ok(());
}
let cell = cell.unwrap();

let source = machine.registers()[A5].to_u64();
let data_piece_id = match DataPieceId::try_from((source, index, 0)) {
Ok(id) => id,
Err(_) => {
machine.set_register(A0, Mac::REG::from_u8(INDEX_OUT_OF_BOUND));
return Ok(());
}
};
let mut sc = self
.snapshot2_context
.lock()
.map_err(|e| VMError::Unexpected(e.to_string()))?;
// We are using 0..u64::max_value() to fetch full cell, there is
// also no need to keep the full length value. Since cell's length
// is already full length.
let (cell, _) = match sc
.data_source()
.load_data(&data_piece_id, 0, u64::max_value())
{
Ok(val) => val,
Err(VMError::External(m)) if m == "INDEX_OUT_OF_BOUND" => {
// This comes from TxData results in an out of bound error, to
// mimic current behavior, we would return INDEX_OUT_OF_BOUND error.
machine.set_register(A0, Mac::REG::from_u8(INDEX_OUT_OF_BOUND));
return Ok(());
}
Err(e) => return Err(e),
};
let content_end = content_offset
.checked_add(content_size)
.ok_or(VMError::MemOutOfBound)?;
if content_offset >= cell.data_bytes
|| content_end > cell.data_bytes
if content_offset >= cell.len() as u64
|| content_end > cell.len() as u64
|| content_size > memory_size
{
machine.set_register(A0, Mac::REG::from_u8(SLICE_OUT_OF_BOUND));
return Ok(());
}
let data = self
.data_loader
.load_cell_data(cell)
.ok_or_else(|| {
VMError::Unexpected(format!(
"Unexpected load_cell_data failed {}",
cell.out_point,
))
})?
.slice((content_offset as usize)..(content_end as usize));
machine.memory_mut().init_pages(
addr,
memory_size,
FLAG_EXECUTABLE | FLAG_FREEZED,
Some(data),
Some(cell.slice((content_offset as usize)..(content_end as usize))),
0,
)?;

sc.track_pages(machine, addr, memory_size, &data_piece_id, content_offset)?;
machine.add_cycles_no_checking(transferred_byte_cycles(memory_size))?;
machine.set_register(A0, Mac::REG::from_u8(SUCCESS));
Ok(())
}

fn load_data<Mac: SupportMachine>(&self, machine: &mut Mac) -> Result<(), VMError> {
let index = machine.registers()[A3].to_u64();
let source = Source::parse_from_u64(machine.registers()[A4].to_u64())?;

let cell = self.fetch_cell(source, index as usize);
if let Err(err) = cell {
machine.set_register(A0, Mac::REG::from_u8(err));
return Ok(());
}
let cell = cell.unwrap();
let data = self.data_loader.load_cell_data(cell).ok_or_else(|| {
VMError::Unexpected(format!(
"Unexpected load_cell_data failed {}",
cell.out_point,
))
})?;

let wrote_size = store_data(machine, &data)?;
machine.add_cycles_no_checking(transferred_byte_cycles(wrote_size))?;
machine.set_register(A0, Mac::REG::from_u8(SUCCESS));
Ok(())
}
}

impl<Mac: SupportMachine, DL: CellDataProvider + Send + Sync> Syscalls<Mac> for LoadCellData<DL> {
impl<Mac: SupportMachine, DL> Syscalls<Mac> for LoadCellData<DL>
where
DL: CellDataProvider + HeaderProvider + ExtensionProvider + Send + Sync + Clone + 'static,
{
fn initialize(&mut self, _machine: &mut Mac) -> Result<(), VMError> {
Ok(())
}
Expand Down
Loading

0 comments on commit d4b52f1

Please sign in to comment.