refactor: Port ContractsIdentifier to Rust

crates/edr_napi/index.d.ts

@@ -347,7 +347,6 @@ export interface SubscriptionEvent {
   result: any
 }
 export function createModelsAndDecodeBytecodes(solcVersion: string, compilerInput: any, compilerOutput: any): Array<Bytecode>
-export function getLibraryAddressPositions(bytecodeOutput: any): Array<number>
 export function linkHexStringBytecode(code: string, address: string, position: number): string
 export const enum ContractFunctionType {
   CONSTRUCTOR = 0,
@@ -706,8 +705,6 @@ export const enum Opcode {
 export function opcodeToString(opcode: Opcode): string
 export function isPush(opcode: Opcode): boolean
 export function isJump(opcode: Opcode): boolean
-export function getPushLength(opcode: Opcode): number
-export function getOpcodeLength(opcode: Opcode): number
 export function isCall(opcode: Opcode): boolean
 export function isCreate(opcode: Opcode): boolean
 export interface TracingMessage {
@@ -841,6 +838,11 @@ export class Contract {
   get receive(): ContractFunction | undefined
   getFunctionFromSelector(selector: Uint8Array): ContractFunction | undefined
 }
+export class ContractsIdentifier {
+  constructor(enableCache?: boolean | undefined | null)
+  addBytecode(bytecode: Bytecode): void
+  getBytecodeForCall(code: Uint8Array, isCreate: boolean): Bytecode | undefined
+}
 export class Exit {
   get kind(): ExitCode
   isError(): boolean

crates/edr_napi/index.js

@@ -310,7 +310,7 @@ if (!nativeBinding) {
   throw new Error(`Failed to load native binding`)
 }
 
-const { SpecId, EdrContext, MineOrdering, Provider, Response, SuccessReason, ExceptionalHalt, createModelsAndDecodeBytecodes, getLibraryAddressPositions, linkHexStringBytecode, SourceFile, SourceLocation, ContractFunctionType, ContractFunctionVisibility, ContractFunction, CustomError, Instruction, JumpType, jumpTypeToString, Bytecode, ContractType, Contract, Exit, ExitCode, Opcode, opcodeToString, isPush, isJump, getPushLength, getOpcodeLength, isCall, isCreate, VmTracer, RawTrace } = nativeBinding
+const { SpecId, EdrContext, MineOrdering, Provider, Response, SuccessReason, ExceptionalHalt, createModelsAndDecodeBytecodes, linkHexStringBytecode, SourceFile, SourceLocation, ContractFunctionType, ContractFunctionVisibility, ContractFunction, CustomError, Instruction, JumpType, jumpTypeToString, Bytecode, ContractType, Contract, ContractsIdentifier, Exit, ExitCode, Opcode, opcodeToString, isPush, isJump, isCall, isCreate, VmTracer, RawTrace } = nativeBinding
 
 module.exports.SpecId = SpecId
 module.exports.EdrContext = EdrContext
@@ -320,7 +320,6 @@ module.exports.Response = Response
 module.exports.SuccessReason = SuccessReason
 module.exports.ExceptionalHalt = ExceptionalHalt
 module.exports.createModelsAndDecodeBytecodes = createModelsAndDecodeBytecodes
-module.exports.getLibraryAddressPositions = getLibraryAddressPositions
 module.exports.linkHexStringBytecode = linkHexStringBytecode
 module.exports.SourceFile = SourceFile
 module.exports.SourceLocation = SourceLocation
@@ -334,14 +333,13 @@ module.exports.jumpTypeToString = jumpTypeToString
 module.exports.Bytecode = Bytecode
 module.exports.ContractType = ContractType
 module.exports.Contract = Contract
+module.exports.ContractsIdentifier = ContractsIdentifier
 module.exports.Exit = Exit
 module.exports.ExitCode = ExitCode
 module.exports.Opcode = Opcode
 module.exports.opcodeToString = opcodeToString
 module.exports.isPush = isPush
 module.exports.isJump = isJump
-module.exports.getPushLength = getPushLength
-module.exports.getOpcodeLength = getOpcodeLength
 module.exports.isCall = isCall
 module.exports.isCreate = isCreate
 module.exports.VmTracer = VmTracer

crates/edr_napi/src/trace.rs

@@ -21,6 +21,7 @@ mod library_utils;
 mod model;
 mod source_map;
 
+mod contracts_identifier;
 mod exit;
 mod message_trace;
 mod opcodes;

crates/edr_napi/src/trace/contracts_identifier.rs

@@ -0,0 +1,311 @@
+use std::{collections::HashMap, rc::Rc};
+
+use edr_eth::Address;
+use edr_evm::hex;
+use napi::{
+    bindgen_prelude::{ClassInstance, Uint8Array, Undefined},
+    Either, Env, JsObject,
+};
+use napi_derive::napi;
+
+use super::{
+    model::{Bytecode, ImmutableReference},
+    opcodes::Opcode,
+};
+use crate::utils::ClassInstanceRef;
+
+/// This class represent a somewhat special Trie of bytecodes.
+///
+/// What makes it special is that every node has a set of all of its descendants
+/// and its depth.
+#[derive(Clone)]
+pub struct BytecodeTrie {
+    child_nodes: HashMap<u8, Box<BytecodeTrie>>,
+    descendants: Vec<Rc<ClassInstanceRef<Bytecode>>>,
+    match_: Option<Rc<ClassInstanceRef<Bytecode>>>,
+    depth: Option<u32>,
+}
+
+impl BytecodeTrie {
+    pub fn new(depth: Option<u32>) -> BytecodeTrie {
+        BytecodeTrie {
+            child_nodes: HashMap::new(),
+            descendants: Vec::new(),
+            match_: None,
+            depth,
+        }
+    }
+
+    pub fn add(&mut self, bytecode: ClassInstance<Bytecode>, env: Env) -> napi::Result<()> {
+        let bytecode = Rc::new(ClassInstanceRef::from_obj(bytecode, env)?);
+
+        let mut cursor = self;
+
+        let bytecode_normalized_code = &bytecode.borrow(env)?.normalized_code;
+        for (index, byte) in bytecode_normalized_code.iter().enumerate() {
+            cursor.descendants.push(bytecode.clone());
+
+            let node = cursor
+                .child_nodes
+                .entry(*byte)
+                .or_insert_with(|| Box::new(BytecodeTrie::new(Some(index as u32))));
+
+            cursor = node;
+        }
+
+        // If multiple contracts with the exact same bytecode are added we keep the last
+        // of them. Note that this includes the metadata hash, so the chances of
+        // happening are pretty remote, except in super artificial cases that we
+        // have in our test suite.
+        cursor.match_ = Some(bytecode.clone());
+
+        Ok(())
+    }
+
+    /// Searches for a bytecode. If it's an exact match, it is returned. If
+    /// there's no match, but a prefix of the code is found in the trie, the
+    /// node of the longest prefix is returned. If the entire code is
+    /// covered by the trie, and there's no match, we return undefined.
+    pub fn search(
+        &self,
+        code: &Uint8Array,
+        current_code_byte: u32,
+    ) -> Option<Either<Rc<ClassInstanceRef<Bytecode>>, &Self>> {
+        if current_code_byte > code.len() as u32 {
+            return None;
+        }
+
+        let mut cursor = self;
+
+        for byte in code.iter().skip(current_code_byte as usize) {
+            let child_node = cursor.child_nodes.get(byte);
+
+            if let Some(node) = child_node {
+                cursor = node;
+            } else {
+                return Some(Either::B(cursor));
+            }
+        }
+
+        cursor
+            .match_
+            .as_ref()
+            .map(|bytecode| Either::A(bytecode.clone()))
+    }
+}
+
+/// Returns true if the lastByte is placed right when the metadata starts or
+/// after it.
+fn is_matching_metadata(code: Uint8Array, last_byte: u32) -> bool {
+    let mut byte = 0;
+    while byte < last_byte {
+        let opcode = Opcode::from_repr(code[byte as usize]).unwrap();
+        let next = code
+            .get(byte as usize + 1)
+            .and_then(|x| Opcode::from_repr(*x));
+
+        if opcode == Opcode::REVERT && next == Some(Opcode::INVALID) {
+            return true;
+        }
+
+        byte += u32::from(opcode.len());
+    }
+
+    false
+}
+
+#[napi]
+pub struct ContractsIdentifier {
+    trie: BytecodeTrie,
+    cache: HashMap<String, Rc<ClassInstanceRef<Bytecode>>>,
+    enable_cache: bool,
+}
+
+#[napi]
+impl ContractsIdentifier {
+    #[napi(constructor)]
+    pub fn new(enable_cache: Option<bool>) -> ContractsIdentifier {
+        let enable_cache = enable_cache.unwrap_or(true);
+
+        ContractsIdentifier {
+            trie: BytecodeTrie::new(None),
+            cache: HashMap::new(),
+            enable_cache,
+        }
+    }
+
+    #[napi]
+    pub fn add_bytecode(
+        &mut self,
+        bytecode: ClassInstance<Bytecode>,
+        env: Env,
+    ) -> napi::Result<()> {
+        self.trie.add(bytecode, env)?;
+        self.cache.clear();
+
+        Ok(())
+    }
+
+    fn search_bytecode(
+        &mut self,
+        is_create: bool,
+        code: Uint8Array,
+        normalize_libraries: Option<bool>,
+        trie: Option<&BytecodeTrie>,
+        first_byte_to_search: Option<u32>,
+        env: Env,
+    ) -> napi::Result<Option<Rc<ClassInstanceRef<Bytecode>>>> {
+        let normalize_libraries = normalize_libraries.unwrap_or(true);
+        let first_byte_to_search = first_byte_to_search.unwrap_or(0);
+        let trie = trie.unwrap_or(&self.trie);
+
+        Self::search_bytecode_inner(
+            is_create,
+            code,
+            normalize_libraries,
+            trie,
+            first_byte_to_search,
+            env,
+        )
+    }
+
+    fn search_bytecode_inner(
+        is_create: bool,
+        code: Uint8Array,
+        normalize_libraries: bool,
+        trie: &BytecodeTrie,
+        first_byte_to_search: u32,
+        env: Env,
+    ) -> napi::Result<Option<Rc<ClassInstanceRef<Bytecode>>>> {
+        let search_result = match trie.search(&code, first_byte_to_search) {
+            None => return Ok(None),
+            Some(Either::A(bytecode)) => return Ok(Some(bytecode.clone())),
+            Some(Either::B(trie)) => trie,
+        };
+
+        // Deployment messages have their abi-encoded arguments at the end of the
+        // bytecode.
+        //
+        // We don't know how long those arguments are, as we don't know which contract
+        // is being deployed, hence we don't know the signature of its
+        // constructor.
+        //
+        // To make things even harder, we can't trust that the user actually passed the
+        // right amount of arguments.
+        //
+        // Luckily, the chances of a complete deployment bytecode being the prefix of
+        // another one are remote. For example, most of the time it ends with
+        // its metadata hash, which will differ.
+        //
+        // We take advantage of this last observation, and just return the bytecode that
+        // exactly matched the search_result (sub)trie that we got.
+        match &search_result.match_ {
+            Some(bytecode) if is_create && bytecode.borrow(env)?.is_deployment => {
+                return Ok(Some(bytecode.clone()));
+            }
+            _ => {}
+        };
+
+        if normalize_libraries {
+            for bytecode_with_libraries in &search_result.descendants {
+                let bytecode_with_libraries = bytecode_with_libraries.borrow(env)?;
+
+                if bytecode_with_libraries.library_address_positions.is_empty()
+                    && bytecode_with_libraries.immutable_references.is_empty()
+                {
+                    continue;
+                }
+
+                let mut normalized_code = code.clone();
+                // zero out addresses
+                for pos in &bytecode_with_libraries.library_address_positions {
+                    normalized_code[*pos as usize..][..Address::len_bytes()].fill(0);
+                }
+                // zero out slices
+                for ImmutableReference { start, length } in
+                    &bytecode_with_libraries.immutable_references
+                {
+                    normalized_code[*start as usize..][..*length as usize].fill(0);
+                }
+
+                let normalized_result = Self::search_bytecode_inner(
+                    is_create,
+                    normalized_code,
+                    false,
+                    search_result,
+                    search_result.depth.map_or(0, |depth| depth + 1),
+                    env,
+                );
+
+                if let Ok(Some(bytecode)) = normalized_result {
+                    return Ok(Some(bytecode.clone()));
+                };
+            }
+        }
+
+        // If we got here we may still have the contract, but with a different metadata
+        // hash.
+        //
+        // We check if we got to match the entire executable bytecode, and are just
+        // stuck because of the metadata. If that's the case, we can assume that
+        // any descendant will be a valid Bytecode, so we just choose the most
+        // recently added one.
+        //
+        // The reason this works is because there's no chance that Solidity includes an
+        // entire bytecode (i.e. with metadata), as a prefix of another one.
+        if let Some(search_depth) = search_result.depth {
+            if is_matching_metadata(code, search_depth) && !search_result.descendants.is_empty() {
+                return Ok(Some(
+                    search_result.descendants[search_result.descendants.len() - 1].clone(),
+                ));
+            }
+        }
+
+        Ok(None)
+    }
+
+    #[napi(ts_return_type = "Bytecode | undefined")]
+    pub fn get_bytecode_for_call(
+        &mut self,
+        code: Uint8Array,
+        is_create: bool,
+        env: Env,
+    ) -> napi::Result<Either<JsObject, Undefined>> {
+        let mut normalized_code = code.clone();
+        normalize_library_runtime_bytecode_if_necessary(&mut normalized_code);
+
+        let normalized_code_hex = hex::encode(normalized_code.as_ref());
+        if self.enable_cache {
+            let cached = self.cache.get(&normalized_code_hex);
+
+            if let Some(cached) = cached {
+                return Ok(Either::A(cached.as_object(env)?));
+            }
+        }
+
+        let result = self.search_bytecode(is_create, normalized_code, None, None, None, env)?;
+
+        if self.enable_cache {
+            if let Some(result) = &result {
+                self.cache.insert(normalized_code_hex, result.clone());
+            }
+        }
+
+        match result {
+            Some(bytecode) => Ok(Either::A(bytecode.as_object(env)?)),
+            None => Ok(Either::B(())),
+        }
+    }
+}
+
+fn normalize_library_runtime_bytecode_if_necessary(code: &mut Uint8Array) {
+    // Libraries' protection normalization:
+    // Solidity 0.4.20 introduced a protection to prevent libraries from being
+    // called directly. This is done by modifying the code on deployment, and
+    // hard-coding the contract address. The first instruction is a PUSH20 of
+    // the address, which we zero-out as a way of normalizing it. Note that it's
+    // also zeroed-out in the compiler output.
+    if code.first().copied() == Some(Opcode::PUSH20 as u8) {
+        code[1..][..Address::len_bytes()].fill(0);
+    }
+}

crates/edr_napi/src/trace/library_utils.rs

@@ -1,17 +1,9 @@
 //! Port of the hardhat-network's `library-utils.ts` to Rust.
 
 use edr_evm::hex;
-use edr_solidity::{artifacts::CompilerOutputBytecode, library_utils};
 use napi::bindgen_prelude::Buffer;
 use napi_derive::napi;
 
-#[napi]
-pub fn get_library_address_positions(bytecode_output: serde_json::Value) -> Vec<u32> {
-    let bytecode_output: CompilerOutputBytecode = serde_json::from_value(bytecode_output).unwrap();
-
-    library_utils::get_library_address_positions(&bytecode_output)
-}
-
 /// Normalizes the compiler output bytecode by replacing the library addresses
 /// with zeros.
 pub fn normalize_compiler_output_bytecode(