risc0_zkvm/host/server/exec/

executor.rs

// Copyright 2024 RISC Zero, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::{
    cell::{Cell, RefCell},
    rc::Rc,
    sync::Arc,
    time::Instant,
};

use anyhow::{bail, Context as _, Result};
use risc0_binfmt::{
    AbiKind, ByteAddr, ExitCode, MemoryImage, Program, ProgramBinary, ProgramBinaryHeader,
    SystemState,
};
use risc0_circuit_rv32im::{
    execute::{
        platform::WORD_SIZE, Executor, Syscall as CircuitSyscall,
        SyscallContext as CircuitSyscallContext, DEFAULT_SEGMENT_LIMIT_PO2,
    },
    MAX_INSN_CYCLES, MAX_INSN_CYCLES_LOWER_PO2,
};
use risc0_core::scope;
use risc0_zkp::core::digest::Digest;
use risc0_zkvm_platform::{align_up, fileno};
use tempfile::tempdir;

use crate::{
    host::{client::env::SegmentPath, server::session::Session},
    receipt_claim::exit_code_from_rv32im_v2_claim,
    Assumptions, ExecutorEnv, FileSegmentRef, Output, Segment, SegmentRef,
};

use super::{
    profiler::{self, Profiler},
    syscall::{SyscallContext, SyscallTable},
    Journal,
};

// The Executor provides an implementation for the execution phase.
///
/// The proving phase uses an execution trace generated by the Executor.
pub struct ExecutorImpl<'a> {
    env: ExecutorEnv<'a>,
    image: MemoryImage,
    pub(crate) syscall_table: SyscallTable<'a>,
    profiler: Option<Rc<RefCell<Profiler>>>,
    return_cache: Cell<(u32, u32)>,
}

/// Check to see if the executor is compatible with the given guest program.
fn check_program_version(header: &ProgramBinaryHeader) -> Result<()> {
    let abi_kind = header.abi_kind;
    let abi_version = &header.abi_version;

    if abi_kind != AbiKind::V1Compat {
        bail!("ProgramBinary abi_kind mismatch {abi_kind:?} != AbiKind::V1Compat");
    }
    if !semver::VersionReq::parse("^1.0.0")
        .unwrap()
        .matches(abi_version)
    {
        bail!("ProgramBinary abi_version mismatch {abi_version} doesn't match ^1.0.0");
    }

    Ok(())
}

impl<'a> ExecutorImpl<'a> {
    /// Construct a new [ExecutorImpl] from a [MemoryImage] and entry point.
    ///
    /// Before a guest program is proven, the [ExecutorImpl] is responsible for
    /// deciding where a zkVM program should be split into [Segment]s and what
    /// work will be done in each segment. This is the execution phase:
    /// the guest program is executed to determine how its proof should be
    /// divided into subparts.
    #[allow(dead_code)]
    pub fn new(env: ExecutorEnv<'a>, image: MemoryImage) -> Result<Self> {
        Self::with_details(env, image, None)
    }

    /// Construct a new [ExecutorImpl] from the ELF binary of the guest program
    /// you want to run and an [ExecutorEnv] containing relevant
    /// environmental configuration details.
    pub fn from_elf(mut env: ExecutorEnv<'a>, elf: &[u8]) -> Result<Self> {
        let binary = ProgramBinary::decode(elf)?;
        check_program_version(&binary.header)?;

        let image = binary.to_image()?;

        let profiler = if env.pprof_out.is_some() {
            let profiler = Rc::new(RefCell::new(Profiler::new(
                &binary,
                None,
                profiler::read_enable_inline_functions_env_var(),
            )?));
            env.trace.push(profiler.clone());
            Some(profiler)
        } else {
            None
        };

        Self::with_details(env, image, profiler)
    }

    /// TODO(flaub)
    #[allow(dead_code)]
    pub(crate) fn from_kernel_elf(env: ExecutorEnv<'a>, elf: &[u8]) -> Result<Self> {
        let kernel = Program::load_elf(elf, u32::MAX)?;
        let image = MemoryImage::new_kernel(kernel);
        Self::with_details(env, image, None)
    }

    fn with_details(
        env: ExecutorEnv<'a>,
        image: MemoryImage,
        profiler: Option<Rc<RefCell<Profiler>>>,
    ) -> Result<Self> {
        let syscall_table = SyscallTable::from_env(&env);
        Ok(Self {
            env,
            image,
            syscall_table,
            profiler,
            return_cache: Cell::new((0, 0)),
        })
    }

    /// This will run the executor to get a [Session] which contain the results
    /// of the execution.
    pub fn run(&mut self) -> Result<Session> {
        if self.env.segment_path.is_none() {
            self.env.segment_path = Some(SegmentPath::TempDir(Arc::new(tempdir()?)));
        }

        let path = self.env.segment_path.clone().unwrap();
        self.run_with_callback(|segment| Ok(Box::new(FileSegmentRef::new(&segment, &path)?)))
    }

    /// Run the executor until [crate::ExitCode::Halted] or
    /// [crate::ExitCode::Paused] is reached, producing a [Session] as a result.
    pub fn run_with_callback<F>(&mut self, mut callback: F) -> Result<Session>
    where
        F: FnMut(Segment) -> Result<Box<dyn SegmentRef>> + Send,
    {
        scope!("execute");

        let journal = Journal::default();
        self.env
            .posix_io
            .borrow_mut()
            .with_write_fd(fileno::JOURNAL, journal.clone());

        let segment_limit_po2 = self
            .env
            .segment_limit_po2
            .unwrap_or(DEFAULT_SEGMENT_LIMIT_PO2 as u32) as usize;

        let mut refs = Vec::new();
        let mut exec = Executor::new(
            self.image.clone(),
            self,
            self.env.input_digest,
            self.env.trace.clone(),
        );

        let max_insn_cycles = if segment_limit_po2 >= 15 {
            MAX_INSN_CYCLES
        } else {
            MAX_INSN_CYCLES_LOWER_PO2
        };

        let start_time = Instant::now();
        let result = exec.run(
            segment_limit_po2,
            max_insn_cycles,
            self.env.session_limit,
            |inner| {
                let output = inner
                    .claim
                    .terminate_state
                    .is_some()
                    .then(|| -> Option<Result<_>> {
                        inner
                            .claim
                            .output
                            .and_then(|digest| {
                                (digest != Digest::ZERO)
                                    .then(|| journal.buf.lock().unwrap().clone())
                            })
                            .map(|journal| {
                                Ok(Output {
                                    journal: journal.into(),
                                    assumptions: Assumptions(
                                        self.syscall_table
                                            .assumptions_used
                                            .lock()
                                            .unwrap()
                                            .iter()
                                            .map(|(a, _)| a.clone().into())
                                            .collect::<Vec<_>>(),
                                    )
                                    .into(),
                                })
                            })
                    })
                    .flatten()
                    .transpose()?;

                let segment = Segment {
                    index: inner.index as u32,
                    inner,
                    output,
                };
                let segment_ref = callback(segment)?;
                refs.push(segment_ref);
                Ok(())
            },
        )?;
        let elapsed = start_time.elapsed();

        tracing::debug!("output_digest: {:?}", result.claim.output);

        let exit_code = exit_code_from_rv32im_v2_claim(&result.claim)?;

        // Set the session_journal to the committed data iff the guest set a non-zero output.
        let session_journal = result.claim.output.and_then(|digest| {
            (digest != Digest::ZERO).then(|| std::mem::take(&mut *journal.buf.lock().unwrap()))
        });
        if !exit_code.expects_output() && session_journal.is_some() {
            tracing::debug!(
                "dropping non-empty journal due to exit code {exit_code:?}: 0x{}",
                hex::encode(journal.buf.lock().unwrap().as_slice())
            );
        };

        let ecall_metrics = exec.take_ecall_metrics();

        // Take (clear out) the list of accessed assumptions.
        // Leave the assumptions cache so it can be used if execution is resumed from pause.
        let assumptions = std::mem::take(&mut *self.syscall_table.assumptions_used.lock().unwrap());
        let mmr_assumptions = self.syscall_table.mmr_assumptions.take();
        let pending_zkrs = self.syscall_table.pending_zkrs.take();
        let pending_keccaks = self.syscall_table.pending_keccaks.take();

        if let Some(profiler) = self.profiler.take() {
            let report = profiler.borrow_mut().finalize_to_vec();
            std::fs::write(self.env.pprof_out.as_ref().unwrap(), report)?;
        }

        self.image = result.post_image.clone();
        let syscall_metrics = self.syscall_table.metrics.borrow().clone();

        // NOTE: When a segment ends in a Halted(_) state, the post_digest will be null.
        let post_digest = match exit_code {
            ExitCode::Halted(_) => Digest::ZERO,
            _ => result.claim.post_state,
        };

        let session = Session {
            segments: refs,
            input: self.env.input_digest.unwrap_or_default(),
            journal: session_journal.map(crate::Journal::new),
            exit_code,
            assumptions,
            mmr_assumptions,
            user_cycles: result.user_cycles,
            paging_cycles: result.paging_cycles,
            reserved_cycles: result.reserved_cycles,
            total_cycles: result.total_cycles,
            pre_state: SystemState {
                pc: 0,
                merkle_root: result.claim.pre_state,
            },
            post_state: SystemState {
                pc: 0,
                merkle_root: post_digest,
            },
            pending_zkrs,
            pending_keccaks,
            syscall_metrics,
            hooks: vec![],
            ecall_metrics: ecall_metrics.into(),
        };

        tracing::info!("execution time: {elapsed:?}");
        session.log();

        assert_eq!(
            session.total_cycles,
            session.user_cycles + session.paging_cycles + session.reserved_cycles
        );

        Ok(session)
    }
}

struct ContextAdapter<'a, 'b> {
    ctx: &'b mut dyn CircuitSyscallContext,
    syscall_table: SyscallTable<'a>,
}

impl<'a> SyscallContext<'a> for ContextAdapter<'a, '_> {
    fn get_pc(&self) -> u32 {
        self.ctx.get_pc()
    }

    fn get_cycle(&self) -> u64 {
        self.ctx.get_cycle()
    }

    fn load_register(&mut self, idx: usize) -> u32 {
        self.ctx.peek_register(idx).unwrap()
    }

    fn load_page(&mut self, page_idx: u32) -> Result<Vec<u8>> {
        self.ctx.peek_page(page_idx)
    }

    fn load_u8(&mut self, addr: ByteAddr) -> Result<u8> {
        self.ctx.peek_u8(addr)
    }

    fn load_u32(&mut self, addr: ByteAddr) -> Result<u32> {
        self.ctx.peek_u32(addr)
    }

    fn syscall_table(&self) -> &SyscallTable<'a> {
        &self.syscall_table
    }
}

impl CircuitSyscall for ExecutorImpl<'_> {
    fn host_read(
        &self,
        ctx: &mut dyn CircuitSyscallContext,
        fd: u32,
        buf: &mut [u8],
    ) -> Result<u32> {
        if fd == 0 {
            let (a0, a1) = self.return_cache.get();
            tracing::trace!("host_read(buf: {}) -> ({a0:#010x}, {a1:#010x})", buf.len());
            let buf: &mut [u32] = bytemuck::cast_slice_mut(buf);
            (buf[0], buf[1]) = (a0, a1);
            return Ok(2 * WORD_SIZE as u32);
        }

        let mut ctx = ContextAdapter {
            ctx,
            syscall_table: self.syscall_table.clone(),
        };

        let name_ptr = ByteAddr(fd);
        let syscall = ctx.peek_string(name_ptr)?;
        tracing::trace!("host_read({syscall}, into_guest: {})", buf.len());

        let words = align_up(buf.len(), WORD_SIZE) / WORD_SIZE;
        let mut to_guest = vec![0u32; words];

        self.return_cache.set(
            self.syscall_table
                .get_syscall(&syscall)
                .context(format!("Unknown syscall: {syscall:?}"))?
                .borrow_mut()
                .syscall(&syscall, &mut ctx, &mut to_guest)?,
        );

        let bytes = bytemuck::cast_slice(to_guest.as_slice());
        let rlen = buf.len();
        buf.copy_from_slice(&bytes[..rlen]);

        Ok(rlen as u32)
    }

    fn host_write(&self, ctx: &mut dyn CircuitSyscallContext, _fd: u32, buf: &[u8]) -> Result<u32> {
        let str = String::from_utf8(buf.to_vec())?;
        tracing::debug!("R0VM[{}] {str}", ctx.get_cycle());
        Ok(buf.len() as u32)
    }
}

impl ContextAdapter<'_, '_> {
    fn peek_string(&mut self, mut addr: ByteAddr) -> Result<String> {
        tracing::trace!("peek_string: {addr:?}");
        let mut buf = Vec::new();
        loop {
            let bytes = self.ctx.peek_u8(addr)?;
            if bytes == 0 {
                break;
            }
            buf.push(bytes);
            addr += 1u32;
        }
        Ok(String::from_utf8(buf)?)
    }
}