Speeding up large data structures in Lean

Question

Consider a simple inductive data structure of binary trees, and a type of pairs of integers:

structure Edge : Type :=
  edge::
  fst : Nat
  snd : Nat

inductive SimpleTree (α : Type) : Type
  | empty : SimpleTree α
  | leaf : α → SimpleTree α
  | node : α → SimpleTree α → SimpleTree α → SimpleTree α

I would like to build trees of edges with a couple of thousand nodes. A small example would look like this:

example : SimpleTree Edge := .node (.edge 5 0) .empty (.leaf (.edge 3 1))

I placed a larger example with around 1200 nodes in example.lean at this gist repository. You will find there comparable Coq, Agda, and OCaml implementations, as well as the variants shown below. Processing these files on my laptop (see time.sh in the gist) gives the following timings:

ocamlc -c example.ml    0.07s user    0.01s system     95% cpu     0.087 total
lean example.lean       5.16s user    0.05s system     99% cpu     5.229 total
agda example.agda       3.63s user    0.22s system    100% cpu     3.847 total
coqc example.v          2.55s user    0.19s system     99% cpu     2.737 total

Lean is the slowest, contrary to its fame. I experimented with @[inline] but that doesn't seem to help with a one-time compilation. Am I doing something wrong?

Things get worse if we use dependent types. Suppose we define edges so that the first component is known to be in Fin n, and the second must be smaller than the first:

structure Edge (n : Nat) : Type :=
  edge::
  fst : Fin n
  snd : Fin fst

-- smart constructor for edges
macro "E[" i:term "," j:term "]" : term =>
 `(.edge $i (Fin.mk $j (by trivial)))

A small example would now look like this:

def test1 : Tree (Edge 100) := .node E[5,0] .empty (.leaf E[3,1])

When we run the same tree, but with the dependent type of edges, we get these timings:

lean exampleDependentNoTactic.lean  18.75s user   0.10s system     99% cpu    18.886 total
lean exampleDependentTactic.lean    15.07s user   0.38s system     99% cpu    15.471 total
agda exampleDependent.agda          5.39s user    0.24s system    100% cpu     5.625 total
coqc exampleDependent.v             7.74s user    0.25s system     99% cpu     7.998 total

Once again, Lean is quite a bit slower than Agda and Coq.

(In above benchmarks there are two versions of Lean, one uses the trivial tactic to constructs proof terms, while the other one avoids doing so by boolean reflection, see exampleDependentNoTactic.lean.)

Eventually I would like to process thousands of such trees, as well as other similar data structures. I would appreciate advice on how to speed things up in Lean. If you have advice on how to speed up Agda and Coq, it's welcome, too.

Answer 1

(This question has since been solved on Zulip, so I'm copying the resolution here for posterity.)

This is an implementation which reads the tree via an external .json file. The main function tree_from_file is a user command that first reads the tree, parses it to Json, then calls mkTree and mkEdge to construct an Expr (using the Qq library to call the appropriate expr constructors), and hand the result directly to the kernel, skipping elaboration and linting (which turn out to be the biggest costs in this example).

The result is something that takes about 0.05 seconds (only counting the tree_from_file line at the end, not the startup costs which are in the neighborhood of 0.35 seconds), competitive with the unverified OCaml version (even though this is constructing the full dependent types version and pushing everything through the kernel at the end).

import Lean
import Qq

structure Edge (n : Nat) : Type :=
  edge ::
  fst : Fin n
  snd : Fin fst
open Edge

inductive Tree (α : Type) : Type
  | empty : Tree α
  | leaf : α → Tree α
  | node : α → Tree α → Tree α → Tree α

def Edge.mk' (n a b : Nat) (H : Nat.blt a n = true) (H2 : Nat.blt b a = true) : Edge n :=
  ⟨⟨a, by simp_all⟩, ⟨b, (by simp_all : b < a)⟩⟩

open Qq Lean Meta Elab

def mkEdge (n : Q(Nat)) (j : Json) : Except String Q(Edge $n) := do
  let arr ← j.getArr?
  have a : Q(Nat) := mkRawNatLit (← arr[0]!.getNat?)
  have b : Q(Nat) := mkRawNatLit (← arr[1]!.getNat?)
  have H : Q(Nat.blt $a $n = true) := (q(Eq.refl true) : Expr)
  have H2 : Q(Nat.blt $b $a = true) := (q(Eq.refl true) : Expr)
  pure q(Edge.mk' $n $a $b $H $H2)

partial def mkTree (n : Q(Nat)) (j : Json) : Except String Q(Tree (Edge $n)) := do
  let arr ← j.getArr?
  match arr with
  | #[e, l, r] => pure q(Tree.node $(← mkEdge n e) $(← mkTree n l) $(← mkTree n r))
  | #[] => pure q(Tree.empty)
  | _ => pure q(Tree.leaf $(← mkEdge n j))

elab "tree_from_file" name:ident " from " file:str : command => do
  let file ← IO.FS.readFile file.getString
  let result : Except String (Q(Nat) × Expr) := do
    let .arr #[n, tree] ← Json.parse file | throw "bad format"
    let n : Q(Nat) := mkRawNatLit (← n.getNat?)
    pure (n, ← mkTree n tree)
  let .ok (n, value) := result | throwError "failed to parse"
  Elab.Command.liftCoreM <| addDecl <| .defnDecl {
    name := name.getId
    levelParams := []
    type := q(Tree (Edge $n))
    value
    hints := .regular 0
    safety := .safe
  }

tree_from_file test from "test.json"
#check test -- test : Tree (Edge 666)

Answer 2

This run-time solution uses Lean's built-in ToJson and FromJson derive handlers to reduce boilerplate. (I added in extra derive handlers that I didn't need but the OP might find useful.) It reads in the structure, one structure/json for each line of the file. Note, I don't think I'm converting from Except to IO in the best way. Also, I haven't timed this on a big tree.

import Lean.Data.Json

structure Edge : Type :=
  edge::
  fst : Nat
  snd : Nat
  deriving Inhabited, BEq, Ord, Hashable, Repr, Lean.ToJson, Lean.FromJson

inductive SimpleTree (α : Type) : Type
  | empty : SimpleTree α
  | leaf : α → SimpleTree α
  | node : α → SimpleTree α → SimpleTree α → SimpleTree α
  deriving Inhabited, BEq, Ord, Hashable, Repr, Lean.ToJson, Lean.FromJson

def example1 : SimpleTree Edge := .node (.edge 5 0) .empty (.leaf (.edge 3 1))

#eval Lean.ToJson.toJson example1 -- {"node": [{"snd": 0, "fst": 5}, "empty", {"leaf": {"snd": 1, "fst": 3}}]}
#eval Lean.Json.parse "{\"node\": [{\"snd\": 0, \"fst\": 5}, \"empty\", {\"leaf\": {\"snd\": 1, \"fst\": 3}}]}"
-- ok: {"node": [{"snd": 0, "fst": 5}, "empty", {"leaf": {"snd": 1, "fst": 3}}]}

def SimpleTree.fromjsonString (s : String) : Except String (SimpleTree Edge) := do 
  let json <- Lean.Json.parse s
  let simpleTree : SimpleTree Edge <- Lean.FromJson.fromJson? json
  return simpleTree

#eval SimpleTree.fromjsonString "{\"node\": [{\"snd\": 0, \"fst\": 5}, \"empty\", {\"leaf\": {\"snd\": 1, \"fst\": 3}}]}"
-- Except.ok (SimpleTree.node { fst := 5, snd := 0 } (SimpleTree.empty) (SimpleTree.leaf { fst := 3, snd := 1 }))

def parseSimpleTree (s : String) : IO (SimpleTree Edge) := do 
  let simpleTree? := (SimpleTree.fromjsonString s)
  match simpleTree? with
  | .ok simpleTree => return simpleTree
  | .error e => panic! e

#eval parseSimpleTree "{\"node\": [{\"snd\": 0, \"fst\": 5}, \"empty\", {\"leaf\": {\"snd\": 1, \"fst\": 3}}]}"
-- SimpleTree.node { fst := 5, snd := 0 } (SimpleTree.empty) (SimpleTree.leaf { fst := 3, snd := 1 })

def main : IO Unit := do
  let trees <- readSimpleTreesFromFile "input.txt"
  for t in trees do
    IO.println (repr t)

Run with lean --run example.lean to run the main function.

I don't think this works with dependent types currently.