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). -/