Merge pull request #235 from alessandrod/hash-map-get-mut

bpf: hash_map: add get_mut() method

bpf/aya-bpf/src/maps/hash_map.rs

@@ -40,6 +40,11 @@ impl<K, V> HashMap<K, V> {
         get(&mut self.def, key)
     }
 
+    #[inline]
+    pub fn get_mut(&mut self, key: &K) -> Option<&mut V> {
+        get_mut(&mut self.def, key)
+    }
+
     #[inline]
     pub fn insert(&mut self, key: &K, value: &V, flags: u64) -> Result<(), c_long> {
         insert(&mut self.def, key, value, flags)
@@ -85,6 +90,11 @@ impl<K, V> LruHashMap<K, V> {
         get(&mut self.def, key)
     }
 
+    #[inline]
+    pub fn get_mut(&mut self, key: &K) -> Option<&mut V> {
+        get_mut(&mut self.def, key)
+    }
+
     #[inline]
     pub fn insert(&mut self, key: &K, value: &V, flags: u64) -> Result<(), c_long> {
         insert(&mut self.def, key, value, flags)
@@ -135,6 +145,11 @@ impl<K, V> PerCpuHashMap<K, V> {
         get(&mut self.def, key)
     }
 
+    #[inline]
+    pub fn get_mut(&mut self, key: &K) -> Option<&mut V> {
+        get_mut(&mut self.def, key)
+    }
+
     #[inline]
     pub fn insert(&mut self, key: &K, value: &V, flags: u64) -> Result<(), c_long> {
         insert(&mut self.def, key, value, flags)
@@ -185,6 +200,11 @@ impl<K, V> LruPerCpuHashMap<K, V> {
         get(&mut self.def, key)
     }
 
+    #[inline]
+    pub fn get_mut(&mut self, key: &K) -> Option<&mut V> {
+        get_mut(&mut self.def, key)
+    }
+
     #[inline]
     pub fn insert(&mut self, key: &K, value: &V, flags: u64) -> Result<(), c_long> {
         insert(&mut self.def, key, value, flags)
@@ -217,6 +237,15 @@ fn get<'a, K, V>(def: &mut bpf_map_def, key: &K) -> Option<&'a V> {
     }
 }
 
+#[inline]
+fn get_mut<'a, K, V>(def: &mut bpf_map_def, key: &K) -> Option<&'a mut V> {
+    unsafe {
+        let value = bpf_map_lookup_elem(def as *mut _ as *mut _, key as *const _ as *const c_void);
+        // FIXME: alignment
+        NonNull::new(value as *mut V).map(|mut p| p.as_mut())
+    }
+}
+
 #[inline]
 fn insert<K, V>(def: &mut bpf_map_def, key: &K, value: &V, flags: u64) -> Result<(), c_long> {
     let ret = unsafe {