sp1-program

main.rs

//! A simple program to be proven inside the zkVM.

#![no_main]
sp1_zkvm::entrypoint!(main);

use wasmi::{Engine, Linker, Module, Store, Caller};
mod zkwasi;

pub fn main() {
    let wasm = sp1_zkvm::io::read::<Vec<u8>>();
    let input = sp1_zkvm::io::read::<String>();
    let input_clone = input.clone();

    let engine = Engine::default();
    println!("cycle-tracker-start: instantiate wasm");
    let module = Module::new(&engine, &mut &wasm[..]).expect("Failed to create module");
    println!("cycle-tracker-end: instantiate wasm");

    let mut linker = <Linker<String>>::new(&engine);
    let mut store = Store::new(&engine, input.clone());

    zkwasi::add_to_linker(&mut linker).expect("Failed to add zkWASI to linker");
    linker.func_wrap(
        "wasi_snapshot_preview1",
        "fd_read",
        move |mut _caller: Caller<'_, String>, fd: i32, iovs: i32, iovs_len: i32, nread: i32| {

            let mem = &_caller.get_export("memory").unwrap().into_memory().unwrap();
            let data = _caller.data();
            if data.is_empty() {
                mem.write(&mut _caller, nread as usize, &(0 as i32).to_le_bytes()).unwrap();
                return Ok(0);
            }


            let mut iovec_pair = [0u8; 8];
            mem.read(&mut _caller, iovs as usize, &mut iovec_pair).unwrap();

            let read_offset = i32::from_le_bytes(iovec_pair[..4].try_into().unwrap());
            let read_len = i32::from_le_bytes(iovec_pair[4..8].try_into().unwrap());


            let stdin_data = input_clone.as_bytes().to_vec();
            mem.write(&mut _caller, read_offset as usize, &stdin_data).unwrap();
            let mut buf = [0u8; 4];
            mem.read(&mut _caller, read_offset as usize, &mut buf).unwrap();

            mem.write(&mut _caller, nread as usize, &(stdin_data.len() as i32).to_le_bytes()).unwrap();

            let mut buf = [0u8; 4];
            mem.read(&mut _caller, nread as usize, &mut buf).unwrap();

            // Make data empty
            *_caller.data_mut() = "".to_string();

            Ok(0)
        },
    ).unwrap();

    linker.func_wrap(
        "wasi_snapshot_preview1",
        "fd_write",
        |mut _caller: Caller<'_, String>, fd: i32, iovs: i32, iovs_len: i32, nwrite: i32| {

            let mem = &_caller.get_export("memory").unwrap().into_memory().unwrap();

            let mut stdout_data_offset = [0u8; 4];
            mem.read(&mut _caller, iovs as usize, &mut stdout_data_offset).unwrap();
            let stdout_data_offset = i32::from_le_bytes(stdout_data_offset.try_into().unwrap());
            let mut stdout_data_len_buf = [0u8; 4];
            mem.read(&mut _caller, iovs as usize + 4, &mut stdout_data_len_buf).unwrap();

            let stdout_data_len: i32 = i32::from_le_bytes(stdout_data_len_buf);

            let mut stdout_data = vec![0u8; stdout_data_len as usize];
            {
                let borrowed_mem = &mut _caller.get_export("memory").unwrap().into_memory().unwrap();
                borrowed_mem.read(&mut _caller, stdout_data_offset as usize, &mut stdout_data).unwrap();
            }
            let data = String::from_utf8(stdout_data).unwrap();
            *_caller.data_mut() = data;
            mem.write(&mut _caller, nwrite as usize, &stdout_data_len_buf).unwrap();
            Ok(0)
        },
    ).unwrap();

    linker.func_wrap(
        "wasi_snapshot_preview1",
        "proc_exit",
        move |mut _caller: Caller<'_, String>, _code: i32| {
            Ok(())
        },
    ).unwrap();

    // Instantiate our first module which only uses WASI, then register that
    // instance with the linker since the next linking will use it.

    let instance = linker
        .instantiate(&mut store, &module)
        .unwrap()
        .start(&mut store)
        .unwrap();

    // let instance2 = linker.instantiate(&mut store, &quick_js_module).unwrap();
    // instance2.start(&mut store).unwrap().

    println!("cycle-tracker-start: call wasm");
    let add_two = instance
        .get_typed_func::<(), ()>(&mut store, "_start")
        .expect("Failed to get typed_func");
    let _res = add_two.call(&mut store, ()).expect("Failed to call");
    let output = store.data();
    println!("cycle-tracker-end: call wasm");
    let output: u32 = output.parse().unwrap();
    println!("fib {} - {}", input, output);

    sp1_zkvm::io::commit(&output);
}