shrubbery/control/
mod.rs

1use crate::{traits::*, ShrubberyError, ShrubberyResult};
2use ahash::HashMap;
3use control_nodes::ControlNode;
4use decorators::StandardDecorator;
5use derive_more::From;
6
7use crate::Status;
8
9pub mod builder;
10pub mod control_nodes;
11pub mod decorators;
12pub mod manipulation;
13pub mod simple_executors;
14
15pub const ROOT_ID: CTreeNodeID = CTreeNodeID(0);
16
17#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
18pub struct ChildUpdate {
19    pub status: Status,
20    pub child_id: CTreeNodeID,
21}
22
23#[derive(Debug, Default, Clone, Copy, PartialEq, Eq, Hash)]
24pub struct CTreeNodeID(usize);
25
26impl CTreeNodeID {
27    pub fn index(&self) -> usize {
28        self.0
29    }
30}
31
32impl From<usize> for CTreeNodeID {
33    fn from(id: usize) -> Self {
34        Self(id)
35    }
36}
37
38#[derive(Debug, Clone)]
39pub struct ControlTree<D: Decorator> {
40    pub(crate) nodes: Vec<CTreeNode<D>>,
41    pub(crate) tree: HashMap<CTreeNodeID, Vec<CTreeNodeID>>,
42}
43
44pub type StdControlTree = ControlTree<StandardDecorator>;
45
46impl<D: Decorator> Default for ControlTree<D> {
47    fn default() -> Self {
48        Self::new()
49    }
50}
51
52impl<D: Decorator> std::ops::Index<CTreeNodeID> for ControlTree<D> {
53    type Output = CTreeNode<D>;
54    fn index(&self, index: CTreeNodeID) -> &Self::Output {
55        &self.nodes[index.0]
56    }
57}
58
59impl<D: Decorator> std::ops::IndexMut<CTreeNodeID> for ControlTree<D> {
60    fn index_mut(&mut self, index: CTreeNodeID) -> &mut Self::Output {
61        &mut self.nodes[index.0]
62    }
63}
64
65impl<D: Decorator> ControlTree<D> {
66    pub fn iter_control_nodes(&self) -> impl Iterator<Item = &ControlNode<D>> + '_ {
67        self.nodes.iter().filter_map(|n| n.try_as_control())
68    }
69    pub fn iter_decorators(&self) -> impl Iterator<Item = &ControlNode<D>> + '_ {
70        self.iter_control_nodes().filter(|c| c.is_decorator())
71    }
72    pub fn iter_tree(&self) -> impl Iterator<Item = (&CTreeNodeID, &Vec<CTreeNodeID>)> + '_ {
73        self.tree.iter()
74    }
75
76    /// The status of the whole control tree (reflected by the status of the root node).
77    pub fn status(&self) -> Status {
78        self[ROOT_ID].status().unwrap_or_default()
79    }
80
81    pub fn run<Hook: ExecutorHook>(&mut self, hook: &mut Hook) -> Status {
82        while self.status() == Status::Running {
83            self.run_from(ROOT_ID, hook);
84        }
85        self.status()
86    }
87
88    pub fn run_with_update_callback<Hook: ExecutorHook, Callback: UpdateCallback<D>>(
89        &mut self,
90        hook: &mut Hook,
91        cb: &mut Callback,
92    ) -> Status {
93        while self.status() == Status::Running {
94            self.run_from_with_update_callback(ROOT_ID, hook, cb);
95        }
96        self.status()
97    }
98
99    pub fn run_from<Hook: ExecutorHook>(
100        &mut self,
101        node_id: CTreeNodeID,
102        hook: &mut Hook,
103    ) -> Status {
104        self.run_from_with_update_callback(node_id, hook, &mut NoCallback)
105    }
106
107    pub fn run_from_with_update_callback<Hook: ExecutorHook, Callback: UpdateCallback<D>>(
108        &mut self,
109        node_id: CTreeNodeID,
110        hook: &mut Hook,
111        cb: &mut Callback,
112    ) -> Status {
113        let mut node_status = self[node_id].tick();
114        cb.callback(self);
115
116        while node_status.is_running() {
117            for child in self.children(&node_id) {
118                // tick the parent node & break if it's finished
119
120                if self[node_id].tick().is_terminal() {
121                    cb.callback(self);
122
123                    break;
124                }
125                if self[child].status().unwrap_or_default().is_success() {
126                    // don't re-run successful nodes
127                    continue;
128                }
129
130                if let CTreeNode::Leaf(leaf) = &self[child] {
131                    // hook the leaf node executor to get the status & update the control node with the
132                    // result
133                    let status = hook.hook(leaf);
134                    self[child].set_status(status); // update the leaf node status from the hook
135
136                    let update = ChildUpdate {
137                        status,
138                        child_id: child,
139                    };
140                    cb.callback(self);
141                    self[node_id].child_updated(update);
142                } else {
143                    // continue down the control tree, updating the control node with the eventual
144                    // result
145                    let status = self[child].tick();
146                    let subtree_status = match status {
147                        Status::Running => self.run_from_with_update_callback(child, hook, cb),
148                        _ => status,
149                    };
150                    let update = ChildUpdate {
151                        status: subtree_status,
152                        child_id: child,
153                    };
154                    self[node_id].child_updated(update);
155                }
156            }
157            // tell the node all the children have run.
158            self[node_id].all_children_seen();
159
160            node_status = self[node_id].tick();
161            self.handle_reset_requests(node_id);
162            cb.callback(self);
163        }
164        node_status
165    }
166
167    fn handle_reset_requests(&mut self, node_id: CTreeNodeID) -> usize {
168        if let Some(reset) = self[node_id]
169            .try_as_control_mut()
170            .map(|c| std::mem::take(&mut c.reset_requests))
171        {
172            reset
173                .into_iter()
174                .map(|id| {
175                    self.reset_branch(id);
176                })
177                .count()
178        } else {
179            0
180        }
181    }
182
183    pub fn reset_branch(&mut self, from: CTreeNodeID) {
184        let mut to_visit = vec![from];
185        while let Some(id) = to_visit.pop() {
186            self[id].reset();
187
188            self.tree[&id]
189                .iter()
190                .for_each(|&child| to_visit.push(child));
191        }
192    }
193
194    pub fn new() -> Self {
195        let root = CTreeNode::root();
196        let mut tree = HashMap::<CTreeNodeID, Vec<CTreeNodeID>>::default();
197        tree.insert(0.into(), vec![]);
198        Self {
199            nodes: vec![root],
200            tree,
201        }
202    }
203
204    /// Validate the tree for the following conditions, error if any are violated.
205    ///
206    /// - No cycles
207    /// - No dangling control nodes
208    /// - Decorators have only one child
209    pub(crate) fn validate_bt_rules(&self) -> ShrubberyResult<()> {
210        self.check_for_cycles()?;
211        self.validate_decorators()?;
212        self.check_for_dangling_control()?;
213        Ok(())
214    }
215
216    /// Look for cycles in the tree, returns an error if any exist.
217    pub(crate) fn check_for_cycles(&self) -> ShrubberyResult<()> {
218        if let Some(err) = self.iter_tree().find_map(|(&parent, children)| {
219            children.iter().find_map(|&child| {
220                if let Err(e) = self.recurse_children_check_cycles(child, vec![parent]) {
221                    Some(e)
222                } else {
223                    None
224                }
225            })
226        }) {
227            Err(err)
228        } else {
229            Ok(())
230        }
231    }
232
233    /// Decorators are only allowed to have a single child
234    pub(crate) fn validate_decorators(&self) -> ShrubberyResult<()> {
235        if let Some(violation) = self
236            .iter_decorators()
237            .flat_map(|d| d.id)
238            .find(|id| self.children(id).len() != 1)
239        {
240            Err(ShrubberyError::InvalidDecorator {
241                decorator: violation,
242                children: self.children(&violation),
243            })
244        } else {
245            Ok(())
246        }
247    }
248
249    /// Control nodes are by definition not leaf nodes so must have at least one child.
250    pub(crate) fn check_for_dangling_control(&self) -> ShrubberyResult<()> {
251        if let Some(dangling) = self
252            .iter_control_nodes()
253            .flat_map(|n| n.id)
254            .find(|id| self.children(id).is_empty())
255        {
256            Err(ShrubberyError::DanglingControlNode(dangling))
257        } else {
258            Ok(())
259        }
260    }
261
262    /// Recursively check for cycles in the tree, returning an error if any are found.
263    fn recurse_children_check_cycles(
264        &self,
265        from: CTreeNodeID,
266        mut history: Vec<CTreeNodeID>,
267    ) -> ShrubberyResult<()> {
268        if let Some(first) = history.first() {
269            if first == &from {
270                history.push(*first);
271                return Err(ShrubberyError::CycleDetected(history));
272            }
273        }
274        history.push(from);
275        let children = self.children(&from);
276        for child in children {
277            self.recurse_children_check_cycles(child, history.clone())?;
278        }
279        Ok(())
280    }
281
282    /// Extract a section of the [`ControlTree`] into a new one.
283    pub fn as_subtree(&self, start_at: CTreeNodeID) -> Self {
284        let mut subtree = Self::new();
285
286        let mut start = self[start_at].clone();
287        let old_id = start.id().unwrap();
288
289        start.unset_id();
290        let new_id = subtree.add_child_unchecked(ROOT_ID, start);
291
292        let mut old_to_new = HashMap::<CTreeNodeID, CTreeNodeID>::default();
293        old_to_new.insert(ROOT_ID, ROOT_ID);
294        old_to_new.insert(old_id, new_id);
295
296        struct Deps<D: Decorator> {
297            old_to_new: HashMap<CTreeNodeID, CTreeNodeID>,
298            subtree: ControlTree<D>,
299        }
300
301        let mut deps = Deps {
302            subtree,
303            old_to_new,
304        };
305
306        self.explore_down_with_deps(start_at, &mut deps, |deps, parent, children| {
307            let old_parent_id = &parent.id().unwrap();
308            let parent_id = deps.old_to_new[old_parent_id];
309
310            children.iter().for_each(|&old_id| {
311                let new_id = deps.subtree.add_floating_node(self[old_id].clone());
312                deps.subtree.tree.entry(parent_id).or_default().push(new_id);
313                deps.old_to_new.insert(old_id, new_id);
314            });
315        });
316
317        deps.subtree
318    }
319
320    fn add_floating_node(&mut self, node: impl Into<CTreeNode<D>>) -> CTreeNodeID {
321        let node = node.into();
322        let id = self.nodes.len().into();
323        self.nodes.push(node);
324        id
325    }
326
327    fn explore_down_with_deps<Deps>(
328        &self,
329        from: CTreeNodeID,
330        deps: &mut Deps,
331        f: impl Fn(&mut Deps, &CTreeNode<D>, &Vec<CTreeNodeID>),
332    ) {
333        let mut to_visit = vec![from];
334        while let Some(from) = to_visit.pop() {
335            let children = self.children(&from);
336            let parent = &self[from];
337            f(deps, parent, &children);
338            for child in children {
339                to_visit.push(child);
340            }
341        }
342    }
343
344    /// Insert `node` between `parent` and `move_down`.
345    ///
346    /// The inserted node with be the **n**'th child, where **n** is the index of the first child
347    /// node in `move_down`. i.e. left->right order is maintained if `move_down` is contiguous with
348    /// the original children.
349    ///
350    /// ## Examples
351    ///
352    /// ```rust,ignore
353    /// insert_between(0, &[2], X)
354    /// ```
355    ///
356    /// ```text
357    ///         0                                  0
358    ///       / | \            ------>           / | \
359    ///      1  2  3                            1  X  3
360    ///                                            |
361    ///                                            2
362    /// ```
363    ///
364    /// ```rust,ignore
365    /// insert_between(0, &[1, 3], X)
366    /// ```
367    ///
368    /// ```text
369    ///         0                                 0
370    ///       / | \            ------>           / \
371    ///      1  2  3                            X   2
372    ///                                        / \
373    ///                                       1   3
374    /// ```
375    pub fn insert_between(
376        &mut self,
377        parent_id: CTreeNodeID,
378        move_down: &[CTreeNodeID],
379        node: impl Into<CTreeNode<D>>,
380    ) -> CTreeNodeID {
381        let node = node.into();
382        let mut i = 0;
383        self.tree
384            .entry(parent_id)
385            .and_modify(|children| {
386                // find the index of the first child getting moved down -- this is where `node`
387                // will be inserted.
388                i = children
389                    .iter()
390                    .enumerate()
391                    .find_map(|(i, c)| if move_down.contains(c) { Some(i) } else { None })
392                    .expect("None of the children are in move_down");
393                children.retain(|v| !move_down.contains(v))
394            })
395            .or_default();
396
397        let new_id = self.add_child_unchecked(parent_id, node);
398
399        self.tree.entry(parent_id).and_modify(|children| {
400            children.pop();
401            children.insert(i, new_id);
402        });
403
404        self.tree
405            .entry(new_id)
406            .or_default()
407            .extend_from_slice(move_down);
408
409        new_id
410    }
411
412    pub fn iter_children_mut<'a, O>(
413        &'a mut self,
414        node_id: &CTreeNodeID,
415        mut f: impl FnMut(&mut CTreeNode<D>) -> O + 'a,
416    ) -> impl Iterator<Item = O> + '_ {
417        self.tree[node_id]
418            .clone()
419            .into_iter()
420            .map(move |id| f(self.node_mut(id)))
421    }
422
423    pub fn node_mut(&mut self, id: CTreeNodeID) -> &mut CTreeNode<D> {
424        &mut self.nodes[id.0]
425    }
426
427    pub fn children(&self, node_id: &CTreeNodeID) -> Vec<CTreeNodeID> {
428        self.tree[node_id].clone()
429    }
430
431    pub fn iter_children(&self, node_id: &CTreeNodeID) -> impl Iterator<Item = &CTreeNode<D>> + '_ {
432        self.tree[node_id].iter().map(|&id| &self[id])
433    }
434
435    pub fn iter_child_ids(&self, node_id: &CTreeNodeID) -> impl Iterator<Item = &CTreeNodeID> + '_ {
436        self.tree[node_id].iter()
437    }
438}
439
440#[derive(Debug, Clone, PartialEq, Eq, From)]
441pub enum CTreeNode<D: Decorator> {
442    Root(RootNode),
443    Control(ControlNode<D>),
444    Leaf(LeafNode),
445}
446
447impl<D: Decorator> CTreeNode<D> {
448    pub fn try_as_leaf_mut(&mut self) -> Option<&mut LeafNode> {
449        match self {
450            CTreeNode::Leaf(c) => Some(c),
451            _ => None,
452        }
453    }
454    pub fn try_as_leaf(&self) -> Option<&LeafNode> {
455        match &self {
456            CTreeNode::Leaf(c) => Some(c),
457            _ => None,
458        }
459    }
460    pub fn is_leaf(&self) -> bool {
461        self.try_as_leaf().is_some()
462    }
463    pub fn try_as_control_mut(&mut self) -> Option<&mut ControlNode<D>> {
464        match self {
465            CTreeNode::Control(c) => Some(c),
466            _ => None,
467        }
468    }
469
470    pub fn try_as_control(&self) -> Option<&ControlNode<D>> {
471        match &self {
472            CTreeNode::Control(c) => Some(c),
473            _ => None,
474        }
475    }
476    pub fn is_control(&self) -> bool {
477        self.try_as_control().is_some()
478    }
479
480    pub fn try_as_root(&self) -> Option<&RootNode> {
481        match &self {
482            CTreeNode::Root(c) => Some(c),
483            _ => None,
484        }
485    }
486
487    pub fn is_root(&self) -> bool {
488        self.try_as_root().is_some()
489    }
490
491    pub fn root() -> Self {
492        CTreeNode::Root(RootNode(ControlNode::sequence()))
493    }
494    pub fn leaf() -> Self {
495        CTreeNode::Leaf(LeafNode::default())
496    }
497    pub fn unset_id(&mut self) {
498        match self {
499            CTreeNode::Root(root) => root.0.id = None,
500            CTreeNode::Control(control) => control.id = None,
501            CTreeNode::Leaf(leaf) => leaf.id = None,
502        };
503    }
504    pub fn set_id(&mut self, id: CTreeNodeID) -> Option<CTreeNodeID> {
505        let old = match self {
506            CTreeNode::Root(root) => &mut root.0.id,
507            CTreeNode::Control(control) => &mut control.id,
508            CTreeNode::Leaf(leaf) => &mut leaf.id,
509        };
510        let mut id = Some(id);
511        std::mem::swap(old, &mut id);
512        id
513    }
514    pub fn id(&self) -> Option<CTreeNodeID> {
515        match self {
516            CTreeNode::Root(root) => root.0.id,
517            CTreeNode::Control(control) => control.id,
518            CTreeNode::Leaf(leaf) => leaf.id,
519        }
520    }
521    pub fn reset(&mut self) {
522        self.clear_status();
523        match self {
524            CTreeNode::Root(root) => root.0.reset(),
525            CTreeNode::Control(control) => control.reset(),
526            CTreeNode::Leaf(leaf) => leaf.reset(),
527        }
528    }
529    pub fn clear_status(&mut self) {
530        match self {
531            CTreeNode::Root(root) => root.0.status = None,
532            CTreeNode::Control(control) => control.status = None,
533            CTreeNode::Leaf(leaf) => leaf.status = None,
534        }
535    }
536    pub fn set_status(&mut self, status: Status) {
537        match self {
538            CTreeNode::Root(root) => root.0.status = Some(status),
539            CTreeNode::Control(control) => control.status = Some(status),
540            CTreeNode::Leaf(leaf) => leaf.status = Some(status),
541        }
542    }
543    pub fn status(&self) -> Option<Status> {
544        match self {
545            CTreeNode::Root(root) => root.0.status,
546            CTreeNode::Control(control) => control.status,
547            CTreeNode::Leaf(leaf) => leaf.status,
548        }
549    }
550}
551
552impl<D: Decorator> Control for CTreeNode<D> {
553    fn tick(&mut self) -> Status {
554        let status = match self {
555            CTreeNode::Root(root) => root.tick(),
556            CTreeNode::Control(control) => control.tick(),
557            CTreeNode::Leaf(leaf) => leaf.tick(),
558        };
559        self.set_status(status);
560        status
561    }
562    fn child_updated(&mut self, update: ChildUpdate) {
563        match self {
564            CTreeNode::Root(root) => root.child_updated(update),
565            CTreeNode::Control(control) => control.child_updated(update),
566            CTreeNode::Leaf(leaf) => leaf.child_updated(update),
567        }
568    }
569
570    fn all_children_seen(&mut self) {
571        match self {
572            CTreeNode::Root(r) => r.all_children_seen(),
573            CTreeNode::Control(c) => c.all_children_seen(),
574            CTreeNode::Leaf(l) => l.all_children_seen(),
575        }
576    }
577}
578
579#[derive(Debug, Clone, PartialEq, Eq)]
580pub struct RootNode(pub ControlNode<StandardDecorator>);
581
582impl Control for RootNode {
583    fn tick(&mut self) -> Status {
584        self.0.tick()
585    }
586    fn child_updated(&mut self, update: ChildUpdate) {
587        self.0.child_updated(update)
588    }
589    fn all_children_seen(&mut self) {
590        self.0.all_children_seen()
591    }
592}
593
594#[derive(Debug, Default, Clone, PartialEq, Eq, Hash)]
595pub struct LeafNode {
596    pub id: Option<CTreeNodeID>,
597    pub status: Option<Status>,
598    pub details: Option<String>,
599    pub name: Option<String>,
600    pub leaf_type: LeafType,
601}
602
603#[derive(Debug, Default, Clone, PartialEq, Eq, Hash)]
604pub enum LeafType {
605    #[default]
606    Unknown,
607    Conditional,
608    Executor,
609}
610
611impl LeafNode {
612    pub fn from_executor<BB: Blackboard, E: Executor<BB>>(executor: &E) -> Self {
613        LeafNode {
614            details: executor.details(),
615            name: executor.name(),
616            leaf_type: LeafType::Executor,
617            ..Default::default()
618        }
619    }
620
621    pub fn from_conditional<BB: Blackboard, C: Conditional<BB>>(conditional: &C) -> Self {
622        LeafNode {
623            details: conditional.details(),
624            name: conditional.name(),
625            leaf_type: LeafType::Conditional,
626            ..Default::default()
627        }
628    }
629    pub fn reset(&mut self) {
630        self.status = None;
631    }
632}
633
634impl Control for LeafNode {
635    fn tick(&mut self) -> Status {
636        self.status.unwrap_or_default()
637    }
638    fn child_updated(&mut self, _: ChildUpdate) {
639        panic!("Leaf nodes should not have children");
640    }
641}
shrubbery/control/mod.rs

shrubbery/control/
mod.rs
