[Haskell-cafe] How to deal with type family applications in type equalities?

[Jakob Brünker]

Imagine you want to have a type level function that calculates the
element-wise OR on two bitstrings, encoded as length-indexed vectors filled
with Bools (this is a simplification of something I need for a project).
The vectors should be "aligned to the right" as it were, such that the new
right-most value is True if the right-most value of the first vector OR the
right-most value of the second vector was True.

Example (replacing True and False with 1 and 0 for readability):
ElementwiseOr [1,0,1,1] [0,1,1,0,0,0,0,1]
= [0,1,1,0,1,0,1,1]

Its type would be something like
ElementwiseOr :: Vec n Bool -> Vec m Bool -> Vec (Max n m) Bool

I have written the following code:

---------------------------------------------------------------------------
{-# LANGUAGE TypeFamilies, GADTs, DataKinds, StandaloneKindSignatures,
             PolyKinds, TypeOperators, RankNTypes, TypeApplications,
             UndecidableInstances #-}

import Data.Kind
import Data.Type.Equality

data Nat = Z | S Nat

infixr 5 :<

type Vec :: Nat -> Type -> Type
data Vec n a where
  Nil :: Vec Z a
  (:<) :: a -> Vec n a -> Vec (S n) a

type Cong :: forall f -> a :~: b -> f a :~: f b
type family Cong f p where
  Cong f Refl = Refl

type Max :: Nat -> Nat -> Nat
type family Max n m where
  Max Z m = m
  Max n Z = n
  Max (S n) (S m) = S (Max n m)

-- The reason for this slightly convoluted Ordering type is that it seems
-- to make the "right-alignment" easier.
type NatOrdering :: Ordering -> Nat -> Nat -> Type
data NatOrdering o n m where
  NOLTE :: NatOrdering EQ n m -> NatOrdering LT n (S m)
  NOLTS :: NatOrdering LT n m -> NatOrdering LT n (S m)
  NOEQZ :: NatOrdering EQ Z Z
  NOEQS :: NatOrdering EQ n m -> NatOrdering EQ (S n) (S m)
  NOGTE :: NatOrdering EQ n m -> NatOrdering GT (S n) m
  NOGTS :: NatOrdering GT n m -> NatOrdering GT (S n) m

type MaxSIsSMaxEQ
  :: forall n m . NatOrdering EQ n m -> Max n (S m) :~: S (Max n m)
type family MaxSIsSMaxEQ o where
  MaxSIsSMaxEQ NOEQZ = Refl
  MaxSIsSMaxEQ (NOEQS o) = Cong S (MaxSIsSMaxEQ o)

type TransportVec :: n :~: m -> Vec n a -> Vec m a
type family TransportVec p v where
  TransportVec Refl v = v

type ElementwiseOr
  :: NatOrdering o n m -> Vec n Bool -> Vec m Bool -> Vec (Max n m) Bool
type family ElementwiseOr o u v where
  ElementwiseOr (NOLTE o) u (s :< v) =
    TransportVec (MaxSIsSMaxEQ o) (s :< ElementwiseOr o u v) -- XXX
  -- other equations ommitted
---------------------------------------------------------------------------

To me, the equation at the end marked with XXX seems like it should work.
However, it produces the following error:

    • Couldn't match kind: Max n ('S m)
                     with: 'S (Max n m)
      Expected kind ‘'S (Max n m) :~: 'S (Max n m)’,
        but ‘MaxSIsSMaxEQ o’ has kind ‘Max n ('S m) :~: 'S (Max n m)’
    • In the first argument of ‘TransportVec’, namely
        ‘(MaxSIsSMaxEQ o)’

So it expects something of kind `S (Max n m) :~: S (Max n m)` - it seems
like Refl fits that bill, but that doesn't work either, because replacing
(MaxSIsSMaxEQ o) with Refl means that TransportVec returns a type of the
wrong kind, producing a very similar error.

I suspect the reason this doesn't work is because type equality isn't
equipped to properly handle type family applications. I could potentially
work around this if I defined Max as

data Max n m r where
  MaxZZ :: Max Z Z Z
  MaxSZ :: Max (S n) Z (S n)
  MaxZS :: Max Z (S m) (S m)
  MaxSS :: Max n m r -> Max (S n) (S m) (S r)

However, dealing with a proof object like this is a lot less convenient
than just being able to use a type family for most purposes, so I'd like to
avoid this if possible.

Is there a way to make this work while sticking with type families?

Thanks

Jakob
-------------- next part --------------
An HTML attachment was scrubbed...
URL: <http://mail.haskell.org/pipermail/haskell-cafe/attachments/20210218/cdaf7151/attachment.html>