go-hep.org/x/hep@v0.38.1/groot/cmd/root-gen-datareader/main.go

// Copyright ©2017 The go-hep Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Command root-gen-datareader generates a Go struct to easily read the
// event data type stored inside a Tree.
//
// Example:
//
//	$> root-gen-datareader -t tree testdata/small-flat-tree.root
//	// automatically generated by root-gen-datareader.
//	// DO NOT EDIT.
//
//	package event
//
//	type Data struct {
//		ROOT_Int32        int32       `groot:"Int32"`
//		ROOT_Int64        int64       `groot:"Int64"`
//		ROOT_UInt32       int32       `groot:"UInt32"`
//		ROOT_UInt64       int64       `groot:"UInt64"`
//		ROOT_Float32      float32     `groot:"Float32"`
//		ROOT_Float64      float64     `groot:"Float64"`
//		ROOT_ArrayInt32   [10]int32   `groot:"ArrayInt32[10]"`
//		ROOT_ArrayInt64   [10]int64   `groot:"ArrayInt64[10]"`
//		ROOT_ArrayInt32   [10]int32   `groot:"ArrayInt32[10]"`
//		ROOT_ArrayInt64   [10]int64   `groot:"ArrayInt64[10]"`
//		ROOT_ArrayFloat32 [10]float32 `groot:"ArrayFloat32[10]"`
//		ROOT_ArrayFloat64 [10]float64 `groot:"ArrayFloat64[10]"`
//		ROOT_N            int32       `groot:"N"`
//		ROOT_SliceInt32   []int32     `groot:"SliceInt32[N]"`
//		ROOT_SliceInt64   []int64     `groot:"SliceInt64[N]"`
//		ROOT_SliceInt32   []int32     `groot:"SliceInt32[N]"`
//		ROOT_SliceInt64   []int64     `groot:"SliceInt64[N]"`
//		ROOT_SliceFloat32 []float32   `groot:"SliceFloat32[N]"`
//		ROOT_SliceFloat64 []float64   `groot:"SliceFloat64[N]"`
//	}
package main // import "go-hep.org/x/hep/groot/cmd/root-gen-datareader"

import (
	"bytes"
	"flag"
	"fmt"
	"go/format"
	"io"
	"log"
	"os"
	"reflect"
	"strings"
	"text/template"

	"go-hep.org/x/hep/groot"
	"go-hep.org/x/hep/groot/riofs"
	"go-hep.org/x/hep/groot/rtree"
)

func main() {

	log.SetPrefix("root-gen-datareader: ")
	log.SetFlags(0)

	var (
		treeName   = flag.String("t", "tree", "name of the tree to inspect")
		pkgName    = flag.String("p", "event", "name of the package where to generate the data model")
		outName    = flag.String("o", "", "name of the file where to store the generated data model (STDOUT)")
		dataReader = flag.Bool("reader", false, "generate data reader code")
		verbose    = flag.Bool("v", false, "enable verbose mode")
	)

	flag.Parse()

	if flag.NArg() <= 0 {
		log.Printf("missing input file name")
		flag.Usage()
		flag.PrintDefaults()
		os.Exit(1)
	}

	ctx := newContext(*pkgName, flag.Arg(0), *treeName, *dataReader, *verbose)

	var o io.Writer = os.Stdout
	if *outName != "" {
		f, err := os.Create(*outName)
		if err != nil {
			log.Fatalf("could not create output file: %+v", err)
		}
		defer f.Close()
		o = f
	}

	err := process(o, ctx)
	if err != nil {
		log.Fatalf("could not generate data-reader: %+v", err)
	}
}

func process(w io.Writer, ctx *Context) error {
	f, err := groot.Open(ctx.File)
	if err != nil {
		log.Fatal(err)
	}
	defer f.Close()

	obj, err := riofs.Dir(f).Get(ctx.Tree)
	if err != nil {
		return fmt.Errorf("could not retrieve tree %q: %w", ctx.Tree, err)
	}
	var (
		tree  = obj.(rtree.Tree)
		rvars = rtree.NewReadVars(tree)
	)
	ctx.printf("entries: %v\n", tree.Entries())

	for i, rvar := range rvars {
		ctx.printf("rvar[%03d]: %q (%v)", i, rvar.Name, reflect.Indirect(reflect.ValueOf(rvar.Value)).Kind())
		rv := reflect.Indirect(reflect.ValueOf(rvar.Value))
		switch rv.Kind() {
		case reflect.Struct:
			def := structDefFrom(ctx, rvar.Name, reflect.Indirect(reflect.ValueOf(rvar.Value)).Type())
			ctx.Defs[rvar.Name] = def
			ctx.DataReader.Fields = append(
				ctx.DataReader.Fields,
				FieldDef{
					Name:    goName(rvar.Name),
					Tag:     rvar.Name,
					VarName: rvar.Name,
					Type:    def.Name,
				},
			)
		case reflect.Array:
			ctx.DataReader.Fields = append(
				ctx.DataReader.Fields,
				FieldDef{
					Name:    goName(rvar.Name),
					Tag:     fmt.Sprintf("%s[%d]", rvar.Name, rv.Type().Len()),
					VarName: rvar.Name,
					Type:    fmt.Sprintf("%T", rv.Interface()),
				},
			)
			ctx.checkType(rv.Type().Elem())

		case reflect.Slice:
			ctx.DataReader.Fields = append(
				ctx.DataReader.Fields,
				FieldDef{
					Name:    goName(rvar.Name),
					Tag:     rvar.Name,
					VarName: rvar.Name,
					Type:    fmt.Sprintf("%T", rv.Interface()),
				},
			)
			ctx.checkType(rv.Type().Elem())

		default:
			ctx.DataReader.Fields = append(
				ctx.DataReader.Fields,
				FieldDef{
					Name:    goName(rvar.Name),
					Tag:     rvar.Name,
					VarName: rvar.Name,
					Type:    fmt.Sprintf("%T", rv.Interface()),
				},
			)
			ctx.checkType(rv.Type())
		}
	}
	delete(ctx.Defs, "DataReader")

	err = ctx.genCode(w)
	if err != nil {
		return fmt.Errorf("could not generate reader code: %w", err)
	}

	return nil
}

func goName(s string) string {
	return "ROOT_" + s
}

// StructDef models a TTree content as a Go struct, where each TBranch
// of the tree is translated as a field of the struct.
type StructDef struct {
	Name   string
	Fields []FieldDef
}

// FieldDef describes a Go struct field, corresponding to a TTree's branch.
type FieldDef struct {
	Name    string
	Type    string
	VarName string
	Tag     string
}

// Context holds together various informations about the TTree being processed.
type Context struct {
	Package       string
	Imports       map[string]int
	DataReader    *StructDef
	Defs          map[string]*StructDef
	GenDataReader bool
	File          string
	Tree          string
	Verbose       bool
}

func newContext(pkg, file, tree string, dataReader, verbose bool) *Context {
	ctx := &Context{
		Package: pkg,
		Imports: make(map[string]int),
		Defs: map[string]*StructDef{
			"DataReader": {
				Name:   "DataReader",
				Fields: nil,
			},
		},
		GenDataReader: dataReader,
		File:          file,
		Tree:          tree,
		Verbose:       verbose,
	}
	ctx.DataReader = ctx.Defs["DataReader"]
	if dataReader {
		ctx.Imports["go-hep.org/x/hep/groot/rtree"]++
	}
	return ctx
}

func (ctx *Context) printf(format string, args ...any) {
	if ctx.Verbose {
		log.Printf(format, args...)
	}
}

func (ctx *Context) genCode(w io.Writer) error {
	t := template.New("top")
	template.Must(t.Parse(codeTmpl))
	buf := new(bytes.Buffer)
	err := t.Execute(buf, ctx)
	if err != nil {
		return err
	}
	src, err := format.Source(buf.Bytes())
	if err != nil {
		log.Printf("source:\n%s\n", buf.Bytes())
		log.Printf("error: %+v", err)
		return err
	}
	_, err = w.Write(src)
	return err
}

const codeTmpl = `// automatically generated by root-gen-datareader.
// DO NOT EDIT.

package {{.Package}}

{{$length := len .Imports}}{{if gt $length 0}}
import (
{{range $key, $value := .Imports}}
"{{$key}}"
{{- end}}
)
{{- end}}

{{range .Defs}}
type {{.Name}} struct {
{{range .Fields}}	{{.Name}} {{.Type}} ` + "`groot:\"{{.Tag}}\"`" + `
{{end}}}
{{end}}


{{with .DataReader}}
// Data is the data contained in a rtree.Tree.
type Data struct {
{{ range .Fields}}	{{.Name}} {{.Type}} ` + "`groot:\"{{.Tag}}\"`" + `
{{ end}}}
{{end}}

{{if .GenDataReader}}
{{with .DataReader}}
type DataReader struct {
	Data   Data
	Tree   rtree.Tree
	Reader *rtree.Reader
}
{{end}}
{{end}}
`

func structDefFrom(ctx *Context, name string, rt reflect.Type) *StructDef {
	def := StructDef{
		Name:   name,
		Fields: make([]FieldDef, rt.NumField()),
	}
	for i := range def.Fields {
		ft := rt.Field(i)
		def.Fields[i] = fieldDefFrom(ctx, ft)
	}

	return &def
}

func fieldDefFrom(ctx *Context, typ reflect.StructField) FieldDef {
	tag := typ.Tag.Get("groot")
	ctx.checkType(typ.Type)

	switch typ.Type.Kind() {
	case reflect.Struct:
		branch := tag
		if i := strings.Index(branch, "["); i > 0 {
			branch = branch[:i]
		}
		ctx.Defs[branch] = structDefFrom(ctx, branch, typ.Type)
		return FieldDef{
			Name:    typ.Name,
			Type:    branch,
			VarName: tag,
			Tag:     tag,
		}
	default:
		return FieldDef{
			Name:    typ.Name,
			Type:    typ.Type.String(),
			VarName: tag,
			Tag:     tag,
		}
	}
}

func (ctx *Context) checkType(typ reflect.Type) {
	pkg := typ.PkgPath()
	if pkg != "" {
		ctx.Imports[pkg]++
	}
}