--- a/FrmMainWindow.cs	Sun Aug 05 21:01:18 2012 +0100
+++ b/FrmMainWindow.cs	Wed Aug 08 20:53:54 2012 +0100
@@ -314,24 +314,99 @@
 		{
 			TreeStore model = (TreeStore)treeUnits.Model;
 			TreeIter iter;
+
+			if (unit.ParentUnit == null)
+			{
+				iter = GetUnitTreeIterForCategory(model, unit.Category);
+			}
+			else
+			{
+				iter = GetUnitTreeIterForUnit(model, unit.ParentUnit);
+			}
+
+			if (!iter.Equals(TreeIter.Zero))
+			{
+				model.AppendValues(iter, unit);
+				TreePath path = model.GetPath(iter);
+				treeUnits.ExpandToPath(path);
+			}
+		}
+
+		private TreeIter GetUnitTreeIterForCategory(TreeStore model, ArmyCategory category)
+		{			
+			TreeIter targetIter = TreeIter.Zero;
+			TreeIter iter;
 			model.GetIterFirst(out iter);
+			bool found = false;
 
 			do
 			{
 				object obj = model.GetValue(iter, 0);
 
-				if (obj is ArmyCategory)
+				if (obj.Equals(category))
 				{
-					ArmyCategory cat = (ArmyCategory)obj;
+					targetIter = iter;
+					break;
+				}
+			} while (model.IterNext(ref iter));
+
+			return targetIter;
+		}
+
+		private TreeIter GetUnitTreeIterForUnit(TreeStore model, WFObjects.Unit unit)
+		{
+			TreeIter targetIter = TreeIter.Zero;
+			TreeIter iter;
+			model.GetIterFirst(out iter);
+			
+			do
+			{
+				TreeIter unitIter = GetUnitTreeIterForUnit(model, iter, unit);
 
-					if (cat.Equals(unit.Category))
+				if (!unitIter.Equals(TreeIter.Zero))
+				{
+					targetIter = unitIter;
+					break;
+				}
+			} while (model.IterNext(ref iter));
+			
+			return targetIter;
+		}
+
+		private TreeIter GetUnitTreeIterForUnit(TreeStore model, TreeIter parentIter, WFObjects.Unit unit)
+		{
+			TreeIter iter;
+			model.IterChildren(out iter, parentIter);
+
+			if (iter.Equals(TreeIter.Zero))
+			{
+				return TreeIter.Zero;
+			}
+
+			TreeIter targetIter = TreeIter.Zero;
+			
+			do
+			{				
+				object obj = model.GetValue(iter, 0);
+				if (obj.Equals(unit))
+				{
+					targetIter = iter;
+					break;
+				}
+				else
+				{
+					TreeIter unitIter = GetUnitTreeIterForUnit(model, iter, unit);
+					
+					if (!unitIter.Equals(TreeIter.Zero))
 					{
-						model.AppendValues(iter, unit);
-						TreePath path = model.GetPath(iter);
-						treeUnits.ExpandToPath(path);
+						targetIter = unitIter;
+						break;
 					}
 				}
 			} while (model.IterNext(ref iter));
+			
+			return targetIter;
+
 		}
 
 		private void OnUnitRemoved(WarFoundryObject obj)
@@ -345,42 +420,12 @@
 		private void RemoveUnitFromTree(WFObjects.Unit unit)
 		{
 			TreeStore model = (TreeStore)treeUnits.Model;
-			TreeIter iter;
-			model.GetIterFirst(out iter);
-			bool removed = false;
-
-			do
-			{
-				object obj = model.GetValue(iter, 0);
-
-				if (obj is ArmyCategory)
-				{
-					ArmyCategory cat = (ArmyCategory)obj;
-
-					if (unit.Category == null || cat.Equals(unit.Category))
-					{
-						TreeIter innerIter;
-						model.IterChildren(out innerIter, iter);
+			TreeIter iter = GetUnitTreeIterForUnit(model, unit);
 
-						do
-						{
-							object innerObj = model.GetValue(innerIter, 0);
-
-							if (unit.Equals(innerObj))
-							{
-								model.Remove(ref innerIter);
-								removed = true;
-								break;
-							}
-						} while (model.IterNext(ref innerIter));
-
-						if (removed)
-						{
-							break;
-						}
-					}
-				}
-			} while (model.IterNext(ref iter));
+			if (!iter.Equals(TreeIter.Zero))
+			{
+				model.Remove(ref iter);
+			}
 		}
 
 		private void RemoveUnitTab(WFObjects.Unit unit)