Skip to main content

kaiju_vm_core/
vm.rs

1use crate::processor::{OpAction, Processor};
2use crate::state::{State, Value};
3use byteorder::{BigEndian, ReadBytesExt};
4use compiler_core::compile_program;
5use compiler_core::module_reader::ModuleReader;
6use core::assembly::{encode_assembly, OpIndex};
7use core::error::*;
8use core::program::OpsDescriptor;
9use core::validator::DeepValidator;
10use core::vm::{Data, VmAssembly};
11use std::ffi::CString;
12use std::io::{Cursor, Read, Seek, SeekFrom};
13use std::mem::size_of;
14
15#[derive(Debug, Clone, Copy, PartialEq)]
16pub enum ExecutionStage {
17    Uninitialized,
18    Running,
19    Complete,
20}
21
22#[derive(Debug, Clone)]
23pub struct CallStackFrame {
24    function: usize,
25    address: usize,
26    params_stackpos: usize,
27    result: Option<Value>,
28    locals: Option<Value>,
29    op_target_addresses: Vec<usize>,
30    op_param_addresses: Vec<usize>,
31    op_stackpos: usize,
32}
33
34impl CallStackFrame {
35    #[inline]
36    pub(crate) fn new(
37        function: usize,
38        address: usize,
39        params_stackpos: usize,
40        result: Option<Value>,
41        locals: Option<Value>,
42    ) -> Self {
43        Self {
44            function,
45            address,
46            params_stackpos,
47            result,
48            locals,
49            op_target_addresses: vec![],
50            op_param_addresses: vec![],
51            op_stackpos: 0,
52        }
53    }
54
55    #[inline]
56    pub fn function(&self) -> usize {
57        self.function
58    }
59
60    #[inline]
61    pub fn address(&self) -> usize {
62        self.address
63    }
64
65    #[inline]
66    pub fn params_stackpos(&self) -> usize {
67        self.params_stackpos
68    }
69
70    #[inline]
71    pub fn result(&self) -> &Option<Value> {
72        &self.result
73    }
74
75    #[inline]
76    pub fn locals(&self) -> &Option<Value> {
77        &self.locals
78    }
79
80    #[inline]
81    pub fn op_target_addresses(&self) -> &[usize] {
82        &self.op_target_addresses
83    }
84
85    #[inline]
86    pub fn op_param_addresses(&self) -> &[usize] {
87        &self.op_param_addresses
88    }
89
90    pub(crate) fn collect_params_targets(&mut self) -> (Vec<usize>, Vec<usize>) {
91        let params = self.op_param_addresses.clone();
92        let targets = self.op_target_addresses.clone();
93        self.op_param_addresses.clear();
94        self.op_target_addresses.clear();
95        (params, targets)
96    }
97
98    pub(crate) fn duplicate(&self) -> Self {
99        Self {
100            function: self.function,
101            address: self.address,
102            params_stackpos: self.params_stackpos,
103            result: self.result,
104            locals: self.locals,
105            op_target_addresses: vec![],
106            op_param_addresses: vec![],
107            op_stackpos: 0,
108        }
109    }
110}
111
112#[derive(Debug, Clone)]
113pub struct Vm {
114    assembly: VmAssembly,
115    state: State,
116    stage: ExecutionStage,
117    callstack: Vec<CallStackFrame>,
118    data: Option<Value>,
119    globals: Option<Value>,
120    pointers: Vec<usize>,
121}
122
123impl Vm {
124    #[allow(clippy::new_ret_no_self)]
125    #[inline]
126    pub fn new(assembly: VmAssembly, stack_size: usize, memory_size: usize) -> SimpleResult<Self> {
127        if stack_size % 4 != 0 {
128            return Err(SimpleError::new(format!(
129                "Stack size is not aligned to 4 bytes: {}",
130                stack_size
131            )));
132        }
133        if memory_size % 4 != 0 {
134            return Err(SimpleError::new(format!(
135                "Memory size is not aligned to 4 bytes: {}",
136                memory_size
137            )));
138        }
139        if memory_size + stack_size > usize::max_value() {
140            return Err(SimpleError::new(format!(
141                "Memory and stack sizes combined cannot exceed machine pointer range: {}",
142                usize::max_value()
143            )));
144        }
145        Ok(Self {
146            assembly,
147            state: State::new(stack_size, memory_size),
148            stage: ExecutionStage::Uninitialized,
149            callstack: vec![],
150            data: None,
151            globals: None,
152            pointers: vec![],
153        })
154    }
155
156    #[inline]
157    pub fn from_bytes(bytes: Vec<u8>, stack_size: usize, memory_size: usize) -> SimpleResult<Self> {
158        Self::new(VmAssembly::new(bytes)?, stack_size, memory_size)
159    }
160
161    #[inline]
162    pub fn from_stream(
163        stream: &mut dyn Read,
164        stack_size: usize,
165        memory_size: usize,
166    ) -> SimpleResult<Self> {
167        let mut bytes = vec![];
168        stream.read_to_end(&mut bytes)?;
169        Self::from_bytes(bytes, stack_size, memory_size)
170    }
171
172    #[inline]
173    pub fn from_source<V, R>(
174        entry_path: &str,
175        module_reader: R,
176        ops_descriptor: &OpsDescriptor,
177        stack_size: usize,
178        memory_size: usize,
179    ) -> SimpleResult<Self>
180    where
181        V: DeepValidator,
182        R: ModuleReader,
183    {
184        let program = compile_program::<V, R>(entry_path, module_reader, ops_descriptor)?;
185        let bytes = encode_assembly(&program, ops_descriptor)?;
186        Self::from_bytes(bytes, stack_size, memory_size)
187    }
188
189    #[inline]
190    pub fn fork(&self) -> SimpleResult<Self> {
191        self.fork_advanced(self.state.stack_size(), self.state.memory_size())
192    }
193
194    #[inline]
195    pub fn fork_advanced(&self, stack_size: usize, memory_size: usize) -> SimpleResult<Self> {
196        if stack_size % 4 != 0 {
197            return Err(SimpleError::new(format!(
198                "Stack size is not aligned to 4 bytes: {}",
199                stack_size
200            )));
201        }
202        if memory_size % 4 != 0 {
203            return Err(SimpleError::new(format!(
204                "Memory size is not aligned to 4 bytes: {}",
205                memory_size
206            )));
207        }
208        if memory_size + stack_size > usize::max_value() {
209            return Err(SimpleError::new(format!(
210                "Memory and stack sizes combined cannot exceed machine pointer range: {}",
211                usize::max_value()
212            )));
213        }
214        Ok(Self {
215            assembly: self.assembly.clone(),
216            state: State::new(stack_size, memory_size),
217            stage: ExecutionStage::Uninitialized,
218            callstack: vec![],
219            data: None,
220            globals: None,
221            pointers: vec![],
222        })
223    }
224
225    #[inline]
226    pub fn assembly(&self) -> &VmAssembly {
227        &self.assembly
228    }
229
230    #[inline]
231    pub fn state(&self) -> &State {
232        &self.state
233    }
234
235    #[inline]
236    pub fn state_mut(&mut self) -> &mut State {
237        &mut self.state
238    }
239
240    #[inline]
241    pub fn stage(&self) -> ExecutionStage {
242        self.stage
243    }
244
245    pub fn start(&mut self, entry: &str) -> SimpleResult<()> {
246        if self.stage == ExecutionStage::Uninitialized {
247            let f = {
248                if let Some(f) = self.assembly.function_by_id(entry) {
249                    if !f.params().is_empty() {
250                        return Err(SimpleError::new(format!(
251                            "Trying to call entry function `{}` that does have params",
252                            entry
253                        )));
254                    }
255                    if f.typeid().is_some() {
256                        return Err(SimpleError::new(
257                            "Trying to run function with return value".to_owned(),
258                        ));
259                    }
260                    f.index()
261                } else {
262                    return Err(SimpleError::new(format!(
263                        "Trying to start non-existing function: {}",
264                        entry
265                    )));
266                }
267            };
268            let vd = self.alloc_data()?;
269            let vg = self
270                .state
271                .alloc_memory_value(self.assembly.globals_size())?;
272            self.data = vd;
273            self.globals = Some(vg);
274            self.stage = ExecutionStage::Running;
275            self.call_function(f)?;
276            Ok(())
277        } else {
278            Err(SimpleError::new(
279                "Trying to start running or complete VM".to_owned(),
280            ))
281        }
282    }
283
284    #[inline]
285    pub fn can_resume(&self) -> bool {
286        self.stage == ExecutionStage::Running && !self.callstack.is_empty()
287    }
288
289    pub fn resume<P>(&mut self) -> SimpleResult<()>
290    where
291        P: Processor,
292    {
293        if self.stage == ExecutionStage::Running {
294            while self.resume_op::<P>()? {}
295            Ok(())
296        } else {
297            Err(SimpleError::new(
298                "Trying to resume uninitialized or complete VM".to_owned(),
299            ))
300        }
301    }
302
303    #[inline]
304    pub fn consume<P>(&mut self) -> SimpleResult<()>
305    where
306        P: Processor,
307    {
308        if self.stage == ExecutionStage::Running {
309            while self.can_resume() {
310                self.resume::<P>()?;
311            }
312            Ok(())
313        } else {
314            Err(SimpleError::new(
315                "Trying to consume uninitialized or complete VM".to_owned(),
316            ))
317        }
318    }
319
320    #[inline]
321    pub fn run<P>(&mut self, entry: &str) -> SimpleResult<()>
322    where
323        P: Processor,
324    {
325        self.start(entry)?;
326        self.consume::<P>()
327    }
328
329    pub fn find_label(&self, id: &str) -> Option<usize> {
330        if let Some((_, func, _)) = self.location() {
331            if let Some(f) = self.assembly.function_body_by_index(func) {
332                f.labels().get(id).cloned()
333            } else {
334                None
335            }
336        } else {
337            None
338        }
339    }
340
341    fn alloc_data(&mut self) -> SimpleResult<Option<Value>> {
342        if self.assembly.data().is_empty() {
343            return Ok(None);
344        }
345        let mut addr = ::std::usize::MAX;
346        let mut size = 0;
347        for d in self.assembly.data() {
348            let v = match d {
349                Data::I8(v) => {
350                    let r = self.state.alloc_memory_value(1)?;
351                    self.state.store_data(r.address, v)?;
352                    r
353                }
354                Data::U8(v) => {
355                    let r = self.state.alloc_memory_value(1)?;
356                    self.state.store_data(r.address, v)?;
357                    r
358                }
359                Data::I16(v) => {
360                    let r = self.state.alloc_memory_value(2)?;
361                    self.state.store_data(r.address, v)?;
362                    r
363                }
364                Data::U16(v) => {
365                    let r = self.state.alloc_memory_value(2)?;
366                    self.state.store_data(r.address, v)?;
367                    r
368                }
369                Data::I32(v) => {
370                    let r = self.state.alloc_memory_value(4)?;
371                    self.state.store_data(r.address, v)?;
372                    r
373                }
374                Data::U32(v) => {
375                    let r = self.state.alloc_memory_value(4)?;
376                    self.state.store_data(r.address, v)?;
377                    r
378                }
379                Data::I64(v) => {
380                    let r = self.state.alloc_memory_value(8)?;
381                    self.state.store_data(r.address, v)?;
382                    r
383                }
384                Data::U64(v) => {
385                    let r = self.state.alloc_memory_value(8)?;
386                    self.state.store_data(r.address, v)?;
387                    r
388                }
389                Data::F32(v) => {
390                    let r = self.state.alloc_memory_value(4)?;
391                    self.state.store_data(r.address, v)?;
392                    r
393                }
394                Data::F64(v) => {
395                    let r = self.state.alloc_memory_value(8)?;
396                    self.state.store_data(r.address, v)?;
397                    r
398                }
399                Data::Isize(v) => {
400                    let r = self.state.alloc_memory_value(size_of::<isize>())?;
401                    self.state.store_data(r.address, v)?;
402                    r
403                }
404                Data::Usize(v) => {
405                    let r = self.state.alloc_memory_value(size_of::<usize>())?;
406                    self.state.store_data(r.address, v)?;
407                    r
408                }
409                Data::String(s) => {
410                    if let Ok(ref s) = CString::new(s.as_str()) {
411                        let bytes = s.as_bytes_with_nul();
412                        let sv = self.state.alloc_memory_value(bytes.len())?;
413                        self.state.store_bytes(sv.address, bytes)?;
414                        let v = self.state.alloc_memory_value(size_of::<usize>())?;
415                        self.state.store_data(v.address, &sv.address)?;
416                        Value::new(sv.address, sv.size + v.size)
417                    } else {
418                        return Err(SimpleError::new(format!(
419                            "Could not store string that is not C-compatible: '{}'",
420                            s
421                        )));
422                    }
423                }
424                Data::None => unreachable!(),
425            };
426            addr = addr.min(v.address);
427            size += v.size;
428        }
429        Ok(Some(Value::new(addr, size)))
430    }
431
432    fn call_function(&mut self, function: usize) -> SimpleResult<()> {
433        if self.stage == ExecutionStage::Running {
434            if let Some(f) = self.assembly.function_by_index(function) {
435                let params_stackpos = self
436                    .state
437                    .stack_pos()
438                    .checked_sub(f.params().iter().map(|l| l.size()).sum())
439                    .unwrap();
440                let r = if let Some(t) = f.typeid() {
441                    Some(self.state.alloc_stack_value(self.assembly.type_size(t))?)
442                } else {
443                    None
444                };
445                let size = f.locals().iter().map(|l| l.size()).sum();
446                let l = if size != 0 {
447                    Some(self.state.alloc_stack_value(size)?)
448                } else {
449                    None
450                };
451                self.callstack
452                    .push(CallStackFrame::new(f.index(), 0, params_stackpos, r, l));
453                Ok(())
454            } else {
455                Err(SimpleError::new(format!(
456                    "Trying to call non-existing function with index: {}",
457                    function
458                )))
459            }
460        } else {
461            Err(SimpleError::new(
462                "Trying to call function on uninitialized or complete VM".to_owned(),
463            ))
464        }
465    }
466
467    fn return_function(&mut self) -> SimpleResult<()> {
468        if self.stage == ExecutionStage::Running {
469            if let Some(frame) = self.callstack.pop() {
470                if let Some(v) = frame.result() {
471                    let bytes = self.state.load_bytes(v.address, v.size)?;
472                    self.state.stack_reset(frame.params_stackpos())?;
473                    let v = self.state.stack_push_bytes(&bytes)?;
474                    self.pointers.push(v.address);
475                } else {
476                    self.state.stack_reset(frame.params_stackpos())?;
477                    self.pointers.push(0usize);
478                }
479                if !self.can_resume() {
480                    self.stage = ExecutionStage::Complete;
481                }
482                Ok(())
483            } else {
484                Err(SimpleError::new(
485                    "Trying to return from no running function".to_owned(),
486                ))
487            }
488        } else {
489            Err(SimpleError::new(
490                "Trying to call function on uninitialized or complete VM".to_owned(),
491            ))
492        }
493    }
494
495    fn location(&self) -> Option<(usize, usize, usize)> {
496        if let Some(f) = self.callstack.last() {
497            Some((self.callstack.len() - 1, f.function(), f.address()))
498        } else {
499            None
500        }
501    }
502
503    fn resume_op<P>(&mut self) -> SimpleResult<bool>
504    where
505        P: Processor,
506    {
507        if let Some((i, func, addr)) = self.location() {
508            let b = self.assembly.function_body_by_index(func).unwrap();
509            let bodysize = b.code().len();
510            if addr >= bodysize {
511                self.return_function()?;
512                return Ok(false);
513            }
514            let mut stream = Cursor::new(b.code());
515            stream.seek(SeekFrom::Start(addr as u64))?;
516            loop {
517                let op = OpIndex::from(stream.read_u8()?);
518                match op {
519                    OpIndex::NoOp => unreachable!(),
520                    OpIndex::DataPointer => {
521                        let offset = stream.read_u64::<BigEndian>()? as usize;
522                        let address = self.data.unwrap().address + offset;
523                        self.callstack[i].address = stream.position() as usize;
524                        self.pointers.push(address);
525                    }
526                    OpIndex::ParamsPointer => {
527                        let offset = stream.read_u64::<BigEndian>()? as usize;
528                        let address = self.callstack[i].params_stackpos + offset;
529                        self.callstack[i].address = stream.position() as usize;
530                        self.pointers.push(address);
531                    }
532                    OpIndex::ResultPointer => {
533                        let address = self.callstack[i].result.unwrap().address;
534                        self.callstack[i].address = stream.position() as usize;
535                        self.pointers.push(address);
536                    }
537                    OpIndex::LocalsPointer => {
538                        let offset = stream.read_u64::<BigEndian>()? as usize;
539                        let address = self.callstack[i].locals.unwrap().address + offset;
540                        self.callstack[i].address = stream.position() as usize;
541                        self.pointers.push(address);
542                    }
543                    OpIndex::GlobalsPointer => {
544                        let offset = stream.read_u64::<BigEndian>()? as usize;
545                        let address = self.globals.unwrap().address + offset;
546                        self.callstack[i].address = stream.position() as usize;
547                        self.pointers.push(address);
548                    }
549                    OpIndex::OffsetPointer => {
550                        let offset = stream.read_u64::<BigEndian>()? as usize;
551                        let address = self.pointers.pop().unwrap() + offset;
552                        self.callstack[i].address = stream.position() as usize;
553                        self.pointers.push(address);
554                    }
555                    OpIndex::ReferencePointer => {
556                        let v = self.state.stack_push_data(&self.pointers.pop().unwrap())?;
557                        self.callstack[i].address = stream.position() as usize;
558                        self.pointers.push(v.address);
559                    }
560                    OpIndex::DereferencePointer => {
561                        let address = self.pointers.pop().unwrap();
562                        let address = self.state.load_data::<usize>(address)?;
563                        self.callstack[i].address = stream.position() as usize;
564                        self.pointers.push(address);
565                    }
566                    OpIndex::StoreTargetAddress => {
567                        let address = self.pointers.pop().unwrap();
568                        self.callstack[i].address = stream.position() as usize;
569                        self.callstack[i].op_target_addresses.push(address);
570                    }
571                    OpIndex::StoreParamAddress => {
572                        let address = self.pointers.pop().unwrap();
573                        self.callstack[i].address = stream.position() as usize;
574                        self.callstack[i].op_param_addresses.push(address);
575                    }
576                    OpIndex::ProduceTuple => {
577                        let count = stream.read_u64::<BigEndian>()? as usize;
578                        let mut addresses_sizes = vec![];
579                        let mut size = 0;
580                        for _ in 0..count {
581                            let s = stream.read_u64::<BigEndian>()? as usize;
582                            let a = self.pointers.pop().unwrap();
583                            addresses_sizes.push((a, s));
584                            size += s;
585                        }
586                        let v = self.state.alloc_stack_value(size)?;
587                        let mut offset = 0;
588                        for (a, s) in addresses_sizes {
589                            self.state.memory_move(a, s, v.address + offset)?;
590                            offset += s;
591                        }
592                        self.callstack[i].address = stream.position() as usize;
593                        self.pointers.push(v.address);
594                    }
595                    OpIndex::ExecuteOpStart => {
596                        self.callstack[i].op_stackpos = self.state.stack_pos();
597                    }
598                    OpIndex::ExecuteOpStop => {
599                        let op = stream.read_u64::<BigEndian>()? as usize;
600                        let op = self.assembly.ops_map().get(op).unwrap().to_owned();
601                        let (params, targets) = self.callstack[i].collect_params_targets();
602                        let addr = stream.position() as usize;
603                        self.callstack[i].address = addr;
604                        let action = { P::process_op(&op, &params, &targets, self)? };
605                        self.state.stack_reset(self.callstack[i].op_stackpos)?;
606                        self.callstack[i].op_stackpos = 0;
607                        match action {
608                            OpAction::None => {
609                                if addr >= bodysize || !self.can_resume() {
610                                    self.return_function()?;
611                                }
612                            }
613                            OpAction::GoTo(a) => {
614                                self.callstack[i].address = a;
615                            }
616                            OpAction::Return => self.return_function()?,
617                        }
618                        break;
619                    }
620                    OpIndex::ExecuteOpInlineStart => {
621                        self.callstack[i].address = stream.position() as usize;
622                        self.callstack.push(self.callstack[i].duplicate());
623                        return Ok(true);
624                    }
625                    OpIndex::ExecuteOpInlineStop => {
626                        let op = stream.read_u64::<BigEndian>()? as usize;
627                        let size = stream.read_u64::<BigEndian>()? as usize;
628                        let op = self.assembly.ops_map().get(op).unwrap().to_owned();
629                        let (params, mut targets) = self.callstack[i].collect_params_targets();
630                        let v = self.state.alloc_stack_value(size)?;
631                        targets.push(v.address);
632                        self.callstack[i - 1].address = stream.position() as usize;
633                        {
634                            P::process_op(&op, &params, &targets, self)?;
635                        }
636                        self.callstack.pop();
637                        self.pointers.push(v.address);
638                        break;
639                    }
640                    OpIndex::CallFunction => {
641                        let f = stream.read_u64::<BigEndian>()? as usize;
642                        if let Some(func) = self.assembly.function_by_index(f) {
643                            let mut addresses_sizes = vec![];
644                            let mut size = 0;
645                            for p in func.params() {
646                                let a = self.pointers.pop().unwrap();
647                                addresses_sizes.push((a, p.size()));
648                                size += p.size();
649                            }
650                            let v = self.state.alloc_stack_value(size)?;
651                            let mut offset = 0;
652                            for (a, s) in addresses_sizes {
653                                self.state.memory_move(a, s, v.address + offset)?;
654                                offset += s;
655                            }
656                            self.callstack[i].address = stream.position() as usize;
657                            self.call_function(f)?;
658                            return Ok(true);
659                        } else {
660                            unreachable!();
661                        }
662                    }
663                }
664            }
665        } else {
666            self.stage = ExecutionStage::Complete;
667        }
668        Ok(false)
669    }
670}