packages feed

tuple-classes-1.0.1: src/Data/Tuple/Classes/TH.hs

{-# LANGUAGE ImplicitParams #-}
{-# LANGUAGE TemplateHaskellQuotes #-}

module Data.Tuple.Classes.TH (
    makeStrictTupleAndInsts,
    makeUncurryNInst,
    makeTupleAtInsts,
) where

import Control.DeepSeq       (NFData(..), NFData1, NFData2(..))
import Control.Monad         (forM)
import Data.Bifoldable       (Bifoldable(..))
import Data.Bifunctor        (Bifunctor(..))
import Data.Bifunctor.Swap   (Swap(..))
import Data.Binary           (Binary(..))
import Data.Bitraversable    (Bitraversable(..))
import Data.Foldable         (foldl')
import Data.Functor.Classes  (Eq1(..), Eq2(..))
import Data.Functor.Identity (Identity(..))
import Data.Hashable         (Hashable)
import Data.Hashable.Lifted  (Hashable1)
import Data.Strict.Classes   (Strict(..))
import Data.Strict.Tuple     (Pair(..))
import GHC.Generics          (Generic, Generic1)
import Language.Haskell.TH   hiding (Strict, strictType)

strictTName :: Int -> Name
strictTName 0 = ''()
strictTName 1 = ''Identity
strictTName 2 = ''Pair
strictTName n = mkName $ "StrictTuple" ++ show n

strictDName :: Int -> Name
strictDName 0 = '()
strictDName 1 = 'Identity
strictDName 2 = '(:!:)
strictDName n = strictTName n

dropLast :: (?n :: Int) => Int -> [a] -> [a]
dropLast u = take $ ?n - u

-- StrictTuple3
mkNames :: (?n :: Int) => (Enum a) => a -> (a -> String) -> [Name]
mkNames x toName = map (mkName . toName) $ take ?n $ enumFrom x

abc, xs, ys :: (?n :: Int) => [Name]
-- a, b, c
abc = mkNames 'a' (: [])
-- x0, x1, x2
xs = mkNames (0 :: Int) $ ('x' :) . show
-- y0, y1, y2
ys = mkNames (0 :: Int) $ ('y' :) . show

fName, gName, cur, unc :: Name
fName = mkName "f"
gName = mkName "g"
cur = mkName "cur"
unc = mkName "unc"

-- !a
banged :: Name -> BangTypeQ
banged a = bangType (bang noSourceUnpackedness sourceStrict) (varT a)

-- StrictTuple3 a
strictTupleUnsatT :: (?n :: Int) => Int -> TypeQ
strictTupleUnsatT u =
    foldl' appT (conT $ strictTName ?n) $ dropLast u $ map varT abc

-- StrictTuple3 a b c
strictTupleSatT :: (?n :: Int) => TypeQ
strictTupleSatT = strictTupleUnsatT 0

-- (a, b, c)
lazyTupleSatT :: (?n :: Int) => TypeQ
lazyTupleSatT = foldl' appT (tupleT ?n) $ map varT abc

-- StrictTuple3 x0 x1 x2
strictTupleP :: (?n :: Int) => [Name] -> PatQ
strictTupleP = conP (strictDName ?n) . map varP

-- StrictTuple3 x0
strictTupleUnsatE :: (?n :: Int) => Int -> ExpQ
strictTupleUnsatE u =
    foldl' appE (conE $ strictDName ?n) $ dropLast u $ map varE xs

-- StrictTuple3 x0 x1 x2
strictTupleSatE :: (?n :: Int) => ExpQ
strictTupleSatE = strictTupleUnsatE 0

{-
instance UncurryN 3 (a -> b -> c -> d) where
    type UncurriedN 3 (a -> b -> c -> d) = (a, b, c) -> d
    type UncurriedN' 3 (a -> b -> c -> d) = StrictTuple3 a b c -> d
    uncurryN cur ~(x0, x1, x2) = cur x0 x1 x2
    curryN unc x0 x1 x2 = unc (x0, x1, x2)
    uncurryN' cur ~(StrictTuple3 x0 x1 x2) = cur x0 x1 x2
    curryN' unc x0 x1 x2 = unc (StrictTuple3 x0 x1 x2)
-}
makeUncurryNInst :: (?n :: Int) => DecQ
makeUncurryNInst =
    let d = mkName [['a' ..] !! ?n]
        nT = litT $ numTyLit $ fromIntegral ?n
        curT = foldr (\a r -> [t| $(varT a) -> $r |]) (varT d) abc
        uncurryClause ctor = clause
            [varP cur, tildeP $ conP ctor $ map varP xs]
            (normalB $ foldl' appE (varE cur) $ map varE xs)
            []
        curryClause ctor = clause
            (map varP $ unc : xs)
            (normalB $ varE unc `appE` foldl' appE (conE ctor) (map varE xs))
            []
    in instanceD
        (pure [])
        [t| $(conT $ mkName "UncurryN") $nT $curT |]
        [
            tySynInstD $ tySynEqn
                Nothing
                [t| $(conT $ mkName "UncurriedN") $nT $curT |]
                [t| $lazyTupleSatT -> $(varT d) |],
            tySynInstD $ tySynEqn
                Nothing
                [t| $(conT $ mkName "UncurriedN'") $nT $curT |]
                [t| $strictTupleSatT -> $(varT d) |],
            funD (mkName "uncurryN") [uncurryClause $ tupleDataName ?n],
            funD (mkName "curryN") [curryClause $ tupleDataName ?n],
            funD (mkName "uncurryN'") [uncurryClause $ strictDName ?n],
            funD (mkName "curryN'") [curryClause $ strictDName ?n]]

makeTupleAtInsts :: (?n :: Int) => DecsQ
makeTupleAtInsts = fmap concat $ forM [1 .. ?n] $ \i -> do
    ((a, x), (abc', xs')) <- case splitAt (i - 1) $ zip abc xs of
        (ls, ax : rs) -> return (ax, unzip $ ls ++ rs)
        _             -> fail $
            "makeTupleAtInsts: no element at index " ++ show i

    let tupleAtInst :: Name -> Name -> Name -> Name -> DecQ
        tupleAtInst removeType removeData insertType insertData = instanceD
            (pure [])
            (foldl' appT (conT $ mkName "TupleAt") [
                litT $ numTyLit $ fromIntegral i,
                varT a,
                foldl' appT (conT removeType) $ map varT abc',
                foldl' appT (conT insertType) $ map varT abc])
            [
                funD (mkName "tupleInsert") [clause
                    [varP x, tildeP $ conP removeData $ map varP xs']
                    (normalB $ foldl' appE (conE insertData) $ map varE xs)
                    []],
                funD (mkName "tupleRemove") [clause
                    [conP insertData $ map varP xs]
                    (normalB $ tupE [
                        varE x,
                        foldl' appE (conE removeData) $ map varE xs'])
                    []]]

    sequence [
        {-
        instance TupleAt 1 b (a, c) (a, b, c) where
            tupleInsert x1 ~(x0, x2) = (x0, x1, x2)
            tupleRemove (x0, x1, x2) = (x1, (x0, x2))
        -}
        tupleAtInst
            (tupleTypeName $ ?n - 1)
            (tupleDataName $ ?n - 1)
            (tupleTypeName ?n)
            (tupleDataName ?n),

        -- TODO This tilde makes implementation simpler and probably has no
        -- impact, but let's confirm that with benchmarks.
        {-
        instance TupleAt 1 b (Pair a c) (StrictTuple3 a b c) where
            tupleInsert x1 ~(x0 :!: x2) = StrictTuple3 x0 x1 x2
            tupleRemove (StrictTuple3 x0 x1 x2) = (x1, x0 :!: x2)
        -}
        tupleAtInst
            (strictTName $ ?n - 1)
            (strictDName $ ?n - 1)
            (strictTName ?n)
            (strictDName ?n)]

makeStrictTupleAndInsts :: (?n :: Int) => DecsQ
makeStrictTupleAndInsts = do
    (x1, x2) <- case drop (?n - 2) xs of
        [x1, x2] -> return (x1, x2)
        _        -> fail "makeStrictTupleAndInsts: ?n < 2"

    {-
    data StrictTuple4 a b c = StrictTuple4 a b c
        deriving (...)
    -}
    strictTupleDef <- dataD
        (pure [])
        (strictTName ?n)
        [PlainTV a BndrReq | a <- abc]
        Nothing
        [normalC (strictDName ?n) $ map banged abc]
        [derivClause Nothing $ map conT [
            ''Eq,
            ''Ord,
            ''Read,
            ''Show,
            ''Foldable,
            ''Functor,
            ''Traversable,
            ''Generic,
            ''Generic1]]

    {-
    instance (Hashable a, Hashable b, Hashable c) => Hashable (StrictTuple3 a b c)
    -}
    hashableInst <- instanceD
        (cxt [[t| Hashable $(varT a) |] | a <- abc])
        [t| Hashable $strictTupleSatT |]
        []

    {-
    instance (Hashable a, Hashable b) => Hashable1 (StrictTuple3 a b)
    -}
    hashable1Inst <- instanceD
        (cxt [[t| Hashable $(varT a) |] | a <- init abc])
        [t| Hashable1 $(strictTupleUnsatT 1) |]
        []

    {-
    instance (Eq a, Eq b) => Eq1 (StrictTuple3 a b) where
        liftEq = liftEq2 (==)
    -}
    eq1Inst <- instanceD
        (cxt [[t| Eq $(varT a) |] | a <- init abc])
        [t| Eq1 $(strictTupleUnsatT 1) |]
        [funD 'liftEq [clause [] (normalB [| liftEq2 (==) |]) []]]

    {-
    instance (Eq a) => Eq2 (StrictTuple3 a) where
        liftEq2 f g (StrictTuple3 x0 x1 x2) (StrictTuple3 y0 y1 y2) =
            x0 == y0 && f x1 y1 && g x2 y2
    -}
    eq2Inst <- instanceD
        (cxt [[t| Eq $(varT a) |] | a <- dropLast 2 abc])
        [t| Eq2 $(strictTupleUnsatT 2) |]
        [funD 'liftEq2 [clause
            [varP fName, varP gName, strictTupleP xs, strictTupleP ys]
            (normalB $ foldr
                (\boolE r -> [| $boolE && $r |])
                [|
                    $(varE fName) $(varE x1) $(varE $ ys !! (?n - 2)) &&
                    $(varE gName) $(varE x2) $(varE $ ys !! (?n - 1)) |]
                (dropLast 2 $
                    zipWith (\x y -> [| $(varE x) == $(varE y) |]) xs ys))
            []]]

    {-
    instance (NFData a, NFData b, NFData c) => NFData (StrictTuple3 a b c)
    -}
    nfDataInst <- instanceD
        (cxt [[t| NFData $(varT a) |] | a <- abc])
        [t| NFData $strictTupleSatT |]
        []

    {-
    instance (NFData a, NFData b) => NFData1 (StrictTuple3 a b)
    -}
    nfData1Inst <- instanceD
        (cxt [[t| NFData $(varT a) |] | a <- init abc])
        [t| NFData1 $(strictTupleUnsatT 1) |]
        []

    {-
    instance (NFData a) => NFData2 (StrictTuple3 a) where
        liftRnf2 r s (StrictTuple3 x0 x1 x2) = x0 `seq` r x1 `seq` s x2
    -}
    nfData2Inst <- instanceD
        (cxt [[t| NFData $(varT a) |] | a <- dropLast 2 abc])
        [t| NFData2 $(strictTupleUnsatT 2) |]
        [funD 'liftRnf2 [clause
            [varP fName, varP gName, strictTupleP xs]
            (normalB $ foldr
                (\x r -> [| rnf $(varE x) `seq` $r |])
                [| $(varE fName) $(varE x1) `seq` $(varE gName) $(varE x2) |]
                (dropLast 2 xs))
            []]]

    {-
    instance (Binary a, Binary b, Binary c) => Binary (StrictTuple3 a b c) where
        put = put . toLazy
        get = toStrict <$> get
    -}
    binaryInst <- instanceD
        (cxt [[t| Binary $(varT a) |] | a <- abc])
        [t| Binary $strictTupleSatT |]
        [
            funD 'put [clause [] (normalB [| put . toLazy |]) []],
            funD 'get [clause [] (normalB [| toStrict <$> get |]) []]]

    {-
    instance (Monoid a, Monoid b) => Applicative (StrictTuple3 a b) where
        pure = StrictTuple3 mempty mempty
        StrictTuple3 x0 x1 x2 <*> StrictTuple3 y0 y1 y2 =
            StrictTuple (x0 <> y0) (x1 <> y1) (x2 y2)
    -}
    applicativeInst <- instanceD
        (cxt [[t| Monoid $(varT a) |] | a <- init abc])
        [t| Applicative $(strictTupleUnsatT 1) |]
        [
            funD 'pure [clause
                [varP $ last xs]
                (normalB $ foldl' appE (conE $ strictDName ?n) $
                    replicate (?n - 1) [| mempty |] ++ [varE $ last xs])
                []],
            funD '(<*>) [clause
                [strictTupleP xs, strictTupleP ys]
                (normalB $ foldl' appE (conE $ strictDName ?n) $
                    let mappendE x y = [| $(varE x) <> $(varE y) |]
                        mappendE's = init $ zipWith mappendE xs ys
                    in mappendE's ++ [varE (last xs) `appE` varE (last ys)])
                []]]

    {-
    instance (Monoid a, Monoid b) => Monad (StrictTuple3 a b) where
        StrictTuple3 x0 x1 x2 >>= f = StrictTuple3 (x0 <> y0) (x1 <> y1) y2 where
            StrictTuple3 y0 y1 y2 = f x2
    -}
    monadInst <- instanceD
        (cxt [[t| Monoid $(varT a) |] | a <- init abc])
        [t| Monad $(strictTupleUnsatT 1) |]
        [funD '(>>=) [clause
            [strictTupleP xs, varP fName]
            (normalB $ foldl' appE (conE $ strictDName ?n) $
                let mappendE x y = [| $(varE x) <> $(varE y) |]
                in init (zipWith mappendE xs ys) ++ [varE $ last ys])
            [valD
                (strictTupleP ys)
                (normalB $ varE fName `appE` varE (last xs))
                []]]]

    otherInsts <-
        let bifoldMapP = tildeP $ conP (strictDName ?n) $
                replicate (?n - 2) wildP ++ [varP x1, varP x2]
        in [d|
            {-
            instance Strict (a, b, c) (StrictTuple3 a b c) where
                toStrict (x0, x1, x2) = StrictTuple3 x0 x1 x2
                toLazy (StrictTuple3 x0 x1 x2) = (x0, x1, x2)
            -}
            instance Strict $lazyTupleSatT $strictTupleSatT where
                toStrict $(tupP $ map varP xs) = $strictTupleSatE
                toLazy $(strictTupleP xs) = $(tupE $ map varE xs)

            {-
            instance Swap (StrictTuple3 a) where
                swap (StrictTuple3 x0 x1 x2) = StrictTuple3 x0 x2 x1
            -}
            instance Swap $(strictTupleUnsatT 2) where
                swap $(strictTupleP xs) =
                    $(strictTupleUnsatE 2) $(varE x2) $(varE x1)

            {-
            instance Bifoldable (StrictTuple3 a) where
                bifoldMap f g ~(StrictTuple3 _ x1 x2) = f x1 <> g x2
            -}
            instance Bifoldable $(strictTupleUnsatT 2) where
                bifoldMap f g $bifoldMapP = f $(varE x1) <> g $(varE x2)

            {-
            instance Bifunctor (StrictTuple3 a) where
                bimap f g ~(StrictTuple3 x0 x1 x2) =
                    StrictTuple3 x0 (f x1) (g x2)
            -}
            instance Bifunctor $(strictTupleUnsatT 2) where
                bimap f g $(tildeP $ strictTupleP xs) =
                    $(strictTupleUnsatE 2) (f $(varE x1)) (g $(varE x2))

            {-
            instance Bitraversable (StrictTuple3 a) where
                bitraverse f g ~(StrictTuple3 x0 x1 x2) =
                    liftA2 (StrictTuple3 x0) (f x1) (g x2)
            -}
            instance Bitraversable $(strictTupleUnsatT 2) where
                bitraverse f g $(tildeP $ strictTupleP xs) =
                    $(strictTupleUnsatE 2) <$> f $(varE x1) <*> g $(varE x2)
            |]

    return $ strictTupleDef
        : hashableInst
        : hashable1Inst
        : eq1Inst
        : eq2Inst
        : nfDataInst
        : nfData1Inst
        : nfData2Inst
        : binaryInst
        : applicativeInst
        : monadInst
        : otherInsts