Conveniently apply function if parameter is `some`

Stream: mathlib4

[Thomas Murrills]

Every so often I find myself in need of something like

def Function.apply? {α} {β} (param? : Option α) (f : α → β → β) (val : β) : β :=
  match param? with
  | some param => f param val
  | none => val

i.e., apply a function if param? is some, and leave the value alone otherwise. It's mainly for code style, to allow something like f.applyIfSome p? val so that I don't need to write out the match and duplicate val, use (if let some p := p? then f p else id) val, or worse.

Does this wrapper exist somewhere? If not, should we have it? And if so, what should it be called, and what should it be dot notation on (Function or Option)? (This seems tied up in the name to me.)

(I think apply? is the best I've come up with so far, but I also considered applyIfSome, through, finallyWith? (but confusing if your value is actually monadic, in which case this is not related to finally)... Technically this is Option.foldr, but that seems quite bad for readability.)

[Robin Arnez]

Is param?.elim id f val fine?

[Thomas Murrills]

I consider using elim pretty unreadable for (meta) code I'd like to commit to mathlib. :) (Having a wrapper also lets us add documentation!)

[Robin Arnez]

Well, in do notation at least the cleanest solution would be

if let some x := x? then
  val := f x val

Or is this for term mode?

[Thomas Murrills]

Having a one-off let mut like that is fine for normal values, but in my case(s) val is sometimes a monad action (such as the motivating one, behind the scenes :grinning_face_with_smiling_eyes:), and

let mut val := do
  <code block>
  ...
if let some x := x? then
  val := f x val
val

isn't great for readability...

[Thomas Murrills]

Whereas

(monadCombinator · foo).applyIfSome param? do
  <code block>
  ...

isn't too bad.

[Thomas Murrills]

(Thinking about it, the same friction also occurs with regular Bools and f : β → β, so it might be nice to have a similar one for Bool. applyIf and applyIfSome, maybe?)

[Robin Arnez]

Now, if it were possible, what about

if flag then
  with_cont cont => monadCombinator cont
<code block>

or

if let some x := x? then
  with_cont cont => monadCombinator cont x
<code block>

[Robin Arnez]

(the do elaborator is getting a makeover anyways, so maybe there's something worth putting in an RFC here)

[Thomas Murrills]

Ooh, interesting! Hmm, I'd be a little concerned about the apparent nonlocality: visually, the action is happening underneath an if branch, but we pop out to the place that it affects. Unlike let mut, we wouldn't have a neon sign earlier in the code to watch out for that. (Or maybe we would?) I also think the naming would need some work; "continuation" isn't the most transparent term! But maybe it can be molded further, it would certainly be nice to bind a variable and take advantage of existing things like if and let...

[Thomas Murrills]

Robin Arnez said:

(the do elaborator is getting a makeover anyways, so maybe there's something worth putting in an RFC here)

Oh, nice! Are there any public discussions on potential feature requests for do? It might be nice to have a central place in which the community can mention ideas and assess their popularity; I've got a couple to toss into the ring. :)

[Robin Arnez]

I don't there is a feature discussion thread yet but we'll at least get do elaborator support, so if you want a feature you can implement it easily! This, together with the fact that the new do elaborator is continuation-based makes it pretty easy to implement a custom with_cont:

import Lean

open Lean Meta Elab Term
elab tk:"with_cont " fn:term : doElem <= dec => do
  let combinator ← elabTerm fn none
  let type ← inferType combinator
  let .forallE _ t b _ ← whnfForall type | throwFunctionExpected combinator
  let u := (← read).monadInfo.u
  let cont ← withLetDecl dec.resultName (.const ``PUnit [u.succ]) (.const ``PUnit.unit [u.succ]) fun e => do
    mkLetFVars #[e] (← dec.k)
  let contTy ← inferType cont
  unless ← isDefEq t contTy do
    throwErrorAt tk "error"
  let res := (← read).doBlockResultType
  let monad := (← read).monadInfo.m
  let expect := .app monad res
  unless ← isDefEq b expect do
    throwErrorAt tk "other error"
  return combinator.beta #[cont]

def test (n : Nat) : ReaderM Nat Nat := do
  for i in *...n do
    with_cont withReader (· + 1)
  read

#guard ((test 0).run 3).run == 3
#guard ((test 5).run 3).run == 8

(fyi this is on https://github.com/leanprover/lean4/tree/sg/newdo)

[Robin Arnez]

And of course, the real reason you'd want this:

/-- info: [a_0, a_1, a_2, a_3, a_4, a_5, a_6, a_7, a_8, a_9] -/
#guard_msgs in
run_meta
  let mut vars : Array Expr := #[]
  for i in *...10 do
    let e ← with_cont withLocalDeclD ((`a).appendIndexAfter i) (mkConst ``Nat)
    vars := vars.push e
  Lean.logInfo m!"{vars}"