ast.go

package roll

import (
	"fmt"
	"sort"
	"strconv"
	"strings"
)

// Opcode identifies a bytecode instruction executed by the roll VM.
type Opcode uint8

const (
	// OpRollDice evaluates a single dice term and pushes its result on the VM stack.
	OpRollDice Opcode = iota
	// OpRollGroup aggregates previously computed child results into a grouped result.
	OpRollGroup
)

func (op Opcode) String() string {
	switch op {
	case OpRollDice:
		return "roll_dice"
	case OpRollGroup:
		return "roll_group"
	default:
		return "unknown"
	}
}

// Instruction is a single bytecode operation.
type Instruction struct {
	Op  Opcode
	Arg int
}

// Program is a compiled roll program ready for VM execution.
type Program struct {
	Code       []Instruction
	DiceTerms  []DiceTerm
	GroupTerms []GroupTerm
	Rendered   string
	MaxDepth   int
}

// String returns the normalized notation generated by the compiler.
func (p *Program) String() string {
	if p == nil {
		return ""
	}
	return p.Rendered
}

// DiceTerm captures the semantics of a single dice instruction.
type DiceTerm struct {
	Multiplier int
	Die        Die
	Modifier   int
	Exploding  *ExplodingOp
	Limit      *LimitOp
	Success    *ComparisonOp
	Failure    *ComparisonOp
	Rerolls    []RerollOp
	Sort       SortType
}

// GroupTerm captures the aggregation semantics of a grouped instruction.
type GroupTerm struct {
	Modifier   int
	Limit      *LimitOp
	Success    *ComparisonOp
	Failure    *ComparisonOp
	Combined   bool
	Negative   bool
	ChildCount int
}

// ComparisonType is the type of comparison.
type ComparisonType int

const (
	// Equals matches only values that are equal to the comparison value.
	Equals ComparisonType = iota
	// GreaterThan matches only values greater than the comparison value.
	GreaterThan
	// LessThan matches only values less than the comparison value.
	LessThan
)

// ComparisonOp is the operation that defines how you compare against a roll.
type ComparisonOp struct {
	Type      ComparisonType
	Value     int
	Inclusive bool
}

// Match returns true if the given value compares positively against the op val.
func (op *ComparisonOp) Match(val int) bool {
	switch op.Type {
	case Equals:
		return val == op.Value
	case GreaterThan:
		if op.Inclusive {
			return val >= op.Value
		}
		return val > op.Value
	case LessThan:
		if op.Inclusive {
			return val <= op.Value
		}
		return val < op.Value
	}
	return false
}

// String returns the string representation of the comparison operator.
func (op *ComparisonOp) String() string {
	switch op.Type {
	case Equals:
		return fmt.Sprintf("=%d", op.Value)
	case GreaterThan:
		if op.Inclusive {
			return fmt.Sprintf(">=%d", op.Value)
		}
		return fmt.Sprintf(">%d", op.Value)
	case LessThan:
		if op.Inclusive {
			return fmt.Sprintf("<=%d", op.Value)
		}
		return fmt.Sprintf("<%d", op.Value)
	}
	return ""
}

// ExplodingType is the type of exploding die.
type ExplodingType int

const (
	// Exploding adds new dice for each roll satisfying the exploding condition.
	Exploding ExplodingType = iota
	// Compounded adds to a single new result for each roll.
	Compounded
	// Penetrating is like Exploding, except each die result has a -1 modifier.
	Penetrating
)

// ExplodingOp is the operation that defines how a dice roll explodes.
type ExplodingOp struct {
	*ComparisonOp
	Type ExplodingType
}

// String returns the string representation of the exploding dice operation.
func (e ExplodingOp) String() (output string) {
	switch e.Type {
	case Exploding:
		output = "!"
	case Compounded:
		output = "!!"
	case Penetrating:
		output = "!p"
	}

	return output + strings.TrimPrefix(e.ComparisonOp.String(), "=")
}

// LimitType is the type of roll limitation.
type LimitType int

const (
	// KeepHighest indicated we should keep the highest results.
	KeepHighest LimitType = iota
	// KeepLowest indicated we should keep the lowest results.
	KeepLowest
	// DropHighest indicated we should drop the highest results.
	DropHighest
	// DropLowest indicated we should drop the lowest results.
	DropLowest
)

// LimitOp is the operation that defines how dice roll results are limited.
type LimitOp struct {
	Amount int
	Type   LimitType
}

func (op LimitOp) String() (output string) {
	switch op.Type {
	case KeepHighest:
		output = "kh"
	case KeepLowest:
		output = "kl"
	case DropHighest:
		output = "dh"
	case DropLowest:
		output = "dl"
	}

	if op.Amount > 1 {
		output += strconv.Itoa(op.Amount)
	}

	return
}

// RerollOp is the operation that defines how dice are rerolled.
type RerollOp struct {
	*ComparisonOp
	Once bool
}

// String returns the string representation of the reroll operation.
func (e RerollOp) String() (output string) {
	output = "r"
	if e.Once {
		output = "ro"
	}

	return output + strings.TrimPrefix(e.ComparisonOp.String(), "=")
}

// SortType is the type of sorting to use for dice roll results.
type SortType int

// String return the string representation of a SortType value.
func (t SortType) String() string {
	switch t {
	case Unsorted:
		return ""
	case Ascending:
		return "s"
	case Descending:
		return "sd"
	}

	return ""
}

const (
	// Unsorted doesn't sort dice rolls.
	Unsorted SortType = iota
	// Ascending sorts dice rolls from lowest to highest.
	Ascending
	// Descending sorts dice rolls from highest to lowest.
	Descending
)

// Limits configures compiler and evaluator safety limits.
type Limits struct {
	MaxDieSize     int
	MaxRollsPerDie int
	MaxRollsTotal  int
	MaxEvalDepth   int
}

// DefaultLimits are the package-level defaults used by Parse and ParseString.
var DefaultLimits = Limits{
	MaxDieSize:     1000000,
	MaxRollsPerDie: 1000000,
	MaxRollsTotal:  1000000,
	MaxEvalDepth:   256,
}

func (l Limits) normalized() Limits {
	if l.MaxDieSize <= 0 {
		l.MaxDieSize = DefaultLimits.MaxDieSize
	}
	if l.MaxRollsPerDie <= 0 {
		l.MaxRollsPerDie = DefaultLimits.MaxRollsPerDie
	}
	if l.MaxRollsTotal <= 0 {
		l.MaxRollsTotal = DefaultLimits.MaxRollsTotal
	}
	if l.MaxEvalDepth <= 0 {
		l.MaxEvalDepth = DefaultLimits.MaxEvalDepth
	}
	return l
}

// ErrUnsafeDie is raised when a die configuration is categorically unsafe.
type ErrUnsafeDie string

func (e ErrUnsafeDie) Error() string {
	return fmt.Sprintf("unsafe die type %q", string(e))
}

// ErrLimitExceeded is raised when compiler or evaluator safety limits are exceeded.
type ErrLimitExceeded string

func (e ErrLimitExceeded) Error() string { return string(e) }

type rollContext struct {
	limits     Limits
	totalRolls int
}

func (ctx *rollContext) recordRoll(perDie *int) error {
	(*perDie)++
	if *perDie > ctx.limits.MaxRollsPerDie {
		return ErrLimitExceeded(fmt.Sprintf("die term exceeded maximum roll count of %d", ctx.limits.MaxRollsPerDie))
	}

	ctx.totalRolls++
	if ctx.totalRolls > ctx.limits.MaxRollsTotal {
		return ErrLimitExceeded(fmt.Sprintf("roll exceeded maximum total roll count of %d", ctx.limits.MaxRollsTotal))
	}
	return nil
}

// Result is a collection of die rolls and a count of successes.
type Result struct {
	Results   []DieRoll
	Total     int
	Successes int
}

// Len is the number of results.
func (r *Result) Len() int {
	return len(r.Results)
}

// Less return true if DieRoll at index i is less than the one at index j.
func (r *Result) Less(i, j int) bool {
	return r.Results[i].Result < r.Results[j].Result
}

// Swap swaps the DieRoll at index i with the one at index j.
func (r *Result) Swap(i, j int) {
	r.Results[i], r.Results[j] = r.Results[j], r.Results[i]
}

type vmValue struct {
	Result   Result
	Modifier int
}

// EvaluateProgram executes a compiled roll program using DefaultLimits.
func EvaluateProgram(program *Program) (Result, error) {
	return EvaluateProgramWithLimits(program, DefaultLimits)
}

// EvaluateProgramWithLimits executes a compiled roll program using explicit safety limits.
func EvaluateProgramWithLimits(program *Program, limits Limits) (Result, error) {
	if program == nil {
		return Result{}, nil
	}

	limits = limits.normalized()
	if program.MaxDepth > limits.MaxEvalDepth {
		return Result{}, ErrLimitExceeded(fmt.Sprintf("roll exceeded maximum evaluation depth of %d", limits.MaxEvalDepth))
	}

	ctx := &rollContext{limits: limits}
	stack := make([]vmValue, 0, len(program.Code))

	for _, instruction := range program.Code {
		switch instruction.Op {
		case OpRollDice:
			if instruction.Arg < 0 || instruction.Arg >= len(program.DiceTerms) {
				return Result{}, fmt.Errorf("invalid dice term index %d", instruction.Arg)
			}
			term := program.DiceTerms[instruction.Arg]
			result, err := evalDiceTerm(ctx, term)
			if err != nil {
				return Result{}, err
			}
			stack = append(stack, vmValue{Result: result, Modifier: term.Modifier})
		case OpRollGroup:
			if instruction.Arg < 0 || instruction.Arg >= len(program.GroupTerms) {
				return Result{}, fmt.Errorf("invalid group term index %d", instruction.Arg)
			}
			term := program.GroupTerms[instruction.Arg]
			if term.ChildCount > len(stack) {
				return Result{}, fmt.Errorf("group term %d requires %d child values, stack has %d", instruction.Arg, term.ChildCount, len(stack))
			}
			children := append([]vmValue(nil), stack[len(stack)-term.ChildCount:]...)
			stack = stack[:len(stack)-term.ChildCount]
			result := evalGroupTerm(term, children)
			stack = append(stack, vmValue{Result: result, Modifier: term.Modifier})
		default:
			return Result{}, fmt.Errorf("unsupported opcode %d", instruction.Op)
		}
	}

	if len(stack) != 1 {
		return Result{}, fmt.Errorf("program left %d results on the VM stack", len(stack))
	}

	return stack[0].Result, nil
}

func evalDiceTerm(ctx *rollContext, term DiceTerm) (result Result, err error) {
	if err = validateDieLimits(term.Die, ctx.limits); err != nil {
		return
	}

	if term.Multiplier == 0 {
		return
	}

	totalMultiplier := 1
	if term.Multiplier < 0 {
		totalMultiplier = -1
	}

	rollCount := term.Multiplier * totalMultiplier
	if rollCount > ctx.limits.MaxRollsPerDie {
		return Result{}, ErrLimitExceeded(fmt.Sprintf("die term exceeded maximum roll count of %d", ctx.limits.MaxRollsPerDie))
	}

	dieRolls := 0
	for i := 0; i < rollCount; i++ {
		if err = ctx.recordRoll(&dieRolls); err != nil {
			return Result{}, err
		}
		result.Results = append(result.Results, term.Die.Roll())
	}

	for i, roll := range result.Results {
	RerollOnce:
		for _, reroll := range term.Rerolls {
			for reroll.Match(roll.Result) {
				if err = ctx.recordRoll(&dieRolls); err != nil {
					return Result{}, err
				}
				roll = term.Die.Roll()
				result.Results[i] = roll
				if reroll.Once {
					break RerollOnce
				}
			}
		}
	}

	if term.Exploding != nil {
		switch term.Exploding.Type {
		case Exploding:
			for _, roll := range result.Results {
				for term.Exploding.Match(roll.Result) {
					if err = ctx.recordRoll(&dieRolls); err != nil {
						return Result{}, err
					}
					roll = term.Die.Roll()
					result.Results = append(result.Results, roll)
				}
			}
		case Compounded:
			compound := 0
			for _, roll := range result.Results {
				for term.Exploding.Match(roll.Result) {
					compound += roll.Result
					if err = ctx.recordRoll(&dieRolls); err != nil {
						return Result{}, err
					}
					roll = term.Die.Roll()
				}
			}
			result.Results = append(result.Results, DieRoll{Result: compound, Symbol: strconv.Itoa(compound)})
		case Penetrating:
			for _, roll := range result.Results {
				for term.Exploding.Match(roll.Result) {
					if err = ctx.recordRoll(&dieRolls); err != nil {
						return Result{}, err
					}
					roll = term.Die.Roll()
					newRoll := roll
					newRoll.Result--
					newRoll.Symbol = strconv.Itoa(newRoll.Result)
					result.Results = append(result.Results, newRoll)
				}
			}
		}
	}

	applyLimit(term.Limit, &result)
	applySuccess(term.Success, term.Modifier, &result)
	applyFailure(term.Failure, term.Modifier, &result)
	applySort(term.Sort, &result)
	finaliseTotals(term.Success, term.Failure, term.Modifier, totalMultiplier, &result)

	return result, nil
}

func evalGroupTerm(term GroupTerm, children []vmValue) (result Result) {
	for _, child := range children {
		if term.Combined {
			for _, res := range child.Result.Results {
				result.Results = append(result.Results, DieRoll{
					Result: res.Result + child.Modifier,
					Symbol: strconv.Itoa(res.Result + child.Modifier),
				})
			}
		} else {
			result.Results = append(result.Results, DieRoll{Result: child.Result.Total, Symbol: strconv.Itoa(child.Result.Total)})
		}
	}

	applyLimit(term.Limit, &result)
	applySuccess(term.Success, term.Modifier, &result)
	applyFailure(term.Failure, term.Modifier, &result)
	finaliseTotals(term.Success, term.Failure, term.Modifier, 1, &result)

	if term.Negative {
		result.Total *= -1
		for i, roll := range result.Results {
			roll.Result *= -1
			if roll.Symbol != "" {
				roll.Symbol = strconv.Itoa(roll.Result)
			}
			result.Results[i] = roll
		}
	}

	return result
}

func validateDieLimits(die Die, limits Limits) error {
	limits = limits.normalized()

	var size int
	switch d := die.(type) {
	case NormalDie:
		size = int(d)
	case PercentileDie:
		size = 100
	case FateDie:
		size = 3
	default:
		return fmt.Errorf("unsupported die type %T for limit validation", die)
	}

	if size < 2 {
		return ErrUnsafeDie(die.String())
	}
	if size > limits.MaxDieSize {
		return ErrLimitExceeded(fmt.Sprintf("die size %d exceeds maximum %d", size, limits.MaxDieSize))
	}
	return nil
}

func applyLimit(limitOp *LimitOp, result *Result) {
	if limitOp != nil {
		var rolls Result
		rolls.Results = result.Results[:]
		sort.Sort(&rolls)

		limit := min(limitOp.Amount, len(rolls.Results))

		switch limitOp.Type {
		case KeepHighest:
			rolls.Results = rolls.Results[len(rolls.Results)-limit:]
		case KeepLowest:
			rolls.Results = rolls.Results[:limit]
		case DropHighest:
			rolls.Results = rolls.Results[:len(rolls.Results)-limit]
		case DropLowest:
			rolls.Results = rolls.Results[limit:]
		}

		m := make(map[int]int, len(rolls.Results))
		for _, r := range rolls.Results {
			m[r.Result]++
		}

		newResults := make([]DieRoll, 0, len(rolls.Results))
		for _, a := range result.Results {
			if b, ok := m[a.Result]; ok {
				newResults = append([]DieRoll{a}, newResults...)
				b--
				if b == 0 {
					delete(m, a.Result)
				}
			}
		}

		result.Results = newResults
	}
}

func applySuccess(successOp *ComparisonOp, modifier int, result *Result) {
	if successOp != nil {
		for _, roll := range result.Results {
			if successOp.Match(roll.Result + modifier) {
				result.Successes++
			}
		}
	}
}

func applyFailure(failureOp *ComparisonOp, modifier int, result *Result) {
	if failureOp != nil {
		for _, roll := range result.Results {
			if failureOp.Match(roll.Result + modifier) {
				result.Successes--
			}
		}
	}
}

func applySort(sortType SortType, result *Result) {
	switch sortType {
	case Unsorted:
		return
	case Ascending:
		sort.Sort(result)
	case Descending:
		sort.Sort(sort.Reverse(result))
	}
}

func finaliseTotals(successOp, failureOp *ComparisonOp, modifier, multiplier int, result *Result) {
	if successOp == nil && failureOp == nil {
		for _, roll := range result.Results {
			result.Total += roll.Result
		}
		result.Total += modifier
		result.Total *= multiplier
	} else {
		result.Total = result.Successes
	}
}