List.map_map How to use optimization

Stream: lean4

[lean_user]

import Lean

#check List.map_map
-- List.map_map.{u_1, u_2, u_3} {β : Type u_1} {γ : Type u_2} {α : Type u_3} {g : β → γ} {f : α → β} {l : List α} :
-- List.map g (List.map f l) = List.map (g ∘ f) l

set_option trace.compiler.ir.result true in
@[simp] def test (l : List Nat) (g : Nat → Nat) (f : Nat → Nat)
  : List Nat :=
  List.map g (List.map f l)

--[result]
--def test (x_1 : obj) (x_2 : obj) (x_3 : obj) : obj :=
--  let x_4 : obj := ctor_0[List.nil];
--  let x_5 : obj := List.mapTR.loop._rarg x_3 x_1 x_4;
--  let x_6 : obj := List.mapTR.loop._rarg x_2 x_5 x_4;
--  ret x_

@[simp] theorem test_optimized_eq (l : List Nat) (g f : Nat → Nat) :
    test l g f = List.map (g ∘ f) l := by
  simp [List.map_map]

#eval do
  Lean.Compiler.compile #[``test]
  Lean.Compiler.LCNF.showDecl .base ``test

--def test l g f : List Nat :=
--  let _x.1 := @List.mapTR _ _ f l;
--  let _x.2 := @List.mapTR _ _ g _x.1;
--  return _x.2

[Aaron Liu]

Can you provide a #mwe?

[Robin Arnez]

This is currently not supported

[Markus Himmel]

If you use iterators on the latest nightly, the two maps are fused, though you get an additional loop turning the array that is used internally back into a list. If you use arrays instead of lists, there will only be a single loop.

import Std.Data.Iterators

set_option trace.compiler.ir.result true in
def test (l : List Nat) (g : Nat → Nat) (f : Nat → Nat) : List Nat :=
  ((l.iter.map f).map g).toList

/-
def Std.Iterators.IterM.DefaultConsumers.toArrayMapped.go._at.test._spec_2 (x_1 : obj) (x_2 : obj) (x_3 : obj) (x_4 : obj) : obj :=
  case x_3 : obj of
  List.nil →
    dec x_2;
    dec x_1;
    ret x_4
  List.cons →
    let x_5 : obj := proj[0] x_3;
    inc x_5;
    let x_6 : obj := proj[1] x_3;
    inc x_6;
    dec x_3;
    inc x_2;
    let x_7 : obj := app x_2 x_5;
    inc x_1;
    let x_8 : obj := app x_1 x_7;
    let x_9 : obj := Array.push ◾ x_4 x_8;
    let x_10 : obj := Std.Iterators.IterM.DefaultConsumers.toArrayMapped.go._at.test._spec_2 x_1 x_2 x_6 x_9;
    ret x_10
def test._closed_1 : obj :=
  let x_1 : obj := ctor_0[List.nil];
  let x_2 : obj := Array.mk ◾ x_1;
  ret x_2
def test (x_1 : obj) (x_2 : obj) (x_3 : obj) : obj :=
  let x_4 : obj := test._closed_1;
  let x_5 : obj := Std.Iterators.IterM.DefaultConsumers.toArrayMapped.go._at.test._spec_2 x_2 x_3 x_1 x_4;
  let x_6 : obj := Array.toList ◾ x_5;
  ret x_6
-/