Unexpected delayed evaluation of higher-order functions

[ohad]

trafficstars

Steps to Reproduce

versionA : Bool -> Nat -> Nat
versionA True  k0 = k0
versionA False k0 = 0

versionB : Bool -> Nat -> Nat
versionB True  = \k0 => k0
versionB False = \k0 => 0

-- This works as expected
Ex1 : Nat -> Nat
Ex1 = versionB True

-- This does not
Ex2 : Bool -> (Nat, Nat, Nat) -> (Nat, Nat, Nat)
Ex2 b = let f = versionB b
        in \(x,y,z) => (f x, f y, f z)

and then :di Ex1 and :di Ex2 in the REPL.

Expected Behavior

λΠ> :di Ex1
Compiled: \ {ext:0} => {ext:0}
λΠ> :di Ex2
Compiled: \ {arg:0} => let f = case {arg:0} of
  { 1 => \{clam:0} => {clam:0}
  ; 0 => \{clam:0} => 0
  } in
\{lamc:0}=> case {lamc:0} of
  { Builtin.MkPair {tag = 0} [cons] {e:2} {e:3} => case {e:3} of  { Builtin.MkPair {tag = 0} [cons] {e:6} {e:7} => Builtin.MkPair {tag = 0} [cons] (f {e:2}) (Builtin.MkPair {tag = 0} [cons] (f {e:6}) (f {e:7}))}
  }

Observed Behavior

λΠ> :di Ex2
Compiled: \ {arg:0} => let f = \{clam:0}=> case {arg:0} of
  { 1 => {clam:0}
  ; 0 => 0
  } in
\{lamc:0}=> case {lamc:0} of
  { Builtin.MkPair {tag = 0} [cons] {e:2} {e:3} => case {e:3} of  { Builtin.MkPair {tag = 0} [cons] {e:6} {e:7} => Builtin.MkPair {tag = 0} [cons] (f {e:2}) (Builtin.MkPair {tag = 0} [cons] (f {e:6}) (f {e:7}))}
  }

Rationale: I'm trying to use DIY-partial evaluation using higher-order-functions. By doing the case-splitting ahead of time, I return the correct function so that when I call this function 3 times at runtime, I'm not repeating the case-splitting.

[Z-snails]

I imagine this is caused by the following optimisation https://github.com/idris-lang/Idris2/blob/main/src/Compiler/CaseOpts.idr#L17-L31

This optimisation can have performance benefits, if for example the function is never applied, but equally, as you have mentioned it can cause slow-downs

[gallais]

This is the exepected behaviour following the case-of-lambdas optimisation:

https://github.com/idris-lang/Idris2/blob/6b38592b5ac12bb45d3fa825e143b3e838c064eb/src/Compiler/CaseOpts.idr#L17-L31

According to Edwin this helps chez producing faster code in general.

[ohad]

Can I disable the case-of-lambdas optimisation?

[gallais]

Oops, did not see zoe's original answer. I don't think you can disable the optimisation atm.

[ohad]

@edwinb suggested as a work-around to return a lazy function in each case. Unfortunately, #1066 blocks this workaround. I suggest we keep this issue open until we resolve #1066, and if the work-around works, close both issues.

Where would be a good place to document this in that case?

[iacore]

I'm using Idris 0.6.0-a84a5a32d (old version). Doesn't seem to be optimization.

Main> :di Ex2
Main.Ex2
Arguments [{arg:0}]
Compile time tree: let f = versionB {arg:0}
                   in \{lamc:0} => let (x, (y, z)) = {lamc:0} in (f x, (f y, f z))
Detaggable arg types: [1]
Compiled: \ {arg:0} => let f = \{eta:0}=> Main.versionB {arg:0} {eta:0} in
\{lamc:0}=> case {lamc:0} of
  { Builtin.MkPair {tag = 0} [cons]
  {e:2}
  {e:3}
  => case {e:3} of
       { Builtin.MkPair {tag = 0} [cons]
       {e:6}
       {e:7}
       => Builtin.MkPair {tag = 0} [cons] (f {e:2}) (Builtin.MkPair {tag = 0} [cons] (f {e:6}) (f {e:7}))
       }
  }
Refers to: Main.case block in Ex2, Main.{_:860}, Main.versionB, Builtin.Pair, Prelude.Types.Nat
Refers to (runtime): Main.versionB, Builtin.MkPair
Flags: covering
Size change: Main.versionB: [Just (0, Same), Nothing]

[ohad]

Thanks, but I'd rather use a new version if possible. According to the blame the actual lifting was introduced 2 years ago. It's good to know we don't have to revert 2-years of commits just to avoid this optimisation.