From 448f7cc69fd89499b01bd0673e16b478ef3cc8da Mon Sep 17 00:00:00 2001
From: Simon Bernier St-Pierre <sbernierstpierre@gmail.com>
Date: Sun, 18 Dec 2016 15:33:48 -0500
Subject: [PATCH] add map

---
 src/net/_session/disk.rs | 12 ++++++------
 src/net/_session/map.rs  | 34 ++++++++++++++++++++++++++++++++++
 src/net/_session/mod.rs  |  7 +------
 3 files changed, 41 insertions(+), 12 deletions(-)
 create mode 100644 src/net/_session/map.rs

diff --git a/src/net/_session/disk.rs b/src/net/_session/disk.rs
index 609170c..38a69d1 100644
--- a/src/net/_session/disk.rs
+++ b/src/net/_session/disk.rs
@@ -3,17 +3,17 @@ use std::fs::File;
 use std::io::{self, Read, Seek, SeekFrom, Write};
 use std::sync::Arc;
 
-use net::_session::TorrentMap;
+use net::_session::map::Map;
 use metainfo::{Hash, Metainfo};
 
 pub struct DiskManager<F> where F: Read + Seek + Write {
-    torrents: TorrentMap<Torrent<F>>,
+    torrents: Map<Hash, Torrent<F>>,
 }
 
 impl<F> DiskManager<F> where F: Read + Seek + Write {
 
     pub fn write_piece(&mut self, info_hash: &Hash, index: u32, mut piece: Vec<u8>) -> io::Result<()> {
-        let torrent = self.torrents.get_mut(&info_hash).unwrap();
+        let torrent = &mut self.torrents[info_hash];
 
         torrent.iter_files_for_piece(index, &mut piece, |file, seek, buf| {
             file.seek(SeekFrom::Start(seek))?;
@@ -25,7 +25,7 @@ impl<F> DiskManager<F> where F: Read + Seek + Write {
     }
 
     pub fn read_piece(&mut self, info_hash: &Hash, index: u32) -> io::Result<Vec<u8>> {
-        let torrent = self.torrents.get_mut(&info_hash).unwrap();
+        let torrent = &mut self.torrents[info_hash];
 
         let len = torrent.metainfo.piece_length as usize;
         let mut piece = Vec::with_capacity(len);
@@ -114,7 +114,7 @@ mod tests {
     use std::path::PathBuf;
     use std::sync::Arc;
 
-    use net::_session::TorrentMap;
+    use net::_session::map::Map;
     use metainfo::{Hash, Metainfo, MetainfoFile};
 
     struct MockFile {
@@ -174,7 +174,7 @@ mod tests {
             length: 20,
         };
 
-        let mut torrents = TorrentMap::new();
+        let mut torrents = Map::new();
         torrents.insert(info_hash, Torrent {
             metainfo: Arc::new(metainfo),
             files: vec![MockFile::new(), MockFile::new(), MockFile::new()],
diff --git a/src/net/_session/map.rs b/src/net/_session/map.rs
new file mode 100644
index 0000000..34f8866
--- /dev/null
+++ b/src/net/_session/map.rs
@@ -0,0 +1,34 @@
+use std::borrow::Borrow;
+use std::collections::HashMap;
+use std::hash::Hash;
+use std::ops::{Deref, DerefMut, Index, IndexMut};
+
+pub struct Map<K, V>(HashMap<K, V>);
+
+impl<K, V> Deref for Map<K, V> {
+    type Target = HashMap<K, V>;
+
+    fn deref(&self) -> &Self::Target {
+        &self.0
+    }
+}
+
+impl<K, V> DerefMut for Map<K, V> {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        &mut self.0
+    }
+}
+
+impl<I, K, V> Index<I> for Map<K, V> where I: Borrow<K>, K: Eq + Hash {
+    type Output = V;
+
+    fn index(&self, index: I) -> &V {
+        self.0.get(index.borrow()).expect("index not found")
+    }
+}
+
+impl<I, K, V> IndexMut<I> for Map<K, V> where I: Borrow<K>, K: Eq + Hash {
+    fn index_mut(&mut self, index: I) -> &mut V {
+        self.0.get_mut(index.borrow()).expect("index not found")
+    }
+}
diff --git a/src/net/_session/mod.rs b/src/net/_session/mod.rs
index 14c7334..3b6bf76 100644
--- a/src/net/_session/mod.rs
+++ b/src/net/_session/mod.rs
@@ -1,7 +1,2 @@
 pub mod disk;
-
-use std::collections::HashMap;
-
-use metainfo::Hash;
-
-type TorrentMap<T> = HashMap<Hash, T>;
+pub mod map;