Haskell record to map/map to record converter

recmap.hs

{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TemplateHaskell #-}

module Routing.Duplex.TH where

import Control.Isomorphism.Partial (Iso(..))
import Data.Dynamic (Dynamic, fromDynamic, toDyn)
import Data.List (foldl')
import Data.Map (Map, fromList, lookup)
import Data.Maybe (mapMaybe)
import Language.Haskell.TH
import Language.Haskell.TH.Syntax
import Prelude hiding (lookup)

defineRecordIsomorphisms :: Name -> Q [Dec]
defineRecordIsomorphisms name = defineRecordIsomorphisms' name ("rec" <>)

defineRecordIsomorphisms' :: Name -> (String -> String) -> Q [Dec]
defineRecordIsomorphisms' name isoRename = do
  info <- reify name
  let ctors = case info of
        TyConI (DataD _ _ _ _ cs _) -> cs
        TyConI (NewtypeD _ _ _ _ c _) -> [c]
        _ -> error $ show name <> " neither denotes a data or newtype declaration. Found: " <> show info
  let recCtors =
        mapMaybe
          ( \case
              RecC name fields -> Just (name, fields)
              _ -> Nothing
          )
          ctors
  traverse (defineRecordIsomorphism isoRename) recCtors

defineRecordIsomorphism :: (String -> String) -> (Name, [VarBangType]) -> Q Dec
defineRecordIsomorphism isoRename (name, fields) = do
  let renamed = mkName $ isoRename $ nameBase name
  body <- [|Iso $(mkToMapFn fields) $(mkFromMapFn (name, fields))|]
  pure $ FunD renamed [Clause [] (NormalB body) []]
  where
    mkToMapFn :: [VarBangType] -> Q Exp
    mkToMapFn fields = do
      let recVarName = mkName "rec"
      LamE [VarP recVarName]
        <$> [|
          Just $
            fromList $(ListE <$> traverse (mkFieldPair recVarName) fields)
          |]
      where
        mkFieldPair :: Name -> VarBangType -> Q Exp
        mkFieldPair recVarName (fieldName, _, _) = do
          toDynFn <- [|toDyn|]
          pure $
            TupE
              [ LitE $ StringL $ nameBase fieldName,
                AppE toDynFn (AppE (VarE fieldName) (VarE recVarName))
              ]
    mkFromMapFn :: (Name, [VarBangType]) -> Q Exp
    mkFromMapFn (name, fields) = do
      let mapVarName = mkName "map"
      body <-
        DoE
          <$> sequence
            ( fmap (mkFieldExtractStmt mapVarName) fields
                <> [ do
                       let resultExp =
                             foldl' AppE (ConE name) $
                               fmap (\(fieldName, _, _) -> VarE (extractedFieldVarName fieldName)) fields
                       NoBindS <$> [|pure $(pure resultExp)|]
                   ]
            )
      pure $ LamE [VarP mapVarName] body
      where
        mkFieldExtractStmt :: Name -> VarBangType -> Q Stmt
        mkFieldExtractStmt mapVarName (fieldName, _, fieldType) =
          BindS (SigP (VarP (extractedFieldVarName fieldName)) fieldType)
            <$> [|
              fromDynamic
                =<< lookup
                  $(pure $ LitE $ StringL $ nameBase fieldName)
                  $(pure $ VarE mapVarName)
              |]
        extractedFieldVarName :: Name -> Name
        extractedFieldVarName fieldName = mkName $ "extracted_" <> nameBase fieldName