Added iteraten and generaten functions
[matthijs/master-project/support/tfvec.git] / Data / Param / TFVec.hs
1 ------------------------------------------------------------------------------
2 -- |
3 -- Module       : Data.Param.TFVec
4 -- Copyright    : (c) 2009 Christiaan Baaij
5 -- Licence      : BSD-style (see the file LICENCE)
6 --
7 -- Maintainer   : christiaan.baaij@gmail.com
8 -- Stability    : experimental
9 -- Portability  : non-portable
10 --
11 -- 'TFVec': Fixed sized vectors. Vectors with numerically parameterized size,
12 --          using type-level numerals from 'tfp' library
13 --
14 ------------------------------------------------------------------------------
15
16 module Data.Param.TFVec
17   ( TFVec
18   , empty
19   , (+>)
20   , singleton
21   , vectorCPS
22   , vectorTH
23   , unsafeVector
24   , readTFVec
25   , length
26   , lengthT
27   , fromVector
28   , null
29   , (!)
30   , replace
31   , head
32   , last
33   , init
34   , tail
35   , take
36   , drop
37   , select
38   , (<+)
39   , (++)
40   , map
41   , zipWith
42   , foldl
43   , foldr
44   , zip
45   , unzip
46   , shiftl
47   , shiftr
48   , rotl
49   , rotr
50   , concat
51   , reverse
52   , iterate
53   , iteraten
54   , generate
55   , generaten
56   , copy
57   , copyn
58   ) where
59     
60
61 import Types
62 import Types.Data.Num.Decimal.Literals.TH
63 import Data.RangedWord
64
65 import Data.Generics (Data, Typeable)
66 import qualified Prelude as P
67 import Prelude hiding (
68   null, length, head, tail, last, init, take, drop, (++), map, foldl, foldr,
69   zipWith, zip, unzip, concat, reverse, iterate )
70 import qualified Data.Foldable as DF (Foldable, foldr)
71 import qualified Data.Traversable as DT (Traversable(traverse))
72 import Language.Haskell.TH
73 import Language.Haskell.TH.Syntax (Lift(..))
74
75 newtype (NaturalT s) => TFVec s a = TFVec {unTFVec :: [a]}
76   deriving (Eq, Typeable)
77
78 deriving instance (NaturalT s, Typeable s, Data s, Typeable a, Data a) => Data (TFVec s a)
79
80 -- ==========================
81 -- = Constructing functions =
82 -- ==========================
83                                                   
84 empty :: TFVec D0 a
85 empty = TFVec []
86
87 (+>) :: a -> TFVec s a -> TFVec (Succ s) a
88 x +> (TFVec xs) = TFVec (x:xs)
89
90 infix 5 +>
91
92 singleton :: a -> TFVec D1 a
93 singleton x = x +> empty
94
95 vectorCPS :: [a] -> (forall s . NaturalT s => TFVec s a -> w) -> w
96 vectorCPS xs = unsafeVectorCPS (toInteger (P.length xs)) xs
97
98 vectorTH :: Lift a => [a] -> ExpQ
99 vectorTH xs = (vectorCPS xs) lift
100
101 unsafeVector :: NaturalT s => s -> [a] -> TFVec s a
102 unsafeVector l xs
103   | fromIntegerT l /= P.length xs =
104     error (show 'unsafeVector P.++ ": dynamic/static lenght mismatch")
105   | otherwise = TFVec xs
106
107 readTFVec :: (Read a, NaturalT s) => String -> TFVec s a
108 readTFVec = read
109
110 readTFVecCPS :: Read a => String -> (forall s . NaturalT s => TFVec s a -> w) -> w
111 readTFVecCPS str = unsafeVectorCPS (toInteger l) xs
112  where fName = show 'readTFVecCPS
113        (xs,l) = case [(xs,l) | (xs,l,rest) <- readTFVecList str,  
114                            ("","") <- lexTFVec rest] of
115                        [(xs,l)] -> (xs,l)
116                        []   -> error (fName P.++ ": no parse")
117                        _    -> error (fName P.++ ": ambiguous parse")
118         
119 -- =======================
120 -- = Observing functions =
121 -- =======================
122 length :: forall s a . NaturalT s => TFVec s a -> Int
123 length _ = fromIntegerT (undefined :: s)
124
125 lengthT :: NaturalT s => TFVec s a -> s
126 lengthT = undefined
127
128 fromVector :: NaturalT s => TFVec s a -> [a]
129 fromVector (TFVec xs) = xs
130
131 null :: TFVec D0 a -> Bool
132 null _ = True
133
134 (!) ::  ( PositiveT s
135         , NaturalT u
136         , (s :>: u) ~ True) => TFVec s a -> RangedWord u -> a
137 (TFVec xs) ! i = xs !! (fromInteger (toInteger i))
138
139 -- ==========================
140 -- = Transforming functions =
141 -- ==========================
142 replace :: (PositiveT s, NaturalT u, (s :>: u) ~ True) =>
143   TFVec s a -> RangedWord u -> a -> TFVec s a
144 replace (TFVec xs) i y = TFVec $ replace' xs (toInteger i) y
145   where replace' []     _ _ = []
146         replace' (_:xs) 0 y = (y:xs)
147         replace' (x:xs) n y = x : (replace' xs (n-1) y)
148   
149 head :: PositiveT s => TFVec s a -> a
150 head = P.head . unTFVec
151
152 tail :: PositiveT s => TFVec s a -> TFVec (Pred s) a
153 tail = liftV P.tail
154
155 last :: PositiveT s => TFVec s a -> a
156 last = P.last . unTFVec
157
158 init :: PositiveT s => TFVec s a -> TFVec (Pred s) a
159 init = liftV P.init
160
161 take :: NaturalT i => i -> TFVec s a -> TFVec (Min s i) a
162 take i = liftV $ P.take (fromIntegerT i)
163
164 drop :: NaturalT i => i -> TFVec s a -> TFVec (s :-: (Min s i)) a
165 drop i = liftV $ P.drop (fromIntegerT i)
166
167 select :: (NaturalT f, NaturalT s, NaturalT n, (f :<: i) ~ True, 
168           (((s :*: n) :+: f) :<=: i) ~ True) => 
169           f -> s -> n -> TFVec i a -> TFVec n a
170 select f s n = liftV (select' f' s' n')
171   where (f', s', n') = (fromIntegerT f, fromIntegerT s, fromIntegerT n)
172         select' f s n = ((selectFirst0 s n).(P.drop f))
173         selectFirst0 :: Int -> Int -> [a] -> [a]
174         selectFirst0 s n l@(x:_)
175           | n > 0 = x : selectFirst0 s (n-1) (P.drop s l)
176           | otherwise = []
177         selectFirst0 _ 0 [] = []
178
179 (<+) :: TFVec s a -> a -> TFVec (Succ s) a
180 (<+) (TFVec xs) x = TFVec (xs P.++ [x])
181
182 (++) :: TFVec s a -> TFVec s2 a -> TFVec (s :+: s2) a
183 (++) = liftV2 (P.++)
184
185 infixl 5 <+
186 infixr 5 ++
187
188 map :: (a -> b) -> TFVec s a -> TFVec s b
189 map f = liftV (P.map f)
190
191 zipWith :: (a -> b -> c) -> TFVec s a -> TFVec s b -> TFVec s c
192 zipWith f = liftV2 (P.zipWith f)
193
194 foldl :: (a -> b -> a) -> a -> TFVec s b -> a
195 foldl f e = (P.foldl f e) . unTFVec
196
197 foldr :: (b -> a -> a) -> a -> TFVec s b -> a
198 foldr f e = (P.foldr f e) . unTFVec
199
200 zip :: TFVec s a -> TFVec s b -> TFVec s (a, b)
201 zip = liftV2 P.zip
202
203 unzip :: TFVec s (a, b) -> (TFVec s a, TFVec s b)
204 unzip (TFVec xs) = let (a,b) = P.unzip xs in (TFVec a, TFVec b)
205
206 shiftl :: (PositiveT s, NaturalT n, n ~ Pred s, s ~ Succ n) => 
207           TFVec s a -> a -> TFVec s a
208 shiftl xs x = x +> init xs
209
210 shiftr :: (PositiveT s, NaturalT n, n ~ Pred s, s ~ Succ n) => 
211           TFVec s a -> a -> TFVec s a
212 shiftr xs x = tail xs <+ x
213   
214 rotl :: forall s a . NaturalT s => TFVec s a -> TFVec s a
215 rotl = liftV rotl'
216   where vlen = fromIntegerT (undefined :: s)
217         rotl' [] = []
218         rotl' xs = let (i,[l]) = splitAt (vlen - 1) xs
219                    in l : i 
220
221 rotr :: NaturalT s => TFVec s a -> TFVec s a
222 rotr = liftV rotr'
223   where
224     rotr' [] = []
225     rotr' (x:xs) = xs P.++ [x] 
226
227 concat :: TFVec s1 (TFVec s2 a) -> TFVec (s1 :*: s2) a
228 concat = liftV (P.foldr ((P.++).unTFVec) [])
229
230 reverse :: TFVec s a -> TFVec s a
231 reverse = liftV P.reverse
232
233 iterate :: NaturalT s => (a -> a) -> a -> TFVec s a
234 iterate = iteraten (undefined :: s)
235
236 iteraten :: NaturalT s => s -> (a -> a) -> a -> TFVec s a
237 iteraten s f x = let s' = fromIntegerT s in TFVec (P.take s' $ P.iterate f x)
238
239 generate :: NaturalT s => (a -> a) -> a -> TFVec s a
240 generate = generaten (undefined :: s)
241
242 generaten :: NaturalT s => s -> (a -> a) -> a -> TFVec s a
243 generaten s f x = let s' = fromIntegerT s in TFVec (P.take s' $ P.tail $ P.iterate f x)
244
245 copy :: NaturalT s => a -> TFVec s a
246 copy x = copyn (undefined :: s) x
247
248 copyn :: NaturalT s => s -> a -> TFVec s a
249 copyn s x = iteraten s id x
250
251 -- =============
252 -- = Instances =
253 -- =============
254 instance Show a => Show (TFVec s a) where
255   showsPrec _ = showV.unTFVec
256     where showV []      = showString "<>"
257           showV (x:xs)  = showChar '<' . shows x . showl xs
258                             where showl []      = showChar '>'
259                                   showl (x:xs)  = showChar ',' . shows x .
260                                                   showl xs
261
262 instance (Read a, NaturalT nT) => Read (TFVec nT a) where
263   readsPrec _ str
264     | all fitsLength possibilities = P.map toReadS possibilities
265     | otherwise = error (fName P.++ ": string/dynamic length mismatch")
266     where 
267       fName = "Data.Param.TFVec.read"
268       expectedL = fromIntegerT (undefined :: nT)
269       possibilities = readTFVecList str
270       fitsLength (_, l, _) = l == expectedL
271       toReadS (xs, _, rest) = (TFVec xs, rest)
272       
273 instance NaturalT s => DF.Foldable (TFVec s) where
274  foldr = foldr
275  
276 instance NaturalT s => Functor (TFVec s) where
277  fmap = map
278
279 instance NaturalT s => DT.Traversable (TFVec s) where 
280   traverse f = (fmap TFVec).(DT.traverse f).unTFVec
281
282 instance (Lift a, NaturalT nT) => Lift (TFVec nT a) where
283   lift (TFVec xs) = [|  unsafeTFVecCoerse 
284                         $(decLiteralV (fromIntegerT (undefined :: nT))) 
285                         (TFVec xs) |]
286
287 -- ======================
288 -- = Internal Functions =
289 -- ======================
290 liftV :: ([a] -> [b]) -> TFVec nT a -> TFVec nT' b
291 liftV f = TFVec . f . unTFVec
292
293 liftV2 :: ([a] -> [b] -> [c]) -> TFVec s a -> TFVec s2 b -> TFVec s3 c
294 liftV2 f a b = TFVec (f (unTFVec a) (unTFVec b))
295
296 splitAtM :: Int -> [a] -> Maybe ([a],[a])
297 splitAtM n xs = splitAtM' n [] xs
298   where splitAtM' 0 xs ys = Just (xs, ys)
299         splitAtM' n xs (y:ys) | n > 0 = do
300           (ls, rs) <- splitAtM' (n-1) xs ys
301           return (y:ls,rs)
302         splitAtM' _ _ _ = Nothing
303
304 unsafeTFVecCoerse :: nT' -> TFVec nT a -> TFVec nT' a
305 unsafeTFVecCoerse _ (TFVec v) = (TFVec v)
306
307 unsafeVectorCPS :: forall a w . Integer -> [a] ->
308                         (forall s . NaturalT s => TFVec s a -> w) -> w
309 unsafeVectorCPS l xs f = reifyNaturalD l 
310                         (\(_ :: lt) -> f ((TFVec xs) :: (TFVec lt a)))
311
312 readTFVecList :: Read a => String -> [([a], Int, String)]
313 readTFVecList = readParen' False (\r -> [pr | ("<",s) <- lexTFVec r,
314                                               pr <- readl s])
315   where
316     readl   s = [([],0,t) | (">",t) <- lexTFVec s] P.++
317                             [(x:xs,1+n,u) | (x,t)       <- reads s,
318                                             (xs, n, u)  <- readl' t]
319     readl'  s = [([],0,t) | (">",t) <- lexTFVec s] P.++
320                             [(x:xs,1+n,v) | (",",t)   <- lex s,
321                                             (x,u)     <- reads t,
322                                             (xs,n,v)  <- readl' u]
323     readParen' b g  = if b then mandatory else optional
324       where optional r  = g r P.++ mandatory r
325             mandatory r = [(x,n,u) | ("(",s)  <- lexTFVec r,
326                                       (x,n,t) <- optional s,
327                                       (")",u) <- lexTFVec t]
328
329 -- Custom lexer for FSVecs, we cannot use lex directly because it considers
330 -- sequences of < and > as unique lexemes, and that breaks nested FSVecs, e.g.
331 -- <<1,2><3,4>>
332 lexTFVec :: ReadS String
333 lexTFVec ('>':rest) = [(">",rest)]
334 lexTFVec ('<':rest) = [("<",rest)]
335 lexTFVec str = lex str
336