`import Lean.Data.HashMap`

# State

In the previous section, 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 Std () abbrevTileIndex :=TileIndex:TypeNat ×Nat:TypeNat -- a 2D index inductiveNat:TypeTileState where |TileState:TypeTileEmpty |TileEmpty:TileStateTileX |TileX:TileStateTileO derivingTileO:TileStateRepr,Repr:Type u → Type uBEq inductiveBEq:Type u → Type uPlayer where |Player:TypeXPlayer |XPlayer:PlayerOPlayer derivingOPlayer:PlayerRepr,Repr:Type u → Type uBEq abbrevBEq:Type u → Type uBoard :=Board:?m.657TileIndex TileState structureTileIndex TileState:?m.657GameState whereGameState:Type u_1board :board:GameState → {Board : Sort u_1} → BoardBoardBoard:Sort u_1currentPlayer :currentPlayer:GameState → PlayerPlayerPlayer:Typegenerator :generator:GameState → StdGenStdGenStdGen:Type

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 deffindOpen :findOpen:StateM GameState (List TileIndex)StateMStateM:Type → Type → TypeGameState (GameState:TypeListList:Type → TypeTileIndex) := do letTileIndex:Typegame ←game:GameStateget returnget:{σ : outParam Type} → {m : Type → Type} → [self : MonadState σ m] → m σdefchooseRandomMove :chooseRandomMove:StateM GameState TileIndexStateMStateM:Type → Type → TypeGameStateGameState:TypeTileIndex := do letTileIndex:Typegame ←game:GameStateget letget:{σ : outParam Type} → {m : Type → Type} → [self : MonadState σ m] → m σopenSpots ←openSpots:List TileIndexfindOpen letfindOpen:StateM GameState (List TileIndex)gen :=gen:StdGengame.game:GameStategenerator let (generator:GameState → StdGeni,i:Natgen') :=gen':StdGenrandNatrandNat:{gen : Type} → [inst : RandomGen gen] → gen → Nat → Nat → Nat × gengengen:StdGen0 (0:NatopenSpots.openSpots:List TileIndexlength -length:{α : Type} → List α → Nat1)1:Natset {set:{σ : semiOutParam Type} → {m : Type → Type} → [self : MonadStateOf σ m] → σ → m PUnitgame with generator :=game:GameStategen' } returngen':StdGenopenSpots[openSpots:List TileIndexi]!i:Nat

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.

Now you can create the function that can make a move:

open Player deftileStateForPlayer :tileStateForPlayer:Player → TileStatePlayer →Player:TypeTileState |TileState:TypeXPlayer =>XPlayer:PlayerTileX |TileX:TileStateOPlayer =>OPlayer:PlayerTileO defTileO:TileStatenextPlayer :nextPlayer:Player → PlayerPlayer →Player:TypePlayer |Player:TypeXPlayer =>XPlayer:PlayerOPlayer |OPlayer:PlayerOPlayer =>OPlayer:PlayerXPlayer defXPlayer:PlayerapplyMove (applyMove:TileIndex → StateM GameState Uniti :i:TileIndexTileIndex):TileIndex:TypeStateMStateM:Type → Type → TypeGameStateGameState:TypeUnit := do letUnit:Typegame ←game:GameStateget letget:{σ : outParam Type} → {m : Type → Type} → [self : MonadState σ m] → m σp :=p:Playergame.game:GameStatecurrentPlayer letcurrentPlayer:GameState → PlayernewBoard :=newBoard:∀ {Board : Prop}, Boardset {set:{σ : semiOutParam Type} → {m : Type → Type} → [self : MonadStateOf σ m] → σ → m PUnitgame with currentPlayer :=game:GameStatenextPlayernextPlayer:Player → Playerp, board :=p:PlayernewBoard }newBoard:∀ {Board : Prop}, Board

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`**: **StateM GameState Bool

isGameDone := do
return `(← findOpen)`**: **List TileIndex

(← `findOpen`**: **StateM GameState (List TileIndex)

findOpen`(← findOpen)`**: **List TileIndex

).`isEmpty`**: **{α : Type} → List α → Bool

isEmpty
def `nextTurn`**: **StateM GameState Bool

nextTurn := do
let `i`**: **TileIndex

i ← `chooseRandomMove`**: **StateM GameState TileIndex

chooseRandomMove
`applyMove`**: **TileIndex → StateM GameState Unit

applyMove `i`**: **TileIndex

i
`isGameDone`**: **StateM GameState Bool

isGameDone

To give you a quick test harness that runs all moves for both players you can run this:

definitBoard :initBoard:{Board : Type u_1} → BoardBoard :=Board:Type u_1Id.run do let mutId.run:{α : Type u_1} → Id α → αboard :=board:Id Boardfori in [i:Nat0:0:Nat3] do for3:Natj in [j:Nat0:0:Nat3] do let3:Natt :t:TileIndexTileIndex := (TileIndex:Typei,i:Natj)j:Natboard :=board:Id Boardboard defboard:Id BoardprintBoard (printBoard:{Board : Sort u_1} → Board → IO Unitboard :board:BoardBoard) :Board:Sort u_1IOIO:Type → TypeUnit := do let mutUnit:Typerow :row:List StringListList:Type → TypeString :=String:Type[][]:List StringdefplayGame := do whileplayGame:StateM GameState PUnittrue do lettrue:Boolfinished ←finished:BoolnextTurn ifnextTurn:StateM GameState Boolfinished thenfinished:Boolreturn defreturn:StateM GameState (ForInStep (MProd (Option PUnit) PUnit))main :main:IO UnitIOIO:Type → TypeUnit := do letUnit:Typegen ←gen:StdGenIO.stdGenRef.IO.stdGenRef:IO.Ref StdGenget let (get:{σ : Type} → {m : Type → Type} → [inst : MonadLiftT (ST σ) m] → {α : Type} → ST.Ref σ α → m αx,x:Natgen') :=gen':StdGenrandNatrandNat:{gen : Type} → [inst : RandomGen gen] → gen → Nat → Nat → Nat × gengengen:StdGen00:Nat1 let1:Natgs := { board :=gs:GameState, currentPlayer := ifx =x:Nat0 then0:NatXPlayer elseXPlayer:PlayerOPlayer, generator :=OPlayer:Playergen' } let (_,gen':StdGeng) :=g:GameStateplayGame |>.playGame:StateM GameState PUnitrunrun:{σ : Type} → {m : Type → Type} → {α : Type} → StateT σ m α → σ → m (α × σ)gsgs:GameStateprintBoardprintBoard:{Board : Prop} → Board → IO Unitg.g:GameStateboardboard:GameState → ∀ {Board : Prop}, Board-- [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`

:

StateMStateM:Type → Type → TypeGameStateGameState:TypeBool -- GameState → Bool × GameStateBool:Type

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

nextTurnManually : `StateM`**: **Type → Type → Type

StateM `GameState`**: **Type

GameState `Bool`**: **Type

Bool
| `state`**: **GameState

state =>
let (`i`**: **TileIndex

i, `gs`**: **GameState

gs) := `chooseRandomMove`**: **StateM GameState TileIndex

chooseRandomMove |>.`run`**: **{σ : Type} → {m : Type → Type} → {α : Type} → StateT σ m α → σ → m (α × σ)

run `state`**: **GameState

state
let (_, `gs'`**: **GameState

gs') := `applyMove`**: **TileIndex → StateM GameState Unit

applyMove `i`**: **TileIndex

i |>.`run`**: **{σ : Type} → {m : Type → Type} → {α : Type} → StateT σ m α → σ → m (α × σ)

run `gs`**: **GameState

gs
let (`result`**: **Bool

result, `gs''`**: **GameState

gs'') := `isGameDone`**: **StateM GameState Bool

isGameDone |>.`run`**: **{σ : Type} → {m : Type → Type} → {α : Type} → StateT σ m α → σ → m (α × σ)

run `gs'`**: **GameState

gs'
(`result`**: **Bool

result, `gs''`**: **GameState

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.