trace: Print arguments and return value

2 files changed, 54 insertions, 25 deletions

diff --git a/src/Data/Array/Nested/Trace.hs b/src/Data/Array/Nested/Trace.hs
index 37c31ea..f89d3e4 100644
--- a/src/Data/Array/Nested/Trace.hs
+++ b/src/Data/Array/Nested/Trace.hs
@@ -5,6 +5,7 @@
 {-# LANGUAGE PatternSynonyms #-}
 {-# LANGUAGE RankNTypes #-}
 {-# LANGUAGE TemplateHaskell #-}
+{-# OPTIONS -Wno-simplifiable-class-constraints #-}
 {-|
 This module is API-compatible with "Data.Array.Nested", except that inputs and
 outputs of the methods are traced using 'Debug.Trace.trace'. Thus the methods
diff --git a/src/Data/Array/Nested/Trace/TH.hs b/src/Data/Array/Nested/Trace/TH.hs
index 4b388e3..6fded46 100644
--- a/src/Data/Array/Nested/Trace/TH.hs
+++ b/src/Data/Array/Nested/Trace/TH.hs
@@ -4,11 +4,11 @@
 module Data.Array.Nested.Trace.TH where
 
 import Control.Monad (zipWithM)
-import Data.List (foldl', intersperse)
+import Data.List (foldl')
 import Data.Maybe (isJust)
 import Language.Haskell.TH hiding (cxt)
-
-import Debug.Trace qualified as Debug
+import System.IO (hPutStr, stderr)
+import System.IO.Unsafe (unsafePerformIO)
 
 import Data.Array.Nested
 
@@ -20,7 +20,7 @@ splitFunTy = \case
     in (vars, cx, t1 : args, ret)
   ForallT vs cx' t ->
     let (vars, cx, args, ret) = splitFunTy t
-    in (vars ++ vs, cx ++ cx', args, ret)
+    in (vs ++ vars, cx' ++ cx, args, ret)
   t -> ([], [], [], t)
 
 data Arg = RRanked Type Arg
@@ -30,17 +30,24 @@ data Arg = RRanked Type Arg
          | ROther Type
   deriving (Show)
 
--- 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
+  | name == ''Ranked = Just (RRanked sht (recogniseElt ty))
+  | name == ''Shaped = Just (RShaped sht (recogniseElt ty))
+  | name == ''Mixed = Just (RMixed sht (recogniseElt ty))
+  | name == ''Conversion = Just (RShowable ty)
 recognise ty@(ConT name `AppT` _)
   | name `elem` [''IShR, ''IIxR, ''ShS, ''IIxS, ''SNat] =
       Just (RShowable ty)
 recognise _ = Nothing
 
+recogniseElt :: Type -> Arg
+recogniseElt (ConT name `AppT` sht `AppT` ty)
+  | name == ''Ranked = RRanked sht (recogniseElt ty)
+  | name == ''Shaped = RShaped sht (recogniseElt ty)
+  | name == ''Mixed = RMixed sht (recogniseElt ty)
+recogniseElt ty = ROther ty
+
 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
@@ -62,37 +69,58 @@ mkShowElt (RMixed sht ty) = [ConT ''Show `AppT` realise (RMixed sht ty), ConT ''
 mkShowElt (RShowable _ty) = []  -- [ConT ''Elt `AppT` ty]
 mkShowElt (ROther ty) = [ConT ''Show `AppT` ty, ConT ''Elt `AppT` ty]
 
-convertType :: Type -> Q (Type, [Bool], Bool)
+-- If you pass a polymorphic function to seq, GHC wants to monomorphise and
+-- doesn't know how to instantiate the type variables. Just don't, I guess.
+isSeqable :: Type -> Bool
+isSeqable ForallT{} = False
+isSeqable (AppT a b) = isSeqable a && isSeqable b
+isSeqable _ = True  -- yolo, I guess
+
+convertType :: Type -> Q (Type, [Bool], [Bool], Bool)
 convertType typ =
   let (tybndrs, cxt, args, ret) = splitFunTy typ
-      argrels = map recognise args
-      retrel = recognise ret
+      argdescrs = map recognise args
+      retdescr = recognise ret
   in return
       (ForallT tybndrs
                (cxt ++ [constr
-                       | Just rel <- retrel : argrels
+                       | Just rel <- retdescr : argdescrs
                        , constr <- mkShow rel])
                (foldr (\a b -> ArrowT `AppT` a `AppT` b) ret args)
-      ,map isJust argrels
-      ,isJust retrel)
+      ,map isJust argdescrs
+      ,map isSeqable args
+      ,isJust retdescr)
 
 convertFun :: Name -> Q [Dec]
 convertFun funname = do
   defname <- newName (nameBase funname)
-  (convty, argarrs, retarr) <- reifyType funname >>= convertType
-  names <- zipWithM (\b i -> newName ((if b then "t" else "x") ++ show i)) argarrs [1::Int ..]
+  -- "ok": whether we understand this type enough to be able to show it
+  (convty, argoks, argsseqable, retok) <- reifyType funname >>= convertType
+  names <- zipWithM (\_ i -> newName ('x' : show i)) argoks [1::Int ..]
+  -- let tracenames = map fst (filter snd (zip (names ++ [resname]) (argarrs ++ [retarr])))
   resname <- newName "res"
-  let tracenames = map fst (filter snd (zip (names ++ [resname]) (argarrs ++ [retarr])))
+  let traceCall str val = VarE 'traceNoNewline `AppE` str `AppE` val
+  let msg1 = [LitE (StringL ("oxtrace: (" ++ nameBase funname ++ " "))] ++
+             [if ok
+                then VarE 'showsPrec `AppE` LitE (IntegerL 11) `AppE` VarE n `AppE` LitE (StringL " ")
+                else LitE (StringL "_ ")
+             | (n, ok) <- zip names argoks] ++
+             [LitE (StringL "...")]
+  let msg2 | retok = [LitE (StringL " = "), VarE 'show `AppE` VarE resname, LitE (StringL ")\n")]
+           | otherwise = [LitE (StringL " = _)\n")]
   let ex = LetE [ValD (VarP resname)
                       (NormalB (foldl' AppE (VarE funname) (map VarE names)))
-                      []]
-                (VarE 'Debug.trace
-                  `AppE` (VarE 'concat `AppE` ListE
-                            ([LitE (StringL ("oxtrace: " ++ nameBase funname ++ " ["))] ++
-                             intersperse (LitE (StringL ", "))
-                                         (map (\n -> VarE 'show `AppE` VarE n) tracenames) ++
-                             [LitE (StringL "]")]))
-                  `AppE` VarE resname)
+                      []] $
+             flip (foldr AppE) [VarE 'seq `AppE` VarE n | (n, True) <- zip names argsseqable] $
+             traceCall (VarE 'concat `AppE` ListE msg1) $
+             VarE 'seq `AppE` VarE resname `AppE`
+             traceCall (VarE 'concat `AppE` ListE msg2) (VarE resname)
   return
     [SigD defname convty
     ,FunD defname [Clause (map VarP names) (NormalB ex) []]]
+
+{-# NOINLINE traceNoNewline #-}
+traceNoNewline :: String -> a -> a
+traceNoNewline str x = unsafePerformIO $ do
+  hPutStr stderr str
+  return x