Fix behavior for alt probe type (should only emit the first in the list)

ejrgilbert · May 10, 2024 · 6dda3f1 · 6dda3f1
1 parent 3e4da68
commit 6dda3f1
Show file tree

Hide file tree

Showing 3 changed files with 39 additions and 9 deletions.
diff --git a/src/behavior/builder_visitor.rs b/src/behavior/builder_visitor.rs
@@ -297,10 +297,16 @@ impl WhammVisitor<()> for BehaviorTreeBuilder {
         trace!("Entering: BehaviorTreeBuilder::visit_probe");
         self.context_name += &format!(":{}", probe.name.clone());
 
-        self.tree.decorator(DecoratorType::ForEach {
+        if probe.name == "alt" {
+            self.tree.decorator(DecoratorType::ForFirstProbe {
                 target: probe.name.clone()
-            })
-            .sequence()
+            });
+        } else {
+            self.tree.decorator(DecoratorType::ForEachProbe {
+                target: probe.name.clone()
+            });
+        }
+        self.tree.sequence()
                 .enter_scope(self.context_name.clone());
 
         // visit globals

diff --git a/src/behavior/tree.rs b/src/behavior/tree.rs
@@ -394,7 +394,12 @@ pub enum DecoratorType {
     PredIs {
         val: bool
     },
-    ForEach {
+    /// Iterates over all probes of the specified name in the list.
+    ForEachProbe {
+        target: String
+    },
+    /// Only pulls the first probe of the specified name from the list.
+    ForFirstProbe {
         target: String
     }
 }
@@ -449,7 +454,8 @@ pub trait BehaviorVisitor<T> {
     fn visit_is_probe_type(&mut self, node: &Node) -> T;
     fn visit_has_params(&mut self, node: &Node) -> T;
     fn visit_pred_is(&mut self, node: &Node) -> T;
-    fn visit_for_each(&mut self, node: &Node) -> T;
+    fn visit_for_each_probe(&mut self, node: &Node) -> T;
+    fn visit_for_first_probe(&mut self, node: &Node) -> T;
 
     // Parameterized action nodes
     fn visit_emit_if_else(&mut self, node: &Node) -> T;

diff --git a/src/behavior/visualize.rs b/src/behavior/visualize.rs
@@ -154,7 +154,8 @@ impl BehaviorVisitor<()> for Visualizer<'_> {
                 DecoratorType::IsProbeType {..} => self.visit_is_probe_type(node),
                 DecoratorType::HasParams {..} => self.visit_has_params(node),
                 DecoratorType::PredIs {..} => self.visit_pred_is(node),
-                DecoratorType::ForEach {..} => self.visit_for_each(node),
+                DecoratorType::ForEachProbe {..} => self.visit_for_each_probe(node),
+                DecoratorType::ForFirstProbe {..} => self.visit_for_first_probe(node),
             }
         } else {
             unreachable!()
@@ -263,10 +264,27 @@ impl BehaviorVisitor<()> for Visualizer<'_> {
         }
     }
 
-    fn visit_for_each(&mut self, node: &TreeNode) -> () {
+    fn visit_for_each_probe(&mut self, node: &TreeNode) -> () {
         if let TreeNode::Decorator { id, ty, parent, child } = node {
-            if let DecoratorType::ForEach{ target } = ty {
-                self.emit_decorator_node(id, &format!("ForEach_{}", target.replace(":", "_")));
+            if let DecoratorType::ForEachProbe { target } = ty {
+                self.emit_decorator_node(id, &format!("ForEachProbe_{}", target.replace(":", "_")));
+                self.emit_edge(parent, id);
+
+                if let Some(node) = self.tree.get_node(child.clone()) {
+                    self.visit_node(node);
+                }
+            } else {
+                unreachable!()
+            }
+        } else {
+            unreachable!()
+        }
+    }
+
+    fn visit_for_first_probe(&mut self, node: &TreeNode) -> () {
+        if let TreeNode::Decorator { id, ty, parent, child } = node {
+            if let DecoratorType::ForFirstProbe { target } = ty {
+                self.emit_decorator_node(id, &format!("ForFirstProbe_{}", target.replace(":", "_")));
                 self.emit_edge(parent, id);
 
                 if let Some(node) = self.tree.get_node(child.clone()) {