import Do.Basic
import Do.LazyList

Local Mutation

open Lean

Disable the automatic monadic lifting feature described in the paper. We want to make it clear that we do not depend on it.

set_option autoLift false

syntax "let" "mut" ident ":=" term ";" stmt : stmt
syntax ident ":=" term : stmt
syntax "if" term "then" stmt "else" stmt:1 : stmt

declare_syntax_cat expander
-- generic syntax for traversal-like functions S_y/R/B/L
syntax "expand!" expander "in" stmt:1 : stmt
syntax "mut" ident : expander  -- corresponds to `S_y`

-- generic traversal rules
macro_rules
  | `(stmt| expand! $
exp: TSyntax `expander
exp in let $
x: TSyntax `ident
x ← $
s: TSyntax `stmt
s; $
s': TSyntax `stmt
s') => `(stmt| let $
x: TSyntax `ident
x ← expand! $
exp: TSyntax `expander
exp in $
s: TSyntax `stmt
s; expand! $
exp: TSyntax `expander
exp in $
s': TSyntax `stmt
s')                -- subsumes (R3, B4, L4)
  | `(stmt| expand! $
exp: TSyntax `expander
exp in let mut $
x: TSyntax `ident
x := $
e: TSyntax `term
e; $
s': TSyntax `stmt
s') => `(stmt| let mut $
x: TSyntax `ident
x := $
e: TSyntax `term
e; expand! $
exp: TSyntax `expander
exp in $
s': TSyntax `stmt
s')                      -- subsumes (R4, B5, L5)
  | `(stmt| expand! $_ in $
x: TSyntax `ident
x:ident := $
e: TSyntax `term
e) => `(stmt| $
x: TSyntax `ident
x:ident := $
e: TSyntax `term
e)                                                      -- subsumes (R5, B6, L6)
  | `(stmt| expand! $
exp: TSyntax `expander
exp in if $
e: TSyntax `term
e then $
s₁: TSyntax `stmt
s₁ else $
s₂: TSyntax `stmt
s₂) => `(stmt| if $
e: TSyntax `term
e then expand! $
exp: TSyntax `expander
exp in $
s₁: TSyntax `stmt
s₁ else expand! $
exp: TSyntax `expander
exp in $
s₂: TSyntax `stmt
s₂)  -- subsumes (S6, R6, B7, L7)
  | `(stmt| expand! $
exp: TSyntax `expander
exp in $
s: TSyntax `stmt
s) => do
    let 
s': TSyntax `stmt
s' ← 
expandStmt: TSyntax `stmt → MacroM (TSyntax `stmt)
expandStmt 
s: TSyntax `stmt
s
    `(stmt| expand! $
exp: TSyntax `expander
exp in $
s': TSyntax `stmt
s')


macro_rules
  | `(d! let mut $
x: TSyntax `ident
x := $
e: TSyntax `term
e; $
s: TSyntax `stmt
s) => `(let $
x: TSyntax `ident
x := $
e: TSyntax `term
e; StateT.run' (d! expand! mut $
x: TSyntax `ident
x in $
s: TSyntax `stmt
s) $
x: TSyntax `ident
x) -- (D3)
  | `(d! $
x: TSyntax `ident
x:ident := $_:term) =>
      -- `s!"..."` is an interpolated string. For more information, see https://leanprover.github.io/lean4/doc/stringinterp.html.
      
throw: {ε : Type} → {m : Type → Type} → [self : MonadExcept ε m] → {α : Type} → ε → m α
throw <| 
Macro.Exception.error: Syntax → String → Macro.Exception
Macro.Exception.error 
x: TSyntax `ident
x s!"variable '{
x: TSyntax `ident
x.
getId: Ident → Name
getId}' is not reassignable in this scope"
  | `(d! if $
e: TSyntax `term
e then $
s₁: TSyntax `stmt
s₁ else $
s₂: TSyntax `stmt
s₂) => `(if $
e: TSyntax `term
e then d! $
s₁: TSyntax `stmt
s₁ else d! $
s₂: TSyntax `stmt
s₂)                       -- (D4)

macro_rules
  | `(stmt| expand! mut $_ in $
e: TSyntax `term
e:term) => `(stmt| StateT.lift $
e: TSyntax `term
e)  -- (S1)
  | `(stmt| expand! mut $
y: TSyntax `ident
y in let $
x: TSyntax `ident
x ← $
s: TSyntax `stmt
s; $
s': TSyntax `stmt
s') =>                -- (S2)
    if 
x: TSyntax `ident
x == 
y: TSyntax `ident
y then
      
throw: {ε : Type} → {m : Type → Type} → [self : MonadExcept ε m] → {α : Type} → ε → m α
throw <| 
Macro.Exception.error: Syntax → String → Macro.Exception
Macro.Exception.error 
x: TSyntax `ident
x s!"cannot shadow 'mut' variable '{
x: TSyntax `ident
x.
getId: Ident → Name
getId}'"
    else
      `(stmt| let $
x: TSyntax `ident
x ← expand! mut $
y: TSyntax `ident
y in $
s: TSyntax `stmt
s; let $
y: TSyntax `ident
y ← get; expand! mut $
y: TSyntax `ident
y in $
s': TSyntax `stmt
s')
  | `(stmt| expand! mut $
y: TSyntax `ident
y in let mut $
x: TSyntax `ident
x := $
e: TSyntax `term
e; $
s': TSyntax `stmt
s') =>           -- (S3)
    if 
x: TSyntax `ident
x == 
y: TSyntax `ident
y then
      
throw: {ε : Type} → {m : Type → Type} → [self : MonadExcept ε m] → {α : Type} → ε → m α
throw <| 
Macro.Exception.error: Syntax → String → Macro.Exception
Macro.Exception.error 
x: TSyntax `ident
x s!"cannot shadow 'mut' variable '{
x: TSyntax `ident
x.
getId: Ident → Name
getId}'"
    else
      `(stmt| let mut $
x: TSyntax `ident
x := $
e: TSyntax `term
e; expand! mut $
y: TSyntax `ident
y in $
s': TSyntax `stmt
s')
  | `(stmt| expand! mut $
y: TSyntax `ident
y in $
x: TSyntax `ident
x:ident := $
e: TSyntax `term
e) =>
    if 
x: TSyntax `ident
x == 
y: TSyntax `ident
y then
      `(stmt| set $
e: TSyntax `term
e)                                             -- (S5)
    else
      `(stmt| $
x: TSyntax `ident
x:ident := $
e: TSyntax `term
e)                                     -- (S4)

macro:0 "let" "mut" 
x: TSyntax `ident
x:ident "←" 
s: TSyntax `stmt
s:stmt:1 ";" 
s': TSyntax `stmt
s':stmt : stmt => `(let y ← $
s: TSyntax `stmt
s; let mut $
x: TSyntax `ident
x := y; $
s': TSyntax `stmt
s') -- (A3)
macro:0 
x: TSyntax `ident
x:ident "←" 
s: TSyntax `stmt
s:stmt:1 : stmt => `(let y ← $
s: TSyntax `stmt
s; $
x: TSyntax `ident
x:ident := y)                                -- (A4)
-- a variant of (A4) since we technically cannot make the above macro a `stmt`
macro:0 
x: TSyntax `ident
x:ident "←" 
s: TSyntax `stmt
s:stmt:1 ";" 
s': TSyntax `stmt
s':stmt : stmt => `(let y ← $
s: TSyntax `stmt
s; $
x: TSyntax `ident
x:ident := y; $
s': TSyntax `stmt
s')
macro "if" 
e: TSyntax `term
e:term "then" 
s₁: TSyntax `stmt
s₁:stmt:1 : stmt => `(if $
e: TSyntax `term
e then $
s₁: TSyntax `stmt
s₁ else pure ())                        -- (A5)
macro "unless" 
e: TSyntax `term
e:term "do'" 
s₂: TSyntax `stmt
s₂:stmt:1 : stmt => `(if $
e: TSyntax `term
e then pure () else $
s₂: TSyntax `stmt
s₂)                     -- (A6)

/-
The `variable` command instructs Lean to insert the declared variables as bound variables
in definitions that refer to them.
-/

variable [
Monad: (Type → Type u_1) → Type (max 1 u_1)
Monad 
m: Type u_1 → Type u_2
m]
variable (
ma: m α
ma 
ma': m α
ma' : 
m: Type u_1 → Type u_2
m 
α: Type ?u.10009
α)

/-
  Mark `map_eq_pure_bind` as a simplification lemma.
  It is a theorem for
  `f <$> x = x >>= pure (f a)`
-/
attribute [local simp] 
map_eq_pure_bind: ∀ {m : Type u_1 → Type u_2} {α β : Type u_1} [inst : Monad m] [inst_1 : LawfulMonad m] (f : α → β) (x : m α),
  f <$> x = do
    let a ← x
    pure (f a)
map_eq_pure_bind

/-
  Remark: an `example` in Lean is like a "nameless" definition, and it does not update the environment.
  It is useful for writing tests.
-/

/-
  The instance `[LawfulMonad m]` contains the monadic laws.
  For more information, see https://github.com/leanprover/lean4/blob/v4.0.0-m4/src/Init/Control/Lawful.lean
-/

example: ∀ {m : Type u_1 → Type u_2} {α : Type u_1} [inst : Monad m] (ma : m α) [inst_1 : LawfulMonad m],
  (do
      let y ← ma
      let x : α := y
      StateT.run' (StateT.lift (pure x)) x) =
    ma
example [
LawfulMonad: (m : Type u_1 → Type u_2) → [inst : Monad m] → Prop
LawfulMonad 
m: Type u_1 → Type u_2
m] :
    (do' let mut 
x: α
x ← 
ma: m α
ma;
         
pure: {f : Type u_1 → Type u_2} → [self : Pure f] → {α : Type u_1} → α → f α
pure 
x: α
x : 
m: Type u_1 → Type u_2
m 
α: Type u_1
α)
    =
    
ma: m α
ma
:= by
m: Type ?u.10102 → Type ?u.10058
α: Type ?u.10102
inst✝¹: Monad m
ma, ma': m α
(do
    let y ← ma
    let x : α := y
    StateT.run' (StateT.lift (pure x)) x) =
  ma simp
Goals accomplished! 🐙

example: ∀ {m : Type u_1 → Type u_2} {α : Type u_1} [inst : Monad m] (ma ma' : m α) [inst_1 : LawfulMonad m],
  (do
      let y ← ma
      let x : α := y
      StateT.run'
          (do
            let y ← StateT.lift ma'
            let _ ← get
            set y
            let x ← get
            StateT.lift (pure x))
          x) =
    do
    let _ ← ma
    ma'
example [
LawfulMonad: (m : Type u_1 → Type u_2) → [inst : Monad m] → Prop
LawfulMonad 
m: Type u_1 → Type u_2
m] :
    (do' let mut 
x: α
x ← 
ma: m α
ma;
         x ← 
ma': m α
ma';
         
pure: {f : Type u_1 → Type u_2} → [self : Pure f] → {α : Type u_1} → α → f α
pure 
x: α
x)
    =
    (
ma: m α
ma >>= fun _ => 
ma': m α
ma')
:= by
m: Type ?u.10970 → Type ?u.10750
α: Type ?u.10970
inst✝¹: Monad m
ma, ma': m α
(do
    let y ← ma
    let x : α := y
    StateT.run'
        (do
          let y ← StateT.lift ma'
          let _ ← get
          set y
          let x ← get
          StateT.lift (pure x))
        x) =
  do
  let _ ← ma
  ma' simp
Goals accomplished! 🐙

/- The command `#check_failure <term>` succeeds only if `<term>` fails to be elaborated. -/

#check_failure do'
  let mut 
x: α
x ← 
ma: m α
ma;
  let
do
  let y ← ma
  let x : α := y
  sorryAx (m ?m.13198) true : m ?m.13198 x
do
  let y ← ma
  let x : α := y
  sorryAx (m ?m.13198) true : m ?m.13198
Warning: cannot shadow 'mut' variable 'x' ← ma';  -- cannot shadow 'mut' variable 'x'
  pure x
do
  let y ← ma
  let x : α := y
  sorryAx (m ?m.13198) true : m ?m.13198

#check_failure do'
do
  let _ ← ma
  sorryAx (m ?m.13260) true : m ?m.13260
  x
do
  let _ ← ma
  sorryAx (m ?m.13260) true : m ?m.13260
Warning: variable 'x' is not reassignable in this scope ← 
ma: m α
ma;  -- variable 'x' is not reassignable in this scope
  pure ()
do
  let _ ← ma
  sorryAx (m ?m.13260) true : m ?m.13260

variable (
b: Bool
b : 
Bool: Type
Bool)

-- The following equivalence is true even if `m` does not satisfy the monadic laws.
example: ∀ {m : Type → Type u_1} {α : Type} [inst : Monad m] (ma : m α) (b : Bool),
  (if b = true then discard ma else pure ()) = if b = true then discard ma else pure ()
example :
    (do' if 
b: Bool
b then {
           
discard: {f : Type → Type u_1} → {α : Type} → [inst : Functor f] → f α → f PUnit
discard 
ma: m α
ma
         })
    =
    (if 
b: Bool
b then 
discard: {f : Type → Type u_1} → {α : Type} → [inst : Functor f] → f α → f PUnit
discard 
ma: m α
ma else 
pure: {f : Type → Type u_1} → [self : Pure f] → {α : Type} → α → f α
pure ())
:= 
rfl: ∀ {α : Type u_1} {a : α}, a = a
rfl

theorem 
simple: ∀ [inst : LawfulMonad m],
  (do
      let y ← ma
      let x : α := y
      StateT.run'
          (do
            if b = true then do
                let y ← StateT.lift ma'
                let _ ← get
                set y
              else StateT.lift (pure ())
            let x ← get
            StateT.lift (pure x))
          x) =
    do
    let x ← ma
    if b = true then ma' else pure x
simple [
LawfulMonad: (m : Type → Type u_1) → [inst : Monad m] → Prop
LawfulMonad 
m: Type → Type u_1
m] :
    (do' let mut 
x: α
x ← 
ma: m α
ma;
         if 
b: Bool
b then {
           x ← 
ma': m α
ma'
         };
         
pure: {f : Type → Type u_1} → [self : Pure f] → {α : Type} → α → f α
pure 
x: α
x)
    =
    (
ma: m α
ma >>= fun 
x: α
x => if 
b: Bool
b then 
ma': m α
ma' else 
pure: {f : Type → Type u_1} → [self : Pure f] → {α : Type} → α → f α
pure 
x: α
x)
:= by
m: Type → Type u_1
α: Type
inst✝¹: Monad m
ma, ma': m α
b: Bool
(do
    let y ← ma
    let x : α := y
    StateT.run'
        (do
          if b = true then do
              let y ← StateT.lift ma'
              let _ ← get
              set y
            else StateT.lift (pure ())
          let x ← get
          StateT.lift (pure x))
        x) =
  do
  let x ← ma
  if b = true then ma' else pure x cases 
b: Bool
b
m: Type → Type u_1
α: Type
inst✝¹: Monad m
ma, ma': m α
false
(do
    let y ← ma
    let x : α := y
    StateT.run'
        (do
          if false = true then do
              let y ← StateT.lift ma'
              let _ ← get
              set y
            else StateT.lift (pure ())
          let x ← get
          StateT.lift (pure x))
        x) =
  do
  let x ← ma
  if false = true then ma' else pure x
m: Type → Type u_1
α: Type
inst✝¹: Monad m
ma, ma': m α
true
(do
    let y ← ma
    let x : α := y
    StateT.run'
        (do
          if true = true then do
              let y ← StateT.lift ma'
              let _ ← get
              set y
            else StateT.lift (pure ())
          let x ← get
          StateT.lift (pure x))
        x) =
  do
  let x ← ma
  if true = true then ma' else pure x <;>
m: Type → Type u_1
α: Type
inst✝¹: Monad m
ma, ma': m α
false
(do
    let y ← ma
    let x : α := y
    StateT.run'
        (do
          if false = true then do
              let y ← StateT.lift ma'
              let _ ← get
              set y
            else StateT.lift (pure ())
          let x ← get
          StateT.lift (pure x))
        x) =
  do
  let x ← ma
  if false = true then ma' else pure x
m: Type → Type u_1
α: Type
inst✝¹: Monad m
ma, ma': m α
true
(do
    let y ← ma
    let x : α := y
    StateT.run'
        (do
          if true = true then do
              let y ← StateT.lift ma'
              let _ ← get
              set y
            else StateT.lift (pure ())
          let x ← get
          StateT.lift (pure x))
        x) =
  do
  let x ← ma
  if true = true then ma' else pure x simp
Goals accomplished! 🐙

example: ∀ {m : Type u_1 → Type u_2} {α : Type u_1} [inst : Monad m] (ma ma' : m α) (b : Bool) [inst_1 : LawfulMonad m]
  (f : α → α → α),
  (do
      let y ← ma
      let x : α := y
      StateT.run'
          (do
            let y ←
              if b = true then do
                  let y ← StateT.lift ma
                  let _ ← get
                  set y
                  let _ ← get
                  StateT.lift ma'
                else StateT.lift ma'
            let x ← get
            StateT.lift (pure (f x y)))
          x) =
    do
    let x ← ma
    if b = true then do
        let x ← ma
        let y ← ma'
        pure (f x y)
      else do
        let y ← ma'
        pure (f x y)
example [
LawfulMonad: (m : Type u_1 → Type u_2) → [inst : Monad m] → Prop
LawfulMonad 
m: Type u_1 → Type u_2
m] (
f: α → α → α
f : 
α: Type u_1
α → 
α: Type u_1
α → 
α: Type u_1
α) :
    (do' let mut 
x: α
x ← 
ma: m α
ma;
         let 
y: α
y ←
           if 
b: Bool
b then {
             x ← 
ma: m α
ma;
             
ma': m α
ma'
           } else {
             
ma': m α
ma'
           };
         
pure: {f : Type u_1 → Type u_2} → [self : Pure f] → {α : Type u_1} → α → f α
pure (
f: α → α → α
f 
x: α
x 
y: α
y))
    =
    (
ma: m α
ma >>= fun 
x: α
x => if 
b: Bool
b then 
ma: m α
ma >>= fun 
x: α
x => 
ma': m α
ma' >>= fun 
y: α
y => 
pure: {f : Type u_1 → Type u_2} → [self : Pure f] → {α : Type u_1} → α → f α
pure (
f: α → α → α
f 
x: α
x 
y: α
y) else 
ma': m α
ma' >>= fun 
y: α
y => 
pure: {f : Type u_1 → Type u_2} → [self : Pure f] → {α : Type u_1} → α → f α
pure (
f: α → α → α
f 
x: α
x 
y: α
y))
:= by
m: Type ?u.17614 → Type ?u.17083
α: Type ?u.17614
inst✝¹: Monad m
ma, ma': m α
b: Bool
f: α → α → α
(do
    let y ← ma
    let x : α := y
    StateT.run'
        (do
          let y ←
            if b = true then do
                let y ← StateT.lift ma
                let _ ← get
                set y
                let _ ← get
                StateT.lift ma'
              else StateT.lift ma'
          let x ← get
          StateT.lift (pure (f x y)))
        x) =
  do
  let x ← ma
  if b = true then do
      let x ← ma
      let y ← ma'
      pure (f x y)
    else do
      let y ← ma'
      pure (f x y) cases 
b: Bool
b
m: Type ?u.17614 → Type ?u.17083
α: Type ?u.17614
inst✝¹: Monad m
ma, ma': m α
f: α → α → α
false
(do
    let y ← ma
    let x : α := y
    StateT.run'
        (do
          let y ←
            if false = true then do
                let y ← StateT.lift ma
                let _ ← get
                set y
                let _ ← get
                StateT.lift ma'
              else StateT.lift ma'
          let x ← get
          StateT.lift (pure (f x y)))
        x) =
  do
  let x ← ma
  if false = true then do
      let x ← ma
      let y ← ma'
      pure (f x y)
    else do
      let y ← ma'
      pure (f x y)
m: Type ?u.17614 → Type ?u.17083
α: Type ?u.17614
inst✝¹: Monad m
ma, ma': m α
f: α → α → α
true
(do
    let y ← ma
    let x : α := y
    StateT.run'
        (do
          let y ←
            if true = true then do
                let y ← StateT.lift ma
                let _ ← get
                set y
                let _ ← get
                StateT.lift ma'
              else StateT.lift ma'
          let x ← get
          StateT.lift (pure (f x y)))
        x) =
  do
  let x ← ma
  if true = true then do
      let x ← ma
      let y ← ma'
      pure (f x y)
    else do
      let y ← ma'
      pure (f x y) <;>
m: Type ?u.17614 → Type ?u.17083
α: Type ?u.17614
inst✝¹: Monad m
ma, ma': m α
f: α → α → α
false
(do
    let y ← ma
    let x : α := y
    StateT.run'
        (do
          let y ←
            if false = true then do
                let y ← StateT.lift ma
                let _ ← get
                set y
                let _ ← get
                StateT.lift ma'
              else StateT.lift ma'
          let x ← get
          StateT.lift (pure (f x y)))
        x) =
  do
  let x ← ma
  if false = true then do
      let x ← ma
      let y ← ma'
      pure (f x y)
    else do
      let y ← ma'
      pure (f x y)
m: Type ?u.17614 → Type ?u.17083
α: Type ?u.17614
inst✝¹: Monad m
ma, ma': m α
f: α → α → α
true
(do
    let y ← ma
    let x : α := y
    StateT.run'
        (do
          let y ←
            if true = true then do
                let y ← StateT.lift ma
                let _ ← get
                set y
                let _ ← get
                StateT.lift ma'
              else StateT.lift ma'
          let x ← get
          StateT.lift (pure (f x y)))
        x) =
  do
  let x ← ma
  if true = true then do
      let x ← ma
      let y ← ma'
      pure (f x y)
    else do
      let y ← ma'
      pure (f x y) simp
Goals accomplished! 🐙

/-
Nondeterminism example from Section 2.
-/
def 
choose: {α : Type u_1} → List α → LazyList α
choose := @
List.toLazy: {α : Type u_1} → List α → LazyList α
List.toLazy

def 
ex: LazyList Nat
ex : 
LazyList: Type → Type
LazyList 
Nat: Type
Nat := do'
  let mut 
x: Nat
x := 
0: Nat
0;
  let 
y: Nat
y ← 
choose: {α : Type} → List α → LazyList α
choose [
0: Nat
0, 
1: Nat
1, 
2: Nat
2, 
3: Nat
3];
  x := 
x: Nat
x + 
1: Nat
1;
  
guard: {f : Type → Type} → [inst : Alternative f] → (p : Prop) → [inst : Decidable p] → f Unit
guard (
x: Nat
x < 
3: Nat
3);
  
pure: {f : Type → Type} → [self : Pure f] → {α : Type} → α → f α
pure (
x: Nat
x + 
y: Nat
y)

-- Generate all solutions
#eval
[1, 2, 3, 4]
 
ex: LazyList Nat
ex.
toList: {α : Type} → LazyList α → List α
toList