stmt.go

/*
 Copyright 2021 The GoPlus Authors (goplus.org)
 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
 You may obtain a copy of the License at
     http://www.apache.org/licenses/LICENSE-2.0
 Unless required by applicable law or agreed to in writing, software
 distributed under the License is distributed on an "AS IS" BASIS,
 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
 limitations under the License.
*/

package gox

import (
	"go/ast"
	"go/token"
	"go/types"
	"log"

	"github.com/goplus/gox/internal"
)

type controlFlow interface {
	Then(cb *CodeBuilder, src ...ast.Node)
}

// ----------------------------------------------------------------------------
//
// block
//
//	...
//
// end
type blockStmt struct {
	old codeBlockCtx
}

func (p *blockStmt) End(cb *CodeBuilder, src ast.Node) {
	stmts, flows := cb.endBlockStmt(&p.old)
	cb.current.flows |= flows
	cb.emitStmt(&ast.BlockStmt{List: stmts})
}

// ----------------------------------------------------------------------------
//
// vblock
//
//	...
//
// end
type vblockStmt struct {
	old vblockCtx
}

func (p *vblockStmt) End(cb *CodeBuilder, src ast.Node) {
	cb.endVBlockStmt(&p.old)
}

// ----------------------------------------------------------------------------
//
// if init; cond then
//
//	...
//
// else
//
//	...
//
// end
type ifStmt struct {
	init ast.Stmt
	cond ast.Expr
	body *ast.BlockStmt
	old  codeBlockCtx
	old2 codeBlockCtx
}

func (p *ifStmt) Then(cb *CodeBuilder, src ...ast.Node) {
	cond := cb.stk.Pop()
	if !types.AssignableTo(cond.Type, types.Typ[types.Bool]) {
		panic("TODO: if statement condition is not a boolean expr")
	}
	p.cond = cond.Val
	switch stmts := cb.clearBlockStmt(); len(stmts) {
	case 0:
		// nothing to do
	case 1:
		p.init = stmts[0]
	default:
		panic("TODO: if statement has too many init statements")
	}
	cb.startBlockStmt(p, src, "if body", &p.old2)
}

func (p *ifStmt) Else(cb *CodeBuilder, src ...ast.Node) {
	if p.body != nil {
		panic("TODO: else statement already exists")
	}

	stmts, flows := cb.endBlockStmt(&p.old2)
	cb.current.flows |= flows

	p.body = &ast.BlockStmt{List: stmts}
	cb.startBlockStmt(p, src, "else body", &p.old2)
}

func (p *ifStmt) End(cb *CodeBuilder, src ast.Node) {
	stmts, flows := cb.endBlockStmt(&p.old2)
	cb.current.flows |= flows

	var blockStmt = &ast.BlockStmt{List: stmts}
	var el ast.Stmt
	if p.body != nil { // if..else
		el = blockStmt
		if len(stmts) == 1 { // if..else if
			if stmt, ok := stmts[0].(*ast.IfStmt); ok {
				el = stmt
			}
		}
	} else { // if without else
		p.body = blockStmt
	}
	cb.endBlockStmt(&p.old)
	cb.emitStmt(&ast.IfStmt{Init: p.init, Cond: p.cond, Body: p.body, Else: el})
}

// ----------------------------------------------------------------------------
//
// switch init; tag then
//
//	expr1, expr2, ..., exprN case(N)
//	  ...
//	  end
//	expr1, expr2, ..., exprM case(M)
//	  ...
//	  end
//
// end
type switchStmt struct {
	init ast.Stmt
	tag  *internal.Elem
	old  codeBlockCtx
}

func (p *switchStmt) Then(cb *CodeBuilder, src ...ast.Node) {
	p.tag = cb.stk.Pop()
	switch stmts := cb.clearBlockStmt(); len(stmts) {
	case 0:
		// nothing to do
	case 1:
		p.init = stmts[0]
	default:
		panic("TODO: switch statement has too many init statements")
	}
}

func (p *switchStmt) Case(cb *CodeBuilder, n int, src ...ast.Node) {
	var list []ast.Expr
	if n > 0 {
		list = make([]ast.Expr, n)
		for i, arg := range cb.stk.GetArgs(n) {
			if p.tag.Val != nil { // switch tag {...}
				if !ComparableTo(cb.pkg, arg, p.tag) {
					src, pos := cb.loadExpr(arg.Src)
					cb.panicCodeErrorf(
						pos, "cannot use %s (type %v) as type %v", src, arg.Type, types.Default(p.tag.Type))
				}
			} else { // switch {...}
				if !types.AssignableTo(arg.Type, types.Typ[types.Bool]) && arg.Type != TyEmptyInterface {
					src, pos := cb.loadExpr(arg.Src)
					cb.panicCodeErrorf(pos, "cannot use %s (type %v) as type bool", src, arg.Type)
				}
			}
			list[i] = arg.Val
		}
		cb.stk.PopN(n)
	}
	stmt := &caseStmt{list: list}
	cb.startBlockStmt(stmt, src, "case statement", &stmt.old)
}

func (p *switchStmt) End(cb *CodeBuilder, src ast.Node) {
	if p.tag == nil {
		return
	}
	stmts, flows := cb.endBlockStmt(&p.old)
	cb.current.flows |= (flows &^ flowFlagBreak)

	body := &ast.BlockStmt{List: stmts}
	cb.emitStmt(&ast.SwitchStmt{Init: p.init, Tag: checkParenExpr(p.tag.Val), Body: body})
}

type caseStmt struct {
	list []ast.Expr
	old  codeBlockCtx
}

func (p *caseStmt) Fallthrough(cb *CodeBuilder) {
	cb.emitStmt(&ast.BranchStmt{Tok: token.FALLTHROUGH})
}

func (p *caseStmt) End(cb *CodeBuilder, src ast.Node) {
	body, flows := cb.endBlockStmt(&p.old)
	cb.current.flows |= flows
	cb.emitStmt(&ast.CaseClause{List: p.list, Body: body})
}

// ----------------------------------------------------------------------------
//
// select
// commStmt1 commCase()
//
//	...
//	end
//
// commStmt2 commCase()
//
//	...
//	end
//
// end
type selectStmt struct {
	old codeBlockCtx
}

func (p *selectStmt) CommCase(cb *CodeBuilder, n int, src ...ast.Node) {
	var comm ast.Stmt
	if n == 1 {
		comm = cb.popStmt()
	}
	stmt := &commCase{comm: comm}
	cb.startBlockStmt(stmt, src, "comm case statement", &stmt.old)
}

func (p *selectStmt) End(cb *CodeBuilder, src ast.Node) {
	stmts, flows := cb.endBlockStmt(&p.old)
	cb.current.flows |= (flows &^ flowFlagBreak)
	cb.emitStmt(&ast.SelectStmt{Body: &ast.BlockStmt{List: stmts}})
}

type commCase struct {
	comm ast.Stmt
	old  codeBlockCtx
}

func (p *commCase) End(cb *CodeBuilder, src ast.Node) {
	body, flows := cb.endBlockStmt(&p.old)
	cb.current.flows |= flows
	cb.emitStmt(&ast.CommClause{Comm: p.comm, Body: body})
}

// ----------------------------------------------------------------------------
//
// typeSwitch(name) init; expr typeAssertThen
// type1, type2, ... typeN typeCase(N)
//
//	...
//	end
//
// type1, type2, ... typeM typeCase(M)
//
//	...
//	end
//
// end
type typeSwitchStmt struct {
	init  ast.Stmt
	name  string
	x     ast.Expr
	xSrc  ast.Node
	xType *types.Interface
	old   codeBlockCtx
}

func (p *typeSwitchStmt) TypeAssertThen(cb *CodeBuilder) {
	switch stmts := cb.clearBlockStmt(); len(stmts) {
	case 0:
		// nothing to do
	case 1:
		p.init = stmts[0]
	default:
		panic("TODO: type switch statement has too many init statements")
	}
	x := cb.stk.Pop()
	xType, ok := cb.checkInterface(x.Type)
	if !ok {
		panic("TODO: can't type assert on non interface expr")
	}
	p.x, p.xSrc, p.xType = x.Val, x.Src, xType
}

func (p *typeSwitchStmt) TypeCase(cb *CodeBuilder, n int, src ...ast.Node) {
	var list []ast.Expr
	var typ types.Type
	if n > 0 {
		list = make([]ast.Expr, n)
		args := cb.stk.GetArgs(n)
		for i, arg := range args {
			typ = arg.Type
			if tt, ok := typ.(*TypeType); ok {
				typ = tt.Type()
				if missing := cb.missingMethod(typ, p.xType); missing != "" {
					xsrc, _ := cb.loadExpr(p.xSrc)
					pos := getSrcPos(arg.Src)
					cb.panicCodeErrorf(
						pos, "impossible type switch case: %s (type %v) cannot have dynamic type %v (missing %s method)",
						xsrc, p.xType, typ, missing)
				}
			} else if typ != types.Typ[types.UntypedNil] {
				src, pos := cb.loadExpr(arg.Src)
				cb.panicCodeErrorf(pos, "%s (type %v) is not a type", src, typ)
			}
			list[i] = arg.Val
		}
		cb.stk.PopN(n)
	}

	stmt := &typeCaseStmt{list: list}
	cb.startBlockStmt(stmt, src, "type case statement", &stmt.old)

	if p.name != "" {
		if n != 1 { // default, or case with multi expr
			typ = p.xType
		}
		name := types.NewParam(token.NoPos, cb.pkg.Types, p.name, typ)
		cb.current.scope.Insert(name)
	}
}

func (p *typeSwitchStmt) End(cb *CodeBuilder, src ast.Node) {
	stmts, flows := cb.endBlockStmt(&p.old)
	cb.current.flows |= (flows &^ flowFlagBreak)

	body := &ast.BlockStmt{List: stmts}
	var assign ast.Stmt
	x := &ast.TypeAssertExpr{X: p.x}
	if p.name != "" {
		assign = &ast.AssignStmt{
			Tok: token.DEFINE,
			Lhs: []ast.Expr{ident(p.name)},
			Rhs: []ast.Expr{x},
		}
	} else {
		assign = &ast.ExprStmt{X: x}
	}
	cb.emitStmt(&ast.TypeSwitchStmt{Init: p.init, Assign: assign, Body: body})
}

type typeCaseStmt struct {
	list []ast.Expr
	old  codeBlockCtx
}

func (p *typeCaseStmt) End(cb *CodeBuilder, src ast.Node) {
	body, flows := cb.endBlockStmt(&p.old)
	cb.current.flows |= flows
	cb.emitStmt(&ast.CaseClause{List: p.list, Body: body})
}

// ----------------------------------------------------------------------------

type loopBodyHandler struct {
	handle func(body *ast.BlockStmt, kind int)
}

func (p *loopBodyHandler) SetBodyHandler(handle func(body *ast.BlockStmt, kind int)) {
	p.handle = handle
}

func (p *loopBodyHandler) handleFor(body *ast.BlockStmt, kind int) *ast.BlockStmt {
	if p.handle != nil {
		p.handle(body, kind)
	}
	return body
}

func InsertStmtFront(body *ast.BlockStmt, stmt ast.Stmt) {
	list := append(body.List, nil)
	copy(list[1:], list)
	list[0] = stmt
	body.List = list
}

// ----------------------------------------------------------------------------
//
// for init; cond then
//
//	body
//	post
//
// end
type forStmt struct {
	init ast.Stmt
	cond ast.Expr
	body *ast.BlockStmt
	old  codeBlockCtx
	old2 codeBlockCtx
	loopBodyHandler
}

func (p *forStmt) Then(cb *CodeBuilder, src ...ast.Node) {
	cond := cb.stk.Pop()
	if cond.Val != nil {
		if !types.AssignableTo(cond.Type, types.Typ[types.Bool]) {
			panic("TODO: for statement condition is not a boolean expr")
		}
		p.cond = cond.Val
	}
	switch stmts := cb.clearBlockStmt(); len(stmts) {
	case 0:
		// nothing to do
	case 1:
		p.init = stmts[0]
	default:
		panic("TODO: for condition has too many init statements")
	}
	cb.startBlockStmt(p, src, "for body", &p.old2)
}

func (p *forStmt) Post(cb *CodeBuilder) {
	stmts, flows := cb.endBlockStmt(&p.old2)
	cb.current.flows |= (flows &^ (flowFlagBreak | flowFlagContinue))
	p.body = &ast.BlockStmt{List: stmts}
}

func (p *forStmt) End(cb *CodeBuilder, src ast.Node) {
	var post ast.Stmt
	if p.body != nil { // has post stmt
		stmts, _ := cb.endBlockStmt(&p.old)
		if len(stmts) != 1 {
			panic("TODO: too many post statements")
		}
		post = stmts[0]
	} else { // no post
		stmts, flows := cb.endBlockStmt(&p.old2)
		cb.current.flows |= (flows &^ (flowFlagBreak | flowFlagContinue))
		p.body = &ast.BlockStmt{List: stmts}
		cb.endBlockStmt(&p.old)
	}
	cb.emitStmt(&ast.ForStmt{
		Init: p.init, Cond: p.cond, Post: post, Body: p.handleFor(p.body, 0),
	})
}

// ----------------------------------------------------------------------------
//
// forRange names... exprX rangeAssignThen
//
//	...
//
// end
//
// forRange exprKey exprVal exprX rangeAssignThen
//
//	...
//
// end
type forRangeStmt struct {
	names []string
	stmt  *ast.RangeStmt
	x     *internal.Elem
	old   codeBlockCtx
	kvt   []types.Type
	udt   int // 0: non-udt, 2: (elem,ok), 3: (key,elem,ok)
	loopBodyHandler
}

func (p *forRangeStmt) RangeAssignThen(cb *CodeBuilder, pos token.Pos) {
	if names := p.names; names != nil { // for k, v := range XXX {
		var val ast.Expr
		switch len(names) {
		case 1:
		case 2:
			val = ident(names[1])
		default:
			cb.panicCodeError(pos, "too many variables in range")
		}
		x := cb.stk.Pop()
		pkg, scope := cb.pkg, cb.current.scope
		typs := p.getKeyValTypes(cb, x.Type)
		if typs == nil {
			src, _ := cb.loadExpr(x.Src)
			cb.panicCodeErrorf(pos, "cannot range over %v (type %v)", src, x.Type)
		}
		if typs[1] == nil { // chan
			if names[0] == "_" && len(names) > 1 {
				names[0], val = names[1], nil
				names = names[:1]
			}
		}
		for i, name := range names {
			if name == "_" {
				continue
			}
			if scope.Insert(types.NewVar(token.NoPos, pkg.Types, name, typs[i])) != nil {
				log.Panicln("TODO: variable already defined -", name)
			}
		}
		if p.udt != 0 {
			p.x = x
		}
		p.stmt = &ast.RangeStmt{
			Key:   ident(names[0]),
			Value: val,
			Tok:   token.DEFINE,
			X:     x.Val,
		}
	} else { // for k, v = range XXX {
		var key, val, x internal.Elem
		n := cb.stk.Len() - cb.current.base
		args := cb.stk.GetArgs(n)
		switch n {
		case 1:
			x = *args[0]
		case 2:
			key, x = *args[0], *args[1]
		case 3:
			key, val, x = *args[0], *args[1], *args[2]
		default:
			cb.panicCodeError(pos, "too many variables in range")
		}
		cb.stk.PopN(n)
		typs := p.getKeyValTypes(cb, x.Type)
		if typs == nil {
			src, _ := cb.loadExpr(x.Src)
			cb.panicCodeErrorf(pos, "cannot range over %v (type %v)", src, x.Type)
		}
		if p.udt != 0 {
			p.x = &x
		}
		p.stmt = &ast.RangeStmt{
			Key:   key.Val,
			Value: val.Val,
			X:     x.Val,
		}
		if n > 1 {
			p.stmt.Tok = token.ASSIGN
			checkAssign(cb.pkg, &key, typs[0], "range")
			if val.Val != nil {
				checkAssign(cb.pkg, &val, typs[1], "range")
			}
		}
	}
	p.stmt.For = pos
}

func (p *forRangeStmt) getKeyValTypes(cb *CodeBuilder, typ types.Type) []types.Type {
retry:
	switch t := typ.(type) {
	case *types.Slice:
		return []types.Type{types.Typ[types.Int], t.Elem()}
	case *types.Map:
		return []types.Type{t.Key(), t.Elem()}
	case *types.Array:
		return []types.Type{types.Typ[types.Int], t.Elem()}
	case *types.Pointer:
		switch e := t.Elem().(type) {
		case *types.Array:
			return []types.Type{types.Typ[types.Int], e.Elem()}
		case *types.Named:
			if kv, ok := p.checkUdt(cb, e); ok {
				return kv
			}
		}
	case *types.Chan:
		return []types.Type{t.Elem(), nil}
	case *types.Basic:
		if (t.Info() & types.IsString) != 0 {
			return []types.Type{types.Typ[types.Int], types.Typ[types.Rune]}
		}
	case *types.Named:
		if kv, ok := p.checkUdt(cb, t); ok {
			return kv
		}
		typ = cb.getUnderlying(t)
		goto retry
	}
	return nil
}

func (p *forRangeStmt) checkUdt(cb *CodeBuilder, o *types.Named) ([]types.Type, bool) {
	if sig := findMethodType(cb, o, nameGopEnum); sig != nil {
		enumRet := sig.Results()
		params := sig.Params()
		switch params.Len() {
		case 0:
			// iter := obj.Gop_Enum()
			// key, val, ok := iter.Next()
		case 1:
			// obj.Gop_Enum(func(key K, val V) { ... })
			if enumRet.Len() != 0 {
				return nil, false
			}
			typ := params.At(0).Type()
			if t, ok := typ.(*types.Signature); ok && t.Results().Len() == 0 {
				kv := t.Params()
				n := kv.Len()
				if n > 0 {
					p.kvt = []types.Type{kv.At(0).Type(), nil}
					if n > 1 {
						n = 2
						p.kvt[1] = kv.At(1).Type()
					}
					p.udt = -n
					return p.kvt, true
				}
			}
			fallthrough
		default:
			return nil, false
		}
		if enumRet.Len() == 1 {
			typ := enumRet.At(0).Type()
			if t, ok := typ.(*types.Pointer); ok {
				typ = t.Elem()
			}
			if it, ok := typ.(*types.Named); ok {
				if next := findMethodType(cb, it, "Next"); next != nil {
					ret := next.Results()
					typs := make([]types.Type, 2)
					n := ret.Len()
					switch n {
					case 2: // elem, ok
						typs[0] = ret.At(0).Type()
					case 3: // key, elem, ok
						typs[0], typs[1] = ret.At(0).Type(), ret.At(1).Type()
					default:
						return nil, false
					}
					if ret.At(n-1).Type() == types.Typ[types.Bool] {
						p.udt = n
						return typs, true
					}
				}
			}
		}
	}
	return nil, false
}

func findMethodType(cb *CodeBuilder, o *types.Named, name string) mthdSignature {
	for i, n := 0, o.NumMethods(); i < n; i++ {
		method := o.Method(i)
		if method.Name() == name {
			return method.Type().(*types.Signature)
		}
	}
	if bti := cb.getBuiltinTI(o); bti != nil {
		return bti.lookupByName(name)
	}
	return nil
}

const (
	cantUseFlows = "can't use return/continue/break/goto in for range of udt.Gop_Enum(callback)"
)

func (p *forRangeStmt) End(cb *CodeBuilder, src ast.Node) {
	if p.stmt == nil {
		return
	}
	stmts, flows := cb.endBlockStmt(&p.old)
	cb.current.flows |= (flows &^ (flowFlagBreak | flowFlagContinue))
	if n := p.udt; n == 0 {
		p.stmt.Body = p.handleFor(&ast.BlockStmt{List: stmts}, 1)
		cb.emitStmt(p.stmt)
	} else if n > 0 {
		cb.stk.Push(p.x)
		cb.MemberVal(nameGopEnum).Call(0)
		callEnum := cb.stk.Pop().Val
		/*
			for _gop_it := X.Gop_Enum();; {
				var _gop_ok bool
				k, v, _gop_ok = _gop_it.Next()
				if !_gop_ok {
					break
				}
				...
			}
		*/
		lhs := make([]ast.Expr, n)
		lhs[0] = p.stmt.Key
		lhs[1] = p.stmt.Value
		lhs[n-1] = identGopOk
		if lhs[0] == nil { // bugfix: for range udt { ... }
			lhs[0] = underscore
			if p.stmt.Tok == token.ILLEGAL {
				p.stmt.Tok = token.ASSIGN
			}
		} else {
			// for _ = range udt { ... }
			if ki, ok := lhs[0].(*ast.Ident); ok && ki.Name == "_" {
				if n == 2 {
					p.stmt.Tok = token.ASSIGN
				} else {
					// for _ , _ = range udt { ... }
					if vi, ok := lhs[1].(*ast.Ident); ok && vi.Name == "_" {
						p.stmt.Tok = token.ASSIGN
					}
				}
			}
		}
		body := make([]ast.Stmt, len(stmts)+3)
		body[0] = stmtGopOkDecl
		body[1] = &ast.AssignStmt{Lhs: lhs, Tok: p.stmt.Tok, Rhs: exprIterNext}
		body[2] = stmtBreakIfNotGopOk
		copy(body[3:], stmts)
		stmt := &ast.ForStmt{
			Init: &ast.AssignStmt{
				Lhs: []ast.Expr{identGopIt},
				Tok: token.DEFINE,
				Rhs: []ast.Expr{callEnum},
			},
			Body: p.handleFor(&ast.BlockStmt{List: body}, 2),
		}
		cb.emitStmt(stmt)
	} else {
		/*
			X.Gop_Enum(func(k K, v V) {
				...
			})
		*/
		if flows != 0 {
			cb.panicCodeError(p.stmt.For, cantUseFlows)
		}
		n = -n
		def := p.stmt.Tok == token.DEFINE
		args := make([]*ast.Field, n)
		if def {
			args[0] = &ast.Field{
				Names: []*ast.Ident{p.stmt.Key.(*ast.Ident)},
				Type:  toType(cb.pkg, p.kvt[0]),
			}
			if n > 1 {
				args[1] = &ast.Field{
					Names: []*ast.Ident{p.stmt.Value.(*ast.Ident)},
					Type:  toType(cb.pkg, p.kvt[1]),
				}
			}
		} else {
			panic("TODO: for range udt assign")
		}
		stmt := &ast.ExprStmt{
			X: &ast.CallExpr{
				Fun: &ast.SelectorExpr{X: p.stmt.X, Sel: identGopEnum},
				Args: []ast.Expr{
					&ast.FuncLit{
						Type: &ast.FuncType{Params: &ast.FieldList{List: args}},
						Body: p.handleFor(&ast.BlockStmt{List: stmts}, -1),
					},
				},
			},
		}
		cb.emitStmt(stmt)
	}
}

var (
	nameGopEnum  = "Gop_Enum"
	identGopOk   = ident("_gop_ok")
	identGopIt   = ident("_gop_it")
	identGopEnum = ident(nameGopEnum)
)

var (
	stmtGopOkDecl = &ast.DeclStmt{
		Decl: &ast.GenDecl{
			Tok: token.VAR,
			Specs: []ast.Spec{
				&ast.ValueSpec{Names: []*ast.Ident{identGopOk}, Type: ident("bool")},
			},
		},
	}
	stmtBreakIfNotGopOk = &ast.IfStmt{
		Cond: &ast.UnaryExpr{Op: token.NOT, X: identGopOk},
		Body: &ast.BlockStmt{List: []ast.Stmt{&ast.BranchStmt{Tok: token.BREAK}}},
	}
	exprIterNext = []ast.Expr{
		&ast.CallExpr{
			Fun: &ast.SelectorExpr{X: identGopIt, Sel: ident("Next")},
		},
	}
)

// ----------------------------------------------------------------------------