-----------------------------------------------------------------------------
-- |
-- Module      :  Data.Singletons.TH.Single.Defun
-- Copyright   :  (C) 2018 Ryan Scott
-- License     :  BSD-style (see LICENSE)
-- Maintainer  :  Ryan Scott
-- Stability   :  experimental
-- Portability :  non-portable
--
-- Creates 'SingI' instances for promoted types' defunctionalization symbols.
--
-----------------------------------------------------------------------------

module Data.Singletons.TH.Single.Defun (singDefuns) where

import Control.Monad
import Data.Foldable
import Data.Singletons.TH.Names
import Data.Singletons.TH.Options
import Data.Singletons.TH.Promote.Defun
import Data.Singletons.TH.Single.Monad
import Data.Singletons.TH.Single.Type
import Data.Singletons.TH.Util
import Language.Haskell.TH.Desugar
import Language.Haskell.TH.Syntax

-- Given the Name of something, take the defunctionalization symbols for its
-- promoted counterpart and create SingI instances for them. As a concrete
-- example, if you have:
--
--   foo :: Eq a => a -> a -> Bool
--
-- Then foo's promoted counterpart, Foo, will have two defunctionalization
-- symbols:
--
--   FooSym0 :: a ~> a ~> Bool
--   FooSym1 :: a -> a ~> Bool
--
-- We can declare SingI instances for these two symbols like so:
--
--   instance SEq a => SingI (FooSym0 :: a ~> a ~> Bool) where
--     sing = singFun2 sFoo
--
--   instance (SEq a, SingI x) => SingI (FooSym1 x :: a ~> Bool) where
--     sing = singFun1 (sFoo (sing @_ @x))
--
-- Note that singDefuns takes Maybe DKinds for the promoted argument and result
-- types, in case we have an entity whose type needs to be inferred.
-- See Note [singDefuns and type inference].
singDefuns :: Name      -- The Name of the thing to promote.
           -> NameSpace -- Whether the above Name is a value, data constructor,
                        -- or a type constructor.
           -> DCxt      -- The type's context.
           -> [Maybe DKind] -- The promoted argument types (if known).
           -> Maybe DKind   -- The promoted result type (if known).
           -> SgM [DDec]
singDefuns :: Name
-> NameSpace -> DCxt -> [Maybe DKind] -> Maybe DKind -> SgM [DDec]
singDefuns Name
n NameSpace
ns DCxt
ty_ctxt [Maybe DKind]
mb_ty_args Maybe DKind
mb_ty_res =
  case [Maybe DKind]
mb_ty_args of
    [] -> [DDec] -> SgM [DDec]
forall (f :: * -> *) a. Applicative f => a -> f a
pure [] -- If a function has no arguments, then it has no
                  -- defunctionalization symbols, so there's nothing to be done.
    [Maybe DKind]
_  -> do Options
opts     <- SgM Options
forall (m :: * -> *). OptionsMonad m => m Options
getOptions
             DCxt
sty_ctxt <- (DKind -> SgM DKind) -> DCxt -> SgM DCxt
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM DKind -> SgM DKind
singPred DCxt
ty_ctxt
             [Name]
names    <- Int -> SgM Name -> SgM [Name]
forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM ([Maybe DKind] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Maybe DKind]
mb_ty_args) (SgM Name -> SgM [Name]) -> SgM Name -> SgM [Name]
forall a b. (a -> b) -> a -> b
$ String -> SgM Name
forall (m :: * -> *). Quasi m => String -> m Name
qNewName String
"d"
             let tvbs :: [DTyVarBndrUnit]
tvbs       = (Name -> Maybe DKind -> DTyVarBndrUnit)
-> [Name] -> [Maybe DKind] -> [DTyVarBndrUnit]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Name -> Maybe DKind -> DTyVarBndrUnit
inferMaybeKindTV [Name]
names [Maybe DKind]
mb_ty_args
                 (Maybe DKind
_, [DDec]
insts) = Options
-> Int
-> DCxt
-> [DTyVarBndrUnit]
-> [DTyVarBndrUnit]
-> (Maybe DKind, [DDec])
go Options
opts Int
0 DCxt
sty_ctxt [] [DTyVarBndrUnit]
tvbs
             [DDec] -> SgM [DDec]
forall (f :: * -> *) a. Applicative f => a -> f a
pure [DDec]
insts
  where
    num_ty_args :: Int
    num_ty_args :: Int
num_ty_args = [Maybe DKind] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Maybe DKind]
mb_ty_args

    -- The inner loop. @go n ctxt arg_tvbs res_tvbs@ returns @(m_result, insts)@.
    -- Using one particular example:
    --
    -- @
    -- instance (SingI a, SingI b, SEq c, SEq d) =>
    --   SingI (ExampleSym2 (x :: a) (y :: b) :: c ~> d ~> Type) where ...
    -- @
    --
    -- We have:
    --
    -- * @n@ is 2. This is incremented in each iteration of `go`.
    --
    -- * @ctxt@ is (SEq c, SEq d). The (SingI a, SingI b) part of the instance
    --   context is added separately.
    --
    -- * @arg_tvbs@ is [(x :: a), (y :: b)].
    --
    -- * @res_tvbs@ is [(z :: c), (w :: d)]. The kinds of these type variable
    --   binders appear in the result kind.
    --
    -- * @m_result@ is `Just (c ~> d ~> Type)`. @m_result@ is returned so
    --   that earlier defunctionalization symbols can build on the result
    --   kinds of later symbols. For instance, ExampleSym1 would get the
    --   result kind `b ~> c ~> d ~> Type` by prepending `b` to ExampleSym2's
    --   result kind `c ~> d ~> Type`.
    --
    -- * @insts@ are all of the instance declarations corresponding to
    --   ExampleSym2 and later defunctionalization symbols. This is the main
    --   payload of the function.
    --
    -- This function is quadratic because it appends a variable at the end of
    -- the @arg_tvbs@ list at each iteration. In practice, this is unlikely
    -- to be a performance bottleneck since the number of arguments rarely
    -- gets to be that large.
    go :: Options -> Int -> DCxt -> [DTyVarBndrUnit] -> [DTyVarBndrUnit]
       -> (Maybe DKind, [DDec])
    go :: Options
-> Int
-> DCxt
-> [DTyVarBndrUnit]
-> [DTyVarBndrUnit]
-> (Maybe DKind, [DDec])
go Options
_    Int
_       DCxt
_        [DTyVarBndrUnit]
_        []                 = (Maybe DKind
mb_ty_res, [])
    go Options
opts Int
sym_num DCxt
sty_ctxt [DTyVarBndrUnit]
arg_tvbs (DTyVarBndrUnit
res_tvb:[DTyVarBndrUnit]
res_tvbs) =
      (Maybe DKind
mb_new_res, DDec
new_instDDec -> [DDec] -> [DDec]
forall a. a -> [a] -> [a]
:[DDec]
insts)
      where
        mb_res :: Maybe DKind
        insts  :: [DDec]
        (Maybe DKind
mb_res, [DDec]
insts) = Options
-> Int
-> DCxt
-> [DTyVarBndrUnit]
-> [DTyVarBndrUnit]
-> (Maybe DKind, [DDec])
go Options
opts (Int
sym_num Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) DCxt
sty_ctxt ([DTyVarBndrUnit]
arg_tvbs [DTyVarBndrUnit] -> [DTyVarBndrUnit] -> [DTyVarBndrUnit]
forall a. [a] -> [a] -> [a]
++ [DTyVarBndrUnit
res_tvb]) [DTyVarBndrUnit]
res_tvbs

        mb_new_res :: Maybe DKind
        mb_new_res :: Maybe DKind
mb_new_res = DTyVarBndrUnit -> Maybe DKind -> Maybe DKind
mk_inst_kind DTyVarBndrUnit
res_tvb Maybe DKind
mb_res

        sing_fun_num :: Int
        sing_fun_num :: Int
sing_fun_num = Int
num_ty_args Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
sym_num

        mk_sing_fun_expr :: DExp -> DExp
        mk_sing_fun_expr :: DExp -> DExp
mk_sing_fun_expr DExp
sing_expr =
          (DExp -> Name -> DExp) -> DExp -> [Name] -> DExp
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\DExp
f Name
tvb_n -> DExp
f DExp -> DExp -> DExp
`DAppE` (Name -> DExp
DVarE Name
singMethName DExp -> DKind -> DExp
`DAppTypeE` Name -> DKind
DVarT Name
tvb_n))
                 DExp
sing_expr
                 ((DTyVarBndrUnit -> Name) -> [DTyVarBndrUnit] -> [Name]
forall a b. (a -> b) -> [a] -> [b]
map DTyVarBndrUnit -> Name
forall flag. DTyVarBndr flag -> Name
extractTvbName [DTyVarBndrUnit]
arg_tvbs)

        singI_ctxt :: DCxt
        singI_ctxt :: DCxt
singI_ctxt = (DTyVarBndrUnit -> DKind) -> [DTyVarBndrUnit] -> DCxt
forall a b. (a -> b) -> [a] -> [b]
map (DKind -> DKind -> DKind
DAppT (Name -> DKind
DConT Name
singIName) (DKind -> DKind)
-> (DTyVarBndrUnit -> DKind) -> DTyVarBndrUnit -> DKind
forall b c a. (b -> c) -> (a -> b) -> a -> c
. DTyVarBndrUnit -> DKind
forall flag. DTyVarBndr flag -> DKind
tvbToType) [DTyVarBndrUnit]
arg_tvbs

        mk_inst_ty :: DType -> DType
        mk_inst_ty :: DKind -> DKind
mk_inst_ty DKind
inst_head
          = case Maybe DKind
mb_new_res of
              Just DKind
inst_kind -> DKind
inst_head DKind -> DKind -> DKind
`DSigT` DKind
inst_kind
              Maybe DKind
Nothing        -> DKind
inst_head

        arg_tvb_tys :: [DType]
        arg_tvb_tys :: DCxt
arg_tvb_tys = (DTyVarBndrUnit -> DKind) -> [DTyVarBndrUnit] -> DCxt
forall a b. (a -> b) -> [a] -> [b]
map DTyVarBndrUnit -> DKind
forall flag. DTyVarBndr flag -> DKind
dTyVarBndrToDType [DTyVarBndrUnit]
arg_tvbs

        -- Construct the arrow kind used to annotate the defunctionalization
        -- symbol (e.g., the `a ~> a ~> Bool` in
        -- `SingI (FooSym0 :: a ~> a ~> Bool)`).
        -- If any of the argument kinds or result kind isn't known (i.e., is
        -- Nothing), then we opt not to construct this arrow kind altogether.
        -- See Note [singDefuns and type inference]
        mk_inst_kind :: DTyVarBndrUnit -> Maybe DKind -> Maybe DKind
        mk_inst_kind :: DTyVarBndrUnit -> Maybe DKind -> Maybe DKind
mk_inst_kind DTyVarBndrUnit
tvb' = Maybe DKind -> Maybe DKind -> Maybe DKind
buildTyFunArrow_maybe (DTyVarBndrUnit -> Maybe DKind
forall flag. DTyVarBndr flag -> Maybe DKind
extractTvbKind DTyVarBndrUnit
tvb')

        new_inst :: DDec
        new_inst :: DDec
new_inst = Maybe Overlap
-> Maybe [DTyVarBndrUnit] -> DCxt -> DKind -> [DDec] -> DDec
DInstanceD Maybe Overlap
forall a. Maybe a
Nothing Maybe [DTyVarBndrUnit]
forall a. Maybe a
Nothing
                              (DCxt
sty_ctxt DCxt -> DCxt -> DCxt
forall a. [a] -> [a] -> [a]
++ DCxt
singI_ctxt)
                              (Name -> DKind
DConT Name
singIName DKind -> DKind -> DKind
`DAppT` DKind -> DKind
mk_inst_ty DKind
defun_inst_ty)
                              [DLetDec -> DDec
DLetDec (DLetDec -> DDec) -> DLetDec -> DDec
forall a b. (a -> b) -> a -> b
$ DPat -> DExp -> DLetDec
DValD (Name -> DPat
DVarP Name
singMethName)
                                       (DExp -> DLetDec) -> DExp -> DLetDec
forall a b. (a -> b) -> a -> b
$ Int -> DKind -> DExp -> DExp
wrapSingFun Int
sing_fun_num DKind
defun_inst_ty
                                       (DExp -> DExp) -> DExp -> DExp
forall a b. (a -> b) -> a -> b
$ DExp -> DExp
mk_sing_fun_expr DExp
sing_exp ]
          where
            defun_inst_ty :: DType
            defun_inst_ty :: DKind
defun_inst_ty = DKind -> DCxt -> DKind
foldType (Name -> DKind
DConT (Options -> Name -> Int -> Name
defunctionalizedName Options
opts Name
n Int
sym_num))
                                     DCxt
arg_tvb_tys

            sing_exp :: DExp
            sing_exp :: DExp
sing_exp = case NameSpace
ns of
                         NameSpace
DataName -> Name -> DExp
DConE (Name -> DExp) -> Name -> DExp
forall a b. (a -> b) -> a -> b
$ Options -> Name -> Name
singledDataConName Options
opts Name
n
                         NameSpace
_        -> Name -> DExp
DVarE (Name -> DExp) -> Name -> DExp
forall a b. (a -> b) -> a -> b
$ Options -> Name -> Name
singledValueName Options
opts Name
n

{-
Note [singDefuns and type inference]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Consider the following function:

  foo :: a -> Bool
  foo _ = True

singDefuns would give the following SingI instance for FooSym0, with an
explicit kind signature:

  instance SingI (FooSym0 :: a ~> Bool) where ...

What happens if we leave off the type signature for foo?

  foo _ = True

Can singDefuns still do its job? Yes! It will simply generate:

  instance SingI FooSym0 where ...

In general, if any of the promoted argument or result types given to singDefun
are Nothing, then we avoid crafting an explicit kind signature. You might worry
that this could lead to SingI instances being generated that GHC cannot infer
the type for, such as:

  bar x = x == x
  ==>
  instance SingI BarSym0 -- Missing an SEq constraint?

This is true, but also not unprecedented, as the singled version of bar, sBar,
will /also/ fail to typecheck due to a missing SEq constraint. Therefore, this
design choice fits within the existing tradition of type inference in
singletons-th.
-}