{-# Language TemplateHaskell #-}
module Text.LLVM.Labels.TH (generateRelabel) where

import Control.Monad (zipWithM)
import Language.Haskell.TH
import Language.Haskell.TH.Datatype

generateRelabel :: Name -> Name -> ExpQ
generateRelabel :: Name -> Name -> ExpQ
generateRelabel Name
relabel Name
dataCon =
  do DatatypeInfo
di <- Name -> Q DatatypeInfo
reifyDatatype Name
dataCon
     DatatypeInfo -> ExpQ -> ExpQ
generateRelabelData DatatypeInfo
di (Name -> ExpQ
forall (m :: * -> *). Quote m => Name -> m Exp
varE Name
relabel)

generateRelabelData :: DatatypeInfo -> ExpQ -> ExpQ
generateRelabelData :: DatatypeInfo -> ExpQ -> ExpQ
generateRelabelData DatatypeInfo
di ExpQ
relabelE =
  [| \f x -> $(ExpQ -> [Q Match] -> ExpQ
forall (m :: * -> *). Quote m => m Exp -> [m Match] -> m Exp
caseE [| x |] (ExpQ -> ConstructorInfo -> Q Match
mkMatch [| f |] (ConstructorInfo -> Q Match) -> [ConstructorInfo] -> [Q Match]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [ConstructorInfo]
cons)) |]
  where
    mkMatch :: ExpQ -> ConstructorInfo -> Q Match
mkMatch = Name -> ExpQ -> ExpQ -> ConstructorInfo -> Q Match
generateRelabelCon Name
lastArg ExpQ
relabelE
    lastArg :: Name
lastArg = TyVarBndr_ () -> Name
forall flag. TyVarBndr_ flag -> Name
tvName ([TyVarBndr_ ()] -> TyVarBndr_ ()
forall a. HasCallStack => [a] -> a
last (DatatypeInfo -> [TyVarBndr_ ()]
datatypeVars DatatypeInfo
di))
    cons :: [ConstructorInfo]
cons    = DatatypeInfo -> [ConstructorInfo]
datatypeCons DatatypeInfo
di

-- | Generates the case arm for the given constructor that
-- relabels values using this constructor given a relabeling
-- function.
generateRelabelCon ::
  Name            {- ^ last type parameter            -} ->
  ExpQ            {- ^ recusive relabel expression    -} ->
  ExpQ            {- ^ function expression            -} ->
  ConstructorInfo {- ^ current constructor            -} ->
  MatchQ          {- ^ match arm for this constructor -}
generateRelabelCon :: Name -> ExpQ -> ExpQ -> ConstructorInfo -> Q Match
generateRelabelCon Name
lastArg ExpQ
relabelE ExpQ
fE ConstructorInfo
ci =
  do [(Name, Type)]
names <- String -> [Type] -> Q [(Name, Type)]
forall a. String -> [a] -> Q [(Name, a)]
nameThings String
"x" (ConstructorInfo -> [Type]
constructorFields ConstructorInfo
ci)
     Q Pat -> Q Body -> [Q Dec] -> Q Match
forall (m :: * -> *).
Quote m =>
m Pat -> m Body -> [m Dec] -> m Match
match
      (Name -> [Q Pat] -> Q Pat
forall (m :: * -> *). Quote m => Name -> [m Pat] -> m Pat
conP Name
cn (((Name, Type) -> Q Pat) -> [(Name, Type)] -> [Q Pat]
forall a b. (a -> b) -> [a] -> [b]
map (Name -> Q Pat
forall (m :: * -> *). Quote m => Name -> m Pat
varP (Name -> Q Pat) -> ((Name, Type) -> Name) -> (Name, Type) -> Q Pat
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Name, Type) -> Name
forall a b. (a, b) -> a
fst) [(Name, Type)]
names))
      (ExpQ -> Q Body
forall (m :: * -> *). Quote m => m Exp -> m Body
normalB (Name -> [Either ExpQ ExpQ] -> ExpQ
bodyExp Name
cn (((Name, Type) -> Either ExpQ ExpQ)
-> [(Name, Type)] -> [Either ExpQ ExpQ]
forall a b. (a -> b) -> [a] -> [b]
map (Name, Type) -> Either ExpQ ExpQ
gen [(Name, Type)]
names)))
      []
  where
    cn :: Name
cn = ConstructorInfo -> Name
constructorName ConstructorInfo
ci

    -- Give a field name and type returns:
    -- Left for a pure field
    -- Right for a field using the Applicative instance
    gen :: (Name, Type) -> Either ExpQ ExpQ
    gen :: (Name, Type) -> Either ExpQ ExpQ
gen (Name
n,Type
t) =
      let nE :: ExpQ
nE = Name -> ExpQ
forall (m :: * -> *). Quote m => Name -> m Exp
varE Name
n in
      case Name -> ExpQ -> ExpQ -> Type -> Maybe ExpQ
generateRelabelField Name
lastArg ExpQ
fE ExpQ
relabelE Type
t of
        Just ExpQ
f  -> ExpQ -> Either ExpQ ExpQ
forall a b. b -> Either a b
Right [| $ExpQ
f $ExpQ
nE |]
        Maybe ExpQ
Nothing -> ExpQ -> Either ExpQ ExpQ
forall a b. a -> Either a b
Left ExpQ
nE

-- | Given a constructor and a list of pure and updated fields,
-- build syntax that rebuilds the expression.
bodyExp ::
  Name               {- ^ constructor                         -} ->
  [Either ExpQ ExpQ] {- ^ list of pure and applicative fields -} ->
  ExpQ               {- ^ applicative result                  -}
bodyExp :: Name -> [Either ExpQ ExpQ] -> ExpQ
bodyExp Name
conname [Either ExpQ ExpQ]
fields = ExpQ -> [ExpQ] -> ExpQ
liftAE ExpQ
conLike [ExpQ]
updates
  where
    updates :: [ExpQ]
updates = [ExpQ
r | Right ExpQ
r <- [Either ExpQ ExpQ]
fields]

    -- Builds a value suitable to be the argument to liftAE that can
    -- combine all of the updated field values
    conLike :: ExpQ
conLike =
      do [Name]
names <- ((Name, ExpQ) -> Name) -> [(Name, ExpQ)] -> [Name]
forall a b. (a -> b) -> [a] -> [b]
map (Name, ExpQ) -> Name
forall a b. (a, b) -> a
fst ([(Name, ExpQ)] -> [Name]) -> Q [(Name, ExpQ)] -> Q [Name]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> String -> [ExpQ] -> Q [(Name, ExpQ)]
forall a. String -> [a] -> Q [(Name, a)]
nameThings String
"y" [ExpQ]
updates
         [Q Pat] -> ExpQ -> ExpQ
forall (m :: * -> *). Quote m => [m Pat] -> m Exp -> m Exp
lamE
           ((Name -> Q Pat) -> [Name] -> [Q Pat]
forall a b. (a -> b) -> [a] -> [b]
map Name -> Q Pat
forall (m :: * -> *). Quote m => Name -> m Pat
varP [Name]
names)
           ([ExpQ] -> ExpQ
forall (m :: * -> *). Quote m => [m Exp] -> m Exp
appsE (Name -> ExpQ
forall (m :: * -> *). Quote m => Name -> m Exp
conE Name
conname ExpQ -> [ExpQ] -> [ExpQ]
forall a. a -> [a] -> [a]
: [ExpQ] -> [Either ExpQ ExpQ] -> [ExpQ]
forall a b. [a] -> [Either a b] -> [a]
replaceRights ((Name -> ExpQ) -> [Name] -> [ExpQ]
forall a b. (a -> b) -> [a] -> [b]
map Name -> ExpQ
forall (m :: * -> *). Quote m => Name -> m Exp
varE [Name]
names) [Either ExpQ ExpQ]
fields))

-- | Replaces all of the 'Right' values in the given list with elements
-- from the first list. The number of replacements must exactly match
-- the number of 'Right' values.
replaceRights ::
  [a]          {- ^ replacements  -} ->
  [Either a b] {- ^ source list   -} ->
  [a]          {- ^ replaced list -}
replaceRights :: forall a b. [a] -> [Either a b] -> [a]
replaceRights [a]
xs     (Left a
y  : [Either a b]
ys) = a
y a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a] -> [Either a b] -> [a]
forall a b. [a] -> [Either a b] -> [a]
replaceRights [a]
xs [Either a b]
ys
replaceRights (a
x:[a]
xs) (Right b
_ : [Either a b]
ys) = a
x a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a] -> [Either a b] -> [a]
forall a b. [a] -> [Either a b] -> [a]
replaceRights [a]
xs [Either a b]
ys
replaceRights []     []             = []
replaceRights [a]
_      [Either a b]
_              = String -> [a]
forall a. HasCallStack => String -> a
error String
"Text.LLVM.Labels.TH.replaceRights: PANIC"

-- | Generate the applicative update value for a field if it
-- has an appropriate type otherwise return nothing if it
-- should be left unchagned.
generateRelabelField ::
  Name       {- ^ last type parameter         -} ->
  ExpQ       {- ^ function expression         -} ->
  ExpQ       {- ^ relabel expression          -} ->
  Type       {- ^ field type                  -} ->
  Maybe ExpQ {- ^ applicative update function -}
generateRelabelField :: Name -> ExpQ -> ExpQ -> Type -> Maybe ExpQ
generateRelabelField Name
lastArg ExpQ
fE ExpQ
relabelE Type
t =
  case Type -> (Int, Type)
typeDepth Type
t of
    (Int
n, VarT Name
tn) | Name
tn Name -> Name -> Bool
forall a. Eq a => a -> a -> Bool
== Name
lastArg -> ExpQ -> Maybe ExpQ
forall a. a -> Maybe a
Just ([ExpQ]
exprs [ExpQ] -> Int -> ExpQ
forall a. HasCallStack => [a] -> Int -> a
!! Int
n)
    (Int, Type)
_                            -> Maybe ExpQ
forall a. Maybe a
Nothing
  where
    exprs :: [ExpQ]
exprs = [| $ExpQ
fE Nothing |] ExpQ -> [ExpQ] -> [ExpQ]
forall a. a -> [a] -> [a]
: (ExpQ -> ExpQ) -> ExpQ -> [ExpQ]
forall a. (a -> a) -> a -> [a]
iterate ExpQ -> ExpQ
traverseE [| $ExpQ
relabelE $ExpQ
fE |]

-- | Figure out the depth of the outer type applications and
-- return the type at the bottom of the stack
typeDepth ::
  Type        {- ^ target type                                     -} ->
  (Int, Type) {- ^ number of type applications and right-most type -}
typeDepth :: Type -> (Int, Type)
typeDepth = Int -> Type -> (Int, Type)
forall {a}. Num a => a -> Type -> (a, Type)
go Int
0
  where
    go :: a -> Type -> (a, Type)
go a
i (AppT Type
_ Type
x) = a -> Type -> (a, Type)
go (a
ia -> a -> a
forall a. Num a => a -> a -> a
+a
1) Type
x
    go a
i Type
t          = (a
i, Type
t)

-- | Associate each element in a list of things with a unique name
-- derived from a given name stem.
nameThings ::
  String        {- ^ base name                       -} ->
  [a]           {- ^ things to name                  -} ->
  Q [(Name, a)] {- ^ things paired with unique names -}
nameThings :: forall a. String -> [a] -> Q [(Name, a)]
nameThings String
base [a]
xs = (Int -> a -> Q (Name, a)) -> [Int] -> [a] -> Q [(Name, a)]
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM Int -> a -> Q (Name, a)
forall {m :: * -> *} {a} {b}.
(Quote m, Show a) =>
a -> b -> m (Name, b)
nameThing [Int
0 :: Int ..] [a]
xs
  where
    nameThing :: a -> b -> m (Name, b)
nameThing a
i b
x = do Name
n <- String -> m Name
forall (m :: * -> *). Quote m => String -> m Name
newName (String
base String -> String -> String
forall a. [a] -> [a] -> [a]
++ a -> String
forall a. Show a => a -> String
show a
i); (Name, b) -> m (Name, b)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Name
n,b
x)

-- | Apply 'traverse' to an expression
traverseE ::
  ExpQ {- ^ f          -} ->
  ExpQ {- ^ traverse f -}
traverseE :: ExpQ -> ExpQ
traverseE ExpQ
e = [| traverse $ExpQ
e |]

-- Applies a pure value to zero or more applicative things to be combined
-- with (<$>) and (<*>)
liftAE :: ExpQ -> [ExpQ] -> ExpQ
liftAE :: ExpQ -> [ExpQ] -> ExpQ
liftAE ExpQ
c []     = [| pure $ExpQ
c |]
liftAE ExpQ
c (ExpQ
x:[ExpQ]
xs) = (ExpQ -> ExpQ -> ExpQ) -> ExpQ -> [ExpQ] -> ExpQ
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl (\ExpQ
f ExpQ
e -> [| $ExpQ
f <*> $ExpQ
e |]) [| $ExpQ
c <$> $ExpQ
x |] [ExpQ]
xs