fix bug that could cause us to dereference uninitialized tree node po…

…inter
interpretml · Nov 18, 2024 · 77b339e · 77b339e
1 parent 4e902d1
commit 77b339e
Showing 1 changed file with 25 additions and 24 deletions.
diff --git a/shared/libebm/PartitionTwoDimensionalBoosting.cpp b/shared/libebm/PartitionTwoDimensionalBoosting.cpp
@@ -385,7 +385,7 @@ template<bool bHessian, size_t cCompilerScores> class PartitionTwoDimensionalBoo
       // TODO: we should calculate the default partial gain before splitting anything. Once we have that
       // number we can calculate the minimum gain we need to reach k_gainMin after the cuts are made
       // which will allow us to avoid some work when the eventual gain will be less than our minimum
-      FloatCalc bestGain = k_gainMin; // do not allow bad cuts that lead to negative gain
+      FloatCalc bestGain = -std::numeric_limits<double>::infinity(); // do not allow bad cuts that lead to negative gain
 
       EBM_ASSERT(std::numeric_limits<FloatCalc>::min() <= hessianMin);
 
@@ -590,6 +590,7 @@ template<bool bHessian, size_t cCompilerScores> class PartitionTwoDimensionalBoo
                         pBinsEndDebug
 #endif // NDEBUG
                   );
+
                   if(Error_None != error) {
                      return error;
                   }
@@ -812,29 +813,6 @@ template<bool bHessian, size_t cCompilerScores> class PartitionTwoDimensionalBoo
       }
    done:;
 
-      auto* pCurTreeNode = pRootTreeNode;
-      EBM_ASSERT(nullptr == pCurTreeNode->GetParent());
-      while(true) {
-         EBM_ASSERT(nullptr != pCurTreeNode->GetChildren());
-         const size_t cBytesOffset1 =
-               reinterpret_cast<char*>(pCurTreeNode->GetChildren()) - reinterpret_cast<char*>(pDeepTreeNode);
-         TreeNodeMulti<bHessian, GetArrayScores(cCompilerScores)>* const pNode1 =
-               IndexTreeNodeMulti(pRootTreeNode, cBytesOffset1);
-         pCurTreeNode->SetChildren(pNode1);
-
-         pCurTreeNode = IndexTreeNodeMulti(pCurTreeNode, cBytesTreeNodeMulti);
-         if(pDeepTreeNode == pCurTreeNode) {
-            break;
-         }
-
-         EBM_ASSERT(nullptr != pCurTreeNode->GetParent());
-         const size_t cBytesOffset2 =
-               reinterpret_cast<char*>(pCurTreeNode->GetParent()) - reinterpret_cast<char*>(pDeepTreeNode);
-         TreeNodeMulti<bHessian, GetArrayScores(cCompilerScores)>* const pNode2 =
-               IndexTreeNodeMulti(pRootTreeNode, cBytesOffset2);
-         pCurTreeNode->SetParent(pNode2);
-      }
-
       EBM_ASSERT(std::isnan(bestGain) || k_illegalGainFloat == bestGain || FloatCalc{0} <= bestGain);
 
       // the bin before the aAuxiliaryBins is the last summation bin of aBinsBase,
@@ -900,6 +878,29 @@ template<bool bHessian, size_t cCompilerScores> class PartitionTwoDimensionalBoo
                if(LIKELY(k_gainMin <= bestGain)) {
                   *pTotalGain = static_cast<double>(bestGain);
 
+                  auto* pCurTreeNode = pRootTreeNode;
+                  EBM_ASSERT(nullptr == pCurTreeNode->GetParent());
+                  while(true) {
+                     EBM_ASSERT(nullptr != pCurTreeNode->GetChildren());
+                     const size_t cBytesOffset1 = reinterpret_cast<char*>(pCurTreeNode->GetChildren()) -
+                           reinterpret_cast<char*>(pDeepTreeNode);
+                     TreeNodeMulti<bHessian, GetArrayScores(cCompilerScores)>* const pNode1 =
+                           IndexTreeNodeMulti(pRootTreeNode, cBytesOffset1);
+                     pCurTreeNode->SetChildren(pNode1);
+
+                     pCurTreeNode = IndexTreeNodeMulti(pCurTreeNode, cBytesTreeNodeMulti);
+                     if(pDeepTreeNode == pCurTreeNode) {
+                        break;
+                     }
+
+                     EBM_ASSERT(nullptr != pCurTreeNode->GetParent());
+                     const size_t cBytesOffset2 =
+                           reinterpret_cast<char*>(pCurTreeNode->GetParent()) - reinterpret_cast<char*>(pDeepTreeNode);
+                     TreeNodeMulti<bHessian, GetArrayScores(cCompilerScores)>* const pNode2 =
+                           IndexTreeNodeMulti(pRootTreeNode, cBytesOffset2);
+                     pCurTreeNode->SetParent(pNode2);
+                  }
+
                   error = MakeTensor<bHessian, cCompilerScores, cCompilerDimensions>(cScores,
                         cRealDimensions,
                         flags,