{-|
Copyright   :  (C) 2019, Myrtle Software Ltd,
                   2021, QBayLogic B.V.
                   2022, Google Inc
License     :  BSD2 (see the file LICENSE)
Maintainer  :  QBayLogic B.V. <devops@qbaylogic.com>

Template Haskell utilities for "Clash.Core.TermLiteral".
-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE TemplateHaskell #-}

module Clash.Core.TermLiteral.TH
  ( deriveTermToData
  , deriveShowsTypePrec
  , deriveTermLiteral
     -- Stop exporting @dcName'@  once `ghcide` stops type-checking expanded
     -- TH splices
  ,  dcName'
  ) where

import           Data.Either
import qualified Data.Text                       as Text
import           Data.List                       (intersperse)
import qualified Data.List.NonEmpty              as NE
import           Data.Proxy
import           Data.Maybe                      (isNothing)
import           Language.Haskell.TH.Syntax
import           Language.Haskell.TH.Lib         hiding (match)

import           Clash.Core.DataCon
import           Clash.Core.Term                 (collectArgs, Term(Data))
import           Clash.Core.Name                 (nameOcc)

-- Workaround for a strange GHC bug, where it complains about Subst only
-- existing as a boot file:
--
-- module Clash.Core.Subst cannot be linked; it is only available as a boot module
import Clash.Core.Subst ()

#if __GLASGOW_HASKELL__ >= 900
type CompatTyVarBndr = TyVarBndr ()
#else
type CompatTyVarBndr = TyVarBndr
#endif

dcName' :: DataCon -> String
dcName' :: DataCon -> String
dcName' = Text -> String
Text.unpack (Text -> String) -> (DataCon -> Text) -> DataCon -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Name DataCon -> Text
forall a. Name a -> Text
nameOcc (Name DataCon -> Text)
-> (DataCon -> Name DataCon) -> DataCon -> Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. DataCon -> Name DataCon
dcName

termToDataName :: Name
termToDataName :: Name
termToDataName =
  -- Note that we can't use a fully qualified name here: GHC disallows fully
  -- qualified names in instance function declarations.
  String -> Name
mkName String
"termToData"

showsTypePrecName :: Name
showsTypePrecName :: Name
showsTypePrecName =
  -- Note that we can't use a fully qualified name here: GHC disallows fully
  -- qualified names in instance function declarations.
  String -> Name
mkName String
"showsTypePrec"

termLiteralName :: Name
termLiteralName :: Name
termLiteralName = String -> Name
mkName String
"Clash.Core.TermLiteral.TermLiteral"

-- | Extracts variable names from a 'TyVarBndr'.
typeVarName :: CompatTyVarBndr -> Q (Name, Maybe Type)
typeVarName :: CompatTyVarBndr -> Q (Name, Maybe Type)
typeVarName = \case
#if __GLASGOW_HASKELL__ >= 900
  PlainTV typVarName ()        -> pure (typVarName, Nothing)
  KindedTV typVarName () StarT -> pure (typVarName, Nothing)
  KindedTV typVarName () kind  -> pure (typVarName, Just kind)
#else
  PlainTV Name
typVarName        -> (Name, Maybe Type) -> Q (Name, Maybe Type)
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (Name
typVarName, Maybe Type
forall a. Maybe a
Nothing)
  KindedTV Name
typVarName Type
StarT -> (Name, Maybe Type) -> Q (Name, Maybe Type)
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (Name
typVarName, Maybe Type
forall a. Maybe a
Nothing)
  KindedTV Name
typVarName Type
kind  -> (Name, Maybe Type) -> Q (Name, Maybe Type)
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (Name
typVarName, Type -> Maybe Type
forall a. a -> Maybe a
Just Type
kind)
#endif

-- | Derive a t'Clash.Core.TermLiteral.TermLiteral' instance for given type
deriveTermLiteral :: Name -> Q [Dec]
deriveTermLiteral :: Name -> Q [Dec]
deriveTermLiteral Name
typName = do
  TyConI (DataD Cxt
_ Name
_ [CompatTyVarBndr]
typeVars Maybe Type
_ [Con]
_ [DerivClause]
_) <- Name -> Q Info
reify Name
typName
#if MIN_VERSION_template_haskell(2,21,0)
  typeVarNames <- mapM (typeVarName . fmap (const ())) typeVars
#else
  [(Name, Maybe Type)]
typeVarNames <- (CompatTyVarBndr -> Q (Name, Maybe Type))
-> [CompatTyVarBndr] -> Q [(Name, Maybe Type)]
forall (t :: Type -> Type) (m :: Type -> Type) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM CompatTyVarBndr -> Q (Name, Maybe Type)
typeVarName [CompatTyVarBndr]
typeVars
#endif
  Dec
showsTypePrec <- Name -> Q Dec
deriveShowsTypePrec Name
typName
  Exp
termToDataBody <- Name -> Q Exp
deriveTermToData Name
typName
  let
    termToData :: Dec
termToData = Name -> [Clause] -> Dec
FunD Name
termToDataName [[Pat] -> Body -> [Dec] -> Clause
Clause [] (Exp -> Body
NormalB Exp
termToDataBody) []]
    innerInstanceType :: Type
innerInstanceType = (Type -> Type -> Type) -> Type -> Cxt -> Type
forall (t :: Type -> Type) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl Type -> Type -> Type
AppT (Name -> Type
ConT Name
typName) (((Name, Maybe Type) -> Type) -> [(Name, Maybe Type)] -> Cxt
forall a b. (a -> b) -> [a] -> [b]
map (Name -> Type
VarT (Name -> Type)
-> ((Name, Maybe Type) -> Name) -> (Name, Maybe Type) -> Type
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Name, Maybe Type) -> Name
forall a b. (a, b) -> a
fst) [(Name, Maybe Type)]
typeVarNames)
    instanceType :: Type
instanceType = Name -> Type
ConT Name
termLiteralName Type -> Type -> Type
`AppT` Type
innerInstanceType
    constraint :: Name -> TypeQ
constraint Name
typVarName = [t| $(conT termLiteralName) $(varT typVarName) |]
  Cxt
constraints <- ((Name, Maybe Type) -> TypeQ) -> [(Name, Maybe Type)] -> Q Cxt
forall (t :: Type -> Type) (m :: Type -> Type) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (Name -> TypeQ
constraint (Name -> TypeQ)
-> ((Name, Maybe Type) -> Name) -> (Name, Maybe Type) -> TypeQ
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Name, Maybe Type) -> Name
forall a b. (a, b) -> a
fst) (((Name, Maybe Type) -> Bool)
-> [(Name, Maybe Type)] -> [(Name, Maybe Type)]
forall a. (a -> Bool) -> [a] -> [a]
filter (Maybe Type -> Bool
forall a. Maybe a -> Bool
isNothing (Maybe Type -> Bool)
-> ((Name, Maybe Type) -> Maybe Type) -> (Name, Maybe Type) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Name, Maybe Type) -> Maybe Type
forall a b. (a, b) -> b
snd) [(Name, Maybe Type)]
typeVarNames)
  [Dec] -> Q [Dec]
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure ([Dec] -> Q [Dec]) -> [Dec] -> Q [Dec]
forall a b. (a -> b) -> a -> b
$ [Maybe Overlap -> Cxt -> Type -> [Dec] -> Dec
InstanceD Maybe Overlap
forall a. Maybe a
Nothing Cxt
constraints Type
instanceType [Dec
showsTypePrec, Dec
termToData]]

-- | For 'Maybe', constructs:
--
-- > showsTypePrec n _
-- >   = let
-- >       showSpace = showChar ' '
-- >       precCalls = [showsTypePrec 11 (Proxy @a)]
-- >       interspersedPrecCalls = intersperse showSpace precCalls
-- >       showType = foldl (.) (showString "Maybe") (showSpace : interspersedPrecCalls)
-- >     in
-- >       showParen (n > 10) showType
--
deriveShowsTypePrec :: Name -> Q Dec
deriveShowsTypePrec :: Name -> Q Dec
deriveShowsTypePrec Name
typName = do
  TyConI (DataD Cxt
_ Name
_ [CompatTyVarBndr]
typeVars Maybe Type
_ [Con]
_ [DerivClause]
_) <- Name -> Q Info
reify Name
typName
#if MIN_VERSION_template_haskell(2,21,0)
  typeVarNames <- mapM (typeVarName . fmap (const ())) typeVars
#else
  [(Name, Maybe Type)]
typeVarNames <- (CompatTyVarBndr -> Q (Name, Maybe Type))
-> [CompatTyVarBndr] -> Q [(Name, Maybe Type)]
forall (t :: Type -> Type) (m :: Type -> Type) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM CompatTyVarBndr -> Q (Name, Maybe Type)
typeVarName [CompatTyVarBndr]
typeVars
#endif
  Exp
showTypeBody <- [(Name, Maybe Type)] -> Q Exp
mkShowTypeBody [(Name, Maybe Type)]
typeVarNames
  Dec -> Q Dec
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (Name -> [Clause] -> Dec
FunD Name
showsTypePrecName [[Pat] -> Body -> [Dec] -> Clause
Clause [Name -> Pat
VarP Name
nName, Pat
WildP] (Exp -> Body
NormalB Exp
showTypeBody) []])
 where
  showTypeName :: Q Exp
showTypeName = [| showString $(litE (StringL (nameBase typName))) |]

  -- Constructs:
  --
  -- > showsTypePrec 11 (Proxy @a)
  --
  -- where the 'a' is given as an argument. The surrounding operator precedence
  -- is set to indicate "function" application. I.e., it instructs the call to
  -- wrap the type string in parentheses.
  --
  mkTypePrecCall :: (Name, Maybe a) -> Q Exp
mkTypePrecCall = \case
    (Name
typVarName, Maybe a
Nothing) ->
      [| $(varE showsTypePrecName) 11 (Proxy @($(varT typVarName))) |]
    (Name
_, Just a
_) ->
      -- XXX: Not sure how to deal with non-Type type variables so we do the dumb
      --      thing and insert an underscore.
      [| showString "_" |]

  -- Constructs:
  --
  -- > showString "Maybe" . showChar ' ' . showsTypePrec 11 (Proxy @a)
  --
  -- This is wrapped in an if-statement wrapping the result in parentheses if the
  -- incoming prec is more than 10 (function application).
  --
  mkShowTypeBody :: [(Name, Maybe Type)] -> Q Exp
  mkShowTypeBody :: [(Name, Maybe Type)] -> Q Exp
mkShowTypeBody [(Name, Maybe Type)]
typeVarNames =
    case [(Name, Maybe Type)]
typeVarNames of
      [] ->
        -- We seq on `n` here to prevent _unused variable_ warnings. This is a
        -- bit of a hack (the real solution would be to selectively pattern
        -- match).
        [| $(varE nName) `seq` $(showTypeName) |]
      [(Name, Maybe Type)]
_  -> [|
        let
          showSpace = showChar ' '
          precCalls = $(listE (map mkTypePrecCall typeVarNames))
          interspersedPrecCalls = intersperse showSpace precCalls
          showType = foldl (.) $(showTypeName) (showSpace : interspersedPrecCalls)
        in
          showParen ($(varE nName) > 10) showType
       |]

  nName :: Name
nName = String -> Name
mkName String
"n"

deriveTermToData :: Name -> Q Exp
deriveTermToData :: Name -> Q Exp
deriveTermToData Name
typName = do
  TyConI (DataD Cxt
_ Name
_ [CompatTyVarBndr]
_ Maybe Type
_ [Con]
constrs [DerivClause]
_) <- Name -> Q Info
reify Name
typName
  Exp -> Q Exp
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure ([(Name, Int)] -> Exp
deriveTermToData1 ((Con -> (Name, Int)) -> [Con] -> [(Name, Int)]
forall a b. (a -> b) -> [a] -> [b]
map Con -> (Name, Int)
toConstr' [Con]
constrs))
 where
  toConstr' :: Con -> (Name, Int)
toConstr' (NormalC Name
cName [BangType]
fields) = (Name
cName, [BangType] -> Int
forall (t :: Type -> Type) a. Foldable t => t a -> Int
length [BangType]
fields)
  toConstr' (RecC Name
cName [VarBangType]
fields) = (Name
cName, [VarBangType] -> Int
forall (t :: Type -> Type) a. Foldable t => t a -> Int
length [VarBangType]
fields)
  toConstr' Con
c = String -> (Name, Int)
forall a. HasCallStack => String -> a
error (String -> (Name, Int)) -> String -> (Name, Int)
forall a b. (a -> b) -> a -> b
$ String
"Unexpected constructor: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Con -> String
forall a. Show a => a -> String
show Con
c

deriveTermToData1 :: [(Name, Int)] -> Exp
deriveTermToData1 :: [(Name, Int)] -> Exp
deriveTermToData1 [(Name, Int)]
constrs =
  [Match] -> Exp
LamCaseE
    [ Pat -> Body -> [Dec] -> Match
Match Pat
pat (Exp -> Body
NormalB (if [Dec] -> Bool
forall (t :: Type -> Type) a. Foldable t => t a -> Bool
null [Dec]
args then Exp
theCase else [Dec] -> Exp -> Exp
LetE [Dec]
args Exp
theCase)) []
    , Pat -> Body -> [Dec] -> Match
Match (Name -> Pat
VarP Name
termName) (Exp -> Body
NormalB ((Name -> Exp
ConE 'Left Exp -> Exp -> Exp
`AppE` Name -> Exp
VarE Name
termName))) []

    ]
 where
  nArgs :: Int
nArgs = [Int] -> Int
forall (t :: Type -> Type) a. (Foldable t, Ord a) => t a -> a
maximum (((Name, Int) -> Int) -> [(Name, Int)] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map (Name, Int) -> Int
forall a b. (a, b) -> b
snd [(Name, Int)]
constrs)

  args :: [Dec]
  args :: [Dec]
args = (Int -> Name -> Dec) -> [Int] -> [Name] -> [Dec]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\Int
n Name
nm -> Pat -> Body -> [Dec] -> Dec
ValD (Name -> Pat
VarP Name
nm) (Exp -> Body
NormalB (Integer -> Exp
arg (Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
n))) []) [Int
0..Int
nArgsInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1] (NonEmpty Name -> [Name]
forall a. NonEmpty a -> [a]
NE.toList NonEmpty Name
argNames)
  arg :: Integer -> Exp
arg Integer
n = Exp -> Exp -> Exp -> Exp
UInfixE (Name -> Exp
VarE Name
argsName) (Name -> Exp
VarE '(!!)) (Lit -> Exp
LitE (Integer -> Lit
IntegerL Integer
n))

  -- case nm of {"ConstrOne" -> ConstOne <$> termToData arg0; "ConstrTwo" -> ...}
  theCase :: Exp
  theCase :: Exp
theCase =
    Exp -> [Match] -> Exp
CaseE
      (Name -> Exp
VarE Name
nameName)
      (((Name, Int) -> Match) -> [(Name, Int)] -> [Match]
forall a b. (a -> b) -> [a] -> [b]
map (Name, Int) -> Match
match [(Name, Int)]
constrs [Match] -> [Match] -> [Match]
forall a. [a] -> [a] -> [a]
++ [Match
emptyMatch])

  emptyMatch :: Match
emptyMatch = Pat -> Body -> [Dec] -> Match
Match Pat
WildP (Exp -> Body
NormalB (Name -> Exp
ConE 'Left Exp -> Exp -> Exp
`AppE` Name -> Exp
VarE Name
termName)) []

  match :: (Name, Int) -> Match
  match :: (Name, Int) -> Match
match (Name
cName, Int
nFields) =
    Pat -> Body -> [Dec] -> Match
Match (Lit -> Pat
LitP (String -> Lit
StringL (Name -> String
forall a. Show a => a -> String
show Name
cName))) (Exp -> Body
NormalB (Name -> Int -> Exp
mkCall Name
cName Int
nFields)) []

  mkCall :: Name -> Int -> Exp
  mkCall :: Name -> Int -> Exp
mkCall Name
cName Int
0  = Name -> Exp
ConE 'Right Exp -> Exp -> Exp
`AppE` Name -> Exp
ConE Name
cName
  mkCall Name
cName Int
1 =
    Exp -> Exp -> Exp -> Exp
UInfixE
      (Name -> Exp
ConE Name
cName)
      (Name -> Exp
VarE '(<$>))
      (Name -> Exp
VarE Name
termToDataName Exp -> Exp -> Exp
`AppE` Name -> Exp
VarE (NonEmpty Name -> Name
forall a. NonEmpty a -> a
NE.head NonEmpty Name
argNames))
  mkCall Name
cName Int
nFields =
    (Exp -> Name -> Exp) -> Exp -> [Name] -> Exp
forall (t :: Type -> Type) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl
      (\Exp
e Name
aName ->
        Exp -> Exp -> Exp -> Exp
UInfixE
          Exp
e
          (Name -> Exp
VarE '(<*>))
          (Name -> Exp
VarE Name
termToDataName Exp -> Exp -> Exp
`AppE` Name -> Exp
VarE Name
aName))
      (Name -> Int -> Exp
mkCall Name
cName Int
1)
      (Int -> [Name] -> [Name]
forall a. Int -> [a] -> [a]
take (Int
nFieldsInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) (NonEmpty Name -> [Name]
forall a. NonEmpty a -> [a]
NE.tail NonEmpty Name
argNames))

  -- term@(collectArgs -> (Data (dcName' -> nm), args))
  pat :: Pat
  pat :: Pat
pat =
    Name -> Pat -> Pat
AsP
      Name
termName
      (Exp -> Pat -> Pat
ViewP
        (Name -> Exp
VarE 'collectArgs)
#if MIN_VERSION_template_haskell(2,18,0)
        (TupP [ ConP 'Data [] [ViewP (VarE 'dcName') (VarP nameName)]
#else
        ([Pat] -> Pat
TupP [ Name -> [Pat] -> Pat
ConP 'Data [Exp -> Pat -> Pat
ViewP (Name -> Exp
VarE 'dcName') (Name -> Pat
VarP Name
nameName)]
#endif
              , Exp -> Pat -> Pat
ViewP
                 (Name -> Exp
VarE 'lefts)
                 (if Int
nArgs Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 then Pat
WildP else Name -> Pat
VarP Name
argsName)]))

  termName :: Name
termName = String -> Name
mkName String
"term"
  argsName :: Name
argsName = String -> Name
mkName String
"args"
  argNames :: NonEmpty Name
argNames = (Word -> Name) -> NonEmpty Word -> NonEmpty Name
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap (String -> Name
mkName (String -> Name) -> (Word -> String) -> Word -> Name
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (String
"arg" String -> String -> String
forall a. Semigroup a => a -> a -> a
<>) (String -> String) -> (Word -> String) -> Word -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Word -> String
forall a. Show a => a -> String
show) ((Word -> Word) -> Word -> NonEmpty Word
forall a. (a -> a) -> a -> NonEmpty a
NE.iterate (Word -> Word -> Word
forall a. Num a => a -> a -> a
+Word
1) (Word
0 :: Word))
  nameName :: Name
nameName = String -> Name
mkName String
"nm"