main.go

package main

import (
	"bytes"
	"errors"
	"flag"
	"fmt"
	"go/ast"
	"go/parser"
	"go/printer"
	"go/token"
	"log"
	"os"
	"strconv"
	"strings"

	"golang.org/x/tools/imports"
)

func main() {
	log.SetPrefix("unroller: ")
	log.SetFlags(0)
	if err := run(); err != nil {
		log.Fatal(err)
	}
}

func run() error {
	var inPath, outPath string
	flag.StringVar(&inPath, "f", getEnv("GOFILE", "-"), "input file path")
	flag.StringVar(&outPath, "o", "", "output file path")
	flag.Parse()

	if outPath == "" {
		if inPath == "-" {
			flag.Usage()
			return fmt.Errorf("-o is required")
		}
		outPath = strings.TrimSuffix(inPath, ".go") + "_unrolled.go"
	}

	fileSet := token.NewFileSet()
	astFile, err := parser.ParseFile(fileSet, inPath, nil, parser.ParseComments)
	if err != nil {
		return fmt.Errorf("parse go file: %w", err)
	}
	commentMap := ast.NewCommentMap(fileSet, astFile, astFile.Comments)

	var buf bytes.Buffer
	decls := astFile.Decls

	astFile.Decls = nil
	astFile.Comments = nil

	if err := printer.Fprint(&buf, token.NewFileSet(), astFile); err != nil {
		return fmt.Errorf("write generated file: %w", err)
	}

	for _, decl := range decls {
		funcDecl, ok := decl.(*ast.FuncDecl)
		if !ok {
			continue
		}
		funcDecl.Doc = nil
		sps, err := specializedParams(commentMap[funcDecl])
		if err != nil {
			return fmt.Errorf("parse specialized param: %w", err)
		}
		for _, sp := range sps {
			var unrolled ast.BlockStmt
			originalFunc := struct {
				Name string
				Body *ast.BlockStmt
				Type ast.FuncType
			}{
				Name: funcDecl.Name.Name,
				Body: funcDecl.Body,
				Type: *funcDecl.Type,
			}
			var rewrote bool
			for i, stmt := range funcDecl.Body.List {
				if len(unrolled.List) > 0 {
					unrolled.List = append(unrolled.List, stmt)
				}
				forStmt, ok := stmt.(*ast.ForStmt)
				if !ok {
					continue
				}
				commentGroups, ok := commentMap[forStmt]
				if !ok {
					continue
				}
				u, err := parseUnroll(commentGroups)
				if err != nil {
					return fmt.Errorf("parse unroll: %w", err)
				}
				if u == nil {
					continue
				}
				r, err := toRange(forStmt)
				if err != nil {
					return fmt.Errorf("to range: %w", err)
				}
				if len(unrolled.List) == 0 {
					unrolled.List = append(unrolled.List, funcDecl.Body.List[:i]...)
				} else {
					unrolled.List = unrolled.List[:len(unrolled.List)-1]
				}
				if r.IsDynamic() {
					if sp == 0 {
						var casesBlock ast.BlockStmt
						m := r.DynamicMax(u.max)
						for i := r.init; i < m; i++ {
							casesBlock.List = append(casesBlock.List, &ast.CaseClause{
								List: []ast.Expr{&ast.BasicLit{Kind: token.INT, Value: strconv.Itoa(m - i)}},
								Body: []ast.Stmt{
									forStmt.Body,
									forStmt.Post,
									&ast.BranchStmt{Tok: token.FALLTHROUGH},
								},
							})
						}
						casesBlock.List = append(casesBlock.List, &ast.CaseClause{
							List: []ast.Expr{&ast.BasicLit{Kind: token.INT, Value: "0"}},
						})
						unrolled.List = append(unrolled.List, &ast.SwitchStmt{
							Init: forStmt.Init,
							Tag:  forStmt.Cond.(*ast.BinaryExpr).Y,
							Body: &casesBlock,
						})
					} else {
						if !rewrote {
							funcDecl.Type.Params = rewriteConstantParam(r.DynamicVar(), funcDecl.Type.Params)
							unrolled.List = append(append([]ast.Stmt{}, defConst(r.DynamicVar(), sp)), unrolled.List...)
							rewrote = true
						}
						m := r.DynamicMax(sp)
						for x := r.init; x < m; x++ {
							unrolled.List = append(unrolled.List, &ast.BlockStmt{
								List: append([]ast.Stmt{defConst(r.name, x)}, forStmt.Body.List...),
							})
						}
					}
				} else {
					m := r.Max()
					for x := r.init; x < m; x++ {
						unrolled.List = append(unrolled.List, &ast.BlockStmt{
							List: append([]ast.Stmt{defConst(r.name, x)}, forStmt.Body.List...),
						})
					}
				}
			}
			if len(unrolled.List) > 0 {
				funcDecl.Name.Name += "Unrolled"
				if sp > 0 {
					funcDecl.Name.Name += strconv.Itoa(sp)
				}
				funcDecl.Body = &unrolled

				fmt.Fprintf(&buf, "// %s is the unrolled version of %s\n", funcDecl.Name.Name, originalFunc.Name)
				if err := printer.Fprint(&buf, token.NewFileSet(), funcDecl); err != nil {
					return fmt.Errorf("write generated file: %w", err)
				}
				buf.WriteString("\n\n")

				funcDecl.Name.Name = originalFunc.Name
				funcDecl.Body = originalFunc.Body
				funcDecl.Type.Params = originalFunc.Type.Params
			}
		}
	}

	b, err := imports.Process(outPath, buf.Bytes(), nil)
	if err != nil {
		return fmt.Errorf("process imports: %w", err)
	}
	f, err := os.Create(outPath)
	if err != nil {
		return err
	}
	defer f.Close()
	fmt.Fprintln(f, "// Code generated by github.com/orisano/unroller. DO NOT EDIT.")
	if _, err := f.Write(b); err != nil {
		return fmt.Errorf("write generated file: %w", err)
	}
	if err := f.Close(); err != nil {
		return fmt.Errorf("close generated file: %w", err)
	}
	return nil
}

func getEnv(key string, defaultValue string) string {
	v := os.Getenv(key)
	if v != "" {
		return v
	} else {
		return defaultValue
	}
}

func specializedParams(cgs []*ast.CommentGroup) ([]int, error) {
	for _, cg := range cgs {
		for _, c := range cg.List {
			if after, found := strings.CutPrefix(c.Text, "// SPECIALIZED:"); found {
				return parseNumbers(after)
			}
		}
	}
	return []int{0}, nil
}

type Unroll struct {
	bunchSize int
	max       int
}

func parseUnroll(cgs []*ast.CommentGroup) (*Unroll, error) {
	for _, cg := range cgs {
		for _, c := range cg.List {
			if after, found := strings.CutPrefix(c.Text, "// UNROLL"); found {
				params, err := parseParams(after)
				if err != nil {
					return nil, fmt.Errorf("parse params: %w", err)
				}
				var u Unroll
				if s, ok := params["bunchSize"]; ok {
					x, err := strconv.Atoi(s)
					if err != nil {
						return nil, fmt.Errorf("parse bunchSize: %w", err)
					}
					u.bunchSize = x
				}
				if s, ok := params["max"]; ok {
					x, err := strconv.Atoi(s)
					if err != nil {
						return nil, fmt.Errorf("parse max: %w", err)
					}
					u.max = x
				}
				return &u, nil
			}
		}
	}
	return nil, nil
}

func isInt(e ast.Expr) bool {
	ident, ok := e.(*ast.Ident)
	if !ok {
		return false
	}
	switch ident.Name {
	case "byte", "rune", "int8", "int16", "int32", "int64", "int", "uint8", "uint16", "uint32", "uint64", "uint":
		return true
	default:
		return false
	}
}

func unwrapCast(e ast.Expr) ast.Expr {
	if ce, ok := e.(*ast.CallExpr); ok && len(ce.Args) == 1 && isInt(ce.Fun) {
		return ce
	}
	return e
}

func litInt(e ast.Expr) (int, bool) {
	lit, ok := unwrapCast(e).(*ast.BasicLit)
	if !ok {
		return 0, false
	}
	if lit.Kind != token.INT {
		return 0, false
	}
	x, err := strconv.ParseInt(lit.Value, 0, 64)
	if err != nil {
		return 0, false
	}
	return int(x), true
}

func parseParams(s string) (map[string]string, error) {
	s = strings.TrimSpace(s)
	if s == "" {
		return nil, nil
	}
	p := map[string]string{}
	for _, t := range strings.Split(s, ",") {
		t = strings.TrimSpace(t)
		sep := strings.IndexByte(t, '=')
		if sep < 0 {
			return nil, fmt.Errorf("parameter seperator not found")
		}
		p[t[:sep]] = t[sep+1:]
	}
	return p, nil
}

func parseNumbers(s string) ([]int, error) {
	var ns []int
	for _, t := range strings.Split(s, ",") {
		n, err := strconv.Atoi(strings.TrimSpace(t))
		if err != nil {
			return nil, err
		}
		ns = append(ns, n)
	}
	return ns, nil
}

func rewriteConstantParam(cParam string, params *ast.FieldList) *ast.FieldList {
	var fl ast.FieldList
	for _, p := range params.List {
		f := &ast.Field{
			Type: p.Type,
		}
		for _, name := range p.Names {
			if name.Name != cParam {
				f.Names = append(f.Names, name)
			}
		}
		if len(f.Names) > 0 {
			fl.List = append(fl.List, f)
		}
	}
	return &fl
}

func defConst(name string, x int) *ast.DeclStmt {
	return &ast.DeclStmt{
		Decl: &ast.GenDecl{Tok: token.CONST, Specs: []ast.Spec{
			&ast.ValueSpec{
				Names:  []*ast.Ident{ast.NewIdent(name)},
				Values: []ast.Expr{&ast.BasicLit{Kind: token.INT, Value: strconv.Itoa(x)}},
			},
		}},
	}
}

type Range struct {
	name string
	init int
	op   token.Token
	cast string
	max  ast.Expr
}

func (r *Range) IsDynamic() bool {
	_, ok := r.max.(*ast.Ident)
	return ok
}

func (r *Range) Max() int {
	x, _ := strconv.ParseInt(r.max.(*ast.BasicLit).Value, 0, 64)
	if r.op == token.LEQ {
		x++
	}
	return int(x)
}

func (r *Range) DynamicVar() string {
	return r.max.(*ast.Ident).Name
}

func (r *Range) DynamicMax(x int) int {
	if r.op == token.LEQ {
		x++
	}
	return x
}

var errUnsupportedStyleForStmt = errors.New("unsupported style for-stmt")

func toRange(s *ast.ForStmt) (*Range, error) {
	assignStmt, ok := s.Init.(*ast.AssignStmt)
	if !ok {
		return nil, errUnsupportedStyleForStmt
	}
	if len(assignStmt.Lhs) != 1 || len(assignStmt.Rhs) != 1 {
		return nil, errUnsupportedStyleForStmt
	}
	name := assignStmt.Lhs[0].(*ast.Ident).Name
	init, ok := litInt(assignStmt.Rhs[0])
	if !ok {
		return nil, errUnsupportedStyleForStmt
	}
	if stmt, ok := s.Post.(*ast.IncDecStmt); !ok || stmt.X.(*ast.Ident).Name != name || stmt.Tok != token.INC {
		return nil, errUnsupportedStyleForStmt
	}
	condExpr, ok := s.Cond.(*ast.BinaryExpr)
	if !ok {
		return nil, errUnsupportedStyleForStmt
	}
	if ident, ok := condExpr.X.(*ast.Ident); !ok || ident.Name != name {
		return nil, errUnsupportedStyleForStmt
	}
	switch condExpr.Op {
	case token.LSS, token.LEQ:
	default:
		return nil, errUnsupportedStyleForStmt
	}
	cast := ""
	y := condExpr.Y
	if ce, ok := y.(*ast.CallExpr); ok && len(ce.Args) == 1 && isInt(ce.Fun) {
		cast = ce.Fun.(*ast.Ident).Name
		y = ce.Args[0]
	}
	switch e := y.(type) {
	case *ast.BasicLit:
		if e.Kind != token.INT {
			return nil, errUnsupportedStyleForStmt
		}
	case *ast.Ident:
	default:
		return nil, errUnsupportedStyleForStmt
	}
	return &Range{
		name: name,
		init: init,
		op:   condExpr.Op,
		cast: cast,
		max:  y,
	}, nil
}