Search code examples
arrayshaskellhaskell-lens

Using lens for array indexing if both array and index are in State


I have an array and an array index in a state monad. I can read idx using use and modify it using += and other similar modifiers:

{-# Language TemplateHaskell #-}
import Control.Lens
import Control.Lens.TH
import Control.Monad.State
import Data.Array

data M = M { _arr :: Array Int Int, _idx :: Int }

$(makeLenses ''M)

foo x = do
    idx += x
    ii <- use idx
    return ii

Now I want to combine arr and idx to form a lens to arr[idx]:

combo arr idx = undefined

bar x = do
    combo arr idx += x
    ii <- combo arr idx
    return ii

How can I do this? Will code be different for Data.Sequence?


Solution

  • The answer turned out to be just

    combo arr idx f m = (arr . ix (m^.idx)) f m
    

    As the index may be out of bounds, ix is a partial lens called Traversal. So bar has to use uses instead of use:

    foo x = do
        combo arr idx += x
        ii <- uses $ combo arr idx
        return ii
    

    Also ii is Monoid m => m Int instead of Int, because of partiality.

    If original unsafe behaviour of returning Int is needed it can be restored by replacing uses with unsafeUses traversal = (^?! traversal) <$> get