dph-test: Start on locked zipWiths
authorBen Lippmeier <benl@ouroborus.net>
Thu, 2 Aug 2012 07:08:34 +0000 (17:08 +1000)
committerBen Lippmeier <benl@ouroborus.net>
Thu, 2 Aug 2012 07:08:34 +0000 (17:08 +1000)
dph-test/fusion/Generic.hs [new file with mode: 0644]
dph-test/fusion/Stream.hs [new file with mode: 0644]
dph-test/fusion/Unboxed.hs [new file with mode: 0644]
dph-test/fusion/ZipWith.hs [new file with mode: 0644]

diff --git a/dph-test/fusion/Generic.hs b/dph-test/fusion/Generic.hs
new file mode 100644 (file)
index 0000000..73f6a71
--- /dev/null
@@ -0,0 +1,44 @@
+{-# LANGUAGE FlexibleContexts #-}
+module Generic where
+import Data.Vector.Generic               as G
+import Data.Vector.Generic.Base          as G
+import Data.Vector.Fusion.Stream.Monadic as S
+import Data.Vector.Fusion.Stream.Size    as S
+import Data.Vector.Fusion.Util           as S
+import Stream                            as S
+
+{-}
+lockedZipWith 
+        :: (Vector v a, Vector v b, Vector v c)
+        -> (a -> b -> c) -> v a -> v b -> v c
+
+lockedZipWith f as bs 
+        = unstream
+        $ S.lockedZipWith f (stream as) (stream bs)
+-}
+
+-- Locked Zips ----------------------------------------------------------------
+-- | Zip two vectors of the same length.
+--   If they do not have the same length then the result is undefined.
+lockedZip
+        :: (Vector v a, Vector v b, Vector v (a, b))
+        => v a -> v b
+        -> v (a, b)
+
+lockedZip aa bb
+        = unstream $ lockedStream2 aa bb 
+{-# INLINE [1] lockedZip #-}
+
+
+-- | Zip three vectors of the same length.
+--   If they do not have the same length then the result is undefined.
+lockedZip3
+        :: (Vector v a, Vector v b, Vector v c, Vector v (a, b, c))
+        => v a -> v b -> v c
+        -> v (a, b, c)
+
+lockedZip3 aa bb cc
+        = unstream $ lockedStream3 aa bb cc
+{-# INLINE [1] lockedZip3 #-}
+
+
diff --git a/dph-test/fusion/Stream.hs b/dph-test/fusion/Stream.hs
new file mode 100644 (file)
index 0000000..7cb695e
--- /dev/null
@@ -0,0 +1,119 @@
+{-# LANGUAGE FlexibleContexts #-}
+module Stream where
+
+import Data.Vector.Generic               as G
+import Data.Vector.Generic.Base          as G
+import Data.Vector.Fusion.Stream.Monadic as S
+import Data.Vector.Fusion.Stream.Size    as S
+import Data.Vector.Fusion.Util           as S
+
+
+-- Locked Streamers -----------------------------------------------------------
+
+-- | Stream two things.
+lockedStream2 
+        :: (Monad m, Vector v a, Vector v b)
+        => v a -> v b 
+        -> Stream m (a, b)
+
+lockedStream2 aa bb
+ = aa `seq` bb `seq` n `seq` (S.unfoldr get 0 `S.sized` Exact n)
+ where  n        = G.length aa
+
+        {-# INLINE [0] get #-}
+        get i   
+         | i >= n       = Nothing
+
+         | Box a        <- basicUnsafeIndexM aa i
+         , Box b        <- basicUnsafeIndexM bb i
+         = Just ((a, b), i + 1)
+{-# INLINE [1] lockedStream2 #-}
+
+
+-- | Stream three things.
+lockedStream3
+        :: (Monad m, Vector v a, Vector v b, Vector v c)
+        => v a -> v b -> v c 
+        -> Stream m (a, b, c)
+
+lockedStream3 aa bb cc
+ = aa `seq` bb `seq` cc
+ `seq` n `seq` (S.unfoldr get 0 `S.sized` Exact n)
+ where  n        = G.length aa
+
+        {-# INLINE [0] get #-}
+        get i   
+         | i >= n       = Nothing
+
+         | Box a        <- basicUnsafeIndexM aa i
+         , Box b        <- basicUnsafeIndexM bb i
+         , Box c        <- basicUnsafeIndexM cc i
+         = Just ((a, b, c), i + 1)
+{-# INLINE [1] lockedStream3 #-}
+
+
+-- Locked Stream Zippers -----------------------------------------------------
+lockedZip2S 
+        :: Monad m 
+        => Stream m a -> Stream m b 
+        -> Stream m (a, b)
+
+lockedZip2S (Stream mkStep1 sa1 size1)
+            (Stream mkStep2 sa2 _)
+ = Stream step (sa1, sa2) size1
+ where 
+        {-# INLINE [0] step #-}
+        step (s1, s2)
+         = do   step1   <- mkStep1 s1
+                step2   <- mkStep2 s2
+                return $ case (step1, step2) of
+                          (Yield x1 s1', Yield x2 s2')  -> Yield (x1, x2) (s1', s2')
+                          _                             -> Done
+{-# INLINE [1] lockedZip2S #-}
+
+
+lockedZip3S 
+        :: Monad m 
+        => Stream m a -> Stream m b -> Stream m c
+        -> Stream m (a, b, c)
+
+lockedZip3S (Stream mkStep1 sa1 size1)
+            (Stream mkStep2 sa2 _)
+            (Stream mkStep3 sa3 _)
+ = Stream step (sa1, sa2, sa3) size1
+ where 
+        {-# INLINE [0] step #-}
+        step (s1, s2, s3)
+         = do   step1   <- mkStep1 s1
+                step2   <- mkStep2 s2
+                step3   <- mkStep3 s3
+                return $ case (step1, step2, step3) of
+                          (Yield x1 s1', Yield x2 s2', Yield x3 s3')  
+                            -> Yield (x1, x2, x3) (s1', s2', s3')
+
+                          _ -> Done
+{-# INLINE [1] lockedZip3S #-}
+
+
+{-}
+lockedZip3 :: Vector a -> Vector b -> Vector c -> Vector (a, b, c)
+lockedZip3 aa bb cc 
+ = unstream $ lockedStream3 aa bb cc
+
+lockedZip3S :: Stream a -> Stream a -> Stremam a -> Stream (a, b, c)
+lockedZip3S
+
+thing aa bb cc
+ = lockedZip3 aa (map double bb) cc
+
+ = unstream $ lockedStream aa (unstream $ replicateS len aa) cc
+
+ = unstream $ lockedZip2S (lockedStream aa cc) (replicateS len aa)
+
+ If we see an unstream as one of the arguments of a 'lockedStreamN' then 
+ use rewrite rules to rotate it away and treat that as a separate stream.
+
+
+-}
+
diff --git a/dph-test/fusion/Unboxed.hs b/dph-test/fusion/Unboxed.hs
new file mode 100644 (file)
index 0000000..dbbd27c
--- /dev/null
@@ -0,0 +1,46 @@
+
+module Unboxed where
+import Data.Vector.Unboxed      as U
+import Stream
+import qualified Generic        as G
+
+
+-- Locked Zips ----------------------------------------------------------------
+lockedZip
+        :: (Unbox a, Unbox b)
+        => Vector a -> Vector b -> Vector (a, b)
+
+lockedZip = G.lockedZip
+{-# INLINE lockedZip #-}
+
+
+lockedZip3
+        :: (Unbox a, Unbox b, Unbox c)
+        => Vector a -> Vector b -> Vector c
+        -> Vector (a, b, c)
+
+lockedZip3 = G.lockedZip3
+{-# INLINE lockedZip3 #-}
+
+
+-- Locked ZipWiths ------------------------------------------------------------
+lockedZipWith2 
+        :: (Unbox a, Unbox b, Unbox c)
+        => (a -> b -> c)
+        -> Vector a -> Vector b -> Vector c
+
+lockedZipWith2 f aa bb
+        = U.map (\(a, b) -> f a b)
+        $ lockedZip aa bb
+{-# INLINE lockedZipWith2 #-}
+
+
+lockedZipWith3
+        :: (Unbox a, Unbox b, Unbox c, Unbox d)
+        => (a -> b -> c -> d)
+        -> Vector a -> Vector b -> Vector c -> Vector d
+
+lockedZipWith3 f aa bb cc
+        = U.map (\(a, b, c) -> f a b c)
+        $ lockedZip3 aa bb cc
+{-# INLINE lockedZipWith3 #-}
diff --git a/dph-test/fusion/ZipWith.hs b/dph-test/fusion/ZipWith.hs
new file mode 100644 (file)
index 0000000..70c751e
--- /dev/null
@@ -0,0 +1,26 @@
+
+
+module Test where
+import Data.Vector.Unboxed              as V
+import Prelude                          as P
+import Unboxed
+
+
+test    :: Vector Int
+        -> Vector Int
+        -> Vector Int
+        -> Vector (Int, Int, Int)
+
+test aa bb cc
+        = V.zipWith3 (,,) aa bb cc
+
+
+
+test2   :: Vector Int
+        -> Vector Int
+        -> Vector Int
+        -> Vector (Int, Int, Int)
+
+test2 aa bb cc
+        = lockedZip3 aa bb cc
+