package debruijn

// Iterator provides depth-first traversal of De Bruijn expressions.
type Iterator struct {
	trace []*Expression
}

// NewIterator creates a new iterator starting at the given expression.
func NewIterator(expr *Expression) *Iterator {
	return &Iterator{[]*Expression{expr}}
}

// Done returns true when the iterator has finished traversal.
func (i *Iterator) Done() bool {
	return len(i.trace) == 0
}

// Current returns a pointer to the current expression.
func (i *Iterator) Current() *Expression {
	if i.Done() {
		return nil
	}

	return i.trace[len(i.trace)-1]
}

// Parent returns a pointer to the parent expression.
func (i *Iterator) Parent() *Expression {
	if len(i.trace) < 2 {
		return nil
	}

	return i.trace[len(i.trace)-2]
}

// Swap replaces the current expression with the given expression.
func (i *Iterator) Swap(with Expression) {
	current := i.Current()
	if current != nil {
		*current = with
	}
}

// Back moves the iterator back to the parent expression.
func (i *Iterator) Back() bool {
	if i.Done() {
		return false
	}

	i.trace = i.trace[:len(i.trace)-1]
	return true
}

// Next advances the iterator to the next expression in leftmost-outermost order.
func (i *Iterator) Next() {
	switch typed := (*i.Current()).(type) {
	case *Abstraction:
		i.trace = append(i.trace, &typed.body)
	case *Application:
		i.trace = append(i.trace, &typed.abstraction)
	case *Variable:
		for len(i.trace) > 1 {
			if app, ok := (*i.Parent()).(*Application); ok {
				if app.abstraction == *i.Current() {
					i.Back()
					i.trace = append(i.trace, &app.argument)
					return
				}
			}

			i.Back()
		}

		i.trace = []*Expression{}
	}
}