github.com/xrash/gopher-lua@v0.0.0-20160304065408-e5faab4db06a/compile.go (about)

     1  package lua
     2  
     3  import (
     4  	"fmt"
     5  	"github.com/yuin/gopher-lua/ast"
     6  	"math"
     7  	"reflect"
     8  )
     9  
    10  /* internal constants & structs  {{{ */
    11  
    12  const maxRegisters = 200
    13  
    14  type expContextType int
    15  
    16  const (
    17  	ecGlobal expContextType = iota
    18  	ecUpvalue
    19  	ecLocal
    20  	ecTable
    21  	ecVararg
    22  	ecMethod
    23  	ecNone
    24  )
    25  
    26  const regNotDefined = opMaxArgsA + 1
    27  const labelNoJump = 0
    28  
    29  type expcontext struct {
    30  	ctype expContextType
    31  	reg   int
    32  	// varargopt >= 0: wants varargopt+1 results, i.e  a = func()
    33  	// varargopt = -1: ignore results             i.e  func()
    34  	// varargopt = -2: receive all results        i.e  a = {func()}
    35  	varargopt int
    36  }
    37  
    38  type assigncontext struct {
    39  	ec       *expcontext
    40  	keyrk    int
    41  	valuerk  int
    42  	keyks    bool
    43  	needmove bool
    44  }
    45  
    46  type lblabels struct {
    47  	t int
    48  	f int
    49  	e int
    50  	b bool
    51  }
    52  
    53  type constLValueExpr struct {
    54  	ast.ExprBase
    55  
    56  	Value LValue
    57  }
    58  
    59  // }}}
    60  
    61  /* utilities {{{ */
    62  var _ecnone0 = &expcontext{ecNone, regNotDefined, 0}
    63  var _ecnonem1 = &expcontext{ecNone, regNotDefined, -1}
    64  var _ecnonem2 = &expcontext{ecNone, regNotDefined, -2}
    65  var ecfuncdef = &expcontext{ecMethod, regNotDefined, 0}
    66  
    67  func ecupdate(ec *expcontext, ctype expContextType, reg, varargopt int) {
    68  	ec.ctype = ctype
    69  	ec.reg = reg
    70  	ec.varargopt = varargopt
    71  }
    72  
    73  func ecnone(varargopt int) *expcontext {
    74  	switch varargopt {
    75  	case 0:
    76  		return _ecnone0
    77  	case -1:
    78  		return _ecnonem1
    79  	case -2:
    80  		return _ecnonem2
    81  	}
    82  	return &expcontext{ecNone, regNotDefined, varargopt}
    83  }
    84  
    85  func sline(pos ast.PositionHolder) int {
    86  	return pos.Line()
    87  }
    88  
    89  func eline(pos ast.PositionHolder) int {
    90  	return pos.LastLine()
    91  }
    92  
    93  func savereg(ec *expcontext, reg int) int {
    94  	if ec.ctype != ecLocal || ec.reg == regNotDefined {
    95  		return reg
    96  	}
    97  	return ec.reg
    98  }
    99  
   100  func raiseCompileError(context *funcContext, line int, format string, args ...interface{}) {
   101  	msg := fmt.Sprintf(format, args...)
   102  	panic(&CompileError{Context: context, Line: line, Message: msg})
   103  }
   104  
   105  func isVarArgReturnExpr(expr ast.Expr) bool {
   106  	switch ex := expr.(type) {
   107  	case *ast.FuncCallExpr:
   108  		return !ex.AdjustRet
   109  	case *ast.Comma3Expr:
   110  		return true
   111  	}
   112  	return false
   113  }
   114  
   115  func lnumberValue(expr ast.Expr) (LNumber, bool) {
   116  	if ex, ok := expr.(*ast.NumberExpr); ok {
   117  		lv, err := parseNumber(ex.Value)
   118  		if err != nil {
   119  			lv = LNumber(math.NaN())
   120  		}
   121  		return lv, true
   122  	} else if ex, ok := expr.(*constLValueExpr); ok {
   123  		return ex.Value.(LNumber), true
   124  	}
   125  	return 0, false
   126  }
   127  
   128  /* utilities }}} */
   129  
   130  type CompileError struct { // {{{
   131  	Context *funcContext
   132  	Line    int
   133  	Message string
   134  }
   135  
   136  func (e *CompileError) Error() string {
   137  	return fmt.Sprintf("compile error near line(%v) %v: %v", e.Line, e.Context.Proto.SourceName, e.Message)
   138  } // }}}
   139  
   140  type codeStore struct { // {{{
   141  	codes []uint32
   142  	lines []int
   143  	pc    int
   144  }
   145  
   146  func (cd *codeStore) Add(inst uint32, line int) {
   147  	if l := len(cd.codes); l <= 0 || cd.pc == l {
   148  		cd.codes = append(cd.codes, inst)
   149  		cd.lines = append(cd.lines, line)
   150  	} else {
   151  		cd.codes[cd.pc] = inst
   152  		cd.lines[cd.pc] = line
   153  	}
   154  	cd.pc++
   155  }
   156  
   157  func (cd *codeStore) AddABC(op int, a int, b int, c int, line int) {
   158  	cd.Add(opCreateABC(op, a, b, c), line)
   159  }
   160  
   161  func (cd *codeStore) AddABx(op int, a int, bx int, line int) {
   162  	cd.Add(opCreateABx(op, a, bx), line)
   163  }
   164  
   165  func (cd *codeStore) AddASbx(op int, a int, sbx int, line int) {
   166  	cd.Add(opCreateASbx(op, a, sbx), line)
   167  }
   168  
   169  func (cd *codeStore) PropagateKMV(top int, save *int, reg *int, inc int) {
   170  	lastinst := cd.Last()
   171  	if opGetArgA(lastinst) >= top {
   172  		switch opGetOpCode(lastinst) {
   173  		case OP_LOADK:
   174  			cindex := opGetArgBx(lastinst)
   175  			if cindex <= opMaxIndexRk {
   176  				cd.Pop()
   177  				*save = opRkAsk(cindex)
   178  				return
   179  			}
   180  		case OP_MOVE:
   181  			cd.Pop()
   182  			*save = opGetArgB(lastinst)
   183  			return
   184  		}
   185  	}
   186  	*save = *reg
   187  	*reg = *reg + inc
   188  }
   189  
   190  func (cd *codeStore) PropagateMV(top int, save *int, reg *int, inc int) {
   191  	lastinst := cd.Last()
   192  	if opGetArgA(lastinst) >= top {
   193  		switch opGetOpCode(lastinst) {
   194  		case OP_MOVE:
   195  			cd.Pop()
   196  			*save = opGetArgB(lastinst)
   197  			return
   198  		}
   199  	}
   200  	*save = *reg
   201  	*reg = *reg + inc
   202  }
   203  
   204  func (cd *codeStore) SetOpCode(pc int, v int) {
   205  	opSetOpCode(&cd.codes[pc], v)
   206  }
   207  
   208  func (cd *codeStore) SetA(pc int, v int) {
   209  	opSetArgA(&cd.codes[pc], v)
   210  }
   211  
   212  func (cd *codeStore) SetB(pc int, v int) {
   213  	opSetArgB(&cd.codes[pc], v)
   214  }
   215  
   216  func (cd *codeStore) SetC(pc int, v int) {
   217  	opSetArgC(&cd.codes[pc], v)
   218  }
   219  
   220  func (cd *codeStore) SetBx(pc int, v int) {
   221  	opSetArgBx(&cd.codes[pc], v)
   222  }
   223  
   224  func (cd *codeStore) SetSbx(pc int, v int) {
   225  	opSetArgSbx(&cd.codes[pc], v)
   226  }
   227  
   228  func (cd *codeStore) At(pc int) uint32 {
   229  	return cd.codes[pc]
   230  }
   231  
   232  func (cd *codeStore) List() []uint32 {
   233  	return cd.codes[:cd.pc]
   234  }
   235  
   236  func (cd *codeStore) PosList() []int {
   237  	return cd.lines[:cd.pc]
   238  }
   239  
   240  func (cd *codeStore) LastPC() int {
   241  	return cd.pc - 1
   242  }
   243  
   244  func (cd *codeStore) Last() uint32 {
   245  	if cd.pc == 0 {
   246  		return opInvalidInstruction
   247  	}
   248  	return cd.codes[cd.pc-1]
   249  }
   250  
   251  func (cd *codeStore) Pop() {
   252  	cd.pc--
   253  } /* }}} Code */
   254  
   255  /* {{{ VarNamePool */
   256  
   257  type varNamePoolValue struct {
   258  	Index int
   259  	Name  string
   260  }
   261  
   262  type varNamePool struct {
   263  	names  []string
   264  	offset int
   265  }
   266  
   267  func newVarNamePool(offset int) *varNamePool {
   268  	return &varNamePool{make([]string, 0, 16), offset}
   269  }
   270  
   271  func (vp *varNamePool) Names() []string {
   272  	return vp.names
   273  }
   274  
   275  func (vp *varNamePool) List() []varNamePoolValue {
   276  	result := make([]varNamePoolValue, len(vp.names), len(vp.names))
   277  	for i, name := range vp.names {
   278  		result[i].Index = i + vp.offset
   279  		result[i].Name = name
   280  	}
   281  	return result
   282  }
   283  
   284  func (vp *varNamePool) LastIndex() int {
   285  	return vp.offset + len(vp.names)
   286  }
   287  
   288  func (vp *varNamePool) Find(name string) int {
   289  	for i := len(vp.names) - 1; i >= 0; i-- {
   290  		if vp.names[i] == name {
   291  			return i + vp.offset
   292  		}
   293  	}
   294  	return -1
   295  }
   296  
   297  func (vp *varNamePool) RegisterUnique(name string) int {
   298  	index := vp.Find(name)
   299  	if index < 0 {
   300  		return vp.Register(name)
   301  	}
   302  	return index
   303  }
   304  
   305  func (vp *varNamePool) Register(name string) int {
   306  	vp.names = append(vp.names, name)
   307  	return len(vp.names) - 1 + vp.offset
   308  }
   309  
   310  /* }}} VarNamePool */
   311  
   312  /* FuncContext {{{ */
   313  
   314  type codeBlock struct {
   315  	LocalVars  *varNamePool
   316  	BreakLabel int
   317  	Parent     *codeBlock
   318  	RefUpvalue bool
   319  	LineStart  int
   320  	LastLine   int
   321  }
   322  
   323  func newCodeBlock(localvars *varNamePool, blabel int, parent *codeBlock, pos ast.PositionHolder) *codeBlock {
   324  	bl := &codeBlock{localvars, blabel, parent, false, 0, 0}
   325  	if pos != nil {
   326  		bl.LineStart = pos.Line()
   327  		bl.LastLine = pos.LastLine()
   328  	}
   329  	return bl
   330  }
   331  
   332  type funcContext struct {
   333  	Proto    *FunctionProto
   334  	Code     *codeStore
   335  	Parent   *funcContext
   336  	Upvalues *varNamePool
   337  	Block    *codeBlock
   338  	Blocks   []*codeBlock
   339  	regTop   int
   340  	labelId  int
   341  	labelPc  map[int]int
   342  }
   343  
   344  func newFuncContext(sourcename string, parent *funcContext) *funcContext {
   345  	fc := &funcContext{
   346  		Proto:    newFunctionProto(sourcename),
   347  		Code:     &codeStore{make([]uint32, 0, 1024), make([]int, 0, 1024), 0},
   348  		Parent:   parent,
   349  		Upvalues: newVarNamePool(0),
   350  		Block:    newCodeBlock(newVarNamePool(0), labelNoJump, nil, nil),
   351  		regTop:   0,
   352  		labelId:  1,
   353  		labelPc:  map[int]int{},
   354  	}
   355  	fc.Blocks = []*codeBlock{fc.Block}
   356  	return fc
   357  }
   358  
   359  func (fc *funcContext) NewLabel() int {
   360  	ret := fc.labelId
   361  	fc.labelId++
   362  	return ret
   363  }
   364  
   365  func (fc *funcContext) SetLabelPc(label int, pc int) {
   366  	fc.labelPc[label] = pc
   367  }
   368  
   369  func (fc *funcContext) GetLabelPc(label int) int {
   370  	return fc.labelPc[label]
   371  }
   372  
   373  func (fc *funcContext) ConstIndex(value LValue) int {
   374  	ctype := value.Type()
   375  	for i, lv := range fc.Proto.Constants {
   376  		if lv.Type() == ctype && lv == value {
   377  			return i
   378  		}
   379  	}
   380  	fc.Proto.Constants = append(fc.Proto.Constants, value)
   381  	v := len(fc.Proto.Constants) - 1
   382  	if v > opMaxArgBx {
   383  		raiseCompileError(fc, fc.Proto.LineDefined, "too many constants")
   384  	}
   385  	return v
   386  }
   387  
   388  func (fc *funcContext) RegisterLocalVar(name string) int {
   389  	ret := fc.Block.LocalVars.Register(name)
   390  	fc.Proto.DbgLocals = append(fc.Proto.DbgLocals, &DbgLocalInfo{Name: name, StartPc: fc.Code.LastPC() + 1})
   391  	fc.SetRegTop(fc.RegTop() + 1)
   392  	return ret
   393  }
   394  
   395  func (fc *funcContext) FindLocalVarAndBlock(name string) (int, *codeBlock) {
   396  	for block := fc.Block; block != nil; block = block.Parent {
   397  		if index := block.LocalVars.Find(name); index > -1 {
   398  			return index, block
   399  		}
   400  	}
   401  	return -1, nil
   402  }
   403  
   404  func (fc *funcContext) FindLocalVar(name string) int {
   405  	idx, _ := fc.FindLocalVarAndBlock(name)
   406  	return idx
   407  }
   408  
   409  func (fc *funcContext) LocalVars() []varNamePoolValue {
   410  	result := make([]varNamePoolValue, 0, 32)
   411  	for _, block := range fc.Blocks {
   412  		result = append(result, block.LocalVars.List()...)
   413  	}
   414  	return result
   415  }
   416  
   417  func (fc *funcContext) EnterBlock(blabel int, pos ast.PositionHolder) {
   418  	fc.Block = newCodeBlock(newVarNamePool(fc.RegTop()), blabel, fc.Block, pos)
   419  	fc.Blocks = append(fc.Blocks, fc.Block)
   420  }
   421  
   422  func (fc *funcContext) CloseUpvalues() int {
   423  	n := -1
   424  	if fc.Block.RefUpvalue {
   425  		n = fc.Block.Parent.LocalVars.LastIndex()
   426  		fc.Code.AddABC(OP_CLOSE, n, 0, 0, fc.Block.LastLine)
   427  	}
   428  	return n
   429  }
   430  
   431  func (fc *funcContext) LeaveBlock() int {
   432  	closed := fc.CloseUpvalues()
   433  	fc.EndScope()
   434  	fc.Block = fc.Block.Parent
   435  	fc.SetRegTop(fc.Block.LocalVars.LastIndex())
   436  	return closed
   437  }
   438  
   439  func (fc *funcContext) EndScope() {
   440  	for _, vr := range fc.Block.LocalVars.List() {
   441  		fc.Proto.DbgLocals[vr.Index].EndPc = fc.Code.LastPC()
   442  	}
   443  }
   444  
   445  func (fc *funcContext) SetRegTop(top int) {
   446  	if top > maxRegisters {
   447  		raiseCompileError(fc, fc.Proto.LineDefined, "too many local variables")
   448  	}
   449  	fc.regTop = top
   450  }
   451  
   452  func (fc *funcContext) RegTop() int {
   453  	return fc.regTop
   454  }
   455  
   456  /* FuncContext }}} */
   457  
   458  func compileChunk(context *funcContext, chunk []ast.Stmt) { // {{{
   459  	for _, stmt := range chunk {
   460  		compileStmt(context, stmt)
   461  	}
   462  } // }}}
   463  
   464  func compileBlock(context *funcContext, chunk []ast.Stmt) { // {{{
   465  	if len(chunk) == 0 {
   466  		return
   467  	}
   468  	ph := &ast.Node{}
   469  	ph.SetLine(sline(chunk[0]))
   470  	ph.SetLastLine(eline(chunk[len(chunk)-1]))
   471  	context.EnterBlock(labelNoJump, ph)
   472  	for _, stmt := range chunk {
   473  		compileStmt(context, stmt)
   474  	}
   475  	context.LeaveBlock()
   476  } // }}}
   477  
   478  func compileStmt(context *funcContext, stmt ast.Stmt) { // {{{
   479  	switch st := stmt.(type) {
   480  	case *ast.AssignStmt:
   481  		compileAssignStmt(context, st)
   482  	case *ast.LocalAssignStmt:
   483  		compileLocalAssignStmt(context, st)
   484  	case *ast.FuncCallStmt:
   485  		compileFuncCallExpr(context, context.RegTop(), st.Expr.(*ast.FuncCallExpr), ecnone(-1))
   486  	case *ast.DoBlockStmt:
   487  		context.EnterBlock(labelNoJump, st)
   488  		compileChunk(context, st.Stmts)
   489  		context.LeaveBlock()
   490  	case *ast.WhileStmt:
   491  		compileWhileStmt(context, st)
   492  	case *ast.RepeatStmt:
   493  		compileRepeatStmt(context, st)
   494  	case *ast.FuncDefStmt:
   495  		compileFuncDefStmt(context, st)
   496  	case *ast.ReturnStmt:
   497  		compileReturnStmt(context, st)
   498  	case *ast.IfStmt:
   499  		compileIfStmt(context, st)
   500  	case *ast.BreakStmt:
   501  		compileBreakStmt(context, st)
   502  	case *ast.NumberForStmt:
   503  		compileNumberForStmt(context, st)
   504  	case *ast.GenericForStmt:
   505  		compileGenericForStmt(context, st)
   506  	}
   507  } // }}}
   508  
   509  func compileAssignStmtLeft(context *funcContext, stmt *ast.AssignStmt) (int, []*assigncontext) { // {{{
   510  	reg := context.RegTop()
   511  	acs := make([]*assigncontext, 0, len(stmt.Lhs))
   512  	for i, lhs := range stmt.Lhs {
   513  		islast := i == len(stmt.Lhs)-1
   514  		switch st := lhs.(type) {
   515  		case *ast.IdentExpr:
   516  			identtype := getIdentRefType(context, context, st)
   517  			ec := &expcontext{identtype, regNotDefined, 0}
   518  			switch identtype {
   519  			case ecGlobal:
   520  				context.ConstIndex(LString(st.Value))
   521  			case ecUpvalue:
   522  				context.Upvalues.RegisterUnique(st.Value)
   523  			case ecLocal:
   524  				if islast {
   525  					ec.reg = context.FindLocalVar(st.Value)
   526  				}
   527  			}
   528  			acs = append(acs, &assigncontext{ec, 0, 0, false, false})
   529  		case *ast.AttrGetExpr:
   530  			ac := &assigncontext{&expcontext{ecTable, regNotDefined, 0}, 0, 0, false, false}
   531  			compileExprWithKMVPropagation(context, st.Object, &reg, &ac.ec.reg)
   532  			compileExprWithKMVPropagation(context, st.Key, &reg, &ac.keyrk)
   533  			if _, ok := st.Key.(*ast.StringExpr); ok {
   534  				ac.keyks = true
   535  			}
   536  			acs = append(acs, ac)
   537  
   538  		default:
   539  			panic("invalid left expression.")
   540  		}
   541  	}
   542  	return reg, acs
   543  } // }}}
   544  
   545  func compileAssignStmtRight(context *funcContext, stmt *ast.AssignStmt, reg int, acs []*assigncontext) (int, []*assigncontext) { // {{{
   546  	lennames := len(stmt.Lhs)
   547  	lenexprs := len(stmt.Rhs)
   548  	namesassigned := 0
   549  
   550  	for namesassigned < lennames {
   551  		ac := acs[namesassigned]
   552  		ec := ac.ec
   553  		var expr ast.Expr = nil
   554  		if namesassigned >= lenexprs {
   555  			expr = &ast.NilExpr{}
   556  			expr.SetLine(sline(stmt.Lhs[namesassigned]))
   557  			expr.SetLastLine(eline(stmt.Lhs[namesassigned]))
   558  		} else if isVarArgReturnExpr(stmt.Rhs[namesassigned]) && (lenexprs-namesassigned-1) <= 0 {
   559  			varargopt := lennames - namesassigned - 1
   560  			regstart := reg
   561  			reginc := compileExpr(context, reg, stmt.Rhs[namesassigned], ecnone(varargopt))
   562  			reg += reginc
   563  			for i := namesassigned; i < namesassigned+int(reginc); i++ {
   564  				acs[i].needmove = true
   565  				if acs[i].ec.ctype == ecTable {
   566  					acs[i].valuerk = regstart + (i - namesassigned)
   567  				}
   568  			}
   569  			namesassigned = lennames
   570  			continue
   571  		}
   572  
   573  		if expr == nil {
   574  			expr = stmt.Rhs[namesassigned]
   575  		}
   576  		idx := reg
   577  		reginc := compileExpr(context, reg, expr, ec)
   578  		if ec.ctype == ecTable {
   579  			if _, ok := expr.(*ast.LogicalOpExpr); !ok {
   580  				context.Code.PropagateKMV(context.RegTop(), &ac.valuerk, &reg, reginc)
   581  			} else {
   582  				ac.valuerk = idx
   583  				reg += reginc
   584  			}
   585  		} else {
   586  			ac.needmove = reginc != 0
   587  			reg += reginc
   588  		}
   589  		namesassigned += 1
   590  	}
   591  
   592  	rightreg := reg - 1
   593  
   594  	// extra right exprs
   595  	for i := namesassigned; i < lenexprs; i++ {
   596  		varargopt := -1
   597  		if i != lenexprs-1 {
   598  			varargopt = 0
   599  		}
   600  		reg += compileExpr(context, reg, stmt.Rhs[i], ecnone(varargopt))
   601  	}
   602  	return rightreg, acs
   603  } // }}}
   604  
   605  func compileAssignStmt(context *funcContext, stmt *ast.AssignStmt) { // {{{
   606  	code := context.Code
   607  	lennames := len(stmt.Lhs)
   608  	reg, acs := compileAssignStmtLeft(context, stmt)
   609  	reg, acs = compileAssignStmtRight(context, stmt, reg, acs)
   610  
   611  	for i := lennames - 1; i >= 0; i-- {
   612  		ex := stmt.Lhs[i]
   613  		switch acs[i].ec.ctype {
   614  		case ecLocal:
   615  			if acs[i].needmove {
   616  				code.AddABC(OP_MOVE, context.FindLocalVar(ex.(*ast.IdentExpr).Value), reg, 0, sline(ex))
   617  				reg -= 1
   618  			}
   619  		case ecGlobal:
   620  			code.AddABx(OP_SETGLOBAL, reg, context.ConstIndex(LString(ex.(*ast.IdentExpr).Value)), sline(ex))
   621  			reg -= 1
   622  		case ecUpvalue:
   623  			code.AddABC(OP_SETUPVAL, reg, context.Upvalues.RegisterUnique(ex.(*ast.IdentExpr).Value), 0, sline(ex))
   624  			reg -= 1
   625  		case ecTable:
   626  			opcode := OP_SETTABLE
   627  			if acs[i].keyks {
   628  				opcode = OP_SETTABLEKS
   629  			}
   630  			code.AddABC(opcode, acs[i].ec.reg, acs[i].keyrk, acs[i].valuerk, sline(ex))
   631  			if !opIsK(acs[i].valuerk) {
   632  				reg -= 1
   633  			}
   634  		}
   635  	}
   636  } // }}}
   637  
   638  func compileRegAssignment(context *funcContext, names []string, exprs []ast.Expr, reg int, nvars int, line int) { // {{{
   639  	lennames := len(names)
   640  	lenexprs := len(exprs)
   641  	namesassigned := 0
   642  	ec := &expcontext{}
   643  
   644  	for namesassigned < lennames && namesassigned < lenexprs {
   645  		if isVarArgReturnExpr(exprs[namesassigned]) && (lenexprs-namesassigned-1) <= 0 {
   646  
   647  			varargopt := nvars - namesassigned
   648  			ecupdate(ec, ecVararg, reg, varargopt-1)
   649  			compileExpr(context, reg, exprs[namesassigned], ec)
   650  			reg += varargopt
   651  			namesassigned = lennames
   652  		} else {
   653  			ecupdate(ec, ecLocal, reg, 0)
   654  			compileExpr(context, reg, exprs[namesassigned], ec)
   655  			reg += 1
   656  			namesassigned += 1
   657  		}
   658  	}
   659  
   660  	// extra left names
   661  	if lennames > namesassigned {
   662  		restleft := lennames - namesassigned - 1
   663  		context.Code.AddABC(OP_LOADNIL, reg, reg+restleft, 0, line)
   664  		reg += restleft
   665  	}
   666  
   667  	// extra right exprs
   668  	for i := namesassigned; i < lenexprs; i++ {
   669  		varargopt := -1
   670  		if i != lenexprs-1 {
   671  			varargopt = 0
   672  		}
   673  		ecupdate(ec, ecNone, reg, varargopt)
   674  		reg += compileExpr(context, reg, exprs[i], ec)
   675  	}
   676  } // }}}
   677  
   678  func compileLocalAssignStmt(context *funcContext, stmt *ast.LocalAssignStmt) { // {{{
   679  	reg := context.RegTop()
   680  	if len(stmt.Names) == 1 && len(stmt.Exprs) == 1 {
   681  		if _, ok := stmt.Exprs[0].(*ast.FunctionExpr); ok {
   682  			context.RegisterLocalVar(stmt.Names[0])
   683  			compileRegAssignment(context, stmt.Names, stmt.Exprs, reg, len(stmt.Names), sline(stmt))
   684  			return
   685  		}
   686  	}
   687  
   688  	compileRegAssignment(context, stmt.Names, stmt.Exprs, reg, len(stmt.Names), sline(stmt))
   689  	for _, name := range stmt.Names {
   690  		context.RegisterLocalVar(name)
   691  	}
   692  } // }}}
   693  
   694  func compileReturnStmt(context *funcContext, stmt *ast.ReturnStmt) { // {{{
   695  	lenexprs := len(stmt.Exprs)
   696  	code := context.Code
   697  	reg := context.RegTop()
   698  	a := reg
   699  	lastisvaarg := false
   700  
   701  	if lenexprs == 1 {
   702  		switch ex := stmt.Exprs[0].(type) {
   703  		case *ast.IdentExpr:
   704  			if idx := context.FindLocalVar(ex.Value); idx > -1 {
   705  				code.AddABC(OP_RETURN, idx, 2, 0, sline(stmt))
   706  				return
   707  			}
   708  		case *ast.FuncCallExpr:
   709  			reg += compileExpr(context, reg, ex, ecnone(-2))
   710  			code.SetOpCode(code.LastPC(), OP_TAILCALL)
   711  			code.AddABC(OP_RETURN, a, 0, 0, sline(stmt))
   712  		}
   713  	}
   714  
   715  	for i, expr := range stmt.Exprs {
   716  		if i == lenexprs-1 && isVarArgReturnExpr(expr) {
   717  			compileExpr(context, reg, expr, ecnone(-2))
   718  			lastisvaarg = true
   719  		} else {
   720  			reg += compileExpr(context, reg, expr, ecnone(0))
   721  		}
   722  	}
   723  	count := reg - a + 1
   724  	if lastisvaarg {
   725  		count = 0
   726  	}
   727  	context.Code.AddABC(OP_RETURN, a, count, 0, sline(stmt))
   728  } // }}}
   729  
   730  func compileIfStmt(context *funcContext, stmt *ast.IfStmt) { // {{{
   731  	thenlabel := context.NewLabel()
   732  	elselabel := context.NewLabel()
   733  	endlabel := context.NewLabel()
   734  
   735  	compileBranchCondition(context, context.RegTop(), stmt.Condition, thenlabel, elselabel, false)
   736  	context.SetLabelPc(thenlabel, context.Code.LastPC())
   737  	compileBlock(context, stmt.Then)
   738  	if len(stmt.Else) > 0 {
   739  		context.Code.AddASbx(OP_JMP, 0, endlabel, sline(stmt))
   740  	}
   741  	context.SetLabelPc(elselabel, context.Code.LastPC())
   742  	if len(stmt.Else) > 0 {
   743  		compileBlock(context, stmt.Else)
   744  		context.SetLabelPc(endlabel, context.Code.LastPC())
   745  	}
   746  
   747  } // }}}
   748  
   749  func compileBranchCondition(context *funcContext, reg int, expr ast.Expr, thenlabel, elselabel int, hasnextcond bool) { // {{{
   750  	// TODO folding constants?
   751  	code := context.Code
   752  	flip := 0
   753  	jumplabel := elselabel
   754  	if hasnextcond {
   755  		flip = 1
   756  		jumplabel = thenlabel
   757  	}
   758  
   759  	switch ex := expr.(type) {
   760  	case *ast.FalseExpr, *ast.NilExpr:
   761  		if !hasnextcond {
   762  			code.AddASbx(OP_JMP, 0, elselabel, sline(expr))
   763  			return
   764  		}
   765  	case *ast.TrueExpr, *ast.NumberExpr, *ast.StringExpr:
   766  		if !hasnextcond {
   767  			return
   768  		}
   769  	case *ast.UnaryNotOpExpr:
   770  		compileBranchCondition(context, reg, ex.Expr, elselabel, thenlabel, !hasnextcond)
   771  		return
   772  	case *ast.LogicalOpExpr:
   773  		switch ex.Operator {
   774  		case "and":
   775  			nextcondlabel := context.NewLabel()
   776  			compileBranchCondition(context, reg, ex.Lhs, nextcondlabel, elselabel, false)
   777  			context.SetLabelPc(nextcondlabel, context.Code.LastPC())
   778  			compileBranchCondition(context, reg, ex.Rhs, thenlabel, elselabel, hasnextcond)
   779  		case "or":
   780  			nextcondlabel := context.NewLabel()
   781  			compileBranchCondition(context, reg, ex.Lhs, thenlabel, nextcondlabel, true)
   782  			context.SetLabelPc(nextcondlabel, context.Code.LastPC())
   783  			compileBranchCondition(context, reg, ex.Rhs, thenlabel, elselabel, hasnextcond)
   784  		}
   785  		return
   786  	case *ast.RelationalOpExpr:
   787  		compileRelationalOpExprAux(context, reg, ex, flip, jumplabel)
   788  		return
   789  	}
   790  
   791  	a := reg
   792  	compileExprWithMVPropagation(context, expr, &reg, &a)
   793  	code.AddABC(OP_TEST, a, 0, 0^flip, sline(expr))
   794  	code.AddASbx(OP_JMP, 0, jumplabel, sline(expr))
   795  } // }}}
   796  
   797  func compileWhileStmt(context *funcContext, stmt *ast.WhileStmt) { // {{{
   798  	thenlabel := context.NewLabel()
   799  	elselabel := context.NewLabel()
   800  	condlabel := context.NewLabel()
   801  
   802  	context.SetLabelPc(condlabel, context.Code.LastPC())
   803  	compileBranchCondition(context, context.RegTop(), stmt.Condition, thenlabel, elselabel, false)
   804  	context.SetLabelPc(thenlabel, context.Code.LastPC())
   805  	context.EnterBlock(elselabel, stmt)
   806  	compileChunk(context, stmt.Stmts)
   807  	context.CloseUpvalues()
   808  	context.Code.AddASbx(OP_JMP, 0, condlabel, eline(stmt))
   809  	context.LeaveBlock()
   810  	context.SetLabelPc(elselabel, context.Code.LastPC())
   811  } // }}}
   812  
   813  func compileRepeatStmt(context *funcContext, stmt *ast.RepeatStmt) { // {{{
   814  	initlabel := context.NewLabel()
   815  	thenlabel := context.NewLabel()
   816  	elselabel := context.NewLabel()
   817  
   818  	context.SetLabelPc(initlabel, context.Code.LastPC())
   819  	context.SetLabelPc(elselabel, context.Code.LastPC())
   820  	context.EnterBlock(thenlabel, stmt)
   821  	compileChunk(context, stmt.Stmts)
   822  	compileBranchCondition(context, context.RegTop(), stmt.Condition, thenlabel, elselabel, false)
   823  
   824  	context.SetLabelPc(thenlabel, context.Code.LastPC())
   825  	n := context.LeaveBlock()
   826  
   827  	if n > -1 {
   828  		label := context.NewLabel()
   829  		context.Code.AddASbx(OP_JMP, 0, label, eline(stmt))
   830  		context.SetLabelPc(elselabel, context.Code.LastPC())
   831  		context.Code.AddABC(OP_CLOSE, n, 0, 0, eline(stmt))
   832  		context.Code.AddASbx(OP_JMP, 0, initlabel, eline(stmt))
   833  		context.SetLabelPc(label, context.Code.LastPC())
   834  	}
   835  
   836  } // }}}
   837  
   838  func compileBreakStmt(context *funcContext, stmt *ast.BreakStmt) { // {{{
   839  	for block := context.Block; block != nil; block = block.Parent {
   840  		if label := block.BreakLabel; label != labelNoJump {
   841  			if block.RefUpvalue {
   842  				context.Code.AddABC(OP_CLOSE, block.Parent.LocalVars.LastIndex(), 0, 0, sline(stmt))
   843  			}
   844  			context.Code.AddASbx(OP_JMP, 0, label, sline(stmt))
   845  			return
   846  		}
   847  	}
   848  	raiseCompileError(context, sline(stmt), "no loop to break")
   849  } // }}}
   850  
   851  func compileFuncDefStmt(context *funcContext, stmt *ast.FuncDefStmt) { // {{{
   852  	if stmt.Name.Func == nil {
   853  		reg := context.RegTop()
   854  		var treg, kreg int
   855  		compileExprWithKMVPropagation(context, stmt.Name.Receiver, &reg, &treg)
   856  		kreg = loadRk(context, &reg, stmt.Func, LString(stmt.Name.Method))
   857  		compileExpr(context, reg, stmt.Func, ecfuncdef)
   858  		context.Code.AddABC(OP_SETTABLE, treg, kreg, reg, sline(stmt.Name.Receiver))
   859  	} else {
   860  		astmt := &ast.AssignStmt{Lhs: []ast.Expr{stmt.Name.Func}, Rhs: []ast.Expr{stmt.Func}}
   861  		astmt.SetLine(sline(stmt.Func))
   862  		astmt.SetLastLine(eline(stmt.Func))
   863  		compileAssignStmt(context, astmt)
   864  	}
   865  } // }}}
   866  
   867  func compileNumberForStmt(context *funcContext, stmt *ast.NumberForStmt) { // {{{
   868  	code := context.Code
   869  	endlabel := context.NewLabel()
   870  	ec := &expcontext{}
   871  
   872  	context.EnterBlock(endlabel, stmt)
   873  	reg := context.RegTop()
   874  	rindex := context.RegisterLocalVar("(for index)")
   875  	ecupdate(ec, ecLocal, rindex, 0)
   876  	compileExpr(context, reg, stmt.Init, ec)
   877  
   878  	reg = context.RegTop()
   879  	rlimit := context.RegisterLocalVar("(for limit)")
   880  	ecupdate(ec, ecLocal, rlimit, 0)
   881  	compileExpr(context, reg, stmt.Limit, ec)
   882  
   883  	reg = context.RegTop()
   884  	rstep := context.RegisterLocalVar("(for step)")
   885  	if stmt.Step == nil {
   886  		stmt.Step = &ast.NumberExpr{Value: "1"}
   887  		stmt.Step.SetLine(sline(stmt.Init))
   888  	}
   889  	ecupdate(ec, ecLocal, rstep, 0)
   890  	compileExpr(context, reg, stmt.Step, ec)
   891  
   892  	code.AddASbx(OP_FORPREP, rindex, 0, sline(stmt))
   893  
   894  	context.RegisterLocalVar(stmt.Name)
   895  
   896  	bodypc := code.LastPC()
   897  	compileChunk(context, stmt.Stmts)
   898  
   899  	context.LeaveBlock()
   900  
   901  	flpc := code.LastPC()
   902  	code.AddASbx(OP_FORLOOP, rindex, bodypc-(flpc+1), sline(stmt))
   903  
   904  	context.SetLabelPc(endlabel, code.LastPC())
   905  	code.SetSbx(bodypc, flpc-bodypc)
   906  
   907  } // }}}
   908  
   909  func compileGenericForStmt(context *funcContext, stmt *ast.GenericForStmt) { // {{{
   910  	code := context.Code
   911  	endlabel := context.NewLabel()
   912  	bodylabel := context.NewLabel()
   913  	fllabel := context.NewLabel()
   914  	nnames := len(stmt.Names)
   915  
   916  	context.EnterBlock(endlabel, stmt)
   917  	rgen := context.RegisterLocalVar("(for generator)")
   918  	context.RegisterLocalVar("(for state)")
   919  	context.RegisterLocalVar("(for control)")
   920  
   921  	compileRegAssignment(context, stmt.Names, stmt.Exprs, context.RegTop()-3, 3, sline(stmt))
   922  
   923  	code.AddASbx(OP_JMP, 0, fllabel, sline(stmt))
   924  
   925  	for _, name := range stmt.Names {
   926  		context.RegisterLocalVar(name)
   927  	}
   928  
   929  	context.SetLabelPc(bodylabel, code.LastPC())
   930  	compileChunk(context, stmt.Stmts)
   931  
   932  	context.LeaveBlock()
   933  
   934  	context.SetLabelPc(fllabel, code.LastPC())
   935  	code.AddABC(OP_TFORLOOP, rgen, 0, nnames, sline(stmt))
   936  	code.AddASbx(OP_JMP, 0, bodylabel, sline(stmt))
   937  
   938  	context.SetLabelPc(endlabel, code.LastPC())
   939  } // }}}
   940  
   941  func compileExpr(context *funcContext, reg int, expr ast.Expr, ec *expcontext) int { // {{{
   942  	code := context.Code
   943  	sreg := savereg(ec, reg)
   944  	sused := 1
   945  	if sreg < reg {
   946  		sused = 0
   947  	}
   948  
   949  	switch ex := expr.(type) {
   950  	case *ast.StringExpr:
   951  		code.AddABx(OP_LOADK, sreg, context.ConstIndex(LString(ex.Value)), sline(ex))
   952  		return sused
   953  	case *ast.NumberExpr:
   954  		num, err := parseNumber(ex.Value)
   955  		if err != nil {
   956  			num = LNumber(math.NaN())
   957  		}
   958  		code.AddABx(OP_LOADK, sreg, context.ConstIndex(num), sline(ex))
   959  		return sused
   960  	case *constLValueExpr:
   961  		code.AddABx(OP_LOADK, sreg, context.ConstIndex(ex.Value), sline(ex))
   962  		return sused
   963  	case *ast.NilExpr:
   964  		code.AddABC(OP_LOADNIL, sreg, sreg, 0, sline(ex))
   965  		return sused
   966  	case *ast.FalseExpr:
   967  		code.AddABC(OP_LOADBOOL, sreg, 0, 0, sline(ex))
   968  		return sused
   969  	case *ast.TrueExpr:
   970  		code.AddABC(OP_LOADBOOL, sreg, 1, 0, sline(ex))
   971  		return sused
   972  	case *ast.IdentExpr:
   973  		switch getIdentRefType(context, context, ex) {
   974  		case ecGlobal:
   975  			code.AddABx(OP_GETGLOBAL, sreg, context.ConstIndex(LString(ex.Value)), sline(ex))
   976  		case ecUpvalue:
   977  			code.AddABC(OP_GETUPVAL, sreg, context.Upvalues.RegisterUnique(ex.Value), 0, sline(ex))
   978  		case ecLocal:
   979  			b := context.FindLocalVar(ex.Value)
   980  			code.AddABC(OP_MOVE, sreg, b, 0, sline(ex))
   981  		}
   982  		return sused
   983  	case *ast.Comma3Expr:
   984  		if context.Proto.IsVarArg == 0 {
   985  			raiseCompileError(context, sline(ex), "cannot use '...' outside a vararg function")
   986  		}
   987  		context.Proto.IsVarArg &= ^VarArgNeedsArg
   988  		code.AddABC(OP_VARARG, sreg, 2+ec.varargopt, 0, sline(ex))
   989  		if context.RegTop() > (sreg+2+ec.varargopt) || ec.varargopt < -1 {
   990  			return 0
   991  		}
   992  		return (sreg + 1 + ec.varargopt) - reg
   993  	case *ast.AttrGetExpr:
   994  		a := sreg
   995  		b := reg
   996  		compileExprWithMVPropagation(context, ex.Object, &reg, &b)
   997  		c := reg
   998  		compileExprWithKMVPropagation(context, ex.Key, &reg, &c)
   999  		opcode := OP_GETTABLE
  1000  		if _, ok := ex.Key.(*ast.StringExpr); ok {
  1001  			opcode = OP_GETTABLEKS
  1002  		}
  1003  		code.AddABC(opcode, a, b, c, sline(ex))
  1004  		return sused
  1005  	case *ast.TableExpr:
  1006  		compileTableExpr(context, reg, ex, ec)
  1007  		return 1
  1008  	case *ast.ArithmeticOpExpr:
  1009  		compileArithmeticOpExpr(context, reg, ex, ec)
  1010  		return sused
  1011  	case *ast.StringConcatOpExpr:
  1012  		compileStringConcatOpExpr(context, reg, ex, ec)
  1013  		return sused
  1014  	case *ast.UnaryMinusOpExpr, *ast.UnaryNotOpExpr, *ast.UnaryLenOpExpr:
  1015  		compileUnaryOpExpr(context, reg, ex, ec)
  1016  		return sused
  1017  	case *ast.RelationalOpExpr:
  1018  		compileRelationalOpExpr(context, reg, ex, ec)
  1019  		return sused
  1020  	case *ast.LogicalOpExpr:
  1021  		compileLogicalOpExpr(context, reg, ex, ec)
  1022  		return sused
  1023  	case *ast.FuncCallExpr:
  1024  		return compileFuncCallExpr(context, reg, ex, ec)
  1025  	case *ast.FunctionExpr:
  1026  		childcontext := newFuncContext(context.Proto.SourceName, context)
  1027  		compileFunctionExpr(childcontext, ex, ec)
  1028  		protono := len(context.Proto.FunctionPrototypes)
  1029  		context.Proto.FunctionPrototypes = append(context.Proto.FunctionPrototypes, childcontext.Proto)
  1030  		code.AddABx(OP_CLOSURE, sreg, protono, sline(ex))
  1031  		for _, upvalue := range childcontext.Upvalues.List() {
  1032  			localidx, block := context.FindLocalVarAndBlock(upvalue.Name)
  1033  			if localidx > -1 {
  1034  				code.AddABC(OP_MOVE, 0, localidx, 0, sline(ex))
  1035  				block.RefUpvalue = true
  1036  			} else {
  1037  				upvalueidx := context.Upvalues.Find(upvalue.Name)
  1038  				if upvalueidx < 0 {
  1039  					upvalueidx = context.Upvalues.RegisterUnique(upvalue.Name)
  1040  				}
  1041  				code.AddABC(OP_GETUPVAL, 0, upvalueidx, 0, sline(ex))
  1042  			}
  1043  		}
  1044  		return sused
  1045  	default:
  1046  		panic(fmt.Sprintf("expr %v not implemented.", reflect.TypeOf(ex).Elem().Name()))
  1047  	}
  1048  
  1049  	panic("should not reach here")
  1050  	return sused
  1051  } // }}}
  1052  
  1053  func compileExprWithPropagation(context *funcContext, expr ast.Expr, reg *int, save *int, propergator func(int, *int, *int, int)) { // {{{
  1054  	reginc := compileExpr(context, *reg, expr, ecnone(0))
  1055  	if _, ok := expr.(*ast.LogicalOpExpr); ok {
  1056  		*save = *reg
  1057  		*reg = *reg + reginc
  1058  	} else {
  1059  		propergator(context.RegTop(), save, reg, reginc)
  1060  	}
  1061  } // }}}
  1062  
  1063  func compileExprWithKMVPropagation(context *funcContext, expr ast.Expr, reg *int, save *int) { // {{{
  1064  	compileExprWithPropagation(context, expr, reg, save, context.Code.PropagateKMV)
  1065  } // }}}
  1066  
  1067  func compileExprWithMVPropagation(context *funcContext, expr ast.Expr, reg *int, save *int) { // {{{
  1068  	compileExprWithPropagation(context, expr, reg, save, context.Code.PropagateMV)
  1069  } // }}}
  1070  
  1071  func constFold(exp ast.Expr) ast.Expr { // {{{
  1072  	switch expr := exp.(type) {
  1073  	case *ast.ArithmeticOpExpr:
  1074  		lvalue, lisconst := lnumberValue(expr.Lhs)
  1075  		rvalue, risconst := lnumberValue(expr.Rhs)
  1076  		if lisconst && risconst {
  1077  			switch expr.Operator {
  1078  			case "+":
  1079  				return &constLValueExpr{Value: lvalue + rvalue}
  1080  			case "-":
  1081  				return &constLValueExpr{Value: lvalue - rvalue}
  1082  			case "*":
  1083  				return &constLValueExpr{Value: lvalue * rvalue}
  1084  			case "/":
  1085  				return &constLValueExpr{Value: lvalue / rvalue}
  1086  			case "%":
  1087  				return &constLValueExpr{Value: luaModulo(lvalue, rvalue)}
  1088  			case "^":
  1089  				return &constLValueExpr{Value: LNumber(math.Pow(float64(lvalue), float64(rvalue)))}
  1090  			default:
  1091  				panic(fmt.Sprintf("unknwon binop: %v", expr.Operator))
  1092  			}
  1093  		} else {
  1094  			retexpr := *expr
  1095  			retexpr.Lhs = constFold(expr.Lhs)
  1096  			retexpr.Rhs = constFold(expr.Rhs)
  1097  			return &retexpr
  1098  		}
  1099  	case *ast.UnaryMinusOpExpr:
  1100  		expr.Expr = constFold(expr.Expr)
  1101  		if value, ok := lnumberValue(expr.Expr); ok {
  1102  			return &constLValueExpr{Value: LNumber(-value)}
  1103  		}
  1104  		return expr
  1105  	default:
  1106  
  1107  		return exp
  1108  	}
  1109  	return exp
  1110  } // }}}
  1111  
  1112  func compileFunctionExpr(context *funcContext, funcexpr *ast.FunctionExpr, ec *expcontext) { // {{{
  1113  	context.Proto.LineDefined = sline(funcexpr)
  1114  	context.Proto.LastLineDefined = eline(funcexpr)
  1115  	if len(funcexpr.ParList.Names) > maxRegisters {
  1116  		raiseCompileError(context, context.Proto.LineDefined, "register overflow")
  1117  	}
  1118  	context.Proto.NumParameters = uint8(len(funcexpr.ParList.Names))
  1119  	if ec.ctype == ecMethod {
  1120  		context.Proto.NumParameters += 1
  1121  		context.RegisterLocalVar("self")
  1122  	}
  1123  	for _, name := range funcexpr.ParList.Names {
  1124  		context.RegisterLocalVar(name)
  1125  	}
  1126  	if funcexpr.ParList.HasVargs {
  1127  		if CompatVarArg {
  1128  			context.Proto.IsVarArg = VarArgHasArg | VarArgNeedsArg
  1129  			if context.Parent != nil {
  1130  				context.RegisterLocalVar("arg")
  1131  			}
  1132  		}
  1133  		context.Proto.IsVarArg |= VarArgIsVarArg
  1134  	}
  1135  
  1136  	compileChunk(context, funcexpr.Stmts)
  1137  
  1138  	context.Code.AddABC(OP_RETURN, 0, 1, 0, eline(funcexpr))
  1139  	context.EndScope()
  1140  	context.Proto.Code = context.Code.List()
  1141  	context.Proto.DbgSourcePositions = context.Code.PosList()
  1142  	context.Proto.DbgUpvalues = context.Upvalues.Names()
  1143  	context.Proto.NumUpvalues = uint8(len(context.Proto.DbgUpvalues))
  1144  	for _, clv := range context.Proto.Constants {
  1145  		sv := ""
  1146  		if slv, ok := clv.(LString); ok {
  1147  			sv = string(slv)
  1148  		}
  1149  		context.Proto.stringConstants = append(context.Proto.stringConstants, sv)
  1150  	}
  1151  	patchCode(context)
  1152  } // }}}
  1153  
  1154  func compileTableExpr(context *funcContext, reg int, ex *ast.TableExpr, ec *expcontext) { // {{{
  1155  	code := context.Code
  1156  	/*
  1157  		tablereg := savereg(ec, reg)
  1158  		if tablereg == reg {
  1159  			reg += 1
  1160  		}
  1161  	*/
  1162  	tablereg := reg
  1163  	reg++
  1164  	code.AddABC(OP_NEWTABLE, tablereg, 0, 0, sline(ex))
  1165  	tablepc := code.LastPC()
  1166  	regbase := reg
  1167  
  1168  	arraycount := 0
  1169  	lastvararg := false
  1170  	for i, field := range ex.Fields {
  1171  		islast := i == len(ex.Fields)-1
  1172  		if field.Key == nil {
  1173  			if islast && isVarArgReturnExpr(field.Value) {
  1174  				reg += compileExpr(context, reg, field.Value, ecnone(-2))
  1175  				lastvararg = true
  1176  			} else {
  1177  				reg += compileExpr(context, reg, field.Value, ecnone(0))
  1178  				arraycount += 1
  1179  			}
  1180  		} else {
  1181  			regorg := reg
  1182  			b := reg
  1183  			compileExprWithKMVPropagation(context, field.Key, &reg, &b)
  1184  			c := reg
  1185  			compileExprWithKMVPropagation(context, field.Value, &reg, &c)
  1186  			opcode := OP_SETTABLE
  1187  			if _, ok := field.Key.(*ast.StringExpr); ok {
  1188  				opcode = OP_SETTABLEKS
  1189  			}
  1190  			code.AddABC(opcode, tablereg, b, c, sline(ex))
  1191  			reg = regorg
  1192  		}
  1193  		flush := arraycount % FieldsPerFlush
  1194  		if (arraycount != 0 && (flush == 0 || islast)) || lastvararg {
  1195  			reg = regbase
  1196  			num := flush
  1197  			if num == 0 {
  1198  				num = FieldsPerFlush
  1199  			}
  1200  			c := (arraycount-1)/FieldsPerFlush + 1
  1201  			b := num
  1202  			if islast && isVarArgReturnExpr(field.Value) {
  1203  				b = 0
  1204  			}
  1205  			line := field.Value
  1206  			if field.Key != nil {
  1207  				line = field.Key
  1208  			}
  1209  			if c > 511 {
  1210  				c = 0
  1211  			}
  1212  			code.AddABC(OP_SETLIST, tablereg, b, c, sline(line))
  1213  			if c == 0 {
  1214  				code.Add(uint32(c), sline(line))
  1215  			}
  1216  		}
  1217  	}
  1218  	code.SetB(tablepc, int2Fb(arraycount))
  1219  	code.SetC(tablepc, int2Fb(len(ex.Fields)-arraycount))
  1220  	if ec.ctype == ecLocal && ec.reg != tablereg {
  1221  		code.AddABC(OP_MOVE, ec.reg, tablereg, 0, sline(ex))
  1222  	}
  1223  } // }}}
  1224  
  1225  func compileArithmeticOpExpr(context *funcContext, reg int, expr *ast.ArithmeticOpExpr, ec *expcontext) { // {{{
  1226  	exp := constFold(expr)
  1227  	if ex, ok := exp.(*constLValueExpr); ok {
  1228  		exp.SetLine(sline(expr))
  1229  		compileExpr(context, reg, ex, ec)
  1230  		return
  1231  	}
  1232  	expr, _ = exp.(*ast.ArithmeticOpExpr)
  1233  	a := savereg(ec, reg)
  1234  	b := reg
  1235  	compileExprWithKMVPropagation(context, expr.Lhs, &reg, &b)
  1236  	c := reg
  1237  	compileExprWithKMVPropagation(context, expr.Rhs, &reg, &c)
  1238  
  1239  	op := 0
  1240  	switch expr.Operator {
  1241  	case "+":
  1242  		op = OP_ADD
  1243  	case "-":
  1244  		op = OP_SUB
  1245  	case "*":
  1246  		op = OP_MUL
  1247  	case "/":
  1248  		op = OP_DIV
  1249  	case "%":
  1250  		op = OP_MOD
  1251  	case "^":
  1252  		op = OP_POW
  1253  	}
  1254  	context.Code.AddABC(op, a, b, c, sline(expr))
  1255  } // }}}
  1256  
  1257  func compileStringConcatOpExpr(context *funcContext, reg int, expr *ast.StringConcatOpExpr, ec *expcontext) { // {{{
  1258  	code := context.Code
  1259  	crange := 1
  1260  	for current := expr.Rhs; current != nil; {
  1261  		if ex, ok := current.(*ast.StringConcatOpExpr); ok {
  1262  			crange += 1
  1263  			current = ex.Rhs
  1264  		} else {
  1265  			current = nil
  1266  		}
  1267  	}
  1268  	a := savereg(ec, reg)
  1269  	basereg := reg
  1270  	reg += compileExpr(context, reg, expr.Lhs, ecnone(0))
  1271  	reg += compileExpr(context, reg, expr.Rhs, ecnone(0))
  1272  	for pc := code.LastPC(); pc != 0 && opGetOpCode(code.At(pc)) == OP_CONCAT; pc-- {
  1273  		code.Pop()
  1274  	}
  1275  	code.AddABC(OP_CONCAT, a, basereg, basereg+crange, sline(expr))
  1276  } // }}}
  1277  
  1278  func compileUnaryOpExpr(context *funcContext, reg int, expr ast.Expr, ec *expcontext) { // {{{
  1279  	opcode := 0
  1280  	code := context.Code
  1281  	var operandexpr ast.Expr
  1282  	switch ex := expr.(type) {
  1283  	case *ast.UnaryMinusOpExpr:
  1284  		exp := constFold(ex)
  1285  		if lvexpr, ok := exp.(*constLValueExpr); ok {
  1286  			exp.SetLine(sline(expr))
  1287  			compileExpr(context, reg, lvexpr, ec)
  1288  			return
  1289  		}
  1290  		ex, _ = exp.(*ast.UnaryMinusOpExpr)
  1291  		operandexpr = ex.Expr
  1292  		opcode = OP_UNM
  1293  	case *ast.UnaryNotOpExpr:
  1294  		switch ex.Expr.(type) {
  1295  		case *ast.TrueExpr:
  1296  			code.AddABC(OP_LOADBOOL, savereg(ec, reg), 0, 0, sline(expr))
  1297  			return
  1298  		case *ast.FalseExpr, *ast.NilExpr:
  1299  			code.AddABC(OP_LOADBOOL, savereg(ec, reg), 1, 0, sline(expr))
  1300  			return
  1301  		default:
  1302  			opcode = OP_NOT
  1303  			operandexpr = ex.Expr
  1304  		}
  1305  	case *ast.UnaryLenOpExpr:
  1306  		opcode = OP_LEN
  1307  		operandexpr = ex.Expr
  1308  	}
  1309  
  1310  	a := savereg(ec, reg)
  1311  	b := reg
  1312  	compileExprWithMVPropagation(context, operandexpr, &reg, &b)
  1313  	code.AddABC(opcode, a, b, 0, sline(expr))
  1314  } // }}}
  1315  
  1316  func compileRelationalOpExprAux(context *funcContext, reg int, expr *ast.RelationalOpExpr, flip int, label int) { // {{{
  1317  	code := context.Code
  1318  	b := reg
  1319  	compileExprWithKMVPropagation(context, expr.Lhs, &reg, &b)
  1320  	c := reg
  1321  	compileExprWithKMVPropagation(context, expr.Rhs, &reg, &c)
  1322  	switch expr.Operator {
  1323  	case "<":
  1324  		code.AddABC(OP_LT, 0^flip, b, c, sline(expr))
  1325  	case ">":
  1326  		code.AddABC(OP_LT, 0^flip, c, b, sline(expr))
  1327  	case "<=":
  1328  		code.AddABC(OP_LE, 0^flip, b, c, sline(expr))
  1329  	case ">=":
  1330  		code.AddABC(OP_LE, 0^flip, c, b, sline(expr))
  1331  	case "==":
  1332  		code.AddABC(OP_EQ, 0^flip, b, c, sline(expr))
  1333  	case "~=":
  1334  		code.AddABC(OP_EQ, 1^flip, b, c, sline(expr))
  1335  	}
  1336  	code.AddASbx(OP_JMP, 0, label, sline(expr))
  1337  } // }}}
  1338  
  1339  func compileRelationalOpExpr(context *funcContext, reg int, expr *ast.RelationalOpExpr, ec *expcontext) { // {{{
  1340  	a := savereg(ec, reg)
  1341  	code := context.Code
  1342  	jumplabel := context.NewLabel()
  1343  	compileRelationalOpExprAux(context, reg, expr, 1, jumplabel)
  1344  	code.AddABC(OP_LOADBOOL, a, 0, 1, sline(expr))
  1345  	context.SetLabelPc(jumplabel, code.LastPC())
  1346  	code.AddABC(OP_LOADBOOL, a, 1, 0, sline(expr))
  1347  } // }}}
  1348  
  1349  func compileLogicalOpExpr(context *funcContext, reg int, expr *ast.LogicalOpExpr, ec *expcontext) { // {{{
  1350  	a := savereg(ec, reg)
  1351  	code := context.Code
  1352  	endlabel := context.NewLabel()
  1353  	lb := &lblabels{context.NewLabel(), context.NewLabel(), endlabel, false}
  1354  	nextcondlabel := context.NewLabel()
  1355  	if expr.Operator == "and" {
  1356  		compileLogicalOpExprAux(context, reg, expr.Lhs, ec, nextcondlabel, endlabel, false, lb)
  1357  		context.SetLabelPc(nextcondlabel, code.LastPC())
  1358  		compileLogicalOpExprAux(context, reg, expr.Rhs, ec, endlabel, endlabel, false, lb)
  1359  	} else {
  1360  		compileLogicalOpExprAux(context, reg, expr.Lhs, ec, endlabel, nextcondlabel, true, lb)
  1361  		context.SetLabelPc(nextcondlabel, code.LastPC())
  1362  		compileLogicalOpExprAux(context, reg, expr.Rhs, ec, endlabel, endlabel, false, lb)
  1363  	}
  1364  
  1365  	if lb.b {
  1366  		context.SetLabelPc(lb.f, code.LastPC())
  1367  		code.AddABC(OP_LOADBOOL, a, 0, 1, sline(expr))
  1368  		context.SetLabelPc(lb.t, code.LastPC())
  1369  		code.AddABC(OP_LOADBOOL, a, 1, 0, sline(expr))
  1370  	}
  1371  
  1372  	lastinst := code.Last()
  1373  	if opGetOpCode(lastinst) == OP_JMP && opGetArgSbx(lastinst) == endlabel {
  1374  		code.Pop()
  1375  	}
  1376  
  1377  	context.SetLabelPc(endlabel, code.LastPC())
  1378  } // }}}
  1379  
  1380  func compileLogicalOpExprAux(context *funcContext, reg int, expr ast.Expr, ec *expcontext, thenlabel, elselabel int, hasnextcond bool, lb *lblabels) { // {{{
  1381  	// TODO folding constants?
  1382  	code := context.Code
  1383  	flip := 0
  1384  	jumplabel := elselabel
  1385  	if hasnextcond {
  1386  		flip = 1
  1387  		jumplabel = thenlabel
  1388  	}
  1389  
  1390  	switch ex := expr.(type) {
  1391  	case *ast.FalseExpr:
  1392  		if elselabel == lb.e {
  1393  			code.AddASbx(OP_JMP, 0, lb.f, sline(expr))
  1394  			lb.b = true
  1395  		} else {
  1396  			code.AddASbx(OP_JMP, 0, elselabel, sline(expr))
  1397  		}
  1398  		return
  1399  	case *ast.NilExpr:
  1400  		if elselabel == lb.e {
  1401  			compileExpr(context, reg, expr, ec)
  1402  			code.AddASbx(OP_JMP, 0, lb.e, sline(expr))
  1403  		} else {
  1404  			code.AddASbx(OP_JMP, 0, elselabel, sline(expr))
  1405  		}
  1406  		return
  1407  	case *ast.TrueExpr:
  1408  		if thenlabel == lb.e {
  1409  			code.AddASbx(OP_JMP, 0, lb.t, sline(expr))
  1410  			lb.b = true
  1411  		} else {
  1412  			code.AddASbx(OP_JMP, 0, thenlabel, sline(expr))
  1413  		}
  1414  		return
  1415  	case *ast.NumberExpr, *ast.StringExpr:
  1416  		if thenlabel == lb.e {
  1417  			compileExpr(context, reg, expr, ec)
  1418  			code.AddASbx(OP_JMP, 0, lb.e, sline(expr))
  1419  		} else {
  1420  			code.AddASbx(OP_JMP, 0, thenlabel, sline(expr))
  1421  		}
  1422  		return
  1423  	case *ast.LogicalOpExpr:
  1424  		switch ex.Operator {
  1425  		case "and":
  1426  			nextcondlabel := context.NewLabel()
  1427  			compileLogicalOpExprAux(context, reg, ex.Lhs, ec, nextcondlabel, elselabel, false, lb)
  1428  			context.SetLabelPc(nextcondlabel, context.Code.LastPC())
  1429  			compileLogicalOpExprAux(context, reg, ex.Rhs, ec, thenlabel, elselabel, hasnextcond, lb)
  1430  		case "or":
  1431  			nextcondlabel := context.NewLabel()
  1432  			compileLogicalOpExprAux(context, reg, ex.Lhs, ec, thenlabel, nextcondlabel, true, lb)
  1433  			context.SetLabelPc(nextcondlabel, context.Code.LastPC())
  1434  			compileLogicalOpExprAux(context, reg, ex.Rhs, ec, thenlabel, elselabel, hasnextcond, lb)
  1435  		}
  1436  		return
  1437  	case *ast.RelationalOpExpr:
  1438  		if thenlabel == elselabel {
  1439  			flip ^= 1
  1440  			jumplabel = lb.t
  1441  			lb.b = true
  1442  		} else if thenlabel == lb.e {
  1443  			jumplabel = lb.t
  1444  			lb.b = true
  1445  		} else if elselabel == lb.e {
  1446  			jumplabel = lb.f
  1447  			lb.b = true
  1448  		}
  1449  		compileRelationalOpExprAux(context, reg, ex, flip, jumplabel)
  1450  		return
  1451  	}
  1452  
  1453  	if !hasnextcond && thenlabel == elselabel {
  1454  		reg += compileExpr(context, reg, expr, ec)
  1455  	} else {
  1456  		a := reg
  1457  		sreg := savereg(ec, a)
  1458  		reg += compileExpr(context, reg, expr, ecnone(0))
  1459  		if sreg == a {
  1460  			code.AddABC(OP_TEST, a, 0, 0^flip, sline(expr))
  1461  		} else {
  1462  			code.AddABC(OP_TESTSET, sreg, a, 0^flip, sline(expr))
  1463  		}
  1464  	}
  1465  	code.AddASbx(OP_JMP, 0, jumplabel, sline(expr))
  1466  } // }}}
  1467  
  1468  func compileFuncCallExpr(context *funcContext, reg int, expr *ast.FuncCallExpr, ec *expcontext) int { // {{{
  1469  	funcreg := reg
  1470  	if ec.ctype == ecLocal && ec.reg == (int(context.Proto.NumParameters)-1) {
  1471  		funcreg = ec.reg
  1472  		reg = ec.reg
  1473  	}
  1474  	argc := len(expr.Args)
  1475  	islastvararg := false
  1476  	name := "(anonymous)"
  1477  
  1478  	if expr.Func != nil { // hoge.func()
  1479  		reg += compileExpr(context, reg, expr.Func, ecnone(0))
  1480  		name = getExprName(context, expr.Func)
  1481  	} else { // hoge:method()
  1482  		b := reg
  1483  		compileExprWithMVPropagation(context, expr.Receiver, &reg, &b)
  1484  		c := loadRk(context, &reg, expr, LString(expr.Method))
  1485  		context.Code.AddABC(OP_SELF, funcreg, b, c, sline(expr))
  1486  		// increments a register for an implicit "self"
  1487  		reg = b + 1
  1488  		reg2 := funcreg + 2
  1489  		if reg2 > reg {
  1490  			reg = reg2
  1491  		}
  1492  		argc += 1
  1493  		name = string(expr.Method)
  1494  	}
  1495  
  1496  	for i, ar := range expr.Args {
  1497  		islastvararg = (i == len(expr.Args)-1) && isVarArgReturnExpr(ar)
  1498  		if islastvararg {
  1499  			compileExpr(context, reg, ar, ecnone(-2))
  1500  		} else {
  1501  			reg += compileExpr(context, reg, ar, ecnone(0))
  1502  		}
  1503  	}
  1504  	b := argc + 1
  1505  	if islastvararg {
  1506  		b = 0
  1507  	}
  1508  	context.Code.AddABC(OP_CALL, funcreg, b, ec.varargopt+2, sline(expr))
  1509  	context.Proto.DbgCalls = append(context.Proto.DbgCalls, DbgCall{Pc: context.Code.LastPC(), Name: name})
  1510  
  1511  	if ec.varargopt == 0 && ec.ctype == ecLocal && funcreg != ec.reg {
  1512  		context.Code.AddABC(OP_MOVE, ec.reg, funcreg, 0, sline(expr))
  1513  		return 1
  1514  	}
  1515  	if context.RegTop() > (funcreg+2+ec.varargopt) || ec.varargopt < -1 {
  1516  		return 0
  1517  	}
  1518  	return ec.varargopt + 1
  1519  } // }}}
  1520  
  1521  func loadRk(context *funcContext, reg *int, expr ast.Expr, cnst LValue) int { // {{{
  1522  	cindex := context.ConstIndex(cnst)
  1523  	if cindex <= opMaxIndexRk {
  1524  		return opRkAsk(cindex)
  1525  	} else {
  1526  		ret := *reg
  1527  		*reg++
  1528  		context.Code.AddABx(OP_LOADK, ret, cindex, sline(expr))
  1529  		return ret
  1530  	}
  1531  } // }}}
  1532  
  1533  func getIdentRefType(context *funcContext, current *funcContext, expr *ast.IdentExpr) expContextType { // {{{
  1534  	if current == nil {
  1535  		return ecGlobal
  1536  	} else if current.FindLocalVar(expr.Value) > -1 {
  1537  		if current == context {
  1538  			return ecLocal
  1539  		}
  1540  		return ecUpvalue
  1541  	}
  1542  	return getIdentRefType(context, current.Parent, expr)
  1543  } // }}}
  1544  
  1545  func getExprName(context *funcContext, expr ast.Expr) string { // {{{
  1546  	switch ex := expr.(type) {
  1547  	case *ast.IdentExpr:
  1548  		return ex.Value
  1549  	case *ast.AttrGetExpr:
  1550  		switch kex := ex.Key.(type) {
  1551  		case *ast.StringExpr:
  1552  			return kex.Value
  1553  		}
  1554  		return "?"
  1555  	}
  1556  	return "?"
  1557  } // }}}
  1558  
  1559  func patchCode(context *funcContext) { // {{{
  1560  	maxreg := 1
  1561  	if np := int(context.Proto.NumParameters); np > 1 {
  1562  		maxreg = np
  1563  	}
  1564  	moven := 0
  1565  	code := context.Code.List()
  1566  	for pc := 0; pc < len(code); pc++ {
  1567  		inst := code[pc]
  1568  		curop := opGetOpCode(inst)
  1569  		switch curop {
  1570  		case OP_CLOSURE:
  1571  			pc += int(context.Proto.FunctionPrototypes[opGetArgBx(inst)].NumUpvalues)
  1572  			moven = 0
  1573  			continue
  1574  		case OP_SETGLOBAL, OP_SETUPVAL, OP_EQ, OP_LT, OP_LE, OP_TEST,
  1575  			OP_TAILCALL, OP_RETURN, OP_FORPREP, OP_FORLOOP, OP_TFORLOOP,
  1576  			OP_SETLIST, OP_CLOSE:
  1577  			/* nothing to do */
  1578  		case OP_CALL:
  1579  			if reg := opGetArgA(inst) + opGetArgC(inst) - 2; reg > maxreg {
  1580  				maxreg = reg
  1581  			}
  1582  		case OP_VARARG:
  1583  			if reg := opGetArgA(inst) + opGetArgB(inst) - 1; reg > maxreg {
  1584  				maxreg = reg
  1585  			}
  1586  		case OP_SELF:
  1587  			if reg := opGetArgA(inst) + 1; reg > maxreg {
  1588  				maxreg = reg
  1589  			}
  1590  		case OP_LOADNIL:
  1591  			if reg := opGetArgB(inst); reg > maxreg {
  1592  				maxreg = reg
  1593  			}
  1594  		case OP_JMP: // jump to jump optimization
  1595  			distance := 0
  1596  			count := 0 // avoiding infinite loops
  1597  			for jmp := inst; opGetOpCode(jmp) == OP_JMP && count < 5; jmp = context.Code.At(pc + distance + 1) {
  1598  				d := context.GetLabelPc(opGetArgSbx(jmp)) - pc
  1599  				if d > opMaxArgSbx {
  1600  					if distance == 0 {
  1601  						raiseCompileError(context, context.Proto.LineDefined, "too long to jump.")
  1602  					}
  1603  					break
  1604  				}
  1605  				distance = d
  1606  				count++
  1607  			}
  1608  			if distance == 0 {
  1609  				context.Code.SetOpCode(pc, OP_NOP)
  1610  			} else {
  1611  				context.Code.SetSbx(pc, distance)
  1612  			}
  1613  		default:
  1614  			if reg := opGetArgA(inst); reg > maxreg {
  1615  				maxreg = reg
  1616  			}
  1617  		}
  1618  
  1619  		// bulk move optimization(reducing op dipatch costs)
  1620  		if curop == OP_MOVE {
  1621  			moven++
  1622  		} else {
  1623  			if moven > 1 {
  1624  				context.Code.SetOpCode(pc-moven, OP_MOVEN)
  1625  				context.Code.SetC(pc-moven, intMin(moven-1, opMaxArgsC))
  1626  			}
  1627  			moven = 0
  1628  		}
  1629  	}
  1630  	maxreg++
  1631  	if maxreg > maxRegisters {
  1632  		raiseCompileError(context, context.Proto.LineDefined, "register overflow(too many local variables)")
  1633  	}
  1634  	context.Proto.NumUsedRegisters = uint8(maxreg)
  1635  } // }}}
  1636  
  1637  func Compile(chunk []ast.Stmt, name string) (proto *FunctionProto, err error) { // {{{
  1638  	defer func() {
  1639  		if rcv := recover(); rcv != nil {
  1640  			if _, ok := rcv.(*CompileError); ok {
  1641  				err = rcv.(error)
  1642  			} else {
  1643  				panic(rcv)
  1644  			}
  1645  		}
  1646  	}()
  1647  	err = nil
  1648  	parlist := &ast.ParList{HasVargs: true, Names: []string{}}
  1649  	funcexpr := &ast.FunctionExpr{ParList: parlist, Stmts: chunk}
  1650  	context := newFuncContext(name, nil)
  1651  	compileFunctionExpr(context, funcexpr, ecnone(0))
  1652  	proto = context.Proto
  1653  	return
  1654  } // }}}