{-# LANGUAGE TemplateHaskell #-}

-- 
-- 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
  , deriveAllButSortInjection
  -- , distributeAnnotation
  -- , declareAnnotatedNames
  -- , declareAnnotated

  , deriveAllButDynCase

  , createSortInclusionType
  , createSortInclusionType'
  , createSortInclusionTypes
  , createSortInclusionInfer
  , createSortInclusionInfer'
  , createSortInclusionInfers
  , sumToNames
  , makeDefaultInstances
  ) where

import Control.Monad ( liftM )
import Language.Haskell.TH.Lib ( conP, conT, varE, varP )
import Language.Haskell.TH.Syntax ( Bang(..), BndrVis(..), Body(..), Con(..), Dec(..), Exp, Info(..), Name, Overlap(..), Pat(..), Q, SourceStrictness(..), SourceUnpackedness(..), TyVarBndr(..), Type(..), mkName, nameBase, newName, reify )

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

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

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

makeIsNotSortInjection :: Name -> Q [Dec]
makeIsNotSortInjection :: Name -> Q [Dec]
makeIsNotSortInjection Name
tName =
  let t :: Q Type
t = Name -> Q Type
forall (m :: * -> *). Quote m => Name -> m Type
conT Name
tName
   in [d| type instance IsSortInjection $Q Type
t = 'False |]

makeIsSortInjection :: Name -> Q [Dec]
makeIsSortInjection :: Name -> Q [Dec]
makeIsSortInjection Name
tName =
  let t :: Q Type
t = Name -> Q Type
forall (m :: * -> *). Quote m => Name -> m Type
conT Name
tName
   in [d| type instance IsSortInjection $Q Type
t = 'True |]


-- | 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]
makeConstrNameHF, Name -> Q [Dec]
smartConstructors, Name -> Q [Dec]
patternSynonyms, Name -> Q [Dec]
smartFConstructors
  , Name -> Q [Dec]
makeDynCase, Name -> Q [Dec]
makeIsNotSortInjection
  ]

deriveAllButSortInjection :: [Name] -> Q [Dec]
deriveAllButSortInjection :: [Name] -> Q [Dec]
deriveAllButSortInjection = [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]
makeConstrNameHF, Name -> Q [Dec]
smartConstructors, Name -> Q [Dec]
patternSynonyms, 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]
makeConstrNameHF, Name -> Q [Dec]
smartConstructors, Name -> Q [Dec]
patternSynonyms, 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 -> Name -> Q [Dec]
createSortInclusionType' :: Name -> Name -> Name -> Q [Dec]
createSortInclusionType' Name
fromNm Name
toNm Name
tName = do
  Name
e <- [Char] -> Q Name
forall (m :: * -> *). Quote m => [Char] -> m Name
newName [Char]
"e"
  Name
i <- [Char] -> Q Name
forall (m :: * -> *). Quote m => [Char] -> m Name
newName [Char]
"i"
  let ctx :: [Type]
ctx = [(Type -> Type -> Type) -> Type -> [Type] -> Type
forall b a. (b -> a -> b) -> b -> [a] -> b
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 Specificity] -> [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 a. a -> Q a
forall (f :: * -> *) a. Applicative f => a -> f a
return ([Dec] -> Q [Dec]) -> [Dec] -> Q [Dec]
forall a b. (a -> b) -> a -> b
$ [[Type]
-> Name
-> [TyVarBndr BndrVis]
-> Maybe Type
-> [Con]
-> [DerivClause]
-> Dec
DataD [] Name
tName [Name -> BndrVis -> Type -> TyVarBndr BndrVis
forall flag. Name -> flag -> Type -> TyVarBndr flag
KindedTV Name
e BndrVis
BndrReq (Type -> Type -> Type
AppT (Type -> Type -> Type
AppT Type
ArrowT Type
StarT) Type
StarT), Name -> BndrVis -> TyVarBndr BndrVis
forall flag. Name -> flag -> TyVarBndr flag
PlainTV Name
i BndrVis
BndrReq] Maybe Type
forall a. Maybe a
Nothing [Con
con] []]

createSortInclusionType :: Name -> Name -> Q [Dec]
createSortInclusionType :: Name -> Name -> Q [Dec]
createSortInclusionType Name
fromNm Name
toNm = do
  let tName :: Name
tName = Name -> Name -> Name
sortInclusionName Name
fromNm Name
toNm
   in Name -> Name -> Name -> Q [Dec]
createSortInclusionType' Name
fromNm Name
toNm Name
tName

createSortInclusionTypes :: [Name] -> [Name] -> Q [Dec]
createSortInclusionTypes :: [Name] -> [Name] -> Q [Dec]
createSortInclusionTypes [Name]
froms [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)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [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 -> Name -> Q [Dec]
createSortInclusionInfer' :: Name -> Name -> Name -> Q [Dec]
createSortInclusionInfer' Name
fromNm Name
toNm Name
tName = do
  let t :: Q Type
t     = Name -> Q Type
forall (m :: * -> *). Quote m => Name -> m Type
conT Name
tName
  let fromT :: Q Type
fromT = Name -> Q Type
forall (m :: * -> *). Quote m => Name -> m Type
conT Name
fromNm
  let toT :: Q Type
toT   = Name -> Q Type
forall (m :: * -> *). Quote m => Name -> m Type
conT Name
toNm
  let smartCon :: Q Exp
smartCon = Name -> Q Exp
forall (m :: * -> *). Quote m => Name -> m Exp
varE (Name -> Q Exp) -> Name -> Q Exp
forall a b. (a -> b) -> a -> b
$ Name -> Name
smartConName Name
tName

  Name
x <- [Char] -> Q Name
forall (m :: * -> *). Quote m => [Char] -> m Name
newName [Char]
"x"
  let p :: Q Pat
p = Name -> [Q Pat] -> Q Pat
forall (m :: * -> *). Quote m => Name -> [m Pat] -> m Pat
conP Name
tName [Name -> Q Pat
forall (m :: * -> *). Quote m => Name -> m Pat
varP Name
x]
  let xe :: Q Exp
xe = Name -> Q Exp
forall (m :: * -> *). Quote m => Name -> m Exp
varE Name
x

  [d|
   type instance IsSortInjection $Q Type
t = True
   type instance SortInjectionSource $Q Type
t = $Q Type
fromT
   type instance SortInjectionTarget $Q Type
t = $Q Type
toT
   instance {-# OVERLAPPING #-} ($Q Type
t :-<: fs, All HFunctor fs) => InjF fs $Q Type
fromT $Q Type
toT where
     injF = $Q Exp
smartCon

     projF' (project' -> Just $Q Pat
p) = Just $Q Exp
xe
     projF' _                     = Nothing
   instance RemoveSortInjectionNode $Q Type
t where
     removeSortInjectionNode $Q Pat
p = $Q Exp
xe
    |]

createSortInclusionInfer :: Name -> Name -> Q [Dec]
createSortInclusionInfer :: Name -> Name -> Q [Dec]
createSortInclusionInfer Name
from Name
to =
  Name -> Name -> Name -> Q [Dec]
createSortInclusionInfer' Name
from Name
to (Name -> Q [Dec]) -> Name -> Q [Dec]
forall a b. (a -> b) -> a -> b
$ Name -> Name -> Name
sortInclusionName Name
from Name
to

createSortInclusionInfers :: [Name] -> [Name] -> Q [Dec]
createSortInclusionInfers :: [Name] -> [Name] -> Q [Dec]
createSortInclusionInfers [Name]
froms [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)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [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 Name
n = [Char] -> Name
mkName ([Char] -> Name) -> [Char] -> Name
forall a b. (a -> b) -> a -> b
$ [Char]
"i" [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Name -> [Char]
nameBase Name
n

fragment :: Name -> String
fragment :: Name -> [Char]
fragment = [Char] -> [Char]
chopL ([Char] -> [Char]) -> (Name -> [Char]) -> Name -> [Char]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Name -> [Char]
nameBase
{-# INLINE fragment #-}

fragmentName :: Name -> Name
fragmentName :: Name -> Name
fragmentName = [Char] -> Name
mkName ([Char] -> Name) -> (Name -> [Char]) -> Name -> Name
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Name -> [Char]
fragment

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

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

-- | Takes something like "type Sig = '[ Add, Mul, Var]" and gives [''Add, ''Mul, ''Var]
sumToNames :: Name -> Q [Name]
sumToNames :: Name -> Q [Name]
sumToNames Name
nm = do
    (TyConI (TySynD Name
_ [TyVarBndr BndrVis]
_ Type
typ)) <- Name -> Q Info
reify Name
nm
    [Name] -> Q [Name]
forall a. a -> Q a
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 Type
t Type
_) = Type -> [Name]
extract Type
t
    extract (AppT (AppT Type
PromotedConsT (ConT Name
n)) Type
res) = Name
n Name -> [Name] -> [Name]
forall a. a -> [a] -> [a]
: Type -> [Name]
extract Type
res
    extract (AppT (AppT Type
PromotedConsT (SigT (ConT Name
n) Type
_)) Type
res) = Name
n Name -> [Name] -> [Name]
forall a. a -> [a] -> [a]
: Type -> [Name]
extract Type
res
    extract Type
PromotedNilT = []
    extract Type
t = [Char] -> [Name]
forall a. HasCallStack => [Char] -> a
error ([Char] -> [Name]) -> [Char] -> [Name]
forall a b. (a -> b) -> a -> b
$
      [Char]
"sumToNames found invalid summand: " [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<>
      Type -> [Char]
forall a. Show a => a -> [Char]
show Type
t [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<>
      [Char]
"; 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 [Type]
typs Name
className Name
methName 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
$ \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 (Overlap -> Maybe Overlap
forall a. a -> Maybe a
Just Overlap
Overlapping) [] Type
instSig [Dec
dec]