adds BFS shortest path
Ziyang Hu
2 years ago
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