Selectors testing.

compiler/Cargo.toml

@@ -10,4 +10,6 @@ edition = "2018"
 vyper-parser = {path = "../parser", version = "0.1.0"}
 hex = "0.4"
 yultsur = { git = "https://github.com/g-r-a-n-t/yultsur" }
-ethabi = "11.0"
+ethabi = "11.0"
+tiny-keccak = { version = "2.0", features = ["keccak"] }
+stringreader = "0.1"

compiler/src/yul/selectors.rs

@@ -1,6 +1,7 @@
 use crate::errors::CompileError;
 use crate::yul::base;
 use yultsur::yul;
+use tiny_keccak::{Hasher, Keccak};
 
 /// Builds a switch statement from the contract ABI.
 /// The switch's expression is the 4 left-most bytes in the calldata and each case is
@@ -18,7 +19,7 @@ pub fn switch(functions: &Vec<ethabi::Function>) -> Result<yul::Switch, CompileE
 }
 
 /// Builds an expression that loads the first 4 bytes of the calldata.
-fn expression() -> yul::Expression {
+pub fn expression() -> yul::Expression {
     yul::Expression::FunctionCall(yul::FunctionCall {
         identifier: base::untyped_identifier("shr"),
         arguments: vec![
@@ -37,11 +38,9 @@ fn expression() -> yul::Expression {
 ///
 /// Currently, this assumes each input and the single output is 256 bits.
 /// TODO: Handle types of different sizes: https://solidity.readthedocs.io/en/v0.6.2/abi-spec.html#types
-fn case(function: &ethabi::Function) -> Result<yul::Case, CompileError> {
-    let selector = vec![0]; // TODO: Load actual selector with ABI module.
-
+pub fn case(function: &ethabi::Function) -> Result<yul::Case, CompileError> {
     Ok(yul::Case {
-        literal: Some(base::untyped_literal(&format!("0x{}", hex::encode(selector)))),
+        literal: Some(selector_literal(function.signature())),
         block: yul::Block {
             statements: vec![
                 yul::Statement::Expression(yul::Expression::FunctionCall(yul::FunctionCall {
@@ -72,11 +71,74 @@ fn case(function: &ethabi::Function) -> Result<yul::Case, CompileError> {
     })
 }
 
-#[test]
-fn test_expression() {
-    assert_eq!(
-        expression().to_string(),
-        "shr(224, calldataload(0))",
-        "Switch expression not correct."
-    )
+/// Computes the keccak-256 value of the input portion of the function signature and returns the
+/// first 4 bytes.
+///
+/// Example: "foo(uint256):(uint256)" => keccak256("foo(uint256)")
+pub fn selector_literal(sig: String) -> yul::Literal {
+    let mut sig_halves = sig.split(":");
+
+    let mut keccak = Keccak::v256();
+    let mut selector = [0u8; 4];
+
+    if let Some(first_half) = sig_halves.next() {
+        keccak.update(first_half.as_bytes());
+        keccak.finalize(&mut selector);
+    }
+
+    base::untyped_literal(&format!("0x{}", hex::encode(selector)))
 }
+
+#[cfg(test)]
+mod tests {
+    use crate::yul::selectors::{expression, selector_literal, case};
+    use stringreader::StringReader;
+
+    #[test]
+    fn test_expression() {
+        assert_eq!(
+            expression().to_string(),
+            "shr(224, calldataload(0))",
+            "Switch expression not correct."
+        )
+    }
+
+    #[test]
+    fn test_selector_literal_basic() {
+        let json_abi = r#"[{"name": "foo", "type": "function", "inputs": [], "outputs": []}]"#;
+        let abi = ethabi::Contract::load(StringReader::new(json_abi)).expect("Unable to load abi.");
+        let ref foo = abi.functions["foo"][0];
+
+        assert_eq!(
+            selector_literal(foo.signature()).to_string(),
+            String::from("0xc2985578"),
+            "Incorrect selector"
+        )
+    }
+
+    #[test]
+    fn test_selector_literal() {
+        let json_abi = r#"[{"name": "foo", "type": "function", "inputs": [{ "name": "bar", "type": "uint256" }], "outputs": []}]"#;
+        let abi = ethabi::Contract::load(StringReader::new(json_abi)).expect("Unable to load abi.");
+        let ref foo = abi.functions["foo"][0];
+
+        assert_eq!(
+            selector_literal(foo.signature()).to_string(),
+            String::from("0x2fbebd38"),
+            "Incorrect selector"
+        )
+    }
+
+    #[test]
+    fn test_case() {
+        let json_abi = r#"[{"name": "foo", "type": "function", "inputs": [{ "name": "bar", "type": "uint256" }], "outputs": []}]"#;
+        let abi = ethabi::Contract::load(StringReader::new(json_abi)).expect("Unable to load abi.");
+        let ref foo = abi.functions["foo"][0];
+
+        assert_eq!(
+            case(foo).expect("Unable to build case.").to_string(),
+            String::from("case 0x2fbebd38 { mstore(0, foo(calldataload(4))) return(0, 32) }"),
+            "Incorrect case"
+        );
+    }
+}