transforms.mx raw
1 package main
2
3 func transformAll(f *File) {
4 for i := int32(0); i < int32(len(f.DeclList)); i++ {
5 f.DeclList[i] = transformDecl(f.DeclList[i])
6 }
7 transformStringsImport(f)
8 transformInit(f)
9 }
10
11 func transformDecl(d Decl) Decl {
12 switch d := d.(type) {
13 case *ConstDecl:
14 if d.Type != nil {
15 d.Type = transformTypeExpr(d.Type)
16 }
17 if d.Values != nil {
18 d.Values = transformExpr(d.Values)
19 }
20 case *TypeDecl:
21 d.Type = transformTypeExpr(d.Type)
22 case *VarDecl:
23 if d.Type != nil {
24 d.Type = transformTypeExpr(d.Type)
25 }
26 if d.Values != nil {
27 d.Values = transformExpr(d.Values)
28 }
29 case *FuncDecl:
30 if d.Recv != nil {
31 transformField(d.Recv)
32 }
33 transformFuncType(d.Type)
34 if d.Body != nil {
35 transformBlock(d.Body)
36 }
37 }
38 return d
39 }
40
41 func transformField(f *Field) {
42 if f.Type != nil {
43 f.Type = transformTypeExpr(f.Type)
44 }
45 }
46
47 func transformFields(fields []*Field) {
48 for _, f := range fields {
49 transformField(f)
50 }
51 }
52
53 func transformFuncType(t *FuncType) {
54 if t == nil {
55 return
56 }
57 transformFields(t.ParamList)
58 transformFields(t.ResultList)
59 }
60
61 func transformBlock(b *BlockStmt) {
62 for i := int32(0); i < int32(len(b.List)); i++ {
63 b.List[i] = transformStmt(b.List[i])
64 }
65 }
66
67 func transformTypeExpr(e Expr) Expr {
68 if e == nil {
69 return nil
70 }
71 switch e := e.(type) {
72 case *Name:
73 if e.Value == "int" {
74 e.Value = "int64"
75 } else if e.Value == "uint" {
76 e.Value = "uint64"
77 }
78 return e
79
80 case *ArrayType:
81 if e.Len != nil {
82 e.Len = transformExpr(e.Len)
83 }
84 e.Elem = transformTypeExpr(e.Elem)
85 return e
86
87 case *SliceType:
88 e.Elem = transformTypeExpr(e.Elem)
89 return e
90
91 case *DotsType:
92 e.Elem = transformTypeExpr(e.Elem)
93 return e
94
95 case *MapType:
96 e.Key = transformTypeExpr(e.Key)
97 e.Value = transformTypeExpr(e.Value)
98 return e
99
100 case *ChanType:
101 e.Elem = transformTypeExpr(e.Elem)
102 return e
103
104 case *StructType:
105 transformFields(e.FieldList)
106 return e
107
108 case *InterfaceType:
109 transformFields(e.MethodList)
110 return e
111
112 case *FuncType:
113 transformFuncType(e)
114 return e
115
116 case *Operation:
117 if e.Y == nil {
118 e.X = transformTypeExpr(e.X)
119 }
120 return e
121
122 case *ParenExpr:
123 e.X = transformTypeExpr(e.X)
124 return e
125
126 case *SelectorExpr:
127 return e
128
129 case *IndexExpr:
130 e.X = transformTypeExpr(e.X)
131 e.Index = transformTypeExpr(e.Index)
132 return e
133 }
134 return e
135 }
136
137 func transformExpr(e Expr) Expr {
138 if e == nil {
139 return nil
140 }
141 switch e := e.(type) {
142 case *Name:
143 if e.Value == "int" {
144 e.Value = "int64"
145 } else if e.Value == "uint" {
146 e.Value = "uint64"
147 }
148 return e
149
150 case *Operation:
151 e.X = transformExpr(e.X)
152 if e.Y != nil {
153 e.Y = transformExpr(e.Y)
154 if e.Op == Add && isStringExpr(e.X, e.Y) {
155 e.Op = Or
156 }
157 }
158 return e
159
160 case *CallExpr:
161 return transformCall(e)
162
163 case *CompositeLit:
164 if e.Type != nil {
165 e.Type = transformTypeExpr(e.Type)
166 }
167 for i := int32(0); i < int32(len(e.ElemList)); i++ {
168 e.ElemList[i] = transformExpr(e.ElemList[i])
169 }
170 return e
171
172 case *KeyValueExpr:
173 e.Key = transformExpr(e.Key)
174 e.Value = transformExpr(e.Value)
175 return e
176
177 case *FuncLit:
178 transformFuncType(e.Type)
179 if e.Body != nil {
180 transformBlock(e.Body)
181 }
182 return e
183
184 case *ParenExpr:
185 e.X = transformExpr(e.X)
186 return e
187
188 case *SelectorExpr:
189 e.X = transformExpr(e.X)
190 if n, ok := e.X.(*Name); ok && n.Value == "strings" {
191 n.Value = "bytes"
192 }
193 return e
194
195 case *IndexExpr:
196 e.X = transformExpr(e.X)
197 e.Index = transformExpr(e.Index)
198 return e
199
200 case *SliceExpr:
201 e.X = transformExpr(e.X)
202 for i := 0; i < 3; i++ {
203 if e.Index[i] != nil {
204 e.Index[i] = transformExpr(e.Index[i])
205 }
206 }
207 return e
208
209 case *AssertExpr:
210 e.X = transformExpr(e.X)
211 if e.Type != nil {
212 e.Type = transformTypeExpr(e.Type)
213 }
214 return e
215
216 case *TypeSwitchGuard:
217 e.X = transformExpr(e.X)
218 return e
219
220 case *ListExpr:
221 for i := int32(0); i < int32(len(e.ElemList)); i++ {
222 e.ElemList[i] = transformExpr(e.ElemList[i])
223 }
224 return e
225
226 case *BasicLit:
227 return e
228 }
229 return e
230 }
231
232 func transformCall(e *CallExpr) Expr {
233 name, ok := e.Fun.(*Name)
234 if !ok {
235 e.Fun = transformExpr(e.Fun)
236 for i := int32(0); i < int32(len(e.ArgList)); i++ {
237 e.ArgList[i] = transformExpr(e.ArgList[i])
238 }
239 return e
240 }
241
242 switch name.Value {
243 case "make":
244 return transformMake(e)
245 case "new":
246 return transformNew(e)
247 case "int":
248 name.Value = "int64"
249 case "uint":
250 name.Value = "uint64"
251 }
252
253 for i := int32(0); i < int32(len(e.ArgList)); i++ {
254 e.ArgList[i] = transformExpr(e.ArgList[i])
255 }
256 return e
257 }
258
259 func transformMake(e *CallExpr) Expr {
260 nargs := int32(len(e.ArgList))
261 if nargs < 1 {
262 return e
263 }
264 typeArg := e.ArgList[0]
265
266 switch t := typeArg.(type) {
267 case *ChanType:
268 t.Elem = transformTypeExpr(t.Elem)
269 lit := &CompositeLit{Type: t}
270 if nargs >= 2 {
271 lit.ElemList = []Expr{transformExpr(e.ArgList[1])}
272 }
273 return lit
274
275 case *MapType:
276 t.Key = transformTypeExpr(t.Key)
277 t.Value = transformTypeExpr(t.Value)
278 return &CompositeLit{Type: t}
279
280 case *SliceType:
281 t.Elem = transformTypeExpr(t.Elem)
282 lit := &CompositeLit{Type: t}
283 bad := &BadExpr{}
284 if nargs == 2 {
285 kv := &KeyValueExpr{Key: bad, Value: transformExpr(e.ArgList[1])}
286 lit.ElemList = []Expr{kv}
287 lit.NKeys = 1
288 } else if nargs == 3 {
289 val := &ListExpr{ElemList: []Expr{
290 transformExpr(e.ArgList[1]),
291 transformExpr(e.ArgList[2]),
292 }}
293 kv := &KeyValueExpr{Key: bad, Value: val}
294 lit.ElemList = []Expr{kv}
295 lit.NKeys = 1
296 }
297 return lit
298
299 case *ArrayType:
300 if t.Len != nil {
301 return e
302 }
303 st := &SliceType{Elem: t.Elem}
304 newArgs := []Expr{st}
305 newArgs = append(newArgs, e.ArgList[1:]...)
306 return transformMake(&CallExpr{Fun: e.Fun, ArgList: newArgs})
307 }
308
309 return e
310 }
311
312 func transformNew(e *CallExpr) Expr {
313 if int32(len(e.ArgList)) != 1 {
314 return e
315 }
316 typeArg := transformTypeExpr(e.ArgList[0])
317 return &Operation{
318 Op: And,
319 X: &CompositeLit{Type: typeArg},
320 }
321 }
322
323 func isStringExpr(x, y Expr) bool {
324 return isStringLit(x) || isStringLit(y)
325 }
326
327 func isStringLit(e Expr) bool {
328 lit, ok := e.(*BasicLit)
329 if !ok {
330 return false
331 }
332 return lit.Kind == StringLit
333 }
334
335 func transformStmt(s Stmt) Stmt {
336 if s == nil {
337 return nil
338 }
339 switch s := s.(type) {
340 case *ExprStmt:
341 s.X = transformExpr(s.X)
342
343 case *SendStmt:
344 s.Chan = transformExpr(s.Chan)
345 s.Value = transformExpr(s.Value)
346
347 case *AssignStmt:
348 if s.Lhs != nil {
349 s.Lhs = transformExpr(s.Lhs)
350 }
351 if s.Rhs != nil {
352 s.Rhs = transformExpr(s.Rhs)
353 }
354
355 case *ReturnStmt:
356 if s.Results != nil {
357 s.Results = transformExpr(s.Results)
358 }
359
360 case *IfStmt:
361 if s.Init != nil {
362 s.Init = transformSimpleStmt(s.Init)
363 }
364 s.Cond = transformExpr(s.Cond)
365 transformBlock(s.Then)
366 if s.Else != nil {
367 s.Else = transformStmt(s.Else)
368 }
369
370 case *ForStmt:
371 if s.Init != nil {
372 s.Init = transformSimpleStmt(s.Init)
373 }
374 if s.Cond != nil {
375 s.Cond = transformExpr(s.Cond)
376 }
377 if s.Post != nil {
378 s.Post = transformSimpleStmt(s.Post)
379 }
380 transformBlock(s.Body)
381
382 case *SwitchStmt:
383 if s.Init != nil {
384 s.Init = transformSimpleStmt(s.Init)
385 }
386 if s.Tag != nil {
387 s.Tag = transformExpr(s.Tag)
388 }
389 for _, cc := range s.Body {
390 if cc.Cases != nil {
391 cc.Cases = transformExpr(cc.Cases)
392 }
393 for i := int32(0); i < int32(len(cc.Body)); i++ {
394 cc.Body[i] = transformStmt(cc.Body[i])
395 }
396 }
397 mergeFallthrough(s)
398
399 case *SelectStmt:
400 for _, cc := range s.Body {
401 if cc.Comm != nil {
402 cc.Comm = transformSimpleStmt(cc.Comm)
403 }
404 for i := int32(0); i < int32(len(cc.Body)); i++ {
405 cc.Body[i] = transformStmt(cc.Body[i])
406 }
407 }
408
409 case *BlockStmt:
410 transformBlock(s)
411
412 case *DeclStmt:
413 for i := int32(0); i < int32(len(s.DeclList)); i++ {
414 s.DeclList[i] = transformDecl(s.DeclList[i])
415 }
416
417 case *CallStmt:
418 s.Call = transformExpr(s.Call)
419
420 case *LabeledStmt:
421 s.Stmt = transformStmt(s.Stmt)
422 }
423 return s
424 }
425
426 func transformSimpleStmt(s SimpleStmt) SimpleStmt {
427 if s == nil {
428 return nil
429 }
430 switch s := s.(type) {
431 case *RangeClause:
432 if s.Lhs != nil {
433 s.Lhs = transformExpr(s.Lhs)
434 }
435 s.X = transformExpr(s.X)
436 case *ExprStmt:
437 s.X = transformExpr(s.X)
438 case *SendStmt:
439 s.Chan = transformExpr(s.Chan)
440 s.Value = transformExpr(s.Value)
441 case *AssignStmt:
442 if s.Lhs != nil {
443 s.Lhs = transformExpr(s.Lhs)
444 }
445 if s.Rhs != nil {
446 s.Rhs = transformExpr(s.Rhs)
447 }
448 }
449 return s
450 }
451
452 func mergeFallthrough(s *SwitchStmt) {
453 for i := int32(len(s.Body)) - 2; i >= 0; i-- {
454 cc := s.Body[i]
455 n := int32(len(cc.Body))
456 if n == 0 {
457 continue
458 }
459 br, ok := cc.Body[n-1].(*BranchStmt)
460 if !ok || br.Tok != Fallthrough {
461 continue
462 }
463 next := s.Body[i+1]
464 if n == 1 {
465 next.Cases = mergeCases(cc.Cases, next.Cases)
466 s.Body = append(s.Body[:i], s.Body[i+1:]...)
467 } else {
468 cc.Body = cc.Body[:n-1]
469 }
470 }
471 }
472
473 func mergeCases(a, b Expr) Expr {
474 if a == nil {
475 return b
476 }
477 if b == nil {
478 return a
479 }
480 var elems []Expr
481 if list, ok := a.(*ListExpr); ok {
482 elems = append(elems, list.ElemList...)
483 } else {
484 elems = append(elems, a)
485 }
486 if list, ok := b.(*ListExpr); ok {
487 elems = append(elems, list.ElemList...)
488 } else {
489 elems = append(elems, b)
490 }
491 return &ListExpr{ElemList: elems}
492 }
493
494 func transformStringsImport(f *File) {
495 for _, d := range f.DeclList {
496 imp, ok := d.(*ImportDecl)
497 if !ok {
498 continue
499 }
500 if imp.Path != nil && imp.Path.Value == "\"strings\"" {
501 imp.Path.Value = "\"bytes\""
502 }
503 }
504 }
505
506 func transformInit(f *File) {
507 initCount := int32(0)
508 var mainDecl *FuncDecl
509 var initCalls []Expr
510
511 for i := int32(0); i < int32(len(f.DeclList)); i++ {
512 fd, ok := f.DeclList[i].(*FuncDecl)
513 if !ok {
514 continue
515 }
516 if fd.Recv != nil {
517 continue
518 }
519 if fd.Name.Value == "main" {
520 mainDecl = fd
521 }
522 if fd.Name.Value == "init" {
523 newName := "init_0"
524 if initCount > 0 {
525 newName = "init_" | intToStr(initCount)
526 }
527 fd.Name.Value = newName
528 initCalls = append(initCalls, &CallExpr{
529 Fun: &Name{Value: newName},
530 })
531 initCount++
532 }
533 }
534
535 if initCount > 0 && mainDecl != nil && mainDecl.Body != nil {
536 var stmts []Stmt
537 for _, call := range initCalls {
538 stmts = append(stmts, &ExprStmt{X: call})
539 }
540 stmts = append(stmts, mainDecl.Body.List...)
541 mainDecl.Body.List = stmts
542 }
543 }
544
545 func intToStr(n int32) string {
546 if n == 0 {
547 return "0"
548 }
549 var buf [10]byte
550 i := int32(len(buf))
551 for n > 0 {
552 i--
553 buf[i] = byte('0' + n%10)
554 n /= 10
555 }
556 return string(buf[i:])
557 }
558