Skip to content

Instantly share code, notes, and snippets.

@tdakkota
Created December 29, 2020 11:20
Show Gist options
  • Save tdakkota/2c08b93ac5348d99a2928b956586c077 to your computer and use it in GitHub Desktop.
Save tdakkota/2c08b93ac5348d99a2928b956586c077 to your computer and use it in GitHub Desktop.
package main
import (
"fmt"
"go/ast"
"go/token"
"os"
"strings"
builders "github.com/tdakkota/astbuilders"
macro "github.com/tdakkota/gomacro"
"github.com/tdakkota/gomacro/runner"
)
func Macro() macro.Handler {
return macro.HandlerFunc(func(ctx macro.Context, node ast.Node) error {
if ctx.Pre { // skip first pass
return nil
}
stmts, ok := node.(*ast.BlockStmt)
if !ok {
return nil
}
newStmts := make([]ast.Stmt, len(stmts.List))
copy(newStmts, stmts.List)
for i, stmt := range stmts.List {
ex, ok := stmt.(*ast.ExprStmt)
if !ok {
continue
}
expr := ex.X
call, ok := expr.(*ast.CallExpr)
if !ok {
continue
}
for paramN, param := range call.Args {
ref, ok := param.(*ast.UnaryExpr)
if !ok || ref.Op != token.AND {
continue
}
paramCall, ok := ref.X.(*ast.CallExpr)
if !ok {
continue
}
paramName := ast.NewIdent(fmt.Sprintf("refcallParam%d%d", i, paramN))
call.Args[paramN] = builders.Ref(paramName)
d := builders.Define(paramName)(paramCall)
newStmts = append(newStmts[:i+1], newStmts[i:]...)
newStmts[i] = d
}
}
stmts.List = newStmts
ctx.Replace(stmts)
return nil
})
}
const src = `
package main
//procm:use=refcall
func main() {
{
println(&time.Now())
}
println("now")
println(&time.Now())
}
`
func main() {
err := runner.Reader("foo", strings.NewReader(src), os.Stdout, macro.Macros{
"refcall": Macro(),
})
if err != nil {
panic(err)
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment