go-hep.org/x/hep@v0.38.1/groot/rtree/rchain.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 rchain struct {
    12  	ch *chain
    13  
    14  	rvs  []ReadVar
    15  	nrab int
    16  	beg  int64
    17  	end  int64
    18  
    19  	ibeg int // first tree to process
    20  	iend int // last-1 tree to process
    21  }
    22  
    23  var (
    24  	_ reader = (*rchain)(nil)
    25  )
    26  
    27  func newRChain(ch *chain, rvars []ReadVar, n int, beg, end int64) *rchain {
    28  	r := &rchain{
    29  		ch:   ch,
    30  		rvs:  rvars,
    31  		nrab: n,
    32  		beg:  beg,
    33  		end:  end,
    34  	}
    35  
    36  	tbeg, tend := r.findTrees(beg, end)
    37  	if tbeg < 0 || tend < 0 {
    38  		panic(fmt.Errorf(
    39  			"rtree: could not find matching trees in chain for [%d, %d) within [%d, %d)",
    40  			beg, end, 0, ch.Entries(),
    41  		))
    42  	}
    43  	r.ibeg = tbeg
    44  	r.iend = tend
    45  
    46  	r.loadRVars()
    47  
    48  	return r
    49  }
    50  
    51  func (r *rchain) Close() error {
    52  	return nil
    53  }
    54  
    55  func (r *rchain) rvars() []ReadVar { return r.rvs }
    56  
    57  func (r *rchain) loadRVars() {
    58  	if len(r.ch.trees) == 0 {
    59  		return
    60  	}
    61  
    62  	rr := newReader(r.ch.trees[0], r.rvs, r.nrab, 0, 1)
    63  	defer rr.Close()
    64  	r.rvs = rr.rvars()
    65  }
    66  
    67  func (r *rchain) run(off, beg, end int64, f func(RCtx) error) error {
    68  	defer r.Close()
    69  
    70  	trees := r.ch.trees[r.ibeg:r.iend]
    71  	if len(trees) == 0 {
    72  		return nil
    73  	}
    74  
    75  	err := r.start()
    76  	if err != nil {
    77  		return err
    78  	}
    79  	defer r.stop()
    80  
    81  	for i := r.ibeg; i < r.iend; i++ {
    82  		var (
    83  			eoff = r.ch.offs[i]
    84  			tots = r.ch.tots[i]
    85  			ibeg = maxI64(beg-eoff, 0)
    86  			iend = minI64(end-eoff, tots-eoff)
    87  			err  = r.runTree(i, eoff+off, ibeg, iend, f)
    88  		)
    89  		if err != nil {
    90  			return fmt.Errorf("rtree: could not process entry %d: %w", i, err)
    91  		}
    92  	}
    93  
    94  	return nil
    95  }
    96  
    97  func (r *rchain) findTrees(beg, end int64) (int, int) {
    98  	var (
    99  		eoff int64
   100  		ibeg = -1
   101  		iend = -1
   102  	)
   103  
   104  	for i, t := range r.ch.trees {
   105  		n := t.Entries()
   106  		if ibeg < 0 && eoff <= beg && beg < eoff+n {
   107  			ibeg = i
   108  		}
   109  		if iend < 0 && end <= eoff+n {
   110  			iend = i + 1
   111  		}
   112  		eoff += n
   113  	}
   114  	if iend < 0 {
   115  		iend = len(r.ch.trees)
   116  	}
   117  
   118  	return ibeg, iend
   119  }
   120  
   121  func (r *rchain) runTree(itree int, off, beg, end int64, f func(RCtx) error) error {
   122  	rr := newReader(r.ch.trees[itree], r.rvs, r.nrab, beg, end)
   123  	return rr.run(off, beg, end, f)
   124  }
   125  
   126  func (r *rchain) start() error { return nil }
   127  func (r *rchain) stop()        {}
   128  func (r *rchain) reset()       {}