expr.mfold with branch pruning

Stream: metaprogramming / tactics

[Eric Wieser]

Is there a variant of expr.mfold that allows me to stop recursing into a particular subexpression?

[Eric Wieser]

I guess I'm asking for this:

meta def expr.pruned_mfold {α : Type} {m : Type → Type} [monad m]
  (f : expr → α → m (α × bool)) : ∀ (init : α) (e : expr), m α
| init e@(expr.var _) := prod.fst <$> f e init
| init e@(expr.sort _) := prod.fst <$> f e init
| init e@(expr.const _ _) := prod.fst <$> f e init
| init e@(expr.mvar unique pretty type) := do
    (val, prune) ← f e init,
    if prune then return val
    else expr.pruned_mfold val type
| init e@(expr.local_const unique pretty bi type) := do
    (val, prune) ← f e init,
    if prune then return val
    else expr.pruned_mfold val type
| init e@(expr.app fn x) := do
    (val, prune) ← f e init,
    if prune then return val
    else do
      val ← expr.pruned_mfold val fn,
      expr.pruned_mfold val x
| init e@(expr.lam var_name bi var_type body) := do
    (val, prune) ← f e init,
    if prune then return val
    else do
      val ← expr.pruned_mfold val var_type,
      expr.pruned_mfold val body
| init e@(expr.pi var_name bi var_type body) := do
    (val, prune) ← f e init,
    if prune then return val
    else do
      val ← expr.pruned_mfold val var_type,
      expr.pruned_mfold val body
| init e@(expr.elet var_name type assignment body) := do
    (val, prune) ← f e init,
    if prune then return val
    else do
      val ← expr.pruned_mfold val type,
      val ← expr.pruned_mfold val assignment,
      expr.pruned_mfold val body
| init e@(expr.macro _ ls) := do
    (val, prune) ← f e init,
    if prune then return val
    else do
      ls.mfoldl expr.pruned_mfold val

[Jannis Limperg]

I'm not aware of an existing implementation of this. I think you can implement it in terms of regular mfold:

meta def expr.pruned_mfold {α : Type} {m : Type → Type} [monad m]
  (f : expr → α → m (α × bool)) (init : α) (e : expr) : m α := do
  (acc, _) ← e.mfold (init, ff) $ λ e' _ ⟨acc, prune⟩,
    if prune
      then pure (acc, tt)
      else f e' acc,
  pure acc

However, I don't think this short-circuits properly, so it'll continue to recurse after you've set prune to tt. However, it'll only do a trivial amount of work for each subexpression.

[Eric Wieser]

That doesn't prune branches though, that prunes the entire recursion

[Eric Wieser]

I want folding over g (f (f a)) (h (f b)) to run on g ..., f (f a), h (f b) and f b, but not visit b or (f a) as those are children of a subexpression that has been "pruned"

[Gabriel Ebner]

Note that the big advantage of expr.fold (which is used to implement expr.mfold) is that it doesn't visit identical subexpressions twice.

[Gabriel Ebner]

This tastes a bit of #xy, what are you actually trying to do?

[Eric Wieser]

This looks like a shorter implementation of pruned_mfold:

/-- get all child expressions of an expression -/
meta def expr.nested_exprs : ∀ (e : expr), list expr
| (expr.var _) := []
| (expr.sort _) := []
| (expr.const _ _) := []
| (expr.mvar unique pretty type) := [type]
| (expr.local_const unique pretty bi type) := [type]
| (expr.app fn x) := [fn, x]
| (expr.lam var_name bi var_type body) := [var_type, body]
| (expr.pi var_name bi var_type body) := [var_type, body]
| (expr.elet var_name type assignment body) := [type, assignment, body]
| (expr.macro _ ls) := ls

/-- mfold over an expression tree, without visiting branches for which `(f e val).snd` is true -/
meta def expr.pruned_mfold {α : Type} {m : Type → Type} [monad m]
  (f : expr → α → m (α × bool)) : ∀ (init : α) (e : expr), m α :=
λ init e, do
  (val, prune) ← f e init,
  ff ← pure prune | return val,
  list.mfoldl expr.pruned_mfold val e.nested_exprs

[Eric Wieser]

what are you actually trying to do?

Find decidable terms, without recursing into the decidable terms that they might be produced by

[Eric Wieser]

I don't know whether I actually need to do that pruning yet - CI seems stuck on #6485, and probably I just need to wait for the queue to clear

[Jannis Limperg]

Eric Wieser said:

That doesn't prune branches though, that prunes the entire recursion

I see. I'm not sure whether there's an mfold/fold incantation that doesn't have this problem, so you might indeed have to roll your own fold.

[Gabriel Ebner]

I don't think expr.mfold is a good idea here. When traversing lambdas/pis/lets, you probably want to instantiate the de Bruijn variables with local constants. Otherwise infer_type, mk_instance, etc. won't work.

[Eric Wieser]

Is it safe if I've done open_pis first?

[Gabriel Ebner]

open_pis will only do the outermost layer. E.g. in Π a b c, foo (λ x, irrational.decidable (x+1)), the x+1 will actually be #0+1 if you traverse it via expr.fold.

[Eric Wieser]

Ah, got it

[Eric Wieser]

That explains why infer_type failed on me sometimes

[Eric Wieser]

But probably my linter is not well thought out for expressions inside binders anyway

[Eric Wieser]

So best to ignore those for now

[Jesse Michael Han]

i use something like this to extract data from proof terms in lean-step-public