step_function.go

package incremental

import (
	"time"

	"github.com/PrositAS/go-incremental/internal/core"
)

// Step is one transition of a StepFunctionValue: its value becomes Value starting at At.
type Step[T any] struct {
	At    time.Time
	Value T
}

// StepFunctionValue is a function from time.Time to T with a finite number of steps, in
// nondecreasing time order - the Go counterpart of Incremental's Step_function.t.
type StepFunctionValue[T any] struct {
	Init  T
	Steps []Step[T]
}

// NewStepFunctionValue creates a StepFunctionValue with value init before the first step,
// stepping to each step's Value at its At time in turn - the Go counterpart of
// Incremental's Step_function.create_exn. It panics if steps isn't in nondecreasing time
// order.
func NewStepFunctionValue[T any](init T, steps []Step[T]) StepFunctionValue[T] {
	for i := 1; i < len(steps); i++ {
		if steps[i].At.Before(steps[i-1].At) {
			panic("incremental: StepFunctionValue steps must be in nondecreasing time order")
		}
	}
	return StepFunctionValue[T]{Init: init, Steps: steps}
}

// ValueAt returns the step function's value at the given time - the Go counterpart of
// Incremental's Step_function.value.
func (sf StepFunctionValue[T]) ValueAt(at time.Time) T {
	v := sf.Init
	for _, step := range sf.Steps {
		if step.At.After(at) {
			break
		}
		v = step.Value
	}
	return v
}

// IncrementalStepFunction creates a node whose value follows child's step function,
// advancing at each step's time as the clock passes it, and re-extracting the step
// function whenever child changes - the Go counterpart of Incremental's
// Clock.incremental_step_function.
func IncrementalStepFunction[T any](clock *Clock, child core.ValueNode[StepFunctionValue[T]]) core.ValueNode[T] {
	n := &stepFunctionNode[T]{
		Node:                        core.New[T](nil, nil),
		child:                       child,
		clock:                       clock,
		extractedFromChildChangedAt: core.NoneStabilizationNum,
	}
	n.SetKind(int(KindStepFunction))
	core.AddEdge(n, child)
	return n
}

// StepFunction creates a node whose value starts at init and follows steps as the clock
// advances - the Go counterpart of Incremental's Clock.step_function. It is a convenience
// over IncrementalStepFunction for a step function fixed at creation.
func StepFunction[T any](clock *Clock, init T, steps []Step[T]) core.ValueNode[T] {
	return IncrementalStepFunction(clock, Const(NewStepFunctionValue(init, steps)))
}

type stepFunctionNode[T any] struct {
	*core.Node[T]
	child                       core.ValueNode[StepFunctionValue[T]]
	clock                       *Clock
	extractedFromChildChangedAt core.StabilizationNum
	current                     T
	upcomingSteps               []Step[T]
	alarm                       *core.Alarm
}

func (n *stepFunctionNode[T]) Recompute(now core.StabilizationNum, _ core.RoundCtx) {
	// extractedFromChildChangedAt starts at NoneStabilizationNum, so the first Recompute
	// always extracts even if child is a Const, whose ChangedAt() never advances past none
	// (unlike OCaml's Const, which goes through one real recompute and so does get a real
	// changed_at the first time - see const.go for why this port skips that step).
	if n.child != nil && (n.extractedFromChildChangedAt.IsNone() || n.child.ChangedAt() > n.extractedFromChildChangedAt) {
		n.extractedFromChildChangedAt = n.child.ChangedAt()
		n.clock.alarms.Remove(n.alarm)
		n.alarm = nil

		sf, _ := n.child.Value()
		n.current = sf.Init
		n.upcomingSteps = sf.Steps

		if n.child.Kind() == int(KindConst) {
			core.RemoveEdge(n, n.child)
			n.child = nil
		}
	}

	n.advanceTo(n.clock.Now())
	n.SetValue(n.current, now)

	if len(n.upcomingSteps) > 0 {
		at := n.upcomingSteps[0].At
		n.alarm = n.clock.alarms.Add(at, core.NewAlarmValue(func(core.StabilizationNum) {
			if n.IsValid() {
				core.MarkStale(n)
			}
		}))
	} else if n.child == nil {
		n.SetKind(int(KindConst))
	}
}

// advanceTo consumes every upcoming step at or before "to", the Go counterpart of
// Incremental's Step_function_node.advance.
func (n *stepFunctionNode[T]) advanceTo(to time.Time) {
	for len(n.upcomingSteps) > 0 && !n.upcomingSteps[0].At.After(to) {
		n.current = n.upcomingSteps[0].Value
		n.upcomingSteps = n.upcomingSteps[1:]
	}
}