Optimizing the computation of `Fin` / `ZMod`

Stream: general

[Yuyang Zhao]

I'm trying to optimize the computation for large Fin. But @[csimp] doesn't seem to work here. I realized that I need to redefine instance. Is there any way to make @[csimp] work automatically for all functions (I know there may be a lot) that use Fin.add? Also this implementation is slower for small numbers, is this because if-then-else in lean is slower?

import Std.Data.Fin.Basic

namespace Fin

def addImpl : Fin n → Fin n → Fin n
  | ⟨a, ha⟩, ⟨b, hb⟩ =>
    let x := a + b
    if hx : x < n then
      ⟨x, hx⟩
    else
      ⟨x - n, Nat.sub_lt_left_of_lt_add (Nat.not_lt.mp hx) (Nat.add_lt_add ha hb)⟩

@[csimp]
theorem add_eq_addImpl : @Fin.add = @addImpl :=
  funext fun n ↦ funext fun ⟨a, ha⟩ ↦ funext fun ⟨b, hb⟩ ↦ by
    dsimp [addImpl, Fin.add]
    split <;> (rename_i hx; congr)
    · exact Nat.mod_eq_of_lt hx
    · rw [Nat.not_lt] at hx
      rw [Nat.mod_eq, if_pos ⟨Nat.lt_of_le_of_lt (Nat.zero_le a) ha, hx⟩]
      exact Nat.mod_eq_of_lt (Nat.sub_lt_left_of_lt_add hx (Nat.add_lt_add ha hb))

-- instance (priority := default + 1) : Add (Fin n) where add := Fin.add

def run (n : Nat) (x : Nat) : IO Unit := do
  let mut t : Fin (n + 1) := 0
  for _ in [0 : x] do
    t := ⟨n, n.lt_succ_self⟩ + ⟨n, n.lt_succ_self⟩

def runAdd (n : Nat) (x : Nat) : IO Unit := do
  let mut t : Fin (n + 1) := 0
  for _ in [0 : x] do
    t := Fin.add ⟨n, n.lt_succ_self⟩ ⟨n, n.lt_succ_self⟩

def runImpl (n : Nat) (x : Nat) : IO Unit := do
  let mut t : Fin (n + 1) := 0
  for _ in [0 : x] do
    t := addImpl ⟨n, n.lt_succ_self⟩ ⟨n, n.lt_succ_self⟩

def n := 10^100000
def x := 100000

#eval timeit "run" (run n x)
#eval timeit "runAdd" (runAdd n x)
#eval timeit "runImpl" (runImpl n x)

[Eric Wieser]

I think the csimp needs to be defined before anything else uses the function

[Eric Wieser]

Otherwise it's too late because Fin.instAdd has already been compiled, and future csimps won't force it to be recompiled