Dynamos possibly miss _apply_iterate?

[femtomc]

I've noticed that I'm missing some calls which I want to grab with my dynamo:

In particular, I have a rand call:

function execute_MA!(step::MentalAct,
                     tree::ActiveBehaviorTree, 
                     env::Environment)
    # Here.
    ret = rand(:mental_computation, step.fn, tree.globals, env, step.args...)
    println("Just finished a mental computation!")

    par = parent(tree, step)
    par.semantics.ret[step.uuid] = ret
    tree.executing_actions[step.uuid] = Completed()
end

which I know gets turned into an _apply_iterate. I have logging setup for these calls, as well as a mechanism which let's me instantiate a new HierarchicalTrace and recurse in.

@inline function (tr::HierarchicalTrace)(fn::typeof(rand), addr::Address, call::Function, args...)
    println("Randomness at $addr, call $call.")
    n_tr = Trace()
    ret = n_tr(call, args...)
    tr.chm[addr] = CallSite(n_tr, call, args, ret)
    return ret
end

but I never hit this call:

Randomness at step_arbiter, call step_arbiter.
Randomness at step_selection.
Randomness at execute_MA, call execute_MA!.
Just finished a mental computation!

I'm not sure how to hit this one?

This call occurs in another call (so it occurs in the n_tr(call, args...) - I'm not sure why that would matter.

[MikeInnes]

Ah, this is one thing that recur! should probably do that it doesn't currently. Basically the rand call gets lowered to

self(Core._apply_iterate, rand, (arg1,), args)

but apply_iterate should be special and instead lower to

Core._apply_iterate(self, (rand,), (arg1,), args)

You can do this as a code transform inside recur but as a workaround, you can also do it by overloading, something like

(tr::HierarchicalTrace)(::typeof(apply_iterate), f, args...) = # concat args and apply f

[femtomc]

Thanks! This was my simple hack for now:

# Fix for _apply_iterate.
function f_push!(arr::Array, t::Tuple{}) end
f_push!(arr::Array, t::Array) = append!(arr, t)
f_push!(arr::Array, t::Tuple) = append!(arr, t)
f_push!(arr, t) = push!(arr, t)
function flatten(t::Tuple)
    arr = Any[]
    for sub in t
        f_push!(arr, sub)
    end
    return arr
end

function (mx::ExecutionContext)(::typeof(Core._apply_iterate), f, c::typeof(rand), args...)
    return mx(c, flatten(args)...)
end

I'm not sure when you want to intercept _apply_iterate? I'm guessing AD - which is a bit more heavyweight than PP.