{-# LANGUAGE TemplateHaskell #-}
module Staged.Commons (
    -- * Types
    TExpQ,
    C (..),
    -- * Constructors
    appTE,
    -- * Function conversions
    toFn, fromFn,
    toFn2, fromFn2,
    -- * Mapping helpers
    -- ** Unary
    mapIC,
    mapKC,
    mapCI,
    mapCK,
    mapCC,
    -- ** Binary
    mapIIC,
    mapIKC,
    mapICI,
    mapICK,
    mapICC,
    mapKIC,
    mapKKC,
    mapKCI,
    mapKCK,
    mapKCC,
    mapCII,
    mapCIK,
    mapCIC,
    mapCKI,
    mapCKK,
    mapCKC,
    mapCCI,
    mapCCK,
    mapCCC,
    ) where

import Language.Haskell.TH (TExpQ)
import Generics.SOP (K (..), I (..))

-- | Code
newtype C a = C (TExpQ a)

-------------------------------------------------------------------------------
-- Constructors
-------------------------------------------------------------------------------

appTE :: TExpQ (a -> b) -> TExpQ a -> TExpQ b
appTE f x = [|| $$f $$x ||]

-------------------------------------------------------------------------------
-- generics-sop inspired functions
-------------------------------------------------------------------------------

-- arity: 2

mapIC :: (x -> TExpQ x) -> I x -> C x
mapIC f (I a) = C (f a)

mapKC :: (a -> TExpQ x) -> K a x -> C x
mapKC f (K a) = C (f a)

mapCI :: (TExpQ x -> x) -> C x -> I x
mapCI f (C a) = I (f a)

mapCK :: (TExpQ x -> a) -> C x -> K a x
mapCK f (C a) = K (f a)

mapCC :: (TExpQ x -> TExpQ x) -> C x -> C x
mapCC f (C a) = C (f a)

-- arity: 3

mapIIC :: ( x ->  x -> TExpQ x) -> I x -> I x -> C x
mapIIC f (I a) (I b) = C (f a b)

mapIKC :: (x -> a -> TExpQ x) -> I x -> K a x -> C x
mapIKC f (I a) (K b) = C (f a b)

mapICI :: (x -> TExpQ x -> x) -> I x -> C x -> I x
mapICI f (I a) (C b) = I (f a b)

mapICK :: (x -> TExpQ x -> a) -> I x -> C x -> K a x
mapICK f (I a) (C b) = K (f a b)

mapICC :: (x -> TExpQ x -> TExpQ x) -> I x -> C x -> C x
mapICC f (I a) (C b) = C (f a b)

mapKIC :: (a -> x -> TExpQ x) -> K a x -> I x -> C x
mapKIC f (K a) (I b) = C (f a b)

mapKKC :: (a -> b -> TExpQ x) -> K a x -> K b x -> C x
mapKKC f (K a) (K b) = C (f a b)

mapKCI :: (a -> TExpQ x -> x) -> K a x -> C x -> I x
mapKCI f (K a) (C b) = I (f a b)

mapKCK :: (a -> TExpQ x -> b) -> K a x -> C x -> K b x
mapKCK f (K a) (C b) = K (f a b)

mapKCC :: (a -> TExpQ x -> TExpQ x) -> K a x -> C x -> C x
mapKCC f (K a) (C b) = C (f a b)

mapCII :: (TExpQ x -> x -> x) -> C x -> I x -> I x
mapCII f (C a) (I b) = I (f a b)

mapCIK :: (TExpQ x ->  x -> a) -> C x -> I x -> K a x
mapCIK f (C a) (I b) = K (f a b)

mapCIC :: (TExpQ x -> x -> TExpQ x) -> C x -> I x -> C x
mapCIC f (C a) (I b) = C (f a b)

mapCKI :: (TExpQ x -> a -> x) -> C x -> K a x -> I x
mapCKI f (C a) (K b) = I (f a b)

mapCKK :: (TExpQ x -> a -> b) -> C x -> K a x -> K b x
mapCKK f (C a) (K b) = K (f a b)

mapCKC :: (TExpQ x -> a -> TExpQ x) -> C x -> K a x -> C x
mapCKC f (C a) (K b) = C (f a b)

mapCCI :: (TExpQ x -> TExpQ x ->  x) -> C x -> C x -> I x
mapCCI f (C a) (C b) = I (f a b)

mapCCK :: (TExpQ x -> TExpQ x -> a) -> C x -> C x -> K a x
mapCCK f (C a) (C b) = K (f a b)

mapCCC :: (TExpQ x -> TExpQ x -> TExpQ x) -> C x -> C x -> C x
mapCCC f (C a) (C b) = C (f a b)

------------------------------------------------------------------------------
-- Function conversions
-------------------------------------------------------------------------------

toFn :: TExpQ (a -> b) -> TExpQ a -> TExpQ b
toFn f x = [|| $$f $$x ||]

toFn2 :: TExpQ (a -> b -> c) -> TExpQ a -> TExpQ b -> TExpQ c
toFn2 f x y = [|| $$f $$x $$y ||]

fromFn :: (TExpQ a -> TExpQ b) -> TExpQ (a -> b)
fromFn f = [|| \x -> $$(f [|| x ||]) ||]

fromFn2 :: (TExpQ a -> TExpQ b -> TExpQ c) -> TExpQ (a -> b -> c)
fromFn2 f = [|| \x y -> $$(f [|| x ||] [|| y ||]) ||]