Refactor - FunctionRegistry (#484)

* Turns FunctionRegistry into a struct. * Use FunctionRegistry in BuiltinProgram too.
solana-labs · Jul 13, 2023 · d5525c2 · d5525c2
1 parent ff96604
commit d5525c2
Show file tree

Hide file tree

Showing 23 changed files with 402 additions and 337 deletions.
diff --git a/benches/elf_loader.rs b/benches/elf_loader.rs
@@ -11,20 +11,23 @@ extern crate test;
 extern crate test_utils;
 
 use solana_rbpf::{
-    elf::Executable,
+    elf::{Executable, FunctionRegistry},
     syscalls,
     verifier::TautologyVerifier,
-    vm::{BuiltinProgram, Config, TestContextObject},
+    vm::{BuiltinFunction, BuiltinProgram, Config, TestContextObject},
 };
 use std::{fs::File, io::Read, sync::Arc};
 use test::Bencher;
 
 fn loader() -> Arc<BuiltinProgram<TestContextObject>> {
-    let mut loader = BuiltinProgram::new_loader(Config::default());
-    loader
-        .register_function(b"log", syscalls::bpf_syscall_string)
+    let mut function_registry = FunctionRegistry::<BuiltinFunction<TestContextObject>>::default();
+    function_registry
+        .register_function_hashed(*b"log", syscalls::bpf_syscall_string)
         .unwrap();
-    Arc::new(loader)
+    Arc::new(BuiltinProgram::new_loader(
+        Config::default(),
+        function_registry,
+    ))
 }
 
 #[bench]

diff --git a/benches/jit_compile.rs b/benches/jit_compile.rs
@@ -12,7 +12,7 @@ extern crate test;
 use solana_rbpf::{
     elf::Executable,
     verifier::{RequisiteVerifier, TautologyVerifier},
-    vm::{BuiltinProgram, Config, TestContextObject},
+    vm::{BuiltinProgram, TestContextObject},
 };
 use std::{fs::File, io::Read, sync::Arc};
 use test::Bencher;
@@ -25,7 +25,7 @@ fn bench_init_vm(bencher: &mut Bencher) {
     file.read_to_end(&mut elf).unwrap();
     let executable = Executable::<TautologyVerifier, TestContextObject>::from_elf(
         &elf,
-        Arc::new(BuiltinProgram::new_loader(Config::default())),
+        Arc::new(BuiltinProgram::new_mock()),
     )
     .unwrap();
     let verified_executable =
@@ -52,7 +52,7 @@ fn bench_jit_compile(bencher: &mut Bencher) {
     file.read_to_end(&mut elf).unwrap();
     let executable = Executable::<TautologyVerifier, TestContextObject>::from_elf(
         &elf,
-        Arc::new(BuiltinProgram::new_loader(Config::default())),
+        Arc::new(BuiltinProgram::new_mock()),
     )
     .unwrap();
     let mut verified_executable =

diff --git a/benches/vm_execution.rs b/benches/vm_execution.rs
@@ -11,7 +11,7 @@ extern crate test;
 
 use solana_rbpf::{
     ebpf,
-    elf::Executable,
+    elf::{Executable, FunctionRegistry},
     memory_region::MemoryRegion,
     verifier::{RequisiteVerifier, TautologyVerifier},
     vm::{BuiltinProgram, Config, TestContextObject},
@@ -27,7 +27,7 @@ fn bench_init_interpreter_start(bencher: &mut Bencher) {
     file.read_to_end(&mut elf).unwrap();
     let executable = Executable::<TautologyVerifier, TestContextObject>::from_elf(
         &elf,
-        Arc::new(BuiltinProgram::new_loader(Config::default())),
+        Arc::new(BuiltinProgram::new_mock()),
     )
     .unwrap();
     let verified_executable =
@@ -56,7 +56,7 @@ fn bench_init_jit_start(bencher: &mut Bencher) {
     file.read_to_end(&mut elf).unwrap();
     let executable = Executable::<TautologyVerifier, TestContextObject>::from_elf(
         &elf,
-        Arc::new(BuiltinProgram::new_loader(Config::default())),
+        Arc::new(BuiltinProgram::new_mock()),
     )
     .unwrap();
     let mut verified_executable =
@@ -88,7 +88,10 @@ fn bench_jit_vs_interpreter(
 ) {
     let executable = solana_rbpf::assembler::assemble::<TestContextObject>(
         assembly,
-        Arc::new(BuiltinProgram::new_loader(config)),
+        Arc::new(BuiltinProgram::new_loader(
+            config,
+            FunctionRegistry::default(),
+        )),
     )
     .unwrap();
     let mut verified_executable =

diff --git a/cli/src/main.rs b/cli/src/main.rs
@@ -3,7 +3,7 @@ use solana_rbpf::{
     aligned_memory::AlignedMemory,
     assembler::assemble,
     ebpf,
-    elf::Executable,
+    elf::{Executable, FunctionRegistry},
     memory_region::{MemoryMapping, MemoryRegion},
     static_analysis::Analysis,
     verifier::{RequisiteVerifier, TautologyVerifier},
@@ -92,11 +92,15 @@ fn main() {
         )
         .get_matches();
 
-    let loader = Arc::new(BuiltinProgram::new_loader(Config {
-        enable_instruction_tracing: matches.is_present("trace") || matches.is_present("profile"),
-        enable_symbol_and_section_labels: true,
-        ..Config::default()
-    }));
+    let loader = Arc::new(BuiltinProgram::new_loader(
+        Config {
+            enable_instruction_tracing: matches.is_present("trace")
+                || matches.is_present("profile"),
+            enable_symbol_and_section_labels: true,
+            ..Config::default()
+        },
+        FunctionRegistry::default(),
+    ));
     let executable = match matches.value_of("assembler") {
         Some(asm_file_name) => {
             let mut file = File::open(Path::new(asm_file_name)).unwrap();

diff --git a/examples/disassemble.rs b/examples/disassemble.rs
@@ -6,10 +6,10 @@
 
 extern crate solana_rbpf;
 use solana_rbpf::{
-    elf::{Executable, SBPFVersion},
+    elf::{Executable, FunctionRegistry, SBPFVersion},
     static_analysis::Analysis,
     verifier::TautologyVerifier,
-    vm::{BuiltinProgram, FunctionRegistry, TestContextObject},
+    vm::{BuiltinProgram, TestContextObject},
 };
 use std::sync::Arc;
 
@@ -31,7 +31,7 @@ fn main() {
         0x00, 0x00, 0x00, 0x00, 0xb7, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x95, 0x00, 0x00,
         0x00, 0x00, 0x00, 0x00, 0x00,
     ];
-    let loader = Arc::new(BuiltinProgram::default());
+    let loader = Arc::new(BuiltinProgram::new_mock());
     let executable = Executable::<TautologyVerifier, TestContextObject>::from_text_bytes(
         program,
         loader,

diff --git a/examples/to_json.rs b/examples/to_json.rs
@@ -12,10 +12,10 @@ use std::path::PathBuf;
 
 extern crate solana_rbpf;
 use solana_rbpf::{
-    elf::{Executable, SBPFVersion},
+    elf::{Executable, FunctionRegistry, SBPFVersion},
     static_analysis::Analysis,
     verifier::TautologyVerifier,
-    vm::{BuiltinProgram, FunctionRegistry, TestContextObject},
+    vm::{BuiltinProgram, TestContextObject},
 };
 use std::sync::Arc;
 // Turn a program into a JSON string.
@@ -30,7 +30,7 @@ use std::sync::Arc;
 fn to_json(program: &[u8]) -> String {
     let executable = Executable::<TautologyVerifier, TestContextObject>::from_text_bytes(
         program,
-        Arc::new(BuiltinProgram::default()),
+        Arc::new(BuiltinProgram::new_mock()),
         SBPFVersion::V2,
         FunctionRegistry::default(),
     )

diff --git a/fuzz/fuzz_targets/dumb.rs b/fuzz/fuzz_targets/dumb.rs
@@ -6,10 +6,10 @@ use libfuzzer_sys::fuzz_target;
 
 use solana_rbpf::{
     ebpf,
-    elf::{Executable, SBPFVersion},
+    elf::{Executable, FunctionRegistry, SBPFVersion},
     memory_region::MemoryRegion,
     verifier::{RequisiteVerifier, TautologyVerifier, Verifier},
-    vm::{BuiltinProgram, FunctionRegistry, TestContextObject},
+    vm::{BuiltinProgram, TestContextObject},
 };
 use test_utils::create_vm;
 
@@ -35,7 +35,7 @@ fuzz_target!(|data: DumbFuzzData| {
     let mut mem = data.mem;
     let executable = Executable::<TautologyVerifier, TestContextObject>::from_text_bytes(
         &prog,
-        std::sync::Arc::new(BuiltinProgram::new_loader(config)),
+        std::sync::Arc::new(BuiltinProgram::new_loader(config, FunctionRegistry::default())),
         SBPFVersion::V2,
         function_registry,
     )

diff --git a/fuzz/fuzz_targets/smart.rs b/fuzz/fuzz_targets/smart.rs
@@ -7,11 +7,11 @@ use libfuzzer_sys::fuzz_target;
 use grammar_aware::*;
 use solana_rbpf::{
     ebpf,
-    elf::{Executable, SBPFVersion},
+    elf::{Executable, FunctionRegistry, SBPFVersion},
     insn_builder::{Arch, IntoBytes},
     memory_region::MemoryRegion,
     verifier::{RequisiteVerifier, TautologyVerifier, Verifier},
-    vm::{BuiltinProgram, FunctionRegistry, TestContextObject},
+    vm::{BuiltinProgram, TestContextObject},
 };
 use test_utils::create_vm;
 
@@ -39,7 +39,7 @@ fuzz_target!(|data: FuzzData| {
     let mut mem = data.mem;
     let executable = Executable::<TautologyVerifier, TestContextObject>::from_text_bytes(
         prog.into_bytes(),
-        std::sync::Arc::new(BuiltinProgram::new_loader(config)),
+        std::sync::Arc::new(BuiltinProgram::new_loader(config, FunctionRegistry::default())),
         SBPFVersion::V2,
         function_registry,
     )

diff --git a/fuzz/fuzz_targets/smart_jit_diff.rs b/fuzz/fuzz_targets/smart_jit_diff.rs
@@ -5,11 +5,11 @@ use libfuzzer_sys::fuzz_target;
 use grammar_aware::*;
 use solana_rbpf::{
     ebpf,
-    elf::{Executable, SBPFVersion},
+    elf::{Executable, FunctionRegistry, SBPFVersion},
     insn_builder::{Arch, Instruction, IntoBytes},
     memory_region::MemoryRegion,
     verifier::{RequisiteVerifier, TautologyVerifier, Verifier},
-    vm::{BuiltinProgram, FunctionRegistry, TestContextObject},
+    vm::{BuiltinProgram, TestContextObject},
 };
 use test_utils::create_vm;
 
@@ -47,7 +47,7 @@ fuzz_target!(|data: FuzzData| {
     let mut jit_mem = data.mem;
     let mut executable = Executable::<TautologyVerifier, TestContextObject>::from_text_bytes(
         prog.into_bytes(),
-        std::sync::Arc::new(BuiltinProgram::new_loader(config)),
+        std::sync::Arc::new(BuiltinProgram::new_loader(config, FunctionRegistry::default())),
         SBPFVersion::V2,
         function_registry,
     )

diff --git a/fuzz/fuzz_targets/smarter_jit_diff.rs b/fuzz/fuzz_targets/smarter_jit_diff.rs
@@ -5,13 +5,13 @@ use libfuzzer_sys::fuzz_target;
 use semantic_aware::*;
 use solana_rbpf::{
     ebpf,
-    elf::{Executable, SBPFVersion},
+    elf::{Executable, FunctionRegistry, SBPFVersion},
     insn_builder::IntoBytes,
     memory_region::MemoryRegion,
     static_analysis::Analysis,
     verifier::{RequisiteVerifier, TautologyVerifier, Verifier},
     vm::{
-        BuiltinProgram, ContextObject, FunctionRegistry, TestContextObject,
+        BuiltinProgram, ContextObject, TestContextObject,
     },
 };
 use test_utils::create_vm;
@@ -46,7 +46,7 @@ fuzz_target!(|data: FuzzData| {
     let mut jit_mem = data.mem;
     let mut executable = Executable::<TautologyVerifier, TestContextObject>::from_text_bytes(
         prog.into_bytes(),
-        std::sync::Arc::new(BuiltinProgram::new_loader(config)),
+        std::sync::Arc::new(BuiltinProgram::new_loader(config, FunctionRegistry::default())),
         SBPFVersion::V2,
         function_registry,
     )

diff --git a/fuzz/fuzz_targets/verify_semantic_aware.rs b/fuzz/fuzz_targets/verify_semantic_aware.rs
@@ -4,10 +4,9 @@ use libfuzzer_sys::fuzz_target;
 
 use semantic_aware::*;
 use solana_rbpf::{
-    elf::SBPFVersion,
+    elf::{FunctionRegistry, SBPFVersion},
     insn_builder::IntoBytes,
     verifier::{RequisiteVerifier, Verifier},
-    vm::FunctionRegistry,
 };
 
 use crate::common::ConfigTemplate;

diff --git a/src/assembler.rs b/src/assembler.rs
@@ -18,9 +18,9 @@ use crate::{
         Statement,
     },
     ebpf::{self, Insn},
-    elf::{register_internal_function, Executable, SBPFVersion},
+    elf::{Executable, FunctionRegistry, SBPFVersion},
     verifier::TautologyVerifier,
-    vm::{BuiltinProgram, ContextObject, FunctionRegistry},
+    vm::{BuiltinProgram, ContextObject},
 };
 use std::{collections::HashMap, sync::Arc};
 
@@ -191,7 +191,7 @@ fn insn(opc: u8, dst: i64, src: i64, off: i64, imm: i64) -> Result<Insn, String>
 ///     be16 r0
 ///     neg64 r2
 ///     exit",
-///     std::sync::Arc::new(BuiltinProgram::new_loader(Config::default())),
+///     std::sync::Arc::new(BuiltinProgram::new_mock()),
 /// ).unwrap();
 /// let program = executable.get_text_bytes().1;
 /// println!("{:?}", program);
@@ -244,14 +244,9 @@ pub fn assemble<C: ContextObject>(
         match statement {
             Statement::Label { name } => {
                 if name.starts_with("function_") || name == "entrypoint" {
-                    register_internal_function(
-                        &mut function_registry,
-                        &loader,
-                        &SBPFVersion::V2,
-                        insn_ptr,
-                        name.as_bytes(),
-                    )
-                    .map_err(|_| format!("Label hash collision {name}"))?;
+                    function_registry
+                        .register_function(insn_ptr as u32, name.as_bytes().to_vec(), insn_ptr)
+                        .map_err(|_| format!("Label hash collision {name}"))?;
                 }
                 labels.insert(name.as_str(), insn_ptr);
             }
@@ -298,14 +293,13 @@ pub fn assemble<C: ContextObject>(
                         (CallImm, [Integer(imm)]) => {
                             let target_pc = *imm + insn_ptr as i64 + 1;
                             let label = format!("function_{}", target_pc as usize);
-                            register_internal_function(
-                                &mut function_registry,
-                                &loader,
-                                &SBPFVersion::V2,
-                                target_pc as usize,
-                                label.as_bytes(),
-                            )
-                            .map_err(|_| format!("Label hash collision {name}"))?;
+                            function_registry
+                                .register_function(
+                                    target_pc as u32,
+                                    label.as_bytes().to_vec(),
+                                    target_pc as usize,
+                                )
+                                .map_err(|_| format!("Label hash collision {name}"))?;
                             insn(opc, 0, 1, 0, target_pc)
                         }
                         (CallReg, [Register(dst)]) => insn(opc, 0, 0, 0, *dst),

diff --git a/src/disassembler.rs b/src/disassembler.rs
@@ -10,9 +10,9 @@
 
 use crate::{
     ebpf,
-    elf::SBPFVersion,
+    elf::{FunctionRegistry, SBPFVersion},
     static_analysis::CfgNode,
-    vm::{BuiltinProgram, ContextObject, FunctionRegistry},
+    vm::{BuiltinProgram, ContextObject},
 };
 use std::collections::BTreeMap;
 
@@ -123,7 +123,7 @@ fn jmp_reg_str(name: &str, insn: &ebpf::Insn, cfg_nodes: &BTreeMap<usize, CfgNod
 pub fn disassemble_instruction<C: ContextObject>(
     insn: &ebpf::Insn, 
     cfg_nodes: &BTreeMap<usize, CfgNode>,
-    function_registry: &FunctionRegistry,
+    function_registry: &FunctionRegistry<usize>,
     loader: &BuiltinProgram<C>,
     sbpf_version: &SBPFVersion,
 ) -> String {
@@ -254,17 +254,17 @@ pub fn disassemble_instruction<C: ContextObject>(
             let mut function_name = None;
             if sbpf_version.static_syscalls() {
                 if insn.src != 0 {
-                    function_name = Some(resolve_label(cfg_nodes, insn.imm as usize));
+                    function_name = Some(resolve_label(cfg_nodes, insn.imm as usize).to_string());
                 }
             } else {
-                function_name = function_registry.get(&(insn.imm as u32)).map(|(_, function_name)| function_name.as_str());
+                function_name = function_registry.lookup_by_key(insn.imm as u32).map(|(function_name, _)| String::from_utf8_lossy(function_name).to_string());
             }
             let function_name = if let Some(function_name) = function_name {
                 name = "call";
                 function_name
             } else {
                 name = "syscall";
-                loader.lookup_function(insn.imm as u32).map(|(function_name, _)| std::str::from_utf8(function_name).unwrap()).unwrap_or("[invalid]")
+                loader.get_function_registry().lookup_by_key(insn.imm as u32).map(|(function_name, _)| String::from_utf8_lossy(function_name).to_string()).unwrap_or_else(|| "[invalid]".to_string())
             };
             desc = format!("{name} {function_name}");
         },