remove potential pointer access issue in Flatten function

interpretml · Dec 25, 2024 · 2f70efa · 2f70efa
1 parent 15cb6fb
commit 2f70efa
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/shared/libebm/PartitionOneDimensionalBoosting.cpp b/shared/libebm/PartitionOneDimensionalBoosting.cpp
@@ -162,6 +162,7 @@ static ErrorEbm Flatten(BoosterShell* const pBoosterShell,
    FloatScore* const aUpdateScore = pInnerTermUpdate->GetTensorScoresPointer();
    FloatScore* pUpdateScore;
 
+   UIntSplit prev = 0;
    const Bin<FloatMain, UIntMain, true, true, bHessian>* const* ppBinCur = nullptr;
    if(bNominal) {
       UIntSplit iSplit = 1;
@@ -177,6 +178,7 @@ static ErrorEbm Flatten(BoosterShell* const pBoosterShell,
          // always put a split on the missing bin
          *pSplit = 1;
          ++pSplit;
+         prev = 1;
 
          // pUpdateScore is overwritten later if bNominal
          pUpdateScore += cScores;
@@ -265,7 +267,8 @@ static ErrorEbm Flatten(BoosterShell* const pBoosterShell,
                   ++iEdge; // missing is at index 0 in the model, so we are offset by one
                   pMissingBin = pTreeNode->GetBin();
                   EBM_ASSERT(iEdge <= cBins + 1);
-                  if(cBins + 1 == iEdge && cBins == *(pSplit - 1)) {
+                  EBM_ASSERT(0 != prev);
+                  if(cBins + 1 == iEdge && cBins == prev) {
                      // this cut would isolate the missing bin, but we handle those scores separately
                      break;
                   }
@@ -368,7 +371,8 @@ static ErrorEbm Flatten(BoosterShell* const pBoosterShell,
                      }
 
                      EBM_ASSERT(!IsConvertError<UIntSplit>(iEdge));
-                     *pSplit = static_cast<UIntSplit>(iEdge);
+                     prev = static_cast<UIntSplit>(iEdge);
+                     *pSplit = prev;
                      ++pSplit;
 
                      break;