func (v *ShortError) VisitExpr(scope *ast.Scope, expr ast.Expr) ScopeVisitor {
if expr, ok := expr.(*ast.CallExpr); ok {
if fun, ok := expr.Fun.(*ast.Ident); ok && fun.Name == MustKeyword {
if len(expr.Args) != 1 {
pos := v.file.Fset.Position(fun.Pos())
fmt.Printf("%s:%d:%d: 'must' builtin must be called with exactly one argument\n", pos.Filename, pos.Line, pos.Column)
return nil
}
tmpVar, tmpErr := v.tempVar("tmp_", scope), v.tempVar("err_", scope)
mustexpr := v.file.Get(expr.Args[0])
if v.block == nil {
// if in top level decleration
v.addToInit("if " + tmpErr + " != nil {panic(" + tmpErr + ")};")
*v.patches = append(*v.patches,
patch.Replace(expr, tmpVar),
patch.Insert(afterImports(v.file.File), ";var "+tmpVar+", "+tmpErr+" = "+mustexpr))
} else {
*v.patches = append(*v.patches, patch.Insert(v.stmt.Pos(),
fmt.Sprint("var ", tmpVar, ", ", tmpErr, " = ", mustexpr, "; ",
"if ", tmpErr, " != nil {panic(", tmpErr, ")};")))
*v.patches = append(*v.patches, patch.Replace(expr, tmpVar))
}
}
}
return v
}
func (v *ShortError) VisitStmt(scope *ast.Scope, stmt ast.Stmt) ScopeVisitor {
v.stmt = stmt
switch stmt := stmt.(type) {
case *ast.BlockStmt:
return &ShortError{v.file, v.patches, v.stmt, stmt, 0, new([]byte)}
case *ast.ExprStmt:
if call := calltomust(stmt.X); call != nil {
// TODO(elazarl): depends on number of variables it returns, currently we assume one
pos := v.file.Fset.Position(stmt.Pos())
fmt.Printf("%s:%d:%d: 'must' builtin must be assigned into variable\n",
pos.Filename, pos.Line, pos.Column)
}
case *ast.AssignStmt:
if len(stmt.Rhs) != 1 {
return v
}
if rhs, ok := stmt.Rhs[0].(*ast.CallExpr); ok {
if fun, ok := rhs.Fun.(*ast.Ident); ok && fun.Name == MustKeyword {
if stmt.Tok == token.DEFINE {
tmpVar := v.tempVar("assignerr_", scope)
*v.patches = append(*v.patches,
patch.Insert(stmt.TokPos, ", "+tmpVar+" "),
patch.Replace(fun, ""),
patch.Insert(stmt.End(),
"; if "+tmpVar+" != nil "+
"{ panic("+tmpVar+") };"),
)
for _, arg := range rhs.Args {
v.VisitExpr(scope, arg)
}
return nil
} else if stmt.Tok == token.ASSIGN {
vars := []string{}
for i := 0; i < len(stmt.Lhs); i++ {
vars = append(vars, v.tempVar(fmt.Sprint("assgn", i, "_"), scope))
}
assgnerr := v.tempVar("assgnErr_", scope)
*v.patches = append(*v.patches,
patch.Insert(stmt.Pos(),
strings.Join(append(vars, assgnerr), ", ")+":="),
patch.InsertNode(stmt.Pos(), rhs.Args[0]),
patch.Insert(stmt.Pos(),
"; if "+assgnerr+" != nil "+
"{ panic("+assgnerr+") };"),
patch.Replace(rhs, strings.Join(vars, ", ")),
)
v.VisitExpr(scope, rhs.Args[0])
return nil
}
}
}
}
return v
}
func (v *ShortError) ExitScope(scope *ast.Scope, node ast.Node, last bool) ScopeVisitor {
if node, ok := node.(*ast.File); ok && len(*v.initTxt) > 0 {
if init := findinit(node); init != nil {
*v.patches = append(*v.patches, patch.Insert(init.Body.Lbrace+1, string(*v.initTxt)))
} else {
*v.patches = append(*v.patches, patch.Insert(afterImports(node),
";func init() {"+string(*v.initTxt)+"}"))
}
}
return v
}
func (p *patchUnused) UnusedObj(obj *ast.Object, parent ast.Node) {
// if the unused variable is a function argument, or TLD - ignore
switch obj.Decl.(type) {
case *ast.Field, *ast.GenDecl, *ast.TypeSpec:
return
case *ast.ValueSpec:
if _, ok := parent.(*ast.File); ok {
return
}
}
if obj.Kind == ast.Fun {
return
}
exempter := "_ = " + obj.Name
switch parent := parent.(type) {
case *ast.ForStmt:
p.patches = append(p.patches, patch.Insert(parent.Body.Lbrace+1, exempter+";"))
case *ast.RangeStmt:
p.patches = append(p.patches, patch.Insert(parent.Body.Lbrace+1, exempter+";"))
case *ast.TypeSwitchStmt:
if len(parent.Body.List) == 0 {
p.patches = append(p.patches, patch.Insert(parent.Body.Lbrace+1, "default: "+exempter))
} else {
// if first statement is not case statement - it won't compile anyhow
if stmt, ok := parent.Body.List[0].(*ast.CaseClause); ok {
p.patches = append(p.patches, patch.Insert(stmt.Colon+1, exempter+";"))
}
}
case *ast.SwitchStmt:
if len(parent.Body.List) == 0 {
p.patches = append(p.patches, patch.Insert(parent.Body.Lbrace+1, "default: "+exempter))
} else {
// if first statement is not case statement - it won't compile anyhow
if stmt, ok := parent.Body.List[0].(*ast.CaseClause); ok {
p.patches = append(p.patches, patch.Insert(stmt.Colon+1, exempter+";"))
}
}
case *ast.CommClause:
if compareAssgn(obj.Decl, parent.Comm) {
p.patches = append(p.patches, patch.Insert(parent.Colon+1, exempter+";"))
} else {
p.patches = append(p.patches, patch.Insert(obj.Decl.(ast.Node).End(), ";"+exempter))
}
case *ast.IfStmt:
p.patches = append(p.patches, patch.Insert(parent.Body.Lbrace+1, exempter+";"))
default:
p.patches = append(p.patches, patch.Insert(obj.Decl.(ast.Node).End(), ";"+exempter))
}
}
func (v *ShortError) VisitDecl(scope *ast.Scope, decl ast.Decl) ScopeVisitor {
if decl, ok := decl.(*ast.GenDecl); ok {
for _, spec := range decl.Specs {
// We'll act only in cases like top level `var a, b, c = must(expr)`
if spec, ok := spec.(*ast.ValueSpec); ok && len(spec.Values) == 1 {
if fun, ok := spec.Values[0].(*ast.CallExpr); ok {
if name, ok := fun.Fun.(*ast.Ident); !ok || name.Name != MustKeyword {
return v
}
if len(fun.Args) != 1 {
pos := v.file.Fset.Position(fun.Pos())
fmt.Printf("%s:%d:%d: 'must' builtin must be called with exactly one argument\n", pos.Filename, pos.Line, pos.Column)
return nil
}
tmpErr := v.tempVar("tlderr_", scope)
*v.patches = append(*v.patches,
patch.Insert(spec.Names[len(spec.Names)-1].End(), ", "+tmpErr),
patch.Replace(fun, v.file.Get(fun.Args[0])))
v.addToInit("if " + tmpErr + " != nil { panic(" + tmpErr + ") }")
}
}
}
return nil
}
return v
}
func (p *patchUnused) UnusedImport(imp *ast.ImportSpec) {
if imp.Name != nil {
p.patches = append(p.patches, patch.Replace(imp.Name, "_"))
} else {
p.patches = append(p.patches, patch.Insert(imp.Pos(), "_ "))
}
}
func main() {
err := instrument.InstrumentCmd(func(p *patch.PatchableFile) (patches patch.Patches) {
for _, dec := range p.File.Decls {
if fun, ok := dec.(*ast.FuncDecl); ok && fun.Name != nil && fun.Body != nil {
patches = append(patches, patch.Insert(fun.Body.Lbrace+1, "println(`"+fun.Name.Name+"`);"))
}
}
return patches
}, os.Args...)
if err != nil {
println(err.Error())
os.Exit(-1)
}
}
func (v *AutoImporter) VisitExpr(scope *ast.Scope, expr ast.Expr) scopes.Visitor {
switch expr := expr.(type) {
case *ast.Ident:
if v.Irrelevant[expr] {
return v
}
if importname, ok := imports.RevStdlib[expr.Name]; ok && len(importname) == 1 &&
!v.m[expr.Name] && scopes.Lookup(scope, expr.Name) == nil {
v.m[expr.Name] = true // don't add it again
v.Patches = append(v.Patches, patch.Insert(v.pkg, "; import "+importname[0]))
}
case *ast.SelectorExpr:
v.Irrelevant[expr.Sel] = true
case *ast.KeyValueExpr:
// if we get a := struct {Count int} {Count: 1}, disregard Count
if id, ok := expr.Key.(*ast.Ident); ok {
v.Irrelevant[id] = true
}
}
return v
}