Update trace wrappers

1 files changed, 38 insertions, 22 deletions

diff --git a/src/Data/Array/Nested/Trace/TH.hs b/src/Data/Array/Nested/Trace/TH.hs
index 47e53cd..4b388e3 100644
--- a/src/Data/Array/Nested/Trace/TH.hs
+++ b/src/Data/Array/Nested/Trace/TH.hs
@@ -10,7 +10,6 @@ import Language.Haskell.TH hiding (cxt)
 
 import Debug.Trace qualified as Debug
 
-import Data.Array.Mixed.Types
 import Data.Array.Nested
 
 
@@ -24,38 +23,55 @@ splitFunTy = \case
     in (vars ++ vs, cx ++ cx', args, ret)
   t -> ([], [], [], t)
 
-data Relevant = RRanked Type Type
-              | RShaped Type Type
-              | RMixed Type Type
-              | RShowable Type
+data Arg = RRanked Type Arg
+         | RShaped Type Arg
+         | RMixed Type Arg
+         | RShowable Type
+         | ROther Type
   deriving (Show)
 
--- | If so, returns the element type
-isRelevant :: Type -> Maybe Relevant
-isRelevant (ConT name `AppT` sht `AppT` ty)
-  | name == ''Ranked = Just (RRanked sht ty)
-  | name == ''Shaped = Just (RShaped sht ty)
-  | name == ''Mixed = Just (RMixed sht ty)
-isRelevant ty@(ConT name `AppT` _)
+-- TODO: always returns Just
+recognise :: Type -> Maybe Arg
+recognise (ConT name `AppT` sht `AppT` ty)
+  | name == ''Ranked = RRanked sht <$> recognise ty
+  | name == ''Shaped = RShaped sht <$> recognise ty
+  | name == ''Mixed = RMixed sht <$> recognise ty
+recognise ty@(ConT name `AppT` _)
   | name `elem` [''IShR, ''IIxR, ''ShS, ''IIxS, ''SNat] =
       Just (RShowable ty)
-isRelevant _ = Nothing
+recognise _ = Nothing
+
+realise :: Arg -> Type
+realise (RRanked sht ty) = ConT ''Ranked `AppT` sht `AppT` realise ty
+realise (RShaped sht ty) = ConT ''Shaped `AppT` sht `AppT` realise ty
+realise (RMixed sht ty) = ConT ''Mixed `AppT` sht `AppT` realise ty
+realise (RShowable ty) = ty
+realise (ROther ty) = ty
+
+mkShow :: Arg -> Cxt
+mkShow (RRanked _ ty) = mkShowElt ty
+mkShow (RShaped _ ty) = mkShowElt ty
+mkShow (RMixed sht ty) = [ConT ''Show `AppT` realise (RMixed sht ty)]
+mkShow (RShowable _) = []
+mkShow (ROther ty) = [ConT ''Show `AppT` ty]
+
+mkShowElt :: Arg -> Cxt
+mkShowElt (RRanked _ ty) = mkShowElt ty
+mkShowElt (RShaped _ ty) = mkShowElt ty
+mkShowElt (RMixed sht ty) = [ConT ''Show `AppT` realise (RMixed sht ty), ConT ''Elt `AppT` realise (RMixed sht ty)]
+mkShowElt (RShowable _ty) = []  -- [ConT ''Elt `AppT` ty]
+mkShowElt (ROther ty) = [ConT ''Show `AppT` ty, ConT ''Elt `AppT` ty]
 
 convertType :: Type -> Q (Type, [Bool], Bool)
 convertType typ =
   let (tybndrs, cxt, args, ret) = splitFunTy typ
-      argrels = map isRelevant args
-      retrel = isRelevant ret
-
-      showhead (RRanked n ty) = [ConT ''Mixed `AppT` (ConT ''Replicate `AppT` n `AppT` ConT 'Nothing) `AppT` ty]
-      showhead (RShaped sh ty) = [ConT ''Mixed `AppT` (ConT ''MapJust `AppT` sh) `AppT` ty]
-      showhead (RMixed sh ty) = [ConT ''Mixed `AppT` sh `AppT` ty]
-      showhead (RShowable _) = []
+      argrels = map recognise args
+      retrel = recognise ret
   in return
       (ForallT tybndrs
-               (cxt ++ [ConT ''Show `AppT` hd
+               (cxt ++ [constr
                        | Just rel <- retrel : argrels
-                       , hd <- showhead rel])
+                       , constr <- mkShow rel])
                (foldr (\a b -> ArrowT `AppT` a `AppT` b) ret args)
       ,map isJust argrels
       ,isJust retrel)