Display traits in tree view with check boxes

xunit.runner.wpf/ViewModel/TraitCollectionView.cs

@@ -1,54 +1,36 @@
-using System;
-using System.Collections.Generic;
-using System.Collections.Immutable;
+using System.Collections.Generic;
 using System.Collections.ObjectModel;
-using System.Collections.Specialized;
 using System.Linq;
-using System.Text;
-using System.Threading.Tasks;
 
 namespace xunit.runner.wpf.ViewModel
 {
     public sealed partial class TraitCollectionView
     {
-        private readonly TraitViewModelComparer _comparer = TraitViewModelComparer.Instance;
-        private readonly ObservableCollection<TraitViewModel> _collection = new ObservableCollection<TraitViewModel>();
+        public ObservableCollection<TraitViewModel> Collection { get; } = new ObservableCollection<TraitViewModel>();
 
-        public ObservableCollection<TraitViewModel> Collection => _collection;
-
-        public TraitCollectionView()
+        public void AddRange(IEnumerable<TraitViewModel> traits)
         {
-
-        }
-
-        public void Add(IEnumerable<TraitViewModel> traitList)
-        {
-            foreach (var traitViewModel in traitList)
+            foreach (var trait in traits)
             {
-                InsertIfNotPresent(traitViewModel);
+                var index = Collection.BinarySearch(trait, TraitViewModel.Comparer);
+                if (index < 0)
+                {
+                    Collection.Insert(~index, trait);
+                }
+                else
+                {
+                    // This trait already exists, add more values.
+                    var originalTrait = Collection[index];
+                    originalTrait.AddValues(trait.Children.Select(x => x.Text));
+                }
             }
         }
 
-        private void InsertIfNotPresent(TraitViewModel trait)
+        public ISet<TraitViewModel> GetCheckedTraits()
         {
-            // TODO: make it a binary search
-            for (int i = 0; i < _collection.Count; i++)
-            {
-                var current = _collection[i];
-                var result = _comparer.Compare(trait, current);
-                if (result < 0)
-                {
-                    _collection.Insert(i, trait);
-                    return;
-                }
-
-                if (result == 0)
-                {
-                    return;
-                }
-            }
-
-            _collection.Add(trait);
+            return new HashSet<TraitViewModel>(
+                Collection.SelectMany(x => x.Children).Where(x => x.IsChecked == true),
+                comparer: TraitViewModel.EqualityComparer);
         }
     }
 }