{-# LANGUAGE BangPatterns        #-}
{-# LANGUAGE CPP                 #-}
{-# LANGUAGE EmptyCase           #-}
{-# LANGUAGE GADTs               #-}
{-# LANGUAGE KindSignatures      #-}
{-# LANGUAGE LambdaCase          #-}
{-# LANGUAGE PatternSynonyms     #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TemplateHaskell     #-}
{-# LANGUAGE TypeOperators       #-}
{-# LANGUAGE ViewPatterns        #-}
{-# OPTIONS_HADDOCK hide #-}
-- |
-- Module      : Data.Array.Accelerate.AST.Idx
-- Copyright   : [2008..2020] The Accelerate Team
-- License     : BSD3
--
-- Maintainer  : Trevor L. McDonell <trevor.mcdonell@gmail.com>
-- Stability   : experimental
-- Portability : non-portable (GHC extensions)
--
-- Typed de Bruijn indices
--

module Data.Array.Accelerate.AST.Idx (

  Idx, pattern ZeroIdx, pattern SuccIdx, pattern VoidIdx,
  idxToInt,
  rnfIdx, liftIdx,

  PairIdx(..)

) where

import Data.Kind
import Language.Haskell.TH.Extra                                    hiding ( Type )

#ifndef ACCELERATE_INTERNAL_CHECKS
import Data.Type.Equality                                           ( (:~:)(Refl) )
import Unsafe.Coerce                                                ( unsafeCoerce )
#endif


#ifdef ACCELERATE_INTERNAL_CHECKS

-- | De Bruijn variable index projecting a specific type from a type
-- environment.  Type environments are nested pairs (..((), t1), t2, ..., tn).
--
data Idx env t where
  ZeroIdx ::              Idx (env, t) t
  SuccIdx :: Idx env t -> Idx (env, s) t

idxToInt :: Idx env t -> Int
idxToInt ZeroIdx       = 0
idxToInt (SuccIdx idx) = 1 + idxToInt idx

rnfIdx :: Idx env t -> ()
rnfIdx ZeroIdx      = ()
rnfIdx (SuccIdx ix) = rnfIdx ix

liftIdx :: Idx env t -> CodeQ (Idx env t)
liftIdx ZeroIdx      = [|| ZeroIdx ||]
liftIdx (SuccIdx ix) = [|| SuccIdx $$(liftIdx ix) ||]

#else

-- | De Bruijn variable index projecting a specific type from a type
-- environment.  Type environments are nested pairs (..((), t1), t2, ..., tn).
--
-- Outside of this file, pretend that this is an ordinary GADT:
-- data Idx env t where
--   ZeroIdx ::              Idx (env, t) t
--   SuccIdx :: Idx env t -> Idx (env, s) t
--
-- For performance, it uses an Int under the hood.
--
newtype Idx :: Type -> Type -> Type where
  UnsafeIdxConstructor :: { forall env t. Idx env t -> Int
unsafeRunIdx :: Int } -> Idx env t

{-# COMPLETE ZeroIdx, SuccIdx #-}

pattern ZeroIdx :: forall envt t. () => forall env. (envt ~ (env, t)) => Idx envt t
pattern $mZeroIdx :: forall {r} {envt} {t}.
Idx envt t
-> (forall {env}. (envt ~ (env, t)) => r) -> ((# #) -> r) -> r
$bZeroIdx :: forall envt t env. (envt ~ (env, t)) => Idx envt t
ZeroIdx <- (\Idx envt t
x -> (Idx envt t -> Int
forall env t. Idx env t -> Int
idxToInt Idx envt t
x, (Any :~: Any) -> envt :~: (Any, t)
forall a b. a -> b
unsafeCoerce Any :~: Any
forall {k} (a :: k). a :~: a
Refl) -> (0, Refl :: envt :~: (env, t)))
  where
    ZeroIdx = Int -> Idx envt t
forall env t. Int -> Idx env t
UnsafeIdxConstructor Int
0

pattern SuccIdx :: forall envs t. () => forall s env. (envs ~ (env, s)) => Idx env t -> Idx envs t
pattern $mSuccIdx :: forall {r} {envs} {t}.
Idx envs t
-> (forall {s} {env}. (envs ~ (env, s)) => Idx env t -> r)
-> ((# #) -> r)
-> r
$bSuccIdx :: forall envs t s env. (envs ~ (env, s)) => Idx env t -> Idx envs t
SuccIdx idx <- (unSucc -> Just (idx, Refl))
  where
    SuccIdx (UnsafeIdxConstructor Int
i) = Int -> Idx envs t
forall env t. Int -> Idx env t
UnsafeIdxConstructor (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)

unSucc :: Idx envs t -> Maybe (Idx env t, envs :~: (env, s))
unSucc :: forall envs t env s.
Idx envs t -> Maybe (Idx env t, envs :~: (env, s))
unSucc (UnsafeIdxConstructor Int
i)
  | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
1     = Maybe (Idx env t, envs :~: (env, s))
forall a. Maybe a
Nothing
  | Bool
otherwise = (Idx env t, envs :~: (env, s))
-> Maybe (Idx env t, envs :~: (env, s))
forall a. a -> Maybe a
Just (Int -> Idx env t
forall env t. Int -> Idx env t
UnsafeIdxConstructor (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1), (Any :~: Any) -> envs :~: (env, s)
forall a b. a -> b
unsafeCoerce Any :~: Any
forall {k} (a :: k). a :~: a
Refl)

idxToInt :: Idx env t -> Int
idxToInt :: forall env t. Idx env t -> Int
idxToInt = Idx env t -> Int
forall env t. Idx env t -> Int
unsafeRunIdx

rnfIdx :: Idx env t -> ()
rnfIdx :: forall env t. Idx env t -> ()
rnfIdx !Idx env t
_ = ()

liftIdx :: Idx env t -> CodeQ (Idx env t)
liftIdx :: forall env t. Idx env t -> CodeQ (Idx env t)
liftIdx (UnsafeIdxConstructor Int
i) = [|| Int -> Idx env t
forall env t. Int -> Idx env t
UnsafeIdxConstructor Int
i ||]

#endif

-- | Despite the 'complete' pragma above, GHC can't infer that there is no
-- pattern possible if the environment is empty. This can be used instead.
--
pattern VoidIdx :: forall env t a. (env ~ ()) => () => a -> Idx env t
pattern $mVoidIdx :: forall {r} {env} {t} {a}.
(env ~ ()) =>
Idx env t -> (a -> r) -> ((# #) -> r) -> r
VoidIdx a <- (\case{} -> a)

{-# COMPLETE VoidIdx #-}

data PairIdx p a where
  PairIdxLeft  :: PairIdx (a, b) a
  PairIdxRight :: PairIdx (a, b) b