From a1c206b774006a63f76b616259710a54860230d9 Mon Sep 17 00:00:00 2001
From: Joe Thornber <ejt@redhat.com>
Date: Thu, 20 Aug 2020 10:55:38 +0100
Subject: [PATCH] [thin_check (rust)] NodeVisitor only needs to see leaf nodes

---
 src/pdata/btree.rs | 40 ++++++++++++-------------
 src/thin/check.rs  | 73 ++++++++++++++++++----------------------------
 2 files changed, 48 insertions(+), 65 deletions(-)

diff --git a/src/pdata/btree.rs b/src/pdata/btree.rs
index c61fa47..38adeae 100644
--- a/src/pdata/btree.rs
+++ b/src/pdata/btree.rs
@@ -147,7 +147,7 @@ pub fn unpack_node<V: Unpack>(
 //------------------------------------------
 
 pub trait NodeVisitor<V: Unpack> {
-    fn visit(&mut self, node: &Node<V>) -> Result<()>;
+    fn visit(&mut self, header: &NodeHeader, keys: &[u64], values: &[V]) -> Result<()>;
 }
 
 #[derive(Clone)]
@@ -219,21 +219,26 @@ impl BTreeWalker {
         NV: NodeVisitor<V>,
         V: Unpack,
     {
+        use Node::*;
+
         let bt = checksum::metadata_block_type(b.get_data());
         if bt != checksum::BT::NODE {
             return Err(anyhow!("checksum failed for node {}, {:?}", b.loc, bt));
         }
 
         let node = unpack_node::<V>(&b.get_data(), self.ignore_non_fatal, is_root)?;
-        visitor.visit(&node)?;
 
-        if let Node::Internal {
-            header: _h,
-            keys: _k,
-            values,
-        } = node
-        {
-            self.walk_nodes(visitor, &values)?;
+        match node {
+            Internal {
+                header: _h,
+                keys: _k,
+                values,
+            } => {
+                self.walk_nodes(visitor, &values)?;
+            }
+            Leaf { header, keys, values } => {
+                visitor.visit(&header, &keys, &values)?;
+            }
         }
 
         Ok(())
@@ -281,18 +286,11 @@ impl<V> ValueCollector<V> {
 }
 
 impl<V: Unpack + Clone> NodeVisitor<V> for ValueCollector<V> {
-    fn visit(&mut self, node: &Node<V>) -> Result<()> {
-        if let Node::Leaf {
-            header: _h,
-            keys,
-            values,
-        } = node
-        {
-            for n in 0..keys.len() {
-                let k = keys[n];
-                let v = values[n].clone();
-                self.values.insert(k, v);
-            }
+    fn visit(&mut self, _h: &NodeHeader, keys: &[u64], values: &[V]) -> Result<()> {
+        for n in 0..keys.len() {
+            let k = keys[n];
+            let v = values[n].clone();
+            self.values.insert(k, v);
         }
 
         Ok(())
diff --git a/src/thin/check.rs b/src/thin/check.rs
index 8d86660..9fcd33f 100644
--- a/src/thin/check.rs
+++ b/src/thin/check.rs
@@ -9,7 +9,7 @@ use threadpool::ThreadPool;
 
 use crate::checksum;
 use crate::io_engine::{AsyncIoEngine, Block, IoEngine, SyncIoEngine};
-use crate::pdata::btree::{btree_to_map, btree_to_map_with_sm, BTreeWalker, Node, NodeVisitor};
+use crate::pdata::btree::{btree_to_map, btree_to_map_with_sm, BTreeWalker, NodeHeader, NodeVisitor};
 use crate::pdata::space_map::*;
 use crate::pdata::unpack::*;
 use crate::report::*;
@@ -50,38 +50,30 @@ struct BottomLevelVisitor {
 //------------------------------------------
 
 impl NodeVisitor<BlockTime> for BottomLevelVisitor {
-    fn visit(&mut self, node: &Node<BlockTime>) -> Result<()> {
+    fn visit(&mut self, _h: &NodeHeader, _k: &[u64], values: &[BlockTime]) -> Result<()> {
         // FIXME: do other checks
 
-        if let Node::Leaf {
-            header: _h,
-            keys: _k,
-            values,
-        } = node
-        {
-            if values.len() == 0 {
-                return Ok(());
-            }
-
-            let mut data_sm = self.data_sm.lock().unwrap();
-
-            let mut start = values[0].block;
-            let mut len = 1;
-
-            for n in 1..values.len() {
-                let block = values[n].block;
-                if block == start + len {
-                    len += 1;
-                } else {
-                    data_sm.inc(start, len)?;
-                    start = block;
-                    len = 1;
-                }
-            }
-
-            data_sm.inc(start, len)?;
+        if values.len() == 0 {
+            return Ok(());
         }
 
+        let mut data_sm = self.data_sm.lock().unwrap();
+
+        let mut start = values[0].block;
+        let mut len = 1;
+
+        for n in 1..values.len() {
+            let block = values[n].block;
+            if block == start + len {
+                len += 1;
+            } else {
+                data_sm.inc(start, len)?;
+                start = block;
+                len = 1;
+            }
+        }
+
+        data_sm.inc(start, len)?;
         Ok(())
     }
 }
@@ -132,21 +124,14 @@ impl<'a> OverflowChecker<'a> {
 }
 
 impl<'a> NodeVisitor<u32> for OverflowChecker<'a> {
-    fn visit(&mut self, node: &Node<u32>) -> Result<()> {
-        if let Node::Leaf {
-            header: _h,
-            keys,
-            values,
-        } = node
-        {
-            for n in 0..keys.len() {
-                let k = keys[n];
-                let v = values[n];
-                let expected = self.data_sm.get(k)?;
-                if expected != v {
-                    return Err(anyhow!("Bad reference count for data block {}.  Expected {}, but space map contains {}.",
-                                      k, expected, v));
-                }
+    fn visit(&mut self, _h: &NodeHeader, keys: &[u64], values: &[u32]) -> Result<()> {
+        for n in 0..keys.len() {
+            let k = keys[n];
+            let v = values[n];
+            let expected = self.data_sm.get(k)?;
+            if expected != v {
+                return Err(anyhow!("Bad reference count for data block {}.  Expected {}, but space map contains {}.",
+                                  k, expected, v));
             }
         }