Deriving Monad Comprehension

Introducing Comprehension

Basics

The notation for list comprehension is the operational interpretation of a set-builder notation.

{(x, y) | (x > 2) ∧ y < 2}

{ xy : ℕ × ℕ // xy.1 > 2 ∧ xy.2 < 2}

[(x, y) | x <- [0..], y <- [0..], x > 2 && y < 2]

(loop for x from 0
  collect
    (loop for y from 0
	    when (and (> x 2) (< y 2)) 
		    nconc (cons x y))
  into ls
  finally (return (remove-if #'null ls)))

[(x, y) for x in [0..] for y in [0..] if x > 2 and y < 2]

• 从上面 type theoretic 的定义可以看出，集合构建式说到底还是一个命题——还不具有 operational semantics. 这要我们去赋予；

• 此外其通常用于描述无穷可数集合，然而在无穷集合/序列上做计算，对于Haskell之类 lazy 的语言是较为简单的；但此外的绝大多数即 strict langauge 是不好操作的，要求一个惰性的数据结构以及相适应的方法，在这里就是我实现的 LazyList.

• 有隐含的 pattern matching (比如 pair 的 destructuring)，虽然这里没有发生：只要语言支持 pattern matching 就没有任何问题，否则应该先去做 pattern matching.

• 区分 sequential / parallel. 即嵌套的 map 和平行 map ；命令式的话就是 for 循环：嵌套其实就是计算笛卡儿积，一般用于 nondeterministic search，后者可以翻译成一个 zip 或 zipWith .

Translating by Examples: Sequential

[(x, y) | x <- [0..], y <- [0..], x > 2 && y < 2] :=
	flatten $ filter (· matches some _) $
	[0..].map fun x => 
		[0..].map fun y => 
			if x > 2 && y < 2 then some (x, y) else none

• 虽然在这里 y <- [0..] 中不含 x ，但是 Haskell 的推导式是有 sequential binding 的——也就是说 x <- e1, y <- e2,.. 中 x 在 e2 中有定义， y 在其之后的表达式中有定义。

• 单纯的嵌套 map 肯定是不行，否则类型就是 List (List (Nat × Nat)) ，需要用 flatten 展平。

• x > 2 && y < 2 这个 predicative filtering 可以把 (x, y) 即我们用来 map 的函数的 codomain 给 lift 到 Option 中，构造一个 anamorphism，说人话就是把 Nat -> Nat × Nat 转换成 Nat -> Option (Nat × Nat) .
正是以下 F-coalgebra :
• 特别地，我们希望直接构造出 ，最好是能够进行一些系统性的句法（syntactic）上的变换，因为 Lean 是 strict 的，采用复合不合适。
• 我们希望能把 和 filterMap fuse 起来，减少一次遍历
• 其中 filterMap id = filter (· matches some _)

• 性能如何？虽然函数式已经没有性能可言，但是能提升的地方当然应该提升，不妨看看遍历性能：
• 嵌套 map 两次，这个不可能没有
• flatten 一次
• filter 一次
我们发现，一个简简单单的东西居然要在一个列表上跑 4 遍，这实在是很难说得过去。
（这是因为 Lean 不是 call-by-name/need，没有办法直接复合这几个操作，Haskell 在这里的翻译就直观得多）。

Translation #1

comp ::=
	  ...
	| [$exp (where $p)? | ($pat <- $exp' in)+]

syntax "[" term " | " withoutPosition(sepBy((term (" ← " <|> " <- ") term), " in ")) "]" : term
syntax "[" term " where " term " | " withoutPosition(sepBy((term (" ← " <|> " <- ") term), " in ")) "]" : ter

性能改进

flatMap

inductive LazyList (α : Type u)
  | nil
  | cons : α -> Thunk (LazyList α) -> LazyList α
-- ... （中略）
def appendH : LazyList α -> Thunk (LazyList α) -> LazyList α
  | [||], l => l.get
  | x ::' xs, l => x ::' appendH xs.get l

def append : LazyList α -> LazyList α -> LazyList α
  | nil, l => l
  | x ::' xs, l => x ::' ↑(append xs.get l)

instance : Append (LazyList α) := ⟨append⟩
instance : HAppend (LazyList α) (Thunk $ LazyList α) (LazyList α) := ⟨appendH⟩

-- ... (中略）
def flatMap (f : α -> LazyList β) : LazyList α -> LazyList β
  | nil => nil
  | x ::' xs => (f x) ++ (Thunk.mk $ fun _ => flatMap f xs.get)

• (++) 同时 dispatch 到 homogeneous 的 Append 和 heterogeneous 的 HAppend

• (s₁ ++ s₂) 会根据 s₂ 的类型从这俩函数中选取一个。static multiple dispatch 确实是很爽。

• f x ++ flatMap f xs.get

• f x ++ (Thunk.mk fun _ => flatMap f xs.get)

• 那么这样就减少了一次遍历（实际上是 次，考虑 个列表）。

[(x, y) where x > 2 && y < 2 | x <- [0..] in y <- [0..]] :=
	filter (· matches some _) $
	[0..].flatMap fun x => 
		[0..].filterMap fun y => 
			if x > 2 && y < 2 then some (x, y) else none

构造 f' 并 fuse filterMap

• exp 会被绑定到 (x, y)

• 出现的两个 binder pat （以后也可能是 pattern）
• x 和 y
• fun pat => e 已经支持 pattern matching，我们不必更改。

• p 会被绑定到 x > 2 && y < 2

local macro f:term:50 " depends_on! " p:term:50 : term =>
  `(if $p then Option.some $f else Option.none)

• e ⇒ `($exp depends_on! Bool.true)

• e where p ⇒ `($exp depends_on! $p)

[(x, y) where x > 2 && y < 2 | x <- [0..] in y <- [0..]] :=
	[0..].flatMap fun x => 
		[0..].filterMap fun y => 
			if x > 2 && y < 2 then some (x, y) else none

Sequential Binding

Generalizing

• 所有的 flatMap 替换成 bind

• 所有的 map 替换成 (· <$> ·)

• (本质上是因为实现 Alternative 相当于定义了某个 Monad 的 zero element)

--   Extends `Functor`, with the addition of `filterMap`, for collections especially.
-- - Any `Monad` that has an `Alternative` instance is `Mappable`.
--
class Mappable (F : Type u -> Type v) extends Functor F where

  filterMap (p : α -> Option β) : F α -> F β
  
  -- Usually if there's `filterMap`, 
  -- there must be a `filter` method as well for collections. 
  -- But if it doesn't then we define `filter` in terms of `filterMap`.
  -- However a direct implementation has, very likely, better performance.
  filter (p : α -> Bool) : F α -> F α := filterMap (fun x => if p x then some x else none)
  
  map f := filterMap (some ∘ f)

@[default_instance] 
instance (priority := low) [Alternative m] [Monad m] : Mappable m where
  filterMap f xs := xs >>= fun x => 
    if let some x := f x then pure x else failure

• 这个依赖是必要的，据我所知， Haskell，至少是 ghc 的翻译，用的是
• guard p *> e 或者 e <$ guard p
• 无论哪个，用到 guard 就需要 Alternative .
• ghc 的翻译是很干净的，为什么我们做不到呢——当然是因为性能问题和 eager eval.

Denotational semantics #1

Implement #1

• 不存在 predicate 的情况（即没有 where 从句）的情况下，没有必要去构造 f' 并 filterMap ，直接 map 即可——省几个 stack frames 嘛——因此我们分开处理两种情况；

• 我们只有两条完整的语法，所以 parse 之后会得到 i, l 两个数组，分别表示 pat 和 e' ；

• 当然我们也可以多加几条语法，然后直接在 macro_rules 上递归（即前面的 denotational semantics），不过性能恐怕不如直接 fold 数组来构造AST来得好。

macro_rules
	...  
  | `([ $f | $[$i <- $l]in* ]) => do
    if h : i.size >= 1 ∧ l.size >= 1 then
      let (iz, lz) := (i[i.size - 1], l[l.size - 1])
      let li := (l.zip i)[:l.size - 1]
      show MacroM Term from
      ``(Functor.map (fun $iz => $f) $lz) >>= li.foldrM fun (t, fb) a =>
        ``(Bind.bind $t (fun $fb => $a))
    else unreachable!
  | `([ $f | $[$i ← $l]in* ]) => ``([ $f | $[$i <- $l]in* ])
  | `([ $f where $p | $[$i <- $l]in* ]) => do
    let mf <- `($f depends_on! $p)
    let (iz, lz) := (i[i.size - 1]!, l[l.size - 1]!)
    let li := (l.zip i)[:l.size - 1]
    show MacroM Term from
    ``(Mappable.filterMap (fun $iz => $mf) $lz) >>=
      li.foldrM fun (t, fb) a => ``(Bind.bind $t (fun $fb => $a))
  | `([ $f where $p | $[$i ← $l]in* ]) => ``([ $f where $p | $[$i <- $l]in* ]

def pyth (n : Nat) : LazyList (Nat × Nat × Nat) :=
  [(x, y, z) where x ^ 2 + y ^ 2 == z ^ 2
    | x <- [|1 ~~ n + 1|]
   in y <- [|x ~~ n + 1|]
   in z <- [|y ~~ n + 1|]]

gen ::= [|num ~~ (num)? (: num)? |]
			| [|num to (num)? (by num)?|]

小缺陷

example : 
  (λ x => λ y => λ z => (x, y, z) : α -> β -> γ -> (α × β × γ)) 
  = 
  λ x y z => (x, y, z) := rfl

(fun (x, _) (_, y) => (x, y)) <$> [(1,2), (3,4)] <*> [(5,6),(7,8)]
-- [(1, 6), (1, 8), (3, 6), (3, 8)]

(fun (x, _) => fun (_, y) => (x, y)) <$> [(1,2), (3,4)] <*> [(5,6),(7,8)]

-- stuck at solving universe constraint
--   max (?u.400287+1) (?u.400288+1) =?= max (?u.400212+1) (?u.400213+1)
-- while trying to unify
--   Prod.{?u.400213, ?u.400212} ?m.400220 ?m.400230 : Type (max ?u.400213 ?u.400212)
-- with
--   Prod.{?u.400213, ?u.400212} ?m.400220 ?m.400230 : Type (max ?u.400213 ?u.400212)

structure NProd (α β : Type u) where
  mkNp ::
  fst : α
  snd : β
  
open NProd in 
#eval 
  (fun {fst,..} => fun {snd,..} => (fst, snd)) 
 <$> [mkNp 1 2, mkNp 3 4] 
 <*> [mkNp 5 6, mkNp 7 8]

[(x, y) | (x, _) <- [(1, 2), (3, 4)] in (_, y) <- [(5, 6), (7, 8)]]
= [(1, 2), (3, 4)] >>= fun (x, _) => 
		[(5, 6), (7, 8)]] >>= fun (_, y) =>
			(x, y)

• 不过这个问题很好解决，显式添加类型标注 (type ascription) 即可。 show List (Nat × Nat) from ...

• 实际使用，toplevel 的定义本就最好要写类型；local 有 local bidirectional type inference, 真正会报错的情况还是少数。

Translation #2

do 
	let x₁' <- x₁
	let x₂' <- x₂
	...
	let xₙ' <- xₙ
	pure (f x₁' x₂' ... xₙ')

• 告诉我们把 中的第 个元素 全部转换成 ，由于其需要多个参数，相当于 partial application, 因此相当于产生一个全是函数的列表，记作 ，

• 接着看 ，两边都是列表，一边是一列函数，另一边是一列参数。

• 固定一个 ，对 ，其结果为一个列表，往复这个过程并把所有的列表连接起来，就是结果。

• 有标准定义 f <*> x = f >>= (· <$> x)

comp ::=
	  ...
	| [$exp (where $p)? | ($pat <- $exp',)+]

syntax "[" term " | " withoutPosition((term (" ← " <|> " <- ") term),+) "]" : term
syntax "[" term " where"  term " | " withoutPosition((term (" ← " <|> " <- ") term),+) "]" : term

macro_rules
	...
  | `([ $f | $[$i <- $l],* ]) => do
    if h : l.size >= 1 then
      let is <- ``(fun $i* => $f)
      let fn <- ``(Functor.map $is)
      let l0 := l[0]
      show MacroM Term from
      ``($fn $l0) >>= l.foldlM (start := 1) fun a s =>
        ``(Seq.seq $a (fun _ : Unit => $s))
    else unreachable!
  | `([ $f | $[$i ← $l],* ]) => ``([ $f | $[$i <- $l],* ])
  | `([ $f where $p | $[$i <- $l],* ]) => do
    let mf <- ``($f depends_on! $p)
    ``(Mappable.filterMap id [ $mf | $[$i <- $l]in* ])
  | `([ $f where $p | $[$i ← $l],* ]) => do
    ``([ $f where $p | $[$i <- $l],* ])

• 我们这里是一个 left fold, <$> 最先求值；

• 和前面的 right fold, map 最后求值（在最内层）不一样。

• 所以此处不能合并，即这种翻译在 predicative filtering 的情况下必须要多遍历一遍。

• 性能上并不好

• 而且显然不支持 sequential binding

• 注意 ``(fun $i* => $f)

Translating by Examples: Parallel

zipWith f xs ys = [| f.body | x <- xs | y <- ys |]

• zip = zipWith (· , ·) ，记为 . right associative, binds tighter than . (因为 (· × ·) 就是右结合的)

• 考虑 zipWithN ，其工作在 个列表上，那么

• 定义很美妙，跑起来像屎一样。

• 因为是平行的，理想情况下的时间复杂度应当是

土办法

• 但我们仍要尽力减少这种 boilerplate 的东西，先考虑 zip zipWith 的关系：

macro:max "prod_of!" n:num : term => 
  match n.getNat with
  | 0 => ``(())
  | 1 => ``(id)
  | n + 2 =>
    ``((· , ·)) >>= n.foldM fun _ _ a =>
      ``((· , $a))

--  Similar to Haskell's `MonadZip`, but for collections especially.
--  But relaxes on the constraint that `m` must be a `Monad` first.
--  - Any `Monad` is `Zippable`.

class Zippable (m : Type u -> Type v) where

  zipWith  (f : α -> β -> γ) : m α -> m β -> m γ
  zipWith3 (f : α -> β -> γ -> δ) : m α -> m β -> m γ -> m δ
  zipWith4 (f : α -> β -> γ -> δ -> ζ) : m α -> m β -> m γ -> m δ -> m ζ  
  zipWith5 (f : α -> β -> γ -> δ -> ζ -> η) : m α -> m β -> m γ -> m δ -> m ζ -> m η
  zipWith6 (f : α -> β -> γ -> δ -> ζ -> η -> θ) : m α -> m β -> m γ -> m δ -> m ζ -> m η -> m θ
  zipWith7 (f : α -> β -> γ -> δ -> ζ -> η -> θ -> ι) : m α -> m β -> m γ -> m δ -> m ζ -> m η -> m θ -> m ι
  
  zip  : m α -> m β -> m (α × β) := zipWith prod_of! 2
  zip3 : m α -> m β -> m γ -> m (α × β × γ) := zipWith3 prod_of! 3
  zip4 : m α -> m β -> m γ -> m δ -> m (α × β × γ × δ) := zipWith4 prod_of! 4
  zip5 : m α -> m β -> m γ -> m δ -> m ε -> m (α × β × γ × δ × ε) := zipWith5 prod_of! 5
  zip6 : m α -> m β -> m γ -> m δ -> m ε -> m ζ -> m (α × β × γ × δ × ε × ζ) := zipWith6 prod_of! 6
  zip7 : m α -> m β -> m γ -> m δ -> m ε -> m ζ -> m η -> m (α × β × γ × δ × ε × ζ × η) := zipWith7 prod_of! 7

Generalizing 土办法

instance (priority := low) [Monad m] : Zippable m where
  zipWith  := (· <$> · <*> ·)
  zipWith3 := (· <$> · <*> · <*> ·)
  zipWith4 := (· <$> · <*> · <*> · <*> ·)
  zipWith5 := (· <$> · <*> · <*> · <*> · <*> ·)
  zipWith6 := (· <$> · <*> · <*> · <*> · <*> · <*> ·)
  zipWith7 := (· <$> · <*> · <*> · <*> · <*> · <*> · <*> ·)

• 其实通篇看下来会发现我们特意不去统一使用Monad 这个API，而且实现几乎都是从序列结构 List 等视角叙述的，这有几个原因：

• 第一是 Lean 的求值语义和 Haskell 不同，上面也强调过很多遍了；

• 第二是这些序列结构确实是有特殊的实现，没有必要为了统一而统一，失去本可以避免失去的性能优势

• 第三是大家都爱 generalizing，monad comprehension 确实是很美妙的统一，但是就我所知很少有人真的用 comprehension 的语法去写 monadic code, 绝大多数都是用 do ; 而 comprehension 则仍然主要用于序列结构，

• 那么很显然，我们的 comprehension 就一定要为序列结构而优化，generalization 都是次要的

• 实际上我们发现，在做足了必要的优化下也没有损失任何一点 generality，使用的 API 虽然较乱，但是这对库使用者来说是不可见的。

• 不仅如此，涉及到的几个 typeclass 本就有很强的关联性，实现起来bing

Zippable List

def zipWith4 (f : α -> β -> γ -> δ -> ζ) : List α -> List β -> List γ -> List δ -> List ζ := go #[] where
  go acc
  | a :: as, b :: bs, c :: cs, d :: ds => go (acc.push $ f a b c d) as bs cs ds
  | _, _, _, _ => acc.toList
def zipWith5 (f : α -> β -> γ -> δ -> ζ -> η) : List α -> List β -> List γ -> List δ -> List ζ -> List η := go #[] where
  go acc
  | a :: as, b :: bs, c :: cs, d :: ds, e :: es => go (acc.push $ f a b c d e) as bs cs ds es
  | _, _, _, _, _ => acc.toList

Zippable LazyList

def zipWith4 (f : α -> β -> γ -> δ -> ζ) : LazyList α -> LazyList β -> LazyList γ -> LazyList δ -> LazyList ζ
  | a ::' as, b ::' bs, c ::' cs, d ::' ds => f a b c d ::' zipWith4 f as.get bs.get cs.get ds.get
  | _, _, _, _ => [||]
def zipWith5 (f : α -> β -> γ -> δ -> ζ -> η) : LazyList α -> LazyList β -> LazyList γ -> LazyList δ -> LazyList ζ -> LazyList η
  | a ::' as, b ::' bs, c ::' cs, d ::' ds, e ::' es => f a b c d e ::' zipWith5 f as.get bs.get cs.get ds.get es.get

Zippable Array

• 在运行时编译器将用一个真正的， primitive的，适合编程的 Array 覆盖之。

• arr[i] 需要证明 i < arr.size，但是最快，没有任何下面提到的 overhead.

• arr[i]? 产生一个 Option α , 但如此，对 n 个数组而言我们就要做 次运行时越界检查,
同时 constructing/destructuring Option 又要做同样多次，i.e.

match as[i]?, bs[i]? .. with
| some a, some b, .. =>
| _, _, .. =>

这么多 overhead 在性能上不可接受。

• arr[i]! 相当于命令式语言的下标访问，若越界运行时会 panic ，但是 Lean 毕竟是一个 theorem prover, 讲求 type soundness, 即便 panic ，我们也需要证明 Inhabited α ；同时这里的 result type 就是 α , 我们也不能通过 Nonempty α 并用选择公理 nonconstructively 直接选一个元素出来，否则整个 zipWithN 就是 noncomputable 的。

证明访问一定不越界