1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
|
{-# LANGUAGE TypeFamilies, TypeOperators, FlexibleInstances #-}
-- | Index space transformation between arrays and slices.
module Data.Array.Repa.Slice
( All (..)
, Any (..)
, FullShape
, SliceShape
, Slice (..))
where
import Data.Array.Repa.Index
import Prelude hiding (replicate, drop)
-- | Select all indices at a certain position.
data All = All
-- | Place holder for any possible shape.
data Any sh = Any
-- | Map a type of the index in the full shape, to the type of the index in the slice.
type family FullShape ss
type instance FullShape Z = Z
type instance FullShape (Any sh) = sh
type instance FullShape (sl :. Int) = FullShape sl :. Int
type instance FullShape (sl :. All) = FullShape sl :. Int
-- | Map the type of an index in the slice, to the type of the index in the full shape.
type family SliceShape ss
type instance SliceShape Z = Z
type instance SliceShape (Any sh) = sh
type instance SliceShape (sl :. Int) = SliceShape sl
type instance SliceShape (sl :. All) = SliceShape sl :. Int
-- | Class of index types that can map to slices.
class Slice ss where
-- | Map an index of a full shape onto an index of some slice.
sliceOfFull :: ss -> FullShape ss -> SliceShape ss
-- | Map an index of a slice onto an index of the full shape.
fullOfSlice :: ss -> SliceShape ss -> FullShape ss
instance Slice Z where
{-# INLINE [1] sliceOfFull #-}
sliceOfFull _ _ = Z
{-# INLINE [1] fullOfSlice #-}
fullOfSlice _ _ = Z
instance Slice (Any sh) where
{-# INLINE [1] sliceOfFull #-}
sliceOfFull _ sh = sh
{-# INLINE [1] fullOfSlice #-}
fullOfSlice _ sh = sh
instance Slice sl => Slice (sl :. Int) where
{-# INLINE [1] sliceOfFull #-}
sliceOfFull (fsl :. _) (ssl :. _)
= sliceOfFull fsl ssl
{-# INLINE [1] fullOfSlice #-}
fullOfSlice (fsl :. n) ssl
= fullOfSlice fsl ssl :. n
instance Slice sl => Slice (sl :. All) where
{-# INLINE [1] sliceOfFull #-}
sliceOfFull (fsl :. All) (ssl :. s)
= sliceOfFull fsl ssl :. s
{-# INLINE [1] fullOfSlice #-}
fullOfSlice (fsl :. All) (ssl :. s)
= fullOfSlice fsl ssl :. s
|
