1 files changed, 185 insertions, 0 deletions
diff --git a/disk/src/qcow/vec_cache.rs b/disk/src/qcow/vec_cache.rs
new file mode 100644
index 0000000..ecd7673
--- /dev/null
+++ b/disk/src/qcow/vec_cache.rs
@@ -0,0 +1,185 @@
+// Copyright 2018 The Chromium OS Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+use std::collections::hash_map::IterMut;
+use std::collections::HashMap;
+use std::io;
+use std::ops::{Index, IndexMut};
+use std::slice::SliceIndex;
+
+/// Trait that allows for checking if an implementor is dirty. Useful for types that are cached so
+/// it can be checked if they need to be committed to disk.
+pub trait Cacheable {
+    /// Used to check if the item needs to be written out or if it can be discarded.
+    fn dirty(&self) -> bool;
+}
+
+#[derive(Debug)]
+/// Represents a vector that implements the `Cacheable` trait so it can be held in a cache.
+pub struct VecCache<T: 'static + Copy + Default> {
+    vec: Box<[T]>,
+    dirty: bool,
+}
+
+impl<T: 'static + Copy + Default> VecCache<T> {
+    /// Creates a `VecCache` that can hold `count` elements.
+    pub fn new(count: usize) -> VecCache<T> {
+        VecCache {
+            vec: vec![Default::default(); count].into_boxed_slice(),
+            dirty: true,
+        }
+    }
+
+    /// Creates a `VecCache` from the passed in `vec`.
+    pub fn from_vec(vec: Vec<T>) -> VecCache<T> {
+        VecCache {
+            vec: vec.into_boxed_slice(),
+            dirty: false,
+        }
+    }
+
+    pub fn get<I>(&self, index: I) -> Option<&<I as SliceIndex<[T]>>::Output>
+    where
+        I: SliceIndex<[T]>,
+    {
+        self.vec.get(index)
+    }
+
+    /// Gets a reference to the underlying vector.
+    pub fn get_values(&self) -> &[T] {
+        &self.vec
+    }
+
+    /// Mark this cache element as clean.
+    pub fn mark_clean(&mut self) {
+        self.dirty = false;
+    }
+
+    /// Returns the number of elements in the vector.
+    pub fn len(&self) -> usize {
+        self.vec.len()
+    }
+}
+
+impl<T: 'static + Copy + Default> Cacheable for VecCache<T> {
+    fn dirty(&self) -> bool {
+        self.dirty
+    }
+}
+
+impl<T: 'static + Copy + Default> Index<usize> for VecCache<T> {
+    type Output = T;
+
+    fn index(&self, index: usize) -> &T {
+        self.vec.index(index)
+    }
+}
+
+impl<T: 'static + Copy + Default> IndexMut<usize> for VecCache<T> {
+    fn index_mut(&mut self, index: usize) -> &mut T {
+        self.dirty = true;
+        self.vec.index_mut(index)
+    }
+}
+
+#[derive(Debug)]
+pub struct CacheMap<T: Cacheable> {
+    capacity: usize,
+    map: HashMap<usize, T>,
+}
+
+impl<T: Cacheable> CacheMap<T> {
+    pub fn new(capacity: usize) -> Self {
+        CacheMap {
+            capacity,
+            map: HashMap::with_capacity(capacity),
+        }
+    }
+
+    pub fn contains_key(&self, key: &usize) -> bool {
+        self.map.contains_key(key)
+    }
+
+    pub fn get(&self, index: &usize) -> Option<&T> {
+        self.map.get(index)
+    }
+
+    pub fn get_mut(&mut self, index: &usize) -> Option<&mut T> {
+        self.map.get_mut(index)
+    }
+
+    pub fn iter_mut(&mut self) -> IterMut<usize, T> {
+        self.map.iter_mut()
+    }
+
+    // Check if the refblock cache is full and we need to evict.
+    pub fn insert<F>(&mut self, index: usize, block: T, write_callback: F) -> io::Result<()>
+    where
+        F: FnOnce(usize, T) -> io::Result<()>,
+    {
+        if self.map.len() == self.capacity {
+            // TODO(dgreid) - smarter eviction strategy.
+            let to_evict = *self.map.iter().nth(0).unwrap().0;
+            if let Some(evicted) = self.map.remove(&to_evict) {
+                if evicted.dirty() {
+                    write_callback(to_evict, evicted)?;
+                }
+            }
+        }
+        self.map.insert(index, block);
+        Ok(())
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    struct NumCache(pub u64);
+    impl Cacheable for NumCache {
+        fn dirty(&self) -> bool {
+            true
+        }
+    }
+
+    #[test]
+    fn evicts_when_full() {
+        let mut cache = CacheMap::<NumCache>::new(3);
+        let mut evicted = None;
+        cache
+            .insert(0, NumCache(5), |index, _| {
+                evicted = Some(index);
+                Ok(())
+            })
+            .unwrap();
+        assert_eq!(evicted, None);
+        cache
+            .insert(1, NumCache(6), |index, _| {
+                evicted = Some(index);
+                Ok(())
+            })
+            .unwrap();
+        assert_eq!(evicted, None);
+        cache
+            .insert(2, NumCache(7), |index, _| {
+                evicted = Some(index);
+                Ok(())
+            })
+            .unwrap();
+        assert_eq!(evicted, None);
+        cache
+            .insert(3, NumCache(8), |index, _| {
+                evicted = Some(index);
+                Ok(())
+            })
+            .unwrap();
+        assert!(evicted.is_some());
+
+        // Check that three of the four items inserted are still there and that the most recently
+        // inserted is one of them.
+        let num_items = (0..=3).filter(|k| cache.contains_key(&k)).count();
+        assert_eq!(num_items, 3);
+        assert!(cache.contains_key(&3));
+    }
+}