diff --git a/src/main/java/net/onrc/onos/ofcontroller/networkgraph/AbstractNetworkGraph.java b/src/main/java/net/onrc/onos/ofcontroller/networkgraph/AbstractNetworkGraph.java
index 71ef16e..5c498b6 100644
--- a/src/main/java/net/onrc/onos/ofcontroller/networkgraph/AbstractNetworkGraph.java
+++ b/src/main/java/net/onrc/onos/ofcontroller/networkgraph/AbstractNetworkGraph.java
@@ -22,7 +22,7 @@
 	protected ConcurrentMap<Long, Switch> switches;
 
 	protected ConcurrentMap<InetAddress, Set<Device>> addr2Device;
-	protected ConcurrentMap<MACAddress, Set<Device>> mac2Device;
+	protected ConcurrentMap<MACAddress, Device> mac2Device;
 
 	public AbstractNetworkGraph() {
 		// TODO: Does these object need to be stored in Concurrent Collection?
@@ -103,11 +103,7 @@
 	}
 
 	@Override
-	public Iterable<Device> getDeviceByMac(MACAddress address) {
-		Set<Device> devices = mac2Device.get(address);
-		if (devices == null) {
-			return Collections.emptyList();
-		}
-		return Collections.unmodifiableCollection(devices);
+	public Device getDeviceByMac(MACAddress address) {
+		return mac2Device.get(address);
 	}
 }
diff --git a/src/main/java/net/onrc/onos/ofcontroller/networkgraph/MutableNetworkGraph.java b/src/main/java/net/onrc/onos/ofcontroller/networkgraph/MutableNetworkGraph.java
index b617e47..ea4b6a9 100644
--- a/src/main/java/net/onrc/onos/ofcontroller/networkgraph/MutableNetworkGraph.java
+++ b/src/main/java/net/onrc/onos/ofcontroller/networkgraph/MutableNetworkGraph.java
@@ -3,6 +3,7 @@
 import java.net.InetAddress;
 import java.util.Collection;
 import java.util.HashMap;
+import java.util.HashSet;
 import java.util.LinkedList;
 
 import net.floodlightcontroller.util.MACAddress;
@@ -12,9 +13,11 @@
  */
 public class MutableNetworkGraph implements NetworkGraph {
 	protected HashMap<Long, SwitchImpl> switches;
+	protected HashMap<MACAddress, Device> devices;
 	
 	public MutableNetworkGraph() {
 		switches = new HashMap<Long, SwitchImpl>();
+		devices = new HashMap<MACAddress, Device>();
 	}
 
 	// Switch operations
@@ -98,13 +101,26 @@
 
 	@Override
 	public Iterable<Device> getDeviceByIp(InetAddress ipAddress) {
-		// TODO Auto-generated method stub
-		return null;
+		// TODO optimize it
+		HashSet<Device> result = new HashSet<Device>();
+		for (Device d: devices.values()) {
+			Collection<InetAddress> addrs = d.getIpAddress();
+			for (InetAddress addr: addrs) {
+				if (addr.equals(ipAddress)) {
+					result.add(d);
+					break;
+				}
+			}
+		}
+		return result;
 	}
 
 	@Override
-	public Iterable<Device> getDeviceByMac(MACAddress address) {
-		// TODO Auto-generated method stub
-		return null;
+	public Device getDeviceByMac(MACAddress address) {
+		return devices.get(address);
+	}
+
+	public void addDevice(Device device) {
+		devices.put(device.getMacAddress(), device);
 	}
 }
diff --git a/src/main/java/net/onrc/onos/ofcontroller/networkgraph/NetworkGraph.java b/src/main/java/net/onrc/onos/ofcontroller/networkgraph/NetworkGraph.java
index 1c4a390..882ebbe 100644
--- a/src/main/java/net/onrc/onos/ofcontroller/networkgraph/NetworkGraph.java
+++ b/src/main/java/net/onrc/onos/ofcontroller/networkgraph/NetworkGraph.java
@@ -19,5 +19,5 @@
 	public Iterable<? extends Link> getIncomingLinksFromSwitch(Long dpid); // Toshi: unnecessary
 	
 	public Iterable<Device> getDeviceByIp(InetAddress ipAddress);
-	public Iterable<Device> getDeviceByMac(MACAddress address);
+	public Device getDeviceByMac(MACAddress address);
 }