Generating RIO-style Has-classes/lenses via template haskell

Example.hs

data WindowSize = MkWindowSize { windowWidth  :: !Natural
                               , windowHeight :: !Natural
                               }
makeRioClassy ''WindowSize
-- This will generate
-- class (HasWindowWidth env, HasWindowHeight env)
--    => HasWindowSize env where windowSizeL :: Lens' env WindowSize
-- class HasWindowWidth env where ... (method signatures omitted for brevity from here on out)
-- class HasWindowHeight env
-- (as well as associated `id` instances, omitted for brevity for this and the other types)

newtype Email = MkEmail { emailString :: String }
makeRioClassy ''Email
-- This will generate
-- class HasEmailString env => HasEmail env
-- class HasEmailString env

data Config = MkConfig { windowSize   :: !WindowSize
                       , verbose      :: !Bool
                       , studentEmail :: !Email
                       , teacherEmail :: !Email
                       }
makeRioClassy ''Config
-- This will *not* generate class HasWindowSize, since that class already exists
-- However, it *will* generate
-- instance HasWindowSize Config -- as well as instances for HasWindowWidth and HasWindowHeight
-- class HasVerbose env
-- class HasStudentEmail env -- since generation is name-driven it will *not* generate a HasEmail instance
-- class HasTeacherEmail env -- ditto
-- class (HasWindowSize env, HasVerbose env, HasStudentEmail env, HasTeacherEmail env)
--    => HasConfig env

data App = MkApp { config  :: !Config
                 , logFunc :: !LogFunc
                 }
makeRioClassy ''App
-- This will generate
-- class (HasConfig env, HasLogFunc env) => HasApp env
-- as well as
-- instance HasConfig App, instance HasWindowSize App, instance HasWindowWidth App, ...
-- instance HasLogFunc App

Generated.hs

class (HasWindowWidth env,
       HasWindowHeight env) => HasWindowSize env
    where windowSizeL :: Lens' env Main.WindowSize
instance HasWindowSize Main.WindowSize
    where windowSizeL = id
class HasWindowWidth env
    where windowWidthL :: Lens' env GHC.Num.Natural.Natural
class HasWindowHeight env
    where windowHeightL :: Lens' env GHC.Num.Natural.Natural
instance HasWindowWidth Main.WindowSize
    where windowWidthL = Lens.Micro.lens Main.windowWidth (\x_0 y_1 -> x_0{Main.windowWidth = y_1})
instance HasWindowHeight Main.WindowSize
    where windowHeightL = Lens.Micro.lens Main.windowHeight (\x_2 y_3 -> x_2{Main.windowHeight = y_3})
class HasEmailString env => HasEmail env
    where emailL :: Lens' env Main.Email
instance HasEmail Main.Email
    where emailL = id
class HasEmailString env
    where emailStringL :: Lens' env GHC.Base.String
instance HasEmailString Main.Email
    where emailStringL = Lens.Micro.lens Main.emailString (\x_0 y_1 -> x_0{Main.emailString = y_1})
class (HasWindowSize env,
       HasVerbose env,
       HasStudentEmail env,
       HasTeacherEmail env) => HasConfig env
    where configL :: Lens' env Main.Config
instance HasConfig Main.Config
    where configL = id
class HasVerbose env
    where verboseL :: Lens' env GHC.Types.Bool
class HasStudentEmail env
    where studentEmailL :: Lens' env Main.Email
class HasTeacherEmail env
    where teacherEmailL :: Lens' env Main.Email
instance HasWindowSize Main.Config
    where windowSizeL = Lens.Micro.lens Main.windowSize (\x_0 y_1 -> x_0{Main.windowSize = y_1})
instance HasVerbose Main.Config
    where verboseL = Lens.Micro.lens Main.verbose (\x_2 y_3 -> x_2{Main.verbose = y_3})
instance HasStudentEmail Main.Config
    where studentEmailL = Lens.Micro.lens Main.studentEmail (\x_4 y_5 -> x_4{Main.studentEmail = y_5})
instance HasTeacherEmail Main.Config
    where teacherEmailL = Lens.Micro.lens Main.teacherEmail (\x_6 y_7 -> x_6{Main.teacherEmail = y_7})
instance Main.HasWindowWidth Main.Config
    where Main.windowWidthL = windowSizeL . Main.windowWidthL
instance Main.HasWindowHeight Main.Config
    where Main.windowHeightL = windowSizeL . Main.windowHeightL
class (HasConfig env, HasLogFunc env) => HasApp env
    where appL :: Lens' env Main.App
instance HasApp Main.App
    where appL = id
instance HasConfig Main.App
    where configL = Lens.Micro.lens Main.config (\x_0 y_1 -> x_0{Main.config = y_1})
instance HasLogFunc Main.App
    where logFuncL = Lens.Micro.lens Main.logFunc (\x_2 y_3 -> x_2{Main.logFunc = y_3})
instance Main.HasWindowSize Main.App
    where Main.windowSizeL = configL . Main.windowSizeL
instance Main.HasWindowWidth Main.App
    where Main.windowWidthL = configL . Main.windowWidthL
instance Main.HasWindowHeight Main.App
    where Main.windowHeightL = configL . Main.windowHeightL
instance Main.HasVerbose Main.App
    where Main.verboseL = configL . Main.verboseL
instance Main.HasStudentEmail Main.App
    where Main.studentEmailL = configL . Main.studentEmailL
instance Main.HasTeacherEmail Main.App
    where Main.teacherEmailL = configL . Main.teacherEmailL
class HasStuffContent env => HasOtherStuff env
    where otherStuffL :: Lens' env Main.OtherStuff
instance HasOtherStuff Main.OtherStuff
    where otherStuffL = id
class HasStuffContent env
    where stuffContentL :: Lens' env GHC.Base.String
instance HasStuffContent Main.OtherStuff
    where stuffContentL = Lens.Micro.lens Main.stuffContent (\x_0 y_1 -> x_0{Main.stuffContent = y_1})
class (HasSaApp env, HasSaOtherStuff env) => HasSuperApp env
    where superAppL :: Lens' env Main.SuperApp
instance HasSuperApp Main.SuperApp
    where superAppL = id
class HasSaApp env
    where saAppL :: Lens' env Main.App
class HasSaOtherStuff env
    where saOtherStuffL :: Lens' env Main.OtherStuff
instance HasSaApp Main.SuperApp
    where saAppL = Lens.Micro.lens Main.saApp (\x_0 y_1 -> x_0{Main.saApp = y_1})
instance HasSaOtherStuff Main.SuperApp
    where saOtherStuffL = Lens.Micro.lens Main.saOtherStuff (\x_2 y_3 -> x_2{Main.saOtherStuff = y_3})

TH.hs

{-# LANGUAGE BlockArguments, LambdaCase, TupleSections, TemplateHaskellQuotes #-}

module TH where

import           RIO
import qualified RIO.Char as C
import           RIO.Lens
import           Language.Haskell.TH

type Field = VarBangType

makeRioClassy :: Name -> DecsQ
makeRioClassy tyConName = do
  fields <- reify tyConName >>= \case
    TyConI (DataD _ _ [] _ [RecC _ fields] _) -> pure fields
    TyConI (NewtypeD _ _ [] _ (RecC _ fields) _) -> pure fields
    _ -> fail "Unsupported declaration"
  (orphanFields, parentedFields) <-
    partitionEithers . map (\(x, my) -> maybe (Left x) (Right . (x,)) my) <$> mapM attachClass fields
  superInsts <- concat <$> mapM mkSuperInsts parentedFields
  insts <- mapM mkInstance fields
  let result = [thisClass fields, thisInstance]
            <> map mkClass orphanFields
            <> insts
            <> superInsts
  -- traceShowM . ppr $ result -- uncomment to print the generated code
  pure result
  where
    tyCon :: Type
    tyCon = ConT tyConName

    attachClass :: Field -> Q (Field, Maybe Name)
    attachClass field = findClass >>= \x -> pure (field, x)
      where
        findClass = lookupTypeName . nameBase $ className field

    thisInstance :: Dec
    thisInstance =
      InstanceD Nothing [] (AppT (ConT thisClassName) tyCon)
        [FunD thisLensName [Clause [] (NormalB (VarE $ mkName "id")) []]]

    thisLensName :: Name
    thisLensName = mapName (lensSuffix . (ix 0 %~ C.toLower)) tyConName

    thisLens :: Dec
    thisLens = SigD thisLensName (envLens' tyCon)

    thisClass :: [Field] -> Dec
    thisClass fields = ClassD ctxt thisClassName envParam [] [thisLens]
      where
        ctxt :: Cxt
        ctxt = map ((\cls -> AppT (ConT cls) (VarT env)) . className) fields

    env :: Name
    env = mkName "env"

    envParam :: [TyVarBndr ()]
    envParam = [PlainTV env ()]

    thisClassName :: Name
    thisClassName = mapName classPrefix tyConName

    classPrefix :: String -> String
    classPrefix = ("Has" <>)

    lensSuffix :: String -> String
    lensSuffix = (<> "L")

    fieldLens :: Field -> Name
    fieldLens (fieldName, _, _) = mapName lensSuffix fieldName

    mapName :: (String -> String) -> Name -> Name
    mapName f = mkName . f . nameBase

    envLens' :: Type -> Type
    envLens' = AppT (AppT (ConT $ mkName "Lens'") (VarT env))

    mkClass :: Field -> Dec
    mkClass field =
      ClassD [] (className field) envParam [] [mkMethod field]

    className :: Field -> Name
    className = mapName (classPrefix . (ix 0 %~ C.toUpper)) . view _1

    mkMethod :: Field -> Dec
    mkMethod (name, _, fieldType) =
      SigD (mapName lensSuffix name) (envLens' fieldType)

    mkInstance :: Field -> Q Dec
    mkInstance field = InstanceD Nothing [] (AppT (ConT $ className field) tyCon) . pure <$>
      mkImpl field
      where
        mkImpl :: Field -> Q Dec
        mkImpl (fieldName, _, _) = do
          b <- body
          pure $ FunD (mapName lensSuffix fieldName) [Clause [] b []]
          where
            body :: Q Body
            body = do
              x <- newName "x"
              y <- newName "y"
              NormalB <$> [| lens $(varE fieldName) \ $(varP x) $(varP y) ->
                $(recUpdE (varE x) [pure (fieldName, VarE y)]) |]

    mkSuperInsts :: (Field, Name) -> DecsQ
    mkSuperInsts (field, cls) = do
      reify cls >>= \case
        ClassI (ClassD ctxt _ _ _ _) _ -> concat <$> mapM superInstsForPred ctxt
        _ -> fail $ "Expected " <> show cls <> " to be a class, but it's not"
      where
        mkSuperInstsRec :: Name -> DecsQ
        mkSuperInstsRec name = reify name >>= \case
          ClassI (ClassD ctxt _ _ _ decs) _ -> do
            inst <- superInstHead . concat <$> mapM mkSuperImpl decs
            (inst :) . concat <$> mapM superInstsForPred ctxt
          _ -> fail $
            "Couldn't make instance for " <> show name <> " - it's not a class"
          where
            superInstHead :: [Dec] -> Dec
            superInstHead = InstanceD Nothing [] (AppT (ConT name) (ConT tyConName))

        superInstsForPred :: Pred -> DecsQ
        superInstsForPred = \case
          AppT (ConT superCls) (VarT _) -> mkSuperInstsRec superCls
          constraint ->
            fail $ "Unsupported superclass constraint " <> show (ppr constraint)

        mkSuperImpl :: Dec -> DecsQ
        mkSuperImpl (SigD methName _) =
          [d| $(varP methName) = $(varE $ fieldLens field) . $(varE methName) |]
        mkSuperImpl _ = pure []