From d3b0d46107462f64c75074a35ada34b0059eefff Mon Sep 17 00:00:00 2001 From: Jordan Bracco Date: Fri, 12 Nov 2021 14:47:02 +0100 Subject: improve rust.. --- native/treebitmap_nif/src/lib.rs | 133 +++++++++++++++++++++------------------ 1 file changed, 73 insertions(+), 60 deletions(-) diff --git a/native/treebitmap_nif/src/lib.rs b/native/treebitmap_nif/src/lib.rs index 2913819..5d1821b 100644 --- a/native/treebitmap_nif/src/lib.rs +++ b/native/treebitmap_nif/src/lib.rs @@ -18,23 +18,16 @@ mod atoms { } trait Address { - fn nibbles(self) -> Nibbles; fn mask(self, masklen: u32) -> Self; } -#[derive(NifUntaggedEnum, Clone)] +#[derive(NifUntaggedEnum, Copy, Clone)] enum AddrTuple { V4(TupleV4), V6(TupleV6) } impl Address for AddrTuple { - fn nibbles(self) -> Nibbles { - match self { - AddrTuple::V4(tuple_v4) => tuple_v4.nibbles(), - AddrTuple::V6(tuple_v6) => tuple_v6.nibbles() - } - } fn mask(self, masklen: u32) -> Self { match self { AddrTuple::V4(tuple_v4) => AddrTuple::V4(tuple_v4.mask(masklen)), @@ -43,7 +36,7 @@ impl Address for AddrTuple { } } -#[derive(Debug, NifRecord, Clone)] +#[derive(Debug, NifRecord, Copy, Clone)] #[tag = "inet4"] struct TupleV4 { pub a: u8, @@ -52,15 +45,33 @@ struct TupleV4 { pub d: u8 } -enum Nibbles { - V4([u8; 8]), - V6([u8; 32]) +struct NibblesV4 { pub n: [u8; 8]} +struct NibblesV6 { pub n: [u8; 32]} + +enum Nibbles { V4(NibblesV4), V6(NibblesV6) } + +impl AsRef<[u8]> for Nibbles { + fn as_ref(&self) -> &[u8] { + match self { + Nibbles::V4(nib4) => nib4.as_ref(), + Nibbles::V6(nib6) => nib6.as_ref() + } + } +} + +impl AsRef<[u8]> for NibblesV4 { + fn as_ref(&self) -> &[u8] { + &self.n + } } -impl TupleV4 { - pub fn new(a1: u8, a2: u8, a3: u8, a4: u8) -> Self { - TupleV4 { a: a1, b: a2, c: a3, d: a4 } +impl AsRef<[u8]> for NibblesV6 { + fn as_ref(&self) -> &[u8] { + &self.n } +} + +impl TupleV4 { pub fn from(num: u32) -> Self { TupleV4 { a: (num >> 24) as u8, @@ -69,19 +80,13 @@ impl TupleV4 { d: num as u8, } } -} -impl Address for TupleV4 { - fn nibbles(self) -> Nibbles { - let mut ret: [u8; 8] = [0; 8]; - let bytes: [u8; 4] = [self.a, self.b, self.c, self.d]; - for (i, byte) in bytes.iter().enumerate() { - ret[i * 2] = byte >> 4; - ret[i * 2 + 1] = byte & 0xf; - } - Nibbles::V4(ret) + pub fn octets(&self) -> [u8; 4] { + [self.a, self.b, self.c, self.d] } +} +impl Address for TupleV4 { fn mask(self, masklen: u32) -> Self { debug_assert!(masklen <= 32); let ip = u32::from(self); @@ -93,6 +98,29 @@ impl Address for TupleV4 { } } + +impl ::std::convert::From for Nibbles { + fn from(a: AddrTuple) -> Nibbles { + match a { + AddrTuple::V4(v4) => Nibbles::from(v4), + AddrTuple::V6(v6) => Nibbles::from(v6) + } + } +} + +impl ::std::convert::From for Nibbles { + fn from(a: TupleV4) -> Nibbles { + let mut ret: [u8; 8] = [0; 8]; + let bytes: [u8; 4] = a.octets(); + for (i, byte) in bytes.iter().enumerate() { + ret[i * 2] = byte >> 4; + ret[i * 2 + 1] = byte & 0xf; + } + Nibbles::V4(NibblesV4 { n: ret }) + } +} + + impl ::std::convert::From for u32 { fn from(a: TupleV4) -> u32 { (a.a as u32) << 24 @@ -102,7 +130,7 @@ impl ::std::convert::From for u32 { } } -#[derive(Debug, NifRecord, Clone)] +#[derive(Debug, NifRecord, Copy, Clone)] #[tag = "inet6"] struct TupleV6 { pub a1: u16, @@ -121,7 +149,7 @@ impl TupleV6 { TupleV6 { a1: a1, a2, a3: a3, a4: a4, a5: a5, a6: a6, a7: a7, a8: a8 } } - fn octets(self) -> [u8; 16] { + fn octets(&self) -> [u8; 16] { [ (self.a1 >> 8) as u8, self.a1 as u8, @@ -142,7 +170,7 @@ impl TupleV6 { ] } - fn segments(self) -> [u16; 8] { + fn segments(&self) -> [u16; 8] { let bytes = self.octets(); [ (bytes[0] as u16) << 8 | (bytes[1] as u16), @@ -160,16 +188,6 @@ impl TupleV6 { impl Address for TupleV6 { - fn nibbles(self) -> Nibbles { - let mut ret: [u8; 32] = [0; 32]; - let bytes = self.octets(); - for (i, byte) in bytes.iter().enumerate() { - ret[i * 2] = byte >> 4; - ret[i * 2 + 1] = byte & 0xf; - } - Nibbles::V6(ret) - } - fn mask(self, masklen: u32) -> Self { debug_assert!(masklen <= 128); let mut ret = self.segments(); @@ -186,6 +204,19 @@ impl Address for TupleV6 { } +impl ::std::convert::From for Nibbles { + fn from(a: TupleV6) -> Nibbles { + let mut ret: [u8; 32] = [0; 32]; + let bytes = a.octets(); + for (i, byte) in bytes.iter().enumerate() { + ret[i * 2] = byte >> 4; + ret[i * 2 + 1] = byte & 0xf; + } + Nibbles::V6(NibblesV6 { n: ret }) + } +} + + #[rustler::nif] fn new() -> NifResult> { let tree = TreeBitmap::new(); @@ -219,11 +250,7 @@ fn add<'a>( value: u32 ) -> Term { let mut tree = table_resource.tree.lock().unwrap(); - let result = match ip.nibbles() { - Nibbles::V4(nibbles) => tree.insert(&nibbles.as_ref(), masklen, value), - Nibbles::V6(nibbles) => tree.insert(&nibbles.as_ref(), masklen, value) - }; - if let Some(value) = result { + if let Some(value) = tree.insert(Nibbles::from(ip).as_ref(), masklen, value) { make_tuple(env, &[atoms::ok().encode(env), value.encode(env)]) } else { make_tuple(env, &[atoms::ok().encode(env), atoms::nil().encode(env)]) @@ -238,11 +265,7 @@ fn remove<'a>( masklen: u32 ) -> Term { let mut tree = table_resource.tree.lock().unwrap(); - let result = match ip.nibbles() { - Nibbles::V4(nibbles) => tree.remove(&nibbles.as_ref(), masklen), - Nibbles::V6(nibbles) => tree.remove(&nibbles.as_ref(), masklen), - }; - if let Some(value) = result { + if let Some(value) = tree.remove(Nibbles::from(ip).as_ref(), masklen) { make_tuple(env, &[atoms::ok().encode(env), value.encode(env)]) } else { make_tuple(env, &[atoms::ok().encode(env), atoms::nil().encode(env)]) @@ -256,12 +279,7 @@ fn longest_match<'a>( ip: AddrTuple ) -> Term { let tree = table_resource.tree.lock().unwrap(); - let ip2 = ip.clone(); - let result = match ip2.nibbles() { - Nibbles::V4(nibbles) => tree.longest_match(&nibbles.as_ref()), - Nibbles::V6(nibbles) => tree.longest_match(&nibbles.as_ref()) - }; - if let Some((bits_matched, value)) = result { + if let Some((bits_matched, value)) = tree.longest_match(Nibbles::from(ip).as_ref()) { let prefix = ip.mask(bits_matched); make_tuple(env, &[atoms::ok().encode(env), prefix.encode(env), bits_matched.encode(env), value.encode(env)]) } else { @@ -277,12 +295,7 @@ fn exact_match<'a>( masklen: u32 ) -> Term { let tree = table_resource.tree.lock().unwrap(); - let ip2 = ip.clone(); - let result = match ip2.nibbles() { - Nibbles::V4(nibbles) => tree.exact_match(&nibbles.as_ref(), masklen), - Nibbles::V6(nibbles) => tree.exact_match(&nibbles.as_ref(), masklen) - }; - if let Some(value) = result { + if let Some(value) = tree.exact_match(Nibbles::from(ip).as_ref(), masklen) { make_tuple(env, &[atoms::ok().encode(env), value.encode(env)]) } else { make_tuple(env, &[atoms::ok().encode(env), atoms::nil().encode(env)]) -- cgit v1.2.3