github.com/yunabe/lgo@v0.0.0-20190709125917-42c42d410fdf/converter/rewriteexpr.go (about)

     1  // This file defines rewriteExpr
     2  
     3  package converter
     4  
     5  import (
     6  	"go/ast"
     7  )
     8  
     9  // rewriteExpr rewrites expressions under node.
    10  func rewriteExpr(node ast.Node, rewrite func(ast.Expr) ast.Expr) {
    11  	ast.Walk(rewiteVisitor{rewrite}, node)
    12  }
    13  
    14  type rewiteVisitor struct {
    15  	rewrite func(ast.Expr) ast.Expr
    16  }
    17  
    18  func rewriteExprList(rewrite func(ast.Expr) ast.Expr, list []ast.Expr) {
    19  	for i, x := range list {
    20  		list[i] = rewrite(x)
    21  	}
    22  }
    23  
    24  func (v rewiteVisitor) Visit(node ast.Node) (w ast.Visitor) {
    25  	switch n := node.(type) {
    26  	case *ast.Field:
    27  		n.Type = v.rewrite(n.Type)
    28  
    29  	case *ast.Ellipsis:
    30  		if n.Elt != nil {
    31  			n.Elt = v.rewrite(n.Elt)
    32  		}
    33  
    34  	case *ast.CompositeLit:
    35  		if n.Type != nil {
    36  			n.Type = v.rewrite(n.Type)
    37  		}
    38  		rewriteExprList(v.rewrite, n.Elts)
    39  
    40  	case *ast.ParenExpr:
    41  		n.X = v.rewrite(n.X)
    42  
    43  	case *ast.SelectorExpr:
    44  		n.X = v.rewrite(n.X)
    45  
    46  	case *ast.IndexExpr:
    47  		n.X = v.rewrite(n.X)
    48  		n.Index = v.rewrite(n.Index)
    49  
    50  	case *ast.SliceExpr:
    51  		n.X = v.rewrite(n.X)
    52  		if n.Low != nil {
    53  			n.Low = v.rewrite(n.Low)
    54  		}
    55  		if n.High != nil {
    56  			n.High = v.rewrite(n.High)
    57  		}
    58  		if n.Max != nil {
    59  			n.Max = v.rewrite(n.Max)
    60  		}
    61  
    62  	case *ast.TypeAssertExpr:
    63  		n.X = v.rewrite(n.X)
    64  		if n.Type != nil {
    65  			n.Type = v.rewrite(n.Type)
    66  		}
    67  
    68  	case *ast.CallExpr:
    69  		n.Fun = v.rewrite(n.Fun)
    70  		rewriteExprList(v.rewrite, n.Args)
    71  
    72  	case *ast.StarExpr:
    73  		n.X = v.rewrite(n.X)
    74  
    75  	case *ast.UnaryExpr:
    76  		n.X = v.rewrite(n.X)
    77  
    78  	case *ast.BinaryExpr:
    79  		n.X = v.rewrite(n.X)
    80  		n.Y = v.rewrite(n.Y)
    81  
    82  	case *ast.KeyValueExpr:
    83  		n.Key = v.rewrite(n.Key)
    84  		n.Value = v.rewrite(n.Value)
    85  
    86  	// Types
    87  	case *ast.ArrayType:
    88  		if n.Len != nil {
    89  			n.Len = v.rewrite(n.Len)
    90  		}
    91  		n.Elt = v.rewrite(n.Elt)
    92  
    93  	case *ast.MapType:
    94  		n.Key = v.rewrite(n.Key)
    95  		n.Value = v.rewrite(n.Value)
    96  
    97  	case *ast.ChanType:
    98  		n.Value = v.rewrite(n.Value)
    99  
   100  	case *ast.ExprStmt:
   101  		n.X = v.rewrite(n.X)
   102  
   103  	case *ast.SendStmt:
   104  		n.Chan = v.rewrite(n.Chan)
   105  		n.Value = v.rewrite(n.Value)
   106  
   107  	case *ast.IncDecStmt:
   108  		n.X = v.rewrite(n.X)
   109  
   110  	case *ast.AssignStmt:
   111  		rewriteExprList(v.rewrite, n.Lhs)
   112  		rewriteExprList(v.rewrite, n.Rhs)
   113  
   114  	case *ast.ReturnStmt:
   115  		rewriteExprList(v.rewrite, n.Results)
   116  
   117  	case *ast.IfStmt:
   118  		n.Cond = v.rewrite(n.Cond)
   119  
   120  	case *ast.CaseClause:
   121  		rewriteExprList(v.rewrite, n.List)
   122  
   123  	case *ast.SwitchStmt:
   124  		if n.Tag != nil {
   125  			n.Tag = v.rewrite(n.Tag)
   126  		}
   127  
   128  	case *ast.ForStmt:
   129  		if n.Cond != nil {
   130  			n.Cond = v.rewrite(n.Cond)
   131  		}
   132  
   133  	case *ast.RangeStmt:
   134  		if n.Key != nil {
   135  			n.Key = v.rewrite(n.Key)
   136  		}
   137  		if n.Value != nil {
   138  			n.Value = v.rewrite(n.Value)
   139  		}
   140  		n.X = v.rewrite(n.X)
   141  
   142  	case *ast.ValueSpec:
   143  		if n.Type != nil {
   144  			n.Type = v.rewrite(n.Type)
   145  		}
   146  		rewriteExprList(v.rewrite, n.Values)
   147  
   148  	case *ast.TypeSpec:
   149  		n.Type = v.rewrite(n.Type)
   150  	}
   151  	return v
   152  }