Skip to content

Commit

Permalink
Implement Decoder/Encoder support for i/u128
Browse files Browse the repository at this point in the history
  • Loading branch information
filmor committed Mar 12, 2024
1 parent f944588 commit d7ded73
Show file tree
Hide file tree
Showing 6 changed files with 176 additions and 0 deletions.
130 changes: 130 additions & 0 deletions rustler/src/types/i128.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,130 @@
use crate::{Decoder, Encoder, Env, Error, NifResult, Term};
use std::convert::TryFrom;

const EXTERNAL_TERM_FORMAT_VERSION: u8 = 131;
const SMALL_BIG_EXT: u8 = 110;

impl Encoder for i128 {
fn encode<'a>(&self, env: Env<'a>) -> Term<'a> {
if let Ok(int) = i64::try_from(*self) {
int.encode(env)
} else {
let mut etf = [0u8; 4 + 16];
etf[0] = EXTERNAL_TERM_FORMAT_VERSION;
etf[1] = SMALL_BIG_EXT;
etf[2] = 16; // length in bytes
if *self < 0 {
etf[3] = 1;
let bytes = (-self).to_le_bytes();
etf[4..].copy_from_slice(&bytes);
} else {
etf[4..].copy_from_slice(&self.to_le_bytes());
}
let (term, _) = env.binary_to_term(&etf).unwrap();
term
}
}
}

impl Encoder for u128 {
fn encode<'a>(&self, env: Env<'a>) -> Term<'a> {
if let Ok(int) = u64::try_from(*self) {
int.encode(env)
} else {
let mut etf = [0u8; 4 + 16];
etf[0] = EXTERNAL_TERM_FORMAT_VERSION;
etf[1] = SMALL_BIG_EXT;
etf[2] = 16; // length in bytes
etf[4..].copy_from_slice(&self.to_le_bytes());
let (term, _) = env.binary_to_term(&etf).unwrap();
term
}
}
}

impl<'a> Decoder<'a> for i128 {
fn decode(term: Term<'a>) -> NifResult<Self> {
if !term.is_integer() {
return Err(Error::BadArg);
}

if let Ok(int) = term.decode::<i64>() {
return Ok(int as i128);
}

let input = term.to_binary();
let input = input.as_slice();
if input.len() < 4 {
return Err(Error::BadArg);
}

if input[0] != EXTERNAL_TERM_FORMAT_VERSION {
return Err(Error::BadArg);
}

if input[1] != SMALL_BIG_EXT {
return Err(Error::BadArg);
}

let n = input[2] as usize;
if n > 16 {
return Err(Error::BadArg);
}

let mut res = [0u8; 16];
res[16 - n..].copy_from_slice(&input[4..4 + n]);
let res = i128::from_le_bytes(res);
if res < 0 {
// The stored data is supposed to be unsigned, so if we interpret it as negative here,
// it was too large.
return Err(Error::BadArg);
}

if input[3] == 0 {
Ok(res)
} else {
Ok(-res)
}
}
}

impl<'a> Decoder<'a> for u128 {
fn decode(term: Term<'a>) -> NifResult<Self> {
if !term.is_integer() {
return Err(Error::BadArg);
}

if let Ok(int) = term.decode::<u64>() {
return Ok(int as u128);
}

let input = term.to_binary();
let input = input.as_slice();

if input.len() < 4 {
return Err(Error::BadArg);
}

if input[0] != EXTERNAL_TERM_FORMAT_VERSION {
return Err(Error::BadArg);
}

if input[1] != SMALL_BIG_EXT {
return Err(Error::BadArg);
}

let n = input[2] as usize;
if n > 16 {
return Err(Error::BadArg);
}

if input[3] == 1 {
// Negative value
return Err(Error::BadArg);
}

let mut res = [0u8; 16];
res[16 - n..].copy_from_slice(&input[4..4 + n]);
Ok(u128::from_le_bytes(res))
}
}
1 change: 1 addition & 0 deletions rustler/src/types/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@ use crate::{Env, Error, NifResult, Term};

#[macro_use]
pub mod atom;
pub mod i128;
pub use crate::types::atom::Atom;

pub mod binary;
Expand Down
2 changes: 2 additions & 0 deletions rustler_tests/lib/rustler_test.ex
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,8 @@ defmodule RustlerTest do
def add_u32(_, _), do: err()
def add_i32(_, _), do: err()
def echo_u8(_), do: err()
def echo_u128(_), do: err()
def echo_i128(_), do: err()
def option_inc(_), do: err()
def erlang_option_inc(_), do: err()
def result_to_int(_), do: err()
Expand Down
2 changes: 2 additions & 0 deletions rustler_tests/native/rustler_test/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,8 @@ rustler::init!(
test_primitives::option_inc,
test_primitives::erlang_option_inc,
test_primitives::result_to_int,
test_primitives::echo_u128,
test_primitives::echo_i128,
test_list::sum_list,
test_list::make_list,
test_term::term_debug,
Expand Down
10 changes: 10 additions & 0 deletions rustler_tests/native/rustler_test/src/test_primitives.rs
Original file line number Diff line number Diff line change
Expand Up @@ -33,3 +33,13 @@ pub fn result_to_int(res: Result<bool, &str>) -> Result<usize, String> {
Err(errstr) => Err(format!("{}{}", errstr, errstr)),
}
}

#[rustler::nif]
pub fn echo_u128(n: u128) -> u128 {
n
}

#[rustler::nif]
pub fn echo_i128(n: i128) -> i128 {
n
}
31 changes: 31 additions & 0 deletions rustler_tests/test/primitives_test.exs
Original file line number Diff line number Diff line change
Expand Up @@ -32,4 +32,35 @@ defmodule RustlerTest.PrimitivesTest do
assert {:error, "watwat"} == RustlerTest.result_to_int({:error, "wat"})
assert_raise ArgumentError, fn -> RustlerTest.result_to_int({:great, true}) end
end

test "i128 support" do
import Bitwise

i = 1 <<< 62
assert i == RustlerTest.echo_i128(i)
assert -i == RustlerTest.echo_i128(-i)

i = 1 <<< 126
assert i == RustlerTest.echo_i128(i)
assert -i == RustlerTest.echo_i128(-i)

assert_raise ArgumentError, fn -> RustlerTest.echo_i128(:non_int) end
assert_raise ArgumentError, fn -> RustlerTest.echo_i128(123.45) end
assert_raise ArgumentError, fn -> RustlerTest.echo_i128(1 <<< 127) end
assert_raise ArgumentError, fn -> RustlerTest.echo_i128(1 <<< 128) end
end

test "u128 support" do
import Bitwise

i = 1 <<< 63
assert i == RustlerTest.echo_u128(i)

i = 1 <<< 127
assert i == RustlerTest.echo_u128(i)

assert_raise ArgumentError, fn -> RustlerTest.echo_u128(:non_int) end
assert_raise ArgumentError, fn -> RustlerTest.echo_u128(123.45) end
assert_raise ArgumentError, fn -> RustlerTest.echo_i128(1 <<< 128) end
end
end

0 comments on commit d7ded73

Please sign in to comment.