humungus - fungo

overview - files - changes

viewing file: fungo/ bango/ main.go

view - history - annotated - download

package main

import (
	"flag"
	"fmt"
	"go/ast"
	"go/format"
	"go/parser"
	"go/token"
	"os"
	"slices"

	"humungus.tedunangst.com/r/fungo/match"
)

func errx(msg string, args ...interface{}) {
	fmt.Fprintf(os.Stderr, msg+"\n", args...)
	os.Exit(1)
}

var wrap bool

func main() {
	var undo bool
	flag.BoolVar(&wrap, "w", false, "wrap")
	flag.BoolVar(&undo, "u", false, "undo")
	flag.Parse()
	args := flag.Args()
	var fset token.FileSet
	pkgs, err := parser.ParseDir(&fset, args[0], nil, parser.ParseComments)
	if err != nil {
		errx("error parsing packages: %s", err)
	}
	fn := rewrite
	if undo {
		fn = unwrite
	}
	for _, pkg := range pkgs {
		process(&fset, pkg, fn)
	}
}

func newcheck(decl *ast.FuncDecl, fail *ast.CallExpr) ast.Stmt {
	stmt := new(ast.IfStmt)
	cond := new(ast.BinaryExpr)
	x := new(ast.Ident)
	x.Name = "err"
	cond.X = x
	cond.Op = token.NEQ
	y := new(ast.Ident)
	y.Name = "nil"
	cond.Y = y
	stmt.Cond = cond
	body := new(ast.BlockStmt)
	ret := new(ast.ReturnStmt)
	ret.Results = append(ret.Results, &ast.Ident{Name: "nil"})
	if wrap {
		call := new(ast.CallExpr)
		call.Fun = &ast.Ident{Name: "fmt.Errorf"}
		var failname string
		if ident, ok := fail.Fun.(*ast.Ident); ok {
			failname = ident.Name
		}
		if sel, ok := fail.Fun.(*ast.SelectorExpr); ok {
			if ident, ok := sel.X.(*ast.Ident); ok {
				failname += ident.Name
			}
			if sel.Sel != nil {
				failname += "." + sel.Sel.Name
			}
		}
		str := &ast.BasicLit{
			Kind:  token.STRING,
			Value: fmt.Sprintf("\"%s %s error: %%w\"", decl.Name.Name, failname),
		}
		call.Args = []ast.Expr{str, &ast.Ident{Name: "err"}}
		ret.Results = append(ret.Results, call)
	} else {
		ret.Results = append(ret.Results, &ast.Ident{Name: "err"})
	}
	body.List = append(body.List, ret)
	stmt.Body = body
	return stmt
}

func rewrite(fset *token.FileSet, decl *ast.FuncDecl) {
	call := match.Call()
	matcher := match.Assign(match.Multi(match.Ident("")), match.Multi(match.Unary(token.XOR, call)))
	walker := func(node ast.Node) bool {
		block, ok := node.(*ast.BlockStmt)
		if !ok {
			return true
		}
		for i := 0; i < len(block.List); i++ {
			if !matcher.Match(block.List[i]) {
				continue
			}
			asgn := matcher.Assign
			asgn.Rhs[0] = call.Call
			err := new(ast.Ident)
			err.Name = "err"
			asgn.Lhs = append(asgn.Lhs, err)
			check := newcheck(decl, call.Call)
			block.List = slices.Insert(block.List, i+1, check)
		}
		return true
	}
	ast.Inspect(decl, walker)
}

func unwrite(fset *token.FileSet, decl *ast.FuncDecl) {
	call := match.Call()
	matcher1 := match.Assign(match.Multi(match.Ident(""), match.Ident("err")), match.Multi(call))
	matcher2 := match.If(match.Binary(match.Ident("err"), token.NEQ, match.Ident("nil")),
		match.Block(match.Multi(match.Return(match.Multi(match.Ident("nil"), match.Ident("err"))))))
	walker := func(node ast.Node) bool {
		block, ok := node.(*ast.BlockStmt)
		if !ok {
			return true
		}
		for i := 0; i < len(block.List)-1; i++ {
			if !matcher1.Match(block.List[i]) {
				continue
			}
			if !matcher2.Match(block.List[i+1]) {
				ast.Print(fset, block.List[i+1])
				continue
			}
			asgn := matcher1.Assign
			op := new(ast.UnaryExpr)
			op.Op = token.XOR
			op.X = asgn.Rhs[0]
			asgn.Lhs = asgn.Lhs[:1]
			asgn.Rhs[0] = op
			block.List = slices.Delete(block.List, i+1, i+2)
		}
		return true
	}
	ast.Inspect(decl, walker)
}

func process(fset *token.FileSet, pkg *ast.Package, rewriter func(*token.FileSet, *ast.FuncDecl)) {
	walker := func(node ast.Node) bool {
		decl, ok := node.(*ast.FuncDecl)
		if !ok {
			return true
		}
		rewriter(fset, decl)
		return false
	}
	for _, file := range pkg.Files {
		ast.Inspect(file, walker)
		w := os.Stdout
		format.Node(w, fset, file)
		fmt.Printf("\n")
	}
}