Skip to content

Commit b2c29cd

Browse files
author
DengYun
committed
简化调用过程，支持函数、参数、返回值不带额外Cost
1 parent 62f08e9 commit b2c29cd

2 files changed

Lines changed: 105 additions & 55 deletions

File tree

‎Algorithm/TimeCost/Def.lean‎

Lines changed: 5 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -13,22 +13,16 @@ namespace Algorithm
1313
def WithCost.wrap (a: α): WithCost α :=
1414
⟨0, a⟩
1515

16+
@[simp]
17+
def WithCost.andThen (a: WithCost α)(f: α → WithCost β): WithCost β :=
18+
let r := f a.val
19+
⟨a.cost + r.cost, r.val⟩
20+
1621
/--
1722
为一个结果附加额外的代价
1823
-/
1924
@[simp]
2025
def WithCost.addCost (a: WithCost α)(cost: Nat) : WithCost α :=
2126
⟨a.cost + cost, a.val⟩
2227

23-
/--
24-
调用一个包装函数
25-
-/
26-
@[simp]
27-
def WithCost.apply (f: WithCost (α → WithCost β))(a: WithCost α): WithCost β :=
28-
(f.val a.val).addCost (f.cost + a.cost + 1)
29-
30-
@[simp]
31-
def WithCost.applyF{α β: Type} (f: WithCost (α → β))(a: WithCost α): WithCost (β) :=
32-
WithCost.mk (f.cost + a.cost + 1) (f.val a.val)
33-
3428
end Algorithm

‎Algorithm/TimeCost/Elab.lean‎

Lines changed: 100 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -7,18 +7,45 @@ open Expr
77

88
def ZeroExpr := Expr.lit (.natVal 0)
99

10+
def Nat.succ.withCost := Nat.succ
11+
12+
def Nat.add.withCost := Nat.add
13+
14+
namespace Algorithm
15+
16+
def isWithCostType(t: Expr): Bool :=
17+
t.isApp ∧ t.appFn!.isConst ∧ t.appFn!.constName! = `Algorithm.WithCost
18+
19+
def getValueType(t: Expr): Expr :=
20+
if isWithCostType t then
21+
t.getArg! 0
22+
else
23+
t
24+
1025
/--
1126
expr: 当前的表达式分支
1227
prevCost: 表示之前let语句中所含所有开销的变量
1328
-/
14-
partial def exprToWithCost(expr prevCost: Expr): MetaM Expr := do
29+
partial def exprToWithCost(expr: Expr)(prevCost: Option Expr): MetaM Expr := do
1530
let type ← inferType expr
1631
logInfo m!"Working: {expr} with type {type}"
1732

1833
let withPrevCost(expr: Expr): MetaM Expr := do
19-
let type ← inferType expr
20-
return Expr.app
21-
(Expr.app (Expr.app (Expr.const `Algorithm.WithCost.mk []) type) prevCost) expr
34+
match prevCost with
35+
| none => return expr
36+
| some prevCost => do
37+
let type ← inferType expr
38+
return Expr.app
39+
(Expr.app (Expr.app (Expr.const `Algorithm.WithCost.mk []) type) prevCost) expr
40+
41+
let addPrevCost(expr: Expr): MetaM Expr := do
42+
match prevCost with
43+
| none => return expr
44+
| some prevCost => do
45+
let type ← inferType expr
46+
let valueType := getValueType type
47+
return Expr.app
48+
(Expr.app (Expr.app (Expr.const `Algorithm.WithCost.addCost []) valueType) expr) prevCost
2249

2350
match expr with
2451
| fvar _ =>
@@ -37,9 +64,7 @@ partial def exprToWithCost(expr prevCost: Expr): MetaM Expr := do
3764
return ← withPrevCost newExpr
3865

3966
-- 已经包装好的函数，附加当前cost
40-
let valueType := newExprType.getArg! 0
41-
let addCost := Expr.app (Expr.const `Algorithm.WithCost.addCost []) valueType
42-
return Expr.app (Expr.app addCost newExpr) prevCost
67+
return ← addPrevCost newExpr
4368
logInfo m!"Ret"
4469
withPrevCost expr
4570
| lam name type body bi =>
@@ -48,7 +73,7 @@ partial def exprToWithCost(expr prevCost: Expr): MetaM Expr := do
4873
(← withLocalDecl name bi type
4974
(fun fvar => do
5075
let body := body.instantiate1 fvar
51-
let ret ← exprToWithCost body ZeroExpr
76+
let ret ← exprToWithCost body none
5277
let ret := ret.liftLooseBVars 0 1
5378
let ret := ret.replaceFVar fvar $ Expr.bvar 0
5479
return ret
@@ -57,36 +82,71 @@ partial def exprToWithCost(expr prevCost: Expr): MetaM Expr := do
5782
bi)
5883
withPrevCost newExpr
5984
| app f a =>
60-
let newF ← exprToWithCost f ZeroExpr
61-
let newA ← exprToWithCost a ZeroExpr
62-
63-
logInfo m!"app {f} {a}"
85+
let newF ← exprToWithCost f none
86+
let newA ← exprToWithCost a none
6487
let newFType ← inferType newF
6588
let newAType ← inferType newA
66-
67-
let fRetType := (newFType.getArg! 0).bindingBody!
68-
69-
-- newF 的返回值类型可能是 WithCost f 也可能是 f（多元函数），支持f的情况
70-
if fRetType.isArrow then
71-
let newRetType := fRetType
72-
logInfo m!"newFType: {newFType}"
73-
let newExpr := Expr.app (Expr.const `Algorithm.WithCost.applyF []) (newAType.getArg! 0)
74-
let newExpr := Expr.app newExpr newRetType
75-
let newExpr := Expr.app newExpr newF
76-
let newExpr := Expr.app newExpr newA
77-
let newExprType ← inferType newExpr
78-
let addCost := Expr.app (Expr.const `Algorithm.WithCost.addCost []) (newExprType.getArg! 0)
79-
return Expr.app (Expr.app addCost newExpr) prevCost
89+
let fValType := getValueType newFType
90+
let retType := fValType.bindingBody!
91+
let retValType := getValueType retType
92+
93+
logInfo m!"retValType {retValType}"
94+
95+
-- 最内层的调用
96+
let genApply(f: Expr)(a: Expr): Expr :=
97+
let app := Expr.app f a
98+
99+
if isWithCostType retType then
100+
-- 带类型调用，增加消耗 ret.addCost 1
101+
let retValueType := getValueType retType
102+
let addCost := Expr.const `Algorithm.WithCost.addCost []
103+
let addCost := Expr.app addCost retValueType
104+
let addCost := Expr.app addCost app
105+
Expr.app addCost (Expr.lit (.natVal 1))
106+
else
107+
-- 不带类型调用，构造消耗 {cost:= 1, val:= ret}
108+
let mkCost := Expr.const `Algorithm.WithCost.mk []
109+
let mkCost := Expr.app mkCost retType
110+
let mkCost := Expr.app mkCost (Expr.lit (.natVal 1))
111+
Expr.app mkCost app
112+
113+
-- 包装Arg
114+
let genApply1(f: Expr): Expr :=
115+
if isWithCostType newAType then
116+
-- a.andThen λa =>...
117+
let AValueType := getValueType newAType
118+
-- f a
119+
let expr := genApply f (Expr.bvar 0)
120+
-- λa => f a
121+
let f := Expr.lam `a AValueType expr .default
122+
-- a.andThen
123+
let expr := Expr.const `Algorithm.WithCost.andThen []
124+
let expr := Expr.app expr AValueType
125+
let andThen := Expr.app expr retValType
126+
127+
let expr := Expr.app andThen newA
128+
Expr.app expr f
129+
else
130+
-- 直接调用
131+
genApply f newA
132+
133+
if isWithCostType newFType then
134+
-- f.andThen λf => ...
135+
-- f a
136+
let expr := genApply1 (Expr.bvar 0)
137+
-- λa => f a
138+
let f := Expr.lam `a fValType expr .default
139+
-- a.andThen
140+
let expr := Expr.const `Algorithm.WithCost.andThen []
141+
let expr := Expr.app expr fValType
142+
let andThen := Expr.app expr retValType
143+
144+
let expr := Expr.app andThen newF
145+
return Expr.app expr f
80146
else
81-
let newRetType := fRetType.getArg! 0
82-
logInfo m!"newFType: {newFType}"
83-
let newExpr := Expr.app (Expr.const `Algorithm.WithCost.apply []) (newAType.getArg! 0)
84-
let newExpr := Expr.app newExpr newRetType
85-
let newExpr := Expr.app newExpr newF
86-
let newExpr := Expr.app newExpr newA
87-
let newExprType ← inferType newExpr
88-
let addCost := Expr.app (Expr.const `Algorithm.WithCost.addCost []) (newExprType.getArg! 0)
89-
return Expr.app (Expr.app addCost newExpr) prevCost
147+
-- 直接调用
148+
return genApply1 newF
149+
90150
| _ =>
91151
throwError m!"不支持的表达式类型：{expr}"
92152

@@ -104,7 +164,7 @@ elab "#autogen_fun_with_cost" declName:ident : command => do
104164
match info with
105165
| ConstantInfo.defnInfo val =>
106166
liftTermElabM <| do
107-
let newDef := ← exprToWithCost val.value ZeroExpr
167+
let newDef := ← exprToWithCost val.value none
108168
let newDefType := ← inferType newDef
109169
logInfo m!"新定义：{newDef}"
110170
let decl := .defnDecl {
@@ -123,16 +183,12 @@ elab "#autogen_fun_with_cost" declName:ident : command => do
123183
catch e =>
124184
throwError m!"失败：{e.toMessageData}"
125185

126-
open Algorithm
127-
128-
def Nat.succ.withCost(a: Nat) := WithCost.wrap $ Nat.succ a
129-
130-
def Nat.add.withCost(a b: Nat) := WithCost.wrap $ Nat.add a b
131-
132186
def test(a b: Nat) := Nat.add a b
133187

134-
#autogen_fun_with_cost test
188+
#autogen_fun_with_cost Algorithm.test
135189

136190
#print test.withCost
137191

138-
#eval (test.withCost.apply (WithCost.wrap 4)).apply (WithCost.wrap 2)
192+
#eval test.withCost 2 4
193+
194+
end Algorithm

0 commit comments

Comments
 (0)