Add Stream instances for L and LN

lehins · Sep 7, 2019 · 22c3776 · 22c3776
1 parent 9b925c8
commit 22c3776
Show file tree

Hide file tree

Showing 3 changed files with 41 additions and 1 deletion.
diff --git a/massiv-test/massiv-test.cabal b/massiv-test/massiv-test.cabal
@@ -51,9 +51,10 @@ test-suite tests
   main-is:            Main.hs
   other-modules:      Spec
                     , Test.Massiv.Core.IndexSpec
+                    , Test.Massiv.Core.ListSpec
                     , Test.Massiv.Core.SchedulerSpec
-                    , Test.Massiv.Array.MutableSpec
                     , Test.Massiv.Array.Delayed.StreamSpec
+                    , Test.Massiv.Array.MutableSpec
                       -- TODO: Below should be moved to Test.Massiv.Array
                     , Data.Massiv.Array.Delayed.InterleavedSpec
                     , Data.Massiv.Array.Delayed.PushSpec

diff --git a/massiv-test/tests/Test/Massiv/Core/ListSpec.hs b/massiv-test/tests/Test/Massiv/Core/ListSpec.hs
@@ -0,0 +1,21 @@
+{-# LANGUAGE AllowAmbiguousTypes #-}
+{-# LANGUAGE ConstraintKinds #-}
+{-# LANGUAGE FlexibleContexts #-}
+{-# LANGUAGE MonoLocalBinds #-}
+{-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE TypeApplications #-}
+{-# LANGUAGE TypeFamilies #-}
+{-# LANGUAGE TypeOperators #-}
+module Test.Massiv.Core.ListSpec (spec) where
+
+import Data.Massiv.Array
+import Test.Massiv.Core
+import Test.Massiv.Array.Delayed
+
+
+spec :: Spec
+spec = do
+  describe "L" $
+    it "toStream" $ property (prop_toStreamIsList @L @Int)
+  describe "LN" $
+    it "toStream" $ property (prop_toStreamIsList @LN @Int)
diff --git a/massiv/src/Data/Massiv/Core/List.hs b/massiv/src/Data/Massiv/Core/List.hs
@@ -32,6 +32,7 @@ import Control.Scheduler
 import Data.Coerce
 import Data.Foldable (foldr')
 import qualified Data.List as L
+import qualified Data.Massiv.Array.Manifest.Vector.Stream as S
 import Data.Massiv.Core.Common
 import Data.Typeable
 import GHC.Exts
@@ -48,6 +49,14 @@ type instance NestedStruct LN ix e = [ListItem ix e]
 newtype instance Array LN ix e = List { unList :: [Elt LN ix e] }
 
 
+instance Construct LN Ix1 e where
+  setComp _ = id
+  {-# INLINE setComp #-}
+  makeArray _ (Sz n) f = coerce (fmap f [0 .. n - 1])
+  {-# INLINE makeArray #-}
+  makeArrayLinear _ (Sz n) f = coerce (fmap f [0 .. n - 1])
+  {-# INLINE makeArrayLinear #-}
+
 instance {-# OVERLAPPING #-} Nested LN Ix1 e where
   fromNested = coerce
   {-# INLINE fromNested #-}
@@ -368,3 +377,12 @@ instance Ragged L ix e => OuterSlice L ix e where
           Just (x, _) | n == i -> x
           Just (_, xs) -> go (n + 1) xs
   {-# INLINE unsafeOuterSlice #-}
+
+
+instance Stream LN Ix1 e where
+  toStream = S.fromList . coerce
+  {-# INLINE toStream #-}
+
+instance Ragged L ix e => Stream L ix e where
+  toStream = S.fromList . coerce . lData . flattenRagged
+  {-# INLINE toStream #-}