Convenient ways to initialize nested fields with lenses

Question

I have some data type which is very similar to ordinary trees, just some specialized form.

data NestedTree = NT
    { _dummy :: Int
    , _tree  :: HashMap String NestedTree
    } deriving (Show)

makeLenses ''NestedTree

I want to initialize instances of my data type imperatively using lenses. Here is what I got now:

example :: NestedTree
example = flip execState (NT 0 mempty) $ do
    dummy .= 3
    tree.at "foo" ?= flip execState (NT 0 mempty) (dummy .= 10)

You can observe in this example that I can replace first (NT 0 mempty) with (NT 3 mempty) but this is not the point. What I want is too be able to initialize nested HashMaps using this nice imperative style. More precise, I want to be able to write something like this:

example :: NestedTree
example = flip execState (NT 0 mempty) $ do
    dummy .= 3
    tree.at "foo" ?= flip execState (NT 0 mempty) $ do
        dummy .= 10
        tree.at "foo nested" ?= NT 5 mempty
    tree.at "bar" ?= flip execState (NT 0 mempty) $ do
        dummy .= 15
        tree.at "bar nested" ?= NT (-3) mempty

My real data structure is more complex and it very soon becomes really ugly to initialize it using just simple records. Thus I want to use some kind of DSL, and lenses fit quite good for my need. But you can notice that code above doesn't compile.

It's because ($) has lowest precedence and I can't just write tree.at "foo" ?= flip execState (NT 0 mempty) $ do. But I really don't want to add () around nested dos.

Is there any nice ways to mix arbitrary operators with $ and do in order to write such functions? I really don't want to introduce some helper like wordsAssign = (?=) and call functions like

wordsAssign (tree.at "foo") $ flip execState (NT 0 mempty) $ do

because I like ?= operator. Maybe I'm doing everything wrong and this kind of things I want to do can be done without lenses with some hand-written operators?

Answer 1

zoom is tailor-made for handling nested state updates. In your case, unfortunately, the Maybe-ness makes using it slightly awkward:

example :: NestedTree
example = flip execState (NT 0 mempty) $ do
    dummy .= 3
    zoom (tree.at "foo") $ do
        put (Just (NT 0 mempty))
        _Just.dummy .= 10
        _Just.tree.at "foo nested" ?= NT 5 mempty
    -- Or, using zoom one more time:
    zoom (tree.at "bar") $ do
        put (Just (NT 0 mempty))
        zoom _Just $ do
            dummy .= 15
            tree.at "bar nested" ?= NT (-3) mempty

For the sake of comparison, if you didn't need to insert new keys at the outer level, you would be able to use ix instead of at and drop all the Maybe-related boilerplate:

    zoom (tree.ix "foo") $ do
        dummy .= 10
        tree.at "foo nested" ?= NT 5 mempty