Use a static script_map per thread

Author: lucidLuckylee

Cargo.toml

@@ -39,6 +39,7 @@ repository = "https://github.com/BitVM/rust-bitcoin-script"
 
 [dependencies]
 bitcoin = { git = "https://github.com/rust-bitcoin/rust-bitcoin", branch = "bitvm", features = ["rand-std"]}
+lazy_static = "1.5.0"
 script-macro = { path = "./macro" }
 stdext = "0.3.3"
 

src/analyzer.rs

@@ -1,13 +1,10 @@
-use crate::builder::{Block, StructuredScript};
-use crate::chunker::Chunk;
+use crate::builder::{thread_get_script, Block, StructuredScript};
 use bitcoin::blockdata::opcodes::Opcode;
 use bitcoin::blockdata::script::{read_scriptint, Instruction};
 use bitcoin::opcodes::all::*;
 use bitcoin::script::PushBytes;
-use bitcoin::ScriptBuf;
-use script_macro::script;
 use std::borrow::BorrowMut;
-use std::cmp::{max, min};
+use std::cmp::min;
 use std::panic;
 
 #[derive(Debug, Clone, Default, PartialEq)]
@@ -149,16 +146,13 @@ impl StackAnalyzer {
         for block in builder.blocks.iter() {
             match block {
                 Block::Call(id) => {
-                    let called_script = builder
-                        .script_map
-                        .get(id)
-                        .expect("Missing entry for a called script");
+                    let called_script = thread_get_script(id);
                     match called_script.stack_hint() {
                         Some(stack_hint) => {
                             self.debug_position += called_script.len();
                             self.stack_change(stack_hint)
                         }
-                        None => self.merge_script(called_script),
+                        None => self.merge_script(&called_script),
                     };
                 }
                 Block::Script(block_script) => {

src/builder.rs

@@ -7,11 +7,36 @@ use std::cmp::min;
 use std::collections::HashMap;
 use std::convert::TryFrom;
 use std::hash::{DefaultHasher, Hash, Hasher};
+use std::sync::RwLock;
 use std::{fs::File, io::Write};
 
 use crate::analyzer::{StackAnalyzer, StackStatus};
 use crate::chunker::Chunker;
 
+// One global script map per thread.
+thread_local! {
+    static SCRIPT_MAP: RwLock<HashMap<u64, Box<StructuredScript>>> =
+        RwLock::new(HashMap::new());
+}
+
+pub(crate) fn thread_add_script(id: u64, script: StructuredScript) {
+    SCRIPT_MAP.with(|script_map| {
+        let mut map = script_map.write().unwrap();
+        if !map.contains_key(&id) {
+            map.insert(id, Box::new(script));
+        }
+    });
+}
+
+pub(crate) fn thread_get_script(id: &u64) -> Box<StructuredScript> {
+    SCRIPT_MAP.with(|script_map| {
+        let map = script_map.read().unwrap();
+        map.get(id)
+            .expect("script id not found in SCRIPT_MAP")
+            .clone()
+    })
+}
+
 #[derive(Clone, Debug, Hash)]
 pub enum Block {
     Call(u64),
@@ -36,8 +61,6 @@ pub struct StructuredScript {
     extra_endif_positions: Vec<usize>,
     max_if_interval: (usize, usize),
     pub blocks: Vec<Block>,
-    // TODO: It may be worth to lazy initialize the script_map
-    pub script_map: HashMap<u64, StructuredScript>,
 }
 
 impl Hash for StructuredScript {
@@ -65,7 +88,6 @@ impl StructuredScript {
             extra_endif_positions: vec![],
             max_if_interval: (0, 0),
             blocks,
-            script_map: HashMap::new(),
         }
     }
 
@@ -87,7 +109,9 @@ impl StructuredScript {
         if self.is_script_buf() {
             match &self.blocks[0] {
                 Block::Call(_) => unreachable!(),
-                Block::Script(block_script) => block_script.instructions().collect::<Vec<_>>().len() == 1,
+                Block::Script(block_script) => {
+                    block_script.instructions().collect::<Vec<_>>().len() == 1
+                }
             }
         } else {
             false
@@ -109,10 +133,7 @@ impl StructuredScript {
             assert!(current_pos <= position, "Target position not found");
             match block {
                 Block::Call(id) => {
-                    let called_script = self
-                        .script_map
-                        .get(id)
-                        .expect("Missing entry for a called script");
+                    let called_script = thread_get_script(id);
                     if position >= current_pos && position < current_pos + called_script.len() {
                         return called_script.debug_info(position - current_pos);
                     }
@@ -259,8 +280,8 @@ impl StructuredScript {
         self.num_unclosed_ifs += data.num_unclosed_ifs;
         let id = calculate_hash(&data);
         self.blocks.push(Block::Call(id));
-        // Register script
-        self.script_map.entry(id).or_insert(data);
+        // Register script in the global script map
+        thread_add_script(id, data);
         self
     }
 
@@ -270,10 +291,7 @@ impl StructuredScript {
         for block in self.blocks.as_slice() {
             match block {
                 Block::Call(id) => {
-                    let called_script = self
-                        .script_map
-                        .get(id)
-                        .expect("Missing entry for a called script");
+                    let called_script = thread_get_script(id);
                     // Check if the script with the hash id is in cache
                     match cache.get(id) {
                         Some(called_start) => {
@@ -329,6 +347,7 @@ impl StructuredScript {
     }
 
     pub fn compile(self) -> ScriptBuf {
+        println!("starting compile step");
         let mut script = Vec::with_capacity(self.size);
         let mut cache = HashMap::new();
         self.compile_to_bytes(&mut script, &mut cache);
@@ -348,12 +367,8 @@ impl StructuredScript {
             File::create("analyzed_chunk_stats.txt").expect("Unable to create stats file");
         let chunk_stats = chunker.chunks.iter().map(|chunk| chunk.stats.clone());
         for entry in chunk_stats {
-            writeln!(
-                stats_file,
-                "{:?}",
-                entry.stack_input_size
-            )
-            .expect("Unable to write to stats file");
+            writeln!(stats_file, "{:?}", entry.stack_input_size)
+                .expect("Unable to write to stats file");
         }
         let mut chunk_sizes_iter = chunk_sizes.iter();
         let mut scripts = vec![];
@@ -397,19 +412,19 @@ impl StructuredScript {
             Some(hint) => {
                 hint.stack_changed = changed;
                 hint.deepest_stack_accessed = access;
-            },
-            None => self.stack_hint = Some(StackAnalyzer::plain_stack_status(access, changed))
+            }
+            None => self.stack_hint = Some(StackAnalyzer::plain_stack_status(access, changed)),
         }
         self
     }
-    
+
     pub fn add_altstack_hint(mut self, access: i32, changed: i32) -> Self {
         match &mut self.stack_hint {
             Some(hint) => {
                 hint.altstack_changed = changed;
                 hint.deepest_altstack_accessed = access;
-            },
-            None => self.stack_hint = Some(StackAnalyzer::plain_stack_status(access, changed))
+            }
+            None => self.stack_hint = Some(StackAnalyzer::plain_stack_status(access, changed)),
         }
         self
     }

src/chunker.rs

@@ -1,14 +1,10 @@
 use core::panic;
 
-use bitcoin::{
-    opcodes::all::{OP_ENDIF, OP_IF, OP_NOTIF},
-    script::Instruction,
-    ScriptBuf,
-};
+use bitcoin::ScriptBuf;
 
 use crate::{
-    analyzer::{self, StackStatus},
-    builder::{Block, StructuredScript},
+    analyzer::StackStatus,
+    builder::{thread_get_script, Block, StructuredScript},
     StackAnalyzer,
 };
 
@@ -203,11 +199,11 @@ impl Chunker {
                     );
                 }
             } else {
-                for block in builder.blocks {
+                for block in &builder.blocks {
                     match block {
                         Block::Call(id) => {
-                            let sub_builder = builder.script_map.get(&id).unwrap();
-                            undo_info.call_stack.push(Box::new(sub_builder.clone()));
+                            let sub_builder = thread_get_script(id);
+                            undo_info.call_stack.push(sub_builder);
                         }
                         Block::Script(script_buf) => {
                             // Split the script_buf
@@ -287,8 +283,8 @@ impl Chunker {
                 for block in builder.blocks.iter().rev() {
                     match block {
                         Block::Call(id) => {
-                            let sub_builder = builder.script_map.get(id).unwrap();
-                            self.call_stack.push(Box::new(sub_builder.clone()));
+                            let sub_builder = thread_get_script(id);
+                            self.call_stack.push(sub_builder);
                             contains_call = true;
                         }
                         Block::Script(script_buf) => {
@@ -322,10 +318,10 @@ impl Chunker {
                 .try_into()
                 .expect("Consuming more stack elements than there are on the stack 3"),
             altstack_input_size: input_altstack_size,
-            altstack_output_size: status
-                .altstack_changed
-                .try_into()
-                .expect(&format!("Consuming more stack elements than there are on the altstack: {:?}", status)),
+            altstack_output_size: status.altstack_changed.try_into().expect(&format!(
+                "Consuming more stack elements than there are on the altstack: {:?}",
+                status
+            )),
         };
 
         Chunk::new(chunk_scripts, chunk_len, chunk_stats, last_constant)

tests/test.rs

@@ -234,15 +234,15 @@ fn test_non_optimal_opcodes() {
         OP_DROP
         OP_DROP
 
-        for i in 0..4 {
-            OP_ROLL
-            { i }
-        }
-
-        for i in 0..4 {
-            { i }
-            OP_ROLL
-        }
+        //for i in 0..4 {
+        //    OP_ROLL
+        //    { i }
+        //}
+
+        //for i in 0..4 {
+        //    { i }
+        //    OP_ROLL
+        //}
 
     };