From aeaa0feb613524587cb8ed8798d594d2056eac98 Mon Sep 17 00:00:00 2001
From: Pankaj Garg <pankaj@solana.com>
Date: Thu, 22 Aug 2019 16:32:38 -0700
Subject: [PATCH] Add range lookups for erasure set indexes (#5612)

---
 core/src/blocktree/meta.rs | 96 +++++++++++++++++++++++++++++++++++++-
 1 file changed, 94 insertions(+), 2 deletions(-)

diff --git a/core/src/blocktree/meta.rs b/core/src/blocktree/meta.rs
index 6e12d215c..dfb70af6f 100644
--- a/core/src/blocktree/meta.rs
+++ b/core/src/blocktree/meta.rs
@@ -1,6 +1,7 @@
 use crate::erasure::ErasureConfig;
 use solana_metrics::datapoint;
-use std::{collections::BTreeSet, ops::RangeBounds};
+use std::cmp::Ordering;
+use std::{collections::BTreeSet, ops::Range, ops::RangeBounds};
 
 #[derive(Clone, Debug, Default, Deserialize, Serialize, Eq, PartialEq)]
 // The Meta column family
@@ -27,6 +28,51 @@ pub struct SlotMeta {
     pub is_connected: bool,
 }
 
+#[derive(Clone, Debug, Default, Deserialize, Serialize, Eq, PartialEq)]
+pub struct ErasureSetRanges {
+    r: Vec<Range<u64>>,
+}
+
+impl ErasureSetRanges {
+    pub fn insert(&mut self, start: u64, end: u64) -> Result<usize, Range<u64>> {
+        let range = if start < end {
+            (start..end)
+        } else {
+            (end..start)
+        };
+
+        match self.pos(range.start) {
+            Ok(pos) => Err(self.r[pos].clone()),
+            Err(pos) => {
+                self.r.insert(pos, range);
+                Ok(pos)
+            }
+        }
+    }
+
+    fn pos(&self, seek: u64) -> Result<usize, usize> {
+        self.r.binary_search_by(|probe| {
+            if probe.contains(&seek) {
+                Ordering::Equal
+            } else {
+                probe.start.cmp(&seek)
+            }
+        })
+    }
+
+    pub fn lookup(&self, seek: u64) -> Result<Range<u64>, usize> {
+        self.pos(seek)
+            .map(|pos| self.r[pos].clone())
+            .or_else(|epos| {
+                if epos < self.r.len() && self.r[epos].contains(&seek) {
+                    Ok(self.r[epos].clone())
+                } else {
+                    Err(epos)
+                }
+            })
+    }
+}
+
 #[derive(Clone, Debug, Default, Deserialize, Serialize, PartialEq)]
 /// Index recording presence/absence of blobs
 pub struct Index {
@@ -56,7 +102,7 @@ pub struct ErasureMeta {
     /// Size of shards in this erasure set
     pub size: usize,
     /// Erasure configuration for this erasure set
-    config: ErasureConfig,
+    pub config: ErasureConfig,
 }
 
 #[derive(Debug, PartialEq)]
@@ -300,4 +346,50 @@ mod test {
             assert_eq!(e_meta.status(&index), DataFull);
         }
     }
+
+    #[test]
+    fn test_erasure_set_ranges() {
+        let mut ranges = ErasureSetRanges::default();
+
+        // Test empty ranges
+        (0..100 as u64).for_each(|i| {
+            assert_eq!(ranges.lookup(i), Err(0));
+        });
+
+        // Test adding one range and all boundary condition lookups
+        assert_eq!(ranges.insert(5, 13), Ok(0));
+        assert_eq!(ranges.lookup(0), Err(0));
+        assert_eq!(ranges.lookup(4), Err(0));
+        assert_eq!(ranges.lookup(5), Ok(5..13));
+        assert_eq!(ranges.lookup(12), Ok(5..13));
+        assert_eq!(ranges.lookup(13), Err(1));
+        assert_eq!(ranges.lookup(100), Err(1));
+
+        // Test adding second range (with backwards values) and all boundary condition lookups
+        assert_eq!(ranges.insert(55, 33), Ok(1));
+        assert_eq!(ranges.lookup(0), Err(0));
+        assert_eq!(ranges.lookup(4), Err(0));
+        assert_eq!(ranges.lookup(5), Ok(5..13));
+        assert_eq!(ranges.lookup(12), Ok(5..13));
+        assert_eq!(ranges.lookup(13), Err(1));
+        assert_eq!(ranges.lookup(32), Err(1));
+        assert_eq!(ranges.lookup(33), Ok(33..55));
+        assert_eq!(ranges.lookup(54), Ok(33..55));
+        assert_eq!(ranges.lookup(55), Err(2));
+
+        // Add a third range between previous two ranges
+        assert_eq!(ranges.insert(23, 30), Ok(1));
+        assert_eq!(ranges.lookup(0), Err(0));
+        assert_eq!(ranges.lookup(4), Err(0));
+        assert_eq!(ranges.lookup(5), Ok(5..13));
+        assert_eq!(ranges.lookup(12), Ok(5..13));
+        assert_eq!(ranges.lookup(13), Err(1));
+        assert_eq!(ranges.lookup(23), Ok(23..30));
+        assert_eq!(ranges.lookup(29), Ok(23..30));
+        assert_eq!(ranges.lookup(30), Err(2));
+        assert_eq!(ranges.lookup(32), Err(2));
+        assert_eq!(ranges.lookup(33), Ok(33..55));
+        assert_eq!(ranges.lookup(54), Ok(33..55));
+        assert_eq!(ranges.lookup(55), Err(3));
+    }
 }