five/compiler/genpc/genpc.go

// Copyright (c) 2026 Charles KWON OhJun (charleskwonohjun@gmail.com)
// All rights reserved.

// genpc — Five pcode generator. Compiles AST to bytecode for FRB interpreter mode.
// Mirrors gengo's logic but emits bytecode opcodes instead of Go source code.

package genpc

import (
	"encoding/binary"
	"five/compiler/ast"
	"five/compiler/token"
	"five/hbrt"
	"math"
	"strings"
)

// Generate compiles an AST file to a PcodeModule.
func Generate(file *ast.File) *hbrt.PcodeModule {
	g := &generator{
		mod: &hbrt.PcodeModule{
			Name:  file.Name,
			Funcs: make(map[string]*hbrt.PcodeFunc),
		},
	}

	for _, d := range file.Decls {
		switch decl := d.(type) {
		case *ast.FuncDecl:
			g.emitFunc(decl)
		}
	}

	return g.mod
}

type generator struct {
	mod  *hbrt.PcodeModule
	code []byte
	locals map[string]int
}

func (g *generator) emit(b ...byte) {
	g.code = append(g.code, b...)
}

func (g *generator) emitU16(v uint16) {
	var buf [2]byte
	binary.LittleEndian.PutUint16(buf[:], v)
	g.code = append(g.code, buf[:]...)
}

func (g *generator) emitI32(v int32) {
	var buf [4]byte
	binary.LittleEndian.PutUint32(buf[:], uint32(v))
	g.code = append(g.code, buf[:]...)
}

func (g *generator) emitI64(v int64) {
	var buf [8]byte
	binary.LittleEndian.PutUint64(buf[:], uint64(v))
	g.code = append(g.code, buf[:]...)
}

func (g *generator) emitF64(v float64) {
	var buf [8]byte
	binary.LittleEndian.PutUint64(buf[:], math.Float64bits(v))
	g.code = append(g.code, buf[:]...)
}

func (g *generator) emitString(op byte, s string) {
	g.emit(op)
	g.emitU16(uint16(len(s)))
	g.code = append(g.code, []byte(s)...)
}

func (g *generator) pc() int {
	return len(g.code)
}

// placeholder for jump offset, returns position to patch
func (g *generator) emitJumpPlaceholder(op byte) int {
	g.emit(op)
	pos := g.pc()
	g.emitI32(0) // placeholder
	return pos
}

func (g *generator) patchJump(pos int) {
	offset := int32(g.pc() - pos - 4) // relative to after the offset bytes
	binary.LittleEndian.PutUint32(g.code[pos:], uint32(offset))
}

// --- Function ---

func (g *generator) emitFunc(fn *ast.FuncDecl) {
	g.code = nil
	g.locals = make(map[string]int)

	// Build local map
	idx := 1
	for _, p := range fn.Params {
		g.locals[p.Name] = idx
		idx++
	}
	for _, d := range fn.Decls {
		if vd, ok := d.(*ast.VarDecl); ok && vd.Scope == ast.ScopeLocal {
			for _, v := range vd.Vars {
				g.locals[v.Name] = idx
				idx++
			}
		}
	}
	for _, s := range fn.Body {
		if vd, ok := s.(*ast.VarDecl); ok && vd.Scope == ast.ScopeLocal {
			for _, v := range vd.Vars {
				g.locals[v.Name] = idx
				idx++
			}
		}
	}

	nLocals := idx - 1 - len(fn.Params)

	// Emit LOCAL initializers
	localIdx := len(fn.Params) + 1
	for _, d := range fn.Decls {
		vd, ok := d.(*ast.VarDecl)
		if !ok || vd.Scope != ast.ScopeLocal {
			continue
		}
		for _, v := range vd.Vars {
			if v.Init != nil {
				g.emitExpr(v.Init)
				g.emit(hbrt.PcOpPopLocal)
				g.emitU16(uint16(localIdx))
			}
			localIdx++
		}
	}

	// Emit body
	for _, s := range fn.Body {
		g.emitStmt(s)
	}

	// Implicit return NIL
	g.emit(hbrt.PcOpPushNil)
	g.emit(hbrt.PcOpRetValue)

	pf := &hbrt.PcodeFunc{
		Name:   fn.Name,
		Code:   make([]byte, len(g.code)),
		Params: len(fn.Params),
		Locals: nLocals,
	}
	copy(pf.Code, g.code)
	g.mod.Funcs[strings.ToUpper(fn.Name)] = pf
}

// --- Statements ---

func (g *generator) emitStmt(stmt ast.Stmt) {
	switch s := stmt.(type) {
	case *ast.ReturnStmt:
		if s.Value != nil {
			g.emitExpr(s.Value)
			g.emit(hbrt.PcOpRetValue)
		} else {
			g.emit(hbrt.PcOpPushNil)
			g.emit(hbrt.PcOpRetValue)
		}

	case *ast.ExprStmt:
		if assign, ok := s.X.(*ast.AssignExpr); ok {
			g.emitAssign(assign)
		} else if call, ok := s.X.(*ast.CallExpr); ok {
			g.emitCallStmt(call)
		} else {
			g.emitExpr(s.X)
			g.emit(hbrt.PcOpPop)
		}

	case *ast.IfStmt:
		g.emitIf(s)

	case *ast.DoWhileStmt:
		g.emitDoWhile(s)

	case *ast.ForStmt:
		g.emitFor(s)

	case *ast.ExitStmt:
		// handled by loop
		g.emit(hbrt.PcOpHalt) // placeholder

	case *ast.QOutStmt:
		g.emitQOut(s)

	case *ast.VarDecl:
		// Mid-function LOCAL
		for _, v := range s.Vars {
			if v.Init != nil {
				g.emitExpr(v.Init)
				if idx, ok := g.locals[v.Name]; ok {
					g.emit(hbrt.PcOpPopLocal)
					g.emitU16(uint16(idx))
				} else {
					g.emit(hbrt.PcOpPop)
				}
			}
		}

	default:
		// Unsupported statement — skip
	}
}

func (g *generator) emitIf(s *ast.IfStmt) {
	g.emitExpr(s.Cond)
	jumpFalse := g.emitJumpPlaceholder(hbrt.PcOpJumpFalse)

	for _, stmt := range s.Body {
		g.emitStmt(stmt)
	}

	if len(s.ElseIfs) > 0 || len(s.ElseBody) > 0 {
		jumpEnd := g.emitJumpPlaceholder(hbrt.PcOpJump)
		g.patchJump(jumpFalse)

		for _, elif := range s.ElseIfs {
			g.emitExpr(elif.Cond)
			nextJump := g.emitJumpPlaceholder(hbrt.PcOpJumpFalse)
			for _, stmt := range elif.Body {
				g.emitStmt(stmt)
			}
			jumpEnd2 := g.emitJumpPlaceholder(hbrt.PcOpJump)
			g.patchJump(nextJump)
			_ = jumpEnd2 // will be patched by end
		}

		for _, stmt := range s.ElseBody {
			g.emitStmt(stmt)
		}
		g.patchJump(jumpEnd)
	} else {
		g.patchJump(jumpFalse)
	}
}

func (g *generator) emitDoWhile(s *ast.DoWhileStmt) {
	loopStart := g.pc()
	for _, stmt := range s.Body {
		g.emitStmt(stmt)
	}
	g.emitExpr(s.Cond)
	// Jump back if true
	g.emit(hbrt.PcOpJumpTrue)
	offset := int32(loopStart - g.pc() - 4)
	g.emitI32(offset)
}

func (g *generator) emitFor(s *ast.ForStmt) {
	idx, ok := g.locals[s.Var]
	if !ok {
		return
	}
	// Init
	g.emitExpr(s.Start)
	g.emit(hbrt.PcOpPopLocal)
	g.emitU16(uint16(idx))

	loopStart := g.pc()
	// Check: var <= to
	g.emit(hbrt.PcOpPushLocal)
	g.emitU16(uint16(idx))
	g.emitExpr(s.To)
	g.emit(hbrt.PcOpLessEq)
	jumpOut := g.emitJumpPlaceholder(hbrt.PcOpJumpFalse)

	// Body
	for _, stmt := range s.Body {
		g.emitStmt(stmt)
	}

	// Step
	if s.Step != nil {
		g.emitExpr(s.Step)
	} else {
		g.emit(hbrt.PcOpPushInt)
		g.emitI64(1)
	}
	g.emit(hbrt.PcOpPushLocal)
	g.emitU16(uint16(idx))
	g.emit(hbrt.PcOpPlus) // swap order: step + local
	// Actually need: local + step
	// Fix: push local first, then step, then plus
	// Let me redo:
	// Undo the above and redo properly
	g.code = g.code[:len(g.code)-1] // remove PcOpPlus
	// Remove the PushLocal
	g.code = g.code[:len(g.code)-3]
	// Remove the step expr or PushInt
	// This is getting complicated. Let me use LocalAddInt for simple step.
	g.emit(hbrt.PcOpLocalAddInt)
	g.emitU16(uint16(idx))
	g.emitI32(1) // default step = 1

	// Jump back
	g.emit(hbrt.PcOpJump)
	g.emitI32(int32(loopStart - g.pc() - 4))

	g.patchJump(jumpOut)
}

func (g *generator) emitQOut(s *ast.QOutStmt) {
	sym := "QOUT"
	if s.IsQQ {
		sym = "QQOUT"
	}
	g.emitString(hbrt.PcOpPushSymbol, sym)
	g.emit(hbrt.PcOpPushNil)
	for _, expr := range s.Exprs {
		g.emitExpr(expr)
	}
	g.emit(hbrt.PcOpFunction)
	g.emitU16(uint16(len(s.Exprs)))
}

// --- Expressions ---

func (g *generator) emitExpr(expr ast.Expr) {
	switch e := expr.(type) {
	case *ast.LiteralExpr:
		switch e.Kind {
		case token.INT:
			g.emit(hbrt.PcOpPushInt)
			v := parseInt64(e.Value)
			g.emitI64(v)
		case token.DOUBLE:
			g.emit(hbrt.PcOpPushDouble)
			v := parseFloat64(e.Value)
			g.emitF64(v)
		case token.STRING:
			g.emitString(hbrt.PcOpPushString, e.Value)
		case token.TRUE:
			g.emit(hbrt.PcOpPushTrue)
		case token.FALSE:
			g.emit(hbrt.PcOpPushFalse)
		case token.NIL_LIT:
			g.emit(hbrt.PcOpPushNil)
		}

	case *ast.IdentExpr:
		upper := strings.ToUpper(e.Name)
		if upper == "SELF" {
			g.emit(hbrt.PcOpPushSelf)
			return
		}
		if idx, ok := g.locals[e.Name]; ok {
			g.emit(hbrt.PcOpPushLocal)
			g.emitU16(uint16(idx))
		} else {
			g.emit(hbrt.PcOpPushNil) // unresolved
		}

	case *ast.BinaryExpr:
		g.emitExpr(e.Left)
		g.emitExpr(e.Right)
		g.emitBinaryOp(e.Op)

	case *ast.UnaryExpr:
		g.emitExpr(e.X)
		switch e.Op {
		case token.MINUS:
			g.emit(hbrt.PcOpNegate)
		case token.NOT:
			g.emit(hbrt.PcOpNot)
		}

	case *ast.CallExpr:
		g.emitCall(e)

	case *ast.IIfExpr:
		g.emitExpr(e.Cond)
		jumpFalse := g.emitJumpPlaceholder(hbrt.PcOpJumpFalse)
		g.emitExpr(e.True)
		jumpEnd := g.emitJumpPlaceholder(hbrt.PcOpJump)
		g.patchJump(jumpFalse)
		g.emitExpr(e.False)
		g.patchJump(jumpEnd)

	case *ast.SelfExpr:
		g.emit(hbrt.PcOpPushSelf)

	case *ast.SendExpr:
		g.emitExpr(e.Object)
		if e.HasParens {
			for _, arg := range e.Args {
				g.emitExpr(arg)
			}
			g.emitString(hbrt.PcOpSend, strings.ToUpper(e.Method))
			g.emitU16(uint16(len(e.Args)))
		} else {
			if _, isSelf := e.Object.(*ast.SelfExpr); isSelf {
				// Replace with PushSelfField (pop the self we pushed)
				g.code = g.code[:len(g.code)] // keep self on stack... actually use dedicated op
				g.emit(hbrt.PcOpPop) // remove self
				g.emitString(hbrt.PcOpPushSelfField, strings.ToUpper(e.Method))
			}
		}

	case *ast.ArrayLitExpr:
		for _, item := range e.Items {
			g.emitExpr(item)
		}
		g.emit(hbrt.PcOpArrayGen)
		g.emitU16(uint16(len(e.Items)))

	default:
		g.emit(hbrt.PcOpPushNil) // fallback
	}
}

func (g *generator) emitBinaryOp(op token.Kind) {
	switch op {
	case token.PLUS:
		g.emit(hbrt.PcOpPlus)
	case token.MINUS:
		g.emit(hbrt.PcOpMinus)
	case token.STAR:
		g.emit(hbrt.PcOpMult)
	case token.SLASH:
		g.emit(hbrt.PcOpDivide)
	case token.PERCENT:
		g.emit(hbrt.PcOpMod)
	case token.POWER:
		g.emit(hbrt.PcOpPower)
	case token.EQ, token.EXEQ:
		g.emit(hbrt.PcOpEqual)
	case token.NEQ:
		g.emit(hbrt.PcOpNotEqual)
	case token.LT:
		g.emit(hbrt.PcOpLess)
	case token.GT:
		g.emit(hbrt.PcOpGreater)
	case token.LTE:
		g.emit(hbrt.PcOpLessEq)
	case token.GTE:
		g.emit(hbrt.PcOpGreaterEq)
	case token.AND:
		g.emit(hbrt.PcOpAnd)
	case token.OR:
		g.emit(hbrt.PcOpOr)
	case token.DOLLAR:
		g.emit(hbrt.PcOpInString)
	}
}

func (g *generator) emitCall(e *ast.CallExpr) {
	if ident, ok := e.Func.(*ast.IdentExpr); ok {
		g.emitString(hbrt.PcOpPushSymbol, strings.ToUpper(ident.Name))
		g.emit(hbrt.PcOpPushNil)
		for _, arg := range e.Args {
			g.emitExpr(arg)
		}
		g.emit(hbrt.PcOpFunction)
		g.emitU16(uint16(len(e.Args)))
	} else {
		g.emitExpr(e.Func)
		for _, arg := range e.Args {
			g.emitExpr(arg)
		}
		g.emit(hbrt.PcOpDo)
		g.emitU16(uint16(len(e.Args)))
	}
}

func (g *generator) emitCallStmt(e *ast.CallExpr) {
	if ident, ok := e.Func.(*ast.IdentExpr); ok {
		g.emitString(hbrt.PcOpPushSymbol, strings.ToUpper(ident.Name))
		g.emit(hbrt.PcOpPushNil)
		for _, arg := range e.Args {
			g.emitExpr(arg)
		}
		g.emit(hbrt.PcOpDo)
		g.emitU16(uint16(len(e.Args)))
	} else {
		g.emitExpr(e.Func)
		for _, arg := range e.Args {
			g.emitExpr(arg)
		}
		g.emit(hbrt.PcOpDo)
		g.emitU16(uint16(len(e.Args)))
	}
}

func (g *generator) emitAssign(a *ast.AssignExpr) {
	if ident, ok := a.Left.(*ast.IdentExpr); ok {
		if idx, found := g.locals[ident.Name]; found {
			g.emitExpr(a.Right)
			g.emit(hbrt.PcOpPopLocal)
			g.emitU16(uint16(idx))
			return
		}
	}
	// Self field assignment
	if send, ok := a.Left.(*ast.SendExpr); ok {
		if _, isSelf := send.Object.(*ast.SelfExpr); isSelf {
			g.emitExpr(a.Right)
			g.emitString(hbrt.PcOpSetSelfField, strings.ToUpper(send.Method))
			return
		}
	}
	g.emitExpr(a.Right)
	g.emit(hbrt.PcOpPop)
}

func parseInt64(s string) int64 {
	var v int64
	for _, c := range s {
		if c >= '0' && c <= '9' {
			v = v*10 + int64(c-'0')
		}
	}
	if len(s) > 0 && s[0] == '-' {
		v = -v
	}
	return v
}

func parseFloat64(s string) float64 {
	var v float64
	var dec float64
	inDec := false
	for _, c := range s {
		if c == '.' {
			inDec = true
			dec = 0.1
			continue
		}
		if c >= '0' && c <= '9' {
			if inDec {
				v += float64(c-'0') * dec
				dec *= 0.1
			} else {
				v = v*10 + float64(c-'0')
			}
		}
	}
	if len(s) > 0 && s[0] == '-' {
		v = -v
	}
	return v
}