printer.mx raw

   1  package main
   2  
   3  import "io"
   4  
   5  type printer struct {
   6  	buf   []byte
   7  	depth int32
   8  }
   9  
  10  func printFile(w io.Writer, f *File) {
  11  	var p printer
  12  	p.file(f)
  13  	w.Write(p.buf)
  14  }
  15  
  16  func (p *printer) w(s string) {
  17  	p.buf = append(p.buf, s...)
  18  }
  19  
  20  func (p *printer) ch(c byte) {
  21  	p.buf = append(p.buf, c)
  22  }
  23  
  24  func (p *printer) nl() {
  25  	p.ch('\n')
  26  }
  27  
  28  func (p *printer) tab() {
  29  	for i := int32(0); i < p.depth; i++ {
  30  		p.ch('\t')
  31  	}
  32  }
  33  
  34  func (p *printer) space() {
  35  	p.ch(' ')
  36  }
  37  
  38  func declGroup(d Decl) *Group {
  39  	switch d := d.(type) {
  40  	case *ImportDecl:
  41  		return d.Group
  42  	case *ConstDecl:
  43  		return d.Group
  44  	case *TypeDecl:
  45  		return d.Group
  46  	case *VarDecl:
  47  		return d.Group
  48  	}
  49  	return nil
  50  }
  51  
  52  func declKeyword(d Decl) string {
  53  	switch d.(type) {
  54  	case *ImportDecl:
  55  		return "import"
  56  	case *ConstDecl:
  57  		return "const"
  58  	case *TypeDecl:
  59  		return "type"
  60  	case *VarDecl:
  61  		return "var"
  62  	case *FuncDecl:
  63  		return "func"
  64  	}
  65  	return ""
  66  }
  67  
  68  func (p *printer) file(f *File) {
  69  	p.w("package ")
  70  	p.w(f.PkgName.Value)
  71  	p.nl()
  72  
  73  	var prevGroup *Group
  74  	prevWasGrouped := false
  75  
  76  	for i := int32(0); i < int32(len(f.DeclList)); i++ {
  77  		d := f.DeclList[i]
  78  		g := declGroup(d)
  79  
  80  		if g != prevGroup {
  81  			if prevGroup != nil {
  82  				p.depth--
  83  				p.w(")")
  84  				p.nl()
  85  			}
  86  			if g != nil {
  87  				p.nl()
  88  				p.w(declKeyword(d))
  89  				p.w(" (")
  90  				p.nl()
  91  				p.depth++
  92  			}
  93  		}
  94  
  95  		if g != nil {
  96  			p.tab()
  97  			p.declBody(d)
  98  			p.nl()
  99  			prevWasGrouped = true
 100  		} else {
 101  			if prevWasGrouped || i > 0 {
 102  				p.nl()
 103  			}
 104  			p.decl(d)
 105  			p.nl()
 106  			prevWasGrouped = false
 107  		}
 108  
 109  		prevGroup = g
 110  	}
 111  
 112  	if prevGroup != nil {
 113  		p.depth--
 114  		p.w(")")
 115  		p.nl()
 116  	}
 117  }
 118  
 119  func (p *printer) decl(d Decl) {
 120  	p.w(declKeyword(d))
 121  	p.space()
 122  	p.declBody(d)
 123  }
 124  
 125  func (p *printer) declBody(d Decl) {
 126  	switch d := d.(type) {
 127  	case *ImportDecl:
 128  		if d.LocalPkgName != nil {
 129  			p.w(d.LocalPkgName.Value)
 130  			p.space()
 131  		}
 132  		p.expr(d.Path)
 133  
 134  	case *ConstDecl:
 135  		p.nameList(d.NameList)
 136  		if d.Type != nil {
 137  			p.space()
 138  			p.expr(d.Type)
 139  		}
 140  		if d.Values != nil {
 141  			p.w(" = ")
 142  			p.expr(d.Values)
 143  		}
 144  
 145  	case *TypeDecl:
 146  		p.w(d.Name.Value)
 147  		if len(d.TParamList) > 0 {
 148  			p.typeParams(d.TParamList)
 149  		}
 150  		if d.Alias {
 151  			p.w(" = ")
 152  		} else {
 153  			p.space()
 154  		}
 155  		p.expr(d.Type)
 156  
 157  	case *VarDecl:
 158  		p.nameList(d.NameList)
 159  		if d.Type != nil {
 160  			p.space()
 161  			p.expr(d.Type)
 162  		}
 163  		if d.Values != nil {
 164  			p.w(" = ")
 165  			p.expr(d.Values)
 166  		}
 167  
 168  	case *FuncDecl:
 169  		p.funcDecl(d)
 170  	}
 171  }
 172  
 173  func (p *printer) funcDecl(d *FuncDecl) {
 174  	if d.Recv != nil {
 175  		p.ch('(')
 176  		p.field(d.Recv)
 177  		p.w(") ")
 178  	}
 179  	p.w(d.Name.Value)
 180  	if len(d.TParamList) > 0 {
 181  		p.typeParams(d.TParamList)
 182  	}
 183  	p.funcType(d.Type)
 184  	if d.Body != nil {
 185  		p.space()
 186  		p.blockStmt(d.Body)
 187  	}
 188  }
 189  
 190  func (p *printer) nameList(names []*Name) {
 191  	for i, n := range names {
 192  		if i > 0 {
 193  			p.w(", ")
 194  		}
 195  		p.w(n.Value)
 196  	}
 197  }
 198  
 199  func (p *printer) typeParams(fields []*Field) {
 200  	p.ch('[')
 201  	for i, f := range fields {
 202  		if i > 0 {
 203  			p.w(", ")
 204  		}
 205  		p.field(f)
 206  	}
 207  	p.ch(']')
 208  }
 209  
 210  func (p *printer) field(f *Field) {
 211  	if f.Name != nil {
 212  		p.w(f.Name.Value)
 213  		if f.Type != nil {
 214  			p.space()
 215  		}
 216  	}
 217  	if f.Type != nil {
 218  		p.expr(f.Type)
 219  	}
 220  }
 221  
 222  func (p *printer) fieldList(fields []*Field, sep string) {
 223  	for i, f := range fields {
 224  		if i > 0 {
 225  			p.w(sep)
 226  		}
 227  		p.field(f)
 228  	}
 229  }
 230  
 231  func (p *printer) funcType(t *FuncType) {
 232  	p.ch('(')
 233  	p.fieldList(t.ParamList, ", ")
 234  	p.ch(')')
 235  	if len(t.ResultList) == 1 && t.ResultList[0].Name == nil {
 236  		p.space()
 237  		p.expr(t.ResultList[0].Type)
 238  	} else if len(t.ResultList) > 0 {
 239  		p.w(" (")
 240  		p.fieldList(t.ResultList, ", ")
 241  		p.ch(')')
 242  	}
 243  }
 244  
 245  func opPrec(op Operator) int32 {
 246  	switch {
 247  	case op == OrOr:
 248  		return 1
 249  	case op == AndAnd:
 250  		return 2
 251  	case op >= Eql && op <= Geq:
 252  		return 3
 253  	case op >= Add && op <= Xor:
 254  		return 4
 255  	case op >= Mul && op <= Shr:
 256  		return 5
 257  	}
 258  	return 6
 259  }
 260  
 261  func (p *printer) expr(e Expr) {
 262  	p.exprPrec(e, 0)
 263  }
 264  
 265  func (p *printer) exprPrec(e Expr, outerPrec int32) {
 266  	if e == nil {
 267  		return
 268  	}
 269  	switch e := e.(type) {
 270  	case *BadExpr:
 271  		p.w("/*BadExpr*/")
 272  
 273  	case *Name:
 274  		p.w(e.Value)
 275  
 276  	case *BasicLit:
 277  		p.w(e.Value)
 278  
 279  	case *CompositeLit:
 280  		if e.Type != nil {
 281  			p.expr(e.Type)
 282  		}
 283  		p.ch('{')
 284  		for i, elem := range e.ElemList {
 285  			if i > 0 {
 286  				p.w(", ")
 287  			}
 288  			if kv, ok := elem.(*KeyValueExpr); ok {
 289  				if _, bad := kv.Key.(*BadExpr); bad {
 290  					p.ch(':')
 291  					if list, ok := kv.Value.(*ListExpr); ok && int32(len(list.ElemList)) == 2 {
 292  						p.expr(list.ElemList[0])
 293  						p.ch(':')
 294  						p.expr(list.ElemList[1])
 295  					} else {
 296  						p.expr(kv.Value)
 297  					}
 298  				} else {
 299  					p.expr(kv.Key)
 300  					p.w(": ")
 301  					p.expr(kv.Value)
 302  				}
 303  			} else {
 304  				p.expr(elem)
 305  			}
 306  		}
 307  		p.ch('}')
 308  
 309  	case *KeyValueExpr:
 310  		p.expr(e.Key)
 311  		p.w(": ")
 312  		p.expr(e.Value)
 313  
 314  	case *FuncLit:
 315  		p.w("func")
 316  		p.funcType(e.Type)
 317  		p.space()
 318  		p.blockStmt(e.Body)
 319  
 320  	case *ParenExpr:
 321  		p.ch('(')
 322  		p.expr(e.X)
 323  		p.ch(')')
 324  
 325  	case *SelectorExpr:
 326  		p.expr(e.X)
 327  		p.ch('.')
 328  		p.w(e.Sel.Value)
 329  
 330  	case *IndexExpr:
 331  		p.expr(e.X)
 332  		p.ch('[')
 333  		p.expr(e.Index)
 334  		p.ch(']')
 335  
 336  	case *SliceExpr:
 337  		p.expr(e.X)
 338  		p.ch('[')
 339  		if e.Index[0] != nil {
 340  			p.expr(e.Index[0])
 341  		}
 342  		p.ch(':')
 343  		if e.Index[1] != nil {
 344  			p.expr(e.Index[1])
 345  		}
 346  		if e.Full {
 347  			p.ch(':')
 348  			if e.Index[2] != nil {
 349  				p.expr(e.Index[2])
 350  			}
 351  		}
 352  		p.ch(']')
 353  
 354  	case *AssertExpr:
 355  		p.expr(e.X)
 356  		p.w(".(")
 357  		if e.Type != nil {
 358  			p.expr(e.Type)
 359  		} else {
 360  			p.w("type")
 361  		}
 362  		p.ch(')')
 363  
 364  	case *TypeSwitchGuard:
 365  		if e.Lhs != nil {
 366  			p.w(e.Lhs.Value)
 367  			p.w(" := ")
 368  		}
 369  		p.expr(e.X)
 370  		p.w(".(type)")
 371  
 372  	case *Operation:
 373  		if e.Y == nil {
 374  			prec := int32(6)
 375  			needParen := prec < outerPrec
 376  			if needParen {
 377  				p.ch('(')
 378  			}
 379  			p.w(e.Op.String())
 380  			p.exprPrec(e.X, prec)
 381  			if needParen {
 382  				p.ch(')')
 383  			}
 384  		} else {
 385  			prec := opPrec(e.Op)
 386  			needParen := prec < outerPrec
 387  			if needParen {
 388  				p.ch('(')
 389  			}
 390  			p.exprPrec(e.X, prec)
 391  			p.space()
 392  			p.w(e.Op.String())
 393  			p.space()
 394  			p.exprPrec(e.Y, prec+1)
 395  			if needParen {
 396  				p.ch(')')
 397  			}
 398  		}
 399  
 400  	case *CallExpr:
 401  		p.expr(e.Fun)
 402  		p.ch('(')
 403  		for i, arg := range e.ArgList {
 404  			if i > 0 {
 405  				p.w(", ")
 406  			}
 407  			p.expr(arg)
 408  		}
 409  		if e.HasDots {
 410  			p.w("...")
 411  		}
 412  		p.ch(')')
 413  
 414  	case *ListExpr:
 415  		for i, elem := range e.ElemList {
 416  			if i > 0 {
 417  				p.w(", ")
 418  			}
 419  			p.expr(elem)
 420  		}
 421  
 422  	case *ArrayType:
 423  		p.ch('[')
 424  		if e.Len != nil {
 425  			p.expr(e.Len)
 426  		}
 427  		p.ch(']')
 428  		p.expr(e.Elem)
 429  
 430  	case *SliceType:
 431  		p.w("[]")
 432  		p.expr(e.Elem)
 433  
 434  	case *DotsType:
 435  		p.w("...")
 436  		p.expr(e.Elem)
 437  
 438  	case *StructType:
 439  		p.w("struct {")
 440  		if len(e.FieldList) > 0 {
 441  			p.nl()
 442  			p.depth++
 443  			for i, f := range e.FieldList {
 444  				p.tab()
 445  				p.field(f)
 446  				if e.TagList != nil && i < int32(len(e.TagList)) && e.TagList[i] != nil {
 447  					p.space()
 448  					p.w(e.TagList[i].Value)
 449  				}
 450  				p.nl()
 451  			}
 452  			p.depth--
 453  			p.tab()
 454  		}
 455  		p.ch('}')
 456  
 457  	case *InterfaceType:
 458  		p.w("interface{")
 459  		if len(e.MethodList) > 0 {
 460  			p.nl()
 461  			p.depth++
 462  			for _, f := range e.MethodList {
 463  				p.tab()
 464  				p.field(f)
 465  				p.nl()
 466  			}
 467  			p.depth--
 468  			p.tab()
 469  		}
 470  		p.ch('}')
 471  
 472  	case *FuncType:
 473  		p.w("func")
 474  		p.funcType(e)
 475  
 476  	case *MapType:
 477  		p.w("map[")
 478  		p.expr(e.Key)
 479  		p.ch(']')
 480  		p.expr(e.Value)
 481  
 482  	case *ChanType:
 483  		switch e.Dir {
 484  		case SendOnly:
 485  			p.w("chan<- ")
 486  		case RecvOnly:
 487  			p.w("<-chan ")
 488  		default:
 489  			p.w("chan ")
 490  		}
 491  		p.expr(e.Elem)
 492  	}
 493  }
 494  
 495  func (p *printer) stmt(s Stmt) {
 496  	if s == nil {
 497  		return
 498  	}
 499  	switch s := s.(type) {
 500  	case *EmptyStmt:
 501  
 502  	case *LabeledStmt:
 503  		oldDepth := p.depth
 504  		p.depth = 0
 505  		p.tab()
 506  		p.w(s.Label.Value)
 507  		p.ch(':')
 508  		p.depth = oldDepth
 509  		p.nl()
 510  		p.tab()
 511  		p.stmt(s.Stmt)
 512  
 513  	case *BlockStmt:
 514  		p.blockStmt(s)
 515  
 516  	case *ExprStmt:
 517  		p.expr(s.X)
 518  
 519  	case *SendStmt:
 520  		p.expr(s.Chan)
 521  		p.w(" <- ")
 522  		p.expr(s.Value)
 523  
 524  	case *DeclStmt:
 525  		for i, d := range s.DeclList {
 526  			if i > 0 {
 527  				p.nl()
 528  				p.tab()
 529  			}
 530  			p.decl(d)
 531  		}
 532  
 533  	case *AssignStmt:
 534  		if s.Rhs == nil {
 535  			p.expr(s.Lhs)
 536  			if s.Op == Add {
 537  				p.w("++")
 538  			} else if s.Op == Sub {
 539  				p.w("--")
 540  			}
 541  		} else {
 542  			p.expr(s.Lhs)
 543  			if s.Op == 0 {
 544  				p.w(" = ")
 545  			} else if s.Op == Def {
 546  				p.w(" := ")
 547  			} else {
 548  				p.w(" ")
 549  				p.w(s.Op.String())
 550  				p.w("= ")
 551  			}
 552  			p.expr(s.Rhs)
 553  		}
 554  
 555  	case *BranchStmt:
 556  		switch s.Tok {
 557  		case Break:
 558  			p.w("break")
 559  		case Continue:
 560  			p.w("continue")
 561  		case Goto:
 562  			p.w("goto")
 563  		case Fallthrough:
 564  			p.w("fallthrough")
 565  		}
 566  		if s.Label != nil {
 567  			p.space()
 568  			p.w(s.Label.Value)
 569  		}
 570  
 571  	case *CallStmt:
 572  		if s.Tok == Go {
 573  			p.w("go ")
 574  		} else if s.Tok == Defer {
 575  			p.w("defer ")
 576  		}
 577  		p.expr(s.Call)
 578  
 579  	case *ReturnStmt:
 580  		p.w("return")
 581  		if s.Results != nil {
 582  			p.space()
 583  			p.expr(s.Results)
 584  		}
 585  
 586  	case *IfStmt:
 587  		p.w("if ")
 588  		if s.Init != nil {
 589  			p.stmt(s.Init)
 590  			p.w("; ")
 591  		}
 592  		p.expr(s.Cond)
 593  		p.space()
 594  		p.blockStmt(s.Then)
 595  		if s.Else != nil {
 596  			p.w(" else ")
 597  			p.stmt(s.Else)
 598  		}
 599  
 600  	case *ForStmt:
 601  		p.w("for ")
 602  		if rc, ok := s.Init.(*RangeClause); ok {
 603  			if rc.Lhs != nil {
 604  				p.expr(rc.Lhs)
 605  				if rc.Def {
 606  					p.w(" := ")
 607  				} else {
 608  					p.w(" = ")
 609  				}
 610  			}
 611  			p.w("range ")
 612  			p.expr(rc.X)
 613  		} else {
 614  			if s.Init != nil {
 615  				p.stmt(s.Init)
 616  			}
 617  			if s.Init != nil || s.Post != nil {
 618  				p.w("; ")
 619  			}
 620  			if s.Cond != nil {
 621  				p.expr(s.Cond)
 622  			}
 623  			if s.Init != nil || s.Post != nil {
 624  				p.w("; ")
 625  			}
 626  			if s.Post != nil {
 627  				p.stmt(s.Post)
 628  			}
 629  		}
 630  		p.space()
 631  		p.blockStmt(s.Body)
 632  
 633  	case *SwitchStmt:
 634  		p.w("switch ")
 635  		if s.Init != nil {
 636  			p.stmt(s.Init)
 637  			p.w("; ")
 638  		}
 639  		if s.Tag != nil {
 640  			p.expr(s.Tag)
 641  			p.space()
 642  		}
 643  		p.ch('{')
 644  		p.nl()
 645  		for _, cc := range s.Body {
 646  			p.tab()
 647  			if cc.Cases != nil {
 648  				p.w("case ")
 649  				p.expr(cc.Cases)
 650  			} else {
 651  				p.w("default")
 652  			}
 653  			p.ch(':')
 654  			p.nl()
 655  			p.depth++
 656  			for _, st := range cc.Body {
 657  				p.tab()
 658  				p.stmt(st)
 659  				p.nl()
 660  			}
 661  			p.depth--
 662  		}
 663  		p.tab()
 664  		p.ch('}')
 665  
 666  	case *SelectStmt:
 667  		p.w("select {")
 668  		p.nl()
 669  		for _, cc := range s.Body {
 670  			p.tab()
 671  			if cc.Comm != nil {
 672  				p.w("case ")
 673  				p.stmt(cc.Comm)
 674  			} else {
 675  				p.w("default")
 676  			}
 677  			p.ch(':')
 678  			p.nl()
 679  			p.depth++
 680  			for _, st := range cc.Body {
 681  				p.tab()
 682  				p.stmt(st)
 683  				p.nl()
 684  			}
 685  			p.depth--
 686  		}
 687  		p.tab()
 688  		p.ch('}')
 689  	}
 690  }
 691  
 692  func (p *printer) blockStmt(b *BlockStmt) {
 693  	p.ch('{')
 694  	p.nl()
 695  	p.depth++
 696  	for _, s := range b.List {
 697  		p.tab()
 698  		p.stmt(s)
 699  		p.nl()
 700  	}
 701  	p.depth--
 702  	p.tab()
 703  	p.ch('}')
 704  }
 705