1 files changed, 73 insertions, 26 deletions

diff --git a/src/org/traccar/database/GroupTree.java b/src/org/traccar/database/GroupTree.java
index 7263a6331..b383b1501 100644
--- a/src/org/traccar/database/GroupTree.java
+++ b/src/org/traccar/database/GroupTree.java
@@ -15,6 +15,7 @@
  */
 package org.traccar.database;
 
+import org.traccar.model.Device;
 import org.traccar.model.Group;
 
 import java.util.ArrayList;
@@ -26,80 +27,126 @@ import java.util.Set;
 
 public class GroupTree {
 
-    private static class GroupNode {
+    private static class TreeNode {
 
         private Group group;
-        private GroupNode parent;
-        private Collection<GroupNode> children = new HashSet<>();
+        private Device device;
+        private Collection<TreeNode> children = new HashSet<>();
 
-        public GroupNode(Group group) {
+        public TreeNode(Group group) {
             this.group = group;
         }
 
+        public TreeNode(Device device) {
+            this.device = device;
+        }
+
         @Override
         public int hashCode() {
-            return (int) group.getId();
+            if (group != null) {
+                return (int) group.getId();
+            } else {
+                return (int) device.getId();
+            }
         }
 
         @Override
         public boolean equals(Object obj) {
-            if (!(obj instanceof GroupNode)) {
+            if (!(obj instanceof TreeNode)) {
                 return false;
             }
-            return obj == this || group.getId() == ((GroupNode) obj).group.getId();
+            TreeNode other = (TreeNode) obj;
+            if (other == this) {
+                return true;
+            }
+            if (group != null) {
+                if (other.group != null) {
+                    return group.getId() == other.group.getId();
+                }
+            } else if (device != null) {
+                if (other.device != null) {
+                    return device.getId() == other.device.getId();
+                }
+            }
+            return false;
         }
 
         public Group getGroup() {
             return group;
         }
 
-        public void setParent(GroupNode parent) {
-            this.parent = parent;
+        public Device getDevice() {
+            return device;
+        }
+
+        public void setParent(TreeNode parent) {
             if (parent != null) {
                 parent.children.add(this);
             }
         }
 
-        public GroupNode getParent() {
-            return parent;
-        }
-
-        public Collection<GroupNode> getChildren() {
+        public Collection<TreeNode> getChildren() {
             return children;
         }
 
     }
 
-    private final Map<Long, GroupNode> groupMap = new HashMap<>();
+    private final Map<Long, TreeNode> groupMap = new HashMap<>();
 
-    public GroupTree(Collection<Group> groups) {
+    public GroupTree(Collection<Group> groups, Collection<Device> devices) {
 
         for (Group group : groups) {
-            groupMap.put(group.getId(), new GroupNode(group));
+            groupMap.put(group.getId(), new TreeNode(group));
         }
 
-        for (GroupNode node : groupMap.values()) {
+        for (TreeNode node : groupMap.values()) {
             if (node.getGroup().getGroupId() != 0) {
                 node.setParent(groupMap.get(node.getGroup().getGroupId()));
             }
         }
 
+        Map<Long, TreeNode> deviceMap = new HashMap<>();
+
+        for (Device device : devices) {
+            deviceMap.put(device.getId(), new TreeNode(device));
+        }
+
+        for (TreeNode node : deviceMap.values()) {
+            if (node.getDevice().getGroupId() != 0) {
+                node.setParent(groupMap.get(node.getDevice().getGroupId()));
+            }
+        }
+
     }
 
-    public Collection<Group> getDescendants(long groupId) {
-        Set<GroupNode> results = new HashSet<>();
-        getDescendants(results, groupMap.get(groupId));
+    public Collection<Group> getGroups(long groupId) {
+        Set<TreeNode> results = new HashSet<>();
+        getNodes(results, groupMap.get(groupId));
         Collection<Group> groups = new ArrayList<>();
-        for (GroupNode node : results) {
-            groups.add(node.getGroup());
+        for (TreeNode node : results) {
+            if (node.getGroup() != null) {
+                groups.add(node.getGroup());
+            }
         }
         return groups;
     }
 
-    private void getDescendants(Set<GroupNode> results, GroupNode node) {
-        for (GroupNode child : node.getChildren()) {
+    public Collection<Device> getDevices(long groupId) {
+        Set<TreeNode> results = new HashSet<>();
+        getNodes(results, groupMap.get(groupId));
+        Collection<Device> devices = new ArrayList<>();
+        for (TreeNode node : results) {
+            if (node.getDevice() != null) {
+                devices.add(node.getDevice());
+            }
+        }
+        return devices;
+    }
+
+    private void getNodes(Set<TreeNode> results, TreeNode node) {
+        for (TreeNode child : node.getChildren()) {
             results.add(child);
-            getDescendants(results, child);
+            getNodes(results, child);
         }
     }