-----------------------------------------------------------------------------
-- |
-- Module      :  Data.Singletons.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.Single.Defun (singDefuns) where

import Data.List
import Data.Singletons.Names
import Data.Singletons.Promote.Defun
import Data.Singletons.Single.Monad
import Data.Singletons.Single.Type
import Data.Singletons.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 n :: Name
n ns :: NameSpace
ns ty_ctxt :: DCxt
ty_ctxt mb_ty_args :: [Maybe DKind]
mb_ty_args mb_ty_res :: 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.
    _  -> do 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
             Int -> DCxt -> [DTyVarBndr] -> [Maybe DKind] -> SgM [DDec]
go 0 DCxt
sty_ctxt [] [Maybe DKind]
mb_ty_args
  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

    -- Sadly, this algorithm is quadratic, because in each iteration of the loop
    -- we must:
    --
    -- * Construct an arrow type of the form (a ~> ... ~> z), using a suffix of
    --   the promoted argument types.
    -- * Append a new type variable to the end of an ordered list.
    --
    -- In practice, this is unlikely to be a bottleneck, as singletons does not
    -- support functions with more than 7 or so arguments anyways.
    go :: Int -> DCxt -> [DTyVarBndr] -> [Maybe DKind] -> SgM [DDec]
    go :: Int -> DCxt -> [DTyVarBndr] -> [Maybe DKind] -> SgM [DDec]
go sym_num :: Int
sym_num sty_ctxt :: DCxt
sty_ctxt tvbs :: [DTyVarBndr]
tvbs mb_tyss :: [Maybe DKind]
mb_tyss
      | Int
sym_num Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
num_ty_args
      , mb_ty :: Maybe DKind
mb_ty:mb_tys :: [Maybe DKind]
mb_tys <- [Maybe DKind]
mb_tyss
      = do Name
new_tvb_name <- String -> SgM Name
forall (m :: * -> *). Quasi m => String -> m Name
qNewName "d"
           let new_tvb :: DTyVarBndr
new_tvb = Name -> Maybe DKind -> DTyVarBndr
inferMaybeKindTV Name
new_tvb_name Maybe DKind
mb_ty
           [DDec]
insts <- Int -> DCxt -> [DTyVarBndr] -> [Maybe DKind] -> SgM [DDec]
go (Int
sym_num Int -> Int -> Int
forall a. Num a => a -> a -> a
+ 1) DCxt
sty_ctxt ([DTyVarBndr]
tvbs [DTyVarBndr] -> [DTyVarBndr] -> [DTyVarBndr]
forall a. [a] -> [a] -> [a]
++ [DTyVarBndr
new_tvb]) [Maybe DKind]
mb_tys
           [DDec] -> SgM [DDec]
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([DDec] -> SgM [DDec]) -> [DDec] -> SgM [DDec]
forall a b. (a -> b) -> a -> b
$ [DDec]
new_insts [DDec] -> [DDec] -> [DDec]
forall a. [a] -> [a] -> [a]
++ [DDec]
insts
      | Bool
otherwise
      = [DDec] -> SgM [DDec]
forall (f :: * -> *) a. Applicative f => a -> f a
pure []
      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

        mk_sing_fun_expr :: DExp -> DExp
        mk_sing_fun_expr :: DExp -> DExp
mk_sing_fun_expr sing_expr :: DExp
sing_expr =
          (DExp -> Name -> DExp) -> DExp -> [Name] -> DExp
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\f :: DExp
f tvb_n :: 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
                 ((DTyVarBndr -> Name) -> [DTyVarBndr] -> [Name]
forall a b. (a -> b) -> [a] -> [b]
map DTyVarBndr -> Name
extractTvbName [DTyVarBndr]
tvbs)

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

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

        tvb_tys :: [DType]
        tvb_tys :: DCxt
tvb_tys = (DTyVarBndr -> DKind) -> [DTyVarBndr] -> DCxt
forall a b. (a -> b) -> [a] -> [b]
map DTyVarBndr -> DKind
dTyVarBndrToDType [DTyVarBndr]
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]
        mb_inst_kind :: Maybe DType
        mb_inst_kind :: Maybe DKind
mb_inst_kind = (Maybe DKind -> Maybe DKind -> Maybe DKind)
-> Maybe DKind -> [Maybe DKind] -> Maybe DKind
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr Maybe DKind -> Maybe DKind -> Maybe DKind
buildTyFunArrow_maybe Maybe DKind
mb_ty_res [Maybe DKind]
mb_tyss

        new_insts :: [DDec]
        new_insts :: [DDec]
new_insts = [Maybe Overlap
-> Maybe [DTyVarBndr] -> DCxt -> DKind -> [DDec] -> DDec
DInstanceD Maybe Overlap
forall a. Maybe a
Nothing Maybe [DTyVarBndr]
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 (Name -> Int -> Name
promoteTySym Name
n Int
sym_num)) DCxt
tvb_tys

            sing_exp :: DExp
            sing_exp :: DExp
sing_exp = case NameSpace
ns of
                         DataName -> Name -> DExp
DConE (Name -> DExp) -> Name -> DExp
forall a b. (a -> b) -> a -> b
$ Name -> Name
singDataConName Name
n
                         _        -> Name -> DExp
DVarE (Name -> DExp) -> Name -> DExp
forall a b. (a -> b) -> a -> b
$ Name -> Name
singValName 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.
-}