ET9_WebGL_Base/CSharp/Game/Model/UnitComponent.cs

using System;
using System.Collections.Generic;
using System.Linq;
using Common.Base;
using MongoDB.Bson;

namespace Model
{
    public class UnitComponent: Component
    {
        private readonly Dictionary<ObjectId, Unit> units =
                new Dictionary<ObjectId, Unit>();

        private readonly Dictionary<UnitType, Dictionary<ObjectId, Unit>> typeUnits =
                new Dictionary<UnitType, Dictionary<ObjectId, Unit>>();

        public void Add(Unit unit)
        {
            this.units.Add(unit.Id, unit);
            if (!this.typeUnits.ContainsKey(unit.Config.Type))
            {
                this.typeUnits.Add(unit.Config.Type, new Dictionary<ObjectId, Unit>());
            }
            this.typeUnits[unit.Config.Type].Add(unit.Id, unit);
        }

        public Unit Get(ObjectId id)
        {
            Unit unit = null;
            this.units.TryGetValue(id, out unit);
            return unit;
        }

        public Unit[] GetOneType(UnitType type)
        {
            Dictionary<ObjectId, Unit> oneTypeUnits = null;
            if (!this.typeUnits.TryGetValue(type, out oneTypeUnits))
            {
                return new Unit[0];
            }
            return oneTypeUnits.Values.ToArray();
        }

        public bool Remove(Unit unit)
        {
            if (unit == null)
            {
                throw new ArgumentNullException("unit");
            }
            if (!this.units.Remove(unit.Id))
            {
                return false;
            }
            if (!this.typeUnits[unit.Config.Type].Remove(unit.Id))
            {
                return false;
            }
            return true;
        }

        public bool Remove(ObjectId id)
        {
            Unit unit = this.Get(id);
            if (unit == null)
            {
                return false;
            }
            return this.Remove(unit);
        }
    }
}