{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE ViewPatterns #-}
-- 
-- A utility wrapper around both Data.Comp.Multi.Derive and
-- the comptrans library, allowing us to easily derive all desired instances
-- for a higher-order functor, including for Cubix's typeclasses

module Cubix.Language.Parametric.Derive
  (
    deriveAll
  -- , distributeAnnotation
  -- , declareAnnotatedNames
  -- , declareAnnotated

  , deriveAllButDynCase

  , createSortInclusionType
  , createSortInclusionTypes
  , createSortInclusionInfer
  , createSortInclusionInfers

  , sumToNames
  , makeDefaultInstances
  ) where

import Control.Monad ( liftM )

import Language.Haskell.TH.Lib
import Language.Haskell.TH.Syntax hiding ( Cxt )

import Data.Comp.Multi ( (:-<:), project', HFunctor, All )
import Data.Comp.Multi.Derive ( derive, makeHFunctor, makeHTraversable, makeHFoldable, makeEqHF, makeShowHF, makeOrdHF )
import Data.Comp.Multi.Strategy.Derive ( makeDynCase )

import Cubix.Language.Parametric.InjF
import Cubix.Sin.Compdata.Derive ( smartFConstructors )

--------------------------------------------------------------------------------

-- | Derives instances of the following for each type in the list:
-- 
-- @
-- 'HFunctor', 'HTraversable', 'HFoldable', 'EqHF', 'ShowHF', 'OrdHF', 'DynCase'
-- @
-- 
-- Additonally, it will create smart constructors for the data type
deriveAll :: [Name] -> Q [Dec]
deriveAll :: [Name] -> Q [Dec]
deriveAll = [Name -> Q [Dec]] -> [Name] -> Q [Dec]
derive [Name -> Q [Dec]
makeHFunctor, Name -> Q [Dec]
makeHTraversable, Name -> Q [Dec]
makeHFoldable, Name -> Q [Dec]
makeEqHF, Name -> Q [Dec]
makeShowHF,
                    Name -> Q [Dec]
makeOrdHF, Name -> Q [Dec]
smartFConstructors, Name -> Q [Dec]
makeDynCase]
  
deriveAllButDynCase :: [Name] -> Q [Dec]
deriveAllButDynCase :: [Name] -> Q [Dec]
deriveAllButDynCase = [Name -> Q [Dec]] -> [Name] -> Q [Dec]
derive [Name -> Q [Dec]
makeHFunctor, Name -> Q [Dec]
makeHTraversable, Name -> Q [Dec]
makeHFoldable, Name -> Q [Dec]
makeEqHF, Name -> Q [Dec]
makeShowHF,
                    Name -> Q [Dec]
makeOrdHF, Name -> Q [Dec]
smartFConstructors]

-- TODO: distributeAnnotation
-- -- | Distributes an annotation over a sum
-- distributeAnnotation :: Type -> Name -> Type
-- distributeAnnotation typ ann = dist typ
--   where
--     dist :: Type -> Type
--     dist (AppT (AppT (ConT c) l) r)
--          | c == ''Sum             = AppT (AppT (ConT ''(:+:)) (dist l)) (dist r) -- parser complained about pattern
--     dist t                       = AppT (AppT (ConT ''(:&:)) t) (ConT ann)

-- declareAnnotatedType :: String -> Name -> Type -> [Dec]
-- declareAnnotatedType s ann typ = [TySynD (mkName s) [] annTyp]
--   where
--     annTyp = distributeAnnotation typ ann

-- -- | @declareAnnotatedNames s l ts@ will declare @s@ to be the sum of the @ts@, each annotated
-- -- with @l@. It corresponds roughly to the following type-level pseudocode:
-- -- 
-- -- @
-- -- type s = fold (:+:) (map (\t -> t :&: l) ts)
-- -- @
-- declareAnnotatedNames :: String -> Name -> [Name] -> Q [Dec]
-- declareAnnotatedNames _ _ []    = fail "declareAnnotatedNames: Name list empty"
-- declareAnnotatedNames s ann nms = return $ declareAnnotatedType s ann $ foldr typeSum (ConT $ last nms) (map ConT $ init nms)
--   where
--     typeSum a b = AppT (AppT (ConT ''(:+:)) a) b

-- declareAnnotated :: String -> Name -> Name -> Q [Dec]
-- declareAnnotated s ann nm = declareAnnotatedType s ann <$> expandSyns (ConT nm)


createSortInclusionType :: Name -> Name -> Q [Dec]
createSortInclusionType :: Name -> Name -> Q [Dec]
createSortInclusionType fromNm :: Name
fromNm toNm :: Name
toNm = do
  let tName :: Name
tName = Name -> Name -> Name
sortInclusionName Name
fromNm Name
toNm
  Name
e <- String -> Q Name
newName "e"
  Name
i <- String -> Q Name
newName "i"
  let ctx :: [Type]
ctx = [(Type -> Type -> Type) -> Type -> [Type] -> Type
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl Type -> Type -> Type
AppT Type
EqualityT [Name -> Type
VarT Name
i, Name -> Type
ConT Name
toNm]]
  let notStrict :: Bang
notStrict = SourceUnpackedness -> SourceStrictness -> Bang
Bang SourceUnpackedness
NoSourceUnpackedness SourceStrictness
NoSourceStrictness
  let con :: Con
con = [TyVarBndr] -> [Type] -> Con -> Con
ForallC [] [Type]
ctx (Con -> Con) -> Con -> Con
forall a b. (a -> b) -> a -> b
$ Name -> [BangType] -> Con
NormalC Name
tName [(Bang
notStrict, Type -> Type -> Type
AppT (Name -> Type
VarT Name
e) (Name -> Type
ConT Name
fromNm))]
  [Dec] -> Q [Dec]
forall (m :: * -> *) a. Monad m => a -> m a
return ([Dec] -> Q [Dec]) -> [Dec] -> Q [Dec]
forall a b. (a -> b) -> a -> b
$ [[Type]
-> Name
-> [TyVarBndr]
-> Maybe Type
-> [Con]
-> [DerivClause]
-> Dec
DataD [] Name
tName [Name -> Type -> TyVarBndr
KindedTV Name
e (Type -> Type -> Type
AppT (Type -> Type -> Type
AppT Type
ArrowT Type
StarT) Type
StarT), Name -> TyVarBndr
PlainTV Name
i] Maybe Type
forall a. Maybe a
Nothing [Con
con] []]

  
createSortInclusionTypes :: [Name] -> [Name] -> Q [Dec]
createSortInclusionTypes :: [Name] -> [Name] -> Q [Dec]
createSortInclusionTypes froms :: [Name]
froms tos :: [Name]
tos = ([[Dec]] -> [Dec]) -> Q [[Dec]] -> Q [Dec]
forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
liftM [[Dec]] -> [Dec]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat (Q [[Dec]] -> Q [Dec]) -> Q [[Dec]] -> Q [Dec]
forall a b. (a -> b) -> a -> b
$ ((Name, Name) -> Q [Dec]) -> [(Name, Name)] -> Q [[Dec]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ((Name -> Name -> Q [Dec]) -> (Name, Name) -> Q [Dec]
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry Name -> Name -> Q [Dec]
createSortInclusionType) ([(Name, Name)] -> Q [[Dec]]) -> [(Name, Name)] -> Q [[Dec]]
forall a b. (a -> b) -> a -> b
$ [Name] -> [Name] -> [(Name, Name)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Name]
froms [Name]
tos

-- | This is separated from createSortInclusionType because of phase limitations.
--  This needs to refer to the smart constructor, and deriving a smart constructor needs to reify
--  the type name.
createSortInclusionInfer :: Name -> Name -> Q [Dec]
createSortInclusionInfer :: Name -> Name -> Q [Dec]
createSortInclusionInfer fromNm :: Name
fromNm toNm :: Name
toNm = do
    let tname :: Name
tname = Name -> Name -> Name
sortInclusionName Name
fromNm Name
toNm
    Name -> Q [Dec]
mkInjF Name
tname
  where
    mkInjF :: Name -> Q [Dec]
    mkInjF :: Name -> Q [Dec]
mkInjF tName :: Name
tName = do
      let t :: TypeQ
t     = Name -> TypeQ
conT Name
tName
      let fromT :: TypeQ
fromT = Name -> TypeQ
conT Name
fromNm
      let toT :: TypeQ
toT   = Name -> TypeQ
conT Name
toNm
      let smartCon :: ExpQ
smartCon = Name -> ExpQ
varE (Name -> ExpQ) -> Name -> ExpQ
forall a b. (a -> b) -> a -> b
$ Name -> Name
smartConName Name
tName

      Name
x <- String -> Q Name
newName "x"
      let p :: PatQ
p = Name -> [PatQ] -> PatQ
conP Name
tName [Name -> PatQ
varP Name
x]
      let xe :: ExpQ
xe = Name -> ExpQ
varE Name
x

      [d| instance ($t :-<: fs, All HFunctor fs) => InjF fs $fromT $toT where
            injF = $smartCon

            projF' (project' -> Just $p) = Just $xe
            projF' _                     = Nothing
        |]


createSortInclusionInfers :: [Name] -> [Name] -> Q [Dec]
createSortInclusionInfers :: [Name] -> [Name] -> Q [Dec]
createSortInclusionInfers froms :: [Name]
froms tos :: [Name]
tos = ([[Dec]] -> [Dec]) -> Q [[Dec]] -> Q [Dec]
forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
liftM [[Dec]] -> [Dec]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat (Q [[Dec]] -> Q [Dec]) -> Q [[Dec]] -> Q [Dec]
forall a b. (a -> b) -> a -> b
$ ((Name, Name) -> Q [Dec]) -> [(Name, Name)] -> Q [[Dec]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ((Name -> Name -> Q [Dec]) -> (Name, Name) -> Q [Dec]
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry Name -> Name -> Q [Dec]
createSortInclusionInfer) ([(Name, Name)] -> Q [[Dec]]) -> [(Name, Name)] -> Q [[Dec]]
forall a b. (a -> b) -> a -> b
$ [Name] -> [Name] -> [(Name, Name)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Name]
froms [Name]
tos

smartConName :: Name -> Name
smartConName :: Name -> Name
smartConName n :: Name
n = String -> Name
mkName (String -> Name) -> String -> Name
forall a b. (a -> b) -> a -> b
$ "i" String -> String -> String
forall a. [a] -> [a] -> [a]
++ (Name -> String
nameBase Name
n)


sortInclusionName :: Name -> Name -> Name
sortInclusionName :: Name -> Name -> Name
sortInclusionName fromNm :: Name
fromNm toNm :: Name
toNm = String -> Name
mkName (String -> Name) -> String -> Name
forall a b. (a -> b) -> a -> b
$ (String -> String
chopL (String -> String) -> String -> String
forall a b. (a -> b) -> a -> b
$ Name -> String
nameBase Name
fromNm) String -> String -> String
forall a. [a] -> [a] -> [a]
++ "Is" String -> String -> String
forall a. [a] -> [a] -> [a]
++ (String -> String
chopL (String -> String) -> String -> String
forall a b. (a -> b) -> a -> b
$ Name -> String
nameBase Name
toNm)

-- Be warned! This is coupled to the type label name scheme in comptrans
chopL :: String -> String
chopL :: String -> String
chopL s :: String
s = if String -> Char
forall a. [a] -> a
last String
s Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
== 'L' then
            String -> String
forall a. [a] -> [a]
init String
s
          else
            String
s

-- | Takes something like "type Sig = '[ Add, Mul, Var]" and gives [''Add, ''Mul, ''Var]
sumToNames :: Name -> Q [Name]
sumToNames :: Name -> Q [Name]
sumToNames nm :: Name
nm = do
    (TyConI (TySynD _ _ typ :: Type
typ)) <- Name -> Q Info
reify Name
nm
    [Name] -> Q [Name]
forall (m :: * -> *) a. Monad m => a -> m a
return ([Name] -> Q [Name]) -> [Name] -> Q [Name]
forall a b. (a -> b) -> a -> b
$ Type -> [Name]
extract Type
typ
  where
    extract :: Type -> [Name]
    extract :: Type -> [Name]
extract (SigT t :: Type
t _) = Type -> [Name]
extract Type
t
    extract (AppT (AppT PromotedConsT (ConT n :: Name
n)) res :: Type
res) = Name
n Name -> [Name] -> [Name]
forall a. a -> [a] -> [a]
: Type -> [Name]
extract Type
res
    extract PromotedNilT = []
    extract _ = String -> [Name]
forall a. HasCallStack => String -> a
error "sumToNames found invalid summand; only recognizes Sum and names"

-- | Used for creating default cases for Untrans instances.
--   There are two default untranslate cases (identity for shared functors, error for unshared),
--   so we need more advanced typeclass machinery to have both defaults. This caused GHC -O2
--   to go quadratic, and once caused a 64 GB server to run out of memory compiling.
--   So, we use TH instead for the error cases.
--
--  Update: Maybe the 64 GB thing was something else (it's still happening at link time),
--          but we still have pretty conclusive tests showing that our pattern of doing this with pure typeclass
--          machinery was making GHC go quadratic, and then switching to TH has substantially improved compilation times
makeDefaultInstances :: [Type] -> Name -> Name -> Exp -> [Dec]
makeDefaultInstances :: [Type] -> Name -> Name -> Exp -> [Dec]
makeDefaultInstances typs :: [Type]
typs className :: Name
className methName :: Name
methName exp :: Exp
exp =
  ((Type -> Dec) -> [Type] -> [Dec])
-> [Type] -> (Type -> Dec) -> [Dec]
forall a b c. (a -> b -> c) -> b -> a -> c
flip (Type -> Dec) -> [Type] -> [Dec]
forall a b. (a -> b) -> [a] -> [b]
map [Type]
typs ((Type -> Dec) -> [Dec]) -> (Type -> Dec) -> [Dec]
forall a b. (a -> b) -> a -> b
$ \t :: Type
t -> let instSig :: Type
instSig = Type -> Type -> Type
AppT (Name -> Type
ConT Name
className) Type
t in
                        let dec :: Dec
dec = Pat -> Body -> [Dec] -> Dec
ValD (Name -> Pat
VarP Name
methName) (Exp -> Body
NormalB Exp
exp) [] in
                        Maybe Overlap -> [Type] -> Type -> [Dec] -> Dec
InstanceD Maybe Overlap
forall a. Maybe a
Nothing [] Type
instSig [Dec
dec]