Merge get_dense and get_sparse registry functions

solana-labs · Oct 21, 2024 · c912e37 · c912e37
1 parent 95a7746
commit c912e37
Show file tree

Hide file tree

Showing 6 changed files with 24 additions and 17 deletions.
diff --git a/src/disassembler.rs b/src/disassembler.rs
@@ -271,7 +271,7 @@ pub fn disassemble_instruction<C: ContextObject>(
                 function_name
             } else {
                 name = "syscall";
-                loader.get_sparse_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())
+                loader.get_function_registry(sbpf_version).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}");
         },

diff --git a/src/elf.rs b/src/elf.rs
@@ -315,7 +315,7 @@ impl<C: ContextObject> Executable<C> {
             self.get_config(),
             self.get_sbpf_version(),
             self.get_function_registry(),
-            self.loader.get_dense_function_registry(),
+            self.loader.get_function_registry(self.get_sbpf_version()),
         )?;
         Ok(())
     }
@@ -1072,7 +1072,7 @@ impl<C: ContextObject> Executable<C> {
                             .or_insert_with(|| ebpf::hash_symbol_name(name));
                         if config.reject_broken_elfs
                             && loader
-                                .get_sparse_function_registry()
+                                .get_function_registry(&SBPFVersion::V1)
                                 .lookup_by_key(hash)
                                 .is_none()
                         {

diff --git a/src/interpreter.rs b/src/interpreter.rs
@@ -536,7 +536,7 @@ impl<'a, 'b, C: ContextObject> Interpreter<'a, 'b, C> {
                 };
 
                 if external {
-                    if let Some((_function_name, function)) = self.executable.get_loader().get_sparse_function_registry().lookup_by_key(insn.imm as u32) {
+                    if let Some((_function_name, function)) = self.executable.get_loader().get_function_registry(self.executable.get_sbpf_version()).lookup_by_key(insn.imm as u32) {
                         resolved = true;
 
                         self.vm.due_insn_count = self.vm.previous_instruction_meter - self.vm.due_insn_count;

diff --git a/src/jit.rs b/src/jit.rs
@@ -717,7 +717,7 @@ impl<'a, C: ContextObject> JitCompiler<'a, C> {
                     };
 
                     if external {
-                        if let Some((_function_name, function)) = self.executable.get_loader().get_sparse_function_registry().lookup_by_key(insn.imm as u32) {
+                        if let Some((_function_name, function)) = self.executable.get_loader().get_function_registry(self.executable.get_sbpf_version()).lookup_by_key(insn.imm as u32) {
                             self.emit_validate_and_profile_instruction_count(false, Some(0));
                             self.emit_ins(X86Instruction::load_immediate(OperandSize::S64, REGISTER_SCRATCH, function as usize as i64));
                             self.emit_ins(X86Instruction::call_immediate(self.relative_to_anchor(ANCHOR_EXTERNAL_FUNCTION_CALL, 5)));

diff --git a/src/program.rs b/src/program.rs
@@ -154,7 +154,7 @@ impl<T: Copy + PartialEq> FunctionRegistry<T> {
                 ebpf::hash_symbol_name(&usize::from(value).to_le_bytes())
             };
             if loader
-                .get_sparse_function_registry()
+                .get_function_registry(&SBPFVersion::V1)
                 .lookup_by_key(hash)
                 .is_some()
             {
@@ -289,14 +289,16 @@ impl<C: ContextObject> BuiltinProgram<C> {
         self.config.as_ref().unwrap()
     }
 
-    /// Get the sparse function registry
-    pub fn get_sparse_function_registry(&self) -> &FunctionRegistry<BuiltinFunction<C>> {
-        &self.sparse_registry
-    }
-
-    /// Get the dense function registry
-    pub fn get_dense_function_registry(&self) -> &FunctionRegistry<BuiltinFunction<C>> {
-        &self.dense_registry
+    /// Get the function registry depending on the SBPF version
+    pub fn get_function_registry(
+        &self,
+        sbpf_version: &SBPFVersion,
+    ) -> &FunctionRegistry<BuiltinFunction<C>> {
+        if sbpf_version == &SBPFVersion::V1 {
+            &self.sparse_registry
+        } else {
+            &self.dense_registry
+        }
     }
 
     /// Calculate memory size

diff --git a/src/static_analysis.rs b/src/static_analysis.rs
@@ -6,6 +6,7 @@ use crate::{
     ebpf,
     elf::Executable,
     error::EbpfError,
+    program::SBPFVersion,
     vm::{ContextObject, DynamicAnalysis, TestContextObject},
 };
 use rustc_demangle::demangle;
@@ -192,7 +193,7 @@ impl<'a> Analysis<'a> {
             dfg_forward_edges: BTreeMap::new(),
             dfg_reverse_edges: BTreeMap::new(),
         };
-        result.split_into_basic_blocks(false);
+        result.split_into_basic_blocks(false, executable.get_sbpf_version());
         result.control_flow_graph_tarjan();
         result.control_flow_graph_dominance_hierarchy();
         result.label_basic_blocks();
@@ -223,7 +224,11 @@ impl<'a> Analysis<'a> {
     /// Splits the sequence of instructions into basic blocks
     ///
     /// Also links the control-flow graph edges between the basic blocks.
-    pub fn split_into_basic_blocks(&mut self, flatten_call_graph: bool) {
+    pub fn split_into_basic_blocks(
+        &mut self,
+        flatten_call_graph: bool,
+        sbpf_version: &SBPFVersion,
+    ) {
         self.cfg_nodes.insert(0, CfgNode::default());
         for pc in self.functions.keys() {
             self.cfg_nodes.entry(*pc).or_default();
@@ -236,7 +241,7 @@ impl<'a> Analysis<'a> {
                     if let Some((function_name, _function)) = self
                         .executable
                         .get_loader()
-                        .get_sparse_function_registry()
+                        .get_function_registry(sbpf_version)
                         .lookup_by_key(insn.imm as u32)
                     {
                         if function_name == b"abort" {