Skip to content

Commit

Permalink
support scope for function call and function lit utils
Browse files Browse the repository at this point in the history
  • Loading branch information
zhenghe3119 committed Nov 6, 2019
1 parent 3a27ff9 commit c638f42
Show file tree
Hide file tree
Showing 9 changed files with 448 additions and 268 deletions.
16 changes: 11 additions & 5 deletions README.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -41,20 +41,26 @@ FprintFile(out io.Writer, df *dst.File) error
HasStmtInsideFuncBody(df *dst.File, funcName string, stmt dst.Stmt) (ret bool)
DeleteStmtFromFuncBody(df *dst.File, funcName string, stmt dst.Stmt) (modified bool)
AddStmtToFuncBody(df *dst.File, funcName string, stmt dst.Stmt, pos int) (modified bool)
AddStmtToFuncLitBody(df *dst.File, stmt dst.Stmt, pos int) (modified bool)
AddStmtToFuncBodyStart(df *dst.File, funcName string, stmt dst.Stmt) (modified bool)
AddStmtToFuncBodyEnd(df *dst.File, funcName string, stmt dst.Stmt) (modified bool)
AddStmtToFuncBodyBefore(df *dst.File, funcName string, stmt, refStmt dst.Stmt) (modified bool)
AddStmtToFuncBodyAfter(df *dst.File, funcName string, stmt, refStmt dst.Stmt) (modified bool)
```

### function lit utilities

```
AddStmtToFuncLitBody(df *dst.File, scope Scope, stmt dst.Stmt, pos int) (modified bool)
AddFieldToFuncLitParams(df *dst.File, scope Scope, field *dst.Field, pos int) (modified bool)
```

### function call utilities

```
HasArgInCallExpr(df *dst.File, funcName string, arg dst.Expr) (ret bool)
DeleteArgFromCallExpr(df *dst.File, funcName string, arg dst.Expr) (modified bool)
AddArgToCallExpr(df *dst.File, funcName string, arg dst.Expr, pos int) (modified bool)
SetMethodOnReceiver(df *dst.File, receiver, oldMethod, newMethod string) (modified bool)
HasArgInCallExpr(df *dst.File, scope Scope, funcName string, arg dst.Expr) (ret bool)
DeleteArgFromCallExpr(df *dst.File, scope Scope, funcName string, arg dst.Expr) (modified bool)
AddArgToCallExpr(df *dst.File, scope Scope, funcName string, arg dst.Expr, pos int) (modified bool)
SetMethodOnReceiver(df *dst.File, scope Scope, receiver, oldMethod, newMethod string) (modified bool)
```

### function declaration utilities
Expand Down
142 changes: 74 additions & 68 deletions examples/rpc_adapter.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,6 @@ package main
import (
"github.com/ZhengHe-MD/gorefactor"
"github.com/dave/dst"
"go/token"
"log"
"os"
)
Expand All @@ -12,9 +11,9 @@ func main() {
var src = `
package main
func rpc(ctx context.Context, hashKey string, timeout time.Duration, fn func(*AccountServiceClient) error) error {
return clientThrift.RpcWithContext(ctx, hashKey, timeout, func(c interface{}) error {
ct, ok := c.(*AccountServiceClient)
func rpc(haskkey string, timeout time.Duration, fn func(*AppServiceClient) error) error {
return clientThrift.Rpc(haskkey, timeout, func(c interface{}) error {
ct, ok := c.(*AppServiceClient)
if ok {
return fn(ct)
} else {
Expand All @@ -23,17 +22,18 @@ func main() {
})
}
func RegisterGuest(ctx context.Context, req *RegisterGuestReq) (r *RegisterGuestRes) {
func VersionCheck(ctx context.Context, req *VersionCheckReq) (r *VersionCheckRes) {
tctx := trace.NewThriftUtilContextFromContext(ctx)
err := rpc(ctx, strconv.Itoa(rand.Int()), time.Millisecond*1000,
func(c *AccountServiceClient) (er error) {
r, er = c.RegisterGuest(req, tctx)
err := rpc(strconv.FormatInt(req.Uid, 10), time.Millisecond*3000,
func(c *AppServiceClient) (er error) {
r, er = c.VersionCheck(req, tctx)
return er
},
)
if err != nil {
r = &RegisterGuestRes{
Errinfo: tiperr.NewInternalErr(-1001, fmt.Sprintf("call serice:%s proc:%s m:RegisterGuest err:%s", service_key, proc_thrift, err)),
r = &VersionCheckRes{
Errinfo: tiperr.NewInternalErr(-1001, fmt.Sprintf("call serice:%s proc:%s m:VersionCheck err:%s", service_key, proc_thrift, err)),
}
}
return
Expand Down Expand Up @@ -77,67 +77,73 @@ func main() {
return
}

gorefactor.SetMethodOnReceiver(df, "clientThrift", "RpcWithContext", "RpcWithContextV2")
gorefactor.AddArgToCallExpr(df, "fn", dst.NewIdent("fctx"), 0)
gorefactor.AddFieldToFuncLitParams(df, &dst.Field{
Names: []*dst.Ident{dst.NewIdent("fctx")},
Type: dst.NewIdent("context.Context"),
}, 0)
gorefactor.DeleteStmtFromFuncBody(df, "RegisterGuest", &dst.AssignStmt{
Lhs: []dst.Expr{
dst.NewIdent("tctx"),
},
Tok: token.DEFINE,
Rhs: []dst.Expr{
&dst.CallExpr{
Fun: &dst.SelectorExpr{
X: dst.NewIdent("trace"),
Sel: dst.NewIdent("NewThriftUtilContextFromContext"),
},
Args: []dst.Expr{
dst.NewIdent("ctx"),
},
},
},
})
gorefactor.AddStmtToFuncLitBody(df, &dst.AssignStmt{
Lhs: []dst.Expr{
dst.NewIdent("tctx"),
},
Tok: token.DEFINE,
Rhs: []dst.Expr{
&dst.CallExpr{
Fun: &dst.SelectorExpr{
X: dst.NewIdent("trace"),
Sel: dst.NewIdent("NewThriftUtilContextFromContext"),
},
Args: []dst.Expr{
dst.NewIdent("tctx"),
},
},
},
}, 0)
gorefactor.DeleteStmtFromFuncBody(df, "rpc", &dst.AssignStmt{
Lhs: []dst.Expr{
dst.NewIdent("tctx"),
},
Tok: token.DEFINE,
Rhs: []dst.Expr{
&dst.CallExpr{
Fun: &dst.SelectorExpr{
X: dst.NewIdent("trace"),
Sel: dst.NewIdent("NewThriftUtilContextFromContext"),
},
Args: []dst.Expr{
dst.NewIdent("tctx"),
},
},
},
})
//gorefactor.AddFieldToFuncDeclParams(df, "rpc", &dst.Field{
// Names: []*dst.Ident{dst.NewIdent("ctx")},
// Type: dst.NewIdent("context.Contex"),
//}, 0)
gorefactor.AddFieldToFuncDeclParams(df, "fn", &dst.Field{Type: dst.NewIdent("context.Context")}, 0)
//gorefactor.SetMethodOnReceiver(df, "clientThrift", "Rpc", "RpcWithContextV2")
//gorefactor.AddArgToCallExpr(df, "RpcWithContextV2", dst.NewIdent("ctx"), 0)
//gorefactor.AddArgToCallExpr(df, "fn", dst.NewIdent("fctx"), 0)
//gorefactor.AddFieldToFuncLitParams(df, &dst.Field{
// Names: []*dst.Ident{dst.NewIdent("fctx")},
// Type: dst.NewIdent("context.Context"),
//}, 0)
//gorefactor.DeleteStmtFromFuncBody(df, "VersionCheck", &dst.AssignStmt{
// Lhs: []dst.Expr{
// dst.NewIdent("tctx"),
// },
// Tok: token.DEFINE,
// Rhs: []dst.Expr{
// &dst.CallExpr{
// Fun: &dst.SelectorExpr{
// X: dst.NewIdent("trace"),
// Sel: dst.NewIdent("NewThriftUtilContextFromContext"),
// },
// Args: []dst.Expr{
// dst.NewIdent("ctx"),
// },
// },
// },
//})
//gorefactor.AddStmtToFuncLitBody(df, &dst.AssignStmt{
// Lhs: []dst.Expr{
// dst.NewIdent("tctx"),
// },
// Tok: token.DEFINE,
// Rhs: []dst.Expr{
// &dst.CallExpr{
// Fun: &dst.SelectorExpr{
// X: dst.NewIdent("trace"),
// Sel: dst.NewIdent("NewThriftUtilContextFromContext"),
// },
// Args: []dst.Expr{
// dst.NewIdent("fctx"),
// },
// },
// },
//}, 0)
//gorefactor.DeleteStmtFromFuncBody(df, "rpc", &dst.AssignStmt{
// Lhs: []dst.Expr{
// dst.NewIdent("tctx"),
// },
// Tok: token.DEFINE,
// Rhs: []dst.Expr{
// &dst.CallExpr{
// Fun: &dst.SelectorExpr{
// X: dst.NewIdent("trace"),
// Sel: dst.NewIdent("NewThriftUtilContextFromContext"),
// },
// Args: []dst.Expr{
// dst.NewIdent("tctx"),
// },
// },
// },
//})

err = gorefactor.FprintFile(os.Stdout, df)
if err != nil {
log.Println(err)
return
}
}
}
22 changes: 0 additions & 22 deletions func_body.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -149,28 +149,6 @@ func AddStmtToFuncBody(df *dst.File, funcName string, stmt dst.Stmt, pos int) (m
return
}

func AddStmtToFuncLitBody(df *dst.File, stmt dst.Stmt, pos int) (modified bool) {
pre := func(c *dstutil.Cursor) bool {
node := c.Node()

switch node.(type) {
case *dst.FuncLit:
nn := node.(*dst.FuncLit)
stmtList := nn.Body.List
pos = normalizePos(pos, len(stmtList))

nn.Body.List = append(
stmtList[:pos],
append([]dst.Stmt{dst.Clone(stmt).(dst.Stmt)}, stmtList[pos:]...)...)
modified = true
}
return true
}

dstutil.Apply(df, pre, nil)
return
}

// AddStmtToFuncBodyStart adds given statement, to the start of function body
func AddStmtToFuncBodyStart(df *dst.File, funcName string, stmt dst.Stmt) (modified bool) {
return AddStmtToFuncBody(df, funcName, stmt, 0)
Expand Down
45 changes: 38 additions & 7 deletions func_call.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -51,12 +51,16 @@ func HasArgInCallExpr(df *dst.File, scope Scope, funcName string, arg dst.Expr)

// DeleteArgFromCallExpr deletes any arg, in the function call's argument list,
// that is semantically equal to the given arg.
func DeleteArgFromCallExpr(df *dst.File, funcName string, arg dst.Expr) (modified bool) {
func DeleteArgFromCallExpr(df *dst.File, scope Scope, funcName string, arg dst.Expr) (modified bool) {
pre := func(c *dstutil.Cursor) bool {
node := c.Node()

switch node.(type) {
case *dst.CallExpr:
if !scope.IsInScope() {
return true
}

var found bool
nn := node.(*dst.CallExpr)
if ie, ok := nn.Fun.(*dst.Ident); ok && ie.Name == funcName {
Expand All @@ -79,23 +83,33 @@ func DeleteArgFromCallExpr(df *dst.File, funcName string, arg dst.Expr) (modifie
nn.Args = newArgs
return false
}
default:
scope.TryEnterScope(node)
}
return true
}

dstutil.Apply(df, pre, nil)
post := func(c *dstutil.Cursor) bool {
scope.TryLeaveScope(c.Node())
return true
}

dstutil.Apply(df, pre, post)
return
}

// AddArgToCallExpr adds given arg, to the function call's argument list, in the given position
func AddArgToCallExpr(df *dst.File, funcName string, arg dst.Expr, pos int) (modified bool) {
func AddArgToCallExpr(df *dst.File, scope Scope, funcName string, arg dst.Expr, pos int) (modified bool) {
pre := func(c *dstutil.Cursor) bool {
node := c.Node()

switch node.(type) {
case *dst.CallExpr:
nn := node.(*dst.CallExpr)
if !scope.IsInScope() {
return true
}

nn := node.(*dst.CallExpr)
var ce *dst.CallExpr

if ie, ok := nn.Fun.(*dst.Ident); ok && ie.Name == funcName {
Expand All @@ -115,20 +129,30 @@ func AddArgToCallExpr(df *dst.File, funcName string, arg dst.Expr, pos int) (mod
modified = true
return false
}
default:
scope.TryEnterScope(node)
}
return true
}

dstutil.Apply(df, pre, nil)
post := func(c *dstutil.Cursor) bool {
scope.TryLeaveScope(c.Node())
return true
}

dstutil.Apply(df, pre, post)
return
}

func SetMethodOnReceiver(df *dst.File, receiver, oldMethod, newMethod string) (modified bool) {
func SetMethodOnReceiver(df *dst.File, scope Scope, receiver, oldMethod, newMethod string) (modified bool) {
pre := func(c *dstutil.Cursor) bool {
node := c.Node()

switch node.(type) {
case *dst.CallExpr:
if !scope.IsInScope() {
return true
}
nn := node.(*dst.CallExpr)

switch nn.Fun.(type) {
Expand All @@ -154,10 +178,17 @@ func SetMethodOnReceiver(df *dst.File, receiver, oldMethod, newMethod string) (m
modified = true
}
}
default:
scope.TryEnterScope(node)
}
return true
}

dstutil.Apply(df, pre, nil)
post := func(c *dstutil.Cursor) bool {
scope.TryLeaveScope(c.Node())
return true
}

dstutil.Apply(df, pre, post)
return
}
Loading

0 comments on commit c638f42

Please sign in to comment.