go-hep.org/x/hep@v0.38.1/groot/rtree/rjoin.go (about)

     1  // Copyright ©2020 The go-hep Authors. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  package rtree
     6  
     7  import (
     8  	"fmt"
     9  )
    10  
    11  type rjoin struct {
    12  	j *join
    13  
    14  	rs []*rtree // FIXME(sbinet): handle join of chains?
    15  
    16  	rvs  []ReadVar
    17  	nrab int
    18  	beg  int64
    19  	end  int64
    20  }
    21  
    22  func newRJoin(t *join, rvars []ReadVar, n int, beg, end int64) *rjoin {
    23  	rvars = bindRVarsTo(t, rvars)
    24  	r := &rjoin{
    25  		j:    t,
    26  		rs:   make([]*rtree, len(t.trees)),
    27  		rvs:  rvars,
    28  		nrab: n,
    29  		beg:  beg,
    30  		end:  end,
    31  	}
    32  	rps := make([][]ReadVar, len(r.rs))
    33  	for i, t := range r.j.trees {
    34  		rps[i] = r.loadRVars(t.(*ttree), rvars)
    35  	}
    36  
    37  	r.rvs = r.rvs[:0]
    38  	for i, tree := range t.trees {
    39  		r.rs[i] = newRTree(tree.(*ttree), rps[i], r.nrab, beg, end)
    40  		r.rvs = append(r.rvs, r.rs[i].rvars()...)
    41  	}
    42  
    43  	return r
    44  }
    45  
    46  func (r *rjoin) loadRVars(t *ttree, rvars []ReadVar) []ReadVar {
    47  	rps := make([]ReadVar, 0, len(rvars))
    48  	for _, rv := range rvars {
    49  		br := asBranch(rv.leaf.Branch())
    50  		if br.tree != t {
    51  			continue
    52  		}
    53  		rps = append(rps, rv)
    54  	}
    55  	return rps
    56  }
    57  
    58  func (r *rjoin) Close() error {
    59  	var err error
    60  	for _, rr := range r.rs {
    61  		e := rr.Close()
    62  		if e != nil && err == nil {
    63  			err = e
    64  		}
    65  	}
    66  	return err
    67  }
    68  
    69  func (r *rjoin) rvars() []ReadVar { return r.rvs }
    70  
    71  func (r *rjoin) reset() {
    72  	for _, rr := range r.rs {
    73  		rr.reset()
    74  	}
    75  }
    76  
    77  func (r *rjoin) run(off, beg, end int64, f func(RCtx) error) error {
    78  	var (
    79  		err  error
    80  		rctx RCtx
    81  	)
    82  	defer r.Close()
    83  
    84  	err = r.start()
    85  	if err != nil {
    86  		return err
    87  	}
    88  	defer r.stop()
    89  
    90  	for i := beg; i < end; i++ {
    91  		err = r.read(i)
    92  		if err != nil {
    93  			return fmt.Errorf("rtree: could not read entry %d: %w", i, err)
    94  		}
    95  		rctx.Entry = i + off
    96  		err = f(rctx)
    97  		if err != nil {
    98  			return fmt.Errorf("rtree: could not process entry %d: %w", i, err)
    99  		}
   100  	}
   101  
   102  	return err
   103  }
   104  
   105  func (r *rjoin) read(ievt int64) error {
   106  	for _, rr := range r.rs {
   107  		err := rr.read(ievt)
   108  		if err != nil {
   109  			return err
   110  		}
   111  	}
   112  	return nil
   113  }
   114  
   115  func (r *rjoin) start() error {
   116  	for _, rr := range r.rs {
   117  		err := rr.start()
   118  		if err != nil {
   119  			return err
   120  		}
   121  	}
   122  	return nil
   123  }
   124  
   125  func (r *rjoin) stop() {
   126  	for _, rr := range r.rs {
   127  		rr.stop()
   128  	}
   129  }
   130  
   131  var (
   132  	_ reader = (*rjoin)(nil)
   133  )