Update BuiltinProgram

solana-labs · Oct 18, 2024 · b3d4cdf · b3d4cdf
1 parent e1a9c32
commit b3d4cdf
Show file tree

Hide file tree

Showing 6 changed files with 68 additions and 22 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_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_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())
             };
             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_function_registry(),
+            self.loader.get_sparse_function_registry(),
         )?;
         Ok(())
     }
@@ -1071,7 +1071,10 @@ impl<C: ContextObject> Executable<C> {
                             .entry(symbol.st_name)
                             .or_insert_with(|| ebpf::hash_symbol_name(name));
                         if config.reject_broken_elfs
-                            && loader.get_function_registry().lookup_by_key(hash).is_none()
+                            && loader
+                                .get_sparse_function_registry()
+                                .lookup_by_key(hash)
+                                .is_none()
                         {
                             return Err(ElfError::UnresolvedSymbol(
                                 String::from_utf8_lossy(name).to_string(),

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_function_registry().lookup_by_key(insn.imm as u32) {
+                    if let Some((_function_name, function)) = self.executable.get_loader().get_sparse_function_registry().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_function_registry().lookup_by_key(insn.imm as u32) {
+                        if let Some((_function_name, function)) = self.executable.get_loader().get_sparse_function_registry().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
@@ -153,7 +153,11 @@ impl<T: Copy + PartialEq> FunctionRegistry<T> {
             } else {
                 ebpf::hash_symbol_name(&usize::from(value).to_le_bytes())
             };
-            if loader.get_function_registry().lookup_by_key(hash).is_some() {
+            if loader
+                .get_sparse_function_registry()
+                .lookup_by_key(hash)
+                .is_some()
+            {
                 return Err(ElfError::SymbolHashCollision(hash));
             }
             hash
@@ -228,13 +232,17 @@ pub type BuiltinFunction<C> = fn(*mut EbpfVm<C>, u64, u64, u64, u64, u64);
 pub struct BuiltinProgram<C: ContextObject> {
     /// Holds the Config if this is a loader program
     config: Option<Box<Config>>,
-    /// Function pointers by symbol
-    functions: FunctionRegistry<BuiltinFunction<C>>,
+    /// Function pointers by symbol with sparse indexing
+    sparse_registry: FunctionRegistry<BuiltinFunction<C>>,
+    /// Function pointers by symbol with dense indexing
+    dense_registry: FunctionRegistry<BuiltinFunction<C>>,
 }
 
 impl<C: ContextObject> PartialEq for BuiltinProgram<C> {
     fn eq(&self, other: &Self) -> bool {
-        self.config.eq(&other.config) && self.functions.eq(&other.functions)
+        self.config.eq(&other.config)
+            && self.sparse_registry.eq(&other.sparse_registry)
+            && self.dense_registry.eq(&other.dense_registry)
     }
 }
 
@@ -243,23 +251,36 @@ impl<C: ContextObject> BuiltinProgram<C> {
     pub fn new_loader(config: Config, functions: FunctionRegistry<BuiltinFunction<C>>) -> Self {
         Self {
             config: Some(Box::new(config)),
-            functions,
+            sparse_registry: functions,
+            dense_registry: FunctionRegistry::default(),
         }
     }
 
     /// Constructs a built-in program
     pub fn new_builtin(functions: FunctionRegistry<BuiltinFunction<C>>) -> Self {
         Self {
             config: None,
-            functions,
+            sparse_registry: functions,
+            dense_registry: FunctionRegistry::default(),
         }
     }
 
     /// Constructs a mock loader built-in program
     pub fn new_mock() -> Self {
         Self {
             config: Some(Box::default()),
-            functions: FunctionRegistry::default(),
+            sparse_registry: FunctionRegistry::default(),
+            dense_registry: FunctionRegistry::default(),
+        }
+    }
+
+    /// Create a new loader with both dense and sparse function registrations
+    /// Use `BuiltinProgram::register_function` to register.
+    pub fn new_loader_with_dense_registration(config: Config) -> Self {
+        Self {
+            config: Some(Box::new(config)),
+            sparse_registry: FunctionRegistry::default(),
+            dense_registry: FunctionRegistry::default(),
         }
     }
 
@@ -269,8 +290,8 @@ impl<C: ContextObject> BuiltinProgram<C> {
     }
 
     /// Get the function registry
-    pub fn get_function_registry(&self) -> &FunctionRegistry<BuiltinFunction<C>> {
-        &self.functions
+    pub fn get_sparse_function_registry(&self) -> &FunctionRegistry<BuiltinFunction<C>> {
+        &self.sparse_registry
     }
 
     /// Calculate memory size
@@ -281,18 +302,40 @@ impl<C: ContextObject> BuiltinProgram<C> {
             } else {
                 0
             })
-            .saturating_add(self.functions.mem_size())
+            .saturating_add(self.sparse_registry.mem_size())
+            .saturating_add(self.dense_registry.mem_size())
+    }
+
+    /// Register a function both in the sparse and dense registries
+    pub fn register_function(
+        &mut self,
+        name: &str,
+        value: BuiltinFunction<C>,
+        dense_key: u32,
+    ) -> Result<(), ElfError> {
+        self.sparse_registry.register_function_hashed(name, value)?;
+        self.dense_registry
+            .register_function(dense_key, name, value)
     }
 }
 
 impl<C: ContextObject> std::fmt::Debug for BuiltinProgram<C> {
     fn fmt(&self, f: &mut std::fmt::Formatter) -> Result<(), std::fmt::Error> {
-        writeln!(f, "{:?}", unsafe {
-            // `derive(Debug)` does not know that `C: ContextObject` does not need to implement `Debug`
-            std::mem::transmute::<&FunctionRegistry<BuiltinFunction<C>>, &FunctionRegistry<usize>>(
-                &self.functions,
-            )
-        })?;
+        unsafe {
+            writeln!(
+                f,
+                "sparse: {:?}\n dense: {:?}",
+                // `derive(Debug)` does not know that `C: ContextObject` does not need to implement `Debug`
+                std::mem::transmute::<
+                    &FunctionRegistry<BuiltinFunction<C>>,
+                    &FunctionRegistry<usize>,
+                >(&self.sparse_registry,),
+                std::mem::transmute::<
+                    &FunctionRegistry<BuiltinFunction<C>>,
+                    &FunctionRegistry<usize>,
+                >(&self.dense_registry,)
+            )?;
+        }
         Ok(())
     }
 }

diff --git a/src/static_analysis.rs b/src/static_analysis.rs
@@ -236,7 +236,7 @@ impl<'a> Analysis<'a> {
                     if let Some((function_name, _function)) = self
                         .executable
                         .get_loader()
-                        .get_function_registry()
+                        .get_sparse_function_registry()
                         .lookup_by_key(insn.imm as u32)
                     {
                         if function_name == b"abort" {