From b21241ec1ceca4221d22593cab59ecc535540fa9 Mon Sep 17 00:00:00 2001
From: Michael Eischer <michael.eischer@fau.de>
Date: Sun, 28 Aug 2022 12:09:28 +0200
Subject: [PATCH] restic: Add CountedBlobSet type

This allows maintaining a usage counter for each blob.
---
 internal/restic/counted_blob_set.go      | 59 ++++++++++++++++++++++++
 internal/restic/counted_blob_set_test.go | 37 +++++++++++++++
 2 files changed, 96 insertions(+)
 create mode 100644 internal/restic/counted_blob_set.go
 create mode 100644 internal/restic/counted_blob_set_test.go

diff --git a/internal/restic/counted_blob_set.go b/internal/restic/counted_blob_set.go
new file mode 100644
index 000000000..29644215b
--- /dev/null
+++ b/internal/restic/counted_blob_set.go
@@ -0,0 +1,59 @@
+package restic
+
+import "sort"
+
+// CountedBlobSet is a set of blobs. For each blob it also stores a uint8 value
+// which can be used to track some information. The CountedBlobSet does not use
+// that value in any way. New entries are created with value 0.
+type CountedBlobSet map[BlobHandle]uint8
+
+// NewCountedBlobSet returns a new CountedBlobSet, populated with ids.
+func NewCountedBlobSet(handles ...BlobHandle) CountedBlobSet {
+	m := make(CountedBlobSet)
+	for _, h := range handles {
+		m[h] = 0
+	}
+
+	return m
+}
+
+// Has returns true iff id is contained in the set.
+func (s CountedBlobSet) Has(h BlobHandle) bool {
+	_, ok := s[h]
+	return ok
+}
+
+// Insert adds id to the set.
+func (s CountedBlobSet) Insert(h BlobHandle) {
+	s[h] = 0
+}
+
+// Delete removes id from the set.
+func (s CountedBlobSet) Delete(h BlobHandle) {
+	delete(s, h)
+}
+
+func (s CountedBlobSet) Len() int {
+	return len(s)
+}
+
+// List returns a sorted slice of all BlobHandle in the set.
+func (s CountedBlobSet) List() BlobHandles {
+	list := make(BlobHandles, 0, len(s))
+	for h := range s {
+		list = append(list, h)
+	}
+
+	sort.Sort(list)
+
+	return list
+}
+
+func (s CountedBlobSet) String() string {
+	str := s.List().String()
+	if len(str) < 2 {
+		return "{}"
+	}
+
+	return "{" + str[1:len(str)-1] + "}"
+}
diff --git a/internal/restic/counted_blob_set_test.go b/internal/restic/counted_blob_set_test.go
new file mode 100644
index 000000000..c5c321503
--- /dev/null
+++ b/internal/restic/counted_blob_set_test.go
@@ -0,0 +1,37 @@
+package restic_test
+
+import (
+	"testing"
+
+	"github.com/restic/restic/internal/restic"
+	"github.com/restic/restic/internal/test"
+)
+
+func TestCountedBlobSet(t *testing.T) {
+	bs := restic.NewCountedBlobSet()
+	test.Equals(t, bs.Len(), 0)
+	test.Equals(t, bs.List(), restic.BlobHandles{})
+
+	bh := restic.NewRandomBlobHandle()
+	// check non existant
+	test.Equals(t, bs.Has(bh), false)
+
+	// test insert
+	bs.Insert(bh)
+	test.Equals(t, bs.Has(bh), true)
+	test.Equals(t, bs.Len(), 1)
+	test.Equals(t, bs.List(), restic.BlobHandles{bh})
+
+	// test remove
+	bs.Delete(bh)
+	test.Equals(t, bs.Len(), 0)
+	test.Equals(t, bs.Has(bh), false)
+	test.Equals(t, bs.List(), restic.BlobHandles{})
+
+	bs = restic.NewCountedBlobSet(bh)
+	test.Equals(t, bs.Len(), 1)
+	test.Equals(t, bs.List(), restic.BlobHandles{bh})
+
+	s := bs.String()
+	test.Assert(t, len(s) > 10, "invalid string: %v", s)
+}