elves/elvish

package eval

import (
    "fmt"
    "math"
    "math/big"
    "strconv"
    "time"

    "src.elv.sh/pkg/eval/errs"
    "src.elv.sh/pkg/eval/vals"
    "src.elv.sh/pkg/parse"
)

func init() {
    addBuiltinFns(map[string]any{
        "sleep":     sleep,
        "time":      timeCmd,
        "benchmark": benchmark,
    })
}

var (
    // Reference to [time.After] that can be mutated for testing. Takes an
    // additional Frame argument to allow inspection of the value of d in tests.
    timeAfter = func(fm *Frame, d time.Duration) <-chan time.Time { return time.After(d) }
    // Reference to [time.Now] that can be overridden in tests.
    timeNow = time.Now
)

func sleep(fm *Frame, duration any) error {
    var f float64
    var d time.Duration

    if err := vals.ScanToGo(duration, &f); err == nil {
        d = time.Duration(f * float64(time.Second))
    } else {
        // See if it is a duration string rather than a simple number.
        switch duration := duration.(type) {
        case string:
            d, err = time.ParseDuration(duration)
            if err != nil {
                return ErrInvalidSleepDuration
            }
        default:
            return ErrInvalidSleepDuration
        }
    }

    if d < 0 {
        return ErrNegativeSleepDuration
    }

    select {
    case <-fm.Context().Done():
        return ErrInterrupted
    case <-timeAfter(fm, d):
        return nil
    }
}

type timeOpt struct{ OnEnd Callable }

func (o *timeOpt) SetDefaultOptions() {}

func timeCmd(fm *Frame, opts timeOpt, f Callable) error {
    t0 := time.Now()
    err := f.Call(fm, NoArgs, NoOpts)
    t1 := time.Now()

    dt := t1.Sub(t0)
    if opts.OnEnd != nil {
        newFm := fm.Fork("on-end callback of time")
        errCb := opts.OnEnd.Call(newFm, []any{dt.Seconds()}, NoOpts)
        if err == nil {
            err = errCb
        }
    } else {
        _, errWrite := fmt.Fprintln(fm.ByteOutput(), dt)
        if err == nil {
            err = errWrite
        }
    }

    return err
}

type benchmarkOpts struct {
    OnEnd    Callable
    OnRunEnd Callable
    MinRuns  int
    MinTime  string
    minTime  time.Duration
}

func (o *benchmarkOpts) SetDefaultOptions() {
    o.MinRuns = 5
    o.minTime = time.Second
}

func (opts *benchmarkOpts) parse() error {
    if opts.MinRuns < 0 {
        return errs.BadValue{What: "min-runs option",
            Valid: "non-negative integer", Actual: strconv.Itoa(opts.MinRuns)}
    }

    if opts.MinTime != "" {
        d, err := time.ParseDuration(opts.MinTime)
        if err != nil {
            return errs.BadValue{What: "min-time option",
                Valid: "duration string", Actual: parse.Quote(opts.MinTime)}
        }
        if d < 0 {
            return errs.BadValue{What: "min-time option",
                Valid: "non-negative duration", Actual: parse.Quote(opts.MinTime)}
        }
        opts.minTime = d
    }

    return nil
}

func benchmark(fm *Frame, opts benchmarkOpts, f Callable) error {
    if err := opts.parse(); err != nil {
        return err
    }

    // Standard deviation is calculated using https://en.wikipedia.org/wiki/Algorithms_for_calculating_variance#Welford's_online_algorithm
    var (
        min   = time.Duration(math.MaxInt64)
        max   = time.Duration(math.MinInt64)
        runs  int64
        total time.Duration
        m2    float64
        err   error
    )
    for {
        t0 := timeNow()
        err = f.Call(fm, NoArgs, NoOpts)
        if err != nil {
            break
        }
        dt := timeNow().Sub(t0)

        if min > dt {
            min = dt
        }
        if max < dt {
            max = dt
        }
        var oldDelta float64
        if runs > 0 {
            oldDelta = float64(dt) - float64(total)/float64(runs)
        }
        runs++
        total += dt
        if runs > 0 {
            newDelta := float64(dt) - float64(total)/float64(runs)
            m2 += oldDelta * newDelta
        }

        if opts.OnRunEnd != nil {
            newFm := fm.Fork("on-run-end callback of benchmark")
            err = opts.OnRunEnd.Call(newFm, []any{dt.Seconds()}, NoOpts)
            if err != nil {
                break
            }
        }

        if runs >= int64(opts.MinRuns) && total >= opts.minTime {
            break
        }
    }

    if runs == 0 {
        return err
    }

    avg := total / time.Duration(runs)
    stddev := time.Duration(math.Sqrt(m2 / float64(runs)))
    if opts.OnEnd == nil {
        _, errOut := fmt.Fprintf(fm.ByteOutput(),
            "%v ± %v (min %v, max %v, %d runs)\n", avg, stddev, min, max, runs)
        if err == nil {
            err = errOut
        }
    } else {
        stats := vals.MakeMap(
            "avg", avg.Seconds(), "stddev", stddev.Seconds(),
            "min", min.Seconds(), "max", max.Seconds(), "runs", int64ToElv(runs))
        newFm := fm.Fork("on-end callback of benchmark")
        errOnEnd := opts.OnEnd.Call(newFm, []any{stats}, NoOpts)
        if err == nil {
            err = errOnEnd
        }
    }
    return err
}

func int64ToElv(i int64) any {
    if i <= int64(math.MaxInt) {
        return int(i)
    } else {
        return big.NewInt(i)
    }
}