package debruijn

// Shift increments all free variable indices by delta when crossing depth abstractions.
func Shift(expr Expression, delta int, depth int) Expression {
	switch e := expr.(type) {
	case *Variable:
		if e.index >= depth {
			return NewVariable(e.index+delta, e.label)
		}
		return e

	case *Abstraction:
		newBody := Shift(e.body, delta, depth+1)
		if newBody == e.body {
			return e
		}
		return NewAbstraction(newBody)

	case *Application:
		newAbs := Shift(e.abstraction, delta, depth)
		newArg := Shift(e.argument, delta, depth)
		if newAbs == e.abstraction && newArg == e.argument {
			return e
		}
		return NewApplication(newAbs, newArg)

	default:
		return expr
	}
}

// Substitute replaces variable at index 0 with replacement in expr.
// This assumes expr is the body of an abstraction being applied.
func Substitute(expr Expression, replacement Expression) Expression {
	return substitute(expr, 0, replacement)
}

// substitute replaces variable at targetIndex with replacement, adjusting indices as needed.
func substitute(expr Expression, targetIndex int, replacement Expression) Expression {
	switch e := expr.(type) {
	case *Variable:
		if e.index == targetIndex {
			return Shift(replacement, targetIndex, 0)
		}
		if e.index > targetIndex {
			return NewVariable(e.index-1, e.label)
		}
		return e

	case *Abstraction:
		newBody := substitute(e.body, targetIndex+1, replacement)
		if newBody == e.body {
			return e
		}
		return NewAbstraction(newBody)

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

	default:
		return expr
	}
}