go-hep.org/x/hep@v0.38.1/rio/scanner.go (about)

     1  // Copyright ©2015 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 rio
     6  
     7  import (
     8  	"fmt"
     9  	"io"
    10  )
    11  
    12  // Selector selects Records based on their name
    13  type Selector struct {
    14  	Name   string // Record name
    15  	Unpack bool   // Whether to unpack the Record
    16  }
    17  
    18  // Scanner provides a convenient interface for reading records of a rio-stream.
    19  type Scanner struct {
    20  	r   *Reader
    21  	err error   // first non-EOF error encountered while reading the rio-stream.
    22  	rec *Record // last record encountered while reading the rio-stream.
    23  
    24  	filter map[string]Selector // records to read. if nil, return everything.
    25  }
    26  
    27  // NewScanner returns a new Scanner to read from r.
    28  func NewScanner(r *Reader) *Scanner {
    29  	scan := &Scanner{
    30  		r:      r,
    31  		err:    nil,
    32  		rec:    newRecord("<N/A>", 0),
    33  		filter: make(map[string]Selector),
    34  	}
    35  	scan.rec.unpack = false
    36  	scan.rec.r = r
    37  	return scan
    38  }
    39  
    40  // Select sets the records selection function.
    41  func (s *Scanner) Select(selectors []Selector) {
    42  	s.filter = make(map[string]Selector, len(selectors))
    43  	for _, sel := range selectors {
    44  		s.filter[sel.Name] = sel
    45  	}
    46  }
    47  
    48  // Scan scans the next Record until io.EOF
    49  func (s *Scanner) Scan() bool {
    50  	if s.err != nil {
    51  		return false
    52  	}
    53  
    54  	for {
    55  		var hdr rioHeader
    56  		err := hdr.RioUnmarshal(s.r.r)
    57  		if err != nil {
    58  			s.err = err
    59  			return false
    60  		}
    61  
    62  		switch hdr.Frame {
    63  		case ftrFrame:
    64  			ftr := rioFooter{Header: hdr}
    65  			err = ftr.unmarshalData(s.r.r)
    66  			if err != nil {
    67  				s.err = err
    68  				return false
    69  			}
    70  			continue
    71  		case recFrame:
    72  			s.rec.raw.Header = hdr
    73  			err := s.rec.raw.unmarshalData(s.r.r)
    74  			if err != nil {
    75  				s.err = err
    76  				return false
    77  			}
    78  
    79  			clen := int64(rioAlignU32(s.rec.raw.CLen))
    80  
    81  			name := s.rec.Name()
    82  			if len(s.filter) > 0 {
    83  				_, ok := s.filter[name]
    84  				if !ok {
    85  					_, err = s.seek(clen, 0)
    86  					if err != nil {
    87  						s.err = err
    88  						return false
    89  					}
    90  					continue
    91  				}
    92  			}
    93  
    94  			s.rec.unpack = s.filter[name].Unpack
    95  
    96  			switch s.rec.unpack {
    97  
    98  			case true:
    99  				err = s.rec.readBlocks(s.r.r)
   100  				if err != nil {
   101  					s.err = err
   102  					return false
   103  				}
   104  				return true
   105  
   106  			case false:
   107  				_, err = s.seek(clen, 0)
   108  				if err != nil {
   109  					s.err = err
   110  					return false
   111  				}
   112  				return true
   113  			}
   114  
   115  		default:
   116  			panic(fmt.Errorf("unknown frame %v", hdr.Frame))
   117  		}
   118  	}
   119  }
   120  
   121  // seek sets the offset for the next Read or Write on file to offset,
   122  // interpreted according to whence: 0 means relative to the origin of the
   123  // file, 1 means relative to the current offset, and 2 means relative to
   124  // the end. It returns the new offset and an error, if any.
   125  func (s *Scanner) seek(offset int64, whence int) (ret int64, err error) {
   126  	switch r := s.r.r.(type) {
   127  	case io.Seeker:
   128  		return r.Seek(offset, whence)
   129  	default:
   130  		if whence != 0 {
   131  			panic("not implemented")
   132  		}
   133  		return io.CopyN(io.Discard, r, offset)
   134  	}
   135  }
   136  
   137  // Err returns the first non-EOF error encountered by the reader.
   138  func (s *Scanner) Err() error {
   139  	if s.err == io.EOF {
   140  		return nil
   141  	}
   142  	return s.err
   143  }
   144  
   145  // Record returns the last Record read by the Scanner.
   146  func (s *Scanner) Record() *Record {
   147  	return s.rec
   148  }