-- SPDX-FileCopyrightText: 2021 Oxhead Alpha
-- SPDX-License-Identifier: LicenseRef-MIT-OA

module Morley.Util.Sing
  ( eqI
  , geqI
  , eqParamSing
  , eqParamSing2
  , eqParamSing3
  , eqParamMixed3
  , castSing
  , SingIOne
  , withSingIOne
  , genSingletonsType
  ) where

import Data.Singletons (KindOf, Sing, SingI, sing)
import Data.Singletons.Decide (SDecide, decideEquality)
import Data.Singletons.TH (genSingletons, singDecideInstance)
import Data.Singletons.TH.Options (Options(..), OptionsMonad, defaultOptions, withOptions)
import Data.Type.Equality (TestEquality(..), (:~:)(..))
import Data.Typeable (eqT)
import Language.Haskell.TH (Dec, Name, mkName, nameBase)

-- | Version of 'testEquality' that uses 'SingI'
eqI :: forall a b. (SingI a, SingI b, TestEquality (Sing @(KindOf a))) => Maybe (a :~: b)
eqI = testEquality (sing @a) (sing @b)

-- | Suppose you have a data type @X@ with parameter @a@ and you have
-- two values: @x1 :: X a1@ and @x2 :: X a2@. You can't compare them
-- using '==', because they have different types. However, you can
-- compare them using 'eqParamSing' as long as both parameters have
-- SingI instances and their kind is `SDecide`.
eqParamSing ::
  forall a1 a2 t.
  ( SingI a1
  , SingI a2
  , SDecide (KindOf a1)
  , Eq (t a1)
  )
  => t a1
  -> t a2
  -> Bool
eqParamSing t1 t2 = isJust @() $ do
  Refl <- sing @a1 `decideEquality` sing @a2
  guard (t1 == t2)

-- | Version of 'eqParamSing' for types with 2 parameters.
eqParamSing2 ::
  forall a1 a2 b1 b2 t.
  ( SingI a1
  , SingI a2
  , SingI b1
  , SingI b2
  , SDecide (KindOf a1)
  , SDecide (KindOf b1)
  , Eq (t a1 b2)
  )
  => t a1 b1
  -> t a2 b2
  -> Bool
eqParamSing2 t1 t2 = isJust @() $ do
  Refl <- sing @a1 `decideEquality` sing @a2
  Refl <- sing @b1 `decideEquality` sing @b2
  guard (t1 == t2)

-- | Version of 'eqParamSing' for types with 3 parameters.
eqParamSing3 ::
  forall a1 a2 b1 b2 c1 c2 t.
  ( SingI a1
  , SingI a2
  , SingI b1
  , SingI b2
  , SingI c1
  , SingI c2
  , SDecide (KindOf a1)
  , SDecide (KindOf b1)
  , SDecide (KindOf c1)
  , Eq (t a1 b1 c1)
  )
  => t a1 b1 c1
  -> t a2 b2 c2
  -> Bool
eqParamSing3 t1 t2 = isJust @() $ do
  Refl <- sing @a1 `decideEquality` sing @a2
  Refl <- sing @b1 `decideEquality` sing @b2
  Refl <- sing @c1 `decideEquality` sing @c2
  guard (t1 == t2)

-- | Version of 'eqParamSing' for types with 3 parameters.
eqParamMixed3 ::
  forall instr1 instr2 a1 a2 b1 b2 t.
  ( Typeable instr1
  , Typeable instr2
  , SingI a1
  , SingI a2
  , SingI b1
  , SingI b2
  , SDecide (KindOf a1)
  , SDecide (KindOf b1)
  , Eq (t instr1 a1 b1)
  )
  => t instr1 a1 b1
  -> t instr2 a2 b2
  -> Bool
eqParamMixed3 t1 t2 = isJust @() $ do
  Refl <- eqT @instr1 @instr2
  Refl <- sing @a1 `decideEquality` sing @a2
  Refl <- sing @b1 `decideEquality` sing @b2
  guard (t1 == t2)

castSing :: forall a b t.
  (SingI a, SingI b, SDecide (KindOf a))
  => t a
  -> Maybe (t b)
castSing ca = do
  Refl <- sing @a `decideEquality` sing @b
  return ca

-- | A second-order analogue of 'SingI'. This serves the same purpose as the
-- @"Data.Singletons".'Data.Singletons.SingIOne'@ class. However, the @singletons@
-- version requires a separate instance for each type in order to support GHC
-- versions that don't offer @QuantifiedConstraints@. We only need one instance
-- for all types.
class (forall x. SingI x => SingI (f x)) => SingIOne f
instance (forall x. SingI x => SingI (f x)) => SingIOne f

withSingIOne :: forall f x r. (SingIOne f, SingI x) => (SingI (f x) => r) -> r
withSingIOne f = f

-- | Helper function to generate 'SingI' and 'Data.Singletons.TH.SDecide' instances
-- using @Sing@ as prefix for the type names and @S@ for constructors'.
genSingletonsType :: OptionsMonad q => Name -> q [Dec]
genSingletonsType name =
  let
    singPrefix, sPrefix :: Name -> Name
    singPrefix nm = mkName ("Sing" ++ nameBase nm)
    sPrefix nm = mkName ("S" ++ nameBase nm)
  in
    withOptions defaultOptions{singledDataConName = sPrefix, singledDataTypeName = singPrefix} $
    concat <$> sequence
      [ genSingletons [name], singDecideInstance name]

-- | Singleton-based implementation of @geq@
geqI
  :: forall f a b. (SingI a, SingI b, TestEquality (Sing @(KindOf a)), Eq (f b))
  => f a -> f b -> Maybe (a :~: b)
geqI l r = case eqI @a @b of
  Just Refl | l == r -> Just Refl
  _ -> Nothing