Expose the id of aggregation types (#51)

policy/aggregation_id_compress.go

@@ -47,7 +47,7 @@ func NewAggregationIDCompressor() AggregationIDCompressor {
 	// NB(cw): If we start to support more than 64 types, the library will
 	// expand the underlying word list itself.
 	return &aggregationIDCompressor{
-		bs: bitset.New(totalAggregationTypes),
+		bs: bitset.New(MaxAggregationTypeID),
 	}
 }
 
@@ -57,7 +57,7 @@ func (c *aggregationIDCompressor) Compress(aggTypes AggregationTypes) (Aggregati
 		if !aggType.IsValid() {
 			return DefaultAggregationID, fmt.Errorf("could not compress invalid AggregationType %v", aggType)
 		}
-		c.bs.Set(uint(aggType))
+		c.bs.Set(uint(aggType.ID()))
 	}
 
 	codes := c.bs.Bytes()
@@ -83,7 +83,7 @@ func NewAggregationIDDecompressor() AggregationIDDecompressor {
 
 // NewPooledAggregationIDDecompressor returns a new pooled AggregationTypeDecompressor.
 func NewPooledAggregationIDDecompressor(pool AggregationTypesPool) AggregationIDDecompressor {
-	bs := bitset.New(totalAggregationTypes)
+	bs := bitset.New(MaxAggregationTypeID)
 	return &aggregationIDDecompressor{
 		bs:   bs,
 		buf:  bs.Bytes(),
@@ -103,7 +103,7 @@ func (c *aggregationIDDecompressor) Decompress(id AggregationID) (AggregationTyp
 
 	var res AggregationTypes
 	if c.pool == nil {
-		res = make(AggregationTypes, 0, totalAggregationTypes)
+		res = make(AggregationTypes, 0, MaxAggregationTypeID)
 	} else {
 		res = c.pool.Get()
 	}

policy/aggregation_id_compress_test.go

@@ -46,7 +46,7 @@ func TestAggregationIDCompressRoundTrip(t *testing.T) {
 
 	p := NewAggregationTypesPool(pool.NewObjectPoolOptions().SetSize(1))
 	p.Init(func() AggregationTypes {
-		return make(AggregationTypes, 0, totalAggregationTypes)
+		return make(AggregationTypes, 0, MaxAggregationTypeID)
 	})
 	compressor, decompressor := NewAggregationIDCompressor(), NewPooledAggregationIDDecompressor(p)
 	for _, test := range testcases {

policy/aggregation_type.go

@@ -54,13 +54,20 @@ const (
 	P999
 	P9999
 
-	totalAggregationTypes = iota
+	nextAggregationTypeID = iota
 )
 
 const (
+	// MaxAggregationTypeID is the largest id of all the valid aggregation types.
+	// NB(cw) MaxAggregationTypeID is guaranteed to be greater or equal
+	// to len(ValidAggregationTypes).
+	// Iff ids of all the valid aggregation types are consecutive,
+	// MaxAggregationTypeID == len(ValidAggregationTypes).
+	MaxAggregationTypeID = nextAggregationTypeID - 1
+
 	// AggregationIDLen is the length of the AggregationID.
-	// The AggregationIDLen will be 1 when totalAggregationTypes <= 64.
-	AggregationIDLen = (totalAggregationTypes-1)/64 + 1
+	// The AggregationIDLen will be 1 when MaxAggregationTypeID <= 63.
+	AggregationIDLen = (MaxAggregationTypeID)/64 + 1
 
 	aggregationTypesSeparator = ","
 )
@@ -104,7 +111,7 @@ var (
 )
 
 func init() {
-	aggregationTypeStringMap = make(map[string]AggregationType, totalAggregationTypes)
+	aggregationTypeStringMap = make(map[string]AggregationType, MaxAggregationTypeID)
 	for aggType := range ValidAggregationTypes {
 		aggregationTypeStringMap[aggType.String()] = aggType
 	}
@@ -122,6 +129,11 @@ func NewAggregationTypeFromSchema(input schema.AggregationType) (AggregationType
 	return aggType, nil
 }
 
+// ID returns the id of the AggregationType.
+func (a AggregationType) ID() int {
+	return int(a)
+}
+
 // IsValid checks if an AggregationType is valid.
 func (a AggregationType) IsValid() bool {
 	_, ok := ValidAggregationTypes[a]

policy/aggregation_type_pool_test.go

@@ -30,18 +30,18 @@ import (
 func TestAggregationTypesPool(t *testing.T) {
 	p := NewAggregationTypesPool(pool.NewObjectPoolOptions().SetSize(1))
 	p.Init(func() AggregationTypes {
-		return make(AggregationTypes, 0, totalAggregationTypes)
+		return make(AggregationTypes, 0, MaxAggregationTypeID)
 	})
 
 	aggTypes := p.Get()
-	require.Equal(t, totalAggregationTypes, cap(aggTypes))
+	require.Equal(t, MaxAggregationTypeID, cap(aggTypes))
 	require.Equal(t, 0, len(aggTypes))
 	aggTypes = append(aggTypes, P9999)
 
 	p.Put(aggTypes)
 
 	aggTypes2 := p.Get()
-	require.Equal(t, totalAggregationTypes, cap(aggTypes2))
+	require.Equal(t, MaxAggregationTypeID, cap(aggTypes2))
 	require.Equal(t, 0, len(aggTypes2))
 
 	aggTypes2 = append(aggTypes2, Last)

policy/aggregation_type_test.go

@@ -34,6 +34,12 @@ func TestAggregationTypeIsValid(t *testing.T) {
 	require.False(t, AggregationType(int(P9999)+1).IsValid())
 }
 
+func TestAggregationTypeMaxID(t *testing.T) {
+	require.Equal(t, MaxAggregationTypeID, P9999.ID())
+	require.Equal(t, P9999, AggregationType(MaxAggregationTypeID))
+	require.Equal(t, MaxAggregationTypeID, len(ValidAggregationTypes))
+}
+
 func TestAggregationTypesIsDefault(t *testing.T) {
 	require.True(t, DefaultAggregationTypes.IsDefault())