groot/rtree: handle parallel-merged TBaskets

This CL adds proper handling of Geant4 files produced in multithreaded
mode (and merged).

Fixes #989.

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

groot/rtree/bkreader.go

@@ -59,8 +59,26 @@ func newBkReader(b Branch, n int, beg, end int64) *bkreader {
 	}
 
 	if len(base.basketEntry) == len(base.basketSeek) {
+		numBaskets := 0
+		for i, v := range base.basketSeek {
+			if v == 0 || i == base.writeBasket {
+				break
+			}
+			numBaskets++
+		}
+		if numBaskets > 0 {
+			base.basketSeek = base.basketSeek[:numBaskets]
+			bkr.spans = bkr.spans[:len(base.basketSeek)]
+		}
+
 		// prepare for recover basket mode.
 		base.basketEntry = append(base.basketEntry, 0)
+		bkr.spans[0] = rspan{
+			pos: base.basketSeek[0],
+			sz:  base.basketBytes[0],
+			beg: base.basketEntry[0],
+			end: base.basketEntry[0+1],
+		}
 	} else {
 		for i, seek := range base.basketSeek {
 			bkr.spans[i] = rspan{

groot/rtree/tree_test.go

@@ -1610,3 +1610,50 @@ func TestUprootTrees(t *testing.T) {
 		})
 	}
 }
+
+func TestReadG4Merge(t *testing.T) {
+	const fname = "../testdata/g4-merge.root"
+
+	f, err := riofs.Open(fname)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer f.Close()
+
+	tree, err := riofs.Get[Tree](f, "Photons")
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	type Data struct {
+		X float64 `groot:"fX"`
+	}
+
+	var (
+		d    Data
+		want = []Data{
+			{X: 0.12123017145693009},
+			{X: 0.11744358487972842},
+			{X: 0.1142964503404715},
+			{X: 0.11702763375596258},
+			{X: 0.1160675474106288},
+		}
+	)
+
+	r, err := NewReader(tree, ReadVarsFromStruct(&d), WithRange(0, 5))
+	if err != nil {
+		t.Fatalf("could not create reader: %+v", err)
+	}
+	defer r.Close()
+
+	err = r.Read(func(ctx RCtx) error {
+		i := int(ctx.Entry)
+		if got, want := d, want[i]; got != want {
+			return fmt.Errorf("invalid entry[%d]: got=%v, want=%v", i, got, want)
+		}
+		return nil
+	})
+	if err != nil {
+		t.Fatalf("error: %+v", err)
+	}
+}