From 13c60990f7336d320cad3ff5a692d25c32db9398 Mon Sep 17 00:00:00 2001
From: Varun Agrawal <varagrawal@gmail.com>
Date: Fri, 18 Mar 2022 00:11:38 -0400
Subject: [PATCH] fix visitWith operator() to be more functional

---
 gtsam/discrete/Assignment.h       | 2 ++
 gtsam/discrete/DecisionTree-inl.h | 7 +++++--
 2 files changed, 7 insertions(+), 2 deletions(-)

diff --git a/gtsam/discrete/Assignment.h b/gtsam/discrete/Assignment.h
index cdbf0a2e9..90e2dbdd8 100644
--- a/gtsam/discrete/Assignment.h
+++ b/gtsam/discrete/Assignment.h
@@ -33,6 +33,8 @@ namespace gtsam {
 template <class L>
 class Assignment : public std::map<L, size_t> {
  public:
+  using std::map<L, size_t>::operator=;
+
   void print(const std::string& s = "Assignment: ") const {
     std::cout << s << ": ";
     for (const typename Assignment::value_type& keyValue : *this)
diff --git a/gtsam/discrete/DecisionTree-inl.h b/gtsam/discrete/DecisionTree-inl.h
index 01c7b689c..c616c0269 100644
--- a/gtsam/discrete/DecisionTree-inl.h
+++ b/gtsam/discrete/DecisionTree-inl.h
@@ -666,8 +666,11 @@ namespace gtsam {
       if (!choice)
         throw std::invalid_argument("DecisionTree::VisitWith: Invalid NodePtr");
       for (size_t i = 0; i < choice->nrChoices(); i++) {
-        choices[choice->label()] = i;    // Set assignment for label to i
-        (*this)(choice->branches()[i]);  // recurse!
+        choices[choice->label()] = i;  // Set assignment for label to i
+
+        VisitWith<L, Y> visit(f);
+        visit.choices = choices;
+        visit(choice->branches()[i]);  // recurse!
       }
     }
   };