polkavm_common/
writer.rs

1use crate::program::{self, Instruction, ProgramExport, ProgramImport};
2use alloc::vec::Vec;
3use core::ops::Range;
4
5#[derive(Default)]
6pub struct ProgramBlobBuilder {
7    ro_data_size: u32,
8    rw_data_size: u32,
9    stack_size: u32,
10    ro_data: Vec<u8>,
11    rw_data: Vec<u8>,
12    imports: Vec<ProgramImport<'static>>,
13    exports: Vec<ProgramExport<'static>>,
14    jump_table: Vec<u8>,
15    code: Vec<u8>,
16    custom: Vec<(u8, Vec<u8>)>,
17    instruction_count: u32,
18    basic_block_count: u32,
19}
20
21impl ProgramBlobBuilder {
22    pub fn new() -> Self {
23        Self::default()
24    }
25
26    pub fn set_ro_data_size(&mut self, size: u32) {
27        self.ro_data_size = size;
28    }
29
30    pub fn set_rw_data_size(&mut self, size: u32) {
31        self.rw_data_size = size;
32    }
33
34    pub fn set_stack_size(&mut self, size: u32) {
35        self.stack_size = size;
36    }
37
38    pub fn set_ro_data(&mut self, data: Vec<u8>) {
39        self.ro_data = data;
40    }
41
42    pub fn set_rw_data(&mut self, data: Vec<u8>) {
43        self.rw_data = data;
44    }
45
46    pub fn add_import(&mut self, import: ProgramImport) {
47        self.imports.push(import.into_owned());
48    }
49
50    pub fn add_export(&mut self, export: ProgramExport) {
51        self.exports.push(export.into_owned());
52    }
53
54    pub fn set_jump_table(&mut self, jump_table: &[u32]) {
55        self.jump_table.clear();
56        let mut writer = Writer::new(&mut self.jump_table);
57        for &target in jump_table {
58            writer.push_varint(target);
59        }
60    }
61
62    pub fn set_code(&mut self, code: &[Instruction]) {
63        self.instruction_count = 0;
64        self.basic_block_count = 0;
65        for instruction in code {
66            let mut buffer = [0; program::MAX_INSTRUCTION_LENGTH];
67            let length = instruction.serialize_into(&mut buffer);
68            self.code.extend_from_slice(&buffer[..length]);
69            self.instruction_count += 1;
70
71            if instruction.opcode().starts_new_basic_block() {
72                self.basic_block_count += 1;
73            }
74        }
75    }
76
77    pub fn add_custom_section(&mut self, section: u8, contents: Vec<u8>) {
78        self.custom.push((section, contents));
79    }
80
81    pub fn into_vec(self) -> Vec<u8> {
82        let mut output = Vec::new();
83        let mut writer = Writer::new(&mut output);
84
85        writer.push_raw_bytes(&program::BLOB_MAGIC);
86        writer.push_byte(program::BLOB_VERSION_V1);
87
88        if self.ro_data_size > 0 || self.rw_data_size > 0 || self.stack_size > 0 {
89            writer.push_section_inplace(program::SECTION_MEMORY_CONFIG, |writer| {
90                writer.push_varint(self.ro_data_size);
91                writer.push_varint(self.rw_data_size);
92                writer.push_varint(self.stack_size);
93            });
94        }
95
96        writer.push_section(program::SECTION_RO_DATA, &self.ro_data);
97        writer.push_section(program::SECTION_RW_DATA, &self.rw_data);
98        if !self.imports.is_empty() {
99            writer.push_section_inplace(program::SECTION_IMPORTS, |writer| {
100                writer.push_varint(self.imports.len().try_into().expect("too many imports"));
101                for import in self.imports {
102                    writer.push_bytes_with_length(import.symbol());
103                }
104            });
105        }
106
107        if !self.exports.is_empty() {
108            writer.push_section_inplace(program::SECTION_EXPORTS, |writer| {
109                writer.push_varint(self.exports.len().try_into().expect("too many exports"));
110                for export in self.exports {
111                    writer.push_varint(export.jump_target());
112                    writer.push_bytes_with_length(export.symbol());
113                }
114            });
115        }
116
117        writer.push_section(program::SECTION_JUMP_TABLE, &self.jump_table);
118        writer.push_section_inplace(program::SECTION_CODE, |writer| {
119            writer.push_varint(self.instruction_count);
120            writer.push_varint(self.basic_block_count);
121            writer.push_raw_bytes(&self.code);
122        });
123
124        for (section, contents) in self.custom {
125            writer.push_section(section, &contents);
126        }
127
128        writer.push_raw_bytes(&[program::SECTION_END_OF_FILE]);
129        output
130    }
131}
132
133pub struct Writer<'a> {
134    buffer: &'a mut Vec<u8>,
135}
136
137impl<'a> Writer<'a> {
138    pub fn new(buffer: &'a mut Vec<u8>) -> Self {
139        Self { buffer }
140    }
141
142    fn push_section_inplace(&mut self, section: u8, callback: impl FnOnce(&mut Self)) -> Range<usize> {
143        let section_position = self.buffer.len();
144        self.buffer.push(section);
145
146        // Reserve the space for the length varint.
147        let length_position = self.buffer.len();
148        self.push_raw_bytes(&[0xff_u8; crate::varint::MAX_VARINT_LENGTH]);
149
150        let payload_position = self.buffer.len();
151        callback(self);
152
153        let payload_length: u32 = (self.buffer.len() - payload_position).try_into().expect("section size overflow");
154        if payload_length == 0 {
155            // Nothing was written by the callback. Skip writing the section.
156            self.buffer.truncate(section_position);
157            return 0..0;
158        }
159
160        // Write the length varint.
161        let length_length = crate::varint::write_varint(payload_length, &mut self.buffer[length_position..]);
162
163        // Drain any excess length varint bytes.
164        self.buffer
165            .drain(length_position + length_length..length_position + crate::varint::MAX_VARINT_LENGTH);
166
167        length_position + length_length..self.buffer.len()
168    }
169
170    fn push_section(&mut self, section: u8, contents: &[u8]) {
171        if contents.is_empty() {
172            return;
173        }
174
175        self.push_byte(section);
176        self.push_varint(contents.len().try_into().expect("section size overflow"));
177        self.push_raw_bytes(contents);
178    }
179
180    pub fn push_raw_bytes(&mut self, slice: &[u8]) {
181        self.buffer.extend_from_slice(slice);
182    }
183
184    pub fn push_byte(&mut self, byte: u8) {
185        self.buffer.push(byte);
186    }
187
188    pub fn push_u32(&mut self, value: u32) {
189        self.push_raw_bytes(&value.to_le_bytes());
190    }
191
192    pub fn push_varint(&mut self, value: u32) {
193        let mut buffer = [0xff_u8; crate::varint::MAX_VARINT_LENGTH];
194        let length = crate::varint::write_varint(value, &mut buffer);
195        self.push_raw_bytes(&buffer[..length]);
196    }
197
198    pub fn push_bytes_with_length(&mut self, slice: &[u8]) {
199        self.push_varint(slice.len().try_into().expect("length overflow"));
200        self.push_raw_bytes(slice);
201    }
202
203    pub fn len(&self) -> usize {
204        self.buffer.len()
205    }
206
207    pub fn is_empty(&self) -> bool {
208        self.buffer.is_empty()
209    }
210}
polkavm_common/writer.rs

polkavm_common/
writer.rs
