From a186db842e26c660f91a2c08e28c1eff28f93c8b Mon Sep 17 00:00:00 2001
From: Ben Selfridge <benselfridge@galois.com>
Date: Tue, 19 Apr 2022 16:21:48 -0700
Subject: [PATCH 1/2] Redefines the EqF type class to be heterogeneous.

---
 src/Data/Parameterized/All.hs              |  2 +-
 src/Data/Parameterized/BoolRepr.hs         |  3 ++
 src/Data/Parameterized/Classes.hs          | 44 ++++++++++------------
 src/Data/Parameterized/Context.hs          |  4 +-
 src/Data/Parameterized/Context/Safe.hs     | 13 +++++--
 src/Data/Parameterized/Context/Unsafe.hs   | 11 ++++--
 src/Data/Parameterized/DataKind.hs         |  8 ++++
 src/Data/Parameterized/HashTable.hs        | 14 +++----
 src/Data/Parameterized/List.hs             | 16 ++++++--
 src/Data/Parameterized/Map.hs              |  2 +-
 src/Data/Parameterized/NatRepr/Internal.hs |  3 ++
 src/Data/Parameterized/Nonce.hs            |  3 ++
 src/Data/Parameterized/Nonce/Unsafe.hs     |  3 ++
 src/Data/Parameterized/Pair.hs             |  8 ++--
 src/Data/Parameterized/Peano.hs            |  3 ++
 src/Data/Parameterized/Some.hs             |  4 +-
 src/Data/Parameterized/SymbolRepr.hs       |  2 +
 17 files changed, 91 insertions(+), 52 deletions(-)

diff --git a/src/Data/Parameterized/All.hs b/src/Data/Parameterized/All.hs
index 0c585f0..dbcde39 100644
--- a/src/Data/Parameterized/All.hs
+++ b/src/Data/Parameterized/All.hs
@@ -63,7 +63,7 @@ instance ShowF f => Show (All f) where
   show (All fa) = showF fa
 
 instance EqF f => Eq (All f) where
-  (All x) == (All y) = eqF x y
+  All x == All y = isJust (x `eqF` y)
 
 allConst :: a -> All (Const a)
 allConst a = All (Const a)
diff --git a/src/Data/Parameterized/BoolRepr.hs b/src/Data/Parameterized/BoolRepr.hs
index 1534c34..13d7f77 100644
--- a/src/Data/Parameterized/BoolRepr.hs
+++ b/src/Data/Parameterized/BoolRepr.hs
@@ -79,6 +79,9 @@ instance TestEquality BoolRepr where
   testEquality FalseRepr FalseRepr = Just Refl
   testEquality _ _ = Nothing
 
+instance EqF BoolRepr where
+  eqF = testEquality
+
 instance DecidableEq BoolRepr where
   decEq TrueRepr  TrueRepr  = Left Refl
   decEq FalseRepr FalseRepr = Left Refl
diff --git a/src/Data/Parameterized/Classes.hs b/src/Data/Parameterized/Classes.hs
index b90754d..28ffe09 100644
--- a/src/Data/Parameterized/Classes.hs
+++ b/src/Data/Parameterized/Classes.hs
@@ -25,7 +25,7 @@ not restricted to '*'.
 {-# LANGUAGE TypeOperators #-}
 module Data.Parameterized.Classes
   ( -- * Equality exports
-    Equality.TestEquality(..)
+    TestEquality(..)
   , (Equality.:~:)(..)
   , EqF(..)
   , PolyEq(..)
@@ -90,27 +90,23 @@ instance CoercibleF (Const x) where
 ------------------------------------------------------------------------
 -- EqF
 
--- | @EqF@ provides a method @eqF@ for testing whether two parameterized
--- types are equal.
---
--- Unlike 'TestEquality', this only works when the type arguments are
--- the same, and does not provide a proof that the types have the same
--- type when they are equal. Thus this can be implemented over
--- parameterized types that are unable to provide evidence that their
--- type arguments are equal.
+-- | @EqF@ allows you to test for type- and value-level equality over a
+-- parameterized type. It is like @TestEquality@, but unlike that class, @EqF@
+-- returns `Nothing` when the type indices are equal but the values are not.
+-- This makes it appropriate to implement on non-singleton types.
 class EqF (f :: k -> Type) where
-  eqF :: f a -> f a -> Bool
-
-instance Eq a => EqF (Const a) where
-  eqF (Const x) (Const y) = x == y
+  eqF :: f a -> f b -> Maybe (a :~: b)
 
-instance EqF Proxy where
-  eqF _ _ = True
+instance (EqF f) => EqF (Compose f g) where
+  eqF (Compose x) (Compose y) =
+    case eqF x y of -- :: Maybe (g x :~: g y)
+      Just Refl -> Just Refl -- :: Maybe (x :~: y)
+      Nothing   -> Nothing
 
 ------------------------------------------------------------------------
 -- PolyEq
 
--- | A polymorphic equality operator that generalizes 'TestEquality'.
+-- | A polymorphic equality operator that generalizes 'EqF'.
 class PolyEq u v where
   polyEqF :: u -> v -> Maybe (u :~: v)
 
@@ -165,30 +161,30 @@ joinOrderingF GTF _ = GTF
 --
 -- [__Transitivity__]: if @leqF x y && leqF y z@ = 'True', then @leqF x = z@ = @True@
 -- [__Reflexivity__]: @leqF x x@ = @True@
--- [__Antisymmetry__]: if @leqF x y && leqF y x@ = 'True', then @testEquality x y@ = @Just Refl@
+-- [__Antisymmetry__]: if @leqF x y && leqF y x@ = 'True', then @eqF x y@ = @Just Refl@
 --
 -- Note that the following operator interactions are expected to hold:
 --
 -- * @geqF x y@ iff @leqF y x@
--- * @ltF  x y@ iff @leqF x y && testEquality x y = Nothing@
+-- * @ltF  x y@ iff @leqF x y && eqF x y = Nothing@
 -- * @gtF  x y@ iff @ltF y x@
 -- * @ltF  x y@ iff @compareF x y == LTF@
 -- * @gtF  x y@ iff @compareF x y == GTF@
--- * @isJust (testEquality x y)@ iff @compareF x y == EQF@
+-- * @isJust (eqF x y)@ iff @compareF x y == EQF@
 --
 -- Furthermore, when @x@ and @y@ both have type @(k tp)@, we expect:
 --
 -- * @toOrdering (compareF x y)@ equals @compare x y@ when @Ord (k tp)@ has an instance.
--- * @isJust (testEquality x y)@ equals @x == y@ when @Eq (k tp)@ has an instance.
+-- * @isJust (eqF x y)@ equals @x == y@ when @Eq (k tp)@ has an instance.
 --
 -- Minimal complete definition: either 'compareF' or 'leqF'.
 -- Using 'compareF' can be more efficient for complex types.
-class TestEquality ktp => OrdF (ktp :: k -> Type) where
+class EqF ktp => OrdF (ktp :: k -> Type) where
   {-# MINIMAL compareF | leqF #-}
 
   compareF :: ktp x -> ktp y -> OrderingF x y
   compareF x y =
-    case testEquality x y of
+    case eqF x y of
       Just Refl -> EQF
       Nothing | leqF x y -> LTF
               | otherwise -> GTF
@@ -342,8 +338,8 @@ instance Hashable a => HashableF (Const a) where
 -- our own new type to avoid orphan instances.
 newtype TypeAp (f :: k -> Type) (tp :: k) = TypeAp (f tp)
 
-instance TestEquality f => Eq (TypeAp f tp) where
-  TypeAp x == TypeAp y = isJust $ testEquality x y
+instance EqF f => Eq (TypeAp f tp) where
+  TypeAp x == TypeAp y = isJust $ eqF x y
 
 instance OrdF f => Ord (TypeAp f tp) where
   compare (TypeAp x) (TypeAp y) = toOrdering (compareF x y)
diff --git a/src/Data/Parameterized/Context.hs b/src/Data/Parameterized/Context.hs
index 900b5bf..9c6c4ce 100644
--- a/src/Data/Parameterized/Context.hs
+++ b/src/Data/Parameterized/Context.hs
@@ -168,7 +168,7 @@ toVector a f = V.create $ do
 --   `prefix` as a prefix, and computing the tail of xs
 --   not in the prefix, if so.
 dropPrefix :: forall f xs prefix a.
-  TestEquality f =>
+  EqF f =>
   Assignment f xs     {- ^ Assignment to split -} ->
   Assignment f prefix {- ^ Expected prefix -} ->
   a {- ^ error continuation -} ->
@@ -189,7 +189,7 @@ dropPrefix xs0 prefix err = go xs0 (sizeInt (size xs0))
     go xs' (sz_x-1) (\zs -> success (zs :> z))
 
   go xs _ success =
-    case testEquality xs prefix of
+    case eqF xs prefix of
       Just Refl -> success Empty
       Nothing   -> err
 
diff --git a/src/Data/Parameterized/Context/Safe.hs b/src/Data/Parameterized/Context/Safe.hs
index 96da3c0..b55c713 100644
--- a/src/Data/Parameterized/Context/Safe.hs
+++ b/src/Data/Parameterized/Context/Safe.hs
@@ -290,6 +290,9 @@ instance TestEquality (Index ctx) where
          Just Refl -> Just Refl
          Nothing -> Nothing
 
+instance EqF (Index ctx) where
+  eqF = testEquality
+
 instance Ord (Index ctx tp) where
   compare i j = toOrdering (compareF i j)
 
@@ -543,8 +546,8 @@ idxlookup _ AssignmentEmpty idx = case idx of {}
 (!^) :: KnownDiff l r => Assignment f r -> Index l tp -> f tp
 a !^ i = a ! extendIndex i
 
-instance TestEquality f => Eq (Assignment f ctx) where
-  x == y = isJust (testEquality x y)
+instance EqF f => Eq (Assignment f ctx) where
+  x == y = isJust (eqF x y)
 
 testEq :: (forall x y. f x -> f y -> Maybe (x :~: y))
        -> Assignment f cxt1 -> Assignment f cxt2 -> Maybe (cxt1 :~: cxt2)
@@ -563,8 +566,10 @@ instance TestEqualityFC Assignment where
    testEqualityFC f = testEq f
 instance TestEquality f => TestEquality (Assignment f) where
    testEquality x y = testEq testEquality x y
-instance TestEquality f => PolyEq (Assignment f x) (Assignment f y) where
-  polyEqF x y = fmap (\Refl -> Refl) $ testEquality x y
+instance EqF f => PolyEq (Assignment f x) (Assignment f y) where
+  polyEqF x y = fmap (\Refl -> Refl) $ eqF x y
+instance EqF f => EqF (Assignment f) where
+  eqF x y = testEq eqF x y
 
 compareAsgn :: (forall x y. f x -> f y -> OrderingF x y)
             -> Assignment f ctx1 -> Assignment f ctx2 -> OrderingF ctx1 ctx2
diff --git a/src/Data/Parameterized/Context/Unsafe.hs b/src/Data/Parameterized/Context/Unsafe.hs
index 609f2af..d0fb903 100644
--- a/src/Data/Parameterized/Context/Unsafe.hs
+++ b/src/Data/Parameterized/Context/Unsafe.hs
@@ -246,8 +246,8 @@ type role Index nominal nominal
 instance Eq (Index ctx tp) where
   Index i == Index j = i == j
 
-instance TestEquality (Index ctx) where
-  testEquality (Index i) (Index j)
+instance EqF (Index ctx) where
+  eqF (Index i) (Index j)
     | i == j = Just unsafeAxiom
     | otherwise = Nothing
 
@@ -832,8 +832,11 @@ instance TestEqualityFC Assignment where
 instance TestEquality f => TestEquality (Assignment f) where
   testEquality = testEqualityFC testEquality
 
-instance TestEquality f => Eq (Assignment f ctx) where
-  x == y = isJust (testEquality x y)
+instance EqF f => Eq (Assignment f ctx) where
+  x == y = isJust (eqF x y)
+
+instance EqF f => EqF (Assignment f) where
+  eqF = testEqualityFC eqF
 
 instance OrdFC Assignment where
   compareFC test (Assignment x) (Assignment y) =
diff --git a/src/Data/Parameterized/DataKind.hs b/src/Data/Parameterized/DataKind.hs
index b826c33..d84dc7b 100644
--- a/src/Data/Parameterized/DataKind.hs
+++ b/src/Data/Parameterized/DataKind.hs
@@ -45,6 +45,14 @@ instance ( TestEquality f, TestEquality g ) => TestEquality (PairRepr f g) where
       , ( TH.DataArg 1 `TH.TypeApp` TH.AnyType, [|testEquality|] )
       ])
 
+instance ( EqF f, EqF g) => EqF (PairRepr f g) where
+  eqF =
+    $(TH.structuralTypeEquality [t|PairRepr|]
+      [
+        ( TH.DataArg 0 `TH.TypeApp` TH.AnyType, [|eqF|] )
+      , ( TH.DataArg 1 `TH.TypeApp` TH.AnyType, [|eqF|] )
+      ])
+
 deriving instance ( Ord (f a), Ord (g b) ) => Ord (PairRepr f g '(a, b))
 instance ( OrdF f, OrdF g ) => OrdF (PairRepr f g) where
   compareF =
diff --git a/src/Data/Parameterized/HashTable.hs b/src/Data/Parameterized/HashTable.hs
index adf52d9..4d0e7cd 100644
--- a/src/Data/Parameterized/HashTable.hs
+++ b/src/Data/Parameterized/HashTable.hs
@@ -9,7 +9,7 @@
 --
 -- NOTE: This API makes use of 'unsafeCoerce' to implement the parameterized
 -- hashtable abstraction.  This should be type-safe provided that the
--- 'TestEquality' instance on the key type is implemented soundly.
+-- 'EqF' instance on the key type is implemented soundly.
 ------------------------------------------------------------------------
 {-# LANGUAGE KindSignatures #-}
 {-# LANGUAGE PolyKinds #-}
@@ -53,7 +53,7 @@ newSized :: Int -> ST s (HashTable s k v)
 newSized n = HashTable <$> H.newSized n
 
 -- | Create a hash table that is a copy of the current one.
-clone :: (HashableF key, TestEquality key)
+clone :: (HashableF key, EqF key)
       => HashTable s key val
       -> ST s (HashTable s key val)
 clone (HashTable tbl) = do
@@ -65,7 +65,7 @@ clone (HashTable tbl) = do
   return $! HashTable r
 
 -- | Lookup value of key in table.
-lookup :: (HashableF key, TestEquality key)
+lookup :: (HashableF key, EqF key)
        => HashTable s key val
        -> key tp
        -> ST s (Maybe (val tp))
@@ -73,7 +73,7 @@ lookup (HashTable h) k = fmap unsafeCoerce <$> H.lookup h (Some k)
 {-# INLINE lookup #-}
 
 -- | Insert new key and value mapping into table.
-insert :: (HashableF key, TestEquality key)
+insert :: (HashableF key, EqF key)
        => HashTable s (key :: k -> Type) (val :: k -> Type)
        -> key tp
        -> val tp
@@ -81,19 +81,19 @@ insert :: (HashableF key, TestEquality key)
 insert (HashTable h) k v = H.insert h (Some k) (unsafeCoerce v)
 
 -- | Return true if the key is in the hash table.
-member :: (HashableF key, TestEquality key)
+member :: (HashableF key, EqF key)
        => HashTable s (key :: k -> Type) (val :: k -> Type)
        -> key (tp :: k)
        -> ST s Bool
 member (HashTable h) k = isJust <$> H.lookup h (Some k)
 
 -- | Delete an element from the hash table.
-delete :: (HashableF key, TestEquality key)
+delete :: (HashableF key, EqF key)
        => HashTable s (key :: k -> Type) (val :: k -> Type)
        -> key (tp :: k)
        -> ST s ()
 delete (HashTable h) k = H.delete h (Some k)
 
-clear :: (HashableF key, TestEquality key)
+clear :: (HashableF key, EqF key)
       => HashTable s (key :: k -> Type) (val :: k -> Type) -> ST s ()
 clear (HashTable h) = H.mapM_ (\(k,_) -> H.delete h k) h
diff --git a/src/Data/Parameterized/List.hs b/src/Data/Parameterized/List.hs
index e9f7432..f5eaf56 100644
--- a/src/Data/Parameterized/List.hs
+++ b/src/Data/Parameterized/List.hs
@@ -214,6 +214,14 @@ instance TestEquality f => TestEquality (List f) where
     pure Refl
   testEquality _ _ = Nothing
 
+instance EqF f => EqF (List f) where
+  eqF Nil Nil = Just Refl
+  eqF (xh :< xl) (yh :< yl) = do
+    Refl <- eqF xh yh
+    Refl <- eqF xl yl
+    pure Refl
+  eqF _ _ = Nothing
+
 instance OrdF f => OrdF (List f) where
   compareF Nil Nil = EQF
   compareF Nil _ = LTF
@@ -267,10 +275,10 @@ deriving instance Show  (Index l x)
 
 instance ShowF (Index l)
 
-instance TestEquality (Index l) where
-  testEquality IndexHere IndexHere = Just Refl
-  testEquality (IndexThere x) (IndexThere y) = testEquality x y
-  testEquality _ _ = Nothing
+instance EqF (Index l) where
+  eqF IndexHere IndexHere = Just Refl
+  eqF (IndexThere x) (IndexThere y) = eqF x y
+  eqF _ _ = Nothing
 
 instance OrdF (Index l) where
   compareF IndexHere IndexHere = EQF
diff --git a/src/Data/Parameterized/Map.hs b/src/Data/Parameterized/Map.hs
index 966df5a..0858d16 100644
--- a/src/Data/Parameterized/Map.hs
+++ b/src/Data/Parameterized/Map.hs
@@ -151,7 +151,7 @@ instance Bin.IsBinTree (MapF k a) (Pair k a) where
   size Tip              = 0
   size (Bin sz _ _ _ _) = sz
 
-instance (TestEquality k, EqF a) => Eq (MapF k a) where
+instance (EqF k, EqF a) => Eq (MapF k a) where
   x == y = size x == size y && toList x == toList y
 
 ------------------------------------------------------------------------
diff --git a/src/Data/Parameterized/NatRepr/Internal.hs b/src/Data/Parameterized/NatRepr/Internal.hs
index fad97e1..07aebaf 100644
--- a/src/Data/Parameterized/NatRepr/Internal.hs
+++ b/src/Data/Parameterized/NatRepr/Internal.hs
@@ -53,6 +53,9 @@ instance TestEquality NatRepr where
     | m == n = Just unsafeAxiom
     | otherwise = Nothing
 
+instance EqF NatRepr where
+  eqF = testEquality
+
 instance DecidableEq NatRepr where
   decEq (NatRepr m) (NatRepr n)
     | m == n    = Left unsafeAxiom
diff --git a/src/Data/Parameterized/Nonce.hs b/src/Data/Parameterized/Nonce.hs
index 5387f3d..98f583f 100644
--- a/src/Data/Parameterized/Nonce.hs
+++ b/src/Data/Parameterized/Nonce.hs
@@ -132,6 +132,9 @@ instance TestEquality (Nonce s) where
   testEquality x y | indexValue x == indexValue y = Just unsafeAxiom
                    | otherwise = Nothing
 
+instance EqF (Nonce s) where
+  eqF = testEquality
+
 instance OrdF (Nonce s) where
   compareF x y =
     case compare (indexValue x) (indexValue y) of
diff --git a/src/Data/Parameterized/Nonce/Unsafe.hs b/src/Data/Parameterized/Nonce/Unsafe.hs
index 1ef11a3..087e830 100644
--- a/src/Data/Parameterized/Nonce/Unsafe.hs
+++ b/src/Data/Parameterized/Nonce/Unsafe.hs
@@ -68,6 +68,9 @@ instance TestEquality Nonce where
   testEquality x y | indexValue x == indexValue y = Just unsafeAxiom
                    | otherwise = Nothing
 
+instance EqF Nonce where
+  eqF = testEquality
+
 instance OrdF Nonce where
   compareF x y =
     case compare (indexValue x) (indexValue y) of
diff --git a/src/Data/Parameterized/Pair.hs b/src/Data/Parameterized/Pair.hs
index 2b8eb41..c211085 100644
--- a/src/Data/Parameterized/Pair.hs
+++ b/src/Data/Parameterized/Pair.hs
@@ -27,10 +27,12 @@ import Data.Parameterized.TraversableF
 data Pair (a :: k -> Type) (b :: k -> Type) where
   Pair :: !(a tp) -> !(b tp) -> Pair a b
 
-instance (TestEquality a, EqF b) => Eq (Pair a b) where
+instance (EqF a, EqF b) => Eq (Pair a b) where
   Pair xa xb == Pair ya yb =
-    case testEquality xa ya of
-      Just Refl -> eqF xb yb
+    case xa `eqF` ya of
+      Just Refl -> case xb `eqF` yb of
+        Just Refl -> True
+        Nothing -> False
       Nothing -> False
 
 instance FunctorF (Pair a) where
diff --git a/src/Data/Parameterized/Peano.hs b/src/Data/Parameterized/Peano.hs
index 942e1c3..4ea505a 100644
--- a/src/Data/Parameterized/Peano.hs
+++ b/src/Data/Parameterized/Peano.hs
@@ -233,6 +233,9 @@ instance TestEquality PeanoRepr where
 
 #endif
 
+instance EqF PeanoRepr where
+  eqF = testEquality
+
 instance DecidableEq PeanoRepr where
 #ifdef UNSAFE_OPS
   decEq (PeanoRepr m) (PeanoRepr n)
diff --git a/src/Data/Parameterized/Some.hs b/src/Data/Parameterized/Some.hs
index 3df9359..6a1aa3c 100644
--- a/src/Data/Parameterized/Some.hs
+++ b/src/Data/Parameterized/Some.hs
@@ -27,8 +27,8 @@ import Data.Parameterized.TraversableF
 
 data Some (f:: k -> Type) = forall x . Some (f x)
 
-instance TestEquality f => Eq (Some f) where
-  Some x == Some y = isJust (testEquality x y)
+instance EqF f => Eq (Some f) where
+  Some x == Some y = isJust (eqF x y)
 
 instance OrdF f => Ord (Some f) where
   compare (Some x) (Some y) = toOrdering (compareF x y)
diff --git a/src/Data/Parameterized/SymbolRepr.hs b/src/Data/Parameterized/SymbolRepr.hs
index 473e4ec..9d69fce 100644
--- a/src/Data/Parameterized/SymbolRepr.hs
+++ b/src/Data/Parameterized/SymbolRepr.hs
@@ -84,6 +84,8 @@ instance TestEquality SymbolRepr where
    testEquality (SymbolRepr x :: SymbolRepr x) (SymbolRepr y)
       | x == y    = Just unsafeAxiom
       | otherwise = Nothing
+instance EqF SymbolRepr where
+  eqF = testEquality
 instance OrdF SymbolRepr where
    compareF (SymbolRepr x :: SymbolRepr x) (SymbolRepr y)
       | x <  y    = LTF

From 2389d52e03da7e97788303dc09b351bc6fba01f5 Mon Sep 17 00:00:00 2001
From: Ben Selfridge <benselfridge@galois.com>
Date: Tue, 19 Apr 2022 16:30:36 -0700
Subject: [PATCH 2/2] Restores deleted (but sensible) TestEquality instance

---
 src/Data/Parameterized/Context/Unsafe.hs | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/src/Data/Parameterized/Context/Unsafe.hs b/src/Data/Parameterized/Context/Unsafe.hs
index d0fb903..145291c 100644
--- a/src/Data/Parameterized/Context/Unsafe.hs
+++ b/src/Data/Parameterized/Context/Unsafe.hs
@@ -246,11 +246,14 @@ type role Index nominal nominal
 instance Eq (Index ctx tp) where
   Index i == Index j = i == j
 
-instance EqF (Index ctx) where
-  eqF (Index i) (Index j)
+instance TestEquality (Index ctx) where
+  testEquality (Index i) (Index j)
     | i == j = Just unsafeAxiom
     | otherwise = Nothing
 
+instance EqF (Index ctx) where
+  eqF = testEquality
+
 instance Ord (Index ctx tp) where
   Index i `compare` Index j = compare i j