atagen

monolith

public
Star0
Watch0
54559daf11c3c52aa4a1a9a9c0916ce6a2511344
1
ssh://brewery@pls.lobotomise.me:33/atagen/monolith.git
crates/planner/src/sat.rsRust159 lines4.2 KB
1
/// A 2-SAT instance over boolean variables, solved through strongly
2
/// connected components of the implication graph.
3
pub(crate) struct TwoSat {
4
    vars: usize,
5
    clauses: Vec<(Lit, Lit)>,
6
}
7
8
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
9
pub(crate) struct Lit {
10
    pub var: usize,
11
    pub value: bool,
12
}
13
14
impl Lit {
15
    pub fn new(var: usize, value: bool) -> Self {
16
        Self { var, value }
17
    }
18
19
    fn node(self) -> usize {
20
        2 * self.var + usize::from(self.value)
21
    }
22
23
    fn negated(self) -> Self {
24
        Self {
25
            var: self.var,
26
            value: !self.value,
27
        }
28
    }
29
}
30
31
impl TwoSat {
32
    pub fn new(vars: usize) -> Self {
33
        Self {
34
            vars,
35
            clauses: Vec::new(),
36
        }
37
    }
38
39
    pub fn either(&mut self, a: Lit, b: Lit) {
40
        self.clauses.push((a, b));
41
    }
42
43
    pub fn require(&mut self, a: Lit) {
44
        self.clauses.push((a, a));
45
    }
46
47
    pub fn solve(&self) -> Option<Vec<bool>> {
48
        let n = 2 * self.vars;
49
        let mut graph = vec![Vec::new(); n];
50
        for &(a, b) in &self.clauses {
51
            graph[a.negated().node()].push(b.node());
52
            graph[b.negated().node()].push(a.node());
53
        }
54
        let comp = tarjan(&graph);
55
        let mut out = Vec::with_capacity(self.vars);
56
        for v in 0..self.vars {
57
            let (f, t) = (comp[2 * v], comp[2 * v + 1]);
58
            if f == t {
59
                return None;
60
            }
61
            out.push(t < f);
62
        }
63
        Some(out)
64
    }
65
66
    /// A solution that takes each variable's preferred value, in order,
67
    /// whenever that stays satisfiable.
68
    pub fn solve_preferring(mut self, preferred: &[bool]) -> Option<Vec<bool>> {
69
        self.solve()?;
70
        for (var, &value) in preferred.iter().enumerate() {
71
            self.require(Lit::new(var, value));
72
            if self.solve().is_none() {
73
                self.clauses.pop();
74
                self.require(Lit::new(var, !value));
75
            }
76
        }
77
        self.solve()
78
    }
79
}
80
81
/// Tarjan's algorithm, iterative. Components come out in reverse
82
/// topological order, so a lower index means later in the implication order.
83
fn tarjan(graph: &[Vec<usize>]) -> Vec<usize> {
84
    struct State {
85
        index: Vec<usize>,
86
        low: Vec<usize>,
87
        on_stack: Vec<bool>,
88
        stack: Vec<usize>,
89
        next: usize,
90
    }
91
    impl State {
92
        fn enter(&mut self, v: usize, work: &mut Vec<(usize, usize)>) {
93
            self.index[v] = self.next;
94
            self.low[v] = self.next;
95
            self.next += 1;
96
            self.stack.push(v);
97
            self.on_stack[v] = true;
98
            work.push((v, 0));
99
        }
100
    }
101
    let n = graph.len();
102
    let mut st = State {
103
        index: vec![usize::MAX; n],
104
        low: vec![0; n],
105
        on_stack: vec![false; n],
106
        stack: Vec::new(),
107
        next: 0,
108
    };
109
    let mut comp = vec![usize::MAX; n];
110
    let mut next_comp = 0;
111
    let mut work = Vec::new();
112
    for root in 0..n {
113
        if st.index[root] != usize::MAX {
114
            continue;
115
        }
116
        st.enter(root, &mut work);
117
        while let Some(&(v, edge)) = work.last() {
118
            if let Some(&w) = graph[v].get(edge) {
119
                if let Some(top) = work.last_mut() {
120
                    top.1 += 1;
121
                }
122
                if st.index[w] == usize::MAX {
123
                    st.enter(w, &mut work);
124
                } else if st.on_stack[w] {
125
                    st.low[v] = st.low[v].min(st.index[w]);
126
                }
127
                continue;
128
            }
129
            work.pop();
130
            if let Some(&(parent, _)) = work.last() {
131
                st.low[parent] = st.low[parent].min(st.low[v]);
132
            }
133
            if st.low[v] == st.index[v] {
134
                while let Some(w) = st.stack.pop() {
135
                    st.on_stack[w] = false;
136
                    comp[w] = next_comp;
137
                    if w == v {
138
                        break;
139
                    }
140
                }
141
                next_comp += 1;
142
            }
143
        }
144
    }
145
    comp
146
}
147
148
#[cfg(test)]
149
mod tests {
150
    use super::*;
151
152
    #[test]
153
    fn unsatisfiable() {
154
        let mut s = TwoSat::new(1);
155
        s.require(Lit::new(0, true));
156
        s.require(Lit::new(0, false));
157
        assert!(s.solve().is_none());
158
    }
159
}