zkvmc_context/
lib.rs

1//! This module provides the context for zkVMc runtime.
2mod compiler;
3mod syscall;
4
5pub use compiler::*;
6use rustc_hash::FxHashMap;
7use std::{marker::PhantomPinned, ptr::NonNull};
8pub use syscall::*;
9use zkvmc_core::{
10    addr2frame::LookupPc,
11    image::Image,
12    memory::{
13        read_bytes, read_word, read_words, write_bytes, write_word, write_words, ForkableMemory,
14        MmapOffset,
15    },
16    traits::Reset,
17};
18
19/// big enough to hold the whole virtual memory
20pub const MEM_SIZE: usize = 0x100000003;
21//TODO: just a temporary value, should be calculated
22// the maximum size of the code block, 100MB is enough?
23const TRACE_BUFF_SIZE: usize = 100 * 1024 * 1024;
24
25/// The Context for the zkVMc runtime.
26#[repr(C)]
27pub struct Context {
28    /// Registers.
29    pub regs: [u32; 32],
30    /// Program Counter.
31    pub pc: u32,
32    /// Next Program Counter.
33    pub next_pc: u32,
34    /// Linear memory between VM and JIT code.
35    pub mem: ForkableMemory,
36    /// Trace buffer for tracing the instructions.
37    pub trace_buf: MmapOffset,
38    /// Handler function for the `ecall` instruction.
39    pub ecall_handler: extern "C" fn(NonNull<Context>),
40    /// Handler function for the `ebreak` instruction.
41    pub ebreak_handler: extern "C" fn(NonNull<Context>),
42    /// Handler function for an undefined instruction.
43    pub undefined_handler: extern "C" fn(NonNull<Context>, u32),
44    /// Handler function for the execution traces.
45    pub trace_handler: extern "C" fn(NonNull<Context>),
46    /// Function called by compiled code or VM to lookup or translate the next compiled code block.
47    pub trampoline: extern "C" fn(NonNull<Context>) -> JitFn,
48
49    /// --------------------------
50    /// Additional context.
51    pub addition: NonNull<dyn Addition>,
52    /// Compile the code block at `pc` and return the compiled function.
53    pub compiler: Box<dyn Compiler>,
54    pub compile_cfg: CompileConfig,
55    /// Exit code of the program.
56    pub exit_code: Option<u32>,
57
58    image: Image,
59    _marker: PhantomPinned,
60}
61
62pub trait Addition: Reset + LookupPc {}
63
64impl Context {
65    // FIXME: use builder pattern
66    pub fn new(
67        compiler: Option<Box<dyn Compiler>>,
68        entry: u32,
69        image: FxHashMap<u32, &[u8]>,
70        addition: NonNull<dyn Addition>,
71        cfg: Option<CompileConfig>,
72        ecall_handler: Option<extern "C" fn(NonNull<Context>)>,
73        ebreak_handler: Option<extern "C" fn(NonNull<Context>)>,
74        undefined_handler: Option<extern "C" fn(NonNull<Context>, u32)>,
75        interrupt_handler: Option<extern "C" fn(NonNull<Context>)>,
76    ) -> Context {
77        let mem = ForkableMemory::new(MEM_SIZE).expect("Failed to create memory");
78        let mut ctx = Context {
79            regs: [0; 32],
80            pc: 0,
81            next_pc: 0,
82            mem,
83            ecall_handler: ecall_handler.unwrap_or(syscall::default_ecall_handler),
84            ebreak_handler: ebreak_handler.unwrap_or(syscall::default_ebreak_handler),
85            undefined_handler: undefined_handler.unwrap_or(syscall::default_undefined_handler),
86            trace_handler: interrupt_handler.unwrap_or(syscall::default_trace_handler),
87            addition,
88            exit_code: None,
89            trampoline,
90            compiler: compiler.unwrap_or_else(|| Box::new(NoopCompiler)),
91            compile_cfg: cfg.unwrap_or_default(),
92            trace_buf: MmapOffset::new(TRACE_BUFF_SIZE).expect("Failed to create trace buffer"),
93            image: Image::new(entry, image),
94            _marker: PhantomPinned,
95        };
96        ctx.setup();
97        ctx
98    }
99
100    fn setup(&mut self) {
101        // load image into memory
102        self.image.load(&mut self.pc, self.mem.as_send_sync_ptr());
103        self.next_pc = self.pc;
104    }
105
106    #[inline]
107    pub fn update_pc(&mut self) {
108        self.pc = self.next_pc;
109    }
110
111    #[inline]
112    pub fn write_bytes(&mut self, addr: u32, bytes: &[u8]) {
113        unsafe { write_bytes(self.mem.as_ptr(), addr, bytes) };
114    }
115
116    #[inline]
117    pub fn write_word(&mut self, addr: u32, word: u32) {
118        unsafe { write_word(self.mem.as_ptr(), addr, word) };
119    }
120
121    #[inline]
122    pub fn write_words(&mut self, addr: u32, words: &[u32]) {
123        unsafe { write_words(self.mem.as_ptr(), addr, words) };
124    }
125
126    #[inline]
127    pub fn read_bytes(&self, addr: u32, len: usize) -> &[u8] {
128        unsafe { read_bytes(self.mem.as_ptr(), addr, len) }
129    }
130
131    #[inline]
132    pub fn read_word(&self, addr: u32) -> u32 {
133        unsafe { read_word(self.mem.as_ptr(), addr) }
134    }
135
136    #[inline]
137    pub fn read_words(&self, addr: u32, len: usize) -> &[u32] {
138        unsafe { read_words(self.mem.as_ptr(), addr, len) }
139    }
140
141    #[inline]
142    pub fn as_addition<T>(&self) -> &T {
143        unsafe { self.addition.cast().as_ref() }
144    }
145
146    #[inline]
147    pub fn as_addition_mut<T>(&mut self) -> &mut T {
148        unsafe { self.addition.cast().as_mut() }
149    }
150}
151
152impl Drop for Context {
153    fn drop(&mut self) {
154        unsafe {
155            let _ = Box::from_raw(self.addition.as_ptr());
156        }
157    }
158}
159
160impl Reset for Context {
161    fn reset(&mut self) {
162        self.regs = [0; 32];
163        self.exit_code = None;
164        self.mem.reset();
165        self.trace_buf.reset();
166        self.setup();
167        unsafe { self.addition.as_mut().reset() }
168    }
169}
zkvmc_context/lib.rs

zkvmc_context/
lib.rs
