{- Data/Singletons/Single/Data.hs

(c) Richard Eisenberg 2013
rae@cs.brynmawr.edu

Singletonizes constructors.
-}

{-# LANGUAGE ParallelListComp, TupleSections, LambdaCase #-}

module Data.Singletons.Single.Data where

import Language.Haskell.TH.Desugar
import Language.Haskell.TH.Syntax
import Data.Singletons.Single.Defun
import Data.Singletons.Single.Monad
import Data.Singletons.Single.Type
import Data.Singletons.Single.Fixity
import Data.Singletons.Promote.Type
import Data.Singletons.Util
import Data.Singletons.Names
import Data.Singletons.Syntax
import Control.Monad
import qualified Data.Set as Set
import Data.Set (Set)

-- We wish to consider the promotion of "Rep" to be *
-- not a promoted data constructor.
singDataD :: DataDecl -> SgM [DDec]
singDataD (DataDecl name tvbs ctors) = do
  let tvbNames = map extractTvbName tvbs
  k <- promoteType (foldType (DConT name) (map DVarT tvbNames))
  ctors' <- mapM singCtor ctors
  ctorFixities <-
    -- try to reify the fixity declarations for the constructors and then
    -- singletonize them. In case the reification fails, we default to an
    -- empty list of singletonized fixity declarations.
    -- why this works:
    -- 1. if we're in a call to 'genSingletons', the data type was defined
    --    earlier and its constructors are in scope, the reification succeeds.
    -- 2. if we're in a call to 'singletons', the reification will fail, but
    --    the fixity declaration will get singletonized by itself (not from
    --    here, look for other invocations of 'singInfixDecl')
    singFixityDeclarations [ n | DCon _ _ n _ _ <- ctors ]
  -- instance for SingKind
  fromSingClauses     <- mapM mkFromSingClause ctors
  emptyFromSingClause <- mkEmptyFromSingClause
  toSingClauses       <- mapM mkToSingClause ctors
  emptyToSingClause   <- mkEmptyToSingClause
  let singKindInst =
        DInstanceD Nothing
                   (map (singKindConstraint . DVarT) tvbNames)
                   (DAppT (DConT singKindClassName) k)
                   [ DTySynInstD demoteName $ DTySynEqn
                      [k]
                      (foldType (DConT name)
                        (map (DAppT demote . DVarT) tvbNames))
                   , DLetDec $ DFunD fromSingName
                               (fromSingClauses `orIfEmpty` [emptyFromSingClause])
                   , DLetDec $ DFunD toSingName
                               (toSingClauses   `orIfEmpty` [emptyToSingClause]) ]

  -- e.g. type SNat = (Sing :: Nat -> Type)
  let kindedSingTy = DArrowT `DAppT` k `DAppT` DConT typeKindName
      kindedSynInst =
        DTySynD (singTyConName name)
                []
                (singFamily `DSigT` kindedSingTy)

  return $ (DDataInstD Data [] singFamilyName [] (Just kindedSingTy) ctors' []) :
           kindedSynInst :
           singKindInst :
           ctorFixities
  where -- in the Rep case, the names of the constructors are in the wrong scope
        -- (they're types, not datacons), so we have to reinterpret them.
        mkConName :: Name -> SgM Name
        mkConName
          | nameBase name == nameBase repName = mkDataName . nameBase
          | otherwise                         = return

        mkFromSingClause :: DCon -> SgM DClause
        mkFromSingClause c = do
          let (cname, numArgs) = extractNameArgs c
          cname' <- mkConName cname
          varNames <- replicateM numArgs (qNewName "b")
          return $ DClause [DConPa (singDataConName cname) (map DVarPa varNames)]
                           (foldExp
                              (DConE cname')
                              (map (DAppE (DVarE fromSingName) . DVarE) varNames))

        mkToSingClause :: DCon -> SgM DClause
        mkToSingClause (DCon _tvbs _cxt cname fields _rty) = do
          let types = tysOfConFields fields
          varNames  <- mapM (const $ qNewName "b") types
          svarNames <- mapM (const $ qNewName "c") types
          promoted  <- mapM promoteType types
          cname' <- mkConName cname
          let varPats        = zipWith mkToSingVarPat varNames promoted
              recursiveCalls = zipWith mkRecursiveCall varNames promoted
          return $
            DClause [DConPa cname' varPats]
                    (multiCase recursiveCalls
                               (map (DConPa someSingDataName . listify . DVarPa)
                                    svarNames)
                               (DAppE (DConE someSingDataName)
                                         (foldExp (DConE (singDataConName cname))
                                                  (map DVarE svarNames))))

        mkToSingVarPat :: Name -> DKind -> DPat
        mkToSingVarPat varName ki =
          DSigPa (DVarPa varName) (DAppT (DConT demoteName) ki)

        mkRecursiveCall :: Name -> DKind -> DExp
        mkRecursiveCall var_name ki =
          DSigE (DAppE (DVarE toSingName) (DVarE var_name))
                (DAppT (DConT someSingTypeName) ki)

        mkEmptyFromSingClause :: SgM DClause
        mkEmptyFromSingClause = do
          x <- qNewName "x"
          pure $ DClause [DVarPa x]
               $ DCaseE (DVarE x) []

        mkEmptyToSingClause :: SgM DClause
        mkEmptyToSingClause = do
          x <- qNewName "x"
          pure $ DClause [DVarPa x]
               $ DConE someSingDataName `DAppE` DCaseE (DVarE x) []

-- refine a constructor.
singCtor :: DCon -> SgM DCon
 -- polymorphic constructors are handled just
 -- like monomorphic ones -- the polymorphism in
 -- the kind is automatic
singCtor (DCon _tvbs cxt name fields rty)
  | not (null cxt)
  = fail "Singling of constrained constructors not yet supported"
  | otherwise
  = do
  let types = tysOfConFields fields
      sName = singDataConName name
      sCon = DConE sName
      pCon = DConT name
  indexNames <- mapM (const $ qNewName "n") types
  let indices = map DVarT indexNames
  kinds <- mapM promoteType types
  let bound_kvs = foldMap fvDType kinds
  args <- zipWithM (buildArgType bound_kvs) types indices
  rty' <- promoteType rty
  let tvbs = map DPlainTV (Set.toList bound_kvs) ++ zipWith DKindedTV indexNames kinds
      kindedIndices = zipWith DSigT indices kinds

  -- SingI instance for data constructor
  emitDecs
    [DInstanceD Nothing
                (map (DAppPr (DConPr singIName)) indices)
                (DAppT (DConT singIName)
                       (foldType pCon kindedIndices))
                [DLetDec $ DValD (DVarPa singMethName)
                       (foldExp sCon (map (const $ DVarE singMethName) types))]]
  -- SingI instances for defunctionalization symbols. Note that we don't
  -- support contexts in constructors at the moment, so it's fine for now to
  -- just assume that the context is always ().
  emitDecs =<< singDefuns name DataName [] (map Just kinds) (Just rty')

  let noBang    = Bang NoSourceUnpackedness NoSourceStrictness
      conFields = case fields of
                    DNormalC dInfix _ -> DNormalC dInfix $ map (noBang,) args
                    DRecC rec_fields ->
                      DRecC [ (singValName field_name, noBang, arg)
                            | (field_name, _, _) <- rec_fields
                            | arg <- args ]
  return $ DCon tvbs
                []
                sName
                conFields
                (DConT singFamilyName `DAppT` foldType pCon indices)
  where buildArgType :: Set Name -> DType -> DType -> SgM DType
        buildArgType bound_kvs ty index = do
          (ty', _, _, _, _, _) <- singType bound_kvs index ty
          return ty'