go-hep.org/x/hep@v0.38.1/groot/rtree/formula.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  	"go-hep.org/x/hep/groot/rtree/rfunc"
    11  )
    12  
    13  func newFormula(r *Reader, f rfunc.Formula) (rfunc.Formula, error) {
    14  	names := f.RVars()
    15  	rvs, missing := formulaAutoLoad(r, names)
    16  	if len(rvs) != len(names) {
    17  		return nil, fmt.Errorf("rtree: could not find all needed ReadVars (missing: %v)", missing)
    18  	}
    19  	args := make([]any, len(rvs))
    20  	for i, rv := range rvs {
    21  		args[i] = rv.Value
    22  	}
    23  	err := f.Bind(args)
    24  	if err != nil {
    25  		return nil, fmt.Errorf("rtree: could not bind formula to rvars: %w", err)
    26  	}
    27  
    28  	return f, nil
    29  }
    30  
    31  func formulaAutoLoad(r *Reader, idents []string) ([]*ReadVar, []string) {
    32  	var (
    33  		loaded  = make(map[string]*ReadVar, len(r.rvars))
    34  		needed  = make([]*ReadVar, 0, len(idents))
    35  		rvars   = NewReadVars(r.tree)
    36  		all     = make(map[string]*ReadVar, len(rvars))
    37  		missing []string
    38  	)
    39  
    40  	for i := range r.rvars {
    41  		rvar := &r.rvars[i]
    42  		loaded[rvar.Name] = rvar
    43  		all[rvar.Name] = rvar
    44  	}
    45  	for i := range rvars {
    46  		rvar := &rvars[i]
    47  		if _, ok := all[rvar.Name]; ok {
    48  			continue
    49  		}
    50  		all[rvar.Name] = rvar
    51  	}
    52  	for _, name := range idents {
    53  		rvar, ok := all[name]
    54  		if !ok {
    55  			missing = append(missing, name)
    56  			continue
    57  		}
    58  		if _, ok := loaded[name]; !ok {
    59  			r.rvars = append(r.rvars, *rvar)
    60  			rvar = &r.rvars[len(r.rvars)-1]
    61  			loaded[name] = rvar
    62  		}
    63  		needed = append(needed, rvar)
    64  	}
    65  
    66  	return needed, missing
    67  }