aboutsummaryrefslogtreecommitdiff
path: root/src/Data/Array
diff options
context:
space:
mode:
authorTom Smeding <tom@tomsmeding.com>2024-06-17 12:00:21 +0200
committerTom Smeding <tom@tomsmeding.com>2024-06-17 12:00:21 +0200
commit63b60c06674127e96cebfc3f1e8710f31df379d7 (patch)
treefd31849b42f7540067521fd20da05a5ccf099b86 /src/Data/Array
parent64f6a5552cffc10b9fbdb812166be8a0317367c0 (diff)
Update trace wrappers
Diffstat (limited to 'src/Data/Array')
-rw-r--r--src/Data/Array/Nested/Trace.hs4
-rw-r--r--src/Data/Array/Nested/Trace/TH.hs60
2 files changed, 40 insertions, 24 deletions
diff --git a/src/Data/Array/Nested/Trace.hs b/src/Data/Array/Nested/Trace.hs
index eadfeeb..6d820e4 100644
--- a/src/Data/Array/Nested/Trace.hs
+++ b/src/Data/Array/Nested/Trace.hs
@@ -55,11 +55,11 @@ module Data.Array.Nested.Trace (
NumElt, FloatElt,
) where
-import Prelude hiding (mappend)
+import Prelude hiding (mappend, mconcat)
import Data.Array.Nested
import Data.Array.Nested.Trace.TH
$(concat <$> mapM convertFun
- ['rshape , 'rrank, 'rindex, 'rindexPartial, 'rgenerate, 'rsumOuter1, 'rtranspose, 'rappend, 'rscalar, 'rfromVector, 'rtoVector, 'runScalar, 'rrerank, 'rreplicate, 'rreplicateScal, 'rfromListOuter, 'rfromList1, 'rfromList1Prim, 'rtoListOuter, 'rtoList1, 'rslice, 'rrev1, 'rreshape, 'riota, 'rlift, 'rlift2, 'rtoXArrayPrim, 'rfromXArrayPrim, 'rcastToShaped, 'rfromOrthotope, 'rtoOrthotope, 'sshape, 'sindex, 'sindexPartial, 'sgenerate, 'ssumOuter1, 'stranspose, 'sappend, 'sscalar, 'sfromVector, 'stoVector, 'sunScalar, 'srerank, 'sreplicate, 'sreplicateScal, 'sfromListOuter, 'sfromList1, 'sfromList1Prim, 'stoListOuter, 'stoList1, 'sslice, 'srev1, 'sreshape, 'siota, 'slift, 'slift2, 'stoXArrayPrim, 'sfromXArrayPrim, 'stoRanked, 'mshape, 'mindex, 'mindexPartial, 'mgenerate, 'msumOuter1, 'mtranspose, 'mappend, 'mscalar, 'mfromVector, 'mtoVector, 'munScalar, 'mrerank, 'mreplicate, 'mreplicateScal, 'mfromListOuter, 'mfromList1, 'mfromList1Prim, 'mtoListOuter, 'mtoList1, 'mslice, 'mrev1, 'mreshape, 'miota, 'mlift, 'mlift2, 'mtoXArrayPrim, 'mfromXArrayPrim, 'mtoRanked, 'mcastToShaped])
+ ['rshape, 'rrank, 'rsize, 'rindex, 'rindexPartial, 'rgenerate, 'rsumOuter1, 'rtranspose, 'rappend, 'rconcat, 'rscalar, 'rfromVector, 'rtoVector, 'runScalar, 'rrerank, 'rreplicate, 'rreplicateScal, 'rfromListOuter, 'rfromList1, 'rfromList1Prim, 'rtoListOuter, 'rtoList1, 'rfromListLinear, 'rtoListLinear, 'rslice, 'rrev1, 'rreshape, 'rflatten, 'riota, 'rminIndexPrim, 'rmaxIndexPrim, 'rdot, 'rdot1, 'rnest, 'runNest, 'rlift, 'rlift2, 'rtoXArrayPrim, 'rfromXArrayPrim, 'rcastToShaped, 'rfromOrthotope, 'rtoOrthotope, 'sshape, 'srank, 'ssize, 'sindex, 'sindexPartial, 'sgenerate, 'ssumOuter1, 'stranspose, 'sappend, 'sscalar, 'sfromVector, 'stoVector, 'sunScalar, 'srerank, 'sreplicate, 'sreplicateScal, 'sfromListOuter, 'sfromList1, 'sfromList1Prim, 'stoListOuter, 'stoList1, 'sfromListLinear, 'stoListLinear, 'sslice, 'srev1, 'sreshape, 'sflatten, 'siota, 'sminIndexPrim, 'smaxIndexPrim, 'sdot, 'sdot1, 'snest, 'sunNest, 'slift, 'slift2, 'stoXArrayPrim, 'sfromXArrayPrim, 'stoRanked, 'sfromOrthotope, 'stoOrthotope, 'mshape, 'mrank, 'msize, 'mindex, 'mindexPartial, 'mgenerate, 'msumOuter1, 'mtranspose, 'mappend, 'mconcat, 'mscalar, 'mfromVector, 'mtoVector, 'munScalar, 'mrerank, 'mreplicate, 'mreplicateScal, 'mfromListOuter, 'mfromList1, 'mfromList1Prim, 'mtoListOuter, 'mtoList1, 'mfromListLinear, 'mtoListLinear, 'mslice, 'mrev1, 'mreshape, 'mflatten, 'miota, 'mminIndexPrim, 'mmaxIndexPrim, 'mdot, 'mdot1, 'mnest, 'munNest, 'mlift, 'mlift2, 'mtoXArrayPrim, 'mfromXArrayPrim, 'mtoRanked, 'mcastToShaped])
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)