asm-calc/codegen/codegen.go

// Package codegen compiles an expression tree to x86-64 assembly.
package codegen

import (
	"fmt"
	"sort"

	"github.com/zg/asm-calc/lexer"
	"github.com/zg/asm-calc/parser"
	"github.com/zg/asm-calc/x64"
)

// Compiler turns expression trees into instructions.
type Compiler struct {
	prog x64.Program
	// vars maps a variable name to its frame offset. Variables live on the
	// stack frame rather than in registers, because the number of them is
	// unbounded while registers are not.
	vars       map[string]int
	nextOffset int
}

// New returns a compiler with no variables defined.
func New() *Compiler {
	return &Compiler{vars: map[string]int{}, nextOffset: -8}
}

// Variables returns the defined variable names, sorted.
func (c *Compiler) Variables() []string {
	out := make([]string, 0, len(c.vars))
	for k := range c.vars {
		out = append(out, k)
	}
	sort.Strings(out)
	return out
}

// Offset returns a variable's frame offset.
func (c *Compiler) Offset(name string) (int, bool) {
	off, ok := c.vars[name]
	return off, ok
}

// Compile emits code leaving the expression's value in RAX.
func (c *Compiler) Compile(node parser.Node) (x64.Program, error) {
	c.prog = nil
	if err := c.emit(node); err != nil {
		return nil, err
	}
	return c.prog, nil
}

func (c *Compiler) add(op x64.Op, comment string, operands ...x64.Operand) {
	c.prog = append(c.prog, x64.Instruction{Op: op, Operands: operands, Comment: comment})
}

func (c *Compiler) emit(node parser.Node) error {
	switch n := node.(type) {
	case *parser.Number:
		c.add(x64.MOV, fmt.Sprintf("literal %d", n.Value), x64.R(x64.RAX), x64.I(n.Value))
		return nil

	case *parser.Variable:
		off, ok := c.vars[n.Name]
		if !ok {
			return fmt.Errorf("undefined variable %q", n.Name)
		}
		c.add(x64.MOV, "load "+n.Name, x64.R(x64.RAX), x64.M(x64.RBP, off))
		return nil

	case *parser.Assignment:
		if err := c.emit(n.Value); err != nil {
			return err
		}
		off, ok := c.vars[n.Name]
		if !ok {
			off = c.nextOffset
			c.vars[n.Name] = off
			c.nextOffset -= 8
		}
		c.add(x64.MOV, "store "+n.Name, x64.M(x64.RBP, off), x64.R(x64.RAX))
		return nil

	case *parser.UnaryOp:
		if err := c.emit(n.Operand); err != nil {
			return err
		}
		if n.Op == lexer.MINUS {
			c.add(x64.NEG, "negate", x64.R(x64.RAX))
			return nil
		}
		return fmt.Errorf("unsupported unary operator %s", n.Op)

	case *parser.BinaryOp:
		return c.emitBinary(n)
	}
	return fmt.Errorf("unsupported node %T", node)
}

// emitBinary compiles a binary operation.
//
// The strategy is: evaluate the left side into RAX, push it, evaluate the
// right side into RAX, then pop the left into RCX. Using the stack rather
// than a second register means arbitrarily deep expressions compile without
// running out of registers — at the cost of a push and pop per operation.
func (c *Compiler) emitBinary(n *parser.BinaryOp) error {
	if n.Op == lexer.CARET {
		return c.emitPower(n)
	}

	if err := c.emit(n.Left); err != nil {
		return err
	}
	c.add(x64.PUSH, "save left operand", x64.R(x64.RAX))

	if err := c.emit(n.Right); err != nil {
		return err
	}
	// The right operand is in RAX and the left is on the stack, so the
	// operands are the wrong way round for a non-commutative operation.
	// Move the right out of the way, then pop the left into RAX.
	c.add(x64.MOV, "right operand", x64.R(x64.RCX), x64.R(x64.RAX))
	c.add(x64.POP, "restore left operand", x64.R(x64.RAX))

	switch n.Op {
	case lexer.PLUS:
		c.add(x64.ADD, "add", x64.R(x64.RAX), x64.R(x64.RCX))
	case lexer.MINUS:
		c.add(x64.SUB, "subtract", x64.R(x64.RAX), x64.R(x64.RCX))
	case lexer.STAR:
		c.add(x64.IMUL, "multiply", x64.R(x64.RAX), x64.R(x64.RCX))
	case lexer.SLASH:
		// idiv divides the 128-bit RDX:RAX by its operand, so RDX must
		// hold the sign extension of RAX first. Omitting cqo gives wrong
		// answers for negative dividends rather than an obvious failure.
		c.add(x64.CQO, "sign-extend rax into rdx:rax", nil...)
		c.add(x64.IDIV, "divide, quotient in rax", x64.R(x64.RCX))
	case lexer.PERCENT:
		c.add(x64.CQO, "sign-extend rax into rdx:rax", nil...)
		c.add(x64.IDIV, "divide, remainder in rdx", x64.R(x64.RCX))
		c.add(x64.MOV, "take the remainder", x64.R(x64.RAX), x64.R(x64.RDX))
	default:
		return fmt.Errorf("unsupported binary operator %s", n.Op)
	}
	return nil
}

// constFold evaluates a node at compile time when it contains nothing but
// literals and arithmetic, reporting false when it depends on a variable.
//
// This is what lets 2^(3^2) and 2^(1+2) compile: the exponent must be known
// to unroll the multiplication, but it need not be written as a literal.
func constFold(node parser.Node) (int64, bool) {
	switch n := node.(type) {
	case *parser.Number:
		return n.Value, true

	case *parser.UnaryOp:
		v, ok := constFold(n.Operand)
		if !ok {
			return 0, false
		}
		if n.Op == lexer.MINUS {
			return -v, true
		}
		return 0, false

	case *parser.BinaryOp:
		l, ok := constFold(n.Left)
		if !ok {
			return 0, false
		}
		r, ok := constFold(n.Right)
		if !ok {
			return 0, false
		}
		switch n.Op {
		case lexer.PLUS:
			return l + r, true
		case lexer.MINUS:
			return l - r, true
		case lexer.STAR:
			return l * r, true
		case lexer.SLASH:
			// A constant division by zero is a compile-time error, not a
			// value; refuse to fold it rather than panicking here.
			if r == 0 {
				return 0, false
			}
			return l / r, true
		case lexer.PERCENT:
			if r == 0 {
				return 0, false
			}
			return l % r, true
		case lexer.CARET:
			if r < 0 || r > 64 {
				return 0, false
			}
			result := int64(1)
			for range r {
				result *= l
			}
			return result, true
		}
	}
	return 0, false
}

// emitPower expands x^n by repeated multiplication. The exponent must be
// known at compile time, since the instruction subset has no jumps to build
// a loop from.
func (c *Compiler) emitPower(n *parser.BinaryOp) error {
	expValue, ok := constFold(n.Right)
	if !ok {
		return fmt.Errorf("exponent must be a constant expression")
	}
	exp := &parser.Number{Value: expValue}
	if exp.Value < 0 {
		// Integer arithmetic has no fractions, so a negative exponent
		// would silently truncate to zero.
		return fmt.Errorf("negative exponent %d has no integer result", exp.Value)
	}
	if exp.Value > 64 {
		return fmt.Errorf("exponent %d is too large", exp.Value)
	}

	if exp.Value == 0 {
		c.add(x64.MOV, "anything to the power 0 is 1", x64.R(x64.RAX), x64.I(1))
		return nil
	}

	if err := c.emit(n.Left); err != nil {
		return err
	}
	if exp.Value == 1 {
		return nil
	}
	// Keep the base in RCX and multiply RAX by it repeatedly. Unrolling
	// avoids needing jumps, which this instruction subset does not have.
	c.add(x64.MOV, "keep the base", x64.R(x64.RCX), x64.R(x64.RAX))
	for range exp.Value - 1 {
		c.add(x64.IMUL, "raise to the power", x64.R(x64.RAX), x64.R(x64.RCX))
	}
	return nil
}

// Assemble wraps a program in a complete, assemblable NASM source file with
// a _start entry point that exits with the result.
func Assemble(prog x64.Program, frameSize int) string {
	// Round the frame up to a multiple of 16: the ABI requires the stack
	// stay 16-byte aligned, and an unaligned call crashes on any function
	// using SSE.
	if frameSize%16 != 0 {
		frameSize += 16 - frameSize%16
	}

	var sb = &sbuilder{}
	sb.line("; generated by asm-calc")
	sb.line("bits 64")
	sb.line("default rel")
	sb.line("")
	sb.line("section .text")
	sb.line("global _start")
	sb.line("")
	sb.line("_start:")
	sb.line("    push rbp")
	sb.line("    mov rbp, rsp")
	if frameSize > 0 {
		sb.line(fmt.Sprintf("    sub rsp, %d", frameSize))
	}
	sb.line("")
	sb.raw(prog.String())
	sb.line("")
	sb.line("    ; exit(rax & 0xff)")
	sb.line("    mov rdi, rax")
	sb.line("    mov rax, 60")
	sb.line("    syscall")
	return sb.String()
}

type sbuilder struct{ parts []string }

func (s *sbuilder) line(l string) { s.parts = append(s.parts, l+"\n") }
func (s *sbuilder) raw(r string)  { s.parts = append(s.parts, r) }
func (s *sbuilder) String() string {
	out := ""
	for _, p := range s.parts {
		out += p
	}
	return out
}

// FrameSize returns the bytes of stack the compiled variables need.
func (c *Compiler) FrameSize() int { return len(c.vars) * 8 }