Skip to content

Commit d782e4b

Browse files
committed
refactor: Remove recursion from script compiler
Compile Bitcoin Script iteratively using a stack of tasks. I use the opportunity to move the `script` and `cache` arguments into the `compile_to_bytes` method to get input sanitization for free.
1 parent cda22ff commit d782e4b

File tree

1 file changed

+62
-24
lines changed

1 file changed

+62
-24
lines changed

src/builder.rs

Lines changed: 62 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -149,26 +149,57 @@ impl StructuredScript {
149149
self
150150
}
151151

152-
// Compiles the builder to bytes using a cache that stores all called_script starting
153-
// positions in script to copy them from script instead of recompiling.
154-
fn compile_to_bytes(&self, script: &mut Vec<u8>, cache: &mut HashMap<u64, usize>) {
155-
for block in self.blocks.as_slice() {
156-
match block {
157-
Block::Call(id) => {
158-
let called_script = self
159-
.script_map
160-
.get(id)
161-
.expect("Missing entry for a called script");
162-
// Check if the script with the hash id is in cache
163-
match cache.get(id) {
152+
/// Compiles the script to bytes.
153+
fn compile_to_bytes(&self) -> Vec<u8> {
154+
#[derive(Debug)]
155+
enum Task<'a> {
156+
CompileCall {
157+
id: u64,
158+
called_script: &'a StructuredScript,
159+
},
160+
PushRaw(&'a ScriptBuf),
161+
UpdateCache {
162+
id: u64,
163+
called_script_start: usize,
164+
},
165+
}
166+
167+
fn push_script<'a>(script: &'a StructuredScript, tasks: &mut Vec<Task<'a>>) {
168+
for block in script.blocks.iter().rev() {
169+
match block {
170+
Block::Call(id) => {
171+
let called_script = script
172+
.script_map
173+
.get(id)
174+
.expect("missing entry for called script");
175+
tasks.push(Task::CompileCall {
176+
id: *id,
177+
called_script,
178+
});
179+
}
180+
Block::Script(buffer) => tasks.push(Task::PushRaw(buffer)),
181+
}
182+
}
183+
}
184+
185+
let mut tasks = Vec::new();
186+
let mut cache = HashMap::new();
187+
let mut script: Vec<u8> = Vec::with_capacity(self.size);
188+
push_script(self, &mut tasks);
189+
190+
while let Some(task) = tasks.pop() {
191+
match task {
192+
Task::CompileCall { id, called_script } => {
193+
match cache.get(&id) {
164194
Some(called_start) => {
165195
// Copy the already compiled called_script from the position it was
166196
// inserted in the compiled script.
167197
let start = script.len();
168198
let end = start + called_script.len();
199+
// TODO: Check if assertion is always true due to code invariants
169200
assert!(
170201
end <= script.capacity(),
171-
"Not enough capacity allocated for compilated script"
202+
"Not enough capacity allocated for compiled script"
172203
);
173204
unsafe {
174205
script.set_len(end);
@@ -184,21 +215,22 @@ impl StructuredScript {
184215
}
185216
}
186217
None => {
187-
// Compile the called_script the first time and add its starting
188-
// position in the compiled script to the cache.
189-
let called_script_start = script.len();
190-
called_script.compile_to_bytes(script, cache);
191-
cache.insert(*id, called_script_start);
218+
tasks.push(Task::UpdateCache {
219+
id,
220+
called_script_start: script.len(),
221+
});
222+
push_script(called_script, &mut tasks);
192223
}
193224
}
194225
}
195-
Block::Script(block_script) => {
196-
let source_script = block_script.as_bytes();
226+
Task::PushRaw(buffer) => {
227+
let source_script = buffer.as_bytes();
197228
let start = script.len();
198229
let end = start + source_script.len();
230+
// TODO: Check if assertion is always true due to code invariants
199231
assert!(
200232
end <= script.capacity(),
201-
"Not enough capacity allocated for compilated script"
233+
"Not enough capacity allocated for compiled script"
202234
);
203235
unsafe {
204236
script.set_len(end);
@@ -209,14 +241,20 @@ impl StructuredScript {
209241
std::ptr::copy_nonoverlapping(src_ptr, dst_ptr, source_script.len());
210242
}
211243
}
244+
Task::UpdateCache {
245+
id,
246+
called_script_start,
247+
} => {
248+
cache.insert(id, called_script_start);
249+
}
212250
}
213251
}
252+
253+
script
214254
}
215255

216256
pub fn compile(self) -> ScriptBuf {
217-
let mut script = Vec::with_capacity(self.size);
218-
let mut cache = HashMap::new();
219-
self.compile_to_bytes(&mut script, &mut cache);
257+
let script = self.compile_to_bytes();
220258
// Ensure that the builder has minimal opcodes:
221259
let script_buf = ScriptBuf::from_bytes(script);
222260
let mut instructions_iter = script_buf.instructions();

0 commit comments

Comments
 (0)