lean4-htt/tests/compiler/t4.lean
2021-01-27 14:45:31 +01:00

116 lines
3.4 KiB
Text

/- Benchmark for new code generator -/
inductive Expr
| Val : Int → Expr
| Var : String → Expr
| Add : Expr → Expr → Expr
| Mul : Expr → Expr → Expr
| Pow : Expr → Expr → Expr
| Ln : Expr → Expr
namespace Expr
protected def Expr.toString : Expr → String
| Val n => toString n
| Var x => x
| Add f g => "(" ++ Expr.toString f ++ " + " ++ Expr.toString g ++ ")"
| Mul f g => "(" ++ Expr.toString f ++ " * " ++ Expr.toString g ++ ")"
| Pow f g => "(" ++ Expr.toString f ++ " ^ " ++ Expr.toString g ++ ")"
| Ln f => "ln(" ++ Expr.toString f ++ ")"
instance : ToString Expr :=
⟨Expr.toString⟩
partial def pown : Int → Int → Int
| a, 0 => 1
| a, 1 => a
| a, n =>
let b := pown a (n / 2);
b * b * (if n % 2 = 0 then 1 else a)
partial def addAux : Expr → Expr → Expr
| Val n, Val m => Val (n + m)
| Val 0, f => f
| f, Val 0 => f
| f, Val n => addAux (Val n) f
| Val n, Add (Val m) f => addAux (Val (n+m)) f
| f, Add (Val n) g => addAux (Val n) (addAux f g)
| Add f g, h => addAux f (addAux g h)
| f, g => Add f g
def add (a b : Expr) : Expr :=
-- dbgTrace (">> add (" ++ toString a ++ ", " ++ toString b ++ ")") $ fun _ =>
addAux a b
-- set_option trace.compiler.borrowed_inference true
partial def mulAux : Expr → Expr → Expr
| Val n, Val m => Val (n*m)
| Val 0, _ => Val 0
| _, Val 0 => Val 0
| Val 1, f => f
| f, Val 1 => f
| f, Val n => mulAux (Val n) f
| Val n, Mul (Val m) f => mulAux (Val (n*m)) f
| f, Mul (Val n) g => mulAux (Val n) (mulAux f g)
| Mul f g, h => mulAux f (mulAux g h)
| f, g => Mul f g
def mul (a b : Expr) : Expr :=
-- dbgTrace (">> mul (" ++ toString a ++ ", " ++ toString b ++ ")") $ fun _ =>
mulAux a b
def pow : Expr → Expr → Expr
| Val m, Val n => Val (pown m n)
| _, Val 0 => Val 1
| f, Val 1 => f
| Val 0, _ => Val 0
| f, g => Pow f g
def ln : Expr → Expr
| Val 1 => Val 0
| f => Ln f
def d (x : String) : Expr → Expr
| Val _ => Val 0
| Var y => if x = y then Val 1 else Val 0
| Add f g => add (d x f) (d x g)
| Mul f g =>
-- dbgTrace (">> d (" ++ toString f ++ ", " ++ toString g ++ ")") $ fun _ =>
add (mul f (d x g)) (mul g (d x f))
| Pow f g => mul (pow f g) (add (mul (mul g (d x f)) (pow f (Val (-1)))) (mul (ln f) (d x g)))
| Ln f => mul (d x f) (pow f (Val (-1)))
def count : Expr → Nat
| Val _ => 1
| Var _ => 1
| Add f g => count f + count g
| Mul f g => count f + count g
| Pow f g => count f + count g
| Ln f => count f
def nestAux (s : Nat) (f : Nat → Expr → IO Expr) : Nat → Expr → IO Expr
| 0, x => pure x
| m@(n+1), x => f (s - m) x >>= nestAux s f n
def nest (f : Nat → Expr → IO Expr) (n : Nat) (e : Expr) : IO Expr :=
nestAux n f n e
def deriv (i : Nat) (f : Expr) : IO Expr :=
do
let d := d "x" f;
IO.println (toString (i+1) ++ " count: " ++ (toString $ count d));
IO.println (toString d);
pure d
end Expr
def main (xs : List String) : IO UInt32 :=
do let x := Expr.Var "x";
let f := Expr.add x (Expr.mul x (Expr.mul x (Expr.add x x)));
IO.println f;
discard <| Expr.nest Expr.deriv 3 f;
pure 0
-- setOption profiler True
-- #eval main []