aggregators/internal/protohash/generate/main.go

// Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
// or more contributor license agreements. Licensed under the Elastic License 2.0;
// you may not use this file except in compliance with the Elastic License 2.0.

package main

import (
	"fmt"
	"go/types"
	"log"
	"os"
	"strings"

	"golang.org/x/tools/go/packages"
)

func main() {
	const pkgpath = "github.com/elastic/apm-aggregation/aggregationpb"
	cfg := &packages.Config{Mode: packages.NeedTypes | packages.NeedTypesInfo}
	pkgs, err := packages.Load(cfg, pkgpath)
	if err != nil {
		log.Fatal(err)
	}

	f, err := os.Create("generated.go")
	if err != nil {
		log.Fatal(err)
	}
	defer f.Close()

	fmt.Fprintln(f, `
// Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
// or more contributor license agreements. Licensed under the Elastic License 2.0;
// you may not use this file except in compliance with the Elastic License 2.0.

// Code generated by protohash/generate. DO NOT EDIT.

package protohash

import (
	"encoding/binary"

	"github.com/cespare/xxhash/v2"

	"github.com/elastic/apm-aggregation/aggregationpb"
)

func writeUint32(h *xxhash.Digest, v uint32) {
	var buf [4]byte
	binary.LittleEndian.PutUint32(buf[:], v)
	h.Write(buf[:])
}

func writeUint64(h *xxhash.Digest, v uint64) {
	var buf [8]byte
	binary.LittleEndian.PutUint64(buf[:], v)
	h.Write(buf[:])
}
`[1:])

	pkg := pkgs[0]
	pkgscope := pkg.Types.Scope()
	for _, name := range pkgscope.Names() {
		if !strings.HasSuffix(name, "Key") {
			continue
		}
		typeName, ok := pkgscope.Lookup(name).(*types.TypeName)
		if !ok || !typeName.Exported() {
			continue
		}
		named := typeName.Type().(*types.Named)
		structType, ok := named.Underlying().(*types.Struct)
		if !ok {
			continue
		}

		fmt.Fprintf(f, "func Hash%[1]s(h xxhash.Digest, k *aggregationpb.%[1]s) xxhash.Digest {\n", name)
		for i := 0; i < structType.NumFields(); i++ {
			field := structType.Field(i)
			if !field.Exported() {
				continue
			}
			var unhandled bool
			switch fieldType := field.Type().(type) {
			case *types.Basic:
				switch kind := fieldType.Kind(); kind {
				case types.Bool:
					fmt.Fprintf(f, "	if k.%s {\n		h.WriteString(\"1\")\n	}\n", field.Name())
				case types.String:
					fmt.Fprintf(f, "	h.WriteString(k.%s)\n", field.Name())
				case types.Uint32:
					fmt.Fprintf(f, "	writeUint32(&h, k.%s)\n", field.Name())
				case types.Uint64:
					fmt.Fprintf(f, "	writeUint64(&h, k.%s)\n", field.Name())
				default:
					unhandled = true
				}
			case *types.Slice:
				switch elemType := fieldType.Elem().(type) {
				case *types.Basic:
					if elemType.Kind() != types.Byte {
						unhandled = true
						break
					}
					fmt.Fprintf(f, "	h.Write(k.%s)\n", field.Name())
				default:
					unhandled = true
				}
			default:
				unhandled = true
			}
			if unhandled {
				panic(fmt.Errorf("unhandled field %s.%s (%v)", name, field.Name(), field.Type()))
			}
		}
		fmt.Fprintln(f, "	return h\n}")
		fmt.Fprintln(f)
	}
}