Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: remove rkyv serialize #288

Closed
wants to merge 1 commit into from
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

62 changes: 8 additions & 54 deletions utils/client/src/oracle/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -11,10 +11,7 @@ use kona_preimage::{
};
use kona_proof::FlushableCache;
use kzg_rs::{get_kzg_settings, Blob as KzgRsBlob, Bytes48};
use rkyv::{
with::{ArchiveWith, DeserializeWith, SerializeWith},
Archive, Archived, Deserialize, Fallible, Infallible, Resolver, Serialize,
};
use rkyv::rancor::Error;
use sha2::{Digest, Sha256};
use spin::mutex::Mutex;
use std::{collections::HashMap, sync::Arc};
Expand All @@ -25,9 +22,8 @@ type LockableMap = Arc<Mutex<HashMap<[u8; 32], Vec<u8>, BytesHasherBuilder>>>;
/// Rather than relying on a trusted host for data, the data in this oracle
/// is verified with the `verify()` function, and then is trusted for
/// the remainder of execution.
#[derive(Debug, Clone, Archive, Serialize, Deserialize)]
#[derive(Debug, Clone)]
pub struct InMemoryOracle {
#[with(Lock)]
cache: LockableMap,
}

Expand All @@ -36,13 +32,14 @@ impl InMemoryOracle {
/// These values are deserialized using rkyv for zero copy deserialization.
pub fn from_raw_bytes(input: Vec<u8>) -> Self {
println!("cycle-tracker-start: in-memory-oracle-from-raw-bytes-archive");
let archived = unsafe { rkyv::archived_root::<Self>(&input) };
let deserialized: HashMap<[u8; 32], Vec<u8>, BytesHasherBuilder> =
rkyv::from_bytes::<HashMap<[u8; 32], Vec<u8>, BytesHasherBuilder>, Error>(&input)
.expect("Failed to deserialize oracle bytes");
println!("cycle-tracker-end: in-memory-oracle-from-raw-bytes-archive");
println!("cycle-tracker-start: in-memory-oracle-from-raw-bytes-deserialize");
let deserialized: Self = archived.deserialize(&mut Infallible).unwrap();
println!("cycle-tracker-end: in-memory-oracle-from-raw-bytes-deserialize");

deserialized
let cache = Arc::new(Mutex::new(deserialized));

Self { cache }
}

/// Creates a new [InMemoryOracle] from a HashMap of B256 keys and Vec<u8> values.
Expand Down Expand Up @@ -199,46 +196,3 @@ impl InMemoryOracle {
Ok(())
}
}

struct Lock;

impl<F: Archive> ArchiveWith<Arc<Mutex<F>>> for Lock {
type Archived = Archived<F>;
type Resolver = Resolver<F>;

unsafe fn resolve_with(
field: &Arc<Mutex<F>>,
pos: usize,
resolver: Self::Resolver,
out: *mut Self::Archived,
) {
field.lock().resolve(pos, resolver, out.cast());
}
}

impl<S, T> SerializeWith<Arc<Mutex<T>>, S> for Lock
where
T: Serialize<S>,
S: Fallible + ?Sized,
{
fn serialize_with(
field: &Arc<Mutex<T>>,
serializer: &mut S,
) -> Result<Self::Resolver, <S as Fallible>::Error> {
field.lock().serialize(serializer)
}
}

impl<D, T> DeserializeWith<Archived<T>, Arc<Mutex<T>>, D> for Lock
where
Archived<T>: Deserialize<T, D>,
T: Archive,
D: Fallible + ?Sized,
{
fn deserialize_with(
field: &Archived<T>,
deserializer: &mut D,
) -> Result<Arc<Mutex<T>>, D::Error> {
Ok(Arc::new(Mutex::new(field.deserialize(deserializer)?)))
}
}
2 changes: 1 addition & 1 deletion utils/host/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,7 @@ kona-host.workspace = true

# general
rkyv.workspace = true
bincode.workspace = true
serde_json.workspace = true
anyhow.workspace = true
cargo_metadata.workspace = true
Expand All @@ -44,4 +45,3 @@ serde.workspace = true
reqwest.workspace = true
log.workspace = true
sysinfo = "0.32.0"

23 changes: 3 additions & 20 deletions utils/host/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -21,13 +21,7 @@ use std::{fs::File, io::Read};

use anyhow::Result;

use rkyv::{
ser::{
serializers::{AlignedSerializer, CompositeSerializer, HeapScratch, SharedSerializeMap},
Serializer,
},
AlignedVec,
};
use rkyv::{api::high::HighSerializer, rancor::Error, ser::Serializer, to_bytes, util::AlignedVec};

sol! {
#[allow(missing_docs)]
Expand Down Expand Up @@ -88,19 +82,8 @@ pub fn get_proof_stdin(host_cli: &HostCli) -> Result<SP1Stdin> {
anyhow::anyhow!("Failed to convert DiskKeyValueStore to MemoryKeyValueStore")
})?;

let mut serializer = CompositeSerializer::new(
AlignedSerializer::new(AlignedVec::new()),
// Note: This value corresponds to the size of the heap needed to serialize the KV store.
// Increase this value if we start running into serialization issues.
HeapScratch::<268435456>::new(),
SharedSerializeMap::new(),
);
// Serialize the underlying KV store.
serializer.serialize_value(&InMemoryOracle::from_b256_hashmap(mem_kv_store.store))?;

let buffer = serializer.into_serializer().into_inner();
let kv_store_bytes = buffer.into_vec();
stdin.write_slice(&kv_store_bytes);
let bytes = bincode::serialize(&mem_kv_store.store).unwrap();
stdin.write_slice(&bytes);

Ok(stdin)
}
Expand Down
Loading