use std::collections::{HashMap, HashSet};
use std::rc::Rc;
use arora_behavior::graph::{
source_port, Graph, GraphDiff, GraphType, Io, Link, LinkSource, Node as GraphNode, Port,
};
use arora_behavior::{
interpreter_module, BehaviorContext, BehaviorError, BehaviorInterpreter, BehaviorStatus,
RunPolicy, TaskHandle, TaskId,
};
use arora_types::call::Call;
use arora_types::data::{DataStore, Key, Slot, StateChange};
use arora_types::value::Value;
use arora_types::value_serde;
use uuid::Uuid;
use crate::arora_generated::behavior_tree::status::Status;
use crate::graph::build_behavior_tree;
use crate::is_native;
use crate::nodes::{
PARALLEL_FUNCTION_ID, RUN_CALL_FUNCTION_ID, RUN_CALL_PARAM_ID, RUN_STATUS_FUNCTION_ID,
RUN_STATUS_LATCH_PARAM_ID, RUN_STATUS_OUT_PARAM_ID,
};
use crate::{lower_behavior_tree, schema_groot, LoweredTree, ModuleFunction};
pub const GRAPH_TYPE_NAME: &str = "behavior-tree";
pub const GRAPH_TYPE_VERSION: semver::Version = semver::Version::new(1, 0, 0);
pub fn graph_type() -> GraphType {
GraphType {
name: GRAPH_TYPE_NAME.to_string(),
version: GRAPH_TYPE_VERSION,
}
}
struct Run {
id: TaskId,
decorator: Uuid,
status_key: Key,
variables: Vec<Uuid>,
}
pub struct BehaviorTreeInterpreter {
graph: Graph,
runner: Uuid,
main_root: Option<Uuid>,
runs: Vec<Run>,
finished: Vec<Run>,
halted: Vec<Key>,
function_index: Rc<HashMap<Uuid, ModuleFunction>>,
lowered: Option<LoweredTree>,
dirty: bool,
}
fn direct_resolver(store: &dyn DataStore) -> impl Fn(&str) -> Option<Box<dyn Slot>> + '_ {
move |name: &str| Some(store.slot(&Key::from(name)))
}
fn subtree(graph: &Graph, from: Uuid) -> HashSet<Uuid> {
reach(graph, [from])
}
fn source_variable(source: &LinkSource) -> Option<Uuid> {
match source {
LinkSource::Variable(id) => Some(*id),
LinkSource::Select { source, .. } => source_variable(source),
_ => None,
}
}
fn check_tree(graph: &Graph, root: Uuid) -> Result<(), BehaviorError> {
let mut parent: HashMap<Uuid, Uuid> = HashMap::new();
for node in graph.nodes.values() {
for child in node.children.iter().flatten() {
if !graph.nodes.contains_key(child) {
return Err(behavior_error(format!(
"node {} has child {child}, which is not one of the graph's nodes",
node.id
)));
}
if *child == root {
return Err(behavior_error(format!(
"node {} has the graph's root {root} as a child",
node.id
)));
}
if let Some(other) = parent.insert(*child, node.id) {
return Err(behavior_error(format!(
"node {child} is a child of both {other} and {}",
node.id
)));
}
}
}
if let Some(stray) = graph
.nodes
.keys()
.find(|id| **id != root && !parent.contains_key(id))
{
return Err(behavior_error(format!(
"node {stray} is not under the graph's root {root}"
)));
}
if subtree(graph, root).len() != graph.nodes.len() {
return Err(behavior_error("the graph's children form a cycle"));
}
Ok(())
}
fn check_acyclic(graph: &Graph, root: Uuid) -> Result<(), BehaviorError> {
let mut done = HashSet::new();
let mut open = HashSet::from([root]);
let mut stack = vec![(root, 0usize)];
while let Some((id, next)) = stack.last_mut() {
let children = graph.nodes[id].children.as_deref().unwrap_or_default();
let Some(child) = children.get(*next).copied() else {
open.remove(id);
done.insert(*id);
stack.pop();
continue;
};
*next += 1;
if !graph.nodes.contains_key(&child) {
return Err(behavior_error(format!(
"node {id} has child {child}, which is not a node"
)));
}
if open.contains(&child) {
return Err(behavior_error(format!(
"node {child} is its own descendant"
)));
}
if !done.contains(&child) {
open.insert(child);
stack.push((child, 0));
}
}
Ok(())
}
fn reach(graph: &Graph, starts: impl IntoIterator<Item = Uuid>) -> HashSet<Uuid> {
let mut found = HashSet::new();
let mut pending: Vec<Uuid> = starts.into_iter().collect();
while let Some(id) = pending.pop() {
if !found.insert(id) {
continue;
}
let children = graph.nodes.get(&id).and_then(|node| node.children.as_ref());
pending.extend(children.into_iter().flatten());
}
found
}
fn behavior_error(message: impl Into<String>) -> BehaviorError {
BehaviorError {
message: message.into(),
}
}
impl BehaviorTreeInterpreter {
pub fn new(function_index: Rc<HashMap<Uuid, ModuleFunction>>) -> Self {
let runner = Uuid::new_v4();
let mut graph = Graph::empty();
graph.nodes.insert(
runner,
GraphNode {
id: runner,
function: PARALLEL_FUNCTION_ID,
children: Some(Vec::new()),
..GraphNode::default()
},
);
graph.root = Some(runner);
Self {
graph,
runner,
main_root: None,
runs: Vec::new(),
finished: Vec::new(),
halted: Vec::new(),
function_index,
lowered: None,
dirty: true,
}
}
pub fn load_groot(&mut self, xml: &str) -> Result<(), crate::error::BehaviorTreeError> {
let groot = schema_groot::BehaviorTree::try_from_groot_xml(xml)?;
let graph = groot.into_graph(self.function_index.as_ref())?;
self.load(graph)
.map_err(|e| crate::error::BehaviorTreeError::InconsistentTreeError {
message: e.to_string(),
})
}
pub fn graph(&self) -> &Graph {
&self.graph
}
fn runner_node(&self) -> GraphNode {
let mut children = Vec::with_capacity(1 + self.runs.len() + self.finished.len());
children.extend(self.main_root);
children.extend(self.runs.iter().map(|run| run.decorator));
children.extend(self.finished.iter().map(|run| run.decorator));
GraphNode {
id: self.runner,
function: PARALLEL_FUNCTION_ID,
children: Some(children),
..GraphNode::default()
}
}
fn all_runs(&self) -> impl Iterator<Item = &Run> {
self.runs.iter().chain(&self.finished)
}
fn decorated_by(&self, id: &Uuid) -> Option<&Run> {
self.all_runs().find(|run| run.decorator == *id)
}
fn run_nodes(&self) -> HashSet<Uuid> {
self.all_runs()
.flat_map(|run| subtree(&self.graph, run.decorator))
.collect()
}
fn run_variables(&self) -> HashSet<Uuid> {
let nodes = self.run_nodes();
self.graph
.links
.iter()
.filter(|link| nodes.contains(&link.target.node))
.filter_map(|link| source_variable(&link.source))
.collect()
}
fn check_renames(
&self,
variables: &HashMap<Uuid, String>,
in_use: &HashSet<Uuid>,
whose: &str,
) -> Result<(), BehaviorError> {
for (id, name) in variables {
match self.graph.variables.get(id) {
Some(declared) if declared != name && in_use.contains(id) => {
return Err(behavior_error(format!(
"variable {id} is declared as '{declared}' and used by {whose}: it \
cannot become '{name}'"
)));
}
_ => {}
}
}
Ok(())
}
fn edit(&mut self, diff: GraphDiff) -> Result<(), BehaviorError> {
let mut graph = self.graph.clone();
graph
.apply(diff)
.map_err(|e| behavior_error(format!("graph diff: {e}")))?;
check_acyclic(&graph, self.runner)?;
self.graph = graph;
self.dirty = true;
Ok(())
}
fn graft(
&mut self,
module: &str,
function: Uuid,
root: Uuid,
mut fragment: GraphDiff,
) -> Result<TaskHandle, BehaviorError> {
let task = TaskId(Uuid::new_v4());
let prefix = format!("arora/tasks/{module}/{function}/{}", task.0);
let status_key = Key::from(format!("{prefix}/status"));
let variables = fragment.variables.keys().copied().collect();
let decorator = Uuid::new_v4();
fragment.add_nodes.push(GraphNode {
id: decorator,
function: RUN_STATUS_FUNCTION_ID,
inputs: vec![Io::new(RUN_STATUS_LATCH_PARAM_ID)],
outputs: vec![Io {
predetermined_key: Some(status_key.path.clone()),
..Io::new(RUN_STATUS_OUT_PARAM_ID)
}],
children: Some(vec![root]),
});
fragment.add_links.push(Link::new(
Port::new(decorator, RUN_STATUS_LATCH_PARAM_ID),
LinkSource::Literal(Value::Unit),
));
self.runs.push(Run {
id: task,
decorator,
status_key: status_key.clone(),
variables,
});
fragment.add_nodes.push(self.runner_node());
self.edit(fragment).inspect_err(|_| {
self.runs.pop();
})?;
Ok(TaskHandle {
id: task,
stop: interpreter_module::encode_halt(task),
status: status_key,
feedback: vec![Key::from(format!("{prefix}/feedback"))],
result: vec![Key::from(format!("{prefix}/result"))],
update: vec![Key::from(format!("{prefix}/update"))],
})
}
fn prune(&mut self, runs: Vec<Run>) -> Result<(), BehaviorError> {
if runs.is_empty() {
return Ok(());
}
let leaving = reach(&self.graph, runs.iter().map(|run| run.decorator));
let staying = reach(
&self.graph,
self.graph
.nodes
.keys()
.filter(|id| **id != self.runner && !leaving.contains(id))
.copied(),
);
let diff = GraphDiff {
remove_nodes: leaving.difference(&staying).copied().collect(),
add_nodes: vec![self.runner_node()],
..GraphDiff::default()
};
self.edit(diff)?;
let read: HashSet<Uuid> = self
.graph
.links
.iter()
.filter_map(|link| source_variable(&link.source))
.collect();
for id in runs.iter().flat_map(|run| &run.variables) {
if !read.contains(id) {
self.graph.variables.remove(id);
}
}
Ok(())
}
fn write_halts(&mut self, store: &dyn DataStore) -> Result<(), BehaviorError> {
for key in std::mem::take(&mut self.halted) {
let failure: Value = Status::Failure.into();
store
.write(StateChange::set(key.clone(), failure))
.map_err(|e| behavior_error(format!("{key}: writing halt status: {e}")))?;
}
Ok(())
}
fn prune_finished(&mut self) -> Result<(), BehaviorError> {
let finished = std::mem::take(&mut self.finished);
self.prune(finished)
}
fn sweep_terminal_runs(&mut self, store: &dyn DataStore) {
let terminal = |run: &Run| {
let value = store.read(std::slice::from_ref(&run.status_key));
matches!(
value.first().and_then(|v| v.clone()).map(Status::try_from),
Some(Ok(Status::Success)) | Some(Ok(Status::Failure))
)
};
let (finished, live) = std::mem::take(&mut self.runs)
.into_iter()
.partition::<Vec<_>, _>(terminal);
self.runs = live;
self.dirty |= !finished.is_empty();
self.finished.extend(finished);
}
}
impl BehaviorInterpreter for BehaviorTreeInterpreter {
fn tick(&mut self, ctx: &mut BehaviorContext) -> Result<BehaviorStatus, BehaviorError> {
self.write_halts(ctx.store)?;
if self.dirty {
self.prune_finished()?;
if let Some(lowered) = self.lowered.take() {
lowered.unregister(ctx.call_bridge);
}
let tree = build_behavior_tree(&self.graph, &direct_resolver(ctx.store))
.map_err(|e| behavior_error(format!("lowering the scaffold: {e:?}")))?;
self.lowered = Some(
lower_behavior_tree(tree, self.function_index.clone(), ctx.call_bridge, false)
.map_err(|e| behavior_error(format!("registering the scaffold: {e:?}")))?,
);
self.dirty = false;
}
if let Some(lowered) = &self.lowered {
lowered
.tick(ctx.call_bridge)
.map_err(|e| behavior_error(format!("behavior tree: {e:?}")))?;
}
self.sweep_terminal_runs(ctx.store);
Ok(BehaviorStatus::Running)
}
fn apply(&mut self, diff: GraphDiff) -> Result<(), BehaviorError> {
self.prune_finished()?;
let mut diff = diff;
let scaffold = |id: &Uuid| *id == self.runner || self.decorated_by(id).is_some();
if let Some(id) = diff
.set_root
.iter()
.chain(
diff.add_nodes
.iter()
.flat_map(|node| node.children.iter().flatten()),
)
.find(|id| scaffold(id))
{
return Err(behavior_error(format!(
"node {id} is the runner or a task run's decorator: an edit cannot place it"
)));
}
let touched = diff
.add_nodes
.iter()
.map(|node| node.id)
.chain(diff.remove_nodes.iter().copied())
.chain(diff.add_links.iter().map(|link| link.target.node))
.chain(diff.remove_links.iter().map(|port| port.node))
.chain(diff.set_predetermined.iter().map(|(port, _)| port.node));
for id in touched {
if let Some(run) = self.decorated_by(&id) {
return Err(behavior_error(format!(
"node {id} is the decorator of task run {}: an edit cannot change it",
run.id.0
)));
}
}
for run in self.all_runs() {
let roots = self.graph.nodes[&run.decorator].children.iter().flatten();
if let Some(root) = roots.into_iter().find(|root| {
diff.remove_nodes.contains(root) && !diff.add_nodes.iter().any(|n| n.id == **root)
}) {
return Err(behavior_error(format!(
"node {root} is the root of task run {}: halt the run to remove it",
run.id.0
)));
}
}
self.check_renames(&diff.variables, &self.run_variables(), "a task run")?;
diff.remove_nodes.retain(|id| *id != self.runner);
if self.main_root.is_some_and(|root| {
diff.remove_nodes.contains(&root) && !diff.add_nodes.iter().any(|n| n.id == root)
}) {
self.main_root = None;
}
if let Some(root) = diff.set_root.take() {
self.main_root = Some(root);
}
diff.add_nodes.push(self.runner_node());
self.edit(diff)
}
fn load(&mut self, graph: Graph) -> Result<(), BehaviorError> {
self.prune_finished()?;
let run_nodes = self.run_nodes();
if let Some(id) = graph
.nodes
.keys()
.find(|id| **id == self.runner || run_nodes.contains(id))
{
return Err(behavior_error(format!(
"node {id} belongs to the runner or a task run: a loaded behavior cannot \
replace it"
)));
}
self.check_renames(&graph.variables, &self.run_variables(), "a task run")?;
let remove_nodes = self
.graph
.nodes
.keys()
.filter(|id| **id != self.runner && !run_nodes.contains(id))
.copied()
.collect();
let mut diff = GraphDiff {
remove_nodes,
..GraphDiff::default()
};
self.main_root = graph.root;
let load = GraphDiff::load(graph);
diff.add_nodes = load.add_nodes;
diff.add_links = load.add_links;
diff.variables = load.variables;
diff.add_nodes.push(self.runner_node());
self.edit(diff)
}
fn spawn(&mut self, call: Call, policy: RunPolicy) -> Result<TaskHandle, BehaviorError> {
let _ = policy;
let module = call
.module_id
.map(|m| m.to_string())
.unwrap_or_else(|| "none".to_string());
let function = call.id;
let call_value = value_serde::to_value(&call).map_err(|e| {
behavior_error(format!("the spawned call does not convert to a value: {e}"))
})?;
self.prune_finished()?;
let call_node = Uuid::new_v4();
let fragment = GraphDiff {
add_nodes: vec![GraphNode {
id: call_node,
function: RUN_CALL_FUNCTION_ID,
inputs: vec![Io::new(RUN_CALL_PARAM_ID)],
..GraphNode::default()
}],
add_links: vec![Link::new(
Port::new(call_node, RUN_CALL_PARAM_ID),
LinkSource::Literal(call_value),
)],
..GraphDiff::default()
};
self.graft(&module, function, call_node, fragment)
}
fn spawn_graph(
&mut self,
graph_type: &GraphType,
graph: Graph,
policy: RunPolicy,
) -> Result<TaskHandle, BehaviorError> {
let _ = policy;
let reads = self::graph_type();
if !reads.reads(graph_type) {
return Err(behavior_error(format!(
"the behavior-tree interpreter reads {reads}, not {graph_type}"
)));
}
let root = graph
.root
.ok_or_else(|| behavior_error("a spawned graph needs a root"))?;
if !graph.nodes.contains_key(&root) {
return Err(behavior_error(format!(
"the spawned graph's root {root} is not one of its nodes"
)));
}
check_tree(&graph, root)?;
if let Some(node) = graph.nodes.values().find(|node| {
!is_native(node.function) && !self.function_index.contains_key(&node.function)
}) {
return Err(behavior_error(format!(
"node {} calls function {}, which the device does not have",
node.id, node.function
)));
}
if let Some(link) = graph.links.iter().find(|link| {
!graph.nodes.contains_key(&link.target.node)
|| source_port(&link.source)
.is_some_and(|port| !graph.nodes.contains_key(&port.node))
}) {
return Err(behavior_error(format!(
"the spawned graph's link onto {}.{} reaches outside its nodes",
link.target.node, link.target.port
)));
}
build_behavior_tree(&graph, &|_| None)
.map_err(|e| behavior_error(format!("the spawned graph does not lower: {e:?}")))?;
self.prune_finished()?;
if let Some(id) = graph
.nodes
.keys()
.find(|id| self.graph.nodes.contains_key(id))
{
return Err(behavior_error(format!(
"node {id} is already part of the device's behavior"
)));
}
let declared = self.graph.variables.keys().copied().collect();
self.check_renames(&graph.variables, &declared, "the device's behavior")?;
let load = GraphDiff::load(graph);
let fragment = GraphDiff {
add_nodes: load.add_nodes,
add_links: load.add_links,
variables: load.variables,
..GraphDiff::default()
};
self.graft(
&interpreter_module::ID.to_string(),
interpreter_module::SPAWN_GRAPH,
root,
fragment,
)
}
fn halt(&mut self, task: TaskId) -> Result<(), BehaviorError> {
let Some(at) = self.runs.iter().position(|run| run.id == task) else {
return Ok(());
};
let run = self.runs.remove(at);
let status_key = run.status_key.clone();
self.prune(vec![run])?;
self.halted.push(status_key);
Ok(())
}
}