Must I cast Nat-kinds every time?

Question

I tried to simulate a quantum computer. Here is the datatype representing qubits:

{-# LANGUAGE DataKinds #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE TypeOperators #-}

import Control.Monad
import Data.Maybe
import Data.Proxy
import Data.Type.Equality
import GHC.TypeNats

import Data.Group.Cyclic

data QBits :: Nat -> * where
    N :: QBits 0
    C :: KnownNat n => Bool -> QBits n -> QBits (n+1)
    S :: KnownNat n => Cyclic 4 -> QBits n -> QBits n -> QBits (n+1)

N represents zero qubits.

C, standing for "classical", assigns the first qubit a boolean value, and specifies the rest.

S, standing for "superposed", states that the first qubit is in superposition, and specifies the rest for each possibility in which the first qubit will fall when measured. It also specifies the phase difference, which is a value in Cyclic 4, which is the ring Z/4Z and has Num instance.

For instance Eq (QBits n), I have a workaround so I won't mess with Nat:

(=?=) :: QBits m -> QBits n -> Bool
N =?= N = True
C b x =?= C c y = b == c && x =?= y
S p x y =?= S q u v = p == q && x =?= u && y =?= v
_ =?= _ = False

instance Eq (QBits n) where
    (==) = (=?=)

Then I implemented swapGate, which swaps first two qubits:

castNat :: forall f m n. (KnownNat m, KnownNat n) => f m -> Maybe (f n)
castNat x = do
    refl <- sameNat (Proxy :: Proxy m) (Proxy :: Proxy n)
    return (castWith (apply Refl refl) x)

swapGate :: KnownNat n => QBits n -> QBits n
swapGate (C b (C c x)) = C c (C b x)
swapGate (C b (S p x y)) = S p (C b x) (C b y)
swapGate (S r (C False x) (C False y)) = let
    Just y' = castNat y
    in C False (S r x y')
swapGate (S r (C False x) (S q u v)) = let
    Just u' = castNat u
    in S (r+q) (S r x u') (C True v)
swapGate (S r (C True y) (C False u)) = S (-r) (C True u) (C False y)
swapGate (S r (C True y) (C True v)) = let
    Just v' = castNat v
    in C True (S r y v')
swapGate (S r (C True y) (S q u v)) = let
    Just v' = castNat v
    in S (-r) (C True u) (S (r+q) y v')
swapGate (S r (S p x y) (C False u)) = let
    Just u' = castNat u
    in S p (S r x u') (C False y)
swapGate (S r (S p x y) (C True v)) = let
    Just v' = castNat v
    in S p (C False x) (S (p-r) y v')
swapGate (S r (S p x y) (S q u v)) = let
    Just u' = castNat u
    Just v' = castNat v
    in S p (S r x u') (S (q-p+r) y v')
swapGate z = z

The fact I must cast Nats is just too annoying. Is castNat truly mandatory?

user_7203016 · Answer

For one thing, to fix the syntactic abomination, you could write:

c :: forall f m n. (KnownNat m, KnownNat n) => f m -> f n
c = fromJust . castNat

and then:

swapGate :: KnownNat n => QBits n -> QBits n
swapGate (C b (C c x)) = C c (C b x)
swapGate (C b (S p x y)) = S p (C b x) (C b y)
swapGate (S r (C False x) (C False y)) = C False (S r x (c y))
swapGate (S r (C False x) (S q u v)) = S (r+q) (S r x (c u)) (C True v)
... etc. ...

As explained in the comments, the underlying "problem" is that GHC's built-in inference for type-level naturals is very limited. Operators will work on concrete types and handle a few special abstract cases, like 0 + m ~ m, but GHC can't make other, even extremely simply inference, like m + 1 - 1 ~ m or "m + 1 ~ n + 1 implies m ~ n".

Your choices are to rewrite using an algebraic Nat type (e.g., Peano naturals) or use a solver plugin.

For this problem, Peano naturals are an (erm...) natural fit, since all your manipulations of type level naturals involve incrementing or decrementing them. After replacing Nat and the + type operator with:

data Nat = ZZ | SS Nat
type family m + n where
  ZZ + n = n
  SS m + n = m + SS n

and adjusting the QBits definition:

data QBits :: Nat -> * where
    N :: QBits ZZ
    C :: Bool -> QBits n -> QBits (SS n)
    S :: Cyclic4 -> QBits n -> QBits n -> QBits (SS n)

the castless definition typechecks fine:

swapGate :: QBits n -> QBits n
swapGate (C b (C c x)) = C c (C b x)
swapGate (C b (S p x y)) = S p (C b x) (C b y)
swapGate (S r (C False x) (C False y)) = C False (S r x y)
swapGate (S r (C False x) (S q u v)) = S (r+q) (S r x u) (C True v)
swapGate (S r (C True y) (C False u)) = S (-r) (C True u) (C False y)
swapGate (S r (C True y) (C True v)) = C True (S r y v)
swapGate (S r (C True y) (S q u v)) = S (-r) (C True u) (S (r+q) y v)
swapGate (S r (S p x y) (C False u)) = S p (S r x u) (C False y)
swapGate (S r (S p x y) (C True v)) = S p (C False x) (S (p-r) y v)
swapGate (S r (S p x y) (S q u v)) = S p (S r x u) (S (q-p+r) y v)
swapGate z = z

Alternatively, you can use a solver plugin. After installing ghc-typelits-natnormalise and adding:

{-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise #-}

to the top of your code, I can get rid of all the casts, and -- again -- it typechecks fine.

By the way, either of these solutions also allow you to drop the KnownNat constraints from the code, too. If performance is a consideration, that may be a win, since you won't have to carry all these dictionaries around.

Must I cast Nat-kinds every time?

1 Answers