Restrict types in Sigma

Question

I have type X indexed by kind S with several functions that work on X. For example, f converts X S1 to X S2 (it doesn't use X S1 in this simplified example, though).

{-# LANGUAGE DataKinds, GADTs, TemplateHaskell, TypeFamilies #-}

import Data.Singletons
import Data.Singletons.Sigma
import Data.Singletons.TH

singletons [d|
    data S = S1 | S2 | S3 | S4
  |]

data X (s :: S) = X

f :: X S1 -> X S2
f x = X

Now I'd like to define functions that may return X S2 or X S3 depending on its argument. The straightforward way would be using Either.

g1 :: Bool -> X S1 -> Either (X S2) (X S3)
g1 True x = Left X
g1 False x = Right X

But I don't want to take this approach because I need to have nested Eithers when a function returns more types.

Another approach would be using Sigma like this.

g2 :: Bool -> X S1 -> Sigma S (TyCon X)
g2 True x = SS2 :&: X
g2 False x = SS3 :&: X

But this doesn't express the idea that g2 returns only X S2 or X S3. I can express this idea by introducing a wrapper around X.

data Y (s :: S) where
    Y2 :: X S2 -> Y S2
    Y3 :: X S3 -> Y S3

singletons [d|
    type Z s = Y s
  |]

g3 :: Bool -> X S1 -> Sigma S ZSym0
g3 True x = SS2 :&: Y2 X
g3 False x = SS3 :&: Y3 X

But it's cumbersome to define these wrappers for each combination and unwrap them on caller sites. It'd be nice if I can directly restrict types using g2 approach, for example, by like applying type constraints, but I'm not sure how I can do it.

How can I directly restrict types using g2 approach?

I'm using GHC 8.8.4 with singletons-2.6.

Answer 1

This question looks too simplified and contrived; it would be nice to have some more concrete motivation. But here is a shot in the dark.

We can define a variant of Sigma with a predicate on the first component:

data SigmaP (i :: Type) (p :: i ~> Constraint) (f :: i -> Type) where
  (:&?:) :: (p @@ x) => Sing x -> f x -> SigmaP i p f

Some code to define the predicate seems unavoidable:

data Y23 :: S ~> Constraint
type instance Apply Y23 x = Y23_ x

type family Y23_ (x :: S) :: Constraint where
  Y23_ S2 = (() :: Constraint)
  Y23_ S3 = (() :: Constraint)
  Y23_ _ = ('True ~ 'False)

But now we can just use X:

g3 :: Bool -> X S1 -> SigmaP S Y23 X
g3 True x = SS2 :&?: X
g3 False x = SS3 :&?: X

Answer 2

Another approach inspired by @luqui's comment. It's similar to @li-yao-xia's answer, but uses an explicit list of types.

{-# LANGUAGE DataKinds, GADTs, EmptyCase, InstanceSigs, PolyKinds, ScopedTypeVariables, TemplateHaskell, TypeApplications, TypeFamilies, TypeOperators, UndecidableInstances #-}

import Data.Kind
import Data.Singletons
import Data.Singletons.Prelude
import Data.Singletons.TH

singletons [d|
    data S = S1 | S2 | S3 | S4 deriving (Show, Eq)
  |]

data X (s :: S) = X

data SigmaL (l :: [s :: Type]) (t :: s ~> Type) where
    (:&!:) :: OneOf fst l => Sing (fst :: s) -> t @@ fst -> SigmaL l t

type family OneOf s l :: Constraint where
    OneOf s l = If (Elem s l) (() :: Constraint) ('True ~ 'False)

g5 :: Bool -> X S1 -> SigmaL '[S2, S3] (TyCon X)
g5 True x = SS2 :&!: X
g5 False x = SS3 :&!: X