TH.hs

{-# LANGUAGE CPP #-}
{-# LANGUAGE TemplateHaskellQuotes #-}

module Data.Recollections.TH
  (  -- * Collections generator
    mkCollection
  , mkIndices

    -- * Distributive
    -- $distributive
  , mkDistributive
  , mkDistributiveFor
  , mkDistribute

    -- * Representable
    -- $representable
  , mkRepresentable
  , mkRepresentableFor
  , mkIndex
  , mkTabulate

    -- * Tag instances
    -- $tagInstances
  , mkBounded
  , mkEnum
  ) where

import Data.Char
import Data.Foldable
import Data.Traversable
import GHC.Enum (boundedEnumFrom, boundedEnumFromThen)
import GHC.Generics (Generic, Generic1, Generically1)
import Language.Haskell.TH
import Language.Haskell.TH.Syntax (mkNameG_tc, mkNameG_v)

{- | Generate a @Collection a@ type from a enum-like type.

Every constructor is represented by a field.
Reserved words like @type@ get a @'@ suffix.'

> data Things = This | That
>   deriving (Eq, Ord, Show, Enum, Bounded)
>
> mkCollection ''Things
>
> -- resulting splice
> data Collection a = Collection { this, that :: a }
>   deriving (Eq, Show, Generic, Generic1, Functor, Foldable, Traversable)
>   deriving Applicative via (Generically1 Collection)

A constructor may wrap the tag type of another collection, generated in the module that defines that tag.
The field then holds that whole collection:

> import Some qualified
>
> data Tag = Something Some.Tag | Local
>
> mkCollection ''Tag
>
> -- resulting splice
> data Collection a = Collection { something :: Some.Collection a, local :: a }

Every other collection generator applied to a nested tag requires the same generator in the nested module.
The instance generators call the nested instances instead, so the nested module needs the same instance generators.
-}
mkCollection :: Name -> Q [Dec]
mkCollection tags = do
  let nothingsBanger = bang noSourceUnpackedness noSourceStrictness

  names <- tagNames tags
  fields <- for names \tag -> do
    fieldType <- case tag of
      Leaf _ -> pure $ varT (mkName "a")
      Nested _ sub -> do
        subCollection <- siblingOf mkNameG_tc "Collection" sub
        pure $ conT subCollection `appT` varT (mkName "a")
    pure $ varBangType (mkName . fieldNameOf $ tagName tag) $ bangType nothingsBanger fieldType
  let constr = recC collectionName fields
  let
    derivs =
      [ derivClause (Just StockStrategy)
          [ conT ''Show
          , conT ''Eq
          , conT ''Generic
          , conT ''Generic1
          , conT ''Functor
          , conT ''Foldable
          , conT ''Traversable
          ]
      , derivClause (Just $ ViaStrategy $ ConT ''Generically1 `AppT` ConT collectionName)
        [ conT ''Applicative
        ]
      ]
  pure <$> dataD mempty collectionName [collectionTyVar] Nothing [constr] derivs

-- | The @a@ binder of @data Collection a@.
--
-- @template-haskell-2.21@ (GHC 9.8) changed the binder flag of 'dataD'
-- from @()@ to 'BndrVis'.
#if MIN_VERSION_template_haskell(2,21,0)
collectionTyVar :: TyVarBndrVis
collectionTyVar = PlainTV (mkName "a") BndrReq
#else
collectionTyVar :: TyVarBndrUnit
collectionTyVar = PlainTV (mkName "a") ()
#endif

{- | Generate a value filled with the indices matching the fields.

> indices :: Collection Things
> indices = Collection { this = This, that = That }

This is useful with the Applicative instance to provide indexed operations:

> indexed c :: Collection a -> Collection (Things, a)
> indexed c = (,) <$> indices <*> c
-}
mkIndices :: Name -> Q [Dec]
mkIndices tags = do
  let indicesName = mkName "indices"
  names <- tagNames tags
  sig <- sigD indicesName $ conT collectionName `appT` conT tags
  fields <- for names \case
    Leaf name -> pure . conE $ mkName name
    Nested name sub -> do
      subIndices <- siblingOf mkNameG_v "indices" sub
      pure $ appE (appE (varE 'fmap) (conE $ mkName name)) (varE subIndices)
  let body = foldl' appE (conE collectionName) fields
  fun <- funD indicesName [ clause [] (normalB body) [] ]
  pure [sig, fun]

-- * Distributive

{- $distributive

<https://hackage-content.haskell.org/package/distributive-0.6.3/docs/Data-Distributive.html>

The instance generators need the @distributive@ package as a dependency of the splicing package,
not of @recollections@.
The class is looked up in the splicing module scope, so it has to be imported there.
-}

{- | Generate a Distributive instance.

The splicing module must import "Data.Distributive", unqualified or qualified without an alias.
Use 'mkDistributiveFor' when the module is imported under an alias.

> instance Distributive Collection where
>   distribute f = Collection
>     { this = this <$> f
>     , that = that <$> f
>     }
-}
mkDistributive :: Name -> Q [Dec]
mkDistributive tags = do
  distributive <- lookupClass "Data.Distributive" "Distributive"
  mkDistributiveFor distributive tags

{- | Generate a Distributive instance for the explicitly provided class name.

> import qualified Data.Distributive as D
>
> mkDistributiveFor ''D.Distributive ''Things
-}
mkDistributiveFor :: Name -> Name -> Q [Dec]
mkDistributiveFor distributive tags = do
  names <- tagNames tags
  distributeName <- classMember mkNameG_v distributive "distribute"
  pure <$> instanceD (cxt []) (conT distributive `appT` conT collectionName)
    [ funD distributeName [distributeClause (const $ pure distributeName) names]
    , inlineP distributeName
    ]

{- | Generate a dual of sequenceA.

Doesn't need the @distributive@ package, but clashes with the "Data.Distributive" import.

> distribute :: Functor f => f (Collection a) -> Collection (f a)
> distribute f = Collection
>   { this      = this      <$> f
>   , that      = that      <$> f
>   , something = something <$> f
>   , else'     = else'     <$> f
>   , entirely  = entirely  <$> f
>   }
-}
mkDistribute :: Name -> Q [Dec]
mkDistribute tags = do
  let distributeName = mkName "distribute"
  names <- tagNames tags
  let f = mkName "f"
  let a = mkName "a"
  sig <- sigD distributeName $
    forallT [] (cxt [conT ''Functor `appT` varT f]) $
      appT (appT arrowT (varT f `appT` (conT collectionName `appT` (varT a)))) $
        (conT collectionName `appT` (varT f `appT` varT a))
  fun <- funD distributeName [distributeClause (siblingOf mkNameG_v "distribute") names]
  inl <- inlineP distributeName
  pure [sig, fun, inl]

distributeClause :: (Name -> Q Name) -> [Tag] -> Q Clause
distributeClause nested names = do
  {- The binder must not shadow any field selector the body refers to by
     'mkName' (those are resolved lexically, by occurrence name). Field names
     always start with a lowercase letter, so a leading underscore is safe;
     plain @f@ would break for a tag named @F@. -}
  arg <- newName "_f"
  fields <- for names \tag -> do
    let projected = appE (varE 'fmap) (varE . mkName . fieldNameOf $ tagName tag) `appE` varE arg
    case tag of
      Leaf _ -> pure projected
      Nested _ sub -> do
        subDistribute <- nested sub
        pure $ appE (varE subDistribute) projected
  let body = foldl' appE (conE collectionName) fields
  clause [varP arg] (normalB body) []

-- * Representable

{- $representable

<https://hackage-content.haskell.org/package/adjunctions-4.4.4/docs/Data-Functor-Rep.html>

The instance generators need the @adjunctions@ package as a dependency of the splicing package,
not of @recollections@.
The class is looked up in the splicing module scope, so it has to be imported there.
The splicing module also needs @TypeFamilies@ for the @Rep@ instance.
-}

{- | Generate a Representable instance.

The splicing module must import "Data.Functor.Rep", unqualified or qualified without an alias.
Use 'mkRepresentableFor' when the module is imported under an alias.

Representable requires a Distributive instance, see 'mkDistributive'.

> instance Representable Collection where
>   type Rep Collection = Things
>   index c = \case This -> this c; That -> that c
>   tabulate k = Collection { this = k This, that = k That }
-}
mkRepresentable :: Name -> Q [Dec]
mkRepresentable tags = do
  representable <- lookupClass "Data.Functor.Rep" "Representable"
  mkRepresentableFor representable tags

{- | Generate a Representable instance for the explicitly provided class name.

> import qualified Data.Functor.Rep as R
>
> mkRepresentableFor ''R.Representable ''Things
-}
mkRepresentableFor :: Name -> Name -> Q [Dec]
mkRepresentableFor representable tags = do
  names <- tagNames tags
  repName <- classMember mkNameG_tc representable "Rep"
  indexName <- classMember mkNameG_v representable "index"
  tabulateName <- classMember mkNameG_v representable "tabulate"
  pure <$> instanceD (cxt []) (conT representable `appT` conT collectionName)
    [ tySynInstD $ tySynEqn Nothing (conT repName `appT` conT collectionName) (conT tags)
    , funD indexName [indexClause (const $ pure indexName) names]
    , inlineP indexName
    , funD tabulateName [tabulateClause (const $ pure tabulateName) names]
    , inlineP tabulateName
    ]

{- | Read a collection field using an index value.

Doesn't need the @adjunctions@ package, but clashes with the "Data.Functor.Rep" import.

> index :: Collection a -> Things -> a
> index c = \case This -> this c; That -> that c; …
-}
mkIndex :: Name -> Q [Dec]
mkIndex tags = do
  let indexName = mkName "index"
  names <- tagNames tags
  let a = mkName "a"
  sig <- sigD indexName $
    appT (appT arrowT (conT collectionName `appT` varT a)) $
      appT (appT arrowT (conT tags)) (varT a)
  fun <- funD indexName [indexClause (siblingOf mkNameG_v "index") names]
  inl <- inlineP indexName
  pure [sig, fun, inl]

indexClause :: (Name -> Q Name) -> [Tag] -> Q Clause
indexClause nested names = do
  bounds <- for names \tag ->
    (,) tag <$> newName (fieldNameOf $ tagName tag)
  matches <- for bounds \case
    (Leaf name, bound) ->
      pure $ match (conP (mkName name) []) (normalB $ varE bound) []
    (Nested name sub, bound) -> do
      subIndex <- nested sub
      subTag <- newName "t"
      pure $ match
        (conP (mkName name) [varP subTag])
        (normalB $ varE subIndex `appE` varE bound `appE` varE subTag)
        []
  let body = lamCaseE matches
  clause
    [ conP collectionName $ map (varP . snd) bounds
    ]
    (normalB body)
    []

{- | Generate a collection with a function from its indices

Doesn't need the @adjunctions@ package, but clashes with the "Data.Functor.Rep" import.

> tabulate :: (Things -> a) -> Collection a
> tabulate k = Collection { this = k This, that = k That }
-}
mkTabulate :: Name -> Q [Dec]
mkTabulate tags = do
  let tabulateName = mkName "tabulate"
  names <- tagNames tags
  let a = mkName "a"
  sig <- sigD tabulateName $
    appT (appT arrowT (appT (appT arrowT (conT tags)) (varT a))) $
      conT collectionName `appT` varT a
  fun <- funD tabulateName [tabulateClause (siblingOf mkNameG_v "tabulate") names]
  inl <- inlineP tabulateName
  pure [sig, fun, inl]

tabulateClause :: (Name -> Q Name) -> [Tag] -> Q Clause
tabulateClause nested names = do
  f <- newName "_f"
  fields <- for names \case
    Leaf name -> pure $ appE (varE f) $ conE (mkName name)
    Nested name sub -> do
      subTabulate <- nested sub
      pure $ appE (varE subTabulate) $ infixE (Just $ varE f) (varE '(.)) (Just . conE $ mkName name)
  let body = foldl' appE (conE collectionName) fields
  clause [varP f] (normalB body) []

-- * Tag instances

{- $tagInstances

Stock 'Bounded' and 'Enum' only work for nullary constructors.
These generators also handle nested tags, using the 'Bounded' and 'Enum' instances of the nested tag types.
Those may be derived or generated.
-}

{- | Generate a 'Bounded' instance for the tag type.

> data Tag = Front Inner.Tag | Local | Back Inner.Tag
>
> instance Bounded Tag where
>   minBound = Front minBound
>   maxBound = Back maxBound
-}
mkBounded :: Name -> Q [Dec]
mkBounded tags = do
  names <- tagNames tags
  (first, final) <- case (names, reverse names) of
    (first : _, final : _) -> pure (first, final)
    _ -> fail $ "Can't bound a type without constructors: " <> show tags
  let
    bound method = \case
      Leaf name -> conE $ mkName name
      Nested name _ -> conE (mkName name) `appE` varE method
  pure <$> instanceD (cxt []) (conT ''Bounded `appT` conT tags)
    [ valD (varP 'minBound) (normalB $ bound 'minBound first) []
    , valD (varP 'maxBound) (normalB $ bound 'maxBound final) []
    ]

{- | Generate an 'Enum' instance for the tag type, numbering the nested tags in place.

> data Tag = Front Inner.Tag | Local | Back Inner.Tag
>
> instance Enum Tag where
>   fromEnum t = case t of
>     Front x -> o0 + (fromEnum x - lo)
>     Local -> o1
>     Back x -> o2 + (fromEnum x - lo)
>     where
>       o0 = 0
>       o1 = o0 + (hi - lo + 1)
>       o2 = o1 + 1
>   toEnum n
>     | n < 0 = errorWithoutStackTrace …
>     | n < o1 = Front (toEnum (n - o0 + lo))
>     | n < o2 = Local
>     | n < o3 = Back (toEnum (n - o2 + lo))
>     | otherwise = errorWithoutStackTrace …
>     where
>       …
>       o3 = o2 + (hi - lo + 1)
>   enumFrom = boundedEnumFrom
>   enumFromThen = boundedEnumFromThen

Here @lo@ and @hi@ stand for the inlined @fromEnum (minBound :: Inner.Tag)@ and @fromEnum (maxBound :: Inner.Tag)@.

The instance requires @Bounded Tag@, see 'mkBounded'.
-}
mkEnum :: Name -> Q [Dec]
mkEnum tags = do
  names <- tagNames tags
  fromEnumClause <- do
    (offsets, bindings) <- runningOffsets . reverse . drop 1 $ reverse names
    t <- newName "t"
    matches <- fromEnumMatches $ zip names offsets
    pure $ clause [varP t] (normalB $ caseE (varE t) matches) bindings
  toEnumClause <- do
    (offsets, bindings) <- runningOffsets names
    n <- newName "n"
    pure $ clause [varP n] (toEnumBody tags n $ zip3 names offsets (drop 1 offsets)) bindings
  pure <$> instanceD (cxt []) (conT ''Enum `appT` conT tags)
    [ funD 'fromEnum [fromEnumClause]
    , funD 'toEnum [toEnumClause]
    , valD (varP 'enumFrom) (normalB $ varE 'boundedEnumFrom) []
    , valD (varP 'enumFromThen) (normalB $ varE 'boundedEnumFromThen) []
    ]

fromEnumMatches :: [(Tag, Name)] -> Q [Q Match]
fromEnumMatches spans =
  for spans \case
    (Leaf name, before) ->
      pure $ match (conP (mkName name) []) (normalB $ varE before) []
    (Nested name sub, before) -> do
      subTag <- newName "t"
      pure $ match
        (conP (mkName name) [varP subTag])
        (normalB $ plusE (varE before) (minusE (varE 'fromEnum `appE` varE subTag) (fromEnumOf 'minBound sub)))
        []

toEnumBody :: Name -> Name -> [(Tag, Name, Name)] -> Q Body
toEnumBody tags n spans =
  guardedB $
    [normalGE (lessThan . litE $ integerL 0) outOfRange]
    <> map pick spans
    <> [normalGE (varE 'otherwise) outOfRange]
  where
    lessThan bound = infixE (Just $ varE n) (varE '(<)) (Just bound)
    pick = \case
      (Leaf name, _, after) ->
        normalGE (lessThan $ varE after) (conE $ mkName name)
      (Nested name sub, before, after) ->
        normalGE (lessThan $ varE after) $
          conE (mkName name) `appE`
            (varE 'toEnum `appE` plusE (minusE (varE n) (varE before)) (fromEnumOf 'minBound sub))
    outOfRange =
      varE 'errorWithoutStackTrace `appE`
        infixE
          (Just . stringE $ "toEnum{" <> nameBase tags <> "}: out of range: ")
          (varE '(<>))
          (Just $ varE 'show `appE` varE n)

runningOffsets :: [Tag] -> Q ([Name], [Q Dec])
runningOffsets names = do
  first <- newName "o0"
  rest <- for [1 .. length names] \i -> newName $ "o" <> show i
  let offsets = first : rest
  let
    sizeE = \case
      Leaf _ -> litE $ integerL 1
      Nested _ sub -> plusE (minusE (fromEnumOf 'maxBound sub) (fromEnumOf 'minBound sub)) (litE $ integerL 1)
    start = valD (varP first) (normalB . litE $ integerL 0) []
    steps = do
      (before, after, tag) <- zip3 offsets (drop 1 offsets) names
      pure $ valD (varP after) (normalB $ plusE (varE before) (sizeE tag)) []
  pure (offsets, start : steps)

fromEnumOf :: Name -> Name -> Q Exp
fromEnumOf bound sub = varE 'fromEnum `appE` sigE (varE bound) (conT sub)

plusE :: Q Exp -> Q Exp -> Q Exp
plusE a b = infixE (Just a) (varE '(+)) (Just b)

minusE :: Q Exp -> Q Exp -> Q Exp
minusE a b = infixE (Just a) (varE '(-)) (Just b)

-- * Utils

collectionName :: Name
collectionName = mkName "Collection"

lookupClass :: String -> String -> Q Name
lookupClass moduleName className = do
  found <- traverse lookupTypeName [moduleName <> "." <> className, className]
  case asum found of
    Just name ->
      pure name
    Nothing ->
      fail $ concat
        [ className, " is not in scope. Import ", moduleName
        , " or use mk", className, "For with an explicit class name."
        ]

classMember :: (String -> String -> String -> Name) -> Name -> String -> Q Name
classMember mkMember className member =
  case (namePackage className, nameModule className) of
    (Just package, Just moduleName) ->
      pure $ mkMember package moduleName member
    _ ->
      fail $ "Expected an imported class name, got: " <> show className

data Tag
  = Leaf String
  | Nested String Name

tagName :: Tag -> String
tagName = \case
  Leaf name -> name
  Nested name _ -> name

tagNames :: Name -> Q [Tag]
tagNames tags =
 reify tags >>= \case
    TyConI (DataD _ _ _ _ constructors _) ->
      foldrM (flip $ extractTags tags) [] constructors
    _ ->
      fail "Expected a type constructor name"

extractTags :: Name -> [Tag] -> Con -> Q [Tag]
extractTags tags acc = \case
  NormalC name [] ->
    pure $ Leaf (nameBase name) : acc
  NormalC name [(_, ConT sub)]
    | sub == tags ->
        fail $ "A tag can't nest its own type: " <> show name
    | otherwise ->
        pure $ Nested (nameBase name) sub : acc
  huh ->
    -- Skipping would produce a 'Collection' with fewer fields than there are
    -- tags, making the generated 'index' silently non-exhaustive.
    fail $ "Expected a nullary constructor or one wrapping a nested tag type, got: " <> show huh

siblingOf :: (String -> String -> String -> Name) -> String -> Name -> Q Name
siblingOf mkNameG occ sub =
  case (namePackage sub, nameModule sub) of
    (Just pkg, Just modName) -> do
      let sibling = mkNameG pkg modName occ
      recover
        (fail $ "Nested tag " <> show sub <> " needs " <> occ <> " generated in " <> modName)
        (sibling <$ reify sibling)
    _ ->
      fail $ "Nested tag " <> show sub <> " must be a top-level type from another module"

-- | Lowercase the leading character and dodge reserved words.
fieldNameOf :: String -> String
fieldNameOf = \case
  [] -> []
  n : ame -> legalizeVariableName $ toLower n : ame

legalizeVariableName :: String -> String
legalizeVariableName name
  | name `elem` illegal = name ++ "'"
  | otherwise = name

-- | Haskell 2010 reserved words, which cannot be used as record field names.
illegal :: [String]
illegal =
  [ "case", "class", "data", "default", "deriving", "do"
  , "else", "foreign", "if", "import", "in"
  , "infix", "infixl", "infixr", "instance"
  , "let", "module", "newtype", "of"
  , "then", "type", "where", "_"
  ]

inlineP :: Name -> Q Dec
inlineP name = pragInlD name Inline FunLike AllPhases