package debruijn

// Substitute replaces the variable at the given index with the replacement expression.
// The replacement is shifted appropriately as we descend into nested abstractions.
func Substitute(expr Expression, index int, replacement Expression) Expression {
	switch e := expr.(type) {
	case *Variable:
		if e.index == index {
			return replacement
		}
		return e

	case *Abstraction:
		// When entering an abstraction, increment the target index and shift the
		// replacement to account for the new binding context.
		shiftedReplacement := Shift(replacement, 1, 0)
		newBody := Substitute(e.body, index+1, shiftedReplacement)
		if newBody == e.body {
			return e
		}
		return NewAbstraction(newBody)

	case *Application:
		newAbs := Substitute(e.abstraction, index, replacement)
		newArg := Substitute(e.argument, index, replacement)
		if newAbs == e.abstraction && newArg == e.argument {
			return e
		}
		return NewApplication(newAbs, newArg)

	default:
		return expr
	}
}