Use native bucket scanner also in EditableMappingService

webknossos-datastore/app/com/scalableminds/webknossos/datastore/helpers/NativeBucketScanner.scala

@@ -4,5 +4,8 @@ import com.github.sbt.jni.nativeLoader
 
 @nativeLoader("webknossosJni0")
 class NativeBucketScanner() {
-  @native def collectSegmentIds(bucketBytes: Array[Byte], bytesPerElement: Int, isSigned: Boolean): Array[Long]
+  @native def collectSegmentIds(bucketBytes: Array[Byte],
+                                bytesPerElement: Int,
+                                isSigned: Boolean,
+                                skipZeroes: Boolean): Array[Long]
 }

webknossos-jni/src/bucketScanner.cpp

@@ -58,7 +58,7 @@ size_t getElementCount(jsize inputLengthBytes, jint bytesPerElement) {
 }
 
 JNIEXPORT jlongArray JNICALL Java_com_scalableminds_webknossos_datastore_helpers_NativeBucketScanner_collectSegmentIds(
-    JNIEnv *env, jobject instance, jbyteArray bucketBytesJavaArray, jint bytesPerElement, jboolean isSigned) {
+    JNIEnv *env, jobject instance, jbyteArray bucketBytesJavaArray, jint bytesPerElement, jboolean isSigned, jboolean skipZeroes) {
 
     const jsize inputLengthBytes = env->GetArrayLength(bucketBytesJavaArray);
     jbyte *bucketBytes = env->GetByteArrayElements(bucketBytesJavaArray, nullptr);
@@ -69,7 +69,7 @@ JNIEXPORT jlongArray JNICALL Java_com_scalableminds_webknossos_datastore_helpers
 
         for (size_t i = 0; i < elementCount; ++i) {
             const int64_t currentValue = segmentIdAtIndex(bucketBytes, i, bytesPerElement, isSigned);
-            if (currentValue != 0) {
+            if (!skipZeroes || currentValue != 0) {
                 uniqueSegmentIds.insert(currentValue);
             }
         }

webknossos-tracingstore/app/com/scalableminds/webknossos/tracingstore/tracings/editablemapping/EditableMappingLayer.scala

@@ -59,14 +59,13 @@ class EditableMappingBucketProvider(layer: EditableMappingLayer)
         additionalCoordinates = readInstruction.bucket.additionalCoordinates
       )
       unmappedData <- editableMappingService.getFallbackBucketFromDataStore(remoteFallbackLayer, dataRequest)(ec, tc)
-      unmappedDataTyped <- editableMappingService.bytesToSegmentInt(unmappedData, layer.tracing.elementClass)
-      segmentIds = editableMappingService.collectSegmentIds(unmappedDataTyped)
+      segmentIds <- editableMappingService.collectSegmentIds(unmappedData, layer.elementClass).toFox
       relevantMapping <- editableMappingService.generateCombinedMappingForSegmentIds(segmentIds,
                                                                                      editableMappingInfo,
                                                                                      version,
                                                                                      layer.tracingId,
                                                                                      remoteFallbackLayer)(tc)
-      mappedData <- editableMappingService.mapData(unmappedDataTyped, relevantMapping, elementClassProto)
+      mappedData <- editableMappingService.mapData(unmappedData, relevantMapping, elementClassProto)
     } yield mappedData
   }
 }

webknossos-tracingstore/app/com/scalableminds/webknossos/tracingstore/tracings/editablemapping/EditableMappingService.scala

@@ -12,9 +12,15 @@ import com.scalableminds.webknossos.datastore.SegmentToAgglomerateProto.SegmentT
 import com.scalableminds.webknossos.datastore.SkeletonTracing.{Edge, Tree, TreeTypeProto}
 import com.scalableminds.webknossos.datastore.VolumeTracing.VolumeTracing
 import com.scalableminds.webknossos.datastore.VolumeTracing.VolumeTracing.ElementClassProto
-import com.scalableminds.webknossos.datastore.helpers.{NodeDefaults, ProtoGeometryImplicits, SkeletonTracingDefaults}
+import com.scalableminds.webknossos.datastore.helpers.{
+  NativeBucketScanner,
+  NodeDefaults,
+  ProtoGeometryImplicits,
+  SkeletonTracingDefaults
+}
 import com.scalableminds.webknossos.datastore.models.DataRequestCollection.DataRequestCollection
 import com.scalableminds.webknossos.datastore.models._
+import com.scalableminds.webknossos.datastore.models.datasource.ElementClass
 import com.scalableminds.webknossos.datastore.models.requests.DataServiceDataRequest
 import com.scalableminds.webknossos.datastore.services.{
   AdHocMeshRequest,
@@ -104,7 +110,7 @@ class EditableMappingService @Inject()(
 
   val defaultSegmentToAgglomerateChunkSize: Int = 64 * 1024 // max. 1 MiB chunks (two 8-byte numbers per element)
 
-  val binaryDataService = new BinaryDataService(Paths.get(""), None, None, None, None)
+  private val binaryDataService = new BinaryDataService(Paths.get(""), None, None, None, None)
 
   adHocMeshServiceHolder.tracingStoreAdHocMeshConfig = (binaryDataService, 30 seconds, 1)
   private val adHocMeshService: AdHocMeshService = adHocMeshServiceHolder.tracingStoreAdHocMeshService
@@ -115,6 +121,8 @@ class EditableMappingService @Inject()(
   private lazy val agglomerateToGraphCache: AlfuCache[(String, Long, Long), AgglomerateGraph] =
     AlfuCache(maxCapacity = 50)
 
+  private lazy val nativeBucketScanner: NativeBucketScanner = new NativeBucketScanner()
+
   def infoJson(tracingId: String, editableMappingInfo: EditableMappingInfo): JsObject =
     Json.obj(
       "tracingId" -> tracingId,
@@ -340,27 +348,31 @@ class EditableMappingService @Inject()(
     } yield segmentIdsOrdered.zip(agglomerateIdsOrdered).toMap
   }
 
-  def collectSegmentIds(data: Array[SegmentInteger]): Set[Long] =
-    data.toSet.map { u: SegmentInteger =>
-      u.toLong
-    }
+  def collectSegmentIds(bytes: Array[Byte], elementClass: ElementClass.Value): Box[Set[Long]] =
+    tryo(
+      nativeBucketScanner
+        .collectSegmentIds(bytes,
+                           ElementClass.bytesPerElement(elementClass),
+                           ElementClass.isSigned(elementClass),
+                           skipZeroes = false)
+        .toSet)
 
-  def mapData(unmappedData: Array[SegmentInteger],
+  def mapData(unmappedData: Array[Byte],
               relevantMapping: Map[Long, Long],
-              elementClass: ElementClassProto): Fox[Array[Byte]] = {
-    val mappedDataLongs = unmappedData.map(element => relevantMapping(element.toLong))
+              elementClass: ElementClassProto): Fox[Array[Byte]] =
     for {
+      unmappedDataTyped <- bytesToSegmentInt(unmappedData, elementClass)
+      mappedDataLongs = unmappedDataTyped.map(element => relevantMapping(element.toLong))
       bytes <- longsToBytes(mappedDataLongs, elementClass)
     } yield bytes
-  }
 
   private def bytesToLongs(bytes: Array[Byte], elementClass: ElementClassProto): Fox[Array[Long]] =
     for {
       _ <- bool2Fox(!elementClass.isuint64)
       segmentIntArray <- tryo(SegmentIntegerArray.fromByteArray(bytes, elementClass)).toFox
     } yield segmentIntArray.map(_.toLong)
 
-  def bytesToSegmentInt(bytes: Array[Byte], elementClass: ElementClassProto): Fox[Array[SegmentInteger]] =
+  private def bytesToSegmentInt(bytes: Array[Byte], elementClass: ElementClassProto): Fox[Array[SegmentInteger]] =
     for {
       _ <- bool2Fox(!elementClass.isuint64)
       segmentIntArray <- tryo(SegmentIntegerArray.fromByteArray(bytes, elementClass)).toFox

webknossos-tracingstore/app/com/scalableminds/webknossos/tracingstore/tracings/volume/VolumeSegmentIndexService.scala

@@ -132,7 +132,10 @@ class VolumeSegmentIndexService @Inject()(val tracingDataStore: TracingDataStore
   private def collectSegmentIds(bytes: Array[Byte], elementClass: ElementClass.Value): Box[Set[Long]] =
     tryo(
       nativeBucketScanner
-        .collectSegmentIds(bytes, ElementClass.bytesPerElement(elementClass), ElementClass.isSigned(elementClass))
+        .collectSegmentIds(bytes,
+                           ElementClass.bytesPerElement(elementClass),
+                           ElementClass.isSigned(elementClass),
+                           skipZeroes = true)
         .toSet)
 
   def getSegmentToBucketIndex(tracing: VolumeTracing,