{-# LANGUAGE CPP #-}
#ifdef TRUSTWORTHY
{-# LANGUAGE Trustworthy #-}
#endif

#include "lens-common.h"

-----------------------------------------------------------------------------
-- |
-- Module      :  Control.Lens.Internal.PrismTH
-- Copyright   :  (C) 2014-2016 Edward Kmett and Eric Mertens
-- License     :  BSD-style (see the file LICENSE)
-- Maintainer  :  Edward Kmett <ekmett@gmail.com>
-- Stability   :  experimental
-- Portability :  non-portable
--
-----------------------------------------------------------------------------

module Control.Lens.Internal.PrismTH
  ( makePrisms
  , makeClassyPrisms
  , makeDecPrisms
  ) where

import Control.Applicative
import Control.Lens.Getter
import Control.Lens.Internal.TH
import Control.Lens.Lens
import Control.Lens.Setter
import Control.Monad
import Data.Char (isUpper)
import qualified Data.List as List
import Data.Set.Lens
import Data.Traversable
import Language.Haskell.TH
import qualified Language.Haskell.TH.Datatype as D
import qualified Language.Haskell.TH.Datatype.TyVarBndr as D
import Language.Haskell.TH.Lens
import qualified Data.Map as Map
import qualified Data.Set as Set
import Data.Set (Set)
import Prelude

-- | Generate a 'Prism' for each constructor of a data type.
-- Isos generated when possible.
-- Reviews are created for constructors with existentially
-- quantified constructors and GADTs.
--
-- /e.g./
--
-- @
-- data FooBarBaz a
--   = Foo Int
--   | Bar a
--   | Baz Int Char
-- makePrisms ''FooBarBaz
-- @
--
-- will create
--
-- @
-- _Foo :: Prism' (FooBarBaz a) Int
-- _Bar :: Prism (FooBarBaz a) (FooBarBaz b) a b
-- _Baz :: Prism' (FooBarBaz a) (Int, Char)
-- @
makePrisms :: Name {- ^ Type constructor name -} -> DecsQ
makePrisms :: Name -> DecsQ
makePrisms = Bool -> Name -> DecsQ
makePrisms' Bool
True


-- | Generate a 'Prism' for each constructor of a data type
-- and combine them into a single class. No Isos are created.
-- Reviews are created for constructors with existentially
-- quantified constructors and GADTs.
--
-- /e.g./
--
-- @
-- data FooBarBaz a
--   = Foo Int
--   | Bar a
--   | Baz Int Char
-- makeClassyPrisms ''FooBarBaz
-- @
--
-- will create
--
-- @
-- class AsFooBarBaz s a | s -> a where
--   _FooBarBaz :: Prism' s (FooBarBaz a)
--   _Foo :: Prism' s Int
--   _Bar :: Prism' s a
--   _Baz :: Prism' s (Int,Char)
--
--   _Foo = _FooBarBaz . _Foo
--   _Bar = _FooBarBaz . _Bar
--   _Baz = _FooBarBaz . _Baz
--
-- instance AsFooBarBaz (FooBarBaz a) a
-- @
--
-- Generate an "As" class of prisms. Names are selected by prefixing the constructor
-- name with an underscore.  Constructors with multiple fields will
-- construct Prisms to tuples of those fields.
--
-- In the event that the name of a data type is also the name of one of its
-- constructors, the name of the 'Prism' generated for the data type will be
-- prefixed with an extra @_@ (if the data type name is prefix) or @.@ (if the
-- name is infix) to disambiguate it from the 'Prism' for the corresponding
-- constructor. For example, this code:
--
-- @
-- data Quux = Quux Int | Fred Bool
-- makeClassyPrisms ''Quux
-- @
--
-- will create:
--
-- @
-- class AsQuux s where
--   __Quux :: Prism' s Quux -- Data type prism
--   _Quux :: Prism' s Int   -- Constructor prism
--   _Fred :: Prism' s Bool
--
--   _Quux = __Quux . _Quux
--   _Fred = __Quux . _Fred
--
-- instance AsQuux Quux
-- @
makeClassyPrisms :: Name {- ^ Type constructor name -} -> DecsQ
makeClassyPrisms :: Name -> DecsQ
makeClassyPrisms = Bool -> Name -> DecsQ
makePrisms' Bool
False


-- | Main entry point into Prism generation for a given type constructor name.
makePrisms' :: Bool -> Name -> DecsQ
makePrisms' :: Bool -> Name -> DecsQ
makePrisms' Bool
normal Name
typeName =
  do DatatypeInfo
info <- Name -> Q DatatypeInfo
D.reifyDatatype Name
typeName
     let cls :: Maybe Name
cls | Bool
normal    = forall a. Maybe a
Nothing
             | Bool
otherwise = forall a. a -> Maybe a
Just (DatatypeInfo -> Name
D.datatypeName DatatypeInfo
info)
         cons :: [ConstructorInfo]
cons = DatatypeInfo -> [ConstructorInfo]
D.datatypeCons DatatypeInfo
info
     Type -> [NCon] -> Maybe Name -> DecsQ
makeConsPrisms (DatatypeInfo -> Type
datatypeTypeKinded DatatypeInfo
info) (forall a b. (a -> b) -> [a] -> [b]
map ConstructorInfo -> NCon
normalizeCon [ConstructorInfo]
cons) Maybe Name
cls


-- | Generate prisms for the given 'Dec'
makeDecPrisms :: Bool {- ^ generate top-level definitions -} -> Dec -> DecsQ
makeDecPrisms :: Bool -> Dec -> DecsQ
makeDecPrisms Bool
normal Dec
dec =
  do DatatypeInfo
info <- Dec -> Q DatatypeInfo
D.normalizeDec Dec
dec
     let cls :: Maybe Name
cls | Bool
normal    = forall a. Maybe a
Nothing
             | Bool
otherwise = forall a. a -> Maybe a
Just (DatatypeInfo -> Name
D.datatypeName DatatypeInfo
info)
         cons :: [ConstructorInfo]
cons = DatatypeInfo -> [ConstructorInfo]
D.datatypeCons DatatypeInfo
info
     Type -> [NCon] -> Maybe Name -> DecsQ
makeConsPrisms (DatatypeInfo -> Type
datatypeTypeKinded DatatypeInfo
info) (forall a b. (a -> b) -> [a] -> [b]
map ConstructorInfo -> NCon
normalizeCon [ConstructorInfo]
cons) Maybe Name
cls


-- | Generate prisms for the given type, normalized constructors, and
-- an optional name to be used for generating a prism class.
-- This function dispatches between Iso generation, normal top-level
-- prisms, and classy prisms.
makeConsPrisms :: Type -> [NCon] -> Maybe Name -> DecsQ

-- special case: single constructor, not classy -> make iso
makeConsPrisms :: Type -> [NCon] -> Maybe Name -> DecsQ
makeConsPrisms Type
t [con :: NCon
con@(NCon Name
_ [] [] [Type]
_)] Maybe Name
Nothing = Type -> NCon -> DecsQ
makeConIso Type
t NCon
con

-- top-level definitions
makeConsPrisms Type
t [NCon]
cons Maybe Name
Nothing =
  forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
t a -> (a -> f b) -> f (t b)
for [NCon]
cons forall a b. (a -> b) -> a -> b
$ \NCon
con ->
    do let conName :: Name
conName = forall s (m :: * -> *) a. MonadReader s m => Getting a s a -> m a
view Lens' NCon Name
nconName NCon
con
       Stab
stab <- Type -> [NCon] -> NCon -> Q Stab
computeOpticType Type
t [NCon]
cons NCon
con
       let n :: Name
n = Name -> Name
prismName Name
conName
       forall (t :: * -> *) (f :: * -> *) a.
(Traversable t, Applicative f) =>
t (f a) -> f (t a)
sequenceA
         ( [ forall (m :: * -> *). Quote m => Name -> m Type -> m Dec
sigD Name
n (forall (m :: * -> *) a. Monad m => a -> m a
return ([Type] -> Type -> Type
quantifyType [] (Set Name -> Stab -> Type
stabToType forall a. Set a
Set.empty Stab
stab)))
           , forall (m :: * -> *).
Quote m =>
m Pat -> m Body -> [m Dec] -> m Dec
valD (forall (m :: * -> *). Quote m => Name -> m Pat
varP Name
n) (forall (m :: * -> *). Quote m => m Exp -> m Body
normalB (Stab -> [NCon] -> NCon -> ExpQ
makeConOpticExp Stab
stab [NCon]
cons NCon
con)) []
           ]
           forall a. [a] -> [a] -> [a]
++ Name -> [Q Dec]
inlinePragma Name
n
         )


-- classy prism class and instance
makeConsPrisms Type
t [NCon]
cons (Just Name
typeName) =
  forall (t :: * -> *) (f :: * -> *) a.
(Traversable t, Applicative f) =>
t (f a) -> f (t a)
sequenceA
    [ Type -> Name -> Name -> [NCon] -> Q Dec
makeClassyPrismClass Type
t Name
className Name
methodName [NCon]
cons
    , Type -> Name -> Name -> [NCon] -> Q Dec
makeClassyPrismInstance Type
t Name
className Name
methodName [NCon]
cons
    ]
  where
  typeNameBase :: String
typeNameBase = Name -> String
nameBase Name
typeName
  className :: Name
className = String -> Name
mkName (String
"As" forall a. [a] -> [a] -> [a]
++ String
typeNameBase)
  sameNameAsCon :: Bool
sameNameAsCon = forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (\NCon
con -> Name -> String
nameBase (forall s (m :: * -> *) a. MonadReader s m => Getting a s a -> m a
view Lens' NCon Name
nconName NCon
con) forall a. Eq a => a -> a -> Bool
== String
typeNameBase) [NCon]
cons
  methodName :: Name
methodName = Bool -> Name -> Name
prismName' Bool
sameNameAsCon Name
typeName


data OpticType = PrismType | ReviewType
data Stab  = Stab Cxt OpticType Type Type Type Type

simplifyStab :: Stab -> Stab
simplifyStab :: Stab -> Stab
simplifyStab (Stab [Type]
cx OpticType
ty Type
_ Type
t Type
_ Type
b) = [Type] -> OpticType -> Type -> Type -> Type -> Type -> Stab
Stab [Type]
cx OpticType
ty Type
t Type
t Type
b Type
b
  -- simplification uses t and b because those types
  -- are interesting in the Review case

stabSimple :: Stab -> Bool
stabSimple :: Stab -> Bool
stabSimple (Stab [Type]
_ OpticType
_ Type
s Type
t Type
a Type
b) = Type
s forall a. Eq a => a -> a -> Bool
== Type
t Bool -> Bool -> Bool
&& Type
a forall a. Eq a => a -> a -> Bool
== Type
b

stabToType :: Set Name -> Stab -> Type
stabToType :: Set Name -> Stab -> Type
stabToType Set Name
clsTVBNames stab :: Stab
stab@(Stab [Type]
cx OpticType
ty Type
s Type
t Type
a Type
b) =
  Set Name -> [Type] -> Type -> Type
quantifyType' Set Name
clsTVBNames [Type]
cx Type
stabTy
  where
  stabTy :: Type
stabTy =
    case OpticType
ty of
      OpticType
PrismType  | Stab -> Bool
stabSimple Stab
stab -> Name
prism'TypeName  Name -> [Type] -> Type
`conAppsT` [Type
t,Type
b]
                 | Bool
otherwise       -> Name
prismTypeName   Name -> [Type] -> Type
`conAppsT` [Type
s,Type
t,Type
a,Type
b]
      OpticType
ReviewType                   -> Name
reviewTypeName  Name -> [Type] -> Type
`conAppsT` [Type
t,Type
b]

stabType :: Stab -> OpticType
stabType :: Stab -> OpticType
stabType (Stab [Type]
_ OpticType
o Type
_ Type
_ Type
_ Type
_) = OpticType
o

computeOpticType :: Type -> [NCon] -> NCon -> Q Stab
computeOpticType :: Type -> [NCon] -> NCon -> Q Stab
computeOpticType Type
t [NCon]
cons NCon
con =
  do let cons' :: [NCon]
cons' = forall a. Eq a => a -> [a] -> [a]
List.delete NCon
con [NCon]
cons
     if forall (t :: * -> *) a. Foldable t => t a -> Bool
null (NCon -> [Name]
_nconVars NCon
con)
         then Type -> [Type] -> [NCon] -> NCon -> Q Stab
computePrismType Type
t (forall s (m :: * -> *) a. MonadReader s m => Getting a s a -> m a
view Lens' NCon [Type]
nconCxt NCon
con) [NCon]
cons' NCon
con
         else Type -> [Type] -> [Type] -> Q Stab
computeReviewType Type
t (forall s (m :: * -> *) a. MonadReader s m => Getting a s a -> m a
view Lens' NCon [Type]
nconCxt NCon
con) (forall s (m :: * -> *) a. MonadReader s m => Getting a s a -> m a
view Lens' NCon [Type]
nconTypes NCon
con)


computeReviewType :: Type -> Cxt -> [Type] -> Q Stab
computeReviewType :: Type -> [Type] -> [Type] -> Q Stab
computeReviewType Type
s' [Type]
cx [Type]
tys =
  do let t :: Type
t = Type
s'
     Type
s <- forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Name -> Type
VarT (forall (m :: * -> *). Quote m => String -> m Name
newName String
"s")
     Type
a <- forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Name -> Type
VarT (forall (m :: * -> *). Quote m => String -> m Name
newName String
"a")
     Type
b <- [TypeQ] -> TypeQ
toTupleT (forall a b. (a -> b) -> [a] -> [b]
map forall (m :: * -> *) a. Monad m => a -> m a
return [Type]
tys)
     forall (m :: * -> *) a. Monad m => a -> m a
return ([Type] -> OpticType -> Type -> Type -> Type -> Type -> Stab
Stab [Type]
cx OpticType
ReviewType Type
s Type
t Type
a Type
b)


-- | Compute the full type-changing Prism type given an outer type,
-- list of constructors, and target constructor name. Additionally
-- return 'True' if the resulting type is a "simple" prism.
computePrismType :: Type -> Cxt -> [NCon] -> NCon -> Q Stab
computePrismType :: Type -> [Type] -> [NCon] -> NCon -> Q Stab
computePrismType Type
t [Type]
cx [NCon]
cons NCon
con =
  do let ts :: [Type]
ts      = forall s (m :: * -> *) a. MonadReader s m => Getting a s a -> m a
view Lens' NCon [Type]
nconTypes NCon
con
         unbound :: Set Name
unbound = forall a s. Getting (Set a) s a -> s -> Set a
setOf forall t. HasTypeVars t => Traversal' t Name
typeVars Type
t forall a. Ord a => Set a -> Set a -> Set a
Set.\\ forall a s. Getting (Set a) s a -> s -> Set a
setOf forall t. HasTypeVars t => Traversal' t Name
typeVars [NCon]
cons
     Map Name Name
sub <- forall (t :: * -> *) (f :: * -> *) a.
(Traversable t, Applicative f) =>
t (f a) -> f (t a)
sequenceA (forall k a. (k -> a) -> Set k -> Map k a
Map.fromSet (forall (m :: * -> *). Quote m => String -> m Name
newName forall b c a. (b -> c) -> (a -> b) -> a -> c
. Name -> String
nameBase) Set Name
unbound)
     Type
b   <- [TypeQ] -> TypeQ
toTupleT (forall a b. (a -> b) -> [a] -> [b]
map forall (m :: * -> *) a. Monad m => a -> m a
return [Type]
ts)
     Type
a   <- [TypeQ] -> TypeQ
toTupleT (forall a b. (a -> b) -> [a] -> [b]
map forall (m :: * -> *) a. Monad m => a -> m a
return (forall t. HasTypeVars t => Map Name Name -> t -> t
substTypeVars Map Name Name
sub [Type]
ts))
     let s :: Type
s = forall t. HasTypeVars t => Map Name Name -> t -> t
substTypeVars Map Name Name
sub Type
t
     forall (m :: * -> *) a. Monad m => a -> m a
return ([Type] -> OpticType -> Type -> Type -> Type -> Type -> Stab
Stab [Type]
cx OpticType
PrismType Type
s Type
t Type
a Type
b)


computeIsoType :: Type -> [Type] -> TypeQ
computeIsoType :: Type -> [Type] -> TypeQ
computeIsoType Type
t' [Type]
fields =
  do Map Name Name
sub <- forall (t :: * -> *) (f :: * -> *) a.
(Traversable t, Applicative f) =>
t (f a) -> f (t a)
sequenceA (forall k a. (k -> a) -> Set k -> Map k a
Map.fromSet (forall (m :: * -> *). Quote m => String -> m Name
newName forall b c a. (b -> c) -> (a -> b) -> a -> c
. Name -> String
nameBase) (forall a s. Getting (Set a) s a -> s -> Set a
setOf forall t. HasTypeVars t => Traversal' t Name
typeVars Type
t'))
     let t :: TypeQ
t = forall (m :: * -> *) a. Monad m => a -> m a
return                    Type
t'
         s :: TypeQ
s = forall (m :: * -> *) a. Monad m => a -> m a
return (forall t. HasTypeVars t => Map Name Name -> t -> t
substTypeVars Map Name Name
sub Type
t')
         b :: TypeQ
b = [TypeQ] -> TypeQ
toTupleT (forall a b. (a -> b) -> [a] -> [b]
map forall (m :: * -> *) a. Monad m => a -> m a
return                    [Type]
fields)
         a :: TypeQ
a = [TypeQ] -> TypeQ
toTupleT (forall a b. (a -> b) -> [a] -> [b]
map forall (m :: * -> *) a. Monad m => a -> m a
return (forall t. HasTypeVars t => Map Name Name -> t -> t
substTypeVars Map Name Name
sub [Type]
fields))

         ty :: TypeQ
ty | forall k a. Map k a -> Bool
Map.null Map Name Name
sub = TypeQ -> [TypeQ] -> TypeQ
appsT (forall (m :: * -> *). Quote m => Name -> m Type
conT Name
iso'TypeName) [TypeQ
t,TypeQ
b]
            | Bool
otherwise    = TypeQ -> [TypeQ] -> TypeQ
appsT (forall (m :: * -> *). Quote m => Name -> m Type
conT Name
isoTypeName) [TypeQ
s,TypeQ
t,TypeQ
a,TypeQ
b]

     [Type] -> Type -> Type
quantifyType [] forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TypeQ
ty



-- | Construct either a Review or Prism as appropriate
makeConOpticExp :: Stab -> [NCon] -> NCon -> ExpQ
makeConOpticExp :: Stab -> [NCon] -> NCon -> ExpQ
makeConOpticExp Stab
stab [NCon]
cons NCon
con =
  case Stab -> OpticType
stabType Stab
stab of
    OpticType
PrismType  -> Stab -> [NCon] -> NCon -> ExpQ
makeConPrismExp Stab
stab [NCon]
cons NCon
con
    OpticType
ReviewType -> NCon -> ExpQ
makeConReviewExp NCon
con


-- | Construct an iso declaration
makeConIso :: Type -> NCon -> DecsQ
makeConIso :: Type -> NCon -> DecsQ
makeConIso Type
s NCon
con =
  do let ty :: TypeQ
ty      = Type -> [Type] -> TypeQ
computeIsoType Type
s (forall s (m :: * -> *) a. MonadReader s m => Getting a s a -> m a
view Lens' NCon [Type]
nconTypes NCon
con)
         defName :: Name
defName = Name -> Name
prismName (forall s (m :: * -> *) a. MonadReader s m => Getting a s a -> m a
view Lens' NCon Name
nconName NCon
con)
     forall (t :: * -> *) (f :: * -> *) a.
(Traversable t, Applicative f) =>
t (f a) -> f (t a)
sequenceA
       ( [ forall (m :: * -> *). Quote m => Name -> m Type -> m Dec
sigD       Name
defName  TypeQ
ty
         , forall (m :: * -> *).
Quote m =>
m Pat -> m Body -> [m Dec] -> m Dec
valD (forall (m :: * -> *). Quote m => Name -> m Pat
varP Name
defName) (forall (m :: * -> *). Quote m => m Exp -> m Body
normalB (NCon -> ExpQ
makeConIsoExp NCon
con)) []
         ] forall a. [a] -> [a] -> [a]
++
         Name -> [Q Dec]
inlinePragma Name
defName
       )


-- | Construct prism expression
--
-- prism <<reviewer>> <<remitter>>
makeConPrismExp ::
  Stab ->
  [NCon] {- ^ constructors       -} ->
  NCon   {- ^ target constructor -} ->
  ExpQ
makeConPrismExp :: Stab -> [NCon] -> NCon -> ExpQ
makeConPrismExp Stab
stab [NCon]
cons NCon
con = forall (m :: * -> *). Quote m => [m Exp] -> m Exp
appsE [forall (m :: * -> *). Quote m => Name -> m Exp
varE Name
prismValName, ExpQ
reviewer, ExpQ
remitter]
  where
  ts :: [Type]
ts = forall s (m :: * -> *) a. MonadReader s m => Getting a s a -> m a
view Lens' NCon [Type]
nconTypes NCon
con
  fields :: Int
fields  = forall (t :: * -> *) a. Foldable t => t a -> Int
length [Type]
ts
  conName :: Name
conName = forall s (m :: * -> *) a. MonadReader s m => Getting a s a -> m a
view Lens' NCon Name
nconName NCon
con

  reviewer :: ExpQ
reviewer                   = Name -> Int -> ExpQ
makeReviewer       Name
conName Int
fields
  remitter :: ExpQ
remitter | Stab -> Bool
stabSimple Stab
stab = Name -> Int -> Int -> ExpQ
makeSimpleRemitter Name
conName (forall (t :: * -> *) a. Foldable t => t a -> Int
length [NCon]
cons) Int
fields
           | Bool
otherwise       = [NCon] -> Name -> ExpQ
makeFullRemitter [NCon]
cons Name
conName


-- | Construct an Iso expression
--
-- iso <<reviewer>> <<remitter>>
makeConIsoExp :: NCon -> ExpQ
makeConIsoExp :: NCon -> ExpQ
makeConIsoExp NCon
con = forall (m :: * -> *). Quote m => [m Exp] -> m Exp
appsE [forall (m :: * -> *). Quote m => Name -> m Exp
varE Name
isoValName, ExpQ
remitter, ExpQ
reviewer]
  where
  conName :: Name
conName = forall s (m :: * -> *) a. MonadReader s m => Getting a s a -> m a
view Lens' NCon Name
nconName NCon
con
  fields :: Int
fields  = forall (t :: * -> *) a. Foldable t => t a -> Int
length (forall s (m :: * -> *) a. MonadReader s m => Getting a s a -> m a
view Lens' NCon [Type]
nconTypes NCon
con)

  reviewer :: ExpQ
reviewer = Name -> Int -> ExpQ
makeReviewer    Name
conName Int
fields
  remitter :: ExpQ
remitter = Name -> Int -> ExpQ
makeIsoRemitter Name
conName Int
fields


-- | Construct a Review expression
--
-- unto (\(x,y,z) -> Con x y z)
makeConReviewExp :: NCon -> ExpQ
makeConReviewExp :: NCon -> ExpQ
makeConReviewExp NCon
con = forall (m :: * -> *). Quote m => m Exp -> m Exp -> m Exp
appE (forall (m :: * -> *). Quote m => Name -> m Exp
varE Name
untoValName) ExpQ
reviewer
  where
  conName :: Name
conName = forall s (m :: * -> *) a. MonadReader s m => Getting a s a -> m a
view Lens' NCon Name
nconName NCon
con
  fields :: Int
fields  = forall (t :: * -> *) a. Foldable t => t a -> Int
length (forall s (m :: * -> *) a. MonadReader s m => Getting a s a -> m a
view Lens' NCon [Type]
nconTypes NCon
con)

  reviewer :: ExpQ
reviewer = Name -> Int -> ExpQ
makeReviewer Name
conName Int
fields


------------------------------------------------------------------------
-- Prism and Iso component builders
------------------------------------------------------------------------


-- | Construct the review portion of a prism.
--
-- (\(x,y,z) -> Con x y z) :: b -> t
makeReviewer :: Name -> Int -> ExpQ
makeReviewer :: Name -> Int -> ExpQ
makeReviewer Name
conName Int
fields =
  do [Name]
xs <- String -> Int -> Q [Name]
newNames String
"x" Int
fields
     forall (m :: * -> *). Quote m => m Pat -> m Exp -> m Exp
lam1E ([PatQ] -> PatQ
toTupleP (forall a b. (a -> b) -> [a] -> [b]
map forall (m :: * -> *). Quote m => Name -> m Pat
varP [Name]
xs))
           (forall (m :: * -> *). Quote m => Name -> m Exp
conE Name
conName ExpQ -> [ExpQ] -> ExpQ
`appsE1` forall a b. (a -> b) -> [a] -> [b]
map forall (m :: * -> *). Quote m => Name -> m Exp
varE [Name]
xs)


-- | Construct the remit portion of a prism.
-- Pattern match only target constructor, no type changing
--
-- (\x -> case s of
--          Con x y z -> Right (x,y,z)
--          _         -> Left x
-- ) :: s -> Either s a
makeSimpleRemitter ::
  Name {- The name of the constructor on which this prism focuses -} ->
  Int  {- The number of constructors the parent data type has     -} ->
  Int  {- The number of fields the constructor has                -} ->
  ExpQ
makeSimpleRemitter :: Name -> Int -> Int -> ExpQ
makeSimpleRemitter Name
conName Int
numCons Int
fields =
  do Name
x  <- forall (m :: * -> *). Quote m => String -> m Name
newName String
"x"
     [Name]
xs <- String -> Int -> Q [Name]
newNames String
"y" Int
fields
     let matches :: [Q Match]
matches =
           [ forall (m :: * -> *).
Quote m =>
m Pat -> m Body -> [m Dec] -> m Match
match (forall (m :: * -> *). Quote m => Name -> [m Pat] -> m Pat
conP Name
conName (forall a b. (a -> b) -> [a] -> [b]
map forall (m :: * -> *). Quote m => Name -> m Pat
varP [Name]
xs))
                   (forall (m :: * -> *). Quote m => m Exp -> m Body
normalB (forall (m :: * -> *). Quote m => m Exp -> m Exp -> m Exp
appE (forall (m :: * -> *). Quote m => Name -> m Exp
conE Name
rightDataName) ([ExpQ] -> ExpQ
toTupleE (forall a b. (a -> b) -> [a] -> [b]
map forall (m :: * -> *). Quote m => Name -> m Exp
varE [Name]
xs))))
                   []
           ] forall a. [a] -> [a] -> [a]
++
           [ forall (m :: * -> *).
Quote m =>
m Pat -> m Body -> [m Dec] -> m Match
match forall (m :: * -> *). Quote m => m Pat
wildP (forall (m :: * -> *). Quote m => m Exp -> m Body
normalB (forall (m :: * -> *). Quote m => m Exp -> m Exp -> m Exp
appE (forall (m :: * -> *). Quote m => Name -> m Exp
conE Name
leftDataName) (forall (m :: * -> *). Quote m => Name -> m Exp
varE Name
x))) []
           | Int
numCons forall a. Ord a => a -> a -> Bool
> Int
1 -- Only generate a catch-all case if there is at least
                         -- one constructor besides the one being focused on.
           ]
     forall (m :: * -> *). Quote m => m Pat -> m Exp -> m Exp
lam1E (forall (m :: * -> *). Quote m => Name -> m Pat
varP Name
x) (forall (m :: * -> *). Quote m => m Exp -> [m Match] -> m Exp
caseE (forall (m :: * -> *). Quote m => Name -> m Exp
varE Name
x) [Q Match]
matches)


-- | Pattern match all constructors to enable type-changing
--
-- (\x -> case s of
--          Con x y z -> Right (x,y,z)
--          Other_n w   -> Left (Other_n w)
-- ) :: s -> Either t a
makeFullRemitter :: [NCon] -> Name -> ExpQ
makeFullRemitter :: [NCon] -> Name -> ExpQ
makeFullRemitter [NCon]
cons Name
target =
  do Name
x <- forall (m :: * -> *). Quote m => String -> m Name
newName String
"x"
     forall (m :: * -> *). Quote m => m Pat -> m Exp -> m Exp
lam1E (forall (m :: * -> *). Quote m => Name -> m Pat
varP Name
x) (forall (m :: * -> *). Quote m => m Exp -> [m Match] -> m Exp
caseE (forall (m :: * -> *). Quote m => Name -> m Exp
varE Name
x) (forall a b. (a -> b) -> [a] -> [b]
map NCon -> Q Match
mkMatch [NCon]
cons))
  where
  mkMatch :: NCon -> Q Match
mkMatch (NCon Name
conName [Name]
_ [Type]
_ [Type]
n) =
    do [Name]
xs <- String -> Int -> Q [Name]
newNames String
"y" (forall (t :: * -> *) a. Foldable t => t a -> Int
length [Type]
n)
       forall (m :: * -> *).
Quote m =>
m Pat -> m Body -> [m Dec] -> m Match
match (forall (m :: * -> *). Quote m => Name -> [m Pat] -> m Pat
conP Name
conName (forall a b. (a -> b) -> [a] -> [b]
map forall (m :: * -> *). Quote m => Name -> m Pat
varP [Name]
xs))
             (forall (m :: * -> *). Quote m => m Exp -> m Body
normalB
               (if Name
conName forall a. Eq a => a -> a -> Bool
== Name
target
                  then forall (m :: * -> *). Quote m => m Exp -> m Exp -> m Exp
appE (forall (m :: * -> *). Quote m => Name -> m Exp
conE Name
rightDataName) ([ExpQ] -> ExpQ
toTupleE (forall a b. (a -> b) -> [a] -> [b]
map forall (m :: * -> *). Quote m => Name -> m Exp
varE [Name]
xs))
                  else forall (m :: * -> *). Quote m => m Exp -> m Exp -> m Exp
appE (forall (m :: * -> *). Quote m => Name -> m Exp
conE Name
leftDataName) (forall (m :: * -> *). Quote m => Name -> m Exp
conE Name
conName ExpQ -> [ExpQ] -> ExpQ
`appsE1` forall a b. (a -> b) -> [a] -> [b]
map forall (m :: * -> *). Quote m => Name -> m Exp
varE [Name]
xs)))
             []


-- | Construct the remitter suitable for use in an 'Iso'
--
-- (\(Con x y z) -> (x,y,z)) :: s -> a
makeIsoRemitter :: Name -> Int -> ExpQ
makeIsoRemitter :: Name -> Int -> ExpQ
makeIsoRemitter Name
conName Int
fields =
  do [Name]
xs <- String -> Int -> Q [Name]
newNames String
"x" Int
fields
     forall (m :: * -> *). Quote m => m Pat -> m Exp -> m Exp
lam1E (forall (m :: * -> *). Quote m => Name -> [m Pat] -> m Pat
conP Name
conName (forall a b. (a -> b) -> [a] -> [b]
map forall (m :: * -> *). Quote m => Name -> m Pat
varP [Name]
xs))
           ([ExpQ] -> ExpQ
toTupleE (forall a b. (a -> b) -> [a] -> [b]
map forall (m :: * -> *). Quote m => Name -> m Exp
varE [Name]
xs))


------------------------------------------------------------------------
-- Classy prisms
------------------------------------------------------------------------


-- | Construct the classy prisms class for a given type and constructors.
--
-- class ClassName r <<vars in type>> | r -> <<vars in Type>> where
--   topMethodName   :: Prism' r Type
--   conMethodName_n :: Prism' r conTypes_n
--   conMethodName_n = topMethodName . conMethodName_n
makeClassyPrismClass ::
  Type   {- Outer type      -} ->
  Name   {- Class name      -} ->
  Name   {- Top method name -} ->
  [NCon] {- Constructors    -} ->
  DecQ
makeClassyPrismClass :: Type -> Name -> Name -> [NCon] -> Q Dec
makeClassyPrismClass Type
t Name
className Name
methodName [NCon]
cons =
  do Name
r <- forall (m :: * -> *). Quote m => String -> m Name
newName String
"r"
     let methodType :: TypeQ
methodType = TypeQ -> [TypeQ] -> TypeQ
appsT (forall (m :: * -> *). Quote m => Name -> m Type
conT Name
prism'TypeName) [forall (m :: * -> *). Quote m => Name -> m Type
varT Name
r,forall (m :: * -> *) a. Monad m => a -> m a
return Type
t]
     [[Dec]]
methodss <- forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (Name -> NCon -> DecsQ
mkMethod Name
r) [NCon]
cons'
     forall (m :: * -> *).
Quote m =>
m [Type] -> Name -> [TyVarBndr ()] -> [FunDep] -> [m Dec] -> m Dec
classD (forall (m :: * -> *). Quote m => [m Type] -> m [Type]
cxt[]) Name
className (Name -> TyVarBndr ()
D.plainTV Name
r forall a. a -> [a] -> [a]
: [TyVarBndr ()]
vs) (Name -> [FunDep]
fds Name
r)
       ( forall (m :: * -> *). Quote m => Name -> m Type -> m Dec
sigD Name
methodName TypeQ
methodType
       forall a. a -> [a] -> [a]
: forall a b. (a -> b) -> [a] -> [b]
map forall (m :: * -> *) a. Monad m => a -> m a
return (forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[Dec]]
methodss)
       )

  where
  mkMethod :: Name -> NCon -> DecsQ
mkMethod Name
r NCon
con =
    do Stab [Type]
cx OpticType
o Type
_ Type
_ Type
_ Type
b <- Type -> [NCon] -> NCon -> Q Stab
computeOpticType Type
t [NCon]
cons NCon
con
       let rTy :: Type
rTy   = Name -> Type
VarT Name
r
           stab' :: Stab
stab' = [Type] -> OpticType -> Type -> Type -> Type -> Type -> Stab
Stab [Type]
cx OpticType
o Type
rTy Type
rTy Type
b Type
b
           defName :: Name
defName = forall s (m :: * -> *) a. MonadReader s m => Getting a s a -> m a
view Lens' NCon Name
nconName NCon
con
           body :: ExpQ
body    = forall (m :: * -> *). Quote m => [m Exp] -> m Exp
appsE [forall (m :: * -> *). Quote m => Name -> m Exp
varE Name
composeValName, forall (m :: * -> *). Quote m => Name -> m Exp
varE Name
methodName, forall (m :: * -> *). Quote m => Name -> m Exp
varE Name
defName]
       forall (t :: * -> *) (f :: * -> *) a.
(Traversable t, Applicative f) =>
t (f a) -> f (t a)
sequenceA
         [ forall (m :: * -> *). Quote m => Name -> m Type -> m Dec
sigD Name
defName        (forall (m :: * -> *) a. Monad m => a -> m a
return (Set Name -> Stab -> Type
stabToType (forall a. Ord a => [a] -> Set a
Set.fromList (Name
rforall a. a -> [a] -> [a]
:[Name]
vNames)) Stab
stab'))
         , forall (m :: * -> *).
Quote m =>
m Pat -> m Body -> [m Dec] -> m Dec
valD (forall (m :: * -> *). Quote m => Name -> m Pat
varP Name
defName) (forall (m :: * -> *). Quote m => m Exp -> m Body
normalB ExpQ
body) []
         ]

  cons' :: [NCon]
cons'         = forall a b. (a -> b) -> [a] -> [b]
map (forall s t a b. ASetter s t a b -> (a -> b) -> s -> t
over Lens' NCon Name
nconName Name -> Name
prismName) [NCon]
cons
  vs :: [TyVarBndr ()]
vs            = forall newFlag oldFlag.
newFlag -> [TyVarBndr_ oldFlag] -> [TyVarBndr_ newFlag]
D.changeTVFlags ()
bndrReq forall a b. (a -> b) -> a -> b
$ [Type] -> [TyVarBndr ()]
D.freeVariablesWellScoped [Type
t]
  vNames :: [Name]
vNames        = forall a b. (a -> b) -> [a] -> [b]
map forall flag. TyVarBndr_ flag -> Name
D.tvName [TyVarBndr ()]
vs
  fds :: Name -> [FunDep]
fds Name
r
    | forall (t :: * -> *) a. Foldable t => t a -> Bool
null [TyVarBndr ()]
vs   = []
    | Bool
otherwise = [[Name] -> [Name] -> FunDep
FunDep [Name
r] [Name]
vNames]



-- | Construct the classy prisms instance for a given type and constructors.
--
-- instance Classname OuterType where
--   topMethodName = id
--   conMethodName_n = <<prism>>
makeClassyPrismInstance ::
  Type ->
  Name     {- Class name      -} ->
  Name     {- Top method name -} ->
  [NCon] {- Constructors    -} ->
  DecQ
makeClassyPrismInstance :: Type -> Name -> Name -> [NCon] -> Q Dec
makeClassyPrismInstance Type
s Name
className Name
methodName [NCon]
cons =
  do let vs :: [TyVarBndr ()]
vs = [Type] -> [TyVarBndr ()]
D.freeVariablesWellScoped [Type
s]
         cls :: Type
cls = Name
className Name -> [Type] -> Type
`conAppsT` (Type
s forall a. a -> [a] -> [a]
: forall a b. (a -> b) -> [a] -> [b]
map forall flag. TyVarBndr_ flag -> Type
tvbToType [TyVarBndr ()]
vs)

     forall (m :: * -> *).
Quote m =>
m [Type] -> m Type -> [m Dec] -> m Dec
instanceD (forall (m :: * -> *). Quote m => [m Type] -> m [Type]
cxt[]) (forall (m :: * -> *) a. Monad m => a -> m a
return Type
cls)
       (   forall (m :: * -> *).
Quote m =>
m Pat -> m Body -> [m Dec] -> m Dec
valD (forall (m :: * -> *). Quote m => Name -> m Pat
varP Name
methodName)
                (forall (m :: * -> *). Quote m => m Exp -> m Body
normalB (forall (m :: * -> *). Quote m => Name -> m Exp
varE Name
idValName)) []
       forall a. a -> [a] -> [a]
: [ do Stab
stab <- Type -> [NCon] -> NCon -> Q Stab
computeOpticType Type
s [NCon]
cons NCon
con
              let stab' :: Stab
stab' = Stab -> Stab
simplifyStab Stab
stab
              forall (m :: * -> *).
Quote m =>
m Pat -> m Body -> [m Dec] -> m Dec
valD (forall (m :: * -> *). Quote m => Name -> m Pat
varP (Name -> Name
prismName Name
conName))
                (forall (m :: * -> *). Quote m => m Exp -> m Body
normalB (Stab -> [NCon] -> NCon -> ExpQ
makeConOpticExp Stab
stab' [NCon]
cons NCon
con)) []
           | NCon
con <- [NCon]
cons
           , let conName :: Name
conName = forall s (m :: * -> *) a. MonadReader s m => Getting a s a -> m a
view Lens' NCon Name
nconName NCon
con
           ]
       )


------------------------------------------------------------------------
-- Utilities
------------------------------------------------------------------------


-- | Normalized constructor
data NCon = NCon
  { NCon -> Name
_nconName :: Name
  , NCon -> [Name]
_nconVars :: [Name]
  , NCon -> [Type]
_nconCxt  :: Cxt
  , NCon -> [Type]
_nconTypes :: [Type]
  }
  deriving (NCon -> NCon -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: NCon -> NCon -> Bool
$c/= :: NCon -> NCon -> Bool
== :: NCon -> NCon -> Bool
$c== :: NCon -> NCon -> Bool
Eq)

instance HasTypeVars NCon where
  typeVarsEx :: Set Name -> Traversal' NCon Name
typeVarsEx Set Name
s Name -> f Name
f (NCon Name
x [Name]
vars [Type]
y [Type]
z) = Name -> [Name] -> [Type] -> [Type] -> NCon
NCon Name
x [Name]
vars forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall t. HasTypeVars t => Set Name -> Traversal' t Name
typeVarsEx Set Name
s' Name -> f Name
f [Type]
y forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall t. HasTypeVars t => Set Name -> Traversal' t Name
typeVarsEx Set Name
s' Name -> f Name
f [Type]
z
    where s' :: Set Name
s' = forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
List.foldl' (forall a b c. (a -> b -> c) -> b -> a -> c
flip forall a. Ord a => a -> Set a -> Set a
Set.insert) Set Name
s [Name]
vars

nconName :: Lens' NCon Name
nconName :: Lens' NCon Name
nconName Name -> f Name
f NCon
x = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\Name
y -> NCon
x {_nconName :: Name
_nconName = Name
y}) (Name -> f Name
f (NCon -> Name
_nconName NCon
x))

nconCxt :: Lens' NCon Cxt
nconCxt :: Lens' NCon [Type]
nconCxt [Type] -> f [Type]
f NCon
x = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\[Type]
y -> NCon
x {_nconCxt :: [Type]
_nconCxt = [Type]
y}) ([Type] -> f [Type]
f (NCon -> [Type]
_nconCxt NCon
x))

nconTypes :: Lens' NCon [Type]
nconTypes :: Lens' NCon [Type]
nconTypes [Type] -> f [Type]
f NCon
x = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\[Type]
y -> NCon
x {_nconTypes :: [Type]
_nconTypes = [Type]
y}) ([Type] -> f [Type]
f (NCon -> [Type]
_nconTypes NCon
x))


-- | Normalize a single 'Con' to its constructor name and field types.
normalizeCon :: D.ConstructorInfo -> NCon
normalizeCon :: ConstructorInfo -> NCon
normalizeCon ConstructorInfo
info = Name -> [Name] -> [Type] -> [Type] -> NCon
NCon (ConstructorInfo -> Name
D.constructorName ConstructorInfo
info)
                         (forall flag. TyVarBndr_ flag -> Name
D.tvName forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ConstructorInfo -> [TyVarBndr ()]
D.constructorVars ConstructorInfo
info)
                         (ConstructorInfo -> [Type]
D.constructorContext ConstructorInfo
info)
                         (ConstructorInfo -> [Type]
D.constructorFields ConstructorInfo
info)


-- | Compute a prism's name by prefixing an underscore for normal
-- constructors and period for operators.
prismName :: Name -> Name
prismName :: Name -> Name
prismName = Bool -> Name -> Name
prismName' Bool
False

-- | Compute a prism's name with a special case for when the type
-- constructor matches one of the value constructors.
--
-- The overlapping flag will be 'True' in the event that:
--
-- 1. We are generating the name of a classy prism for a
--    data type, and
-- 2. The data type shares a name with one of its
--    constructors (e.g., @data A = A@).
--
-- In such a scenario, we take care not to generate the same
-- prism name that the constructor receives (e.g., @_A@).
-- For prefix names, we accomplish this by adding an extra
-- underscore; for infix names, an extra dot.
prismName' ::
  Bool {- ^ overlapping constructor -} ->
  Name {- ^ type constructor        -} ->
  Name {- ^ prism name              -}
prismName' :: Bool -> Name -> Name
prismName' Bool
sameNameAsCon Name
n =
  case Name -> String
nameBase Name
n of
    [] -> forall a. HasCallStack => String -> a
error String
"prismName: empty name base?"
    nb :: String
nb@(Char
x:String
_) | Char -> Bool
isUpper Char
x -> String -> Name
mkName (Char -> String -> String
prefix Char
'_' String
nb)
             | Bool
otherwise -> String -> Name
mkName (Char -> String -> String
prefix Char
'.' String
nb) -- operator
  where
    prefix :: Char -> String -> String
    prefix :: Char -> String -> String
prefix Char
char String
str | Bool
sameNameAsCon = Char
charforall a. a -> [a] -> [a]
:Char
charforall a. a -> [a] -> [a]
:String
str
                    | Bool
otherwise     =      Char
charforall a. a -> [a] -> [a]
:String
str