NextVarRange

crates/pindakaas-derive/src/lib.rs

@@ -299,10 +299,9 @@ pub fn ipasir_solver_derive(input: TokenStream) -> TokenStream {
 			}
 		}
 
-		impl #ident {
-			/// Return the next [`size`] variables that can be stored as an inclusive range.
-			pub fn new_var_range(&mut self, size: usize) -> crate::helpers::VarRange {
-				#vars .new_var_range(size).expect("variable pool exhaused")
+		impl crate::solver::NextVarRange for #ident {
+			fn next_var_range(&mut self, size: usize) -> Option<crate::helpers::VarRange> {
+				#vars .next_var_range(size)
 			}
 		}
 

crates/pindakaas/src/lib.rs

@@ -22,7 +22,7 @@ use std::{
 
 use helpers::VarRange;
 use itertools::{Itertools, Position};
-use solver::VarFactory;
+use solver::{NextVarRange, VarFactory};
 
 mod cardinality;
 mod cardinality_one;
@@ -382,11 +382,11 @@ impl Cnf {
 	pub fn literals(&self) -> usize {
 		self.size.iter().sum()
 	}
+}
 
-	pub fn new_var_range(&mut self, size: usize) -> VarRange {
-		self.nvar
-			.new_var_range(size)
-			.expect("exhausted variable pool")
+impl NextVarRange for Cnf {
+	fn next_var_range(&mut self, size: usize) -> Option<VarRange> {
+		self.nvar.next_var_range(size)
 	}
 }
 
@@ -656,10 +656,6 @@ impl Wcnf {
 	pub fn iter(&self) -> impl Iterator<Item = (&[Lit], &Option<Coeff>)> {
 		self.cnf.iter().zip(self.weights.iter())
 	}
-
-	pub fn new_var_range(&mut self, size: usize) -> VarRange {
-		self.cnf.new_var_range(size)
-	}
 }
 
 impl ClauseDatabase for Wcnf {
@@ -671,6 +667,12 @@ impl ClauseDatabase for Wcnf {
 	}
 }
 
+impl NextVarRange for Wcnf {
+	fn next_var_range(&mut self, size: usize) -> Option<VarRange> {
+		self.cnf.next_var_range(size)
+	}
+}
+
 impl Cnf {
 	pub fn iter(&self) -> CnfIterator<'_> {
 		CnfIterator {

crates/pindakaas/src/solver.rs

@@ -251,15 +251,49 @@ impl VarFactory {
 			Self::MAX_VARS
 		}
 	}
+}
+
+impl Default for VarFactory {
+	fn default() -> Self {
+		Self {
+			next_var: Some(Var(NonZeroI32::new(1).unwrap())),
+		}
+	}
+}
+
+impl Iterator for VarFactory {
+	type Item = Var;
+
+	fn next(&mut self) -> Option<Self::Item> {
+		let var = self.next_var;
+		if let Some(var) = var {
+			self.next_var = var.next_var();
+		}
+		var
+	}
+}
+
+/// Allow request for sequential ranges of variables.
+pub trait NextVarRange {
+	/// Request the next sequential range of variables.
+	///
+	/// The method is can return [`None`] if the range of the requested [`size`] is not
+	/// available.
+	fn next_var_range(&mut self, size: usize) -> Option<VarRange>;
+}
 
-	/// Return the next [`size`] variables that can be stored as an inclusive range.
-	pub fn new_var_range(&mut self, size: usize) -> Option<VarRange> {
+impl NextVarRange for VarFactory {
+	fn next_var_range(&mut self, size: usize) -> Option<VarRange> {
 		let start = self.next_var?;
 		match size {
 			0 => Some(VarRange::new(
 				Var(NonZeroI32::new(2).unwrap()),
 				Var(NonZeroI32::new(1).unwrap()),
 			)),
+			1 => {
+				self.next_var = start.next_var();
+				Some(VarRange::new(start, start))
+			}
 			_ if size > NonZeroI32::MAX.get() as usize => None,
 			_ => {
 				// Size is reduced by 1 since it includes self.next_var
@@ -276,22 +310,33 @@ impl VarFactory {
 	}
 }
 
-impl Default for VarFactory {
-	fn default() -> Self {
-		Self {
-			next_var: Some(Var(NonZeroI32::new(1).unwrap())),
-		}
-	}
-}
+#[cfg(test)]
+mod tests {
+	use std::num::NonZeroI32;
 
-impl Iterator for VarFactory {
-	type Item = Var;
+	use crate::{
+		solver::{NextVarRange, VarFactory},
+		Var,
+	};
 
-	fn next(&mut self) -> Option<Self::Item> {
-		let var = self.next_var;
-		if let Some(var) = var {
-			self.next_var = var.next_var();
-		}
-		var
+	#[test]
+	fn test_var_range() {
+		let mut factory = VarFactory::default();
+
+		let range = factory.next_var_range(0).unwrap();
+		assert_eq!(range.len(), 0);
+		assert_eq!(factory.next_var, Some(Var(NonZeroI32::new(1).unwrap())));
+
+		let range = factory.next_var_range(1).unwrap();
+		assert_eq!(range.len(), 1);
+		assert_eq!(factory.next_var, Some(Var(NonZeroI32::new(2).unwrap())));
+
+		let range = factory.next_var_range(2).unwrap();
+		assert_eq!(range.len(), 2);
+		assert_eq!(factory.next_var, Some(Var(NonZeroI32::new(4).unwrap())));
+
+		let range = factory.next_var_range(100).unwrap();
+		assert_eq!(range.len(), 100);
+		assert_eq!(factory.next_var, Some(Var(NonZeroI32::new(104).unwrap())));
 	}
 }

crates/pindakaas/src/solver/cadical.rs

@@ -100,12 +100,12 @@ mod tests {
 
 		use crate::{
 			helpers::tests::lits,
-			solver::{PropagatingSolver, Propagator, VarRange},
+			solver::{NextVarRange, PropagatingSolver, Propagator, VarRange},
 		};
 
 		let mut slv = Cadical::default();
 
-		let vars = slv.new_var_range(5);
+		let vars = slv.next_var_range(5).unwrap();
 
 		struct Dist2 {
 			vars: VarRange,