use crate::factory::ConstraintFactory;
use crate::stream_def::{Stream, Arity1, ConstraintRecipe};
use crate::{GreynetFact, Score, constraint::{ConstraintWeights, ConstraintId}, Result, ResourceLimits};
use crate::session::Session;
use std::rc::Rc;
use std::cell::RefCell;
use std::marker::PhantomData;
pub struct ConstraintStreamBuilder<'a, S: Score> {
factory: Rc<RefCell<ConstraintFactory<S>>>,
constraint_id: ConstraintId,
_phantom: PhantomData<&'a S>,
}
impl<'a, S: Score + 'static> ConstraintStreamBuilder<'a, S> {
pub fn new(
factory: Rc<RefCell<ConstraintFactory<S>>>,
constraint_id: ConstraintId,
) -> Self {
Self {
factory,
constraint_id,
_phantom: PhantomData,
}
}
pub fn for_each<T: GreynetFact + 'static>(self) -> Stream<Arity1, S> {
let mut stream = ConstraintFactory::from::<T>(&self.factory);
stream.constraint_id_context = Some(self.constraint_id);
stream
}
}
pub struct ConstraintBuilder<S: Score> {
pub factory: Rc<RefCell<ConstraintFactory<S>>>,
pub weights: Rc<RefCell<ConstraintWeights>>,
pub _phantom: PhantomData<S>,
}
impl<S: Score + 'static> ConstraintBuilder<S> {
pub fn new() -> Self {
Self::with_limits(ResourceLimits::default())
}
pub fn with_limits(limits: ResourceLimits) -> Self {
let weights = Rc::new(RefCell::new(ConstraintWeights::new()));
let factory = Rc::new(RefCell::new(ConstraintFactory::with_limits(Rc::clone(&weights), limits)));
Self {
factory,
weights,
_phantom: PhantomData,
}
}
#[inline]
pub fn for_each<T: GreynetFact + 'static>(&self) -> Stream<Arity1, S> {
ConstraintFactory::from::<T>(&self.factory)
}
pub fn build(self) -> Result<Session<S>> {
let factory = Rc::try_unwrap(self.factory)
.map_err(|_| crate::GreynetError::constraint_builder_error("ConstraintFactory has multiple owners"))?
.into_inner();
factory.build_session()
}
pub fn add_constraint(&mut self, id: &str, weight: f64) -> ConstraintStreamBuilder<S> {
self.weights.borrow().set_weight(id, weight);
let constraint_id = self.weights.borrow().get_or_create_id(id);
ConstraintStreamBuilder::new(
self.factory.clone(),
constraint_id,
)
}
pub fn constraint(mut self, id: &str, weight: f64, recipe_fn: impl Fn(ConstraintId) -> ConstraintRecipe<S>) -> Self {
self.weights.borrow().set_weight(id, weight);
let constraint_id = self.weights.borrow().get_or_create_id(id);
let recipe = recipe_fn(constraint_id);
self.factory.borrow_mut().add_constraint_def(recipe);
self
}
pub fn bulk_constraints(
mut self,
constraints: Vec<(&str, f64, Box<dyn Fn(ConstraintId) -> ConstraintRecipe<S>>)>,
) -> Self {
for (id, weight, recipe_fn) in constraints {
self.weights.borrow().set_weight(id, weight);
let constraint_id = self.weights.borrow().get_or_create_id(id);
let recipe = recipe_fn(constraint_id);
self.factory.borrow_mut().add_constraint_def(recipe);
}
self
}
}
impl<S: Score> Default for ConstraintBuilder<S> {
fn default() -> Self {
Self::new()
}
}