From c181170caef298b46d306aaba1d572c1fae27fee Mon Sep 17 00:00:00 2001
From: Bill Katz <billkatz@gmail.com>
Date: Fri, 6 May 2022 23:23:41 -0400
Subject: [PATCH] add /supervoxel-sizes and improve /sizes; fixes #262

---
 datatype/labelmap/equiv.go       |  2 +
 datatype/labelmap/labelidx.go    | 62 ++++++++++++++++++++-----------
 datatype/labelmap/labelmap.go    | 64 ++++++++++++++++++++++++++++++++
 datatype/labelmap/mutate_test.go | 31 +++++++++++++++-
 4 files changed, 137 insertions(+), 22 deletions(-)

diff --git a/datatype/labelmap/equiv.go b/datatype/labelmap/equiv.go
index 07a49455..833a0d0d 100644
--- a/datatype/labelmap/equiv.go
+++ b/datatype/labelmap/equiv.go
@@ -474,6 +474,7 @@ func (d *Data) GetMapStats(ctx *datastore.VersionedCtx) (jsonBytes []byte, err e
 		if ds, err = datastore.GetDataByDataUUID(dataUUID); err != nil {
 			return
 		}
+		svm.RLock()
 		maxVersion := 0
 		for _, v := range svm.versions {
 			if int(v) > maxVersion {
@@ -486,6 +487,7 @@ func (d *Data) GetMapStats(ctx *datastore.VersionedCtx) (jsonBytes []byte, err e
 			NumVersions: len(svm.versions),
 			MaxVersion:  maxVersion,
 		}
+		svm.RUnlock()
 	}
 	return json.Marshal(stats)
 }
diff --git a/datatype/labelmap/labelidx.go b/datatype/labelmap/labelidx.go
index fc982670..a1ec74a2 100644
--- a/datatype/labelmap/labelidx.go
+++ b/datatype/labelmap/labelidx.go
@@ -437,28 +437,52 @@ func GetLabelSize(d dvid.Data, v dvid.VersionID, label uint64, isSupervoxel bool
 	return idx.NumVoxels(), nil
 }
 
-// GetLabelSizes returns the # of voxels in the given labels.  If isSupervoxel = true, the given
-// labels are interpreted as supervoxel ids and the sizes are of a supervoxel.  If a label doesn't
-// exist, a zero (not error) is returned.
-func GetLabelSizes(d dvid.Data, v dvid.VersionID, labels []uint64, isSupervoxel bool) ([]uint64, error) {
-	var supervoxels []uint64
-	if isSupervoxel {
-		svmap, err := getMapping(d, v)
-		if err != nil {
-			return nil, fmt.Errorf("couldn't get mapping for data %q, version %d: %v", d.DataName(), v, err)
-		}
-		supervoxels = make([]uint64, len(labels))
-		copy(supervoxels, labels)
+func getSupervoxelSizes(d dvid.Data, v dvid.VersionID, supervoxels []uint64) ([]uint64, error) {
+	svmap, err := getMapping(d, v)
+	if err != nil {
+		return nil, fmt.Errorf("couldn't get mapping for data %q, version %d: %v", d.DataName(), v, err)
+	}
+	labelsets := make(map[uint64][]uint64) // maps labels -> set of supervoxels in it.
+	labels, _, err := svmap.MappedLabels(v, supervoxels)
+	if err != nil {
+		return nil, err
+	}
+	for i, label := range labels {
+		labelsets[label] = append(labelsets[label], supervoxels[i])
+	}
 
-		labels, _, err = svmap.MappedLabels(v, labels)
+	sizemap := make(map[uint64]uint64, len(supervoxels))
+	for label, svlist := range labelsets {
+		idx, err := GetLabelIndex(d, v, label, false)
 		if err != nil {
 			return nil, err
 		}
+		if idx == nil {
+			for _, sv := range svlist {
+				sizemap[sv] = 0
+			}
+		} else {
+			svcounts := idx.GetSupervoxelCounts()
+			for _, sv := range svlist {
+				sizemap[sv] = svcounts[sv]
+			}
+		}
 	}
-	// TODO -- could optimize by doing unique set of labels if supervoxels, since multiple supervoxels
-	// may be in same label.  However, caching might simply remove this optimization issue since label
-	// index will already be cached.
-	sizes := make([]uint64, len(labels))
+	sizes := make([]uint64, len(supervoxels))
+	for i, sv := range supervoxels {
+		sizes[i] = sizemap[sv]
+	}
+	return sizes, nil
+}
+
+// GetLabelSizes returns the # of voxels in the given labels.  If isSupervoxel = true, the given
+// labels are interpreted as supervoxel ids and the sizes are of a supervoxel.  If a label doesn't
+// exist, a zero (not error) is returned.
+func GetLabelSizes(d dvid.Data, v dvid.VersionID, labels []uint64, isSupervoxel bool) (sizes []uint64, err error) {
+	if isSupervoxel {
+		return getSupervoxelSizes(d, v, labels)
+	}
+	sizes = make([]uint64, len(labels))
 	for i, label := range labels {
 		idx, err := GetLabelIndex(d, v, label, false)
 		if err != nil {
@@ -466,10 +490,6 @@ func GetLabelSizes(d dvid.Data, v dvid.VersionID, labels []uint64, isSupervoxel
 		}
 		if idx == nil {
 			sizes[i] = 0
-			continue
-		}
-		if isSupervoxel {
-			sizes[i] = idx.GetSupervoxelCount(supervoxels[i])
 		} else {
 			sizes[i] = idx.NumVoxels()
 		}
diff --git a/datatype/labelmap/labelmap.go b/datatype/labelmap/labelmap.go
index b177c647..c0c1d330 100644
--- a/datatype/labelmap/labelmap.go
+++ b/datatype/labelmap/labelmap.go
@@ -951,6 +951,23 @@ GET <api URL>/node/<UUID>/<data name>/sizes[?supervoxels=true]
 	supervoxels   If "true", interprets the given labels as a supervoxel ids.
     hash          MD5 hash of request body content in hexidecimal string format.
 
+GET <api URL>/node/<UUID>/<data name>/supervoxel-sizes/<label>
+
+	Returns the supervoxels and their sizes for the given label in JSON:
+
+	{
+		"supervoxels": [1,2,3,4,...],
+		"sizes": [100,200,300,400,...]
+	}
+
+	Returns a status code 404 (Not Found) if label does not exist.
+	
+    Arguments:
+    UUID          Hexadecimal string with enough characters to uniquely identify a version node.
+    data name     Name of labelmap instance.
+    label     	  A 64-bit integer label id
+
+
 GET  <api URL>/node/<UUID>/<data name>/sparsevol-size/<label>[?supervoxels=true]
 
 	Returns JSON giving the number of voxels, number of native blocks and the coarse bounding box in DVID
@@ -3810,6 +3827,9 @@ func (d *Data) ServeHTTP(uuid dvid.UUID, ctx *datastore.VersionedCtx, w http.Res
 	case "supervoxels":
 		d.handleSupervoxels(ctx, w, r, parts)
 
+	case "supervoxel-sizes":
+		d.handleSupervoxelSizes(ctx, w, r, parts)
+
 	case "size":
 		d.handleSize(ctx, w, r, parts)
 
@@ -4771,6 +4791,50 @@ func (d *Data) handleSupervoxels(ctx *datastore.VersionedCtx, w http.ResponseWri
 	timedLog.Infof("HTTP GET supervoxels for label %d (%s)", label, r.URL)
 }
 
+func (d *Data) handleSupervoxelSizes(ctx *datastore.VersionedCtx, w http.ResponseWriter, r *http.Request, parts []string) {
+	// GET <api URL>/node/<UUID>/<data name>/supervoxel-sizes/<label>
+	if len(parts) < 5 {
+		server.BadRequest(w, r, "DVID requires label to follow 'supervoxels' command")
+		return
+	}
+	timedLog := dvid.NewTimeLog()
+
+	label, err := strconv.ParseUint(parts[4], 10, 64)
+	if err != nil {
+		server.BadRequest(w, r, err)
+		return
+	}
+	if label == 0 {
+		server.BadRequest(w, r, "Label 0 is protected background value and cannot be queried as body.\n")
+		return
+	}
+
+	idx, err := GetLabelIndex(d, ctx.VersionID(), label, false)
+	if err != nil {
+		server.BadRequest(w, r, err)
+		return
+	}
+	if idx == nil || len(idx.Blocks) == 0 {
+		w.WriteHeader(http.StatusNotFound)
+		return
+	}
+
+	counts := idx.GetSupervoxelCounts()
+	var supervoxels_json, sizes_json string
+	for sv, count := range counts {
+		supervoxels_json += strconv.FormatUint(sv, 10) + ","
+		sizes_json += strconv.FormatUint(count, 10) + ","
+	}
+	w.Header().Set("Content-type", "application/json")
+	fmt.Fprintf(w, "{\n  ")
+	fmt.Fprintf(w, `"supervoxels": [%s],`, supervoxels_json[:len(supervoxels_json)-1])
+	fmt.Fprintf(w, "\n  ")
+	fmt.Fprintf(w, `"sizes": [%s]`, sizes_json[:len(sizes_json)-1])
+	fmt.Fprintf(w, "\n}")
+
+	timedLog.Infof("HTTP GET supervoxel-sizes for label %d (%s)", label, r.URL)
+}
+
 func (d *Data) handleLabelmod(ctx *datastore.VersionedCtx, w http.ResponseWriter, r *http.Request, parts []string) {
 	// GET <api URL>/node/<UUID>/<data name>/lastmod/<label>
 	if len(parts) < 5 {
diff --git a/datatype/labelmap/mutate_test.go b/datatype/labelmap/mutate_test.go
index cb46770b..e4cc3063 100644
--- a/datatype/labelmap/mutate_test.go
+++ b/datatype/labelmap/mutate_test.go
@@ -749,10 +749,27 @@ func TestMergeLabels(t *testing.T) {
 	if err := json.Unmarshal(r, &supervoxels); err != nil {
 		t.Errorf("Unable to parse supervoxels from server.  Got: %v\n", supervoxels)
 	}
-	if len(supervoxels) != 1 && supervoxels[0] != 2 {
+	if len(supervoxels) != 1 || supervoxels[0] != 2 {
 		t.Errorf("expected [2] for supervoxels in body 2, got %v\n", supervoxels)
 	}
 
+	// Check /supervoxel-sizes
+	reqStr = fmt.Sprintf("%snode/%s/labels/supervoxel-sizes/2", server.WebAPIPath, uuid)
+	r = server.TestHTTP(t, "GET", reqStr, nil)
+	svsizes := struct {
+		Supervoxels []uint64
+		Sizes       []uint64
+	}{}
+	if err := json.Unmarshal(r, &svsizes); err != nil {
+		t.Errorf("Unable to parse supervoxel-sizes from server.  Got: %s\n", string(r))
+	}
+	if len(svsizes.Supervoxels) != 1 || svsizes.Supervoxels[0] != 2 {
+		t.Errorf("expected [2] for supervoxels in body 2, got %v\n", svsizes.Supervoxels)
+	}
+	if len(svsizes.Sizes) != 1 || svsizes.Sizes[0] != 50000 {
+		t.Errorf("expected [2] for supervoxel sizes in body 2, got %v\n", svsizes.Sizes)
+	}
+
 	// Make sure /label and /labels endpoints work.
 	apiStr := fmt.Sprintf("%snode/%s/%s/label/94_58_89", server.WebAPIPath, uuid, "labels")
 	jsonResp := server.TestHTTP(t, "GET", apiStr, nil)
@@ -813,6 +830,18 @@ func TestMergeLabels(t *testing.T) {
 		t.Errorf("expected supervoxel 3 within body 3 but didn't find it\n")
 	}
 
+	reqStr = fmt.Sprintf("%snode/%s/labels/supervoxel-sizes/2", server.WebAPIPath, uuid)
+	r = server.TestHTTP(t, "GET", reqStr, nil)
+	if err := json.Unmarshal(r, &svsizes); err != nil {
+		t.Errorf("Unable to parse supervoxel-sizes from server.  Got: %s\n", string(r))
+	}
+	if len(svsizes.Supervoxels) != 2 || svsizes.Supervoxels[0] != 2 || svsizes.Supervoxels[1] != 3 {
+		t.Errorf("expected [2,3] for supervoxels in body 2, got %v\n", svsizes.Supervoxels)
+	}
+	if len(svsizes.Sizes) != 2 || svsizes.Sizes[0] != 50000 || svsizes.Sizes[1] != 12000 {
+		t.Errorf("expected [50000, 12000] for supervoxel sizes in body 2, got %v\n", svsizes.Sizes)
+	}
+
 	// Make sure label changes are correct after completion
 	if err := datastore.BlockOnUpdating(uuid, "labels"); err != nil {
 		t.Fatalf("Error blocking on sync of labels: %v\n", err)