didn't actually remove state on widget remove (mem leak)

src/default/mod.rs

@@ -143,6 +143,7 @@ impl<State> UiRsc for DefaultRsc<State> {
 
     fn on_remove(&mut self, id: WidgetId) {
         self.events.remove(id);
+        self.state.remove(id);
     }
 }
 

src/default/state.rs

@@ -1,4 +1,7 @@
-use iris_core::{WidgetId, util::HashMap};
+use iris_core::{
+    WidgetId,
+    util::{HashMap, HashSet},
+};
 use std::{
     any::{Any, TypeId},
     marker::PhantomData,
@@ -13,6 +16,7 @@ struct Key {
 
 #[derive(Default)]
 pub struct WidgetState {
+    widgets: HashMap<WidgetId, HashSet<(TypeId, usize)>>,
     counts: HashMap<(WidgetId, TypeId), usize>,
     map: HashMap<Key, Box<dyn Any>>,
 }
@@ -24,16 +28,20 @@ impl WidgetState {
     pub fn add<T: 'static>(&mut self, id: WidgetId, data: T) -> WeakState<T> {
         let ty = TypeId::of::<T>();
         let count = self.counts.entry((id, ty)).or_default();
-        let key = Key { ty, i: *count, id };
+        let i = *count;
+        let key = Key { ty, i, id };
         self.map.insert(key, Box::new(data));
+        self.widgets.entry(id).or_default().insert((ty, i));
         *count += 1;
         WeakState {
             key,
             _pd: PhantomData,
         }
     }
-    pub fn remove<T>(&mut self, state: WeakState<T>) {
+    pub fn remove(&mut self, id: WidgetId) {
-        self.map.remove(&state.key);
+        for &(ty, i) in self.widgets.get(&id).into_iter().flatten() {
+            self.map.remove(&Key { id, ty, i });
+        }
     }
     pub fn get<T: 'static>(&self, state: WeakState<T>) -> &T {
         self.map.get(&state.key).unwrap().downcast_ref().unwrap()