Lists with predicates and proof erasure

Stream: lean4

[Wojciech Nawrocki]

I am wondering about the overhead of proving things about the elements of a list in a context where we have a list l and know ∀ x ∈ l, P x. There are some examples of transformations on those in mathlib4. Supposing we have such a predicate, there are at least two ways to use it on every element in a map operation. One, we can wrap the membership proof in a Sigma-type with

def List.withMems : (l : List α) → List (Σ' (x : α), x ∈ l)
  | [] => []
  | a :: as =>
    ⟨a, by simp⟩ :: as.withMems.map' fun ⟨x, hX⟩ => ⟨x, by simp [hX]⟩

and then use it in standard List.map. For example if we know that every element is green and want to pass it through a function that has greenness as a precondition

def isGreen (n : Nat) : Prop := n = 3
def onlyGreen (n : Nat) (_ : isGreen n) : Nat := n

def greens₁ (l : List Nat) (H : ∀ x ∈ l, isGreen x) : List Nat :=
  l.withMems.map fun ⟨n, hN⟩ => onlyGreen n (H n hN)

However this seems to generate suboptimal code since the Σ's have a different runtime representation, so withMems is not a noop and is called first followed by map, resulting in two loops.

def greens₁ (x_1 : obj) (x_2 : ◾) : obj :=
  let x_3 : obj := List.withMems._rarg x_1;
  let x_4 : obj := greens₁._closed_1;
  let x_5 : obj := List.map'._rarg x_4 x_3;
  ret x_5

Note that I used a custom List.map' to avoid the tail-recursive specialization mapTR, but this does not seem to matter.

Second, we could write a custom membership-proving map like

def List.mapMems (l : List α) (f : (x : α) → x ∈ l → β) : List β :=
  go l (fun _ => id)
where go : (l' : List α) → (∀ x ∈ l', x ∈ l) → List β
  | [],    _ => []
  | a::as, H =>
    have hA : a ∈ l := H a <| List.mem_cons_self _ _
    have hAs : ∀ x ∈ as, x ∈ l := List.forall_mem_of_forall_mem_cons H
    f a hA :: go as hAs

and then use it directly

def greens₂ (l : List Nat) (H : ∀ x ∈ l, isGreen x) : List Nat :=
  l.mapMems fun n hN => onlyGreen n (H n hN)

The code generated here seems very close to standard List.map and more optimal, although I didn't look at it in too much detail.

Is this problem explored in the literature on compiling dependently-typed languages? How do people deal with it? Is the mapMems variant (perhaps with a csimp lemma to compile it to a tail-recursive version) the right solution here? Should the Lean compiler (eventually) be able to compile withMems to the identity function, assuming that the runtime representation of PSigma with Prop can be made just the representation of the value?

[Leonardo de Moura]

Should the Lean compiler (eventually) be able to compile withMems to the identity function, assuming that the runtime

Yes, this is on our to-do list, but for Subtype. That is, one would have to write

def List.withMems : (l : List α) → List { x : α // x ∈ l }
  | [] => []
  | a :: as =>
    ⟨a, by simp⟩ :: as.withMems.map' fun ⟨x, hX⟩ => ⟨x, by simp [hX]⟩

Recall that { x : α // x ∈ l } and α already have the same representation in our runtime.

[Leonardo de Moura]

@Wojciech Nawrocki I just noticed you used a nested map in your definition. When we considered a similar example, we used the following alternative definition which is easier to detect that withMems is the identity function at the IR level.

set_option trace.compiler.ir.init true
def List.withMems (as : List α) (as' : List α) (h : (a : α) → a ∈ as' → a ∈ as) : List { x : α // x ∈ as } :=
  match as' with
  | [] => []
  | a :: as' => ⟨a, h a (List.Mem.head ..)⟩ :: withMems as as' fun a h' => h a (List.Mem.tail _ h')

We can also write it as

def List.withMems (as : List α) : List { x : α // x ∈ as } :=
  go as as fun _ h => h
where
  go (as : List α) (as' : List α) (h : (a : α) → a ∈ as' → a ∈ as) : List { x : α // x ∈ as } :=
    match as' with
    | [] => []
    | a :: as' => ⟨a, h a (List.Mem.head ..)⟩ :: go as as' fun a h' => h a (List.Mem.tail _ h')

[Mario Carneiro]

that go function looks pretty similar to mathlib's docs#list.pmap which generalizes the predicate x \in as to an arbitrary proposition (which also makes it easier to see that as is not a computationally relevant input to the function)

[Leonardo de Moura]

assuming that the runtime representation of PSigma with Prop can be made just the representation of the value?

This part is not on our to-do list. We currently don't have plans for specializing parametric types. That is, PSigma is always a pair in our runtime, even if the second argument is a Prop.

[Mario Carneiro]

although I guess the problem with pmap is that it looks like map and is not specialized to subtype.mk meaning that it can be recognized as the identity function

[Mario Carneiro]

That is, I would guess that this one would be a good candidate to be recognized as the identity

def pmap' (P : α → Prop) : (as : List α) → ((a : α) → a ∈ as → P a) → List { x : α // P x }
  | [], h => []
  | a :: as', h => ⟨a, h a (List.Mem.head ..)⟩ :: pmap' P as' fun a h' => h a (List.Mem.tail _ h')

[Leonardo de Moura]

Yes, it is.

[Leonardo de Moura]

@Wojciech Nawrocki Are you planning to use withMems to justify termination when using map and foldl? If yes, here is another approach we have been considering https://github.com/leanprover/lean4/blob/master/tests/lean/run/combinatorsAndWF.lean
In this approach, we only need to modify the WF module, and the user would be able to just write

inductive Expr where
  | app (f : String) (args : List Expr)
  | var (n : String)

def Expr.numVars : Expr → Nat
  | app f args => args.foldl (init := 0) fun sum arg =>  sum + numVars arg
  | var _ => 1

[Wojciech Nawrocki]

Mario Carneiro said:

that go function looks pretty similar to mathlib's docs#list.pmap which generalizes the predicate x \in as to an arbitrary proposition (which also makes it easier to see that as is not a computationally relevant input to the function)

Oh indeed, and withMems is list.attach.

Leonardo de Moura said:

Wojciech Nawrocki Are you planning to use withMems to justify termination when using map and foldl?

No, I needed it to typecheck some array accesses using indices obtained from List.range and a lemma ∀ x ∈ List. range n, x < n. In this case I could also go with an alternative List.range' : (n : Nat) → List (Fin n). It looks like Fin n gets the subtype treatment and is represented as just the Nat value.

The WF feature would be great to have regardless, I have ran into the map/fold termination problem many times.

[Wojciech Nawrocki]

Thank you for your answers!