From 84037e173f5e276929e50253729206db2f7a6a9f Mon Sep 17 00:00:00 2001
From: Frank Dellaert <dellaert@gmail.com>
Date: Sat, 7 Jan 2023 08:11:44 -0800
Subject: [PATCH] Add all discrete keys

---
 gtsam/hybrid/HybridJunctionTree.cpp | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/gtsam/hybrid/HybridJunctionTree.cpp b/gtsam/hybrid/HybridJunctionTree.cpp
index 573df7eca..a463c625f 100644
--- a/gtsam/hybrid/HybridJunctionTree.cpp
+++ b/gtsam/hybrid/HybridJunctionTree.cpp
@@ -62,8 +62,12 @@ struct HybridConstructorTraversalData {
 
     // Add all the discrete keys in the hybrid factors to the current data
     for (const auto& f : node->factors) {
-      if (auto p = boost::dynamic_pointer_cast<HybridFactor>(f)) {
-        for (auto& k : p->discreteKeys()) {
+      if (auto hf = boost::dynamic_pointer_cast<HybridFactor>(f)) {
+        for (auto& k : hf->discreteKeys()) {
+          data.discreteKeys.insert(k.first);
+        }
+      } else if (auto hf = boost::dynamic_pointer_cast<DecisionTreeFactor>(f)) {
+        for (auto& k : hf->discreteKeys()) {
           data.discreteKeys.insert(k.first);
         }
       }