How do I implement a partially injective type family in Haskell?

Question

I am implementing a varity of functions that require type safe natural numbers in Haskell, and have recently needed an Exponential type to represent a new type.

Below are the three type families that I have made up to this point for ease of reference.

type family Add n m where
  Add 'One n = 'Succ n
  Add ('Succ n) m = 'Succ (Add n m)

-- Multiplication, allowing ever more powerful operators
type family Mul n m where
  Mul 'One m = m
  Mul ('Succ n) m = Add m (Mul n m)

-- Exponentiation, allowing even even more powerful operators
type family Exp n m where
  Exp n 'One = n
  Exp n ('Succ m) = Mul n (Exp n m)

However, when using this type I came across the issue that it wasn't injective; this meant that some of the type inferences I kinda wanted didn't exist. (The error was NB: ‘Exp’ is a non-injective type family). I can ignore the issue by using -XAllowAmbiguousTypes, but I would prefer not to use this extension, so all the types can checked where the function is defined.

I think that Exp n m should be injective when m is constant, so I wanted to try to implement that, but I'm unsure as to how to do that after a lot of trial and error. Even if it doesn't solve my problem currently, it might be useful in future. Alternatively, Exp n m is injective for a given n where m changes, and n is not One.

Upon asking other people, they suggested that something like type family Exp n m = inj | inj, n -> m where but that doesn't work, giving a syntax error on the comma if it's there and a parse error on the final n if it isn't. This was intended to allow inj and n to uniquely identify a given m.

The function I am trying to implement but am having trouble with currently has a signature as follows.

tensorPower :: forall i a n m . (Num a, KnownNat i) => Matrix n m a -> Matrix (Exp n i) (Exp m i) a

This function can be called with tensorPower @Three a (when -XAllowAmbiguousTypes is set), but I'd like GHC to be able to determine the i value by itself if possible. For the purposes of this question it's fine to assume that a given a matrix isn't polymorphic.

Adjusting the constraint to the following doesn't work either; this was an attempt at creating injectivity in the type of the above function instead of where the type family is defined

forall i a n m
   . ( Num a
     , KnownNat i
     , Exp n ( 'Succ i) ~ Mul n (Exp n i)
     , Exp m ( 'Succ i) ~ Mul m (Exp m i)
     , Exp n One ~ n
     , Exp m One ~ m
     )

So, is it possible to implement injectivity for this function, and if so, how do I implement it?

(To see more of the code in action, please visit the repository. The src folder has the majority of the source of the code, with the main areas being fiddled with in this question belonging to Lib.hs and Quantum.hs. The extensions used can (mostly) be found in the package.yaml)

Answer 1

There's actually a surprisingly simple way to get this to work in at least one way; the following type family, when used appropriately in a constraint, allows tensorPower to be used without an annotation.

-- Reverse the exponent - if it can't match then it goes infinitely
type family RLog n m x c where
  RLog m n n i = i
  RLog m n x i = RLog m n (Mul m x) ('Succ i)

type ReverseLog n m = RLog n m n 'One
type GetExp n i = ReverseLog n (Exp n i)
----------------
-- adjusted constraint for tensorPower
forall i a n m . (Num a, KnownNat i, i ~ GetExp n i, i ~ GetExp m i)

For example, now one can type (tensorPower hadamard) *.* (zero .*. zero .*. one) (where hadamard is Matrix Two Two Double, both zero and one are Matrix Two One Double, (*.*) is matrix multiplication, (.*.) is the tensor product and the type of i is completely inferred).

The way this type family works is that it has four parameters: the base, the target, the accumulator, and the current exponent. If the target and the accumulator are equal, then the current exponent is "returned". If they are not equal, we recurse by multiplying the current accumulator by the base, and increment the current exponent.

There is one issue with this solution that I can see: if it can't match up the "bases", it has a really horribly long error message as it recurses as deep as it can into types. This can be fixed by doing some other type trickery which is out of scope for this question, but can be seen in this commit on my project's repository.

In conclusion: introducing some abstract injectivity seemed to be a no-go, but implementing a sort of reversal of the exponent resulted in clean, simple, and functioning code - this is effectively injectivity, by proving that there is a reversible function for the Exp.

(One note is that this solution needs a bit more fiddling to work fully, since GetExp n i doesn't really work for n=='One; I've gotten around this by never having GetExp ('One) i in the first place)