1 files changed, 98 insertions, 0 deletions
diff --git a/src/Data/Array/Nested/Trace/TH.hs b/src/Data/Array/Nested/Trace/TH.hs
new file mode 100644
index 0000000..4b388e3
--- /dev/null
+++ b/src/Data/Array/Nested/Trace/TH.hs
@@ -0,0 +1,98 @@
+{-# LANGUAGE ImportQualifiedPost #-}
+{-# LANGUAGE LambdaCase #-}
+{-# LANGUAGE TemplateHaskellQuotes #-}
+module Data.Array.Nested.Trace.TH where
+
+import Control.Monad (zipWithM)
+import Data.List (foldl', intersperse)
+import Data.Maybe (isJust)
+import Language.Haskell.TH hiding (cxt)
+
+import Debug.Trace qualified as Debug
+
+import Data.Array.Nested
+
+
+splitFunTy :: Type -> ([TyVarBndr Specificity], Cxt, [Type], Type)
+splitFunTy = \case
+  ArrowT `AppT` t1 `AppT` t2 ->
+    let (vars, cx, args, ret) = splitFunTy t2
+    in (vars, cx, t1 : args, ret)
+  ForallT vs cx' t ->
+    let (vars, cx, args, ret) = splitFunTy t
+    in (vars ++ vs, cx ++ cx', args, ret)
+  t -> ([], [], [], t)
+
+data Arg = RRanked Type Arg
+         | RShaped Type Arg
+         | RMixed Type Arg
+         | RShowable Type
+         | 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
+recognise ty@(ConT name `AppT` _)
+  | name `elem` [''IShR, ''IIxR, ''ShS, ''IIxS, ''SNat] =
+      Just (RShowable ty)
+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 recognise args
+      retrel = recognise ret
+  in return
+      (ForallT tybndrs
+               (cxt ++ [constr
+                       | Just rel <- retrel : argrels
+                       , constr <- mkShow rel])
+               (foldr (\a b -> ArrowT `AppT` a `AppT` b) ret args)
+      ,map isJust argrels
+      ,isJust retrel)
+
+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 ..]
+  resname <- newName "res"
+  let tracenames = map fst (filter snd (zip (names ++ [resname]) (argarrs ++ [retarr])))
+  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)
+  return
+    [SigD defname convty
+    ,FunD defname [Clause (map VarP names) (NormalB ex) []]]