/
Control.lean
198 lines (167 loc) · 7.45 KB
/
Control.lean
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
/-
Copyright (c) 2019 Microsoft Corporation. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Author: Leonardo de Moura
-/
prelude
import Init.Control.Basic
import Init.Data.List.Basic
namespace List
universe u v w u₁ u₂
/-!
Remark: we can define `mapM`, `mapM₂` and `forM` using `Applicative` instead of `Monad`.
Example:
```
def mapM {m : Type u → Type v} [Applicative m] {α : Type w} {β : Type u} (f : α → m β) : List α → m (List β)
| [] => pure []
| a::as => List.cons <$> (f a) <*> mapM as
```
However, we consider `f <$> a <*> b` an anti-idiom because the generated code
may produce unnecessary closure allocations.
Suppose `m` is a `Monad`, and it uses the default implementation for `Applicative.seq`.
Then, the compiler expands `f <$> a <*> b <*> c` into something equivalent to
```
(Functor.map f a >>= fun g_1 => Functor.map g_1 b) >>= fun g_2 => Functor.map g_2 c
```
In an ideal world, the compiler may eliminate the temporary closures `g_1` and `g_2` after it inlines
`Functor.map` and `Monad.bind`. However, this can easily fail. For example, suppose
`Functor.map f a >>= fun g_1 => Functor.map g_1 b` expanded into a match-expression.
This is not unreasonable and can happen in many different ways, e.g., we are using a monad that
may throw exceptions. Then, the compiler has to decide whether it will create a join-point for
the continuation of the match or float it. If the compiler decides to float, then it will
be able to eliminate the closures, but it may not be feasible since floating match expressions
may produce exponential blowup in the code size.
Finally, we rarely use `mapM` with something that is not a `Monad`.
Users that want to use `mapM` with `Applicative` should use `mapA` instead.
-/
@[inline]
def mapM {m : Type u → Type v} [Monad m] {α : Type w} {β : Type u} (f : α → m β) (as : List α) : m (List β) :=
let rec @[specialize] loop
| [], bs => pure bs.reverse
| a :: as, bs => do loop as ((← f a)::bs)
loop as []
@[specialize]
def mapA {m : Type u → Type v} [Applicative m] {α : Type w} {β : Type u} (f : α → m β) : List α → m (List β)
| [] => pure []
| a::as => List.cons <$> f a <*> mapA f as
@[specialize]
protected def forM {m : Type u → Type v} [Monad m] {α : Type w} (as : List α) (f : α → m PUnit) : m PUnit :=
match as with
| [] => pure ⟨⟩
| a :: as => do f a; List.forM as f
@[specialize]
def forA {m : Type u → Type v} [Applicative m] {α : Type w} (as : List α) (f : α → m PUnit) : m PUnit :=
match as with
| [] => pure ⟨⟩
| a :: as => f a *> forA as f
@[specialize]
def filterAuxM {m : Type → Type v} [Monad m] {α : Type} (f : α → m Bool) : List α → List α → m (List α)
| [], acc => pure acc
| h :: t, acc => do
let b ← f h
filterAuxM f t (cond b (h :: acc) acc)
@[inline]
def filterM {m : Type → Type v} [Monad m] {α : Type} (f : α → m Bool) (as : List α) : m (List α) := do
let as ← filterAuxM f as []
pure as.reverse
@[inline]
def filterRevM {m : Type → Type v} [Monad m] {α : Type} (f : α → m Bool) (as : List α) : m (List α) :=
filterAuxM f as.reverse []
@[inline]
def filterMapM {m : Type u → Type v} [Monad m] {α β : Type u} (f : α → m (Option β)) (as : List α) : m (List β) :=
let rec @[specialize] loop
| [], bs => pure bs
| a :: as, bs => do
match (← f a) with
| none => loop as bs
| some b => loop as (b::bs)
loop as.reverse []
@[specialize]
protected def foldlM {m : Type u → Type v} [Monad m] {s : Type u} {α : Type w} : (f : s → α → m s) → (init : s) → List α → m s
| _, s, [] => pure s
| f, s, a :: as => do
let s' ← f s a
List.foldlM f s' as
@[inline]
def foldrM {m : Type u → Type v} [Monad m] {s : Type u} {α : Type w} (f : α → s → m s) (init : s) (l : List α) : m s :=
l.reverse.foldlM (fun s a => f a s) init
@[specialize]
def firstM {m : Type u → Type v} [Alternative m] {α : Type w} {β : Type u} (f : α → m β) : List α → m β
| [] => failure
| a::as => f a <|> firstM f as
@[specialize]
def anyM {m : Type → Type u} [Monad m] {α : Type v} (f : α → m Bool) : List α → m Bool
| [] => pure false
| a::as => do
match (← f a) with
| true => pure true
| false => anyM f as
@[specialize]
def allM {m : Type → Type u} [Monad m] {α : Type v} (f : α → m Bool) : List α → m Bool
| [] => pure true
| a::as => do
match (← f a) with
| true => allM f as
| false => pure false
@[specialize]
def findM? {m : Type → Type u} [Monad m] {α : Type} (p : α → m Bool) : List α → m (Option α)
| [] => pure none
| a::as => do
match (← p a) with
| true => pure (some a)
| false => findM? p as
@[specialize]
def findSomeM? {m : Type u → Type v} [Monad m] {α : Type w} {β : Type u} (f : α → m (Option β)) : List α → m (Option β)
| [] => pure none
| a::as => do
match (← f a) with
| some b => pure (some b)
| none => findSomeM? f as
@[inline] protected def forIn {α : Type u} {β : Type v} {m : Type v → Type w} [Monad m] (as : List α) (init : β) (f : α → β → m (ForInStep β)) : m β :=
let rec @[specialize] loop
| [], b => pure b
| a::as, b => do
match (← f a b) with
| ForInStep.done b => pure b
| ForInStep.yield b => loop as b
loop as init
instance : ForIn m (List α) α where
forIn := List.forIn
@[simp] theorem forIn_nil [Monad m] (f : α → β → m (ForInStep β)) (b : β) : forIn [] b f = pure b :=
rfl
@[simp] theorem forIn_cons [Monad m] (f : α → β → m (ForInStep β)) (a : α) (as : List α) (b : β)
: forIn (a::as) b f = f a b >>= fun | ForInStep.done b => pure b | ForInStep.yield b => forIn as b f :=
rfl
@[inline] protected def forIn' {α : Type u} {β : Type v} {m : Type v → Type w} [Monad m] (as : List α) (init : β) (f : (a : α) → a ∈ as → β → m (ForInStep β)) : m β :=
let rec @[specialize] loop : (as' : List α) → (b : β) → Exists (fun bs => bs ++ as' = as) → m β
| [], b, _ => pure b
| a::as', b, h => do
have : a ∈ as := by
have ⟨bs, h⟩ := h
subst h
exact mem_append_of_mem_right _ (Mem.head ..)
match (← f a this b) with
| ForInStep.done b => pure b
| ForInStep.yield b =>
have : Exists (fun bs => bs ++ as' = as) := have ⟨bs, h⟩ := h; ⟨bs ++ [a], by rw [← h, append_cons bs a as']⟩
loop as' b this
loop as init ⟨[], rfl⟩
instance : ForIn' m (List α) α inferInstance where
forIn' := List.forIn'
@[simp] theorem forIn'_eq_forIn {α : Type u} {β : Type v} {m : Type v → Type w} [Monad m] (as : List α) (init : β) (f : α → β → m (ForInStep β)) : forIn' as init (fun a _ b => f a b) = forIn as init f := by
simp [forIn', forIn, List.forIn, List.forIn']
have : ∀ cs h, List.forIn'.loop cs (fun a _ b => f a b) as init h = List.forIn.loop f as init := by
intro cs h
induction as generalizing cs init with
| nil => intros; rfl
| cons a as ih => intros; simp [List.forIn.loop, List.forIn'.loop, ih]
apply this
instance : ForM m (List α) α where
forM := List.forM
@[simp] theorem forM_nil [Monad m] (f : α → m PUnit) : forM [] f = pure ⟨⟩ :=
rfl
@[simp] theorem forM_cons [Monad m] (f : α → m PUnit) (a : α) (as : List α) : forM (a::as) f = f a >>= fun _ => forM as f :=
rfl
instance : Functor List where
map := List.map
end List