Fixed root context not working.

src/context.rs

@@ -227,10 +227,8 @@ impl KayakRootContext {
         parent_id: Option<Entity>,
         context_entity: Entity,
     ) {
-        if let Some(parent_id) = parent_id {
-            self.context_entities
-                .add_context_entity::<T>(parent_id, context_entity);
-        }
+        self.context_entities
+            .add_context_entity::<T>(parent_id, context_entity);
     }
 
     /// Returns a new/existing widget entity.

src/context_entities.rs

@@ -8,7 +8,7 @@ use dashmap::DashMap;
 
 #[derive(Debug, Clone)]
 pub struct ContextEntities {
-    ce: Arc<DashMap<Entity, DashMap<TypeId, Entity>>>,
+    ce: Arc<DashMap<Option<Entity>, DashMap<TypeId, Entity>>>,
 }
 
 impl ContextEntities {
@@ -20,7 +20,7 @@ impl ContextEntities {
 
     pub fn add_context_entity<T: Default + 'static>(
         &self,
-        parent_id: Entity,
+        parent_id: Option<Entity>,
         context_entity: Entity,
     ) {
         if !self.ce.contains_key(&parent_id) {
@@ -30,7 +30,7 @@ impl ContextEntities {
         inner.insert(T::default().type_id(), context_entity);
     }
 
-    pub fn get_context_entity<T: Default + 'static>(&self, parent_id: Entity) -> Option<Entity> {
+    pub fn get_context_entity<T: Default + 'static>(&self, parent_id: Option<Entity>) -> Option<Entity> {
         if !self.ce.contains_key(&parent_id) {
             return None;
         }

src/widget_context.rs

@@ -72,10 +72,8 @@ impl KayakWidgetContext {
         parent_id: Option<Entity>,
         context_entity: Entity,
     ) {
-        if let Some(parent_id) = parent_id {
-            self.context_entities
-                .add_context_entity::<T>(parent_id, context_entity);
-        }
+        self.context_entities
+            .add_context_entity::<T>(parent_id, context_entity);
     }
 
     /// Finds the closest matching context entity by traversing up the tree.
@@ -86,7 +84,7 @@ impl KayakWidgetContext {
         // Check self first..
         if let Some(entity) = self
             .context_entities
-            .get_context_entity::<T>(current_entity)
+            .get_context_entity::<T>(Some(current_entity))
         {
             return Some(entity);
         }
@@ -97,12 +95,20 @@ impl KayakWidgetContext {
             while parent.is_some() {
                 if let Some(entity) = self
                     .context_entities
-                    .get_context_entity::<T>(parent.unwrap().0)
+                    .get_context_entity::<T>(Some(parent.unwrap().0))
                 {
                     return Some(entity);
                 }
                 parent = tree.get_parent(parent.unwrap());
             }
+
+            // Finally check root AKA no parent.
+            if let Some(entity) = self
+                .context_entities
+                .get_context_entity::<T>(None)
+            {
+                return Some(entity);
+            }
         }
 
         None