Is there a flexible way of specifying return types for Haskell's type families?

Question

Recently I started using Haskell's Repa library, which relies heavily on type families and associated types. I can define a Repa array like this:

ghci> let x = fromListUnboxed (Z :. (5 :: Int) :. (2 :: Int)) [1..10]

and operate on it like this:

ghci> computeP $ R.map id x :: IO (Array U DIM2 Double)
AUnboxed ((Z :. 5) :. 2) (fromList [1.0,2.0,3.0,4.0,5.0,6.0,7.0,8.0,9.0,10.0])

where U is an associated data type that will cause the result to be represented as an unboxed array; DIM2 is array dimension. I don't like that I have to specify a concrete dimension even though it could be deduced. Instead I would want to write something like this:

ghci> computeP $ R.map id x :: Shape sh => IO (Array U sh Double)

This is not valid, but my intention is to be able to specify array type by passing appropriate associated data type (U in this example), but leave the shape unchanged. Is something like that possible?

Answer 1

Does something like this work?

asArrayU :: Array U sh a -> Array U sh a
asArrayU = id

computeP $ asArrayU <$> R.map id x

Answer 2

Your example

ghci> computeP $ R.map id x :: Shape sh => IO (Array U sh Double)

doesn't work because here you declare that the resulting value is polymorphic and can have any possible shape, but this is clearly untrue since x has a specific shape, and hence the result can only have the same specific shape, namely DIM2.

What you can say is that the combination of computeP and R.map id retains the same shape, no matter what the shape originally is. I.e.

ghci> (computeP . R.map id :: Shape sh => Array U sh Double -> IO (Array U sh Double)) x
AUnboxed ((Z :. 5) :. 2) (fromList [1.0,2.0,3.0,4.0,5.0,6.0,7.0,8.0,9.0,10.0])

The important part here is that we lock the representation to unboxed values, because that is the part that cannot be automatically determined by the compiler.

But without knowing more about your real use-case it's difficult to say what is the nicest option syntactically to express what you want. Are you mainly interested in testing stuff in GHCi without specifying explicit types or do you want to avoid explicit typing in somewhere in a real program?

For example, you could define something like

computeMap :: (Unbox a, Unbox b, Shape sh) => (a -> b) -> Array U sh a -> IO (Array U sh b)
computeMap f = computeP . R.map f

and then you can say

ghci> computeMap id x
AUnboxed ((Z :. 5) :. 2) (fromList [1.0,2.0,3.0,4.0,5.0,6.0,7.0,8.0,9.0,10.0])