265 lines
11 KiB
Text
265 lines
11 KiB
Text
import Lean.Data.HashMap
|
||
/-!
|
||
# State
|
||
|
||
In the [previous section](readers.lean.md), you learned about the `ReaderM` monad. Hopefully this gave you
|
||
a new perspective on Lean. It showed that, in fact, you _can_ have global variables of some sort;
|
||
you just need to encode them in the type signature somehow, and this is what monads are for! In this
|
||
part, you will explore the `StateM` monad, which is like a `ReaderM` only the state can also be updated.
|
||
|
||
## Motivating example: Tic Tac Toe
|
||
|
||
For this section, let's build a simple model for a Tic Tace Toe game. The main object is the `GameState`
|
||
data type containing several important pieces of information. First and foremost, it has the
|
||
"board", a map from 2D tile indices to the "Tile State" (X, O or empty). Then it also knows the
|
||
current player, and it has a random generator.
|
||
-/
|
||
|
||
open Batteries (HashMap)
|
||
abbrev TileIndex := Nat × Nat -- a 2D index
|
||
|
||
inductive TileState where
|
||
| TileEmpty | TileX | TileO
|
||
deriving Repr, BEq
|
||
|
||
inductive Player where
|
||
| XPlayer | OPlayer
|
||
deriving Repr, BEq
|
||
|
||
abbrev Board := HashMap TileIndex TileState
|
||
|
||
structure GameState where
|
||
board : Board
|
||
currentPlayer : Player
|
||
generator : StdGen
|
||
|
||
/-!
|
||
Let's think at a high level about how some of the game functions would work. You could, for
|
||
instance, have a function for selecting a random move. This would output a `TileIndex` to play and
|
||
alter the game's number generator. You would then make a move based on the selected move and the
|
||
current player. This would change the board state as well as swap the current player. In other
|
||
words, you have operations that depend on the current state of the game, but also need to **update
|
||
that state**.
|
||
|
||
## The StateM Monad to the Rescue
|
||
|
||
This is exactly the situation the `StateM` monad deals with. The `StateM` monad wraps computations in
|
||
the context of reading and modifying a global state object.
|
||
|
||
It is parameterized by a single type parameter `s`, the state type in use. So just like the `ReaderM`
|
||
has a single type you read from, the `StateM` has a single type you can both **read from and write
|
||
to**. There are three primary actions you can take within the `StateM`monad:
|
||
|
||
- **get** - retrieves the state, like Reader.read
|
||
- **set** - updates the state
|
||
- **modifyGet** - retrieves the state, then updates it
|
||
|
||
There is also a `run` function, similar to `run` on `ReaderM`. Like the `ReaderM` monad, you must
|
||
provide an initial state, in addition to the computation to run. `StateM` then produces two outputs:
|
||
the result of the computation combined with the final updated state.
|
||
|
||
If you wish to discard the final state and just get the computation's result, you can use
|
||
`run'` method instead. Yes in Lean, the apostrophe can be part of a name, you read this "run
|
||
prime", and the general naming convention is that the prime method discards something.
|
||
|
||
So for your Tic Tac Toe game, many of your functions will have a signature like `State GameState a`.
|
||
|
||
## Stateful Functions
|
||
|
||
Now you can examine some of the different functions mentioned above and determine their types.
|
||
You can, for instance, pick a random move:
|
||
|
||
-/
|
||
open TileState
|
||
|
||
def findOpen : StateM GameState (List TileIndex) := do
|
||
let game ← get
|
||
return game.board.toList.filterMap fun (i, x) => guard (x == TileEmpty) *> pure i
|
||
|
||
def chooseRandomMove : StateM GameState TileIndex := do
|
||
let game ← get
|
||
let openSpots ← findOpen
|
||
let gen := game.generator
|
||
let (i, gen') := randNat gen 0 (openSpots.length - 1)
|
||
set { game with generator := gen' }
|
||
return openSpots[i]!
|
||
|
||
/-!
|
||
This returns a `TileIndex` and modifies the random number generator stored in the `GameState`!
|
||
Notice you have a fun little use of the `Applicative.seqRight` operator `*>` in `findOpen`
|
||
as described in [Applicatives](applicatives.lean.md).
|
||
|
||
Now you can create the function that can make a move:
|
||
-/
|
||
open Player
|
||
|
||
def tileStateForPlayer : Player → TileState
|
||
| XPlayer => TileX
|
||
| OPlayer => TileO
|
||
|
||
def nextPlayer : Player → Player
|
||
| XPlayer => OPlayer
|
||
| OPlayer => XPlayer
|
||
|
||
def applyMove (i : TileIndex): StateM GameState Unit := do
|
||
let game ← get
|
||
let p := game.currentPlayer
|
||
let newBoard := game.board.insert i (tileStateForPlayer p)
|
||
set { game with currentPlayer := nextPlayer p, board := newBoard }
|
||
|
||
/-!
|
||
This updates the board in the `GameState` with the new tile, and then changes the current player,
|
||
providing no output (`Unit` return type).
|
||
|
||
So finally, you can combine these functions together with `do` notation, and it actually looks quite
|
||
clean! You don't need to worry about the side effects. The different monadic functions handle them.
|
||
Here's a sample of what your function might look like to play one turn of the game. At the end, it
|
||
returns a boolean determining if all the spaces have been filled.
|
||
|
||
Notice in `isGameDone` and `nextTurn` we have stopped providing the full return type
|
||
`StateM GameState Unit`. This is because Lean is able to infer the correct monadic return type
|
||
from the context and as a result the code is now looking really clean.
|
||
-/
|
||
|
||
def isGameDone := do
|
||
return (← findOpen).isEmpty
|
||
|
||
def nextTurn := do
|
||
let i ← chooseRandomMove
|
||
applyMove i
|
||
isGameDone
|
||
|
||
/-!
|
||
To give you a quick test harness that runs all moves for both players you can run this:
|
||
-/
|
||
|
||
def initBoard : Board := Id.run do
|
||
let mut board := HashMap.empty
|
||
for i in [0:3] do
|
||
for j in [0:3] do
|
||
let t : TileIndex := (i, j)
|
||
board := board.insert t TileEmpty
|
||
board
|
||
|
||
def printBoard (board : Board) : IO Unit := do
|
||
let mut row : List String := []
|
||
for i in board.toList do
|
||
let s := match i.2 with
|
||
| TileEmpty => " "
|
||
| TileX => "X"
|
||
| TileO => "O"
|
||
row := row.append [s]
|
||
if row.length == 3 then
|
||
IO.println row
|
||
row := []
|
||
|
||
def playGame := do
|
||
while true do
|
||
let finished ← nextTurn
|
||
if finished then return
|
||
|
||
def main : IO Unit := do
|
||
let gen ← IO.stdGenRef.get
|
||
let (x, gen') := randNat gen 0 1
|
||
let gs := {
|
||
board := initBoard,
|
||
currentPlayer := if x = 0 then XPlayer else OPlayer,
|
||
generator := gen' }
|
||
let (_, g) := playGame |>.run gs
|
||
printBoard g.board
|
||
|
||
#eval main
|
||
-- [X, X, O]
|
||
-- [X, O, O]
|
||
-- [O, O, X]
|
||
|
||
/-!
|
||
|
||
Note that when you run the above code interactively the random number generator always starts in the
|
||
same place. But if you run `lean --run states.lean` then you will see randomness in the result.
|
||
|
||
## Implementation
|
||
|
||
It may be helpful to see how the `StateM` monad adds the input state and output state. If you look
|
||
at the reduced Type for `nextTurn`:
|
||
-/
|
||
#reduce StateM GameState Bool
|
||
-- GameState → Bool × GameState
|
||
/-!
|
||
|
||
So a function like `nextTurn` that might have just returned a `Bool` has been modified by the
|
||
`StateM` monad such that the initial `GameState` is passed in as a new input argument, and the output
|
||
value has been changed to the pair `Bool × GameState` so that it can return the pure `Bool` and the
|
||
updated `GameState`. So `playGame` then is automatically saving that updated game state so that each
|
||
time around the `while` loop it is acting on the new state, otherwise that would be an infinite loop!
|
||
|
||
It is also interesting to see how much work the `do` and `←` notation are doing for you. To
|
||
implement the `nextTurn` function without these you would have to write this, manually plumbing
|
||
the state all the way through:
|
||
-/
|
||
def nextTurnManually : StateM GameState Bool
|
||
| state =>
|
||
let (i, gs) := chooseRandomMove |>.run state
|
||
let (_, gs') := applyMove i |>.run gs
|
||
let (result, gs'') := isGameDone |>.run gs'
|
||
(result, gs'')
|
||
|
||
/-!
|
||
|
||
This expression `let (i, gs)` conveniently breaks a returned pair up into 2 variables.
|
||
In the expression `let (_, gs')` we didn't care what the first value was so we used underscore.
|
||
Notice that nextTurn is capturing the updated game state from `chooseRandomMove` in the variable
|
||
`gs`, which it is then passing to `applyMove` which returns `gs'` which is passed to `isGameDone`
|
||
and that function returns `gs''` which we then return from `nextTurnManually`. Phew, what a lot
|
||
of work you don't have to do when you use `do` notation!
|
||
|
||
## StateM vs ReaderM
|
||
|
||
While `ReaderM` functions can use `withReader` to modify the context before calling another function,
|
||
`StateM` functions are a little more powerful, let's look at this function again:
|
||
```
|
||
def nextTurn : StateM GameState Bool := do
|
||
let i ← chooseRandomMove
|
||
applyMove i
|
||
isGameDone
|
||
```
|
||
|
||
In this function `chooseRandomMove` is modifying the state that `applyMove` is getting
|
||
and `chooseRandomMove` knows nothing about `applyMove`. So `StateM` functions can have this
|
||
kind of downstream effect outside their own scope, whereas, `withReader` cannot do that.
|
||
|
||
So there is no equivalent to `withReader` for `StateM`, besides you can always use the `StateM`
|
||
`set` function to modify the state before calling the next function anyway. You could however,
|
||
manually call a `StateM` function like you see in `nextTurnManually` and completely override
|
||
the state at any point that way.
|
||
|
||
## State, IO and other languages
|
||
|
||
When thinking about Lean, it is often seen as a restriction that you can't have global variables or
|
||
`static` variables like you can with other languages like Python or C++. However, hopefully you see
|
||
now this isn't true. You can have a data type with exactly the same functionality as a Python class.
|
||
You would simply have many functions that can modify some global state using the `StateM` monad.
|
||
|
||
The difference is in Lean you simply put a label on these types of functions. You don't allow it to
|
||
happen for free anywhere in an uncontrolled fashion because that results in too many sleepless
|
||
nights debugging nasty code. You want to know when side effects can potentially happen, because
|
||
knowing when they can happen makes your code easier to reason about. In a Python class, many of the
|
||
methods won't actually need to modify the global state. But they could, which makes it harder to
|
||
debug them. In Lean you can simply make these pure functions, and the compiler will ensure they stay
|
||
pure and cannot modify any global state.
|
||
|
||
IO is the same way. It's not like you can't perform IO in Lean. Instead, you want to label the areas
|
||
where you can, to increase your certainty about the areas where you don't need to. When you know part of
|
||
your code cannot communicate with the outside world, you can be far more certain of its behavior.
|
||
|
||
The `StateM` monad is also a more disciplined way of managing side effects. Top level code could
|
||
call a `StateM` function multiple times with different independent initial states, even doing that
|
||
across multiple tasks in parallel and each of these cannot clobber the state belonging to other
|
||
tasks. Monadic code is more predictable and reusable than code that uses global variables.
|
||
|
||
## Summary
|
||
|
||
That wraps it up for the `StateM` monad! There is one more very useful monad that can be used to do
|
||
exception handling which will be covered in the [next section](except.lean.md).
|
||
|
||
-/
|