Rust NIFs

Introduction

Elixir and Erlang are not the most performant options for number-intensive calculations. Fortunately, they provide NIFs (Native Implemented Functions) for delegating such tasks to languages like C and Rust that good at number crunching.

Recently, I had an opportunity to work with / write Rust NIFs. Surprisingly, it’s not as hard as I thought it would be. The only skill you need to do that is the basic knowledge of Rust. In this post, I’ll explain simple steps for writing Rust NIFs and I’ll describe a couple of real-world NIFs as examples.

Steps

1. Add rustler to your project.

2. Run mix rustler.new to generate a rustler project.

3. Create a wrapper function in the elixir module created in the previous step. For example:

  def my_wrapper_function(_data), do: :erlang.nif_error(:nif_not_loaded)

1. Select a scheduler type and create the implementaion for my_wrapper_function:

rustler::rustler_export_nifs! {
    "Elixir.MyModule",
    [
        ("my_wrapper_function", 1, my_wrapper_function, rustler::SchedulerFlags::DirtyCpu)
    ],
    None
}


fn my_wrapper_function<'a>(env: Env<'a>, args: &[Term<'a>]) -> Term<'a> {

    ...

    atoms::ok()
}

Examples

Let’s see a couple of examples.

ex_keccak

ex_keccak is a NIF that wraps the KECCAK-256 function from the tiny-keccak Rust library:

{
  :ok,
  <<28, 138, 255, 149, 6, 133, 194, 237, 75, 195, 23, 79, 52, 114, 40, 123, 86, 217, 81, 123, 156, 148, 129, 39, 49, 154, 9, 167, 163, 109, 234, 200>>
} = ExKeccak.hash_256("hello")

It implements a single function wrapper for hash:

defmodule ExKeccak do
  use Rustler, otp_app: :ex_keccak, crate: :exkeccak

  def hash_256(_data), do: :erlang.nif_error(:nif_not_loaded)
end

The rust module calls a function in the tiny-keccak library:

use rustler::types::binary::{Binary, OwnedBinary};
use rustler::types::Encoder;
use rustler::{Env, Term};
use tiny_keccak::{Hasher, Keccak};

mod atoms {
    rustler::rustler_atoms! {
        atom ok;
        atom error;
        atom invalid_type;
    }
}

rustler::rustler_export_nifs! {
    "Elixir.ExKeccak",
    [
        ("hash_256", 1, hash_256, rustler::SchedulerFlags::DirtyCpu)
    ],
    None
}

fn hash_256<'a>(env: Env<'a>, args: &[Term<'a>]) -> Term<'a> {
    let data: Binary = match args[0].decode() {
        Ok(binary) => binary,
        Err(_error) => return (atoms::error(), atoms::invalid_type()).encode(env),
    };
    let data_slice = data.as_slice();

    let mut keccak = Keccak::v256();

    keccak.update(data_slice);

    let mut result: [u8; 32] = [0; 32];
    keccak.finalize(&mut result);

    let mut erl_bin: OwnedBinary = OwnedBinary::new(32).unwrap();
    erl_bin.as_mut_slice().copy_from_slice(&result);

    (atoms::ok(), erl_bin.release(env)).encode(env)
}

atoms defines possible atoms that can be returned. hash_256 looks more complex than our initial example because we have to convert data to the format accepted by Keccak::v256 and convert it again to return the result.

ex_secp256k1

ex_secp256k1 wraps a couple functions from the libsecp256k1 Rust library:

message =
  <<0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
  0, 0, 0, 0, 0, 0, 2>>

private_key =
  <<120, 128, 174, 201, 52, 19, 241, 23, 239, 20, 189, 78, 109, 19, 8, 117,
  171, 44, 125, 125, 85, 160, 100, 250, 195, 194, 247, 189, 81, 81, 99, 128>>

{:ok {r_binary, s_binary, recovery_id_int}} = ExSecp256k1.sign(message, private_key)

{:ok, _binary} = ExSecp256k1.create_public_key(<<120, 128, 174, 201, 52, 19,
  241, 23, 239, 20, 189, 78, 109, 19, 8, 117, 171, 44, 125, 125, 85, 160, 100,
  250, 195, 194, 247, 189, 81, 81, 99, 128>>)

I think you see a pattern, all functions that will be called in rust should be defined in the elixir wrapper module:

defmodule ExSecp256k1 do
  use Rustler, otp_app: :ex_secp256k1, crate: "exsecp256k1"

  def sign(_message, _private_key), do: :erlang.nif_error(:nif_not_loaded)

  def sign_compact(_message, _private_key), do: :erlang.nif_error(:nif_not_loaded)

  def recover(_hash, _r, _s, _recovery_id), do: :erlang.nif_error(:nif_not_loaded)

  def recover_compact(_hash, _signature, _recovery_id), do: :erlang.nif_error(:nif_not_loaded)

  def create_public_key(_private_key), do: :erlang.nif_error(:nif_not_loaded)
end

Let’s see one function from the rust module:

fn sign<'a>(env: Env<'a>, args: &[Term<'a>]) -> Term<'a> {
    let (Signature { s, r }, recid) = match secp256k1_sign(env, args) {
        Ok(result) => result,
        Err(error) => return error,
    };

    let mut r_bin: OwnedBinary = OwnedBinary::new(32).unwrap();
    let mut s_bin: OwnedBinary = OwnedBinary::new(32).unwrap();

    r_bin.as_mut_slice().copy_from_slice(&r.b32());
    s_bin.as_mut_slice().copy_from_slice(&s.b32());
    let recid_u8: u8 = recid.into();

    (
        atoms::ok(),
        (r_bin.release(env), s_bin.release(env), recid_u8),
    )
        .encode(env)
}

fn secp256k1_sign<'a>(
    env: Env<'a>,
    args: &[Term<'a>],
) -> Result<(Signature, RecoveryId), Term<'a>> {
    let message_bin: Binary = match args[0].decode() {
        Ok(binary) => binary,
        Err(_error) => return Err((atoms::error(), atoms::message_not_binary()).encode(env)),
    };

    let private_key_bin: Binary = match args[1].decode() {
        Ok(binary) => binary,
        Err(_error) => return Err((atoms::error(), atoms::private_key_not_binary()).encode(env)),
    };

    if message_bin.len() != 32 {
        return Err((atoms::error(), atoms::wrong_message_size()).encode(env));
    }

    if private_key_bin.len() != 32 {
        return Err((atoms::error(), atoms::wrong_private_key_size()).encode(env));
    }

    let mut private_key_fixed: [u8; 32] = [0; 32];
    private_key_fixed.copy_from_slice(&private_key_bin.as_slice()[..32]);

    let mut message_fixed: [u8; 32] = [0; 32];
    message_fixed.copy_from_slice(&message_bin.as_slice()[..32]);

    let private_key = SecretKey::parse(&private_key_fixed).unwrap();
    let message = Message::parse(&message_fixed);

    Ok(secp256k1::sign(&message, &private_key))
}

The code here is just checking of various conditions before passing parameters to the secp256k1::sign and converting types.

Conclusion

As you can see rust nifs are easy to write. All you need is the basic knowledge of Rust.