atagen

monolith

public
Star0
Watch0
e13e3c02ad0c6449311cbc953672be591643b8e7
1
ssh://brewery@pls.lobotomise.me:33/atagen/monolith.git

Blame: crates/planner/src/sat.rs

b563d2da1c80atagen2026-10-16 08:25 UTC1
/// A 2-SAT instance over boolean variables, solved through strongly
b563d2da1c80atagen2026-10-16 08:25 UTC2
/// connected components of the implication graph.
b563d2da1c80atagen2026-10-16 08:25 UTC3
pub(crate) struct TwoSat {
b563d2da1c80atagen2026-10-16 08:25 UTC4
    vars: usize,
b563d2da1c80atagen2026-10-16 08:25 UTC5
    clauses: Vec<(Lit, Lit)>,
b563d2da1c80atagen2026-10-16 08:25 UTC6
}
b563d2da1c80atagen2026-10-16 08:25 UTC7
b563d2da1c80atagen2026-10-16 08:25 UTC8
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
b563d2da1c80atagen2026-10-16 08:25 UTC9
pub(crate) struct Lit {
b563d2da1c80atagen2026-10-16 08:25 UTC10
    pub var: usize,
b563d2da1c80atagen2026-10-16 08:25 UTC11
    pub value: bool,
b563d2da1c80atagen2026-10-16 08:25 UTC12
}
b563d2da1c80atagen2026-10-16 08:25 UTC13
b563d2da1c80atagen2026-10-16 08:25 UTC14
impl Lit {
b563d2da1c80atagen2026-10-16 08:25 UTC15
    pub fn new(var: usize, value: bool) -> Self {
b563d2da1c80atagen2026-10-16 08:25 UTC16
        Self { var, value }
b563d2da1c80atagen2026-10-16 08:25 UTC17
    }
b563d2da1c80atagen2026-10-16 08:25 UTC18
b563d2da1c80atagen2026-10-16 08:25 UTC19
    fn node(self) -> usize {
b563d2da1c80atagen2026-10-16 08:25 UTC20
        2 * self.var + usize::from(self.value)
b563d2da1c80atagen2026-10-16 08:25 UTC21
    }
b563d2da1c80atagen2026-10-16 08:25 UTC22
b563d2da1c80atagen2026-10-16 08:25 UTC23
    fn negated(self) -> Self {
b563d2da1c80atagen2026-10-16 08:25 UTC24
        Self {
b563d2da1c80atagen2026-10-16 08:25 UTC25
            var: self.var,
b563d2da1c80atagen2026-10-16 08:25 UTC26
            value: !self.value,
b563d2da1c80atagen2026-10-16 08:25 UTC27
        }
b563d2da1c80atagen2026-10-16 08:25 UTC28
    }
b563d2da1c80atagen2026-10-16 08:25 UTC29
}
b563d2da1c80atagen2026-10-16 08:25 UTC30
b563d2da1c80atagen2026-10-16 08:25 UTC31
impl TwoSat {
b563d2da1c80atagen2026-10-16 08:25 UTC32
    pub fn new(vars: usize) -> Self {
b563d2da1c80atagen2026-10-16 08:25 UTC33
        Self {
b563d2da1c80atagen2026-10-16 08:25 UTC34
            vars,
b563d2da1c80atagen2026-10-16 08:25 UTC35
            clauses: Vec::new(),
b563d2da1c80atagen2026-10-16 08:25 UTC36
        }
b563d2da1c80atagen2026-10-16 08:25 UTC37
    }
b563d2da1c80atagen2026-10-16 08:25 UTC38
b563d2da1c80atagen2026-10-16 08:25 UTC39
    pub fn either(&mut self, a: Lit, b: Lit) {
b563d2da1c80atagen2026-10-16 08:25 UTC40
        self.clauses.push((a, b));
b563d2da1c80atagen2026-10-16 08:25 UTC41
    }
b563d2da1c80atagen2026-10-16 08:25 UTC42
b563d2da1c80atagen2026-10-16 08:25 UTC43
    pub fn require(&mut self, a: Lit) {
b563d2da1c80atagen2026-10-16 08:25 UTC44
        self.clauses.push((a, a));
b563d2da1c80atagen2026-10-16 08:25 UTC45
    }
b563d2da1c80atagen2026-10-16 08:25 UTC46
b563d2da1c80atagen2026-10-16 08:25 UTC47
    pub fn solve(&self) -> Option<Vec<bool>> {
b563d2da1c80atagen2026-10-16 08:25 UTC48
        let n = 2 * self.vars;
b563d2da1c80atagen2026-10-16 08:25 UTC49
        let mut graph = vec![Vec::new(); n];
b563d2da1c80atagen2026-10-16 08:25 UTC50
        for &(a, b) in &self.clauses {
b563d2da1c80atagen2026-10-16 08:25 UTC51
            graph[a.negated().node()].push(b.node());
b563d2da1c80atagen2026-10-16 08:25 UTC52
            graph[b.negated().node()].push(a.node());
b563d2da1c80atagen2026-10-16 08:25 UTC53
        }
b563d2da1c80atagen2026-10-16 08:25 UTC54
        let comp = tarjan(&graph);
b563d2da1c80atagen2026-10-16 08:25 UTC55
        let mut out = Vec::with_capacity(self.vars);
b563d2da1c80atagen2026-10-16 08:25 UTC56
        for v in 0..self.vars {
b563d2da1c80atagen2026-10-16 08:25 UTC57
            let (f, t) = (comp[2 * v], comp[2 * v + 1]);
b563d2da1c80atagen2026-10-16 08:25 UTC58
            if f == t {
b563d2da1c80atagen2026-10-16 08:25 UTC59
                return None;
b563d2da1c80atagen2026-10-16 08:25 UTC60
            }
b563d2da1c80atagen2026-10-16 08:25 UTC61
            out.push(t < f);
b563d2da1c80atagen2026-10-16 08:25 UTC62
        }
b563d2da1c80atagen2026-10-16 08:25 UTC63
        Some(out)
b563d2da1c80atagen2026-10-16 08:25 UTC64
    }
b563d2da1c80atagen2026-10-16 08:25 UTC65
b563d2da1c80atagen2026-10-16 08:25 UTC66
    /// A solution that takes each variable's preferred value, in order,
b563d2da1c80atagen2026-10-16 08:25 UTC67
    /// whenever that stays satisfiable.
b563d2da1c80atagen2026-10-16 08:25 UTC68
    pub fn solve_preferring(mut self, preferred: &[bool]) -> Option<Vec<bool>> {
b563d2da1c80atagen2026-10-16 08:25 UTC69
        self.solve()?;
b563d2da1c80atagen2026-10-16 08:25 UTC70
        for (var, &value) in preferred.iter().enumerate() {
b563d2da1c80atagen2026-10-16 08:25 UTC71
            self.require(Lit::new(var, value));
b563d2da1c80atagen2026-10-16 08:25 UTC72
            if self.solve().is_none() {
b563d2da1c80atagen2026-10-16 08:25 UTC73
                self.clauses.pop();
b563d2da1c80atagen2026-10-16 08:25 UTC74
                self.require(Lit::new(var, !value));
b563d2da1c80atagen2026-10-16 08:25 UTC75
            }
b563d2da1c80atagen2026-10-16 08:25 UTC76
        }
b563d2da1c80atagen2026-10-16 08:25 UTC77
        self.solve()
b563d2da1c80atagen2026-10-16 08:25 UTC78
    }
b563d2da1c80atagen2026-10-16 08:25 UTC79
}
b563d2da1c80atagen2026-10-16 08:25 UTC80
b563d2da1c80atagen2026-10-16 08:25 UTC81
/// Tarjan's algorithm, iterative. Components come out in reverse
b563d2da1c80atagen2026-10-16 08:25 UTC82
/// topological order, so a lower index means later in the implication order.
b563d2da1c80atagen2026-10-16 08:25 UTC83
fn tarjan(graph: &[Vec<usize>]) -> Vec<usize> {
b563d2da1c80atagen2026-10-16 08:25 UTC84
    struct State {
b563d2da1c80atagen2026-10-16 08:25 UTC85
        index: Vec<usize>,
b563d2da1c80atagen2026-10-16 08:25 UTC86
        low: Vec<usize>,
b563d2da1c80atagen2026-10-16 08:25 UTC87
        on_stack: Vec<bool>,
b563d2da1c80atagen2026-10-16 08:25 UTC88
        stack: Vec<usize>,
b563d2da1c80atagen2026-10-16 08:25 UTC89
        next: usize,
b563d2da1c80atagen2026-10-16 08:25 UTC90
    }
b563d2da1c80atagen2026-10-16 08:25 UTC91
    impl State {
b563d2da1c80atagen2026-10-16 08:25 UTC92
        fn enter(&mut self, v: usize, work: &mut Vec<(usize, usize)>) {
b563d2da1c80atagen2026-10-16 08:25 UTC93
            self.index[v] = self.next;
b563d2da1c80atagen2026-10-16 08:25 UTC94
            self.low[v] = self.next;
b563d2da1c80atagen2026-10-16 08:25 UTC95
            self.next += 1;
b563d2da1c80atagen2026-10-16 08:25 UTC96
            self.stack.push(v);
b563d2da1c80atagen2026-10-16 08:25 UTC97
            self.on_stack[v] = true;
b563d2da1c80atagen2026-10-16 08:25 UTC98
            work.push((v, 0));
b563d2da1c80atagen2026-10-16 08:25 UTC99
        }
b563d2da1c80atagen2026-10-16 08:25 UTC100
    }
b563d2da1c80atagen2026-10-16 08:25 UTC101
    let n = graph.len();
b563d2da1c80atagen2026-10-16 08:25 UTC102
    let mut st = State {
b563d2da1c80atagen2026-10-16 08:25 UTC103
        index: vec![usize::MAX; n],
b563d2da1c80atagen2026-10-16 08:25 UTC104
        low: vec![0; n],
b563d2da1c80atagen2026-10-16 08:25 UTC105
        on_stack: vec![false; n],
b563d2da1c80atagen2026-10-16 08:25 UTC106
        stack: Vec::new(),
b563d2da1c80atagen2026-10-16 08:25 UTC107
        next: 0,
b563d2da1c80atagen2026-10-16 08:25 UTC108
    };
b563d2da1c80atagen2026-10-16 08:25 UTC109
    let mut comp = vec![usize::MAX; n];
b563d2da1c80atagen2026-10-16 08:25 UTC110
    let mut next_comp = 0;
b563d2da1c80atagen2026-10-16 08:25 UTC111
    let mut work = Vec::new();
b563d2da1c80atagen2026-10-16 08:25 UTC112
    for root in 0..n {
b563d2da1c80atagen2026-10-16 08:25 UTC113
        if st.index[root] != usize::MAX {
b563d2da1c80atagen2026-10-16 08:25 UTC114
            continue;
b563d2da1c80atagen2026-10-16 08:25 UTC115
        }
b563d2da1c80atagen2026-10-16 08:25 UTC116
        st.enter(root, &mut work);
b563d2da1c80atagen2026-10-16 08:25 UTC117
        while let Some(&(v, edge)) = work.last() {
b563d2da1c80atagen2026-10-16 08:25 UTC118
            if let Some(&w) = graph[v].get(edge) {
b563d2da1c80atagen2026-10-16 08:25 UTC119
                if let Some(top) = work.last_mut() {
b563d2da1c80atagen2026-10-16 08:25 UTC120
                    top.1 += 1;
b563d2da1c80atagen2026-10-16 08:25 UTC121
                }
b563d2da1c80atagen2026-10-16 08:25 UTC122
                if st.index[w] == usize::MAX {
b563d2da1c80atagen2026-10-16 08:25 UTC123
                    st.enter(w, &mut work);
b563d2da1c80atagen2026-10-16 08:25 UTC124
                } else if st.on_stack[w] {
b563d2da1c80atagen2026-10-16 08:25 UTC125
                    st.low[v] = st.low[v].min(st.index[w]);
b563d2da1c80atagen2026-10-16 08:25 UTC126
                }
b563d2da1c80atagen2026-10-16 08:25 UTC127
                continue;
b563d2da1c80atagen2026-10-16 08:25 UTC128
            }
b563d2da1c80atagen2026-10-16 08:25 UTC129
            work.pop();
b563d2da1c80atagen2026-10-16 08:25 UTC130
            if let Some(&(parent, _)) = work.last() {
b563d2da1c80atagen2026-10-16 08:25 UTC131
                st.low[parent] = st.low[parent].min(st.low[v]);
b563d2da1c80atagen2026-10-16 08:25 UTC132
            }
b563d2da1c80atagen2026-10-16 08:25 UTC133
            if st.low[v] == st.index[v] {
b563d2da1c80atagen2026-10-16 08:25 UTC134
                while let Some(w) = st.stack.pop() {
b563d2da1c80atagen2026-10-16 08:25 UTC135
                    st.on_stack[w] = false;
b563d2da1c80atagen2026-10-16 08:25 UTC136
                    comp[w] = next_comp;
b563d2da1c80atagen2026-10-16 08:25 UTC137
                    if w == v {
b563d2da1c80atagen2026-10-16 08:25 UTC138
                        break;
b563d2da1c80atagen2026-10-16 08:25 UTC139
                    }
b563d2da1c80atagen2026-10-16 08:25 UTC140
                }
b563d2da1c80atagen2026-10-16 08:25 UTC141
                next_comp += 1;
b563d2da1c80atagen2026-10-16 08:25 UTC142
            }
b563d2da1c80atagen2026-10-16 08:25 UTC143
        }
b563d2da1c80atagen2026-10-16 08:25 UTC144
    }
b563d2da1c80atagen2026-10-16 08:25 UTC145
    comp
b563d2da1c80atagen2026-10-16 08:25 UTC146
}
b563d2da1c80atagen2026-10-16 08:25 UTC147
b563d2da1c80atagen2026-10-16 08:25 UTC148
#[cfg(test)]
b563d2da1c80atagen2026-10-16 08:25 UTC149
mod tests {
b563d2da1c80atagen2026-10-16 08:25 UTC150
    use super::*;
b563d2da1c80atagen2026-10-16 08:25 UTC151
b563d2da1c80atagen2026-10-16 08:25 UTC152
    #[test]
b563d2da1c80atagen2026-10-16 08:25 UTC153
    fn unsatisfiable() {
b563d2da1c80atagen2026-10-16 08:25 UTC154
        let mut s = TwoSat::new(1);
b563d2da1c80atagen2026-10-16 08:25 UTC155
        s.require(Lit::new(0, true));
b563d2da1c80atagen2026-10-16 08:25 UTC156
        s.require(Lit::new(0, false));
b563d2da1c80atagen2026-10-16 08:25 UTC157
        assert!(s.solve().is_none());
b563d2da1c80atagen2026-10-16 08:25 UTC158
    }
b563d2da1c80atagen2026-10-16 08:25 UTC159
}