Ergonomic generic way to add scip objects & variable builder

mmghannam · mmghannam · commit 083171ee11ff · 2025-01-26T15:18:15.000+01:00
diff --git a/src/builder.rs b/src/builder.rs
@@ -0,0 +1,128 @@
+use crate::{Model, ModelWithProblem, ProblemCreated, VarType, Variable};
+
+/// A trait for adding two values together.
+pub trait CanBeAddedToModel {
+    /// The return type after adding to the model (e.g. `Variable` / `Constraint` ).
+    type Return;
+    /// How to add the value to the model.
+    fn add(&self, model: &mut Model<ProblemCreated>) -> Self::Return;
+}
+
+/// A builder for variables.
+pub struct VarBuilder {
+    name: Option<String>,
+    obj: f64,
+    lb: f64,
+    ub: f64,
+    var_type: VarType,
+}
+
+impl Default for VarBuilder {
+    fn default() -> Self {
+        VarBuilder {
+            name: None,
+            obj: 0.0,
+            lb: 0.0,
+            ub: f64::INFINITY,
+            var_type: VarType::Continuous,
+        }
+    }
+}
+
+impl VarBuilder {
+    /// Sets the variable to be an integer variable.
+    pub fn integer(mut self, lb: isize, ub: isize) -> Self {
+        self.lb = lb as f64;
+        self.ub = ub as f64;
+        self.var_type = VarType::Integer;
+        self
+    }
+
+    /// Sets the variable to be a binary variable.
+    pub fn binary(mut self) -> Self {
+        self.lb = 0.0;
+        self.ub = 1.0;
+        self.var_type = VarType::Binary;
+        self
+    }
+
+    /// Sets the variable to be a continuous variable.
+    pub fn continuous(mut self, lb: f64, ub: f64) -> Self {
+        self.lb = lb;
+        self.ub = ub;
+        self.var_type = VarType::Continuous;
+        self
+    }
+
+    /// Sets the variable to be an implicit integer variable.
+    pub fn impl_int(mut self, lb: f64, ub: f64) -> Self {
+        self.lb = lb;
+        self.ub = ub;
+        self.var_type = VarType::ImplInt;
+        self
+    }
+
+    /// Sets the name of the variable.
+    pub fn name(mut self, name: String) -> Self {
+        self.name = Some(name);
+        self
+    }
+
+    /// Sets the objective coefficient of the variable.
+    pub fn obj(mut self, obj: f64) -> Self {
+        self.obj = obj;
+        self
+    }
+}
+
+impl CanBeAddedToModel for VarBuilder {
+    type Return = Variable;
+    fn add(&self, model: &mut Model<ProblemCreated>) -> Variable {
+        let name = self.name.clone().unwrap_or_else(|| {
+            let n_vars = model.n_vars();
+            format!("x{}", n_vars)
+        });
+
+        model.add_var(self.lb, self.ub, self.obj, &name,  self.var_type)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_var_builder() {
+        let var = VarBuilder::default()
+            .name("x".to_string())
+            .obj(1.0)
+            .continuous(0.0, 1.0);
+
+        assert_eq!(var.name, Some("x".to_string()));
+        assert_eq!(var.obj, 1.0);
+        assert_eq!(var.lb, 0.0);
+        assert_eq!(var.ub, 1.0);
+    }
+
+    #[test]
+    fn test_var_builder_add() {
+        let mut model = Model::default()
+            .set_obj_sense(crate::ObjSense::Maximize);
+        let var = VarBuilder::default()
+            .name("x".to_string())
+            .obj(1.0)
+            .continuous(0.0, 1.0);
+
+        let var = model.add(var);
+
+        assert_eq!(model.n_vars(), 1);
+        assert_eq!(var.name(), "x");
+        assert_eq!(var.obj(), 1.0);
+        assert_eq!(var.lb(), 0.0);
+        assert_eq!(var.ub(), 1.0);
+
+        let solved = model.solve();
+        assert_eq!(solved.status(), crate::Status::Optimal);
+        assert_eq!(solved.obj_val(), 1.0);
+    }
+}
diff --git a/src/lib.rs b/src/lib.rs
@@ -70,6 +70,9 @@ mod param;
 /// Contains the `Row` struct, which represents a row in an LP relaxation.
 pub mod row;
 
+/// Contains methods for creating scip objects in an ergonomic way.
+pub mod builder;
+
 pub use row::*;
 
 /// A macro for calling a `SCIP` function and returning an error if the return code is not `SCIP_OKAY`.
diff --git a/src/model.rs b/src/model.rs
@@ -11,6 +11,7 @@ use crate::{ffi, Row, Separator};
 use crate::{BranchRule, HeurTiming, Heuristic, Pricer};
 use scip_sys::SCIP;
 use std::rc::Rc;
+use crate::builder::CanBeAddedToModel;
 
 /// Represents an optimization model.
 #[non_exhaustive]
@@ -171,6 +172,11 @@ impl Model<ProblemCreated> {
         }
     }
 
+    /// Adds anything that could be added to the model (variables, constraints, etc).
+    pub fn add<R, O: CanBeAddedToModel<Return=R>>(&mut self, object: O) -> R {
+        object.add(self)
+    }
+
     /// Includes a new branch rule in the model with the given name, description, priority, maximum depth, maximum bound distance, and implementation.
     ///
     /// # Arguments
@@ -551,7 +557,8 @@ pub trait ModelWithProblem {
     fn write(&self, path: &str, ext: &str) -> Result<(), Retcode>;
 }
 
-trait ModelStageWithProblem {}
+/// A trait for model stages that have a problem.
+pub trait ModelStageWithProblem {}
 impl ModelStageWithProblem for ProblemCreated {}
 impl ModelStageWithProblem for Solved {}
 impl ModelStageWithProblem for Solving {}
@@ -803,7 +810,9 @@ pub trait ProblemOrSolving {
         name: &str,
     ) -> Constraint;
 }
-trait ModelStageProblemOrSolving {}
+
+/// A trait for model stages that have a problem or are during solving.
+pub trait ModelStageProblemOrSolving {}
 impl ModelStageProblemOrSolving for ProblemCreated {}
 impl ModelStageProblemOrSolving for Solving {}
 
