Skip to content

Commit

Permalink
fix: solve deep recursion issues on some ZDag functions
Browse files Browse the repository at this point in the history
  • Loading branch information
arthurpaulino committed Mar 6, 2024
1 parent fe52127 commit c72f63d
Show file tree
Hide file tree
Showing 2 changed files with 93 additions and 10 deletions.
94 changes: 90 additions & 4 deletions src/cli/zstore.rs
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
use anyhow::{bail, Result};
use indexmap::IndexSet;
use serde::{Deserialize, Serialize};
use std::collections::{BTreeMap, HashMap, HashSet};

Expand Down Expand Up @@ -37,7 +38,7 @@ impl<F: LurkField> ZDag<F> {
store: &Store<F>,
cache: &mut HashMap<Ptr, ZPtr<F>>,
) -> ZPtr<F> {
let mut recurse = |ptr: &Ptr| -> ZPtr<F> {
let mut recurse = |ptr: &Ptr, cache: &mut HashMap<_, _>| -> ZPtr<F> {
if let Some(z_ptr) = cache.get(ptr) {
*z_ptr
} else {
Expand Down Expand Up @@ -135,21 +136,68 @@ impl<F: LurkField> ZDag<F> {
z_ptr
}
};
recurse(ptr)
let mut dag = IndexSet::from([*ptr]);
let mut stack = vec![*ptr];
macro_rules! feed_loop {
($x:expr) => {
if $x.raw().is_hash() {
if !cache.contains_key(&$x) {
if dag.insert($x) {
stack.push($x);
}
}
}
};
}
while let Some(ptr) = stack.pop() {
match ptr.raw() {
RawPtr::Atom(..) => (),
RawPtr::Hash4(idx) => {
for ptr in expect_ptrs!(store, 2, *idx) {
feed_loop!(ptr)
}
}
RawPtr::Hash6(idx) => {
for ptr in expect_ptrs!(store, 3, *idx) {
feed_loop!(ptr)
}
}
RawPtr::Hash8(idx) => {
for ptr in expect_ptrs!(store, 4, *idx) {
feed_loop!(ptr)
}
}
}
}
dag.iter()
.rev()
.map(|ptr| recurse(ptr, cache))
.last()
.expect("dag doesn't start empty")
}

#[inline]
fn get_type(&self, z_ptr: &ZPtr<F>) -> Option<&ZPtrType<F>> {
self.0.get(z_ptr)
}

fn get_children(&self, z_ptr: &ZPtr<F>) -> Result<Vec<&ZPtr<F>>> {
match self.get_type(z_ptr) {
None => bail!("Couldn't find ZPtr on ZStore"),
Some(ZPtrType::Atom) => Ok(vec![]),
Some(ZPtrType::Tuple2(z1, z2)) => Ok(vec![z1, z2]),
Some(ZPtrType::Tuple3(z1, z2, z3) | ZPtrType::Env(z1, z2, z3)) => Ok(vec![z1, z2, z3]),
Some(ZPtrType::Tuple4(z1, z2, z3, z4)) => Ok(vec![z1, z2, z3, z4]),
}
}

pub(crate) fn populate_store(
&self,
z_ptr: &ZPtr<F>,
store: &Store<F>,
cache: &mut HashMap<ZPtr<F>, Ptr>,
) -> Result<Ptr> {
let mut recurse = |z_ptr: &ZPtr<F>| -> Result<Ptr> {
let recurse = |z_ptr: &ZPtr<F>, cache: &mut HashMap<_, _>| -> Result<Ptr> {
if let Some(z_ptr) = cache.get(z_ptr) {
Ok(*z_ptr)
} else {
Expand Down Expand Up @@ -191,10 +239,34 @@ impl<F: LurkField> ZDag<F> {
Ok(ptr)
}
};
recurse(z_ptr)
let mut dag = IndexSet::from([z_ptr]);
let mut stack = vec![z_ptr];
macro_rules! feed_loop {
($x:expr) => {
if !cache.contains_key($x) {
if dag.insert($x) {
stack.push($x);
}
}
};
}
while let Some(z_ptr) = stack.pop() {
for z_ptr in self.get_children(z_ptr)? {
feed_loop!(z_ptr)
}
}
dag.iter()
.rev()
.map(|z_ptr| recurse(z_ptr, cache))
.last()
.expect("dag doesn't start empty")
}

/// Populates a `ZDag` with data from self
///
/// # Warning
/// This function is recursive and might reach Rust's recursion depth limit
// TODO: solve this issue if we actually use this function
#[allow(dead_code)]
pub(crate) fn populate_z_dag(
&self,
Expand Down Expand Up @@ -425,4 +497,18 @@ mod tests {
// but not in `z_dag_new`
assert!(z_dag_new.get_type(&z_two_thr).is_none());
}

#[test]
fn test_deep_lurk_data_roundtrip() {
let string = String::from_utf8(vec![b'0'; 8192]).unwrap();
let store = Store::<Bn>::default();
let ptr = store.intern_string(&string);
let mut z_dag = ZDag::default();
let z_ptr = z_dag.populate_with(&ptr, &store, &mut HashMap::default());

let store = Store::<Bn>::default();
z_dag
.populate_store(&z_ptr, &store, &mut HashMap::default())
.unwrap();
}
}
9 changes: 3 additions & 6 deletions src/lem/store.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1076,20 +1076,17 @@ impl<F: LurkField> Store<F> {
match ptr {
RawPtr::Atom(..) => (),
RawPtr::Hash4(idx) => {
let ptrs = self.expect_raw_ptrs::<4>(*idx);
for ptr in ptrs {
for ptr in self.expect_raw_ptrs::<4>(*idx) {
feed_loop!(ptr)
}
}
RawPtr::Hash6(idx) => {
let ptrs = self.expect_raw_ptrs::<6>(*idx);
for ptr in ptrs {
for ptr in self.expect_raw_ptrs::<6>(*idx) {
feed_loop!(ptr)
}
}
RawPtr::Hash8(idx) => {
let ptrs = self.expect_raw_ptrs::<8>(*idx);
for ptr in ptrs {
for ptr in self.expect_raw_ptrs::<8>(*idx) {
feed_loop!(ptr)
}
}
Expand Down

0 comments on commit c72f63d

Please sign in to comment.