groot/rtree: support reading multi-leaves branches

Fixes #156.

groot/rtree/branch.go

@@ -707,21 +707,62 @@ func (b *tbranch) setupBasket(bk *Basket, ib int, entry int64) error {
 }
 
 func (b *tbranch) scan(ptr interface{}) error {
-	for _, leaf := range b.leaves {
+	switch len(b.leaves) {
+	case 1:
+		leaf := b.leaves[0]
 		err := leaf.scan(b.basket.rbuf, ptr)
 		if err != nil {
 			return err
 		}
+	default:
+		rv := reflect.ValueOf(ptr).Elem()
+		rt := rv.Type()
+		for i := 0; i < rt.NumField(); i++ {
+			var (
+				leaf = b.leaves[i]
+				fv   = rv.Field(i)
+				ptr  = fv.Addr().Interface()
+			)
+			err := leaf.scan(b.basket.rbuf, ptr)
+			if err != nil {
+				return err
+			}
+		}
 	}
 	return b.basket.rbuf.Err()
 }
 
 func (b *tbranch) setAddress(ptr interface{}) error {
-	for i, leaf := range b.leaves {
-		// FIXME(sbinet): adjust ptr for e.g. a "f1/F;f2/I;f3/i" branch
-		err := leaf.setAddress(ptr)
+	switch len(b.leaves) {
+	case 0:
+		return xerrors.Errorf("rtree: can not set address for a leaf-less branch (name=%q)", b.Name())
+
+	case 1:
+		err := b.leaves[0].setAddress(ptr)
 		if err != nil {
-			return xerrors.Errorf("rtree: could not set address for leaf[%d][%s]: %w", i, leaf.Name(), err)
+			return xerrors.Errorf("rtree: could not set address for leaf[%d][%s]: %w", 0, b.leaves[0].Name(), err)
+		}
+
+	default:
+		rv := reflect.ValueOf(ptr).Elem()
+		rt := rv.Type()
+		switch kind := rv.Kind(); kind {
+		case reflect.Struct:
+			if len(b.leaves) != rt.NumField() {
+				// FIXME(sbinet): be more lenient and clever about this?
+				return xerrors.Errorf("rtree: fields/leaves number mismatch (name=%q, fields=%d, leaves=%d)", b.Name(), rt.NumField(), len(b.leaves))
+			}
+			for i, leaf := range b.leaves {
+				fv := rv.Field(i)
+				err := leaf.setAddress(fv.Addr().Interface())
+				if err != nil {
+					return xerrors.Errorf("rtree: could not set address for leaf[%d][%s]: %w", i, leaf.Name(), err)
+				}
+			}
+		default:
+			// TODO(sbinet): also support map[string]*T ?
+			// TODO(sbinet): also support []*T ?
+			return xerrors.Errorf("rtree: multi-leaf branches need a pointer-to-struct (got=%T)", ptr)
 		}
 	}
 	return nil

groot/rtree/scanner.go

@@ -278,10 +278,6 @@ func NewTreeScannerVars(t Tree, vars ...ScanVar) (*TreeScanner, error) {
 			return nil, xerrors.Errorf("rtree: Tree %q has no leaf named %q", t.Name(), sv.Leaf)
 		}
 		lset[leaf] = struct{}{}
-		err := br.setAddress(sv.Value)
-		if err != nil {
-			return nil, xerrors.Errorf("rtree: could not set branch address for %q: %w", br.Name(), err)
-		}
 		if lcnt := leaf.LeafCount(); lcnt != nil {
 			lbr := t.Leaf(lcnt.Name())
 			if lbr == nil {
@@ -301,9 +297,9 @@ func NewTreeScannerVars(t Tree, vars ...ScanVar) (*TreeScanner, error) {
 		if rv := reflect.ValueOf(arg); rv.Kind() != reflect.Ptr {
 			return nil, xerrors.Errorf("rtree: ScanVar %d (name=%v) has non pointer Value", i, sv.Name)
 		}
-		err = br.setAddress(arg)
+		err := br.setAddress(arg)
 		if err != nil {
-			panic(err)
+			return nil, xerrors.Errorf("rtree: could not set branch address for %q: %w", br.Name(), err)
 		}
 	}
 
@@ -525,7 +521,15 @@ func NewScannerVars(t Tree, vars ...ScanVar) (*Scanner, error) {
 		if rv := reflect.ValueOf(arg); rv.Kind() != reflect.Ptr {
 			return nil, xerrors.Errorf("rtree: ScanVar %d (name=%v) has non pointer Value", i, sv.Name)
 		}
-		err := br.setAddress(arg)
+		var err error
+		switch br := br.(type) {
+		case *tbranchElement:
+			err = br.setAddress(arg)
+		case *tbranch:
+			err = leaf.setAddress(arg)
+		default:
+			panic(xerrors.Errorf("rtree: unknown Branch type %T", br))
+		}
 		if err != nil {
 			panic(err)
 		}

groot/rtree/scanner_test.go

@@ -1423,3 +1423,136 @@ func TestG4LikeTree(t *testing.T) {
 		t.Fatal(err)
 	}
 }
+
+func TestMultiLeafBranchWithScanVars(t *testing.T) {
+	t.Parallel()
+
+	f, err := riofs.Open("../testdata/root_numpy_struct.root")
+	if err != nil {
+		t.Fatalf("%+v", err)
+	}
+	defer f.Close()
+
+	obj, err := f.Get("test")
+	if err != nil {
+		t.Fatalf("%+v", err)
+	}
+
+	tree := obj.(Tree)
+
+	type Data struct {
+		b1l1 int32
+		b1l2 float32
+		b2l1 int32
+		b2l2 float32
+	}
+
+	var (
+		data Data
+		want = []Data{
+			{10, 15.5, 20, 781.2},
+		}
+	)
+
+	rvars := []ScanVar{
+		{
+			Name:  "branch1",
+			Leaf:  "intleaf",
+			Value: &data.b1l1,
+		},
+		{
+			Name:  "branch1",
+			Leaf:  "floatleaf",
+			Value: &data.b1l2,
+		},
+		{
+			Name:  "branch2",
+			Leaf:  "intleaf",
+			Value: &data.b2l1,
+		},
+		{
+			Name:  "branch2",
+			Leaf:  "floatleaf",
+			Value: &data.b2l2,
+		},
+	}
+
+	sc, err := NewScannerVars(tree, rvars...)
+	if err != nil {
+		t.Fatalf("could not create scanner: %+v", err)
+	}
+	defer sc.Close()
+
+	for sc.Next() {
+		err = sc.Scan()
+		if err != nil {
+			t.Fatalf("could not scan entry %d: %+v", sc.Entry(), err)
+		}
+
+		if got, want := data, want[sc.Entry()]; !reflect.DeepEqual(got, want) {
+			t.Fatalf("invalid entry %d:\ngot= %#v\nwant=%#v", sc.Entry(), got, want)
+		}
+	}
+}
+
+func TestMultiLeafBranchWithTreeScanVars(t *testing.T) {
+	t.Parallel()
+
+	f, err := riofs.Open("../testdata/root_numpy_struct.root")
+	if err != nil {
+		t.Fatalf("%+v", err)
+	}
+	defer f.Close()
+
+	obj, err := f.Get("test")
+	if err != nil {
+		t.Fatalf("%+v", err)
+	}
+
+	tree := obj.(Tree)
+
+	type B struct {
+		L1 int32   `groot:"intleaf"`
+		L2 float32 `groot:"floatleaf"`
+	}
+
+	type Data struct {
+		B1 B `groot:"branch1"`
+		B2 B `groot:"branch2"`
+	}
+
+	var (
+		data Data
+		want = []Data{
+			{B{10, 15.5}, B{20, 781.2}},
+		}
+	)
+
+	rvars := []ScanVar{
+		{
+			Name:  "branch1",
+			Value: &data.B1,
+		},
+		{
+			Name:  "branch2",
+			Value: &data.B2,
+		},
+	}
+
+	sc, err := NewTreeScannerVars(tree, rvars...)
+	if err != nil {
+		t.Fatalf("could not create scanner: %+v", err)
+	}
+	defer sc.Close()
+
+	for sc.Next() {
+		err = sc.Scan(&data.B1, &data.B2)
+		if err != nil {
+			t.Fatalf("could not scan entry %d: %+v", sc.Entry(), err)
+		}
+
+		if got, want := data, want[sc.Entry()]; !reflect.DeepEqual(got, want) {
+			t.Fatalf("invalid entry %d:\ngot= %#v\nwant=%#v", sc.Entry(), got, want)
+		}
+	}
+}