groot/{riofs,rtree}: expose API to customize key-compression

Signed-off-by: Sebastien Binet <[email protected]>

groot/riofs/file.go

@@ -574,7 +574,7 @@ func (f *File) writeStreamerInfo() error {
 		return fmt.Errorf("riofs: could not write StreamerInfo list: %w", err)
 	}
 
-	key, err = newKeyFromBuf(&f.dir, "StreamerInfo", sinfos.Title(), sinfos.Class(), 1, buf.Bytes(), f)
+	key, err = newKeyFromBuf(&f.dir, "StreamerInfo", sinfos.Title(), sinfos.Class(), 1, buf.Bytes(), f, nil)
 	if err != nil {
 		return fmt.Errorf("riofs: could not create StreamerInfo key: %w", err)
 	}

groot/riofs/key.go

@@ -18,6 +18,27 @@ import (
 	"go-hep.org/x/hep/groot/rvers"
 )
 
+// KeyOption configures how a Key may be created.
+type KeyOption func(cfg *keyCfg) error
+
+// WithKeyCompression configures a Key to use the provided compression scheme.
+func WithKeyCompression(compression int32) KeyOption {
+	return func(cfg *keyCfg) error {
+		cfg.compr = rcompress.SettingsFrom(compression)
+		return nil
+	}
+}
+
+type keyCfg struct {
+	compr rcompress.Settings
+}
+
+func newKeyCfg() *keyCfg {
+	return &keyCfg{
+		compr: rcompress.Settings{Alg: rcompress.Inherit},
+	}
+}
+
 // noKeyError is the error returned when a riofs.Key could not be found.
 type noKeyError struct {
 	key string
@@ -122,12 +143,12 @@ func newKey(dir *tdirectoryFile, name, title, class string, objlen int32, f *Fil
 // NewKey creates a new key from the provided serialized object buffer.
 // NewKey puts the key and its payload at the end of the provided file f.
 // Depending on the file configuration, NewKey may compress the provided object buffer.
-func NewKey(dir Directory, name, title, class string, cycle int16, obj []byte, f *File) (Key, error) {
+func NewKey(dir Directory, name, title, class string, cycle int16, obj []byte, f *File, kopts ...KeyOption) (Key, error) {
 	var d *tdirectoryFile
 	if dir != nil {
 		d = dir.(*tdirectoryFile)
 	}
-	return newKeyFromBuf(d, name, title, class, cycle, obj, f)
+	return newKeyFromBuf(d, name, title, class, cycle, obj, f, kopts)
 }
 
 func newKeyFrom(dir *tdirectoryFile, name, title, class string, obj root.Object, f *File) (Key, error) {
@@ -185,12 +206,20 @@ func newKeyFrom(dir *tdirectoryFile, name, title, class string, obj root.Object,
 	return k, nil
 }
 
-func newKeyFromBuf(dir *tdirectoryFile, name, title, class string, cycle int16, buf []byte, f *File) (Key, error) {
+func newKeyFromBuf(dir *tdirectoryFile, name, title, class string, cycle int16, buf []byte, f *File, kopts []KeyOption) (Key, error) {
 	var err error
 	if dir == nil {
 		dir = &f.dir
 	}
 
+	kcfg := newKeyCfg()
+	for _, opt := range kopts {
+		err = opt(kcfg)
+		if err != nil {
+			return Key{}, fmt.Errorf("riofs: could not setup Key option: %w", err)
+		}
+	}
+
 	keylen := keylenFor(name, title, class, dir, f.end)
 	objlen := int32(len(buf))
 	k := Key{
@@ -212,7 +241,12 @@ func newKeyFromBuf(dir *tdirectoryFile, name, title, class string, cycle int16,
 		k.rvers += 1000
 	}
 
-	k.buf, err = rcompress.Compress(nil, buf, k.f.compression)
+	compress := k.f.compression
+	if kcfg.compr.Alg != rcompress.Inherit {
+		compress = kcfg.compr.Compression()
+	}
+
+	k.buf, err = rcompress.Compress(nil, buf, compress)
 	if err != nil {
 		return k, fmt.Errorf("riofs: could not compress object %s for key %q: %w", class, name, err)
 	}

groot/rtree/basket.go

@@ -306,7 +306,7 @@ func (b *Basket) grow(n int) {
 	b.offsets = append(b.offsets, make([]int32, delta)...)
 }
 
-func (b *Basket) writeFile(f *riofs.File) (totBytes int64, zipBytes int64, err error) {
+func (b *Basket) writeFile(f *riofs.File, compr int32) (totBytes int64, zipBytes int64, err error) {
 	header := b.header
 	b.header = true
 	defer func() {
@@ -332,7 +332,7 @@ func (b *Basket) writeFile(f *riofs.File) (totBytes int64, zipBytes int64, err e
 		b.wbuf.WriteArrayI32(b.offsets[:b.nevbuf])
 		b.wbuf.WriteI32(0)
 	}
-	b.key, err = riofs.NewKey(nil, b.key.Name(), b.key.Title(), b.Class(), int16(b.key.Cycle()), b.wbuf.Bytes(), f)
+	b.key, err = riofs.NewKey(nil, b.key.Name(), b.key.Title(), b.Class(), int16(b.key.Cycle()), b.wbuf.Bytes(), f, riofs.WithKeyCompression(compr))
 	if err != nil {
 		return 0, 0, fmt.Errorf("rtree: could not create basket-key: %w", err)
 	}

groot/rtree/branch.go

@@ -749,7 +749,7 @@ func (b *tbranch) flush() error {
 	}
 
 	f := b.tree.getFile()
-	totBytes, zipBytes, err := b.ctx.bk.writeFile(f)
+	totBytes, zipBytes, err := b.ctx.bk.writeFile(f, int32(b.compress))
 	if err != nil {
 		return fmt.Errorf("could not marshal basket[%d] (branch=%q): %w", b.writeBasket, b.Name(), err)
 	}