{- |
Module      : Control.Lens.Internal.NestedPrismTH
Description : nested pair prisms
Copyright   : (C) 2025 - Eitan Chatav
License     : BSD-style (see the file LICENSE)
Maintainer  : Eitan Chatav <eitan.chatav@gmail.com>
Stability   : provisional
Portability : non-portable

Code is duplicated from `Control.Lens.Internal.PrismTH`,
with small tweaks to support nested pairs.
-}

module Control.Lens.Internal.NestedPrismTH
  ( -- * Nested Prisms
    makeNestedPrisms
  ) where

import Control.Applicative
import Control.Lens.Getter
import Control.Lens.Internal.TH
import Control.Lens.Lens
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 Language.Haskell.TH.Lens
import qualified Data.Map as Map
import qualified Data.Set as Set
import Data.Set (Set)
import Prelude

-- | Generate a `Control.Lens.Prism.Prism`
-- for each constructor of a data type.
-- `Control.Lens.Iso.Iso`s generated when possible.
-- `Control.Lens.Review.Review`s are created for constructors with existentially
-- quantified constructors and GADTs.
--
-- See `Control.Lens.Internal.PrismTH.makePrisms` for details and examples.
-- The difference in `makeNestedPrisms`
-- is that constructors with @n > 2@ arguments
-- will use right-nested pairs, rather than a flat @n@-tuple.
-- This makes them suitable for use on the left-hand-side of
-- `Control.Lens.PartialIso.>?` and `Control.Lens.PartialIso.>?<`;
-- with repeated use of `Data.Profunctor.Distributor.>*<`
-- on the right-hand-side, resulting in right-nested pairs.
makeNestedPrisms :: Name -> DecsQ
makeNestedPrisms :: Name -> DecsQ
makeNestedPrisms Name
typeName =
  do DatatypeInfo
info <- Name -> Q DatatypeInfo
D.reifyDatatype Name
typeName
     let cons :: [ConstructorInfo]
cons = DatatypeInfo -> [ConstructorInfo]
D.datatypeCons DatatypeInfo
info
     Type -> [NCon] -> DecsQ
makeConsPrisms (DatatypeInfo -> Type
datatypeTypeKinded DatatypeInfo
info) ((ConstructorInfo -> NCon) -> [ConstructorInfo] -> [NCon]
forall a b. (a -> b) -> [a] -> [b]
map ConstructorInfo -> NCon
normalizeCon [ConstructorInfo]
cons)

-- Generate prisms for the given type, and normalized constructors.
-- This function dispatches between Iso generation, and normal top-level
makeConsPrisms :: Type -> [NCon] -> DecsQ
-- special case: single constructor -> make iso
makeConsPrisms :: Type -> [NCon] -> DecsQ
makeConsPrisms Type
t [con :: NCon
con@(NCon Name
_ [] [] [Type]
_)] = Type -> NCon -> DecsQ
makeConIso Type
t NCon
con
-- top-level definitions
makeConsPrisms Type
t [NCon]
cons =
  ([[Dec]] -> [Dec]) -> Q [[Dec]] -> DecsQ
forall a b. (a -> b) -> Q a -> Q b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [[Dec]] -> [Dec]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat (Q [[Dec]] -> DecsQ) -> Q [[Dec]] -> DecsQ
forall a b. (a -> b) -> a -> b
$ [NCon] -> (NCon -> DecsQ) -> Q [[Dec]]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
t a -> (a -> f b) -> f (t b)
for [NCon]
cons ((NCon -> DecsQ) -> Q [[Dec]]) -> (NCon -> DecsQ) -> Q [[Dec]]
forall a b. (a -> b) -> a -> b
$ \NCon
con ->
    do let conName :: Name
conName = Getting Name NCon Name -> NCon -> Name
forall s (m :: * -> *) a. MonadReader s m => Getting a s a -> m a
view Getting Name NCon Name
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
       [Q Dec] -> DecsQ
forall (t :: * -> *) (f :: * -> *) a.
(Traversable t, Applicative f) =>
t (f a) -> f (t a)
forall (f :: * -> *) a. Applicative f => [f a] -> f [a]
sequenceA
         ( [ Name -> Q Type -> Q Dec
forall (m :: * -> *). Quote m => Name -> m Type -> m Dec
sigD Name
n (Type -> Q Type
forall a. a -> Q a
forall (m :: * -> *) a. Monad m => a -> m a
return ([Type] -> Type -> Type
quantifyType [] (Set Name -> Stab -> Type
stabToType Set Name
forall a. Set a
Set.empty Stab
stab)))
           , Q Pat -> Q Body -> [Q Dec] -> Q Dec
forall (m :: * -> *).
Quote m =>
m Pat -> m Body -> [m Dec] -> m Dec
valD (Name -> Q Pat
forall (m :: * -> *). Quote m => Name -> m Pat
varP Name
n) (Q Exp -> Q Body
forall (m :: * -> *). Quote m => m Exp -> m Body
normalB (Stab -> [NCon] -> NCon -> Q Exp
makeConOpticExp Stab
stab [NCon]
cons NCon
con)) []
           ]
           [Q Dec] -> [Q Dec] -> [Q Dec]
forall a. [a] -> [a] -> [a]
++ Name -> [Q Dec]
inlinePragma Name
n
         )

data OpticType = PrismType | ReviewType

data Stab  = Stab Cxt OpticType Type Type Type Type

stabSimple :: Stab -> Bool
stabSimple :: Stab -> Bool
stabSimple (Stab [Type]
_ OpticType
_ Type
s Type
t Type
a Type
b) = Type
s Type -> Type -> Bool
forall a. Eq a => a -> a -> Bool
== Type
t Bool -> Bool -> Bool
&& Type
a Type -> Type -> Bool
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' = NCon -> [NCon] -> [NCon]
forall a. Eq a => a -> [a] -> [a]
List.delete NCon
con [NCon]
cons
     if [Name] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null (NCon -> [Name]
_nconVars NCon
con)
         then Type -> [Type] -> [NCon] -> NCon -> Q Stab
computePrismType Type
t (Getting [Type] NCon [Type] -> NCon -> [Type]
forall s (m :: * -> *) a. MonadReader s m => Getting a s a -> m a
view Getting [Type] NCon [Type]
Lens' NCon [Type]
nconCxt NCon
con) [NCon]
cons' NCon
con
         else Type -> [Type] -> [Type] -> Q Stab
computeReviewType Type
t (Getting [Type] NCon [Type] -> NCon -> [Type]
forall s (m :: * -> *) a. MonadReader s m => Getting a s a -> m a
view Getting [Type] NCon [Type]
Lens' NCon [Type]
nconCxt NCon
con) (Getting [Type] NCon [Type] -> NCon -> [Type]
forall s (m :: * -> *) a. MonadReader s m => Getting a s a -> m a
view Getting [Type] NCon [Type]
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 <- (Name -> Type) -> Q Name -> Q Type
forall a b. (a -> b) -> Q a -> Q b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Name -> Type
VarT ([Char] -> Q Name
forall (m :: * -> *). Quote m => [Char] -> m Name
newName [Char]
"s")
     Type
a <- (Name -> Type) -> Q Name -> Q Type
forall a b. (a -> b) -> Q a -> Q b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Name -> Type
VarT ([Char] -> Q Name
forall (m :: * -> *). Quote m => [Char] -> m Name
newName [Char]
"a")
     Type
b <- [Q Type] -> Q Type
toNestedPairT ((Type -> Q Type) -> [Type] -> [Q Type]
forall a b. (a -> b) -> [a] -> [b]
map Type -> Q Type
forall a. a -> Q a
forall (m :: * -> *) a. Monad m => a -> m a
return [Type]
tys)
     Stab -> Q Stab
forall a. a -> Q a
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      = Getting [Type] NCon [Type] -> NCon -> [Type]
forall s (m :: * -> *) a. MonadReader s m => Getting a s a -> m a
view Getting [Type] NCon [Type]
Lens' NCon [Type]
nconTypes NCon
con
         unbound :: Set Name
unbound = Getting (Set Name) Type Name -> Type -> Set Name
forall a s. Getting (Set a) s a -> s -> Set a
setOf Getting (Set Name) Type Name
forall t. HasTypeVars t => Traversal' t Name
Traversal' Type Name
typeVars Type
t Set Name -> Set Name -> Set Name
forall a. Ord a => Set a -> Set a -> Set a
Set.\\ Getting (Set Name) [NCon] Name -> [NCon] -> Set Name
forall a s. Getting (Set a) s a -> s -> Set a
setOf Getting (Set Name) [NCon] Name
forall t. HasTypeVars t => Traversal' t Name
Traversal' [NCon] Name
typeVars [NCon]
cons
     Map Name Name
sub <- Map Name (Q Name) -> Q (Map Name Name)
forall (t :: * -> *) (f :: * -> *) a.
(Traversable t, Applicative f) =>
t (f a) -> f (t a)
forall (f :: * -> *) a.
Applicative f =>
Map Name (f a) -> f (Map Name a)
sequenceA ((Name -> Q Name) -> Set Name -> Map Name (Q Name)
forall k a. (k -> a) -> Set k -> Map k a
Map.fromSet ([Char] -> Q Name
forall (m :: * -> *). Quote m => [Char] -> m Name
newName ([Char] -> Q Name) -> (Name -> [Char]) -> Name -> Q Name
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Name -> [Char]
nameBase) Set Name
unbound)
     Type
b   <- [Q Type] -> Q Type
toNestedPairT ((Type -> Q Type) -> [Type] -> [Q Type]
forall a b. (a -> b) -> [a] -> [b]
map Type -> Q Type
forall a. a -> Q a
forall (m :: * -> *) a. Monad m => a -> m a
return [Type]
ts)
     Type
a   <- [Q Type] -> Q Type
toNestedPairT ((Type -> Q Type) -> [Type] -> [Q Type]
forall a b. (a -> b) -> [a] -> [b]
map Type -> Q Type
forall a. a -> Q a
forall (m :: * -> *) a. Monad m => a -> m a
return (Map Name Name -> [Type] -> [Type]
forall t. HasTypeVars t => Map Name Name -> t -> t
substTypeVars Map Name Name
sub [Type]
ts))
     let s :: Type
s = Map Name Name -> Type -> Type
forall t. HasTypeVars t => Map Name Name -> t -> t
substTypeVars Map Name Name
sub Type
t
     Stab -> Q Stab
forall a. a -> Q a
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] -> Q Type
computeIsoType Type
t' [Type]
fields =
  do Map Name Name
sub <- Map Name (Q Name) -> Q (Map Name Name)
forall (t :: * -> *) (f :: * -> *) a.
(Traversable t, Applicative f) =>
t (f a) -> f (t a)
forall (f :: * -> *) a.
Applicative f =>
Map Name (f a) -> f (Map Name a)
sequenceA ((Name -> Q Name) -> Set Name -> Map Name (Q Name)
forall k a. (k -> a) -> Set k -> Map k a
Map.fromSet ([Char] -> Q Name
forall (m :: * -> *). Quote m => [Char] -> m Name
newName ([Char] -> Q Name) -> (Name -> [Char]) -> Name -> Q Name
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Name -> [Char]
nameBase) (Getting (Set Name) Type Name -> Type -> Set Name
forall a s. Getting (Set a) s a -> s -> Set a
setOf Getting (Set Name) Type Name
forall t. HasTypeVars t => Traversal' t Name
Traversal' Type Name
typeVars Type
t'))
     let t :: Q Type
t = Type -> Q Type
forall a. a -> Q a
forall (m :: * -> *) a. Monad m => a -> m a
return                    Type
t'
         s :: Q Type
s = Type -> Q Type
forall a. a -> Q a
forall (m :: * -> *) a. Monad m => a -> m a
return (Map Name Name -> Type -> Type
forall t. HasTypeVars t => Map Name Name -> t -> t
substTypeVars Map Name Name
sub Type
t')
         b :: Q Type
b = [Q Type] -> Q Type
toNestedPairT ((Type -> Q Type) -> [Type] -> [Q Type]
forall a b. (a -> b) -> [a] -> [b]
map Type -> Q Type
forall a. a -> Q a
forall (m :: * -> *) a. Monad m => a -> m a
return                    [Type]
fields)
         a :: Q Type
a = [Q Type] -> Q Type
toNestedPairT ((Type -> Q Type) -> [Type] -> [Q Type]
forall a b. (a -> b) -> [a] -> [b]
map Type -> Q Type
forall a. a -> Q a
forall (m :: * -> *) a. Monad m => a -> m a
return (Map Name Name -> [Type] -> [Type]
forall t. HasTypeVars t => Map Name Name -> t -> t
substTypeVars Map Name Name
sub [Type]
fields))
         ty :: Q Type
ty | Map Name Name -> Bool
forall k a. Map k a -> Bool
Map.null Map Name Name
sub = Q Type -> [Q Type] -> Q Type
appsT (Name -> Q Type
forall (m :: * -> *). Quote m => Name -> m Type
conT Name
iso'TypeName) [Q Type
t,Q Type
b]
            | Bool
otherwise    = Q Type -> [Q Type] -> Q Type
appsT (Name -> Q Type
forall (m :: * -> *). Quote m => Name -> m Type
conT Name
isoTypeName) [Q Type
s,Q Type
t,Q Type
a,Q Type
b]
     [Type] -> Type -> Type
quantifyType [] (Type -> Type) -> Q Type -> Q Type
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Q Type
ty

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

-- Construct an iso declaration
makeConIso :: Type -> NCon -> DecsQ
makeConIso :: Type -> NCon -> DecsQ
makeConIso Type
s NCon
con =
  do let ty :: Q Type
ty      = Type -> [Type] -> Q Type
computeIsoType Type
s (Getting [Type] NCon [Type] -> NCon -> [Type]
forall s (m :: * -> *) a. MonadReader s m => Getting a s a -> m a
view Getting [Type] NCon [Type]
Lens' NCon [Type]
nconTypes NCon
con)
         defName :: Name
defName = Name -> Name
prismName (Getting Name NCon Name -> NCon -> Name
forall s (m :: * -> *) a. MonadReader s m => Getting a s a -> m a
view Getting Name NCon Name
Lens' NCon Name
nconName NCon
con)
     [Q Dec] -> DecsQ
forall (t :: * -> *) (f :: * -> *) a.
(Traversable t, Applicative f) =>
t (f a) -> f (t a)
forall (f :: * -> *) a. Applicative f => [f a] -> f [a]
sequenceA
       ( [ Name -> Q Type -> Q Dec
forall (m :: * -> *). Quote m => Name -> m Type -> m Dec
sigD       Name
defName  Q Type
ty
         , Q Pat -> Q Body -> [Q Dec] -> Q Dec
forall (m :: * -> *).
Quote m =>
m Pat -> m Body -> [m Dec] -> m Dec
valD (Name -> Q Pat
forall (m :: * -> *). Quote m => Name -> m Pat
varP Name
defName) (Q Exp -> Q Body
forall (m :: * -> *). Quote m => m Exp -> m Body
normalB (NCon -> Q Exp
makeConIsoExp NCon
con)) []
         ] [Q Dec] -> [Q Dec] -> [Q Dec]
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 -> Q Exp
makeConPrismExp Stab
stab [NCon]
cons NCon
con = [Q Exp] -> Q Exp
forall (m :: * -> *). Quote m => [m Exp] -> m Exp
appsE [Name -> Q Exp
forall (m :: * -> *). Quote m => Name -> m Exp
varE Name
prismValName, Q Exp
reviewer, Q Exp
remitter]
  where
  ts :: [Type]
ts = Getting [Type] NCon [Type] -> NCon -> [Type]
forall s (m :: * -> *) a. MonadReader s m => Getting a s a -> m a
view Getting [Type] NCon [Type]
Lens' NCon [Type]
nconTypes NCon
con
  fields :: Int
fields  = [Type] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Type]
ts
  conName :: Name
conName = Getting Name NCon Name -> NCon -> Name
forall s (m :: * -> *) a. MonadReader s m => Getting a s a -> m a
view Getting Name NCon Name
Lens' NCon Name
nconName NCon
con
  reviewer :: Q Exp
reviewer                   = Name -> Int -> Q Exp
makeReviewer       Name
conName Int
fields
  remitter :: Q Exp
remitter | Stab -> Bool
stabSimple Stab
stab = Name -> Int -> Int -> Q Exp
makeSimpleRemitter Name
conName ([NCon] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [NCon]
cons) Int
fields
           | Bool
otherwise       = [NCon] -> Name -> Q Exp
makeFullRemitter [NCon]
cons Name
conName

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

-- Construct a Review expression
--
-- unto (\(x,y,z) -> Con x y z)
makeConReviewExp :: NCon -> ExpQ
makeConReviewExp :: NCon -> Q Exp
makeConReviewExp NCon
con = Q Exp -> Q Exp -> Q Exp
forall (m :: * -> *). Quote m => m Exp -> m Exp -> m Exp
appE (Name -> Q Exp
forall (m :: * -> *). Quote m => Name -> m Exp
varE Name
untoValName) Q Exp
reviewer
  where
  conName :: Name
conName = Getting Name NCon Name -> NCon -> Name
forall s (m :: * -> *) a. MonadReader s m => Getting a s a -> m a
view Getting Name NCon Name
Lens' NCon Name
nconName NCon
con
  fields :: Int
fields  = [Type] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length (Getting [Type] NCon [Type] -> NCon -> [Type]
forall s (m :: * -> *) a. MonadReader s m => Getting a s a -> m a
view Getting [Type] NCon [Type]
Lens' NCon [Type]
nconTypes NCon
con)
  reviewer :: Q Exp
reviewer = Name -> Int -> Q Exp
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 -> Q Exp
makeReviewer Name
conName Int
fields =
  do [Name]
xs <- [Char] -> Int -> Q [Name]
newNames [Char]
"x" Int
fields
     Q Pat -> Q Exp -> Q Exp
forall (m :: * -> *). Quote m => m Pat -> m Exp -> m Exp
lam1E ([Q Pat] -> Q Pat
toNestedPairP ((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]
xs))
           (Name -> Q Exp
forall (m :: * -> *). Quote m => Name -> m Exp
conE Name
conName Q Exp -> [Q Exp] -> Q Exp
`appsE1` (Name -> Q Exp) -> [Name] -> [Q Exp]
forall a b. (a -> b) -> [a] -> [b]
map Name -> Q Exp
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 -> Q Exp
makeSimpleRemitter Name
conName Int
numCons Int
fields =
  do Name
x  <- [Char] -> Q Name
forall (m :: * -> *). Quote m => [Char] -> m Name
newName [Char]
"x"
     [Name]
xs <- [Char] -> Int -> Q [Name]
newNames [Char]
"y" Int
fields
     let matches :: [Q Match]
matches =
           [ 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
conName ((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]
xs))
                   (Q Exp -> Q Body
forall (m :: * -> *). Quote m => m Exp -> m Body
normalB (Q Exp -> Q Exp -> Q Exp
forall (m :: * -> *). Quote m => m Exp -> m Exp -> m Exp
appE (Name -> Q Exp
forall (m :: * -> *). Quote m => Name -> m Exp
conE Name
rightDataName) ([Q Exp] -> Q Exp
toNestedPairE ((Name -> Q Exp) -> [Name] -> [Q Exp]
forall a b. (a -> b) -> [a] -> [b]
map Name -> Q Exp
forall (m :: * -> *). Quote m => Name -> m Exp
varE [Name]
xs))))
                   []
           ] [Q Match] -> [Q Match] -> [Q Match]
forall a. [a] -> [a] -> [a]
++
           [ Q Pat -> Q Body -> [Q Dec] -> Q Match
forall (m :: * -> *).
Quote m =>
m Pat -> m Body -> [m Dec] -> m Match
match Q Pat
forall (m :: * -> *). Quote m => m Pat
wildP (Q Exp -> Q Body
forall (m :: * -> *). Quote m => m Exp -> m Body
normalB (Q Exp -> Q Exp -> Q Exp
forall (m :: * -> *). Quote m => m Exp -> m Exp -> m Exp
appE (Name -> Q Exp
forall (m :: * -> *). Quote m => Name -> m Exp
conE Name
leftDataName) (Name -> Q Exp
forall (m :: * -> *). Quote m => Name -> m Exp
varE Name
x))) []
           | Int
numCons Int -> Int -> Bool
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.
           ]
     Q Pat -> Q Exp -> Q Exp
forall (m :: * -> *). Quote m => m Pat -> m Exp -> m Exp
lam1E (Name -> Q Pat
forall (m :: * -> *). Quote m => Name -> m Pat
varP Name
x) (Q Exp -> [Q Match] -> Q Exp
forall (m :: * -> *). Quote m => m Exp -> [m Match] -> m Exp
caseE (Name -> Q Exp
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 -> Q Exp
makeFullRemitter [NCon]
cons Name
target =
  do Name
x <- [Char] -> Q Name
forall (m :: * -> *). Quote m => [Char] -> m Name
newName [Char]
"x"
     Q Pat -> Q Exp -> Q Exp
forall (m :: * -> *). Quote m => m Pat -> m Exp -> m Exp
lam1E (Name -> Q Pat
forall (m :: * -> *). Quote m => Name -> m Pat
varP Name
x) (Q Exp -> [Q Match] -> Q Exp
forall (m :: * -> *). Quote m => m Exp -> [m Match] -> m Exp
caseE (Name -> Q Exp
forall (m :: * -> *). Quote m => Name -> m Exp
varE Name
x) ((NCon -> Q Match) -> [NCon] -> [Q Match]
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 <- [Char] -> Int -> Q [Name]
newNames [Char]
"y" ([Type] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Type]
n)
       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
conName ((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]
xs))
             (Q Exp -> Q Body
forall (m :: * -> *). Quote m => m Exp -> m Body
normalB
               (if Name
conName Name -> Name -> Bool
forall a. Eq a => a -> a -> Bool
== Name
target
                  then Q Exp -> Q Exp -> Q Exp
forall (m :: * -> *). Quote m => m Exp -> m Exp -> m Exp
appE (Name -> Q Exp
forall (m :: * -> *). Quote m => Name -> m Exp
conE Name
rightDataName) ([Q Exp] -> Q Exp
toNestedPairE ((Name -> Q Exp) -> [Name] -> [Q Exp]
forall a b. (a -> b) -> [a] -> [b]
map Name -> Q Exp
forall (m :: * -> *). Quote m => Name -> m Exp
varE [Name]
xs))
                  else Q Exp -> Q Exp -> Q Exp
forall (m :: * -> *). Quote m => m Exp -> m Exp -> m Exp
appE (Name -> Q Exp
forall (m :: * -> *). Quote m => Name -> m Exp
conE Name
leftDataName) (Name -> Q Exp
forall (m :: * -> *). Quote m => Name -> m Exp
conE Name
conName Q Exp -> [Q Exp] -> Q Exp
`appsE1` (Name -> Q Exp) -> [Name] -> [Q Exp]
forall a b. (a -> b) -> [a] -> [b]
map Name -> Q Exp
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 -> Q Exp
makeIsoRemitter Name
conName Int
fields =
  do [Name]
xs <- [Char] -> Int -> Q [Name]
newNames [Char]
"x" Int
fields
     Q Pat -> Q Exp -> Q Exp
forall (m :: * -> *). Quote m => m Pat -> m Exp -> m Exp
lam1E (Name -> [Q Pat] -> Q Pat
forall (m :: * -> *). Quote m => Name -> [m Pat] -> m Pat
conP Name
conName ((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]
xs))
           ([Q Exp] -> Q Exp
toNestedPairE ((Name -> Q Exp) -> [Name] -> [Q Exp]
forall a b. (a -> b) -> [a] -> [b]
map Name -> Q Exp
forall (m :: * -> *). Quote m => Name -> m Exp
varE [Name]
xs))

------------------------------------------------------------------------
-- 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
(NCon -> NCon -> Bool) -> (NCon -> NCon -> Bool) -> Eq NCon
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: NCon -> NCon -> Bool
== :: NCon -> NCon -> Bool
$c/= :: NCon -> NCon -> Bool
/= :: 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 ([Type] -> [Type] -> NCon) -> f [Type] -> f ([Type] -> NCon)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Set Name -> Traversal' [Type] Name
forall t. HasTypeVars t => Set Name -> Traversal' t Name
typeVarsEx Set Name
s' Name -> f Name
f [Type]
y f ([Type] -> NCon) -> f [Type] -> f NCon
forall a b. f (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Set Name -> Traversal' [Type] Name
forall t. HasTypeVars t => Set Name -> Traversal' t Name
typeVarsEx Set Name
s' Name -> f Name
f [Type]
z
    where s' :: Set Name
s' = (Set Name -> Name -> Set Name) -> Set Name -> [Name] -> Set Name
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
List.foldl' ((Name -> Set Name -> Set Name) -> Set Name -> Name -> Set Name
forall a b c. (a -> b -> c) -> b -> a -> c
flip Name -> Set Name -> Set Name
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 = (Name -> NCon) -> f Name -> f NCon
forall a b. (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\Name
y -> NCon
x {_nconName = 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 = ([Type] -> NCon) -> f [Type] -> f NCon
forall a b. (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\[Type]
y -> NCon
x {_nconCxt = 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 = ([Type] -> NCon) -> f [Type] -> f NCon
forall a b. (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\[Type]
y -> NCon
x {_nconTypes = 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)
                         (TyVarBndr_ () -> Name
forall flag. TyVarBndr_ flag -> Name
D.tvName (TyVarBndr_ () -> Name) -> [TyVarBndr_ ()] -> [Name]
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 {- type constructor        -} ->
  Name {- prism name              -}
prismName :: Name -> Name
prismName Name
n =
  case Name -> [Char]
nameBase Name
n of
    [] -> [Char] -> Name
forall a. HasCallStack => [Char] -> a
error [Char]
"prismName: empty name base?"
    nb :: [Char]
nb@(Char
x:[Char]
_) | Char -> Bool
isUpper Char
x -> [Char] -> Name
mkName (Char -> [Char] -> [Char]
prefix Char
'_' [Char]
nb)
             | Bool
otherwise -> [Char] -> Name
mkName (Char -> [Char] -> [Char]
prefix Char
'.' [Char]
nb) -- operator
  where
    prefix :: Char -> String -> String
    prefix :: Char -> [Char] -> [Char]
prefix Char
char [Char]
str = Char
charChar -> [Char] -> [Char]
forall a. a -> [a] -> [a]
:[Char]
str

-- Construct a tuple type given a list of types.
toNestedPairT :: [TypeQ] -> TypeQ
toNestedPairT :: [Q Type] -> Q Type
toNestedPairT [] = Q Type -> [Q Type] -> Q Type
appsT (Int -> Q Type
forall (m :: * -> *). Quote m => Int -> m Type
tupleT Int
0) []
toNestedPairT [Q Type
x] = Q Type
x
toNestedPairT (Q Type
x:[Q Type]
xs) = Q Type -> [Q Type] -> Q Type
appsT (Int -> Q Type
forall (m :: * -> *). Quote m => Int -> m Type
tupleT Int
2) [Q Type
x, [Q Type] -> Q Type
toNestedPairT [Q Type]
xs]

-- Construct a tuple value given a list of expressions.
toNestedPairE :: [ExpQ] -> ExpQ
toNestedPairE :: [Q Exp] -> Q Exp
toNestedPairE [] = [Q Exp] -> Q Exp
forall (m :: * -> *). Quote m => [m Exp] -> m Exp
tupE []
toNestedPairE [Q Exp
x] = Q Exp
x
toNestedPairE (Q Exp
x:[Q Exp]
xs) = [Q Exp] -> Q Exp
forall (m :: * -> *). Quote m => [m Exp] -> m Exp
tupE [Q Exp
x, [Q Exp] -> Q Exp
toNestedPairE [Q Exp]
xs]

-- Construct a tuple pattern given a list of patterns.
toNestedPairP :: [PatQ] -> PatQ
toNestedPairP :: [Q Pat] -> Q Pat
toNestedPairP [] = [Q Pat] -> Q Pat
forall (m :: * -> *). Quote m => [m Pat] -> m Pat
tupP []
toNestedPairP [Q Pat
x] = Q Pat
x
toNestedPairP (Q Pat
x:[Q Pat]
xs) = [Q Pat] -> Q Pat
forall (m :: * -> *). Quote m => [m Pat] -> m Pat
tupP [Q Pat
x, [Q Pat] -> Q Pat
toNestedPairP [Q Pat]
xs]