runtime: Simplify the LRU-backed consensus verifier light store

.changelog/5149.internal.md

@@ -0,0 +1 @@
+runtime: Simplify the LRU-backed consensus verifier light store

runtime/src/consensus/tendermint/verifier/store/lru.rs

@@ -20,47 +20,33 @@ impl std::fmt::Debug for LruStore {
     }
 }
 
-struct SubStore {
-    lowest_height: Option<Height>,
-    blocks: lru::LruCache<Height, LightBlock>,
+#[derive(Clone, Debug, PartialEq)]
+struct StoreEntry {
+    light_block: LightBlock,
+    status: Status,
 }
 
-impl SubStore {
-    fn new(capacity: usize) -> Self {
+impl StoreEntry {
+    fn new(light_block: LightBlock, status: Status) -> Self {
         Self {
-            lowest_height: None,
-            blocks: lru::LruCache::new(NonZeroUsize::new(capacity).unwrap()),
+            light_block,
+            status,
         }
     }
 }
 
 struct Inner {
-    unverified: SubStore,
-    verified: SubStore,
-    trusted: SubStore,
-    failed: SubStore,
-}
-
-impl Inner {
-    fn store(&mut self, status: Status) -> &mut SubStore {
-        match status {
-            Status::Unverified => &mut self.unverified,
-            Status::Verified => &mut self.verified,
-            Status::Trusted => &mut self.trusted,
-            Status::Failed => &mut self.failed,
-        }
-    }
+    lowest_height: Option<Height>,
+    blocks: lru::LruCache<Height, StoreEntry>,
 }
 
 impl LruStore {
     /// Create a new, empty, in-memory LRU store of the given capacity.
     pub fn new(capacity: usize) -> Self {
         Self {
             inner: Mutex::new(Inner {
-                unverified: SubStore::new(capacity),
-                verified: SubStore::new(capacity),
-                trusted: SubStore::new(capacity),
-                failed: SubStore::new(capacity),
+                lowest_height: None,
+                blocks: lru::LruCache::new(NonZeroUsize::new(capacity).unwrap()),
             }),
         }
     }
@@ -76,11 +62,16 @@ impl LruStore {
 
 impl LightStore for LruStore {
     fn get(&self, height: Height, status: Status) -> Option<LightBlock> {
-        self.inner().store(status).blocks.get(&height).cloned()
+        self.inner()
+            .blocks
+            .get(&height)
+            .filter(|e| e.status == status)
+            .cloned()
+            .map(|e| e.light_block)
     }
 
     fn insert(&mut self, light_block: LightBlock, status: Status) {
-        let mut store = self.inner_mut().store(status);
+        let mut store = self.inner_mut();
 
         // Promote lowest to prevent it from being evicted.
         let height = light_block.height();
@@ -94,12 +85,22 @@ impl LightStore for LruStore {
             store.lowest_height = Some(height);
         }
 
-        store.blocks.put(height, light_block);
+        store
+            .blocks
+            .put(height, StoreEntry::new(light_block, status));
     }
 
     fn remove(&mut self, height: Height, status: Status) {
-        let mut store = self.inner_mut().store(status);
+        let mut store = self.inner_mut();
 
+        if store
+            .blocks
+            .get(&height)
+            .map(|e| e.status != status)
+            .unwrap_or(false)
+        {
+            return;
+        }
         store.blocks.pop(&height);
 
         // Check if this is the lowest height being explicitly removed.
@@ -114,40 +115,44 @@ impl LightStore for LruStore {
 
     fn highest(&self, status: Status) -> Option<LightBlock> {
         self.inner()
-            .store(status)
             .blocks
             .iter()
+            .filter(|(_, e)| e.status == status)
             .max_by_key(|(&height, _)| height)
-            .map(|(_, lb)| lb.clone())
+            .map(|(_, e)| e.light_block.clone())
     }
 
     fn highest_before(&self, height: Height, status: Status) -> Option<LightBlock> {
         self.inner()
-            .store(status)
             .blocks
             .iter()
+            .filter(|(_, e)| e.status == status)
             .filter(|(h, _)| h <= &&height)
             .max_by_key(|(&height, _)| height)
-            .map(|(_, lb)| lb.clone())
+            .map(|(_, e)| e.light_block.clone())
     }
 
     fn lowest(&self, status: Status) -> Option<LightBlock> {
-        let mut inner = self.inner();
-        let store = inner.store(status);
+        let mut store = self.inner();
 
-        store
-            .lowest_height
-            .and_then(|lowest_height| store.blocks.get(&lowest_height).cloned())
+        store.lowest_height.and_then(|lowest_height| {
+            store
+                .blocks
+                .get(&lowest_height)
+                .filter(|e| e.status == status)
+                .cloned()
+                .map(|e| e.light_block)
+        })
     }
 
     #[allow(clippy::needless_collect)]
     fn all(&self, status: Status) -> Box<dyn Iterator<Item = LightBlock>> {
         let light_blocks: Vec<_> = self
             .inner()
-            .store(status)
             .blocks
             .iter()
-            .map(|(_, lb)| lb.clone())
+            .filter(|(_, e)| e.status == status)
+            .map(|(_, e)| e.light_block.clone())
             .collect();
 
         Box::new(light_blocks.into_iter())