From d7ce52fd52b7c0c08fedf61b7a19da47b81e57e5 Mon Sep 17 00:00:00 2001 From: Izumi Hoshino Date: Thu, 4 Apr 2024 22:02:38 +0900 Subject: [PATCH] Updated code according to review --- src/lib.rs | 3 --- src/serde/ser_br.rs | 44 ++++++++++++++++++++++++++++++++++---- {src => wasm/src}/flags.rs | 0 wasm/src/lazy_node.rs | 10 ++++----- wasm/src/lib.rs | 1 + wasm/src/run_program.rs | 2 +- wasm/src/serialize.rs | 19 +++++++--------- 7 files changed, 55 insertions(+), 24 deletions(-) rename {src => wasm/src}/flags.rs (100%) diff --git a/src/lib.rs b/src/lib.rs index 25862f95..f233c48f 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -6,7 +6,6 @@ pub mod cost; pub mod dialect; pub mod err_utils; pub mod f_table; -pub mod flags; pub mod more_ops; pub mod number; pub mod op_utils; @@ -26,8 +25,6 @@ pub use chia_dialect::{ ENABLE_BLS_OPS_OUTSIDE_GUARD, ENABLE_FIXED_DIV, LIMIT_HEAP, MEMPOOL_MODE, NO_UNKNOWN_OPS, }; -pub use flags::ALLOW_BACKREFS; - #[cfg(feature = "counters")] pub use run_program::run_program_with_counters; diff --git a/src/serde/ser_br.rs b/src/serde/ser_br.rs index 22e1ea8e..24d7ab7c 100644 --- a/src/serde/ser_br.rs +++ b/src/serde/ser_br.rs @@ -7,6 +7,7 @@ use super::object_cache::{serialized_length, treehash, ObjectCache}; use super::read_cache_lookup::ReadCacheLookup; use super::write_atom::write_atom; use crate::allocator::{Allocator, NodePtr, SExp}; +use crate::serde::ser::LimitedWriter; const BACK_REFERENCE: u8 = 0xfe; const CONS_BOX_MARKER: u8 = 0xff; @@ -77,10 +78,45 @@ pub fn node_to_stream_backrefs( Ok(()) } +pub fn node_to_bytes_backrefs_limit( + a: &Allocator, + node: NodePtr, + limit: usize, +) -> io::Result> { + let buffer = Cursor::new(Vec::new()); + let mut writer = LimitedWriter::new(buffer, limit); + node_to_stream_backrefs(a, node, &mut writer)?; + let vec = writer.into_inner().into_inner(); + Ok(vec) +} + pub fn node_to_bytes_backrefs(a: &Allocator, node: NodePtr) -> io::Result> { - let mut buffer = Cursor::new(Vec::new()); + node_to_bytes_backrefs_limit(a, node, 2000000) +} - node_to_stream_backrefs(a, node, &mut buffer)?; - let vec = buffer.into_inner(); - Ok(vec) +#[test] +fn test_serialize_limit() { + let mut a = Allocator::new(); + + let leaf = a.new_atom(&[1, 2, 3, 4, 5]).unwrap(); + let l1 = a.new_pair(leaf, leaf).unwrap(); + let l2 = a.new_pair(l1, l1).unwrap(); + let l3 = a.new_pair(l2, l2).unwrap(); + + let expected = &[255, 255, 255, 133, 1, 2, 3, 4, 5, 254, 2, 254, 2, 254, 2]; + + { + assert_eq!(node_to_bytes_backrefs(&a, l3).unwrap(), expected); + } + + { + assert_eq!(node_to_bytes_backrefs_limit(&a, l3, 15).unwrap(), expected); + } + + { + assert_eq!( + node_to_bytes_backrefs_limit(&a, l3, 14).unwrap_err().kind(), + io::ErrorKind::OutOfMemory + ); + } } diff --git a/src/flags.rs b/wasm/src/flags.rs similarity index 100% rename from src/flags.rs rename to wasm/src/flags.rs diff --git a/wasm/src/lazy_node.rs b/wasm/src/lazy_node.rs index 36f33828..5c50968e 100644 --- a/wasm/src/lazy_node.rs +++ b/wasm/src/lazy_node.rs @@ -1,10 +1,10 @@ -use clvmr::allocator::{Allocator, NodePtr, SExp}; +use js_sys::Array; use std::rc::Rc; +use wasm_bindgen::prelude::*; +use crate::flags::ALLOW_BACKREFS; +use clvmr::allocator::{Allocator, NodePtr, SExp}; use clvmr::serde::{node_to_bytes, node_to_bytes_backrefs}; -use clvmr::ALLOW_BACKREFS; -use js_sys::Array; -use wasm_bindgen::prelude::*; #[wasm_bindgen] #[derive(Clone)] @@ -45,7 +45,7 @@ impl LazyNode { } else { node_to_bytes }; -serializer(&self.allocator, self.node).ok() + serializer(&self.allocator, self.node).ok() } } diff --git a/wasm/src/lib.rs b/wasm/src/lib.rs index 7a54e2ae..2a1a5143 100644 --- a/wasm/src/lib.rs +++ b/wasm/src/lib.rs @@ -1,3 +1,4 @@ +pub mod flags; pub mod lazy_node; pub mod run_program; pub mod serialize; diff --git a/wasm/src/run_program.rs b/wasm/src/run_program.rs index 786ca527..0863abf4 100644 --- a/wasm/src/run_program.rs +++ b/wasm/src/run_program.rs @@ -2,6 +2,7 @@ use js_sys::Array; use std::rc::Rc; use wasm_bindgen::prelude::*; +use crate::flags::ALLOW_BACKREFS; use crate::lazy_node::LazyNode; use clvmr::allocator::Allocator; use clvmr::chia_dialect::ChiaDialect; @@ -9,7 +10,6 @@ use clvmr::chia_dialect::NO_UNKNOWN_OPS as _no_unknown_ops; use clvmr::cost::Cost; use clvmr::run_program::run_program; use clvmr::serde::{node_from_bytes, node_from_bytes_backrefs, node_to_bytes}; -use clvmr::ALLOW_BACKREFS; // When the `wee_alloc` feature is enabled, use `wee_alloc` as the global // allocator. diff --git a/wasm/src/serialize.rs b/wasm/src/serialize.rs index cf0285ef..f49ab01e 100644 --- a/wasm/src/serialize.rs +++ b/wasm/src/serialize.rs @@ -1,15 +1,14 @@ -use crate::lazy_node::LazyNode; -use clvmr::serde::{node_from_bytes, node_from_bytes_backrefs, serialized_length_from_bytes}; -use clvmr::{Allocator, ALLOW_BACKREFS}; use std::rc::Rc; use wasm_bindgen::prelude::wasm_bindgen; +use crate::flags::ALLOW_BACKREFS; +use crate::lazy_node::LazyNode; +use clvmr::serde::{node_from_bytes, node_from_bytes_backrefs, serialized_length_from_bytes}; +use clvmr::Allocator; + #[wasm_bindgen] pub fn serialized_length(program: &[u8]) -> Result { - match serialized_length_from_bytes(program) { - Ok(length) => Ok(length), - Err(err) => Err(err.to_string()), - } + serialized_length_from_bytes(program).map_err(|x| x.to_string()) } #[wasm_bindgen] @@ -20,8 +19,6 @@ pub fn sexp_from_bytes(b: &[u8], flag: u32) -> Result { } else { node_from_bytes }; - match deserializer(&mut allocator, b) { - Ok(node) => Ok(LazyNode::new(Rc::new(allocator), node)), - Err(err) => Err(err.to_string()), - } + let node = deserializer(&mut allocator, b).map_err(|e| e.to_string())?; + Ok(LazyNode::new(Rc::new(allocator), node)) }