1package lua
2
3import (
4	"fmt"
5	"github.com/yuin/gopher-lua/ast"
6	"math"
7	"reflect"
8)
9
10/* internal constants & structs  {{{ */
11
12const maxRegisters = 200
13
14type expContextType int
15
16const (
17	ecGlobal expContextType = iota
18	ecUpvalue
19	ecLocal
20	ecTable
21	ecVararg
22	ecMethod
23	ecNone
24)
25
26const regNotDefined = opMaxArgsA + 1
27const labelNoJump = 0
28
29type 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
38type assigncontext struct {
39	ec       *expcontext
40	keyrk    int
41	valuerk  int
42	keyks    bool
43	needmove bool
44}
45
46type lblabels struct {
47	t int
48	f int
49	e int
50	b bool
51}
52
53type constLValueExpr struct {
54	ast.ExprBase
55
56	Value LValue
57}
58
59// }}}
60
61/* utilities {{{ */
62var _ecnone0 = &expcontext{ecNone, regNotDefined, 0}
63var _ecnonem1 = &expcontext{ecNone, regNotDefined, -1}
64var _ecnonem2 = &expcontext{ecNone, regNotDefined, -2}
65var ecfuncdef = &expcontext{ecMethod, regNotDefined, 0}
66
67func ecupdate(ec *expcontext, ctype expContextType, reg, varargopt int) {
68	ec.ctype = ctype
69	ec.reg = reg
70	ec.varargopt = varargopt
71}
72
73func 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
85func sline(pos ast.PositionHolder) int {
86	return pos.Line()
87}
88
89func eline(pos ast.PositionHolder) int {
90	return pos.LastLine()
91}
92
93func savereg(ec *expcontext, reg int) int {
94	if ec.ctype != ecLocal || ec.reg == regNotDefined {
95		return reg
96	}
97	return ec.reg
98}
99
100func 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
105func 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
115func 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
130type CompileError struct { // {{{
131	context *funcContext
132	Line    int
133	Message string
134}
135
136func (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
140type codeStore struct { // {{{
141	codes []uint32
142	lines []int
143	pc    int
144}
145
146func (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
157func (cd *codeStore) AddABC(op int, a int, b int, c int, line int) {
158	cd.Add(opCreateABC(op, a, b, c), line)
159}
160
161func (cd *codeStore) AddABx(op int, a int, bx int, line int) {
162	cd.Add(opCreateABx(op, a, bx), line)
163}
164
165func (cd *codeStore) AddASbx(op int, a int, sbx int, line int) {
166	cd.Add(opCreateASbx(op, a, sbx), line)
167}
168
169func (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
190func (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
204func (cd *codeStore) SetOpCode(pc int, v int) {
205	opSetOpCode(&cd.codes[pc], v)
206}
207
208func (cd *codeStore) SetA(pc int, v int) {
209	opSetArgA(&cd.codes[pc], v)
210}
211
212func (cd *codeStore) SetB(pc int, v int) {
213	opSetArgB(&cd.codes[pc], v)
214}
215
216func (cd *codeStore) SetC(pc int, v int) {
217	opSetArgC(&cd.codes[pc], v)
218}
219
220func (cd *codeStore) SetBx(pc int, v int) {
221	opSetArgBx(&cd.codes[pc], v)
222}
223
224func (cd *codeStore) SetSbx(pc int, v int) {
225	opSetArgSbx(&cd.codes[pc], v)
226}
227
228func (cd *codeStore) At(pc int) uint32 {
229	return cd.codes[pc]
230}
231
232func (cd *codeStore) List() []uint32 {
233	return cd.codes[:cd.pc]
234}
235
236func (cd *codeStore) PosList() []int {
237	return cd.lines[:cd.pc]
238}
239
240func (cd *codeStore) LastPC() int {
241	return cd.pc - 1
242}
243
244func (cd *codeStore) Last() uint32 {
245	if cd.pc == 0 {
246		return opInvalidInstruction
247	}
248	return cd.codes[cd.pc-1]
249}
250
251func (cd *codeStore) Pop() {
252	cd.pc--
253} /* }}} Code */
254
255/* {{{ VarNamePool */
256
257type varNamePoolValue struct {
258	Index int
259	Name  string
260}
261
262type varNamePool struct {
263	names  []string
264	offset int
265}
266
267func newVarNamePool(offset int) *varNamePool {
268	return &varNamePool{make([]string, 0, 16), offset}
269}
270
271func (vp *varNamePool) Names() []string {
272	return vp.names
273}
274
275func (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
284func (vp *varNamePool) LastIndex() int {
285	return vp.offset + len(vp.names)
286}
287
288func (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
297func (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
305func (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
314type codeBlock struct {
315	LocalVars  *varNamePool
316	BreakLabel int
317	Parent     *codeBlock
318	RefUpvalue bool
319	LineStart  int
320	LastLine   int
321}
322
323func 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
332type 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
344func 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
359func (fc *funcContext) NewLabel() int {
360	ret := fc.labelId
361	fc.labelId++
362	return ret
363}
364
365func (fc *funcContext) SetLabelPc(label int, pc int) {
366	fc.labelPc[label] = pc
367}
368
369func (fc *funcContext) GetLabelPc(label int) int {
370	return fc.labelPc[label]
371}
372
373func (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
388func (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
395func (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
404func (fc *funcContext) FindLocalVar(name string) int {
405	idx, _ := fc.FindLocalVarAndBlock(name)
406	return idx
407}
408
409func (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
417func (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
422func (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
431func (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
439func (fc *funcContext) EndScope() {
440	for _, vr := range fc.Block.LocalVars.List() {
441		fc.Proto.DbgLocals[vr.Index].EndPc = fc.Code.LastPC()
442	}
443}
444
445func (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
452func (fc *funcContext) RegTop() int {
453	return fc.regTop
454}
455
456/* FuncContext }}} */
457
458func compileChunk(context *funcContext, chunk []ast.Stmt) { // {{{
459	for _, stmt := range chunk {
460		compileStmt(context, stmt)
461	}
462} // }}}
463
464func 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
478func 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
509func 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
545func 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
605func 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
638func 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
678func 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
694func 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			return
713		}
714	}
715
716	for i, expr := range stmt.Exprs {
717		if i == lenexprs-1 && isVarArgReturnExpr(expr) {
718			compileExpr(context, reg, expr, ecnone(-2))
719			lastisvaarg = true
720		} else {
721			reg += compileExpr(context, reg, expr, ecnone(0))
722		}
723	}
724	count := reg - a + 1
725	if lastisvaarg {
726		count = 0
727	}
728	context.Code.AddABC(OP_RETURN, a, count, 0, sline(stmt))
729} // }}}
730
731func compileIfStmt(context *funcContext, stmt *ast.IfStmt) { // {{{
732	thenlabel := context.NewLabel()
733	elselabel := context.NewLabel()
734	endlabel := context.NewLabel()
735
736	compileBranchCondition(context, context.RegTop(), stmt.Condition, thenlabel, elselabel, false)
737	context.SetLabelPc(thenlabel, context.Code.LastPC())
738	compileBlock(context, stmt.Then)
739	if len(stmt.Else) > 0 {
740		context.Code.AddASbx(OP_JMP, 0, endlabel, sline(stmt))
741	}
742	context.SetLabelPc(elselabel, context.Code.LastPC())
743	if len(stmt.Else) > 0 {
744		compileBlock(context, stmt.Else)
745		context.SetLabelPc(endlabel, context.Code.LastPC())
746	}
747
748} // }}}
749
750func compileBranchCondition(context *funcContext, reg int, expr ast.Expr, thenlabel, elselabel int, hasnextcond bool) { // {{{
751	// TODO folding constants?
752	code := context.Code
753	flip := 0
754	jumplabel := elselabel
755	if hasnextcond {
756		flip = 1
757		jumplabel = thenlabel
758	}
759
760	switch ex := expr.(type) {
761	case *ast.FalseExpr, *ast.NilExpr:
762		if !hasnextcond {
763			code.AddASbx(OP_JMP, 0, elselabel, sline(expr))
764			return
765		}
766	case *ast.TrueExpr, *ast.NumberExpr, *ast.StringExpr:
767		if !hasnextcond {
768			return
769		}
770	case *ast.UnaryNotOpExpr:
771		compileBranchCondition(context, reg, ex.Expr, elselabel, thenlabel, !hasnextcond)
772		return
773	case *ast.LogicalOpExpr:
774		switch ex.Operator {
775		case "and":
776			nextcondlabel := context.NewLabel()
777			compileBranchCondition(context, reg, ex.Lhs, nextcondlabel, elselabel, false)
778			context.SetLabelPc(nextcondlabel, context.Code.LastPC())
779			compileBranchCondition(context, reg, ex.Rhs, thenlabel, elselabel, hasnextcond)
780		case "or":
781			nextcondlabel := context.NewLabel()
782			compileBranchCondition(context, reg, ex.Lhs, thenlabel, nextcondlabel, true)
783			context.SetLabelPc(nextcondlabel, context.Code.LastPC())
784			compileBranchCondition(context, reg, ex.Rhs, thenlabel, elselabel, hasnextcond)
785		}
786		return
787	case *ast.RelationalOpExpr:
788		compileRelationalOpExprAux(context, reg, ex, flip, jumplabel)
789		return
790	}
791
792	a := reg
793	compileExprWithMVPropagation(context, expr, &reg, &a)
794	code.AddABC(OP_TEST, a, 0, 0^flip, sline(expr))
795	code.AddASbx(OP_JMP, 0, jumplabel, sline(expr))
796} // }}}
797
798func compileWhileStmt(context *funcContext, stmt *ast.WhileStmt) { // {{{
799	thenlabel := context.NewLabel()
800	elselabel := context.NewLabel()
801	condlabel := context.NewLabel()
802
803	context.SetLabelPc(condlabel, context.Code.LastPC())
804	compileBranchCondition(context, context.RegTop(), stmt.Condition, thenlabel, elselabel, false)
805	context.SetLabelPc(thenlabel, context.Code.LastPC())
806	context.EnterBlock(elselabel, stmt)
807	compileChunk(context, stmt.Stmts)
808	context.CloseUpvalues()
809	context.Code.AddASbx(OP_JMP, 0, condlabel, eline(stmt))
810	context.LeaveBlock()
811	context.SetLabelPc(elselabel, context.Code.LastPC())
812} // }}}
813
814func compileRepeatStmt(context *funcContext, stmt *ast.RepeatStmt) { // {{{
815	initlabel := context.NewLabel()
816	thenlabel := context.NewLabel()
817	elselabel := context.NewLabel()
818
819	context.SetLabelPc(initlabel, context.Code.LastPC())
820	context.SetLabelPc(elselabel, context.Code.LastPC())
821	context.EnterBlock(thenlabel, stmt)
822	compileChunk(context, stmt.Stmts)
823	compileBranchCondition(context, context.RegTop(), stmt.Condition, thenlabel, elselabel, false)
824
825	context.SetLabelPc(thenlabel, context.Code.LastPC())
826	n := context.LeaveBlock()
827
828	if n > -1 {
829		label := context.NewLabel()
830		context.Code.AddASbx(OP_JMP, 0, label, eline(stmt))
831		context.SetLabelPc(elselabel, context.Code.LastPC())
832		context.Code.AddABC(OP_CLOSE, n, 0, 0, eline(stmt))
833		context.Code.AddASbx(OP_JMP, 0, initlabel, eline(stmt))
834		context.SetLabelPc(label, context.Code.LastPC())
835	}
836
837} // }}}
838
839func compileBreakStmt(context *funcContext, stmt *ast.BreakStmt) { // {{{
840	for block := context.Block; block != nil; block = block.Parent {
841		if label := block.BreakLabel; label != labelNoJump {
842			if block.RefUpvalue {
843				context.Code.AddABC(OP_CLOSE, block.Parent.LocalVars.LastIndex(), 0, 0, sline(stmt))
844			}
845			context.Code.AddASbx(OP_JMP, 0, label, sline(stmt))
846			return
847		}
848	}
849	raiseCompileError(context, sline(stmt), "no loop to break")
850} // }}}
851
852func compileFuncDefStmt(context *funcContext, stmt *ast.FuncDefStmt) { // {{{
853	if stmt.Name.Func == nil {
854		reg := context.RegTop()
855		var treg, kreg int
856		compileExprWithKMVPropagation(context, stmt.Name.Receiver, &reg, &treg)
857		kreg = loadRk(context, &reg, stmt.Func, LString(stmt.Name.Method))
858		compileExpr(context, reg, stmt.Func, ecfuncdef)
859		context.Code.AddABC(OP_SETTABLE, treg, kreg, reg, sline(stmt.Name.Receiver))
860	} else {
861		astmt := &ast.AssignStmt{Lhs: []ast.Expr{stmt.Name.Func}, Rhs: []ast.Expr{stmt.Func}}
862		astmt.SetLine(sline(stmt.Func))
863		astmt.SetLastLine(eline(stmt.Func))
864		compileAssignStmt(context, astmt)
865	}
866} // }}}
867
868func compileNumberForStmt(context *funcContext, stmt *ast.NumberForStmt) { // {{{
869	code := context.Code
870	endlabel := context.NewLabel()
871	ec := &expcontext{}
872
873	context.EnterBlock(endlabel, stmt)
874	reg := context.RegTop()
875	rindex := context.RegisterLocalVar("(for index)")
876	ecupdate(ec, ecLocal, rindex, 0)
877	compileExpr(context, reg, stmt.Init, ec)
878
879	reg = context.RegTop()
880	rlimit := context.RegisterLocalVar("(for limit)")
881	ecupdate(ec, ecLocal, rlimit, 0)
882	compileExpr(context, reg, stmt.Limit, ec)
883
884	reg = context.RegTop()
885	rstep := context.RegisterLocalVar("(for step)")
886	if stmt.Step == nil {
887		stmt.Step = &ast.NumberExpr{Value: "1"}
888		stmt.Step.SetLine(sline(stmt.Init))
889	}
890	ecupdate(ec, ecLocal, rstep, 0)
891	compileExpr(context, reg, stmt.Step, ec)
892
893	code.AddASbx(OP_FORPREP, rindex, 0, sline(stmt))
894
895	context.RegisterLocalVar(stmt.Name)
896
897	bodypc := code.LastPC()
898	compileChunk(context, stmt.Stmts)
899
900	context.LeaveBlock()
901
902	flpc := code.LastPC()
903	code.AddASbx(OP_FORLOOP, rindex, bodypc-(flpc+1), sline(stmt))
904
905	context.SetLabelPc(endlabel, code.LastPC())
906	code.SetSbx(bodypc, flpc-bodypc)
907
908} // }}}
909
910func compileGenericForStmt(context *funcContext, stmt *ast.GenericForStmt) { // {{{
911	code := context.Code
912	endlabel := context.NewLabel()
913	bodylabel := context.NewLabel()
914	fllabel := context.NewLabel()
915	nnames := len(stmt.Names)
916
917	context.EnterBlock(endlabel, stmt)
918	rgen := context.RegisterLocalVar("(for generator)")
919	context.RegisterLocalVar("(for state)")
920	context.RegisterLocalVar("(for control)")
921
922	compileRegAssignment(context, stmt.Names, stmt.Exprs, context.RegTop()-3, 3, sline(stmt))
923
924	code.AddASbx(OP_JMP, 0, fllabel, sline(stmt))
925
926	for _, name := range stmt.Names {
927		context.RegisterLocalVar(name)
928	}
929
930	context.SetLabelPc(bodylabel, code.LastPC())
931	compileChunk(context, stmt.Stmts)
932
933	context.LeaveBlock()
934
935	context.SetLabelPc(fllabel, code.LastPC())
936	code.AddABC(OP_TFORLOOP, rgen, 0, nnames, sline(stmt))
937	code.AddASbx(OP_JMP, 0, bodylabel, sline(stmt))
938
939	context.SetLabelPc(endlabel, code.LastPC())
940} // }}}
941
942func compileExpr(context *funcContext, reg int, expr ast.Expr, ec *expcontext) int { // {{{
943	code := context.Code
944	sreg := savereg(ec, reg)
945	sused := 1
946	if sreg < reg {
947		sused = 0
948	}
949
950	switch ex := expr.(type) {
951	case *ast.StringExpr:
952		code.AddABx(OP_LOADK, sreg, context.ConstIndex(LString(ex.Value)), sline(ex))
953		return sused
954	case *ast.NumberExpr:
955		num, err := parseNumber(ex.Value)
956		if err != nil {
957			num = LNumber(math.NaN())
958		}
959		code.AddABx(OP_LOADK, sreg, context.ConstIndex(num), sline(ex))
960		return sused
961	case *constLValueExpr:
962		code.AddABx(OP_LOADK, sreg, context.ConstIndex(ex.Value), sline(ex))
963		return sused
964	case *ast.NilExpr:
965		code.AddABC(OP_LOADNIL, sreg, sreg, 0, sline(ex))
966		return sused
967	case *ast.FalseExpr:
968		code.AddABC(OP_LOADBOOL, sreg, 0, 0, sline(ex))
969		return sused
970	case *ast.TrueExpr:
971		code.AddABC(OP_LOADBOOL, sreg, 1, 0, sline(ex))
972		return sused
973	case *ast.IdentExpr:
974		switch getIdentRefType(context, context, ex) {
975		case ecGlobal:
976			code.AddABx(OP_GETGLOBAL, sreg, context.ConstIndex(LString(ex.Value)), sline(ex))
977		case ecUpvalue:
978			code.AddABC(OP_GETUPVAL, sreg, context.Upvalues.RegisterUnique(ex.Value), 0, sline(ex))
979		case ecLocal:
980			b := context.FindLocalVar(ex.Value)
981			code.AddABC(OP_MOVE, sreg, b, 0, sline(ex))
982		}
983		return sused
984	case *ast.Comma3Expr:
985		if context.Proto.IsVarArg == 0 {
986			raiseCompileError(context, sline(ex), "cannot use '...' outside a vararg function")
987		}
988		context.Proto.IsVarArg &= ^VarArgNeedsArg
989		code.AddABC(OP_VARARG, sreg, 2+ec.varargopt, 0, sline(ex))
990		if context.RegTop() > (sreg+2+ec.varargopt) || ec.varargopt < -1 {
991			return 0
992		}
993		return (sreg + 1 + ec.varargopt) - reg
994	case *ast.AttrGetExpr:
995		a := sreg
996		b := reg
997		compileExprWithMVPropagation(context, ex.Object, &reg, &b)
998		c := reg
999		compileExprWithKMVPropagation(context, ex.Key, &reg, &c)
1000		opcode := OP_GETTABLE
1001		if _, ok := ex.Key.(*ast.StringExpr); ok {
1002			opcode = OP_GETTABLEKS
1003		}
1004		code.AddABC(opcode, a, b, c, sline(ex))
1005		return sused
1006	case *ast.TableExpr:
1007		compileTableExpr(context, reg, ex, ec)
1008		return 1
1009	case *ast.ArithmeticOpExpr:
1010		compileArithmeticOpExpr(context, reg, ex, ec)
1011		return sused
1012	case *ast.StringConcatOpExpr:
1013		compileStringConcatOpExpr(context, reg, ex, ec)
1014		return sused
1015	case *ast.UnaryMinusOpExpr, *ast.UnaryNotOpExpr, *ast.UnaryLenOpExpr:
1016		compileUnaryOpExpr(context, reg, ex, ec)
1017		return sused
1018	case *ast.RelationalOpExpr:
1019		compileRelationalOpExpr(context, reg, ex, ec)
1020		return sused
1021	case *ast.LogicalOpExpr:
1022		compileLogicalOpExpr(context, reg, ex, ec)
1023		return sused
1024	case *ast.FuncCallExpr:
1025		return compileFuncCallExpr(context, reg, ex, ec)
1026	case *ast.FunctionExpr:
1027		childcontext := newFuncContext(context.Proto.SourceName, context)
1028		compileFunctionExpr(childcontext, ex, ec)
1029		protono := len(context.Proto.FunctionPrototypes)
1030		context.Proto.FunctionPrototypes = append(context.Proto.FunctionPrototypes, childcontext.Proto)
1031		code.AddABx(OP_CLOSURE, sreg, protono, sline(ex))
1032		for _, upvalue := range childcontext.Upvalues.List() {
1033			localidx, block := context.FindLocalVarAndBlock(upvalue.Name)
1034			if localidx > -1 {
1035				code.AddABC(OP_MOVE, 0, localidx, 0, sline(ex))
1036				block.RefUpvalue = true
1037			} else {
1038				upvalueidx := context.Upvalues.Find(upvalue.Name)
1039				if upvalueidx < 0 {
1040					upvalueidx = context.Upvalues.RegisterUnique(upvalue.Name)
1041				}
1042				code.AddABC(OP_GETUPVAL, 0, upvalueidx, 0, sline(ex))
1043			}
1044		}
1045		return sused
1046	default:
1047		panic(fmt.Sprintf("expr %v not implemented.", reflect.TypeOf(ex).Elem().Name()))
1048	}
1049
1050	panic("should not reach here")
1051	return sused
1052} // }}}
1053
1054func compileExprWithPropagation(context *funcContext, expr ast.Expr, reg *int, save *int, propergator func(int, *int, *int, int)) { // {{{
1055	reginc := compileExpr(context, *reg, expr, ecnone(0))
1056	if _, ok := expr.(*ast.LogicalOpExpr); ok {
1057		*save = *reg
1058		*reg = *reg + reginc
1059	} else {
1060		propergator(context.RegTop(), save, reg, reginc)
1061	}
1062} // }}}
1063
1064func compileExprWithKMVPropagation(context *funcContext, expr ast.Expr, reg *int, save *int) { // {{{
1065	compileExprWithPropagation(context, expr, reg, save, context.Code.PropagateKMV)
1066} // }}}
1067
1068func compileExprWithMVPropagation(context *funcContext, expr ast.Expr, reg *int, save *int) { // {{{
1069	compileExprWithPropagation(context, expr, reg, save, context.Code.PropagateMV)
1070} // }}}
1071
1072func constFold(exp ast.Expr) ast.Expr { // {{{
1073	switch expr := exp.(type) {
1074	case *ast.ArithmeticOpExpr:
1075		lvalue, lisconst := lnumberValue(expr.Lhs)
1076		rvalue, risconst := lnumberValue(expr.Rhs)
1077		if lisconst && risconst {
1078			switch expr.Operator {
1079			case "+":
1080				return &constLValueExpr{Value: lvalue + rvalue}
1081			case "-":
1082				return &constLValueExpr{Value: lvalue - rvalue}
1083			case "*":
1084				return &constLValueExpr{Value: lvalue * rvalue}
1085			case "/":
1086				return &constLValueExpr{Value: lvalue / rvalue}
1087			case "%":
1088				return &constLValueExpr{Value: luaModulo(lvalue, rvalue)}
1089			case "^":
1090				return &constLValueExpr{Value: LNumber(math.Pow(float64(lvalue), float64(rvalue)))}
1091			default:
1092				panic(fmt.Sprintf("unknown binop: %v", expr.Operator))
1093			}
1094		} else {
1095			retexpr := *expr
1096			retexpr.Lhs = constFold(expr.Lhs)
1097			retexpr.Rhs = constFold(expr.Rhs)
1098			return &retexpr
1099		}
1100	case *ast.UnaryMinusOpExpr:
1101		expr.Expr = constFold(expr.Expr)
1102		if value, ok := lnumberValue(expr.Expr); ok {
1103			return &constLValueExpr{Value: LNumber(-value)}
1104		}
1105		return expr
1106	default:
1107
1108		return exp
1109	}
1110	return exp
1111} // }}}
1112
1113func compileFunctionExpr(context *funcContext, funcexpr *ast.FunctionExpr, ec *expcontext) { // {{{
1114	context.Proto.LineDefined = sline(funcexpr)
1115	context.Proto.LastLineDefined = eline(funcexpr)
1116	if len(funcexpr.ParList.Names) > maxRegisters {
1117		raiseCompileError(context, context.Proto.LineDefined, "register overflow")
1118	}
1119	context.Proto.NumParameters = uint8(len(funcexpr.ParList.Names))
1120	if ec.ctype == ecMethod {
1121		context.Proto.NumParameters += 1
1122		context.RegisterLocalVar("self")
1123	}
1124	for _, name := range funcexpr.ParList.Names {
1125		context.RegisterLocalVar(name)
1126	}
1127	if funcexpr.ParList.HasVargs {
1128		if CompatVarArg {
1129			context.Proto.IsVarArg = VarArgHasArg | VarArgNeedsArg
1130			if context.Parent != nil {
1131				context.RegisterLocalVar("arg")
1132			}
1133		}
1134		context.Proto.IsVarArg |= VarArgIsVarArg
1135	}
1136
1137	compileChunk(context, funcexpr.Stmts)
1138
1139	context.Code.AddABC(OP_RETURN, 0, 1, 0, eline(funcexpr))
1140	context.EndScope()
1141	context.Proto.Code = context.Code.List()
1142	context.Proto.DbgSourcePositions = context.Code.PosList()
1143	context.Proto.DbgUpvalues = context.Upvalues.Names()
1144	context.Proto.NumUpvalues = uint8(len(context.Proto.DbgUpvalues))
1145	for _, clv := range context.Proto.Constants {
1146		sv := ""
1147		if slv, ok := clv.(LString); ok {
1148			sv = string(slv)
1149		}
1150		context.Proto.stringConstants = append(context.Proto.stringConstants, sv)
1151	}
1152	patchCode(context)
1153} // }}}
1154
1155func compileTableExpr(context *funcContext, reg int, ex *ast.TableExpr, ec *expcontext) { // {{{
1156	code := context.Code
1157	/*
1158		tablereg := savereg(ec, reg)
1159		if tablereg == reg {
1160			reg += 1
1161		}
1162	*/
1163	tablereg := reg
1164	reg++
1165	code.AddABC(OP_NEWTABLE, tablereg, 0, 0, sline(ex))
1166	tablepc := code.LastPC()
1167	regbase := reg
1168
1169	arraycount := 0
1170	lastvararg := false
1171	for i, field := range ex.Fields {
1172		islast := i == len(ex.Fields)-1
1173		if field.Key == nil {
1174			if islast && isVarArgReturnExpr(field.Value) {
1175				reg += compileExpr(context, reg, field.Value, ecnone(-2))
1176				lastvararg = true
1177			} else {
1178				reg += compileExpr(context, reg, field.Value, ecnone(0))
1179				arraycount += 1
1180			}
1181		} else {
1182			regorg := reg
1183			b := reg
1184			compileExprWithKMVPropagation(context, field.Key, &reg, &b)
1185			c := reg
1186			compileExprWithKMVPropagation(context, field.Value, &reg, &c)
1187			opcode := OP_SETTABLE
1188			if _, ok := field.Key.(*ast.StringExpr); ok {
1189				opcode = OP_SETTABLEKS
1190			}
1191			code.AddABC(opcode, tablereg, b, c, sline(ex))
1192			reg = regorg
1193		}
1194		flush := arraycount % FieldsPerFlush
1195		if (arraycount != 0 && (flush == 0 || islast)) || lastvararg {
1196			reg = regbase
1197			num := flush
1198			if num == 0 {
1199				num = FieldsPerFlush
1200			}
1201			c := (arraycount-1)/FieldsPerFlush + 1
1202			b := num
1203			if islast && isVarArgReturnExpr(field.Value) {
1204				b = 0
1205			}
1206			line := field.Value
1207			if field.Key != nil {
1208				line = field.Key
1209			}
1210			if c > 511 {
1211				c = 0
1212			}
1213			code.AddABC(OP_SETLIST, tablereg, b, c, sline(line))
1214			if c == 0 {
1215				code.Add(uint32(c), sline(line))
1216			}
1217		}
1218	}
1219	code.SetB(tablepc, int2Fb(arraycount))
1220	code.SetC(tablepc, int2Fb(len(ex.Fields)-arraycount))
1221	if ec.ctype == ecLocal && ec.reg != tablereg {
1222		code.AddABC(OP_MOVE, ec.reg, tablereg, 0, sline(ex))
1223	}
1224} // }}}
1225
1226func compileArithmeticOpExpr(context *funcContext, reg int, expr *ast.ArithmeticOpExpr, ec *expcontext) { // {{{
1227	exp := constFold(expr)
1228	if ex, ok := exp.(*constLValueExpr); ok {
1229		exp.SetLine(sline(expr))
1230		compileExpr(context, reg, ex, ec)
1231		return
1232	}
1233	expr, _ = exp.(*ast.ArithmeticOpExpr)
1234	a := savereg(ec, reg)
1235	b := reg
1236	compileExprWithKMVPropagation(context, expr.Lhs, &reg, &b)
1237	c := reg
1238	compileExprWithKMVPropagation(context, expr.Rhs, &reg, &c)
1239
1240	op := 0
1241	switch expr.Operator {
1242	case "+":
1243		op = OP_ADD
1244	case "-":
1245		op = OP_SUB
1246	case "*":
1247		op = OP_MUL
1248	case "/":
1249		op = OP_DIV
1250	case "%":
1251		op = OP_MOD
1252	case "^":
1253		op = OP_POW
1254	}
1255	context.Code.AddABC(op, a, b, c, sline(expr))
1256} // }}}
1257
1258func compileStringConcatOpExpr(context *funcContext, reg int, expr *ast.StringConcatOpExpr, ec *expcontext) { // {{{
1259	code := context.Code
1260	crange := 1
1261	for current := expr.Rhs; current != nil; {
1262		if ex, ok := current.(*ast.StringConcatOpExpr); ok {
1263			crange += 1
1264			current = ex.Rhs
1265		} else {
1266			current = nil
1267		}
1268	}
1269	a := savereg(ec, reg)
1270	basereg := reg
1271	reg += compileExpr(context, reg, expr.Lhs, ecnone(0))
1272	reg += compileExpr(context, reg, expr.Rhs, ecnone(0))
1273	for pc := code.LastPC(); pc != 0 && opGetOpCode(code.At(pc)) == OP_CONCAT; pc-- {
1274		code.Pop()
1275	}
1276	code.AddABC(OP_CONCAT, a, basereg, basereg+crange, sline(expr))
1277} // }}}
1278
1279func compileUnaryOpExpr(context *funcContext, reg int, expr ast.Expr, ec *expcontext) { // {{{
1280	opcode := 0
1281	code := context.Code
1282	var operandexpr ast.Expr
1283	switch ex := expr.(type) {
1284	case *ast.UnaryMinusOpExpr:
1285		exp := constFold(ex)
1286		if lvexpr, ok := exp.(*constLValueExpr); ok {
1287			exp.SetLine(sline(expr))
1288			compileExpr(context, reg, lvexpr, ec)
1289			return
1290		}
1291		ex, _ = exp.(*ast.UnaryMinusOpExpr)
1292		operandexpr = ex.Expr
1293		opcode = OP_UNM
1294	case *ast.UnaryNotOpExpr:
1295		switch ex.Expr.(type) {
1296		case *ast.TrueExpr:
1297			code.AddABC(OP_LOADBOOL, savereg(ec, reg), 0, 0, sline(expr))
1298			return
1299		case *ast.FalseExpr, *ast.NilExpr:
1300			code.AddABC(OP_LOADBOOL, savereg(ec, reg), 1, 0, sline(expr))
1301			return
1302		default:
1303			opcode = OP_NOT
1304			operandexpr = ex.Expr
1305		}
1306	case *ast.UnaryLenOpExpr:
1307		opcode = OP_LEN
1308		operandexpr = ex.Expr
1309	}
1310
1311	a := savereg(ec, reg)
1312	b := reg
1313	compileExprWithMVPropagation(context, operandexpr, &reg, &b)
1314	code.AddABC(opcode, a, b, 0, sline(expr))
1315} // }}}
1316
1317func compileRelationalOpExprAux(context *funcContext, reg int, expr *ast.RelationalOpExpr, flip int, label int) { // {{{
1318	code := context.Code
1319	b := reg
1320	compileExprWithKMVPropagation(context, expr.Lhs, &reg, &b)
1321	c := reg
1322	compileExprWithKMVPropagation(context, expr.Rhs, &reg, &c)
1323	switch expr.Operator {
1324	case "<":
1325		code.AddABC(OP_LT, 0^flip, b, c, sline(expr))
1326	case ">":
1327		code.AddABC(OP_LT, 0^flip, c, b, sline(expr))
1328	case "<=":
1329		code.AddABC(OP_LE, 0^flip, b, c, sline(expr))
1330	case ">=":
1331		code.AddABC(OP_LE, 0^flip, c, b, sline(expr))
1332	case "==":
1333		code.AddABC(OP_EQ, 0^flip, b, c, sline(expr))
1334	case "~=":
1335		code.AddABC(OP_EQ, 1^flip, b, c, sline(expr))
1336	}
1337	code.AddASbx(OP_JMP, 0, label, sline(expr))
1338} // }}}
1339
1340func compileRelationalOpExpr(context *funcContext, reg int, expr *ast.RelationalOpExpr, ec *expcontext) { // {{{
1341	a := savereg(ec, reg)
1342	code := context.Code
1343	jumplabel := context.NewLabel()
1344	compileRelationalOpExprAux(context, reg, expr, 1, jumplabel)
1345	code.AddABC(OP_LOADBOOL, a, 0, 1, sline(expr))
1346	context.SetLabelPc(jumplabel, code.LastPC())
1347	code.AddABC(OP_LOADBOOL, a, 1, 0, sline(expr))
1348} // }}}
1349
1350func compileLogicalOpExpr(context *funcContext, reg int, expr *ast.LogicalOpExpr, ec *expcontext) { // {{{
1351	a := savereg(ec, reg)
1352	code := context.Code
1353	endlabel := context.NewLabel()
1354	lb := &lblabels{context.NewLabel(), context.NewLabel(), endlabel, false}
1355	nextcondlabel := context.NewLabel()
1356	if expr.Operator == "and" {
1357		compileLogicalOpExprAux(context, reg, expr.Lhs, ec, nextcondlabel, endlabel, false, lb)
1358		context.SetLabelPc(nextcondlabel, code.LastPC())
1359		compileLogicalOpExprAux(context, reg, expr.Rhs, ec, endlabel, endlabel, false, lb)
1360	} else {
1361		compileLogicalOpExprAux(context, reg, expr.Lhs, ec, endlabel, nextcondlabel, true, lb)
1362		context.SetLabelPc(nextcondlabel, code.LastPC())
1363		compileLogicalOpExprAux(context, reg, expr.Rhs, ec, endlabel, endlabel, false, lb)
1364	}
1365
1366	if lb.b {
1367		context.SetLabelPc(lb.f, code.LastPC())
1368		code.AddABC(OP_LOADBOOL, a, 0, 1, sline(expr))
1369		context.SetLabelPc(lb.t, code.LastPC())
1370		code.AddABC(OP_LOADBOOL, a, 1, 0, sline(expr))
1371	}
1372
1373	lastinst := code.Last()
1374	if opGetOpCode(lastinst) == OP_JMP && opGetArgSbx(lastinst) == endlabel {
1375		code.Pop()
1376	}
1377
1378	context.SetLabelPc(endlabel, code.LastPC())
1379} // }}}
1380
1381func compileLogicalOpExprAux(context *funcContext, reg int, expr ast.Expr, ec *expcontext, thenlabel, elselabel int, hasnextcond bool, lb *lblabels) { // {{{
1382	// TODO folding constants?
1383	code := context.Code
1384	flip := 0
1385	jumplabel := elselabel
1386	if hasnextcond {
1387		flip = 1
1388		jumplabel = thenlabel
1389	}
1390
1391	switch ex := expr.(type) {
1392	case *ast.FalseExpr:
1393		if elselabel == lb.e {
1394			code.AddASbx(OP_JMP, 0, lb.f, sline(expr))
1395			lb.b = true
1396		} else {
1397			code.AddASbx(OP_JMP, 0, elselabel, sline(expr))
1398		}
1399		return
1400	case *ast.NilExpr:
1401		if elselabel == lb.e {
1402			compileExpr(context, reg, expr, ec)
1403			code.AddASbx(OP_JMP, 0, lb.e, sline(expr))
1404		} else {
1405			code.AddASbx(OP_JMP, 0, elselabel, sline(expr))
1406		}
1407		return
1408	case *ast.TrueExpr:
1409		if thenlabel == lb.e {
1410			code.AddASbx(OP_JMP, 0, lb.t, sline(expr))
1411			lb.b = true
1412		} else {
1413			code.AddASbx(OP_JMP, 0, thenlabel, sline(expr))
1414		}
1415		return
1416	case *ast.NumberExpr, *ast.StringExpr:
1417		if thenlabel == lb.e {
1418			compileExpr(context, reg, expr, ec)
1419			code.AddASbx(OP_JMP, 0, lb.e, sline(expr))
1420		} else {
1421			code.AddASbx(OP_JMP, 0, thenlabel, sline(expr))
1422		}
1423		return
1424	case *ast.LogicalOpExpr:
1425		switch ex.Operator {
1426		case "and":
1427			nextcondlabel := context.NewLabel()
1428			compileLogicalOpExprAux(context, reg, ex.Lhs, ec, nextcondlabel, elselabel, false, lb)
1429			context.SetLabelPc(nextcondlabel, context.Code.LastPC())
1430			compileLogicalOpExprAux(context, reg, ex.Rhs, ec, thenlabel, elselabel, hasnextcond, lb)
1431		case "or":
1432			nextcondlabel := context.NewLabel()
1433			compileLogicalOpExprAux(context, reg, ex.Lhs, ec, thenlabel, nextcondlabel, true, lb)
1434			context.SetLabelPc(nextcondlabel, context.Code.LastPC())
1435			compileLogicalOpExprAux(context, reg, ex.Rhs, ec, thenlabel, elselabel, hasnextcond, lb)
1436		}
1437		return
1438	case *ast.RelationalOpExpr:
1439		if thenlabel == elselabel {
1440			flip ^= 1
1441			jumplabel = lb.t
1442			lb.b = true
1443		} else if thenlabel == lb.e {
1444			jumplabel = lb.t
1445			lb.b = true
1446		} else if elselabel == lb.e {
1447			jumplabel = lb.f
1448			lb.b = true
1449		}
1450		compileRelationalOpExprAux(context, reg, ex, flip, jumplabel)
1451		return
1452	}
1453
1454	if !hasnextcond && thenlabel == elselabel {
1455		reg += compileExpr(context, reg, expr, ec)
1456	} else {
1457		a := reg
1458		sreg := savereg(ec, a)
1459		reg += compileExpr(context, reg, expr, ecnone(0))
1460		if sreg == a {
1461			code.AddABC(OP_TEST, a, 0, 0^flip, sline(expr))
1462		} else {
1463			code.AddABC(OP_TESTSET, sreg, a, 0^flip, sline(expr))
1464		}
1465	}
1466	code.AddASbx(OP_JMP, 0, jumplabel, sline(expr))
1467} // }}}
1468
1469func compileFuncCallExpr(context *funcContext, reg int, expr *ast.FuncCallExpr, ec *expcontext) int { // {{{
1470	funcreg := reg
1471	if ec.ctype == ecLocal && ec.reg == (int(context.Proto.NumParameters)-1) {
1472		funcreg = ec.reg
1473		reg = ec.reg
1474	}
1475	argc := len(expr.Args)
1476	islastvararg := false
1477	name := "(anonymous)"
1478
1479	if expr.Func != nil { // hoge.func()
1480		reg += compileExpr(context, reg, expr.Func, ecnone(0))
1481		name = getExprName(context, expr.Func)
1482	} else { // hoge:method()
1483		b := reg
1484		compileExprWithMVPropagation(context, expr.Receiver, &reg, &b)
1485		c := loadRk(context, &reg, expr, LString(expr.Method))
1486		context.Code.AddABC(OP_SELF, funcreg, b, c, sline(expr))
1487		// increments a register for an implicit "self"
1488		reg = b + 1
1489		reg2 := funcreg + 2
1490		if reg2 > reg {
1491			reg = reg2
1492		}
1493		argc += 1
1494		name = string(expr.Method)
1495	}
1496
1497	for i, ar := range expr.Args {
1498		islastvararg = (i == len(expr.Args)-1) && isVarArgReturnExpr(ar)
1499		if islastvararg {
1500			compileExpr(context, reg, ar, ecnone(-2))
1501		} else {
1502			reg += compileExpr(context, reg, ar, ecnone(0))
1503		}
1504	}
1505	b := argc + 1
1506	if islastvararg {
1507		b = 0
1508	}
1509	context.Code.AddABC(OP_CALL, funcreg, b, ec.varargopt+2, sline(expr))
1510	context.Proto.DbgCalls = append(context.Proto.DbgCalls, DbgCall{Pc: context.Code.LastPC(), Name: name})
1511
1512	if ec.varargopt == 0 && ec.ctype == ecLocal && funcreg != ec.reg {
1513		context.Code.AddABC(OP_MOVE, ec.reg, funcreg, 0, sline(expr))
1514		return 1
1515	}
1516	if context.RegTop() > (funcreg+2+ec.varargopt) || ec.varargopt < -1 {
1517		return 0
1518	}
1519	return ec.varargopt + 1
1520} // }}}
1521
1522func loadRk(context *funcContext, reg *int, expr ast.Expr, cnst LValue) int { // {{{
1523	cindex := context.ConstIndex(cnst)
1524	if cindex <= opMaxIndexRk {
1525		return opRkAsk(cindex)
1526	} else {
1527		ret := *reg
1528		*reg++
1529		context.Code.AddABx(OP_LOADK, ret, cindex, sline(expr))
1530		return ret
1531	}
1532} // }}}
1533
1534func getIdentRefType(context *funcContext, current *funcContext, expr *ast.IdentExpr) expContextType { // {{{
1535	if current == nil {
1536		return ecGlobal
1537	} else if current.FindLocalVar(expr.Value) > -1 {
1538		if current == context {
1539			return ecLocal
1540		}
1541		return ecUpvalue
1542	}
1543	return getIdentRefType(context, current.Parent, expr)
1544} // }}}
1545
1546func getExprName(context *funcContext, expr ast.Expr) string { // {{{
1547	switch ex := expr.(type) {
1548	case *ast.IdentExpr:
1549		return ex.Value
1550	case *ast.AttrGetExpr:
1551		switch kex := ex.Key.(type) {
1552		case *ast.StringExpr:
1553			return kex.Value
1554		}
1555		return "?"
1556	}
1557	return "?"
1558} // }}}
1559
1560func patchCode(context *funcContext) { // {{{
1561	maxreg := 1
1562	if np := int(context.Proto.NumParameters); np > 1 {
1563		maxreg = np
1564	}
1565	moven := 0
1566	code := context.Code.List()
1567	for pc := 0; pc < len(code); pc++ {
1568		inst := code[pc]
1569		curop := opGetOpCode(inst)
1570		switch curop {
1571		case OP_CLOSURE:
1572			pc += int(context.Proto.FunctionPrototypes[opGetArgBx(inst)].NumUpvalues)
1573			moven = 0
1574			continue
1575		case OP_SETGLOBAL, OP_SETUPVAL, OP_EQ, OP_LT, OP_LE, OP_TEST,
1576			OP_TAILCALL, OP_RETURN, OP_FORPREP, OP_FORLOOP, OP_TFORLOOP,
1577			OP_SETLIST, OP_CLOSE:
1578			/* nothing to do */
1579		case OP_CALL:
1580			if reg := opGetArgA(inst) + opGetArgC(inst) - 2; reg > maxreg {
1581				maxreg = reg
1582			}
1583		case OP_VARARG:
1584			if reg := opGetArgA(inst) + opGetArgB(inst) - 1; reg > maxreg {
1585				maxreg = reg
1586			}
1587		case OP_SELF:
1588			if reg := opGetArgA(inst) + 1; reg > maxreg {
1589				maxreg = reg
1590			}
1591		case OP_LOADNIL:
1592			if reg := opGetArgB(inst); reg > maxreg {
1593				maxreg = reg
1594			}
1595		case OP_JMP: // jump to jump optimization
1596			distance := 0
1597			count := 0 // avoiding infinite loops
1598			for jmp := inst; opGetOpCode(jmp) == OP_JMP && count < 5; jmp = context.Code.At(pc + distance + 1) {
1599				d := context.GetLabelPc(opGetArgSbx(jmp)) - pc
1600				if d > opMaxArgSbx {
1601					if distance == 0 {
1602						raiseCompileError(context, context.Proto.LineDefined, "too long to jump.")
1603					}
1604					break
1605				}
1606				distance = d
1607				count++
1608			}
1609			if distance == 0 {
1610				context.Code.SetOpCode(pc, OP_NOP)
1611			} else {
1612				context.Code.SetSbx(pc, distance)
1613			}
1614		default:
1615			if reg := opGetArgA(inst); reg > maxreg {
1616				maxreg = reg
1617			}
1618		}
1619
1620		// bulk move optimization(reducing op dipatch costs)
1621		if curop == OP_MOVE {
1622			moven++
1623		} else {
1624			if moven > 1 {
1625				context.Code.SetOpCode(pc-moven, OP_MOVEN)
1626				context.Code.SetC(pc-moven, intMin(moven-1, opMaxArgsC))
1627			}
1628			moven = 0
1629		}
1630	}
1631	maxreg++
1632	if maxreg > maxRegisters {
1633		raiseCompileError(context, context.Proto.LineDefined, "register overflow(too many local variables)")
1634	}
1635	context.Proto.NumUsedRegisters = uint8(maxreg)
1636} // }}}
1637
1638func Compile(chunk []ast.Stmt, name string) (proto *FunctionProto, err error) { // {{{
1639	defer func() {
1640		if rcv := recover(); rcv != nil {
1641			if _, ok := rcv.(*CompileError); ok {
1642				err = rcv.(error)
1643			} else {
1644				panic(rcv)
1645			}
1646		}
1647	}()
1648	err = nil
1649	parlist := &ast.ParList{HasVargs: true, Names: []string{}}
1650	funcexpr := &ast.FunctionExpr{ParList: parlist, Stmts: chunk}
1651	context := newFuncContext(name, nil)
1652	compileFunctionExpr(context, funcexpr, ecnone(0))
1653	proto = context.Proto
1654	return
1655} // }}}
1656