Search code examples
haskellfunctional-programminghaskell-lens

Functionally updating arbitrarily nested data structures with lenses


Say I have a data structure representing a Bag of Holding, which can hold multiple items. The user could place another Bag of Holding in this bag, and that bag could contain other bags, or even bags containing bags. Is there a lens for functionally updating arbitrarily nested bags, e.g. removing item foo from a bag inside a bag inside a bag inside a bag? Note that level of nesting, as well as the total depth of the tree, is dynamic, not known at compile time. Other questions like this and this seem to only deal with statically-known levels of nesting.

What I'm looking for can be done in Clojure with the update-in function, by generating a vector of accessors dynamically to pass to that function.


Solution

  • Your description of "Bag of Holding" wasn't precise but I think this is close to what you meant. The basic idea is to traverse into a child bag using a [Int] (similar to the Ixed instance for Tree) and use the At instance for Map to edit the items.

    {-# LANGUAGE FlexibleInstances #-}
    {-# LANGUAGE OverloadedLists   #-}
    {-# LANGUAGE RankNTypes        #-}
    {-# LANGUAGE TypeFamilies      #-}
    
    import           Control.Lens
    import qualified Data.Map     as M
    
    data Bag k a = Bag (M.Map k a) [Bag k a]
      deriving (Show)
    
    -- | Lens onto top level items of a bag.
    items :: Lens' (Bag k a) (M.Map k a)
    items f (Bag k a) = f k <&> \k' -> Bag k' a
    
    -- | Use 'At' instance for 'M.Map' to edit top level items.
    atItem :: Ord k => k -> Lens' (Bag k a) (Maybe a)
    atItem k = items . at k
    
    type instance Index (Bag k a)   = [Int]
    type instance IxValue (Bag k a) = Bag k a
    instance Ixed (Bag k a) where
      ix is0 f = go is0 where
        -- Use the `Ixed` instance for lists to traverse over
        -- item `i` in the list of bags.
        go (i:is) (Bag m bs) = Bag m <$> ix i (go is) bs
        go _      b          = f b
      {-# INLINE ix #-}
    
    mybag :: Bag String Char
    mybag =
      Bag [("a1",'a')] -- ix []
       [ Bag [] []     -- ix [0]
       , Bag []        -- ix [1]
         [ Bag [("foo", 'x'), ("bar",'y')] [] -- ix [1,0]
         , Bag [("FOO", 'X'), ("BAR",'Y')] [] -- ix [1,1]
         ]
      ]
    

    So now if we want to delete the "FOO" item from bag [1,1]:

    > mybag & ix [1,1] . atItem "FOO" .~ Nothing
    Bag (fromList [("a1",'a')])
      [Bag (fromList []) []
      ,Bag (fromList [])
         [Bag (fromList [("bar",'y'),("foo",'x')]) []
         ,Bag (fromList [("BAR",'Y')]) []]]
    

    or insert "foobar" into bag [1,0]:

    > mybag & ix [1,0] . atItem "foobar" ?~ 'z'
    Bag (fromList [("a1",'a')])
      [Bag (fromList []) []
      ,Bag (fromList [])
        [Bag (fromList [("bar",'y'),("foo",'x'),("foobar",'z')]) []
        ,Bag (fromList [("BAR",'Y'),("FOO",'X')]) []]]
    

    Actually my definition of a Bag was just a specialised Tree:

    import Data.Tree
    import Data.Tree.Lens
    
    type Bag k a = Tree (M.Map k a)
    
    atItem :: Ord k => k -> Lens' (Bag k a) (Maybe a)
    atItem k = root . at k
    
    subBag :: [Int] -> Traversal' (Bag k a) (Bag k a)
    subBag (i:is) = branches . ix i . subBag is
    subBag _      = id
    

    This can be used the same as before expect use subBag instead of ix. The definition of subBag is probably clearer written this way.

    In fact you don't need to write any new functions because the Ixed instance for Tree is the same as subBag is . root, so editing can be done by:

    > mybag & ix [1,1] . at "FOO" .~ Nothing