Skip to content

Commit

Permalink
feat(cell): remove cell hashes indirection
Browse files Browse the repository at this point in the history
  • Loading branch information
@Rexagon
Rexagon committed Nov 21, 2024
1 parent 92d5a8b commit 724cbc8
Show file tree
Hide file tree
Showing 6 changed files with 267 additions and 167 deletions.
14 changes: 6 additions & 8 deletions src/cell/cell_context.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -78,7 +78,7 @@ pub struct CellParts<'a> {

impl<'a> CellParts<'a> {
/// Validates cell and computes all hashes.
pub fn compute_hashes(&self) -> Result<Vec<(HashBytes, u16)>, Error> {
pub fn compute_hashes(&self) -> Result<ArrayVec<(HashBytes, u16), 4>, Error> {
const HASH_BITS: usize = 256;
const DEPTH_BITS: usize = 16;

Expand All @@ -89,9 +89,6 @@ impl<'a> CellParts<'a> {

let references = self.references.as_ref();

// `hashes_len` is guaranteed to be in range 1..4
let mut hashes_len = level + 1;

let (cell_type, computed_level_mask) = if unlikely(descriptor.is_exotic()) {
let Some(&first_byte) = self.data.first() else {
return Err(Error::InvalidCell);
Expand All @@ -114,7 +111,6 @@ impl<'a> CellParts<'a> {
return Err(Error::InvalidCell);
}

hashes_len = 1;
(CellType::PrunedBranch, level_mask)
}
// 8 bits type, hash, depth
Expand Down Expand Up @@ -157,7 +153,7 @@ impl<'a> CellParts<'a> {
let level_offset = cell_type.is_merkle() as u8;
let is_pruned = cell_type.is_pruned_branch();

let mut hashes = Vec::<(HashBytes, u16)>::with_capacity(hashes_len);
let mut hashes = ArrayVec::<(HashBytes, u16), 4>::new();
for level in 0..4 {
// Skip non-zero levels for pruned branches and insignificant hashes for other cells
if level != 0 && (is_pruned || !level_mask.contains(level)) {
Expand All @@ -181,7 +177,7 @@ impl<'a> CellParts<'a> {
} else {
// SAFETY: new hash is added on each iteration, so there will
// definitely be a hash, when level>0
let prev_hash = unsafe { hashes.last().unwrap_unchecked() };
let prev_hash = unsafe { hashes.last_unchecked() };
hasher.update(prev_hash.0.as_slice());
}

Expand All @@ -203,7 +199,9 @@ impl<'a> CellParts<'a> {
}

let hash = hasher.finalize().into();
hashes.push((hash, depth));

// SAFETY: iteration count is in range 0..4
unsafe { hashes.push((hash, depth)) };
}

Ok(hashes)
Expand Down
132 changes: 111 additions & 21 deletions src/cell/cell_impl/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@ use super::{
Cell, CellDescriptor, CellFamily, CellImpl, CellInner, DynCell, HashBytes, EMPTY_CELL_HASH,
MAX_REF_COUNT,
};
use crate::util::TryAsMut;
use crate::util::{ArrayVec, TryAsMut};

macro_rules! define_gen_vtable_ptr {
(($($param:tt)*) => $($type:tt)*) => {
Expand Down Expand Up @@ -223,25 +223,25 @@ const ALL_ONES_CELL_HASH: [u8; 32] = [
0x66, 0x12, 0x81, 0x70, 0x30, 0x1a, 0x7b, 0xec, 0xc2, 0x7a, 0xf1, 0xad, 0xbe, 0x6a, 0x31, 0xc9,
];

type OrdinaryCell<const N: usize> = HeaderWithData<OrdinaryCellHeader, N>;

struct OrdinaryCellHeader {
bit_len: u16,
#[cfg(feature = "stats")]
stats: CellTreeStats,
hashes: Vec<(HashBytes, u16)>,
descriptor: CellDescriptor,
references: [MaybeUninit<Cell>; MAX_REF_COUNT],
without_first: bool,
}

impl OrdinaryCellHeader {
fn level_descr(&self, level: u8) -> &(HashBytes, u16) {
let hash_index = hash_index(self.descriptor, level);
debug_assert!((hash_index as usize) < self.hashes.len());

// SAFETY: hash index is in range 0..=3
unsafe { self.hashes.get_unchecked(hash_index as usize) }
#[inline(always)]
unsafe fn into_full<const H: usize>(
self,
hashes: ArrayVec<(HashBytes, u16), 4>,
) -> FullOrdinaryCellHeader<H> {
FullOrdinaryCellHeader {
base: self,
hashes: hashes.into_plain::<H>(),
}
}

fn reference(&self, i: u8) -> Option<&Cell> {
Expand Down Expand Up @@ -387,39 +387,56 @@ impl Drop for OrdinaryCellHeader {
}
}

struct FullOrdinaryCellHeader<const H: usize> {
base: OrdinaryCellHeader,
hashes: [(HashBytes, u16); H],
}

impl<const H: usize> FullOrdinaryCellHeader<H> {
fn level_descr(&self, level: u8) -> &(HashBytes, u16) {
let hash_index = hash_index(self.base.descriptor, level);
debug_assert!((hash_index as usize) < H);

// SAFETY: hash index is in range 0..=3
unsafe { self.hashes.get_unchecked(hash_index as usize) }
}
}

type OrdinaryCell<const H: usize, const N: usize> = HeaderWithData<FullOrdinaryCellHeader<H>, N>;

// TODO: merge VTables for different data array sizes

impl<const N: usize> CellImpl for OrdinaryCell<N> {
impl<const H: usize, const N: usize> CellImpl for OrdinaryCell<H, N> {
#[inline]
fn untrack(self: CellInner<Self>) -> Cell {
Cell(self)
}

fn descriptor(&self) -> CellDescriptor {
self.header.descriptor
self.header.base.descriptor
}

fn data(&self) -> &[u8] {
let data_ptr = std::ptr::addr_of!(self.data) as *const u8;
let data_len = self.header.descriptor.byte_len() as usize;
let data_len = self.header.base.descriptor.byte_len() as usize;
// SAFETY: header is initialized
unsafe { std::slice::from_raw_parts(data_ptr, data_len) }
}

fn bit_len(&self) -> u16 {
self.header.bit_len
self.header.base.bit_len
}

fn reference(&self, index: u8) -> Option<&DynCell> {
Some(self.header.reference(index)?.as_ref())
Some(self.header.base.reference(index)?.as_ref())
}

fn reference_cloned(&self, index: u8) -> Option<Cell> {
Some(self.header.reference(index)?.clone())
Some(self.header.base.reference(index)?.clone())
}

fn virtualize(&self) -> &DynCell {
if self.header.descriptor.level_mask().is_empty() {
if self.header.base.descriptor.level_mask().is_empty() {
self
} else {
VirtualCellWrapper::wrap(self)
Expand All @@ -435,23 +452,78 @@ impl<const N: usize> CellImpl for OrdinaryCell<N> {
}

fn take_first_child(&mut self) -> Option<Cell> {
self.header.take_first_child()
self.header.base.take_first_child()
}

fn replace_first_child(&mut self, parent: Cell) -> ReplacedChild {
self.header.replace_first_child_with_parent(parent)
self.header.base.replace_first_child_with_parent(parent)
}

fn take_next_child(&mut self) -> Option<Cell> {
self.header.take_next_child()
self.header.base.take_next_child()
}

#[cfg(feature = "stats")]
fn stats(&self) -> CellTreeStats {
self.header.stats
self.header.base.stats
}
}

mod ordinary {
use super::*;

define_gen_vtable_ptr!((const H: usize, const N: usize) => OrdinaryCell<H, N>);

pub const VTABLES_GROUP_LEN: usize = 9;

pub const VTABLES: [[*const (); VTABLES_GROUP_LEN]; 4] = [
[
gen_vtable_ptr::<1, 0>(),
gen_vtable_ptr::<1, 8>(), // 1, aligned to 8
gen_vtable_ptr::<1, 8>(), // 2, aligned to 8
gen_vtable_ptr::<1, 8>(), // 4, aligned to 8
gen_vtable_ptr::<1, 8>(),
gen_vtable_ptr::<1, 16>(),
gen_vtable_ptr::<1, 32>(),
gen_vtable_ptr::<1, 64>(),
gen_vtable_ptr::<1, 128>(),
],
[
gen_vtable_ptr::<2, 0>(),
gen_vtable_ptr::<2, 8>(), // 1, aligned to 8
gen_vtable_ptr::<2, 8>(), // 2, aligned to 8
gen_vtable_ptr::<2, 8>(), // 4, aligned to 8
gen_vtable_ptr::<2, 8>(),
gen_vtable_ptr::<2, 16>(),
gen_vtable_ptr::<2, 32>(),
gen_vtable_ptr::<2, 64>(),
gen_vtable_ptr::<2, 128>(),
],
[
gen_vtable_ptr::<3, 0>(),
gen_vtable_ptr::<3, 8>(), // 1, aligned to 8
gen_vtable_ptr::<3, 8>(), // 2, aligned to 8
gen_vtable_ptr::<3, 8>(), // 4, aligned to 8
gen_vtable_ptr::<3, 8>(),
gen_vtable_ptr::<3, 16>(),
gen_vtable_ptr::<3, 32>(),
gen_vtable_ptr::<3, 64>(),
gen_vtable_ptr::<3, 128>(),
],
[
gen_vtable_ptr::<4, 0>(),
gen_vtable_ptr::<4, 8>(), // 1, aligned to 8
gen_vtable_ptr::<4, 8>(), // 2, aligned to 8
gen_vtable_ptr::<4, 8>(), // 4, aligned to 8
gen_vtable_ptr::<4, 8>(),
gen_vtable_ptr::<4, 16>(),
gen_vtable_ptr::<4, 32>(),
gen_vtable_ptr::<4, 64>(),
gen_vtable_ptr::<4, 128>(),
],
];
}

struct LibraryReference {
repr_hash: HashBytes,
descriptor: CellDescriptor,
Expand Down Expand Up @@ -617,6 +689,24 @@ impl<const N: usize> CellImpl for PrunedBranch<N> {
}
}

mod pruned {
use super::*;

define_gen_vtable_ptr!((const N: usize) => PrunedBranch<N>);

pub const LENGTHS: [usize; 3] = [
PrunedBranchHeader::cell_data_len(1),
PrunedBranchHeader::cell_data_len(2),
PrunedBranchHeader::cell_data_len(3),
];

pub const VTABLES: [*const (); 3] = [
gen_vtable_ptr::<{ LENGTHS[0] }>(),
gen_vtable_ptr::<{ LENGTHS[1] }>(),
gen_vtable_ptr::<{ LENGTHS[2] }>(),
];
}

#[repr(transparent)]
pub struct VirtualCell<T>(T);

Expand Down
Loading

0 comments on commit 724cbc8

Please sign in to comment.