{-# LANGUAGE DataKinds              #-}
{-# LANGUAGE GADTs                  #-}
{-# LANGUAGE PolyKinds              #-}
{-# LANGUAGE RankNTypes             #-}
{-# LANGUAGE ScopedTypeVariables    #-}
{-# LANGUAGE StandaloneDeriving     #-}
{-# LANGUAGE TypeApplications       #-}
{-# LANGUAGE TypeFamilies           #-}
{-# LANGUAGE TypeFamilyDependencies #-}
{-# LANGUAGE TypeInType             #-}
{-# LANGUAGE TypeOperators          #-}
{-# LANGUAGE UndecidableInstances   #-}

-- | Testing dependent functions
module Test.DepFn (
    -- * Infrastructure: singletons
    Sing(..)
  , SingI(..)
  , SomeSing(..)
  , withSomeSing
  , SingKind(..)
    -- * Infrastructure: generic
  , All
    -- * Testing dependent functions
  , DepArgs(..)
  , DepFn(..)
  ) where

import           Data.Kind (Constraint, Type)
import           Test.QuickCheck

{-------------------------------------------------------------------------------
  Excerpt from singletons
-------------------------------------------------------------------------------}

data family Sing (a :: k)

-- | Singletons for lists
--
-- NOTE: Unlike the singletons library, we do /not/ require instances for the
-- elements of the list.
data instance Sing (xs :: [k]) where
  SNil  :: Sing '[]
  SCons :: Sing xs -> Sing (x ': xs)

class SingI (a :: k) where
  sing :: Sing a

instance             SingI '[]       where sing = SNil
instance SingI xs => SingI (x ': xs) where sing = SCons sing

data SomeSing k where
  SomeSing :: Sing (a :: k) -> SomeSing k

withSomeSing :: forall k r. SingKind k
             => Demote k
             -> (forall (a :: k). Sing a -> r)
             -> r
withSomeSing d k = case toSing @k d :: SomeSing k of
                     SomeSing s -> k s

class SingKind k where
  type Demote k = (r :: Type) | r -> k

  fromSing :: Sing (a :: k) -> Demote k
  toSing :: Demote k -> SomeSing k

{-------------------------------------------------------------------------------
  Generic auxiliary
-------------------------------------------------------------------------------}

type family All (f :: k -> Constraint) (xs :: [k]) :: Constraint where
  All _ '[]       = ()
  All f (x ': xs) = (f x, All f xs)

{-------------------------------------------------------------------------------
  'Testable' support for dependently typed functions
-------------------------------------------------------------------------------}

data DepArgs :: [k -> Type] -> k -> Type where
  Nil  :: DepArgs '[] t
  (:*) :: f t -> DepArgs fs t -> DepArgs (f ': fs) t

infixr :*

newtype DepFn (fs :: [k -> Type]) a = DepFn {
    applyDepFn :: forall (t :: k). Sing (t :: k) -> DepArgs fs t -> a
  }

class SingArbitrary (f :: k -> Type) where
  singArbitrary :: Sing t -> Gen (f t)

class SingShow (f :: k -> Type) where
  singShow :: Sing t -> f t -> String

instance (All SingArbitrary fs, SingI fs) => SingArbitrary (DepArgs fs) where
  singArbitrary (st :: Sing t) = go sing
    where
      go :: All SingArbitrary fs' => Sing fs' -> Gen (DepArgs fs' t)
      go SNil       = return Nil
      go (SCons ss) = (:*) <$> singArbitrary st <*> go ss

instance All SingShow fs => SingShow (DepArgs fs) where
  singShow _              Nil       = "[]"
  singShow (st :: Sing t) (x :* xs) = "[" ++ singShow st x ++ go xs
    where
      go :: All SingShow fs' => DepArgs fs' t -> String
      go Nil         = "]"
      go (x' :* xs') = "," ++ singShow st x' ++ go xs'

instance ( SingKind k
         , Arbitrary (Demote k)
         , All SingArbitrary fs
         , All SingShow fs
         , SingI fs
         , Testable a
         ) => Testable (DepFn (fs :: [k -> Type]) a) where
  property (DepFn f) = property $ do
    t :: Demote k <- arbitrary
    withSomeSing t $ \st -> do
      args <- singArbitrary st
      return $ counterexample (singShow st args) $ property (f st args)

{-------------------------------------------------------------------------------
  Example
-------------------------------------------------------------------------------}

data Univ = Two | Many

instance Arbitrary Univ where
  arbitrary = elements [Two, Many]

data instance Sing (k :: Univ) where
  SingTwo  :: Sing ('Two  :: Univ)
  SingMany :: Sing ('Many :: Univ)

instance SingI ('Two  :: Univ) where sing = SingTwo
instance SingI ('Many :: Univ) where sing = SingMany

instance SingKind Univ where
  type Demote Univ = Univ

  fromSing SingTwo  = Two
  fromSing SingMany = Many

  toSing Two  = SomeSing SingTwo
  toSing Many = SomeSing SingMany

data Val :: Univ -> Type where
  ValTwo  :: Bool -> Val 'Two
  ValMany :: Int  -> Val 'Many

deriving instance Show (Val t)

instance Eq (Val t) where
  ValTwo  x == ValTwo  y = x == y
  ValMany x == ValMany y = x == y

instance SingShow Val where
  singShow _ = show

instance SingArbitrary Val where
  singArbitrary SingTwo  = ValTwo  <$> arbitrary
  singArbitrary SingMany = ValMany <$> arbitrary

prop_example :: DepFn '[Val, Val] Bool
prop_example = DepFn $ \_sing (x :* y :* Nil) -> x == y

_example :: IO ()
_example = quickCheck prop_example