Move entid allocation logic into Partition r=nalexander

db/src/db.rs

@@ -1089,13 +1089,8 @@ impl PartitionMap {
     /// Allocate `n` fresh entids in the given `partition`.
     pub(crate) fn allocate_entids(&mut self, partition: &str, n: usize) -> Range<i64> {
         match self.get_mut(partition) {
-            Some(partition) => {
+            Some(partition) => partition.allocate_entids(n),
-                let idx = partition.next_entid();
+            None => panic!("Cannot allocate entid from unknown partition: {}", partition)
-                partition.set_next_entid(idx + n as i64);
-                idx..partition.next_entid()
-            },
-            // This is a programming error.
-            None => panic!("Cannot allocate entid from unknown partition: {}", partition),
         }
     }
 

db/src/types.rs

@@ -21,6 +21,7 @@ use std::iter::{
 use std::ops::{
     Deref,
     DerefMut,
+    Range,
 };
 
 extern crate mentat_core;
@@ -83,6 +84,12 @@ impl Partition {
         assert!(self.allows_entid(e), "Partition index must be within its allocated space.");
         self.next_entid_to_allocate = e;
     }
+
+    pub fn allocate_entids(&mut self, n: usize) -> Range<i64> {
+        let idx = self.next_entid();
+        self.set_next_entid(idx + n as i64);
+        idx..self.next_entid()
+    }
 }
 
 /// Map partition names to `Partition` instances.
@@ -205,10 +212,20 @@ mod tests {
     }
 
     #[test]
+    #[should_panic(expected = "Partition index must be within its allocated space.")]
     fn test_partition_limits_boundary5() {
+        let mut part = Partition::new(100, 1000, 100, true);
+        part.allocate_entids(901); // One more than allowed.
+    }
+
+    #[test]
+    fn test_partition_limits_boundary6() {
         let mut part = Partition::new(100, 1000, 100, true);
         part.set_next_entid(100); // First entid that's allowed.
         part.set_next_entid(101); // Just after first.
+
+        assert_eq!(101..111, part.allocate_entids(10));
+
         part.set_next_entid(1000); // Last entid that's allowed.
         part.set_next_entid(999); // Just before last.
     }