adds BFS shortest path
parent
5c7bff171c
commit
03d23173cb
@ -0,0 +1,176 @@
|
|||||||
|
/*
|
||||||
|
* Copyright 2022, The Cozo Project Authors.
|
||||||
|
*
|
||||||
|
* This Source Code Form is subject to the terms of the Mozilla Public License, v. 2.0.
|
||||||
|
* If a copy of the MPL was not distributed with this file,
|
||||||
|
* You can obtain one at https://mozilla.org/MPL/2.0/.
|
||||||
|
*/
|
||||||
|
|
||||||
|
/*
|
||||||
|
* Copyright 2022, The Cozo Project Authors.
|
||||||
|
*
|
||||||
|
* This Source Code Form is subject to the terms of the Mozilla Public License, v. 2.0.
|
||||||
|
* If a copy of the MPL was not distributed with this file,
|
||||||
|
* You can obtain one at https://mozilla.org/MPL/2.0/.
|
||||||
|
*/
|
||||||
|
|
||||||
|
use std::collections::{BTreeMap, BTreeSet, VecDeque};
|
||||||
|
|
||||||
|
use itertools::Itertools;
|
||||||
|
use miette::Result;
|
||||||
|
use smartstring::{LazyCompact, SmartString};
|
||||||
|
|
||||||
|
use crate::data::expr::Expr;
|
||||||
|
use crate::data::symb::Symbol;
|
||||||
|
use crate::data::value::DataValue;
|
||||||
|
use crate::fixed_rule::{FixedRule, FixedRulePayload};
|
||||||
|
use crate::parse::SourceSpan;
|
||||||
|
use crate::runtime::db::Poison;
|
||||||
|
use crate::runtime::temp_store::RegularTempStore;
|
||||||
|
|
||||||
|
pub(crate) struct ShortestPathBFS;
|
||||||
|
|
||||||
|
impl FixedRule for ShortestPathBFS {
|
||||||
|
fn run(
|
||||||
|
&self,
|
||||||
|
payload: FixedRulePayload<'_, '_>,
|
||||||
|
out: &mut RegularTempStore,
|
||||||
|
poison: Poison,
|
||||||
|
) -> Result<()> {
|
||||||
|
let edges = payload.get_input(0)?.ensure_min_len(2)?;
|
||||||
|
let starting_nodes: Vec<_> = payload
|
||||||
|
.get_input(1)?
|
||||||
|
.ensure_min_len(1)?
|
||||||
|
.iter()?
|
||||||
|
.map_ok(|n| n.into_iter().next().unwrap())
|
||||||
|
.try_collect()?;
|
||||||
|
let ending_nodes: BTreeSet<_> = payload
|
||||||
|
.get_input(2)?
|
||||||
|
.ensure_min_len(1)?
|
||||||
|
.iter()?
|
||||||
|
.map_ok(|n| n.into_iter().next().unwrap())
|
||||||
|
.try_collect()?;
|
||||||
|
|
||||||
|
for starting_node in starting_nodes.iter() {
|
||||||
|
let mut pending: BTreeSet<_> = ending_nodes.clone();
|
||||||
|
let mut visited: BTreeSet<DataValue> = Default::default();
|
||||||
|
let mut backtrace: BTreeMap<DataValue, DataValue> = Default::default();
|
||||||
|
|
||||||
|
visited.insert(starting_node.clone());
|
||||||
|
|
||||||
|
let mut queue: VecDeque<DataValue> = VecDeque::default();
|
||||||
|
queue.push_front(starting_node.clone());
|
||||||
|
|
||||||
|
while let Some(candidate) = queue.pop_back() {
|
||||||
|
for edge in edges.prefix_iter(&candidate)? {
|
||||||
|
let edge = edge?;
|
||||||
|
let to_node = &edge[1];
|
||||||
|
if visited.contains(to_node) {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
|
||||||
|
visited.insert(to_node.clone());
|
||||||
|
backtrace.insert(to_node.clone(), candidate.clone());
|
||||||
|
|
||||||
|
pending.remove(to_node);
|
||||||
|
|
||||||
|
if pending.is_empty() {
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
|
||||||
|
queue.push_front(to_node.clone());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
for ending_node in ending_nodes.iter() {
|
||||||
|
if backtrace.contains_key(ending_node) {
|
||||||
|
let mut route = vec![];
|
||||||
|
let mut current = ending_node.clone();
|
||||||
|
while current != *starting_node {
|
||||||
|
route.push(current.clone());
|
||||||
|
current = backtrace.get(¤t).unwrap().clone();
|
||||||
|
}
|
||||||
|
route.push(starting_node.clone());
|
||||||
|
route.reverse();
|
||||||
|
let tuple = vec![
|
||||||
|
starting_node.clone(),
|
||||||
|
ending_node.clone(),
|
||||||
|
DataValue::List(route),
|
||||||
|
];
|
||||||
|
out.put(tuple);
|
||||||
|
} else {
|
||||||
|
out.put(vec![
|
||||||
|
starting_node.clone(),
|
||||||
|
ending_node.clone(),
|
||||||
|
DataValue::Null,
|
||||||
|
])
|
||||||
|
}
|
||||||
|
}
|
||||||
|
poison.check()?;
|
||||||
|
}
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
fn arity(
|
||||||
|
&self,
|
||||||
|
_options: &BTreeMap<SmartString<LazyCompact>, Expr>,
|
||||||
|
_rule_head: &[Symbol],
|
||||||
|
_span: SourceSpan,
|
||||||
|
) -> Result<usize> {
|
||||||
|
Ok(3)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use serde_json::json;
|
||||||
|
|
||||||
|
use crate::new_cozo_mem;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_bfs_path() {
|
||||||
|
let db = new_cozo_mem().unwrap();
|
||||||
|
let res = db
|
||||||
|
.run_script(
|
||||||
|
r#"
|
||||||
|
love[loving, loved] <- [['alice', 'eve'],
|
||||||
|
['bob', 'alice'],
|
||||||
|
['eve', 'alice'],
|
||||||
|
['eve', 'bob'],
|
||||||
|
['eve', 'charlie'],
|
||||||
|
['charlie', 'eve'],
|
||||||
|
['david', 'george'],
|
||||||
|
['george', 'george']]
|
||||||
|
start[] <- [['alice']]
|
||||||
|
end[] <- [['bob']]
|
||||||
|
?[fr, to, path] <~ ShortestPathBFS(love[], start[], end[])
|
||||||
|
"#,
|
||||||
|
Default::default(),
|
||||||
|
)
|
||||||
|
.unwrap()
|
||||||
|
.rows;
|
||||||
|
println!("{:?}", res);
|
||||||
|
assert_eq!(res[0][2].as_array().unwrap().len(), 3);
|
||||||
|
let res = db
|
||||||
|
.run_script(
|
||||||
|
r#"
|
||||||
|
love[loving, loved] <- [['alice', 'eve'],
|
||||||
|
['bob', 'alice'],
|
||||||
|
['eve', 'alice'],
|
||||||
|
['eve', 'bob'],
|
||||||
|
['eve', 'charlie'],
|
||||||
|
['charlie', 'eve'],
|
||||||
|
['david', 'george'],
|
||||||
|
['george', 'george']]
|
||||||
|
start[] <- [['alice']]
|
||||||
|
end[] <- [['george']]
|
||||||
|
?[fr, to, path] <~ ShortestPathBFS(love[], start[], end[])
|
||||||
|
"#,
|
||||||
|
Default::default(),
|
||||||
|
)
|
||||||
|
.unwrap()
|
||||||
|
.rows;
|
||||||
|
assert_eq!(res[0][2], json!(null));
|
||||||
|
println!("{:?}", res);
|
||||||
|
}
|
||||||
|
}
|
Loading…
Reference in New Issue