wacko/structs/
function_body.rs

1use crate::Error;
2use crate::Instruction;
3use crate::ValType;
4use leb128::write;
5use std::io::Write;
6
7#[derive(Clone)]
8pub struct FnBody<'a> {
9    fn_type: (Vec<ValType>, Vec<ValType>),
10    /// `Vec<(count, ValType)>`
11    locals: Vec<(u32, ValType)>,
12    instructions: Vec<Instruction<'a>>,
13}
14
15impl<'a> FnBody<'a> {
16    pub fn new(arguments: Vec<ValType>, return_type: Vec<ValType>) -> Self {
17        Self {
18            fn_type: (arguments, return_type),
19            locals: Vec::new(),
20            instructions: Vec::new(),
21        }
22    }
23
24    pub fn add_local(&mut self, kind: ValType) {
25        if self.locals.is_empty() {
26            self.locals.push((1, kind));
27            return;
28        }
29
30        let index = self.locals.len() - 1;
31        if self.locals[index].1 == kind {
32            self.locals[index].0 += 1;
33            return;
34        }
35
36        self.locals.push((1, kind));
37    }
38
39    pub fn add_instruction(&mut self, instruction: Instruction<'a>) {
40        self.instructions.push(instruction);
41    }
42
43    pub fn add_instructions<T>(&mut self, instructions: T)
44    where
45        T: IntoIterator<Item = Instruction<'a>>,
46    {
47        self.instructions.extend(instructions)
48    }
49
50    pub(crate) fn get_fn_type(&self) -> (&[ValType], &[ValType]) {
51        (&self.fn_type.0, &self.fn_type.1)
52    }
53
54    pub(crate) fn compile(mut self, writer: &mut impl Write) -> Result<(), Error> {
55        let mut buff = Vec::new();
56
57        write::unsigned(&mut buff, self.locals.len() as u64)?;
58        for local in self.locals {
59            write::unsigned(&mut buff, local.0 as u64)?;
60            (&mut buff).write_all(&[local.1.into()])?;
61        }
62        self.instructions.push(Instruction::End);
63
64        for x in self.instructions {
65            x.encode(&mut buff)?;
66        }
67
68        write::unsigned(writer, buff.len() as u64)?;
69        writer.write_all(&buff)?;
70        Ok(())
71    }
72
73    pub(crate) fn optimize(&mut self) {}
74}
75
76#[cfg(test)]
77mod test {
78    use crate::*;
79    #[test]
80    fn encode_fn_body() {
81        let mut buff = Vec::new();
82        let args = vec![ValType::I32, ValType::I32];
83        let return_type = vec![ValType::I32];
84        let mut fn_body = FnBody::new(args, return_type);
85        fn_body.add_instructions(vec![
86            Instruction::LocalGet(0),
87            Instruction::LocalGet(1),
88            Instruction::I32Add,
89        ]);
90        fn_body.compile(&mut buff).unwrap();
91        assert_eq!(buff, vec![0x07, 0x00, 0x20, 0x00, 0x20, 0x01, 0x6A, 0x0B]);
92    }
93}