go-hep.org/x/hep@v0.38.1/hepmc/rootcnv/reader.go (about)

     1  // Copyright ©2022 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 rootcnv
     6  
     7  import (
     8  	"errors"
     9  	"fmt"
    10  	"io"
    11  
    12  	"go-hep.org/x/hep/groot/rtree"
    13  	"go-hep.org/x/hep/hepmc"
    14  )
    15  
    16  type rstream struct {
    17  	evt hepmc.Event
    18  	err error
    19  }
    20  
    21  type FlatTreeReader struct {
    22  	r *rtree.Reader
    23  
    24  	evt   event
    25  	rvars []rtree.ReadVar
    26  
    27  	evts chan rstream
    28  }
    29  
    30  func NewFlatTreeReader(t rtree.Tree, opts ...rtree.ReadOption) (*FlatTreeReader, error) {
    31  	r := FlatTreeReader{
    32  		evts: make(chan rstream),
    33  	}
    34  
    35  	r.rvars = rtree.ReadVarsFromStruct(&r.evt)
    36  
    37  	rr, err := rtree.NewReader(t, r.rvars, opts...)
    38  	if err != nil {
    39  		return nil, fmt.Errorf("hepmc: could not create ROOT Tree reader: %w", err)
    40  	}
    41  	r.r = rr
    42  
    43  	go func() {
    44  		defer close(r.evts)
    45  		err := rr.Read(func(ctx rtree.RCtx) error {
    46  			defer r.evt.reset()
    47  			var o rstream
    48  			err := r.evt.write(&o.evt)
    49  			if err != nil {
    50  				return err
    51  			}
    52  			r.evts <- o
    53  			return nil
    54  		})
    55  		if err != nil {
    56  			r.evts <- rstream{err: err}
    57  			return
    58  		}
    59  		r.evts <- rstream{err: io.EOF}
    60  	}()
    61  
    62  	return &r, nil
    63  }
    64  
    65  func (r *FlatTreeReader) Close() error {
    66  	return r.r.Close()
    67  }
    68  
    69  func (r *FlatTreeReader) Read(evt *hepmc.Event) error {
    70  	o := <-r.evts
    71  	if o.err != nil {
    72  		if errors.Is(o.err, io.EOF) {
    73  			return io.EOF
    74  		}
    75  		return fmt.Errorf("hepmc: could not read event from ROOT: %w", o.err)
    76  	}
    77  	*evt = o.evt
    78  	return nil
    79  }