-----------------------------------------------------------------------------
-- |
-- 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.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{,1,2} 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 and SingI1 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))
--
--   instance SEq a => SingI1 (FooSym1 :: a -> a ~> Bool) where
--     liftSing s = singFun1 (sFoo s)
--
-- Only FooSym1 will have a SingI1 instance, as unlike FooSym0, it is able to
-- be partially applied (using normal function application) to a single
-- argument. Neither FooSym0 nor FooSym1 can be partially applied to two
-- arguments, so neither will receive a SingI2 instance.
--
-- 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 a. a -> SgM a
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 opts     <- SgM Options
forall (m :: * -> *). OptionsMonad m => m Options
getOptions
             sty_ctxt <- mapM singPred ty_ctxt
             names    <- replicateM (length mb_ty_args) $ qNewName "d"
             let 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
             (_, insts) <- go opts 0 sty_ctxt [] tvbs
             pure insts
  where
    num_ty_args :: Int
    num_ty_args :: Int
num_ty_args = [Maybe DKind] -> Int
forall a. [a] -> 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]
       -> SgM (Maybe DKind, [DDec])
    go :: Options
-> Int
-> DCxt
-> [DTyVarBndrUnit]
-> [DTyVarBndrUnit]
-> SgM (Maybe DKind, [DDec])
go Options
_    Int
_       DCxt
_        [DTyVarBndrUnit]
_        []                 = (Maybe DKind, [DDec]) -> SgM (Maybe DKind, [DDec])
forall a. a -> SgM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe DKind
mb_ty_res, [])
    go Options
opts Int
sym_num DCxt
sty_ctxt [DTyVarBndrUnit]
arg_tvbs (DTyVarBndrUnit
res_tvb:[DTyVarBndrUnit]
res_tvbs) = do
      (mb_res, insts) <- Options
-> Int
-> DCxt
-> [DTyVarBndrUnit]
-> [DTyVarBndrUnit]
-> SgM (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
      new_insts <- mapMaybeM (mb_new_inst mb_res) [0, 1, 2]
      pure (mk_inst_kind [] res_tvb mb_res, new_insts ++ insts)
      where
        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

        -- Construct the arrow kind used to annotate the defunctionalization
        -- symbol. For example, this constructs the `a -> b -> c ~> Bool` in
        -- `SingI1 (FooSym1 :: a -> b -> c ~> Bool)`, where:
        --
        -- * The first argument to `mk_inst_kind` gives the kinds [a, b], which
        --   are used with normal function arrows.
        -- * The second argumen to `mk_inst_kind` gives the kind `c`, which is
        --   used with a defunctionalized function arrow.
        --
        -- 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] -> DTyVarBndrUnit -> Maybe DKind -> Maybe DKind
        mk_inst_kind :: [DTyVarBndrUnit] -> DTyVarBndrUnit -> Maybe DKind -> Maybe DKind
mk_inst_kind [DTyVarBndrUnit]
funTvbs DTyVarBndrUnit
defunTvb Maybe DKind
mbKind =
          (Maybe DKind -> Maybe DKind -> Maybe DKind)
-> Maybe DKind -> [Maybe DKind] -> Maybe DKind
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr Maybe DKind -> Maybe DKind -> Maybe DKind
buildFunArrow_maybe
                (Maybe DKind -> Maybe DKind -> Maybe DKind
buildTyFunArrow_maybe (DTyVarBndrUnit -> Maybe DKind
forall flag. DTyVarBndr flag -> Maybe DKind
extractTvbKind DTyVarBndrUnit
defunTvb) Maybe DKind
mbKind)
                ((DTyVarBndrUnit -> Maybe DKind)
-> [DTyVarBndrUnit] -> [Maybe DKind]
forall a b. (a -> b) -> [a] -> [b]
map DTyVarBndrUnit -> Maybe DKind
forall flag. DTyVarBndr flag -> Maybe DKind
extractTvbKind [DTyVarBndrUnit]
funTvbs)

        -- @mb_new_inst mb_res k@ returns 'Just' an instance of @SingI<k>@ if
        -- @k@ is less than or equal to the number of arguments to which the
        -- defunctionalization symbol can be partially applied using normal
        -- function application. Otherwise, it returns 'Nothing'.
        mb_new_inst :: Maybe DKind -> Int -> SgM (Maybe DDec)
        mb_new_inst :: Maybe DKind -> Int -> SgM (Maybe DDec)
mb_new_inst Maybe DKind
mb_res Int
k
          | Int
k Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
sym_num
          = do vs <- Int -> SgM Name -> SgM [Name]
forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM Int
k (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
"s"
               let sing_vs = (Name -> DTyVarBndrUnit -> DPat)
-> [Name] -> [DTyVarBndrUnit] -> [DPat]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\Name
v DTyVarBndrUnit
arg_tvb ->
                                       DPat -> DKind -> DPat
DSigP (Name -> DPat
DVarP Name
v)
                                             (DKind
singFamily DKind -> DKind -> DKind
`DAppT` DTyVarBndrUnit -> DKind
forall flag. DTyVarBndr flag -> DKind
dTyVarBndrToDType DTyVarBndrUnit
arg_tvb))
                                     [Name]
vs [DTyVarBndrUnit]
last_arg_tvbs
               pure $ Just $
                 DInstanceD Nothing Nothing
                   (sty_ctxt ++ singI_ctxt)
                   (DConT (mkSingIName k) `DAppT` mk_inst_ty (mk_defun_inst_ty init_arg_tvbs))
                   [ DLetDec $ DFunD (mkSingMethName k)
                      [ DClause sing_vs
                         $ wrapSingFun sing_fun_num (mk_defun_inst_ty arg_tvbs)
                         $ mk_sing_fun_expr sing_exp vs
                      ]
                   ]
          | Bool
otherwise
          = Maybe DDec -> SgM (Maybe DDec)
forall a. a -> SgM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe DDec
forall a. Maybe a
Nothing
          where
            init_arg_tvbs, last_arg_tvbs :: [DTyVarBndrUnit]
            ([DTyVarBndrUnit]
init_arg_tvbs, [DTyVarBndrUnit]
last_arg_tvbs) = Int -> [DTyVarBndrUnit] -> ([DTyVarBndrUnit], [DTyVarBndrUnit])
forall a. Int -> [a] -> ([a], [a])
splitAt (Int
sym_num Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
k) [DTyVarBndrUnit]
arg_tvbs

            mk_defun_inst_ty :: [DTyVarBndrUnit] -> DType
            mk_defun_inst_ty :: [DTyVarBndrUnit] -> DKind
mk_defun_inst_ty [DTyVarBndrUnit]
tvbs =
              DKind -> DCxt -> DKind
foldType (Name -> DKind
DConT (Options -> Name -> Int -> Name
defunctionalizedName Options
opts Name
n Int
sym_num))
                       ((DTyVarBndrUnit -> DKind) -> [DTyVarBndrUnit] -> DCxt
forall a b. (a -> b) -> [a] -> [b]
map DTyVarBndrUnit -> DKind
forall flag. DTyVarBndr flag -> DKind
dTyVarBndrToDType [DTyVarBndrUnit]
tvbs)

            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

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

            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]
init_arg_tvbs

            mk_inst_ty :: DType -> DType
            mk_inst_ty :: DKind -> DKind
mk_inst_ty DKind
inst_head
              = case [DTyVarBndrUnit] -> DTyVarBndrUnit -> Maybe DKind -> Maybe DKind
mk_inst_kind [DTyVarBndrUnit]
last_arg_tvbs DTyVarBndrUnit
res_tvb Maybe DKind
mb_res of
                  Just DKind
inst_kind -> DKind
inst_head DKind -> DKind -> DKind
`DSigT` DKind
inst_kind
                  Maybe DKind
Nothing        -> DKind
inst_head

-- Shorthand for building (k1 -> k2)
buildFunArrow :: DKind -> DKind -> DKind
buildFunArrow :: DKind -> DKind -> DKind
buildFunArrow DKind
k1 DKind
k2 = DKind
DArrowT DKind -> DKind -> DKind
`DAppT` DKind
k1 DKind -> DKind -> DKind
`DAppT` DKind
k2

buildFunArrow_maybe :: Maybe DKind -> Maybe DKind -> Maybe DKind
buildFunArrow_maybe :: Maybe DKind -> Maybe DKind -> Maybe DKind
buildFunArrow_maybe Maybe DKind
m_k1 Maybe DKind
m_k2 = DKind -> DKind -> DKind
buildFunArrow (DKind -> DKind -> DKind) -> Maybe DKind -> Maybe (DKind -> DKind)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe DKind
m_k1 Maybe (DKind -> DKind) -> Maybe DKind -> Maybe DKind
forall a b. Maybe (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Maybe DKind
m_k2

{-
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.
-}