docs(zkvms/zkwasm): Add detailed documentation comments

3 files changed, 65 insertions, 1 deletions

diff --git a/zkvms/zkwasm/guest/Cargo.toml b/zkvms/zkwasm/guest/Cargo.toml
index 79a9216..5adaea6 100644
--- a/zkvms/zkwasm/guest/Cargo.toml
+++ b/zkvms/zkwasm/guest/Cargo.toml
@@ -1,5 +1,6 @@
 [package]
 name = "guest"
+description = "Jolt's specific guest crate, which includes the chosen guest in guests"
 version = "0.1.0"
 edition = "2021"
 
@@ -14,3 +15,6 @@ zkwasm-rust-sdk = { git = "https://github.com/DelphinusLab/zkWasm-rust.git" }
 wasm-bindgen = "=0.2.95"
 
 wrapper_macro = { version = "0.1.0", path = "../wrapper_macro" }
+
+# The zkp dependency references a chosen guest in guests. It is included
+# (inserted here) by Nix. See zkvmLib.nix
diff --git a/zkvms/zkwasm/host/src/main.rs b/zkvms/zkwasm/host/src/main.rs
index e480ef5..7ada9e0 100644
--- a/zkvms/zkwasm/host/src/main.rs
+++ b/zkvms/zkwasm/host/src/main.rs
@@ -5,6 +5,12 @@ use std::process::{Command, Stdio};
 static PUBLIC_INPUT_PATH: &str = "public_input.bin";
 static PRIVATE_INPUT_PATH: &str = "private_input.bin";
 
+/// Inserts array sizes before every square bracket
+///
+/// # Example
+///
+/// If `flat` is "[[0,1], [2,3,4], []]"
+/// Output will be "3[2[0,1], 3[2,3,4], 0[]]"
 fn get_with_sizes(flat: &str) -> String {
     let mut values = flat
         .split('[')
@@ -47,11 +53,21 @@ fn get_with_sizes(flat: &str) -> String {
     }
 }
 
+/// Creates an anonymous function which takes `run_info`, "serializes" the
+/// specified input, outputs it into a file and returns a "path:<PATH>"
+/// argument, ready to be passed to zkWasm.
+///
+/// The macro takes three arguments: run_info input expression, path for file
+/// output and the name of a foreach macro.
+///
+/// For collection types, first the size is emitted and afterwards its actual
+/// values.
 macro_rules! build_input {
     ($input:expr , $path:ident , $type:ident) => {
         |run_info: &RunWith| {
             let mut ret: Vec<u64> = Vec::new();
             $type!{
+                // Simplify input string
                 let flat = format!("{:?}", $input.yield)
                     .replace("false", "0")
                     .replace("true",  "1")
diff --git a/zkvms/zkwasm/wrapper_macro/src/lib.rs b/zkvms/zkwasm/wrapper_macro/src/lib.rs
index 301ba60..8bdb581 100644
--- a/zkvms/zkwasm/wrapper_macro/src/lib.rs
+++ b/zkvms/zkwasm/wrapper_macro/src/lib.rs
@@ -5,6 +5,12 @@ mod parse_fn;
 use crate::parse_fn::{ split_fn, args_divide_grouped, args_divide_public, group_streams };
 use toml::Table;
 
+/// Extends an out TokenStream with `let` directives for all patterns (and
+/// types). readfn specifies wether the input is public or private.
+///
+/// Each `let` binding calls the read! macro (defined in zkWasm wrapper_macro
+/// crate), feeding it an autogenerated "type note", which is then used to
+/// "deserialize" the input.
 fn insert_reads(out: &mut TokenStream, patterns: &Vec<TokenStream>, types: &Vec<TokenStream>, readfn: &str) {
     for i in 0..patterns.len() {
         let type_note: String = format!("{}", types[i])
@@ -14,6 +20,32 @@ fn insert_reads(out: &mut TokenStream, patterns: &Vec<TokenStream>, types: &Vec<
     }
 }
 
+/// Creates a body, which reads all inputs, stores them in variables, then
+/// executes the entrypoint function with those arguments and writes the
+/// result.
+///
+/// Inputs are read via the read! macro (defined in the zkWasm wrapper_macro
+/// crate). Their public status is dependent on `default_public_input.toml`.
+///
+/// # Usage
+///
+/// Inside zkWasm's guest (excluding the `entrypoint_expr` call):
+///
+/// ```rust
+/// make_wrapper!{fn main(...) -> ...}
+/// ```
+///
+/// # Example output
+///
+/// ```rust
+/// {
+///     let ... : ... = read!(... ...);
+///     let ... : ... = read!(... ...);
+///     ...
+///     let result = zkp::main(..., ..., ...);
+///     write(result as u64);
+/// }
+/// ```
 #[proc_macro]
 pub fn make_wrapper(item: TokenStream) -> TokenStream {
     let (name, args, ret) = split_fn(&item);
@@ -33,7 +65,6 @@ pub fn make_wrapper(item: TokenStream) -> TokenStream {
 
     out.extend(format!("
         let result = zkp::{}{};
-        assert(result);
         write(result as u64);
     ", name, ts_patterns).parse::<TokenStream>());
 
@@ -129,6 +160,19 @@ fn return_tuple(readfn: &TokenTree, inner: &TokenStream) -> TokenStream {
     ").parse().unwrap()
 }
 
+/// Creates a body which returns a value of the type, defined as a "type note"
+/// argument.
+///
+/// The host "serializes" all input data by flattening it as a series of
+/// integers. This function, in turn, unflattens the input, by reading integers
+/// multiple times and combining them in the appropriate structures.
+///
+/// It takes two arguments, separated by a space. The first is a string "type
+/// note" and the second is the name of the read function (either read_private
+/// or read_public).
+///
+/// The type note is similar to a print!("{:?}") output, however angled braces
+/// are square.
 #[proc_macro]
 pub fn read(item: TokenStream) -> TokenStream {
     let mut parts = item.clone().into_iter();