{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE DefaultSignatures #-}
{-# LANGUAGE DerivingVia #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE ScopedTypeVariables #-}
module Barbies.TH (FieldNamesB(..)
  , LensB(..)
  , getLensB
  , AccessorsB(..)
  , declareBareB
  ) where

import Language.Haskell.TH hiding (cxt)
import Language.Haskell.TH.Syntax (VarBangType, Name(..), mkOccName, occString)
import Data.String
import Data.Foldable (foldl')
import Barbies
import Barbies.Constraints
import Barbies.Bare
import Data.Functor.Product
import GHC.Generics (Generic)
import Control.Applicative
import Data.Functor.Identity (Identity(..))
import Data.Functor.Compose (Compose(..))
import Data.List.Split

-- | A pair of a getter and a setter
-- Not van Laarhoven to avoid dictionary passing
data LensB b a = LensB
  { LensB b a -> forall (h :: k -> *). b h -> h a
viewB :: forall h. b h -> h a
  , LensB b a -> forall (h :: k -> *). h a -> b h -> b h
setB :: forall h. h a -> b h -> b h
  }

getLensB :: Functor f => LensB b a -> (h a -> f (h a)) -> b h -> f (b h)
getLensB :: LensB b a -> (h a -> f (h a)) -> b h -> f (b h)
getLensB (LensB forall (h :: k -> *). b h -> h a
v forall (h :: k -> *). h a -> b h -> b h
s) h a -> f (h a)
f b h
b = (\h a
x -> h a -> b h -> b h
forall (h :: k -> *). h a -> b h -> b h
s h a
x b h
b) (h a -> b h) -> f (h a) -> f (b h)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> h a -> f (h a)
f (b h -> h a
forall (h :: k -> *). b h -> h a
v b h
b)
{-# INLINE getLensB #-}

class AccessorsB b where
  -- | A collection of lenses (getter-setter pairs)
  baccessors :: b (LensB b)

-- | barbies doesn't care about field names, but they are useful in many use cases
class FieldNamesB b where
  -- | A collection of field names.
  bfieldNames :: IsString a => b (Const a)

-- | Transform a regular Haskell record declaration into HKD form.
-- 'BareB', 'FieldNamesB', 'FunctorB', 'DistributiveB',
-- 'TraversableB', 'ApplicativeB' and 'ConstraintsB' instances are
-- derived.
--
-- For example,
--
-- @declareBareB [d|data User = User { uid :: Int, name :: String}|]@
--
-- becomes
--
-- @data User t f = User { uid :: Wear t f Int, name :: Wear t f String }@
--
declareBareB :: DecsQ -> DecsQ
declareBareB :: DecsQ -> DecsQ
declareBareB DecsQ
decsQ = do
  [Dec]
decs <- DecsQ
decsQ
  [[Dec]]
decs' <- (Dec -> DecsQ) -> [Dec] -> Q [[Dec]]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse Dec -> DecsQ
go [Dec]
decs
  [Dec] -> DecsQ
forall (m :: * -> *) a. Monad m => a -> m a
return ([Dec] -> DecsQ) -> [Dec] -> DecsQ
forall a b. (a -> b) -> a -> b
$ [[Dec]] -> [Dec]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[Dec]]
decs'
  where
    go :: Dec -> DecsQ
go (DataD Cxt
_ Name
dataName [TyVarBndr]
tvbs Maybe Kind
_ [con :: Con
con@(RecC Name
conName [VarBangType]
fields)] [DerivClause]
classes) = do
      Name
varS <- String -> Q Name
newName String
"sw"
      Name
varW <- String -> Q Name
newName String
"h"
      let xs :: [Name]
xs = String -> [VarBangType] -> [Name]
varNames String
"x" [VarBangType]
fields
      let ys :: [Name]
ys = String -> [VarBangType] -> [Name]
varNames String
"y" [VarBangType]
fields
      Name
varB <- String -> Q Name
newName String
"b"
      let transformed :: Con
transformed = Name -> Name -> Con -> Con
transformCon Name
varS Name
varW Con
con
      let names :: Exp
names = (Exp -> Exp -> Exp) -> Exp -> [Exp] -> Exp
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' Exp -> Exp -> Exp
AppE (Name -> Exp
ConE Name
conName) [Exp -> Exp -> Exp
AppE (Name -> Exp
ConE 'Const) (Exp -> Exp) -> Exp -> Exp
forall a b. (a -> b) -> a -> b
$ Exp -> Exp -> Exp
AppE (Name -> Exp
VarE 'fromString) (Exp -> Exp) -> Exp -> Exp
forall a b. (a -> b) -> a -> b
$ Lit -> Exp
LitE (Lit -> Exp) -> Lit -> Exp
forall a b. (a -> b) -> a -> b
$ String -> Lit
StringL (String -> Lit) -> String -> Lit
forall a b. (a -> b) -> a -> b
$ Name -> String
nameBase Name
name | (Name
name, Bang
_, Kind
_) <- [VarBangType]
fields]
          accessors :: ExpQ
accessors = (ExpQ -> ExpQ -> ExpQ) -> ExpQ -> [ExpQ] -> ExpQ
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' ExpQ -> ExpQ -> ExpQ
appE (Name -> ExpQ
conE Name
conName)
            [ [|LensB
                $(varE name)
                (\ $(varP varW) $(varP varB) -> $(recUpdE (varE varB) [pure (name, VarE varW)])) |]
            | (Name
name, Bang
_, Kind
_) <- [VarBangType]
fields]

          -- Turn TyVarBndr into just a Name such that we can
          -- reconstruct the constructor applied to already-present
          -- type variables below.
          varName :: TyVarBndr -> Name
varName (PlainTV Name
n) = Name
n
          varName (KindedTV Name
n Kind
_) = Name
n

          -- The type name as present originally along with its type
          -- variables.
          vanillaType :: TypeQ
vanillaType = (TypeQ -> TypeQ -> TypeQ) -> TypeQ -> [TypeQ] -> TypeQ
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' TypeQ -> TypeQ -> TypeQ
appT (Name -> TypeQ
conT Name
dataName) (Name -> TypeQ
varT (Name -> TypeQ) -> (TyVarBndr -> Name) -> TyVarBndr -> TypeQ
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TyVarBndr -> Name
varName (TyVarBndr -> TypeQ) -> [TyVarBndr] -> [TypeQ]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [TyVarBndr]
tvbs)

          types :: Cxt
types = [Kind
t | (Name
_, Bang
_, Kind
t) <- [VarBangType]
fields]

      Name
varCstr <- String -> Q Name
newName String
"c"

      let datC :: TypeQ
datC = TypeQ
vanillaType TypeQ -> TypeQ -> TypeQ
`appT` Name -> TypeQ
conT ''Covered
      [Dec]
decs <- [d|
        instance BareB $(vanillaType) where
          bcover $(conP conName $ map varP xs) = $(foldl'
              appE
              (conE conName)
              (appE (conE 'Identity) . varE <$> xs)
            )
          {-# INLINE bcover #-}
          bstrip $(conP conName $ map varP xs) = $(foldl'
              appE
              (conE conName)
              (appE (varE 'runIdentity) . varE <$> xs)
            )
          {-# INLINE bstrip #-}
        instance FieldNamesB $(datC) where bfieldNames = $(pure names)
        instance AccessorsB $(datC) where baccessors = $(accessors)
        instance FunctorB $(datC) where
          bmap f $(conP conName $ map varP xs) = $(foldl'
              appE
              (conE conName)
              (appE (varE 'f) . varE <$> xs)
            )
        instance DistributiveB $(datC) where
          bdistribute fb = $(foldl'
              appE
              (conE conName)
              [ [| Compose ($(varE (unmangle fd)) <$> fb) |] | (fd, _, _) <- fields ]
            )
        instance TraversableB $(datC) where
          btraverse f $(conP conName $ map varP xs) = $(fst $ foldl'
              (\(l, op) r -> (infixE (Just l) (varE op) (Just r), '(<*>)))
              (conE conName, '(<$>))
              (appE (varE 'f) . varE <$> xs)
            )
          {-# INLINE btraverse #-}
        instance ConstraintsB $(datC) where
          type AllB $(varT varCstr) $(datC) = $(foldl appT (tupleT (length types)) [varT varCstr `appT` pure t | t <- types])
          baddDicts $(conP conName $ map varP xs) = $(foldl'
            (\r x -> [|$(r) (Pair Dict $(varE x))|])
            (conE conName) xs)
        instance ApplicativeB $(datC) where
          bpure x = $(foldl'
            (\r _ -> [|$(r) x|])
            (conE conName) xs)
          bprod $(conP conName $ map varP xs) $(conP conName $ map varP ys) = $(foldl'
            (\r (x, y) -> [|$(r) (Pair $(varE x) $(varE y))|])
            (conE conName) (zip xs ys))
        |]
      [[Dec]]
drvs <- (TypeQ -> DecsQ) -> [TypeQ] -> Q [[Dec]]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (\TypeQ
cls ->
        [d|deriving via Barbie $(datC) $(varT varW)
            instance ($(cls) (Barbie $(datC) $(varT varW))) => $(cls) ($(datC) $(varT varW))|])
        [ Kind -> TypeQ
forall (f :: * -> *) a. Applicative f => a -> f a
pure Kind
t | DerivClause Maybe DerivStrategy
_ Cxt
preds <- [DerivClause]
classes, Kind
t <- Cxt
preds ]
      [Dec] -> DecsQ
forall (m :: * -> *) a. Monad m => a -> m a
return ([Dec] -> DecsQ) -> [Dec] -> DecsQ
forall a b. (a -> b) -> a -> b
$ Cxt
-> Name
-> [TyVarBndr]
-> Maybe Kind
-> [Con]
-> [DerivClause]
-> Dec
DataD [] Name
dataName
        ([TyVarBndr]
tvbs [TyVarBndr] -> [TyVarBndr] -> [TyVarBndr]
forall a. [a] -> [a] -> [a]
++ [Name -> TyVarBndr
PlainTV Name
varS, Name -> TyVarBndr
PlainTV Name
varW])
        Maybe Kind
forall a. Maybe a
Nothing
        [Con
transformed]
        [Maybe DerivStrategy -> Cxt -> DerivClause
DerivClause Maybe DerivStrategy
forall a. Maybe a
Nothing [Name -> Kind
ConT ''Generic]]
        Dec -> [Dec] -> [Dec]
forall a. a -> [a] -> [a]
: [Dec]
decs [Dec] -> [Dec] -> [Dec]
forall a. [a] -> [a] -> [a]
++ [[Dec]] -> [Dec]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[Dec]]
drvs
    go Dec
d = [Dec] -> DecsQ
forall (f :: * -> *) a. Applicative f => a -> f a
pure [Dec
d]

varNames :: String -> [VarBangType] -> [Name]
varNames :: String -> [VarBangType] -> [Name]
varNames String
p [VarBangType]
vbt = [String -> Name
mkName (String
p String -> String -> String
forall a. [a] -> [a] -> [a]
++ Name -> String
nameBase (Name -> Name
unmangle Name
v)) | (Name
v, Bang
_, Kind
_) <- [VarBangType]
vbt]

transformCon :: Name -- ^ switch variable
  -> Name -- ^ wrapper variable
  -> Con -- ^ original constructor
  -> Con
transformCon :: Name -> Name -> Con -> Con
transformCon Name
switchName Name
wrapperName (RecC Name
name [VarBangType]
xs) = Name -> [VarBangType] -> Con
RecC Name
name
  [(Name -> Name
unmangle Name
v, Bang
b, Name -> Kind
ConT ''Wear
    Kind -> Kind -> Kind
`AppT` Name -> Kind
VarT Name
switchName
    Kind -> Kind -> Kind
`AppT` Name -> Kind
VarT Name
wrapperName
    Kind -> Kind -> Kind
`AppT` Kind
t)
  | (Name
v, Bang
b, Kind
t) <- [VarBangType]
xs
  ]
transformCon Name
var Name
w (ForallC [TyVarBndr]
tvbs Cxt
cxt Con
con) = [TyVarBndr] -> Cxt -> Con -> Con
ForallC [TyVarBndr]
tvbs Cxt
cxt (Con -> Con) -> Con -> Con
forall a b. (a -> b) -> a -> b
$ Name -> Name -> Con -> Con
transformCon Name
var Name
w Con
con
transformCon Name
_ Name
_ Con
con = String -> Con
forall a. HasCallStack => String -> a
error (String -> Con) -> String -> Con
forall a b. (a -> b) -> a -> b
$ String
"transformCon: unsupported " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Con -> String
forall a. Show a => a -> String
show Con
con

-- | Unmangle record field names
--
-- When 'DuplicateRecordFields' is turned on, record field names are mangled.
-- (see https://gitlab.haskell.org/ghc/ghc/-/wikis/records/overloaded-record-fields/duplicate-record-fields#mangling-selector-names)
-- We undo that because these mangled field names don't round-trip through TH splices.
unmangle :: Name -> Name
unmangle :: Name -> Name
unmangle (Name OccName
occ NameFlavour
flavour) = OccName -> NameFlavour -> Name
Name OccName
occ' NameFlavour
flavour
  where
    occ' :: OccName
occ' = case (Char -> Bool) -> String -> [String]
forall a. (a -> Bool) -> [a] -> [[a]]
wordsBy (Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
== Char
':') (OccName -> String
occString OccName
occ) of
        [String
"$sel", String
fd, String
_qual] -> String -> OccName
mkOccName String
fd
        [String]
_ -> OccName
occ