Merge pull request #8 from parquet-go/generics

use generics to enforce type safety

Author: achille-roussel

bitpack.go

@@ -2,6 +2,12 @@
 // integers of various bit widths.
 package bitpack
 
+// Int is a type constraint representing the integer types that this package
+// supports.
+type Int interface {
+	~int32 | ~uint32 | ~int64 | ~uint64 | ~int | ~uintptr
+}
+
 // ByteCount returns the number of bytes needed to hold the given bit count.
 func ByteCount(bitCount uint) int {
 	return int((bitCount + 7) / 8)

examples/example_test.go

@@ -0,0 +1,78 @@
+package examples
+
+import (
+	"fmt"
+	"github.com/parquet-go/bitpack"
+)
+
+func ExamplePack_int32() {
+	// Pack int32 values
+	values := []int32{1, 2, 3, 4, 5}
+	bitWidth := uint(3)
+
+	packedSize := bitpack.ByteCount(uint(len(values)) * bitWidth)
+	// Allocate buffer with extra space for pack algorithm
+	// (needs 3 bytes per value + padding for safe bit operations)
+	dst := make([]byte, packedSize+bitpack.PaddingInt32)
+
+	// Pack the values
+	bitpack.Pack(dst, values, bitWidth)
+
+	// Calculate actual packed size
+	fmt.Printf("Packed %d values into %d bytes\n", len(values), packedSize)
+	// Output: Packed 5 values into 2 bytes
+}
+
+func ExampleUnpack_int32() {
+	// First, pack some values
+	values := []int32{10, 20, 30, 40, 50}
+	bitWidth := uint(6)
+
+	packedSize := bitpack.ByteCount(uint(len(values)) * bitWidth)
+	// Allocate buffer for packing
+	packed := make([]byte, packedSize+bitpack.PaddingInt32)
+	bitpack.Pack(packed, values, bitWidth)
+
+	// Now unpack them
+	dst := make([]int32, len(values))
+	bitpack.Unpack(dst, packed, bitWidth)
+
+	fmt.Printf("Unpacked values: %v\n", dst)
+	// Output: Unpacked values: [10 20 30 40 50]
+}
+
+func ExamplePack_int64() {
+	// Pack int64 values
+	values := []int64{100, 200, 300, 400, 500}
+	bitWidth := uint(9)
+
+	packedSize := bitpack.ByteCount(uint(len(values)) * bitWidth)
+	// Allocate buffer with extra space for pack algorithm
+	// (needs packed size + padding for safe bit operations)
+	dst := make([]byte, packedSize+bitpack.PaddingInt64)
+
+	// Pack the values
+	bitpack.Pack(dst, values, bitWidth)
+
+	// Calculate actual packed size
+	fmt.Printf("Packed %d values into %d bytes\n", len(values), packedSize)
+	// Output: Packed 5 values into 6 bytes
+}
+
+func ExampleUnpack_int64() {
+	// First, pack some values
+	values := []int64{1000, 2000, 3000, 4000, 5000}
+	bitWidth := uint(13)
+
+	packedSize := bitpack.ByteCount(uint(len(values)) * bitWidth)
+	// Allocate buffer for packing
+	packed := make([]byte, packedSize+bitpack.PaddingInt64)
+	bitpack.Pack(packed, values, bitWidth)
+
+	// Now unpack them
+	dst := make([]int64, len(values))
+	bitpack.Unpack(dst, packed, bitWidth)
+
+	fmt.Printf("Unpacked values: %v\n", dst)
+	// Output: Unpacked values: [1000 2000 3000 4000 5000]
+}

examples/int32_test.go

@@ -1,8 +1,5 @@
 module github.com/parquet-go/bitpack
 
-go 1.25.1
+go 1.24.0
 
-require (
-	github.com/parquet-go/parquet-go v0.25.1
-	golang.org/x/sys v0.37.0
-)
+require golang.org/x/sys v0.37.0

examples/int64_test.go

@@ -1,4 +1,2 @@
-github.com/parquet-go/parquet-go v0.25.1 h1:l7jJwNM0xrk0cnIIptWMtnSnuxRkwq53S+Po3KG8Xgo=
-github.com/parquet-go/parquet-go v0.25.1/go.mod h1:AXBuotO1XiBtcqJb/FKFyjBG4aqa3aQAAWF3ZPzCanY=
 golang.org/x/sys v0.37.0 h1:fdNQudmxPjkdUTPnLn5mdQv7Zwvbvpaxqs831goi9kQ=
 golang.org/x/sys v0.37.0/go.mod h1:OgkHotnGiDImocRcuBABYBEXf8A9a87e/uXjp9XT3ks=

go.mod

@@ -1,23 +1,19 @@
 package bitpack
 
-// PackInt32 packs values from src to dst, each value is packed into the given
-// bit width regardless of how many bits are needed to represent it.
-//
-// The function panics if dst is too short to hold the bit packed values.
-func PackInt32(dst []byte, src []int32, bitWidth uint) {
-	assertPack(dst, len(src), bitWidth)
-	packInt32(dst, src, bitWidth)
-}
+import (
+	"unsafe"
 
-// PackInt64 packs values from src to dst, each value is packed into the given
-// bit width regardless of how many bits are needed to represent it.
-//
-// The function panics if dst is too short to hold the bit packed values.
-func PackInt64(dst []byte, src []int64, bitWidth uint) {
-	assertPack(dst, len(src), bitWidth)
-	packInt64(dst, src, bitWidth)
-}
+	"github.com/parquet-go/bitpack/unsafecast"
+)
 
-func assertPack(dst []byte, count int, bitWidth uint) {
-	_ = dst[:ByteCount(bitWidth*uint(count))]
+// Pack packs values from src to dst, each value is packed into the given
+// bit width regardless of how many bits are needed to represent it.
+func Pack[T Int](dst []byte, src []T, bitWidth uint) {
+	_ = dst[:ByteCount(bitWidth*uint(len(src)))]
+	switch unsafe.Sizeof(T(0)) {
+	case 4:
+		packInt32(dst, unsafecast.Slice[int32](src), bitWidth)
+	default:
+		packInt64(dst, unsafecast.Slice[int64](src), bitWidth)
+	}
 }

go.sum

@@ -25,11 +25,11 @@ func TestPackInt32(t *testing.T) {
 				// Pack the values
 				size := bitpack.ByteCount(uint(n) * bitWidth)
 				packed := make([]byte, size+bitpack.PaddingInt32)
-				bitpack.PackInt32(packed, block[:n], bitWidth)
+				bitpack.Pack(packed, block[:n], bitWidth)
 
 				// Unpack and verify
 				unpacked := make([]int32, n)
-				bitpack.UnpackInt32(unpacked, packed, bitWidth)
+				bitpack.Unpack(unpacked, packed, bitWidth)
 
 				if !slices.Equal(block[:n], unpacked) {
 					t.Fatalf("values mismatch for length=%d\nwant: %v\ngot:  %v", n, block[:n], unpacked)
@@ -55,11 +55,11 @@ func TestPackInt64(t *testing.T) {
 				// Pack the values
 				size := bitpack.ByteCount(uint(n) * bitWidth)
 				packed := make([]byte, size+bitpack.PaddingInt64)
-				bitpack.PackInt64(packed, block[:n], bitWidth)
+				bitpack.Pack(packed, block[:n], bitWidth)
 
 				// Unpack and verify
 				unpacked := make([]int64, n)
-				bitpack.UnpackInt64(unpacked, packed, bitWidth)
+				bitpack.Unpack(unpacked, packed, bitWidth)
 
 				if !slices.Equal(block[:n], unpacked) {
 					t.Fatalf("values mismatch for length=%d\nwant: %v\ngot:  %v", n, block[:n], unpacked)
@@ -86,7 +86,7 @@ func BenchmarkPackInt32(b *testing.B) {
 			src := block[:]
 
 			for i := 0; i < b.N; i++ {
-				bitpack.PackInt32(dst, src, bitWidth)
+				bitpack.Pack(dst, src, bitWidth)
 			}
 
 			b.SetBytes(4 * blockSize)
@@ -114,7 +114,7 @@ func BenchmarkPackInt64(b *testing.B) {
 			src := block[:]
 
 			for i := 0; i < b.N; i++ {
-				bitpack.PackInt64(dst, src, bitWidth)
+				bitpack.Pack(dst, src, bitWidth)
 			}
 
 			b.SetBytes(8 * blockSize)

pack.go

@@ -1,5 +1,11 @@
 package bitpack
 
+import (
+	"unsafe"
+
+	"github.com/parquet-go/bitpack/unsafecast"
+)
+
 // PaddingInt32 is the padding expected to exist after the end of input buffers
 // for the UnpackInt32 algorithm to avoid reading beyond the end of the input.
 const PaddingInt32 = 16
@@ -8,20 +14,16 @@ const PaddingInt32 = 16
 // for the UnpackInt32 algorithm to avoid reading beyond the end of the input.
 const PaddingInt64 = 32
 
-// UnpackInt32 unpacks 32 bit integers from src to dst.
-//
-// The function unpacked len(dst) integers, it panics if src is too short to
-// contain len(dst) values of the given bit width.
-func UnpackInt32(dst []int32, src []byte, bitWidth uint) {
-	_ = src[:ByteCount(bitWidth*uint(len(dst))+8*PaddingInt32)]
-	unpackInt32(dst, src, bitWidth)
-}
-
-// UnpackInt64 unpacks 64 bit integers from src to dst.
-//
-// The function unpacked len(dst) integers, it panics if src is too short to
-// contain len(dst) values of the given bit width.
-func UnpackInt64(dst []int64, src []byte, bitWidth uint) {
-	_ = src[:ByteCount(bitWidth*uint(len(dst))+8*PaddingInt64)]
-	unpackInt64(dst, src, bitWidth)
+// Unpack unpacks values from src to dst, each value is unpacked from the given
+// bit width regardless of how many bits are needed to represent it.
+func Unpack[T Int](dst []T, src []byte, bitWidth uint) {
+	sizeofT := uint(unsafe.Sizeof(T(0)))
+	padding := (8 * sizeofT) / 2 // 32 bits => 16, 64 bits => 32
+	_ = src[:ByteCount(bitWidth*uint(len(dst))+8*padding)]
+	switch sizeofT {
+	case 4:
+		unpackInt32(unsafecast.Slice[int32](dst), src, bitWidth)
+	default:
+		unpackInt64(unsafecast.Slice[int64](dst), src, bitWidth)
+	}
 }