From 29599cfb021df26ddce72e7dd77f84c8c5d5fe44 Mon Sep 17 00:00:00 2001
From: Simon Bernier St-Pierre <sbernierstpierre@gmail.com>
Date: Sat, 17 Dec 2016 00:31:49 -0500
Subject: [PATCH] bitfield: keep track of # of set bit

---
 src/net/bitfield.rs | 43 ++++++++++++++++++++++++++++++++++---------
 1 file changed, 34 insertions(+), 9 deletions(-)

diff --git a/src/net/bitfield.rs b/src/net/bitfield.rs
index bead130..5933e78 100644
--- a/src/net/bitfield.rs
+++ b/src/net/bitfield.rs
@@ -3,13 +3,16 @@ use std::fmt;
 pub struct BitField {
     bits: Vec<u8>,
     len: u32,
+    num_set: u32,
 }
 
 impl BitField {
     pub fn new(bits: Vec<u8>, len: u32) -> BitField {
+        let num_set = bits.iter().map(|b| b.count_ones()).sum::<u32>();
         BitField {
             bits: bits,
             len: len,
+            num_set: num_set,
         }
     }
 
@@ -17,41 +20,57 @@ impl BitField {
         BitField {
             bits: vec![0u8; f64::ceil(len as f64 / 8.) as usize],
             len: len,
+            num_set: 0,
         }
     }
 
+    #[inline]
+    fn calc_positions(&self, index: u32) -> (u32, u32) {
+        (index / 8, index % 8)
+    }
+
     #[inline]
     pub fn len(&self) -> u32 {
         self.len
     }
 
+    #[inline]
+    pub fn num_set(&self) -> u32 {
+        self.num_set
+    }
+
     #[inline]
     pub fn is_set(&self, index: u32) -> bool {
-        let idx = index / 8;
-        let offset = index % 8;
+        let (idx, offset) = self.calc_positions(index);
         (self.bits[idx as usize] & (1 << 7 - offset)) != 0
     }
 
     #[inline]
     pub fn set(&mut self, index: u32) {
-        let idx = index / 8;
-        let offset = index % 8;
-        self.bits[idx as usize] |= 1 << 7 - offset;
+        if !self.is_set(index) {
+            let (idx, offset) = self.calc_positions(index);
+            self.bits[idx as usize] |= 1 << 7 - offset;
+            self.num_set += 1;
+        }
     }
 
     #[inline]
     pub fn unset(&mut self, index: u32) {
-        let idx = index / 8;
-        let offset = index % 8;
-        self.bits[idx as usize] ^= 1 << 7 - offset;
+        if self.is_set(index) {
+            let (idx, offset) = self.calc_positions(index);
+            self.bits[idx as usize] ^= 1 << 7 - offset;
+            self.num_set -= 1;
+        }
     }
 
+    #[inline]
     pub fn as_bytes(&self) -> &[u8] {
         &self.bits
     }
 
+    #[inline]
     pub fn set_ratio(&self) -> f64 {
-        self.bits.iter().map(|b| b.count_ones()).sum::<u32>() as f64 / self.len as f64
+        self.num_set as f64 / self.len as f64
     }
 }
 
@@ -84,6 +103,8 @@ fn test_bitfield_is_set() {
     assert!(bf.is_set(0xd));
     assert!(bf.is_set(0xe));
     assert!(bf.is_set(0xf));
+
+    assert_eq!(bf.num_set(), 12);
 }
 
 #[test]
@@ -97,6 +118,8 @@ fn test_bitfield_set() {
     assert_eq!(bf.bits[1], 128);
     bf.set(15);
     assert_eq!(bf.bits[1], 129);
+
+    assert_eq!(bf.num_set(), 4);
 }
 
 #[test]
@@ -110,6 +133,8 @@ fn test_bitfield_unset() {
     assert_eq!(bf.bits[1], 255-128);
     bf.unset(15);
     assert_eq!(bf.bits[1], 255-128-1);
+
+    assert_eq!(bf.num_set(), 12);
 }
 
 #[test]