make this faster)

[github-actions]

with a memory estimate of 388.02 MiB, over 8512429 allocations.

retries[] # 30

Range (min … max): 2.029 ms … 137.030 ms ┊ GC (min … max): 0.00% … 98.14%

Time (median): 2.879 ms ┊ GC (median): 0.00%

Time (mean ± σ): 4.377 ms ± 6.119 ms ┊ GC (mean ± σ): 4.36% ± 4.07%

██████████████████▅▇▆▄▆▄▆▇▄▆▄▅▁▆▁▅▇▁▄▄▁▁▁▄▁▁▅▄▆▁▄▄▁▁▁▄▁▁▁▅▅ █

2.03 ms Histogram: log(frequency) by time 22.6 ms <

https://github.com/Tractables/Dice.jl/blob/daa0fc51f50e9ffbb8cf0be440618d8af92247cd/qc/benchmarks/stlc_faster_10samples.jl#L24

using Revise
using Dice
using BenchmarkTools
using ProfileView

include("benchmarks.jl")

generation_params = LangSiblingDerivedGenerator{STLC}(
    root_ty=Expr.t,
    ty_sizes=[Expr.t=>5, Typ.t=>2],
    stack_size=2,
    intwidth=3,
)

SEED = 0
out_dir="/tmp"
log_path="/dev/null"
rs = RunState(Valuation(), Dict{String,ADNode}(), open(log_path, "w"), out_dir, MersenneTwister(SEED), nothing,generation_params)

generation::Generation = generate(rs, generation_params)

g::Dist = generation.value

# Sample some tree until it's valid (TODO: make this faster)
a = ADComputer(rs.var_vals)

NUM_SAMPLES = 10

function wellTyped(e)
    @assert isdeterministic(e)
    @match typecheck(e) [
        Some(_) -> true,
        None() -> false,
    ]
end

retries = Ref(0)
#== @benchmark ==# @time begin
    samples = []
    while length(samples) < NUM_SAMPLES
        retries[] += 1
        s = sample_as_dist(rs.rng, a, g)
        if wellTyped(s)
            push!(samples, s)
        end
    end
end
#  Single result which took 26.426 s (3.00% GC) to evaluate, (7s, 26s, 30s, 40s)
#  with a memory estimate of 388.02 MiB, over 8512429 allocations.
retries[] # 30 

l = Dice.LogPrExpander(WMC(BDDCompiler([
    prob_equals(g, sample)
    for sample in samples
])))
@time begin
    loss, actual_loss = sum(
        begin
            lpr_eq = Dice.expand_logprs(l, LogPr(prob_equals(g, sample)))
            [lpr_eq * compute(a, lpr_eq), lpr_eq]
        end
        for sample in samples
    )
end
# 5.3s first run, 1.4s rest

length(l.cache) # 331

@benchmark vals, derivs = differentiate(
    rs.var_vals,
    Derivs([loss => 1.])
)
# BenchmarkTools.Trial: 1060 samples with 1 evaluation.
#  Range (min … max):  2.029 ms … 137.030 ms  ┊ GC (min … max): 0.00% … 98.14%
#  Time  (median):     2.879 ms               ┊ GC (median):    0.00%
#  Time  (mean ± σ):   4.377 ms ±   6.119 ms  ┊ GC (mean ± σ):  4.36% ±  4.07%

#   ██▇▆▅▃▃▂▃▁▁▂▂  ▁                                             
#   ██████████████████▅▇▆▄▆▄▆▇▄▆▄▅▁▆▁▅▇▁▄▄▁▁▁▄▁▁▅▄▆▁▄▄▁▁▁▄▁▁▁▅▅ █
#   2.03 ms      Histogram: log(frequency) by time      22.6 ms <

#  Memory estimate: 292.17 KiB, allocs estimate: 8034.

ct = Ref(0)
Dice.foreach_down(loss) do _ ct[] += 1 end
ct[] # 350