use crate::common::{KeyedTriple, URI};
use datafrog::Iteration;
use disjoint_sets::UnionFind;
use std::collections::HashMap;
pub struct DisjointSets {
lists: HashMap<URI, Vec<URI>>,
uri2idx: HashMap<URI, usize>,
idx2uri: Vec<URI>,
ds: UnionFind,
}
impl DisjointSets {
pub fn new(input: &[KeyedTriple], rdffirst_node: URI, rdfrest_node: URI, rdfnil_node: URI) -> Self {
let mut lists: HashMap<URI, Vec<URI>> = HashMap::new();
let mut values: HashMap<URI, URI> = HashMap::new();
let mut uri2idx: HashMap<URI, usize> = HashMap::new();
let mut idx2uri: Vec<URI> = Vec::with_capacity(input.len());
let mut counter: usize = 0;
for &(a, (b, c)) in input.iter() {
for e in [a, b, c] {
let index = uri2idx.entry(e).or_insert(counter);
if *index == counter {
idx2uri.push(e);
counter += 1;
}
}
}
let mut ds = UnionFind::new(counter);
let mut list_iter = Iteration::new();
let rdffirst = list_iter.variable::<(URI, ())>("rdffirst");
rdffirst.extend(vec![(rdffirst_node, ())]);
let rdfrest = list_iter.variable::<(URI, ())>("rdfrest");
rdfrest.extend(vec![(rdfrest_node, ())]);
let list_triples = list_iter.variable::<KeyedTriple>("list_triples");
list_triples.extend(input.iter().map(|&(s, (p, o))| (p, (s, o))));
let rdf_firsts = list_iter.variable::<()>("rdf_firsts");
let rdf_rests = list_iter.variable::<()>("rdf_rests");
while list_iter.changed() {
rdf_firsts.from_join(&list_triples, &rdffirst, |&_, &(key, value), &_| {
values.insert(key, value);
});
rdf_rests.from_join(&list_triples, &rdfrest, |&_, &(head, tail), &_| {
if tail == rdfnil_node {
return;
}
if let (Some(&hn), Some(&tn)) = (uri2idx.get(&head), uri2idx.get(&tail)) {
ds.union(hn, tn);
}
});
}
for (key, value) in values.into_iter() {
if let Some(idx) = uri2idx.get(&key) {
let set = ds.find(*idx);
let listname = idx2uri[set];
let list = lists.entry(listname).or_insert_with(Vec::new);
list.push(value);
}
}
DisjointSets {
lists,
uri2idx,
idx2uri,
ds,
}
}
pub fn get_list_values(&self, head: URI) -> Option<&[URI]> {
if let Some(idx) = self.uri2idx.get(&head) {
let realhead = self.idx2uri[self.ds.find(*idx)];
if let Some(v) = self.lists.get(&realhead) {
return Some(v.as_slice());
}
}
None
}
}