1 files changed, 15 insertions, 3 deletions
diff --git a/client/src/render/scene/demand_map.rs b/client/src/render/scene/demand_map.rs
index 16fa181..c27eaac 100644
--- a/client/src/render/scene/demand_map.rs
+++ b/client/src/render/scene/demand_map.rs
@@ -25,7 +25,7 @@ pub struct DemandMap<K, V> {
     inner: RwLock<DemandMapState<K, V>>,
 }
 struct DemandMapState<K, V> {
-    values: HashMap<K, V>,
+    values: HashMap<K, (V, usize)>,
     needed: HashSet<K>,
     size_metric: usize,
 }
@@ -46,18 +46,30 @@ impl<K: Hash + Eq + Clone, V: Clone> DemandMap<K, V> {
     pub fn insert(&self, key: K, value: V, size: usize) {
         let mut s = self.inner.write().unwrap();
         s.needed.remove(&key);
-        s.values.insert(key, value);
+        if let Some((_, old_size)) = s.values.insert(key, (value, size)) {
+            s.size_metric -= old_size;
+        }
         s.size_metric += size;
     }
     pub fn try_get(&self, key: K) -> Option<V> {
         let mut s = self.inner.write().unwrap();
-        if let Some(k) = s.values.get(&key) {
+        if let Some((k, _)) = s.values.get(&key) {
             Some(k.to_owned())
         } else {
             s.needed.insert(key);
             None
         }
     }
+    pub fn regenerate_all(&self) {
+        let mut s = self.inner.write().unwrap();
+        let keys = s.values.keys().cloned().collect::<Vec<_>>();
+        s.needed.extend(keys);
+    }
+    pub fn clear(&self) {
+        let mut s = self.inner.write().unwrap();
+        s.values.clear();
+        s.size_metric = 0;
+    }
 }
 
 impl<K, V> Widget for &DemandMap<K, V> {