diff --git a/db/src/db.rs b/db/src/db.rs index cd52b3f1..93094eba 100644 --- a/db/src/db.rs +++ b/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 { match self.get_mut(partition) { - Some(partition) => { - let idx = partition.next_entid(); - 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), + Some(partition) => partition.allocate_entids(n), + None => panic!("Cannot allocate entid from unknown partition: {}", partition) } } diff --git a/db/src/types.rs b/db/src/types.rs index 7d6c01a4..9c326e48 100644 --- a/db/src/types.rs +++ b/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 { + 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. }