; Bench fixture: Collatz step-counter, pure-compute integer math.
;
; For each starting value n in [1..N], iteratively count the number of
; Collatz steps to reach 1. Sum all step counts.
;
; Distinct from bench_compute_intsum:
;   - Branchy: each step does an `n % 2 == 0` check and either halves n
;     or computes 3n+1. Tests branch-prediction friendliness of the
;     codegen.
;   - Two nested tail-recursions: outer (sum over starting values) and
;     inner (count steps for one value). Both must lower to musttail
;     loops or the bench segfaults at scale.
;   - No heap, no closure, no pattern match — pure integer + branch.
;
; Sizes (small because Collatz step counts grow logarithmically; the
; cost is dominated by the per-step overhead, ~30ns each):
;   N =     10_000   sum_steps =       849666
;   N =    100_000   sum_steps =     10753840
;   N =    500_000   sum_steps =     62134795
;
; Step counts cross-validated against a Python reference; deterministic
; across allocators.

(module bench_compute_collatz

  (fn collatz_steps
    (doc "Tail-recursive: count Collatz steps from n to 1, accumulating in acc.")
    (type
      (fn-type
        (params (own (con Int)) (own (con Int)))
        (ret (own (con Int)))))
    (params n acc)
    (body
      (if (app eq n 1)
        acc
        (if (app eq (app % n 2) 0)
          (tail-app collatz_steps (app / n 2)            (app + acc 1))
          (tail-app collatz_steps (app + (app * n 3) 1)  (app + acc 1))))))

  (fn sum_steps_loop
    (doc "Tail-recursive: sum collatz_steps(i) for i in [n, n-1, ..., 1].")
    (type
      (fn-type
        (params (own (con Int)) (own (con Int)))
        (ret (own (con Int)))))
    (params i total)
    (body
      (if (app eq i 0)
        total
        (tail-app sum_steps_loop
          (app - i 1)
          (app + total (app collatz_steps i 0))))))

  (fn run_one
    (type (fn-type (params (own (con Int))) (ret (own (con Unit))) (effects IO)))
    (params n)
    (body (app print (app sum_steps_loop n 0))))

  (fn main
    (type (fn-type (params) (ret (own (con Unit))) (effects IO)))
    (params)
    (body
      (seq (app run_one 10000)
      (seq (app run_one 100000)
           (app run_one 500000))))))
