Add size return to ReadMap and ReadSlice methods

Update ReadMap and ReadSlice to return collection size along with
iterators, enabling efficient pre-allocation of maps and slices.
Iterator remains the primary return value for natural usage patterns.

Author: oschwald

example_test.go

@@ -183,23 +183,35 @@ type CustomCity struct {
 // This provides custom decoding logic, similar to how json.Unmarshaler works
 // with encoding/json, allowing fine-grained control over data processing.
 func (c *CustomCity) UnmarshalMaxMindDB(d *mmdbdata.Decoder) error {
-	for key, err := range d.ReadMap() {
+	mapIter, _, err := d.ReadMap()
+	if err != nil {
+		return err
+	}
+	for key, err := range mapIter {
 		if err != nil {
 			return err
 		}
 
 		switch string(key) {
 		case "city":
 			// Decode nested city structure
-			for cityKey, cityErr := range d.ReadMap() {
+			cityMapIter, _, err := d.ReadMap()
+			if err != nil {
+				return err
+			}
+			for cityKey, cityErr := range cityMapIter {
 				if cityErr != nil {
 					return cityErr
 				}
 				switch string(cityKey) {
 				case "names":
 					// Decode nested map[string]string for localized names
 					names := make(map[string]string)
-					for nameKey, nameErr := range d.ReadMap() {
+					nameMapIter, _, err := d.ReadMap()
+					if err != nil {
+						return err
+					}
+					for nameKey, nameErr := range nameMapIter {
 						if nameErr != nil {
 							return nameErr
 						}

internal/decoder/decoder.go

@@ -214,20 +214,20 @@ func (d *Decoder) ReadUInt128() (hi, lo uint64, err error) {
 	return hi, lo, nil
 }
 
-// ReadMap returns an iterator to read the map. The first value from the
-// iterator is the key. Please note that this byte slice is only valid during
-// the iteration. This is done to avoid an unnecessary allocation. You must
-// make a copy of it if you are storing it for later use. The second value is
-// an error indicating that the database is malformed or that the pointed
-// value is not a map.
-func (d *Decoder) ReadMap() iter.Seq2[[]byte, error] {
-	return func(yield func([]byte, error) bool) {
-		size, offset, err := d.decodeCtrlDataAndFollow(KindMap)
-		if err != nil {
-			yield(nil, d.wrapError(err))
-			return
-		}
+// ReadMap returns an iterator to read the map along with the map size. The
+// size can be used to pre-allocate a map with the correct capacity for better
+// performance. The first value from the iterator is the key. Please note that
+// this byte slice is only valid during the iteration. This is done to avoid
+// an unnecessary allocation. You must make a copy of it if you are storing it
+// for later use. The second value is an error indicating that the database is
+// malformed or that the pointed value is not a map.
+func (d *Decoder) ReadMap() (iter.Seq2[[]byte, error], uint, error) {
+	size, offset, err := d.decodeCtrlDataAndFollow(KindMap)
+	if err != nil {
+		return nil, 0, d.wrapError(err)
+	}
 
+	iterator := func(yield func([]byte, error) bool) {
 		currentOffset := offset
 
 		for range size {
@@ -257,19 +257,21 @@ func (d *Decoder) ReadMap() iter.Seq2[[]byte, error] {
 		// Set the final offset after map iteration
 		d.reset(currentOffset)
 	}
+
+	return iterator, size, nil
 }
 
-// ReadSlice returns an iterator over the values of the slice. The iterator
-// returns an error if the database is malformed or if the pointed value is
-// not a slice.
-func (d *Decoder) ReadSlice() iter.Seq[error] {
-	return func(yield func(error) bool) {
-		size, offset, err := d.decodeCtrlDataAndFollow(KindSlice)
-		if err != nil {
-			yield(d.wrapError(err))
-			return
-		}
+// ReadSlice returns an iterator over the values of the slice along with the
+// slice size. The size can be used to pre-allocate a slice with the correct
+// capacity for better performance. The iterator returns an error if the
+// database is malformed or if the pointed value is not a slice.
+func (d *Decoder) ReadSlice() (iter.Seq[error], uint, error) {
+	size, offset, err := d.decodeCtrlDataAndFollow(KindSlice)
+	if err != nil {
+		return nil, 0, d.wrapError(err)
+	}
 
+	iterator := func(yield func(error) bool) {
 		currentOffset := offset
 
 		for i := range size {
@@ -301,6 +303,8 @@ func (d *Decoder) ReadSlice() iter.Seq[error] {
 		// Set final offset after slice iteration
 		d.reset(currentOffset)
 	}
+
+	return iterator, size, nil
 }
 
 // SkipValue skips over the current value without decoding it.

internal/decoder/decoder_test.go

@@ -146,8 +146,9 @@ func TestDecodeMap(t *testing.T) {
 	for hexStr, expected := range tests {
 		t.Run(hexStr, func(t *testing.T) {
 			decoder := newDecoderFromHex(t, hexStr)
-			resultMap := make(map[string]any)
-			mapIter := decoder.ReadMap() // [cite: 53]
+			mapIter, size, err := decoder.ReadMap() // [cite: 53]
+			require.NoError(t, err, "ReadMap failed")
+			resultMap := make(map[string]any, size) // Pre-allocate with correct capacity
 
 			// Iterate through the map [cite: 54]
 			for keyBytes, err := range mapIter { // [cite: 50]
@@ -178,8 +179,9 @@ func TestDecodeSlice(t *testing.T) {
 	for hexStr, expected := range tests {
 		t.Run(hexStr, func(t *testing.T) {
 			decoder := newDecoderFromHex(t, hexStr)
-			results := make([]any, 0)
-			sliceIter := decoder.ReadSlice() // [cite: 56]
+			sliceIter, size, err := decoder.ReadSlice() // [cite: 56]
+			require.NoError(t, err, "ReadSlice failed")
+			results := make([]any, 0, size) // Pre-allocate with correct capacity
 
 			// Iterate through the slice [cite: 57]
 			for err := range sliceIter {
@@ -359,7 +361,9 @@ func TestPointersInDecoder(t *testing.T) {
 			actualValue := make(map[string]string)
 
 			// Expecting a map at the target offset (may be behind a pointer)
-			mapIter := decoder.ReadMap()
+			mapIter, size, err := decoder.ReadMap()
+			require.NoError(t, err, "ReadMap failed")
+			_ = size // Use size if needed for pre-allocation
 			for keyBytes, errIter := range mapIter {
 				require.NoError(t, errIter)
 				key := string(keyBytes)

internal/decoder/error_context_test.go

@@ -171,7 +171,8 @@ func TestContextualErrorIntegration(t *testing.T) {
 		// Test new Decoder API - no automatic path tracking
 		dd := NewDataDecoder(buffer)
 		decoder := NewDecoder(dd, 0)
-		mapIter := decoder.ReadMap()
+		mapIter, _, err := decoder.ReadMap()
+		require.NoError(t, err, "ReadMap failed")
 
 		var mapErr error
 		for _, iterErr := range mapIter {
@@ -230,7 +231,8 @@ func TestContextualErrorIntegration(t *testing.T) {
 		decoder := NewDecoder(dd, 0)
 
 		// Navigate through the nested structure manually
-		mapIter := decoder.ReadMap()
+		mapIter, _, err := decoder.ReadMap()
+		require.NoError(t, err, "ReadMap failed")
 		var mapErr error
 
 		for key, iterErr := range mapIter {
@@ -241,7 +243,8 @@ func TestContextualErrorIntegration(t *testing.T) {
 			require.Equal(t, "list", string(key))
 
 			// Read the array
-			sliceIter := decoder.ReadSlice()
+			sliceIter, _, err := decoder.ReadSlice()
+			require.NoError(t, err, "ReadSlice failed")
 			sliceIndex := 0
 			for sliceIterErr := range sliceIter {
 				if sliceIterErr != nil {
@@ -251,7 +254,8 @@ func TestContextualErrorIntegration(t *testing.T) {
 				require.Equal(t, 0, sliceIndex) // Should be first element
 
 				// Read the nested map (array element)
-				innerMapIter := decoder.ReadMap()
+				innerMapIter, _, err := decoder.ReadMap()
+				require.NoError(t, err, "ReadMap failed")
 				for innerKey, innerIterErr := range innerMapIter {
 					if innerIterErr != nil {
 						mapErr = innerIterErr

reader_test.go

@@ -1119,16 +1119,24 @@ type TestCity struct {
 // UnmarshalMaxMindDB implements the Unmarshaler interface for TestCity.
 // This demonstrates custom decoding that avoids reflection for better performance.
 func (c *TestCity) UnmarshalMaxMindDB(d *mmdbdata.Decoder) error {
-	for key, err := range d.ReadMap() {
+	mapIter, _, err := d.ReadMap()
+	if err != nil {
+		return err
+	}
+	for key, err := range mapIter {
 		if err != nil {
 			return err
 		}
 
 		switch string(key) {
 		case "names":
 			// Decode nested map[string]string for localized names
-			names := make(map[string]string)
-			for nameKey, nameErr := range d.ReadMap() {
+			nameMapIter, size, err := d.ReadMap()
+			if err != nil {
+				return err
+			}
+			names := make(map[string]string, size) // Pre-allocate with correct capacity
+			for nameKey, nameErr := range nameMapIter {
 				if nameErr != nil {
 					return nameErr
 				}
@@ -1163,7 +1171,11 @@ type TestASN struct {
 
 // UnmarshalMaxMindDB implements the Unmarshaler interface for TestASN.
 func (a *TestASN) UnmarshalMaxMindDB(d *mmdbdata.Decoder) error {
-	for key, err := range d.ReadMap() {
+	mapIter, _, err := d.ReadMap()
+	if err != nil {
+		return err
+	}
+	for key, err := range mapIter {
 		if err != nil {
 			return err
 		}