use crate::Sampler;
use indexmap::IndexSet;
use serde::{Deserialize, Serialize};
#[derive(Debug, PartialEq, Eq, Clone, Copy, Hash, Serialize, Deserialize)]
pub struct Edge {
pub variable: usize,
pub constraint: usize,
}
impl Edge {
pub fn new(variable: usize, constraint: usize) -> Self {
Self {
variable,
constraint,
}
}
}
#[derive(Debug, PartialEq, Eq, Clone, Serialize, Deserialize)]
pub struct Graph {
variable_neighbors: Vec<IndexSet<usize>>,
constraint_neighbors: Vec<IndexSet<usize>>,
edges: IndexSet<Edge>,
}
impl Graph {
pub fn new() -> Self {
Self {
variable_neighbors: Vec::new(),
constraint_neighbors: Vec::new(),
edges: IndexSet::new(),
}
}
pub fn complete_graph(number_of_variables: usize, number_of_constraints: usize) -> Self {
let mut graph = Self::new();
for variable in 0..number_of_variables {
for constraint in 0..number_of_constraints {
graph.insert_edge(Edge::new(variable, constraint));
}
}
graph
}
pub fn contains_edge(&self, edge: Edge) -> bool {
self.edges.contains(&edge)
}
pub fn insert_edge(&mut self, edge: Edge) -> bool {
if self.edges.insert(edge) {
self.insert_variable(edge);
self.insert_constraint(edge);
true
} else {
false
}
}
fn insert_variable(&mut self, edge: Edge) {
if edge.variable >= self.number_of_variables() {
self.variable_neighbors.extend(vec![
IndexSet::new();
edge.variable - self.number_of_variables() + 1
]);
}
self.variable_neighbors[edge.variable].insert(edge.constraint);
}
fn insert_constraint(&mut self, edge: Edge) {
if edge.constraint >= self.number_of_constraints() {
self.constraint_neighbors.extend(vec![
IndexSet::new();
edge.constraint - self.number_of_constraints()
+ 1
]);
}
self.constraint_neighbors[edge.constraint].insert(edge.variable);
}
pub fn remove_edge(&mut self, edge: Edge) -> bool {
if self.edges.remove(&edge) {
self.variable_neighbors[edge.variable].remove(&edge.constraint);
self.constraint_neighbors[edge.constraint].remove(&edge.variable);
true
} else {
false
}
}
pub fn edges(&self) -> impl Iterator<Item = Edge> + '_ {
self.edges.iter().cloned()
}
pub fn number_of_variables(&self) -> usize {
self.variable_neighbors.len()
}
pub fn number_of_constraints(&self) -> usize {
self.constraint_neighbors.len()
}
pub fn number_of_edges(&self) -> usize {
self.edges.len()
}
pub fn variables(&self) -> Nodes {
Nodes {
iter: self.variable_neighbors.iter().enumerate(),
kind: NodeKind::Variable,
}
}
pub fn constraints(&self) -> Nodes {
Nodes {
iter: self.constraint_neighbors.iter().enumerate(),
kind: NodeKind::Constraint,
}
}
pub(crate) fn from_sampler(sampler: &Sampler) -> Self {
Self {
variable_neighbors: vec![
IndexSet::with_capacity(sampler.variable_degree());
sampler.number_of_variables()
],
constraint_neighbors: vec![
IndexSet::with_capacity(sampler.constraint_degree());
sampler.number_of_constraints()
],
edges: IndexSet::with_capacity(sampler.number_of_edges()),
}
}
}
#[derive(Debug, Clone)]
pub struct Nodes<'g> {
iter: std::iter::Enumerate<std::slice::Iter<'g, IndexSet<usize>>>,
kind: NodeKind,
}
impl<'g> Iterator for Nodes<'g> {
type Item = Node<'g>;
fn next(&mut self) -> Option<Self::Item> {
self.iter.next().map(|(label, neighbors)| Node {
neighbors,
label,
kind: self.kind,
})
}
}
#[derive(Debug, PartialEq, Eq, Clone)]
pub struct Node<'g> {
neighbors: &'g IndexSet<usize>,
label: usize,
kind: NodeKind,
}
impl<'g> Node<'g> {
pub fn neighbors(&self) -> &IndexSet<usize> {
self.neighbors
}
pub fn label(&self) -> usize {
self.label
}
pub fn degree(&self) -> usize {
self.neighbors.len()
}
pub fn has_neighbor(&self, label: usize) -> bool {
self.neighbors.contains(&label)
}
pub fn is_variable(&self) -> bool {
self.kind == NodeKind::Variable
}
pub fn is_constraint(&self) -> bool {
self.kind == NodeKind::Constraint
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum NodeKind {
Variable,
Constraint,
}