From 4eb35ca9b82c346bf37e891112d5ada9aa50fcb9 Mon Sep 17 00:00:00 2001
From: Tony Garnock-Jones <tonyg@leastfixedpoint.com>
Date: Sun, 20 Oct 2019 22:24:44 +0100
Subject: [PATCH] Refinements to BTreeBag

---
 src/bag.rs | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/src/bag.rs b/src/bag.rs
index 23c8d62..6b0ab11 100644
--- a/src/bag.rs
+++ b/src/bag.rs
@@ -1,6 +1,5 @@
 use std::collections::BTreeMap;
-use std::collections::btree_map::Iter;
-use std::collections::btree_map::Keys;
+use std::collections::btree_map::{Iter, Keys, Entry};
 use std::iter::{FromIterator, IntoIterator};
 
 type Count = i32;
@@ -12,6 +11,7 @@ pub enum Net {
     PresentToPresent,
 }
 
+#[derive(Debug, PartialEq, Eq, PartialOrd, Ord)]
 // Allows negative counts - a "delta"
 pub struct BTreeBag<V: std::cmp::Ord> {
     counts: BTreeMap<V, Count>,
@@ -63,6 +63,10 @@ impl<V: std::cmp::Ord> BTreeBag<V> {
     pub fn keys(&self) -> Keys<V, Count> {
         self.counts.keys()
     }
+
+    pub fn entry(&mut self, key: V) -> Entry<V, Count> {
+        self.counts.entry(key)
+    }
 }
 
 impl<'a, V: std::cmp::Ord> IntoIterator for &'a BTreeBag<V> {