rustsat_tools/encodings/pb/
assignment.rs

1//! # PB (Multi-Criteria) Assignment Problem Encoding
2
3use rustsat::{
4    instances::fio::opb,
5    lit,
6    types::{constraints::CardConstraint, Lit},
7};
8
9use crate::encodings::assignment::Assignment;
10
11#[derive(Default)]
12enum Line {
13    /// Hint line
14    #[default]
15    Hint,
16    /// Description
17    Description,
18    /// An objective with given index
19    Objective(usize),
20    /// Constraints stating that exactly one task per agent needs to be selected
21    OneTask(usize),
22    /// Constraints stating that exactly one agent per task needs to be selected
23    OneAgent(usize),
24}
25
26pub struct Encoding {
27    data: Assignment,
28    next_line: Option<Line>,
29}
30
31impl Encoding {
32    pub fn new(data: Assignment) -> Self {
33        Self {
34            data,
35            next_line: Some(Line::default()),
36        }
37    }
38}
39
40impl Iterator for Encoding {
41    type Item = opb::FileLine<<Vec<(Lit, usize)> as IntoIterator>::IntoIter>;
42
43    fn next(&mut self) -> Option<Self::Item> {
44        let selector = |task: usize, agent: usize| lit![(self.data.n_tasks * task + agent) as u32];
45        match self.next_line.take() {
46            Some(line) => Some(match line {
47                Line::Hint => {
48                    self.next_line = Some(Line::Description);
49                    opb::FileLine::Comment(format!(
50                        "#variable= {} #constraint= {}",
51                        self.data.n_tasks * self.data.n_tasks,
52                        2 * self.data.n_tasks
53                    ))
54                }
55                Line::Description => {
56                    self.next_line = Some(Line::Objective(0));
57                    opb::FileLine::Comment(
58                        "MO Assignment instance generated by RustSAT".to_string(),
59                    )
60                }
61                Line::Objective(objidx) => {
62                    let data = &self.data;
63                    let obj: Vec<_> = (0..self.data.n_tasks)
64                        .flat_map(|task| {
65                            (0..self.data.n_tasks).map(move |agent| {
66                                (selector(task, agent), data.cost(task, agent, objidx))
67                            })
68                        })
69                        .collect();
70                    self.next_line = Some(if objidx + 1 < self.data.n_objs {
71                        Line::Objective(objidx + 1)
72                    } else {
73                        Line::OneTask(0)
74                    });
75                    opb::FileLine::Objective(obj.into_iter())
76                }
77                Line::OneTask(agent) => {
78                    let constr = CardConstraint::new_eq(
79                        (0..self.data.n_tasks).map(|task| selector(task, agent)),
80                        1,
81                    );
82                    self.next_line = Some(if agent + 1 < self.data.n_tasks {
83                        Line::OneTask(agent + 1)
84                    } else {
85                        Line::OneAgent(0)
86                    });
87                    opb::FileLine::Card(constr)
88                }
89                Line::OneAgent(task) => {
90                    let constr = CardConstraint::new_eq(
91                        (0..self.data.n_tasks).map(|agent| selector(task, agent)),
92                        1,
93                    );
94                    self.next_line = if task + 1 < self.data.n_tasks {
95                        Some(Line::OneAgent(task + 1))
96                    } else {
97                        None
98                    };
99                    opb::FileLine::Card(constr)
100                }
101            }),
102            None => None,
103        }
104    }
105}