Skip to content

Commit 4c5a43b

Browse files
author
DengYun
committed
时间复杂度计算框架
1 parent 2457ddc commit 4c5a43b

4 files changed

Lines changed: 93 additions & 48 deletions

File tree

‎Algorithm/TimeCost.lean‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1 +1,2 @@
1+
import Algorithm.TimeCost.Def
12
import Algorithm.TimeCost.Elab

‎Algorithm/TimeCost/Def.lean‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,14 @@
1+
import Mathlib.Data.FunLike.Basic
2+
3+
namespace Algorithm
4+
5+
structure WithCost(α: Type) where
6+
cost: Nat
7+
val: α
8+
9+
@[simp]
10+
def WithCost.apply {f}[FunLike f α (WithCost β)] (self: WithCost f)(arg: α): WithCost β :=
11+
let ⟨cost1, ret⟩:= self.val arg
12+
⟨self.cost + cost1, ret⟩
13+
14+
end Algorithm

‎Algorithm/TimeCost/Elab.lean‎

Lines changed: 70 additions & 47 deletions
Original file line numberDiff line numberDiff line change
@@ -1,57 +1,80 @@
11
import Lean
2+
import Algorithm.TimeCost.Def
23

3-
open Lean Meta Elab Term
4-
5-
syntax (name := time_cost_term)(priority := 0) "time_cost" term : term
6-
7-
mutual
8-
9-
partial def time_cost_for_syntax: Syntax → TermElabM Expr := fun stx =>
10-
withIncRecDepth do
11-
match stx with
12-
| `($_:num) => return mkNatLit 0
13-
| `($_:ident) => return mkNatLit 0
14-
| `($a + $b) =>
15-
let acost ← time_cost_for_syntax a
16-
let bcost ← time_cost_for_syntax b
17-
return ← mkAdd (mkNatLit 1) (← mkAdd acost bcost)
18-
| `($f $args*) =>
19-
-- 如果函数本身是计算得出的，也需要考虑合成过程的开销
20-
let fcost ← time_cost_for_syntax f
21-
-- 函数展开的计算开销
22-
-- let callcostFun ← time_cost_fun_for_syntax f
23-
-- let argsExpr ← args.mapM (λ stx => elabTerm stx none)
24-
-- let callcost := mkAppN callcostFun argsExpr
25-
-- 参数的计算开销
26-
let argscost ← args.foldlM (init:=mkNatLit 0) (fun sum arg => do
27-
let argCost ← time_cost_for_syntax arg
28-
mkAdd sum argCost
4+
open Lean Meta Elab Term Command
5+
6+
open Expr
7+
8+
partial def exprToWithCost(expr prevCost: Expr): MetaM Expr := do
9+
let type ← inferType expr
10+
logInfo m!"Working: {expr} with type {type}"
11+
12+
match expr with
13+
| fvar _ =>
14+
let ret := Expr.app
15+
(Expr.app (Expr.app (Expr.const `Algorithm.WithCost.mk []) type) prevCost)
16+
expr
17+
logInfo m!"Ret: {ret}"
18+
return ret
19+
| lam name type body bi =>
20+
logInfo m!"lam {name} {type} {body}"
21+
return (Expr.lam name type
22+
(← withLocalDecl name bi type
23+
(fun fvar => do
24+
let body := body.instantiate1 fvar
25+
let ret ← exprToWithCost body prevCost
26+
let ret := ret.liftLooseBVars 0 1
27+
let ret := ret.replaceFVar fvar $ Expr.bvar 0
28+
return ret
29+
)
2930
)
31+
bi)
32+
-- | app a b =>
33+
-- logInfo m!"app {a} {b}"
34+
-- match a with
35+
-- | _ => throwError m!"TODO"
36+
| _ =>
37+
throwError m!"不支持的表达式类型：{expr}"
38+
39+
elab "#autogen_fun_with_cost" declName:ident : command => do
40+
let decl := declName.getId
41+
let newName := Name.str decl "withCost"
3042

31-
-- return ← mkAdd (mkNatLit 1) (← mkAdd fcost (← mkAdd argscost callcost))
32-
return ← mkAdd (mkNatLit 1) (← mkAdd fcost argscost)
33-
| `(($stx)) => time_cost_for_syntax stx
34-
| _ => throwError "Syntax {stx} not supported for time_cost"
43+
try
44+
let env ← getEnv
45+
let info := env.find? decl
46+
match info with
47+
| none =>
48+
throwError m!"失败：定义未找到"
49+
| some info =>
50+
match info with
51+
| ConstantInfo.defnInfo val =>
52+
liftTermElabM <| do
53+
let newDef := ← exprToWithCost val.value (Expr.lit (.natVal 0))
54+
let newDefType := ← inferType newDef
55+
logInfo m!"新定义：{newDef}"
56+
let decl := .defnDecl {
57+
name := newName
58+
levelParams := val.levelParams
59+
type := newDefType
60+
value := newDef
61+
safety := .safe
62+
hints := val.hints
63+
}
64+
addDecl decl
65+
compileDecl decl
66+
logInfo m!"定义：{newName}"
67+
| _ => throwError m!"失败：定义不是函数定义"
3568

36-
partial def time_cost_fun_for_syntax: Syntax → TermElabM Expr := fun stx =>
37-
withIncRecDepth do
38-
match stx with
39-
-- | _ => return mkNatLit 0
40-
| _ => throwError "Syntax {stx} not supported for time_cost_fun"
69+
catch e =>
70+
throwError m!"失败：{e.toMessageData}"
4171

42-
end
72+
def test(a b: Nat) := a
4373

44-
@[term_elab time_cost_term]
45-
def time_cost_term_impl : TermElab := fun stx type? => do
46-
match stx with
47-
| `(time_cost $stx1) => return ← time_cost_for_syntax stx1
48-
| _ => throwError "Bad syntax"
74+
#eval Algorithm.WithCost.mk 2 2
4975

50-
syntax (name := time_cost_fun_term)(priority := 0) "time_cost_fun" term : term
76+
#autogen_fun_with_cost test
5177

78+
#print test.withCost
5279

53-
@[term_elab time_cost_fun_term]
54-
def time_cost_fun_term_impl : TermElab := fun stx type? => do
55-
match stx with
56-
| `(time_cost_fun $stx1) => return ← time_cost_for_syntax stx1
57-
| _ => throwError "Bad syntax"
80+
#eval test.withCost 1 2

‎Examples/Basic/Fibonacci.lean‎

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,13 @@ def fib: ℕ → ℕ
99
| 1 => 1
1010
| n+2 => fib (n+1) + fib n
1111

12+
-- 期待通过Meta Programming自动生成这个函数：
13+
def fib.withCost : ℕ → WithCost ℕ
14+
| 0 => ⟨1, 0⟩
15+
| 1 => ⟨1, 1⟩
16+
| n+2 => fib.withCost (n+1) + fib.withCost n
17+
18+
1219
-- 正确性证明：
1320
-- f_0 = 0, f_1 = 1, f_n = f_{n-1}+f_{n-2} for n ≥ 2
1421
theorem fib_is_correct:
@@ -27,5 +34,5 @@ theorem fib_is_correct:
2734

2835
-- 时间复杂度证明
2936
theorem fib_time_complexity:
30-
(λ n => ((time_cost (fib n)) : ℝ)) =O[Filter.atTop] (λ n => ((√5 + 1) / 2) ^ n) := by
37+
(λ n => ((fib.withCost n).cost : ℝ)) =O[Filter.atTop] (λ n => ((√5 + 1) / 2) ^ n) := by
3138
sorry

0 commit comments

Comments
 (0)