{-# LANGUAGE DataKinds           #-}
{-# LANGUAGE LambdaCase          #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TupleSections       #-}

-- | Additional utilities for working with @haskell-src-exts@
module Data.Record.QQ.CodeGen.HSE (
    -- * Language extensions
    extensionFromTH
  , processRecordPuns
    -- * Naming
  , fromHseName
  , resolveHseName
  , resolveKnownHseName
  ) where

import Data.Generics
import Language.Haskell.Exts
import Language.Haskell.TH.Syntax (Quasi)

import qualified Language.Haskell.TH.Syntax as TH

import qualified Data.Record.Internal.TH.Name as N

{-------------------------------------------------------------------------------
  Language extensions
-------------------------------------------------------------------------------}

-- | Translate TH extension into HSE extension
--
-- Useful in combination with 'extsEnabled'.
extensionFromTH :: TH.Extension -> Extension
extensionFromTH :: Extension -> Extension
extensionFromTH = \case
    Extension
TH.DataKinds        -> KnownExtension -> Extension
EnableExtension (KnownExtension -> Extension) -> KnownExtension -> Extension
forall a b. (a -> b) -> a -> b
$ KnownExtension
DataKinds
    Extension
TH.RecordPuns       -> KnownExtension -> Extension
EnableExtension (KnownExtension -> Extension) -> KnownExtension -> Extension
forall a b. (a -> b) -> a -> b
$ KnownExtension
NamedFieldPuns
    Extension
TH.TypeApplications -> KnownExtension -> Extension
EnableExtension (KnownExtension -> Extension) -> KnownExtension -> Extension
forall a b. (a -> b) -> a -> b
$ KnownExtension
TypeApplications
    Extension
TH.ViewPatterns     -> KnownExtension -> Extension
EnableExtension (KnownExtension -> Extension) -> KnownExtension -> Extension
forall a b. (a -> b) -> a -> b
$ KnownExtension
ViewPatterns

    -- We don't care about all extensions; there are many of them, and they vary
    -- from ghc version to ghc version. Treating them all would be a lot of work
    -- for little benefit. We assume that calling @show@ gives us a valid
    -- extension name; by and large this seems to be true (though for instance
    -- it will give us 'RecordPuns' rather than 'NamedFieldPuns', which although
    -- valid, is deprecated).
    Extension
e -> String -> Extension
UnknownExtension (String -> Extension) -> String -> Extension
forall a b. (a -> b) -> a -> b
$ Extension -> String
forall a. Show a => a -> String
show Extension
e

processRecordPuns :: forall l. Data l => Pat l -> Pat l
processRecordPuns :: Pat l -> Pat l
processRecordPuns = (forall a. Data a => a -> a) -> forall a. Data a => a -> a
everywhere ((PatField l -> PatField l) -> a -> a
forall a b. (Typeable a, Typeable b) => (b -> b) -> a -> a
mkT PatField l -> PatField l
go)
  where
    go :: PatField l -> PatField l
    go :: PatField l -> PatField l
go (PFieldPun l
pLoc n :: QName l
n@(UnQual l
nLoc Name l
n')) = l -> QName l -> Pat l -> PatField l
forall l. l -> QName l -> Pat l -> PatField l
PFieldPat l
pLoc QName l
n (l -> Name l -> Pat l
forall l. l -> Name l -> Pat l
PVar l
nLoc Name l
n')
    go PatField l
p = PatField l
p

{-------------------------------------------------------------------------------
  Naming
-------------------------------------------------------------------------------}

-- | HSE generated names are always dynamically bound
fromHseName :: TH.Name -> N.Name flavour 'N.Dynamic
fromHseName :: Name -> Name flavour 'Dynamic
fromHseName = Name -> Name flavour 'Dynamic
forall (ns :: NameSpace) (flavour :: Flavour).
IsFlavour flavour =>
Name -> Name ns flavour
N.fromTH'

-- | Resolve HSE generated name
--
-- As mentioned in 'fromHseName', HSE generated names are always dynamically
-- bound, and we therefore need to do a "renaming pass": we need to resolve the
-- name. However, the exact name we want to lookup might not be the name as it
-- appears in the QQ place; for if the user writes @MkR@, the name we actually
-- want to look up might be, say, @LR__MkR@.
resolveHseName :: (Quasi m, N.LookupName ns')
  => (String -> String)
  ->           N.Name ns  'N.Dynamic
  -> m (Maybe (N.Name ns' 'N.Global))
resolveHseName :: (String -> String)
-> Name ns 'Dynamic -> m (Maybe (Name ns' 'Global))
resolveHseName String -> String
f = Name ns' 'Dynamic -> m (Maybe (Name ns' 'Global))
forall (ns :: NameSpace) (m :: Type -> Type).
(LookupName ns, Quasi m) =>
Name ns 'Dynamic -> m (Maybe (Name ns 'Global))
N.lookupName (Name ns' 'Dynamic -> m (Maybe (Name ns' 'Global)))
-> (Name ns 'Dynamic -> Name ns' 'Dynamic)
-> Name ns 'Dynamic
-> m (Maybe (Name ns' 'Global))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (String -> String) -> Name ns 'Dynamic -> Name ns' 'Dynamic
forall (ns :: NameSpace) (flavour :: Flavour) (ns' :: NameSpace).
(String -> String) -> Name ns flavour -> Name ns' flavour
N.mapNameBase String -> String
f

-- | Variation on 'resolveHseName' that fails if the name is not known
resolveKnownHseName :: (Quasi m, N.LookupName ns')
  => (String -> String)
  ->    N.Name ns  'N.Dynamic
  -> m (N.Name ns' 'N.Global)
resolveKnownHseName :: (String -> String) -> Name ns 'Dynamic -> m (Name ns' 'Global)
resolveKnownHseName String -> String
f Name ns 'Dynamic
n = do
    Maybe (Name ns' 'Global)
mn' <- (String -> String)
-> Name ns 'Dynamic -> m (Maybe (Name ns' 'Global))
forall (m :: Type -> Type) (ns' :: NameSpace) (ns :: NameSpace).
(Quasi m, LookupName ns') =>
(String -> String)
-> Name ns 'Dynamic -> m (Maybe (Name ns' 'Global))
resolveHseName String -> String
f Name ns 'Dynamic
n
    case Maybe (Name ns' 'Global)
mn' of
      Just Name ns' 'Global
n' -> Name ns' 'Global -> m (Name ns' 'Global)
forall (m :: Type -> Type) a. Monad m => a -> m a
return Name ns' 'Global
n'
      Maybe (Name ns' 'Global)
Nothing -> String -> m (Name ns' 'Global)
forall (m :: Type -> Type) a. MonadFail m => String -> m a
fail (String -> m (Name ns' 'Global)) -> String -> m (Name ns' 'Global)
forall a b. (a -> b) -> a -> b
$ String
"resolveKnownHseName: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Name ns 'Dynamic -> String
forall (ns :: NameSpace) (flavour :: Flavour).
Name ns flavour -> String
N.nameBase Name ns 'Dynamic
n String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
" not in scope"