diff --git a/Modules/Segmentation/LevelSets/include/itkParallelSparseFieldLevelSetImageFilter.hxx b/Modules/Segmentation/LevelSets/include/itkParallelSparseFieldLevelSetImageFilter.hxx index 28e77d0e7d99..987ab3995414 100644 --- a/Modules/Segmentation/LevelSets/include/itkParallelSparseFieldLevelSetImageFilter.hxx +++ b/Modules/Segmentation/LevelSets/include/itkParallelSparseFieldLevelSetImageFilter.hxx @@ -772,8 +772,14 @@ ParallelSparseFieldLevelSetImageFilter::ThreadedAlloc m_Data[ThreadId].m_Semaphore[0] = 0; m_Data[ThreadId].m_Semaphore[1] = 0; + const std::size_t bufferLayerSize = 2 * m_NumberOfLayers + 1; + if (bufferLayerSize > PTRDIFF_MAX) + { + itkExceptionMacro("Number of requested layers is too large " << m_NumberOfLayers + << ". Cannot allocate enough buffer memory"); + } // Allocate the layers for the sparse field. - m_Data[ThreadId].m_Layers.reserve(2 * m_NumberOfLayers + 1); + m_Data[ThreadId].m_Layers.reserve(bufferLayerSize); for (unsigned int i = 0; i < 2 * static_cast(m_NumberOfLayers) + 1; ++i) { m_Data[ThreadId].m_Layers.push_back(LayerType::New()); @@ -785,7 +791,8 @@ ParallelSparseFieldLevelSetImageFilter::ThreadedAlloc } // Layers used as buffers for transferring pixels during load balancing - m_Data[ThreadId].m_LoadTransferBufferLayers = new LayerListType[2 * m_NumberOfLayers + 1]; + + m_Data[ThreadId].m_LoadTransferBufferLayers = new LayerListType[bufferLayerSize]; for (unsigned int i = 0; i < 2 * static_cast(m_NumberOfLayers) + 1; ++i) { m_Data[ThreadId].m_LoadTransferBufferLayers[i].reserve(m_NumOfWorkUnits); @@ -802,8 +809,7 @@ ParallelSparseFieldLevelSetImageFilter::ThreadedAlloc // The SAFETY_FACTOR simple ensures that the number of nodes created // is larger than those required to start with for each thread. - auto nodeNum = - static_cast(SAFETY_FACTOR * m_Layers[0]->Size() * (2 * m_NumberOfLayers + 1) / m_NumOfWorkUnits); + auto nodeNum = static_cast(SAFETY_FACTOR * m_Layers[0]->Size() * (bufferLayerSize) / m_NumOfWorkUnits); m_Data[ThreadId].m_LayerNodeStore->Reserve(nodeNum); m_Data[ThreadId].m_RMSChange = m_ValueZero;