Using lens, cosmosOf, uniplate, and the State monad to extract info about an AST

Question

I have the following code that traverses an AST using cosmosOf and uniplate looking for nodes of a certain type. For any that it finds, it sets a Bool flag in a record that is propagated using a State monad with the help of the lens package.

This all works, but feels pretty heavy-handed. It feels like lenses, the State monad, and possibly cosmosOf/uniplate may all be overkill here. Is there a better or more idiomatic way to do this?

{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE TemplateHaskell #-}

module Docvim.Visitor.Section (getSectionInfo) where

import Control.Lens
import Control.Monad.State
import Data.Data.Lens (uniplate)
import Docvim.AST

data SectionInfo = SectionInfo { _hasCommand :: Bool
                               , _hasCommands :: Bool
                               , _hasFunction :: Bool
                               , _hasFunctions :: Bool
                               , _hasMapping :: Bool
                               , _hasMappings :: Bool
                               , _hasOption :: Bool
                               , _hasOptions :: Bool
                               } deriving (Show)

type Env = State SectionInfo

makeLenses ''SectionInfo

defaultSectionInfo :: SectionInfo
defaultSectionInfo = SectionInfo { _hasCommand = False
                                 , _hasCommands = False
                                 , _hasFunction = False
                                 , _hasFunctions = False
                                 , _hasMapping = False
                                 , _hasMappings = False
                                 , _hasOption = False
                                 , _hasOptions = False
                                 }

getSectionInfo :: Node -> SectionInfo
getSectionInfo n = execState (mapMOf_ (cosmosOf uniplate) check n) defaultSectionInfo
  where
    check (CommandAnnotation {}) = hasCommand .= True
    check CommandsAnnotation     = hasCommands .= True
    check (FunctionAnnotation _) = hasFunction .= True
    check FunctionsAnnotation    = hasFunctions .= True
    check (MappingAnnotation _)  = hasMapping .= True
    check MappingsAnnotation     = hasMappings .= True
    check (OptionAnnotation {})  = hasOption .= True
    check OptionsAnnotation      = hasOptions .= True
    check _                      = modify id

Answer 1

What you want to do can be accomplished by para from the Uniplate module.

Basically para aggregates information collected from a node and its children and passed it up to the node's parent for further aggregation.

Here's a simplified version of your example - we determine whether or not a Node contains CommandAnnotation and/or FunctionAnnotation nodes

import Data.Monoid
import qualified Data.Set as Set
import qualified Data.Generics.Uniplate.Data as Uniplate
import Data.Data

...

data HasSection = HasCommandAnnotation | HasFunction | HasOther
  deriving (Show,Read,Enum,Bounded,Ord,Eq)

toHas :: Node -> HasSection  
toHas (CommandAnnotation {})   = HasCommandAnnotation
toHas (FunctionsAnnotation {}) = HasFunction
toHas _                        = HasOther

getSectionInfo :: Node -> Set.Set HasSection
getSectionInfo n = Uniplate.para visit n
  where visit n res = Set.singleton (toHas n) <> mconcat res

The README.md at the uniplate github repo has a good overview of the library with examples.

For better efficiency you could use the bitset package for the sets.