{-|
Module     : Lang.Crucible.Utils.Structural
Copyright  : (c) Galois, Inc 2013-2016
License    : BSD3
Maintainer : Joe Hendrix <jhendrix@galois.com>

This module declares template Haskell primitives so that it is easier
to work with GADTs that have many constructors.
-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE TypeOperators #-}
module Lang.Crucible.Utils.Structural
  ( structuralPretty
  ) where

import Data.Char (toLower)
import Language.Haskell.TH
import Language.Haskell.TH.Datatype
import Prettyprinter (brackets)

import Data.Parameterized.TH.GADT
import Data.Parameterized.TraversableFC

import Lang.Crucible.Utils.PrettyPrint (ppFn, commas)

------------------------------------------------------------------------
-- Contructor cases

-- | @structuralPretty tp@ generates a function with the type
--   @forall f ann. (forall x. f x -> Doc ann) -> (forall x. tp f x -> Doc ann)@
--   suitable for instantiating the @PrettyApp@ class.
structuralPretty :: TypeQ -> [(TypePat, ExpQ)] -> ExpQ
structuralPretty :: TypeQ -> [(TypePat, ExpQ)] -> ExpQ
structuralPretty TypeQ
tpq [(TypePat, ExpQ)]
pats0 = do
  DatatypeInfo
d <- Name -> Q DatatypeInfo
lookupDataType' (Name -> Q DatatypeInfo) -> Q Name -> Q DatatypeInfo
forall (m :: Type -> Type) a b. Monad m => (a -> m b) -> m a -> m b
=<< String -> Type -> Q Name
asTypeCon String
"structuralPretty" (Type -> Q Name) -> TypeQ -> Q Name
forall (m :: Type -> Type) a b. Monad m => (a -> m b) -> m a -> m b
=<< TypeQ
tpq
  Name
pp <- String -> Q Name
forall (m :: Type -> Type). Quote m => String -> m Name
newName String
"pp"
  Name
a <- String -> Q Name
forall (m :: Type -> Type). Quote m => String -> m Name
newName String
"a"

  let pats :: Type -> Q (Maybe ExpQ)
pats = [Type] -> [(TypePat, ExpQ)] -> Type -> Q (Maybe ExpQ)
forall v. [Type] -> [(TypePat, v)] -> Type -> Q (Maybe v)
assocTypePats (DatatypeInfo -> [Type]
dataParamTypes DatatypeInfo
d) [(TypePat, ExpQ)]
pats0
  [Q Pat] -> ExpQ -> ExpQ
forall (m :: Type -> Type). Quote m => [m Pat] -> m Exp -> m Exp
lamE [Name -> Q Pat
forall (m :: Type -> Type). Quote m => Name -> m Pat
varP Name
pp, Name -> Q Pat
forall (m :: Type -> Type). Quote m => Name -> m Pat
varP Name
a] (ExpQ -> ExpQ) -> ExpQ -> ExpQ
forall a b. (a -> b) -> a -> b
$
      ExpQ -> [Q Match] -> ExpQ
forall (m :: Type -> Type). Quote m => m Exp -> [m Match] -> m Exp
caseE (Name -> ExpQ
forall (m :: Type -> Type). Quote m => Name -> m Exp
varE Name
a) ((Type -> Q (Maybe ExpQ)) -> ExpQ -> ConstructorInfo -> Q Match
matchPretty Type -> Q (Maybe ExpQ)
pats (Name -> ExpQ
forall (m :: Type -> Type). Quote m => Name -> m Exp
varE Name
pp) (ConstructorInfo -> Q Match) -> [ConstructorInfo] -> [Q Match]
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> DatatypeInfo -> [ConstructorInfo]
datatypeCons DatatypeInfo
d)

matchPretty :: (Type -> Q (Maybe ExpQ))  -- ^ Pattern match functions
            -> ExpQ
            -> ConstructorInfo
            -> MatchQ
matchPretty :: (Type -> Q (Maybe ExpQ)) -> ExpQ -> ConstructorInfo -> Q Match
matchPretty Type -> Q (Maybe ExpQ)
matchPat ExpQ
pp ConstructorInfo
con = do
  let nm :: Name
nm  = ConstructorInfo -> Name
constructorName ConstructorInfo
con
      tps :: [Type]
tps = ConstructorInfo -> [Type]
constructorFields ConstructorInfo
con
  (Pat
pat,[Name]
nms) <- ConstructorInfo -> String -> Q (Pat, [Name])
conPat ConstructorInfo
con String
"x"
  let vars :: [ExpQ]
vars = Name -> ExpQ
forall (m :: Type -> Type). Quote m => Name -> m Exp
varE (Name -> ExpQ) -> [Name] -> [ExpQ]
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> [Name]
nms
  let nm' :: String
nm' = case Name -> String
nameBase Name
nm of
              Char
c:String
r -> Char -> Char
toLower Char
c Char -> String -> String
forall a. a -> [a] -> [a]
: String
r
              [] -> String -> String
forall a. HasCallStack => String -> a
error String
"matchPretty given constructor with empty name."
  let mkPP0 :: ExpQ -> Type -> ExpQ
mkPP0 ExpQ
v Type
tp = do
        Maybe ExpQ
me <- Type -> Q (Maybe ExpQ)
matchPat Type
tp
        case Maybe ExpQ
me of
          Maybe ExpQ
Nothing -> ExpQ -> Type -> ExpQ
mkPP ExpQ
v Type
tp
          Just ExpQ
f -> [| $(ExpQ
f) $(ExpQ
pp) $(ExpQ
v)|]
      mkPP :: ExpQ -> Type -> ExpQ
mkPP ExpQ
v ConT{} = [| viaShow $(ExpQ
v) |]
      mkPP ExpQ
v (AppT VarT{} Type
_) = ExpQ -> ExpQ -> ExpQ
forall (m :: Type -> Type). Quote m => m Exp -> m Exp -> m Exp
appE ExpQ
pp ExpQ
v
      mkPP ExpQ
v (AppT (ConT Name
cnm) Type
_)
       | Name -> String
nameBase Name
cnm String -> [String] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: Type -> Type) a.
(Foldable t, Eq a) =>
a -> t a -> Bool
`elem` [ String
"Vector" ]
       = [| brackets (commas (fmap $(ExpQ
pp) $(ExpQ
v))) |]
      mkPP ExpQ
v (AppT (AppT (ConT Name
cnm) Type
_) Type
_)
       | Name -> String
nameBase Name
cnm String -> [String] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: Type -> Type) a.
(Foldable t, Eq a) =>
a -> t a -> Bool
`elem` [ String
"Assignment" ]
       = [| brackets (commas (toListFC $(ExpQ
pp) $(ExpQ
v))) |]
      mkPP ExpQ
v Type
_ = [| viaShow $(ExpQ
v) |]
      --mkPP _ tp = error $ "Unsupported type " ++ show tp ++ " with " ++ nameBase nm
  let rhs :: ExpQ
rhs = [| ppFn $(Lit -> ExpQ
forall (m :: Type -> Type). Quote m => Lit -> m Exp
litE (String -> Lit
stringL String
nm')) $([ExpQ] -> ExpQ
forall (m :: Type -> Type). Quote m => [m Exp] -> m Exp
listE ((ExpQ -> Type -> ExpQ) -> [ExpQ] -> [Type] -> [ExpQ]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith ExpQ -> Type -> ExpQ
mkPP0 [ExpQ]
vars [Type]
tps)) |]
  Q Pat -> Q Body -> [Q Dec] -> Q Match
forall (m :: Type -> Type).
Quote m =>
m Pat -> m Body -> [m Dec] -> m Match
match (Pat -> Q Pat
forall a. a -> Q a
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure Pat
pat) (ExpQ -> Q Body
forall (m :: Type -> Type). Quote m => m Exp -> m Body
normalB ExpQ
rhs) []