74be5ad · streamly · git.iko.soy

git.iko.soy/streamly

on master help sign in

Move lastN from StreamD to Array.Storable.Foreign

74be5adAdithya Kumar6 years ago

src/Streamly/Internal/Data/Array/Storable/Foreign.hs+30−3

compare with master

unifiedsplit

@@ -135,7 +135,7 @@ module Streamly.Internal.Data.Array.Storable.Foreign

135 , fold

136

137 -- * Folds with Array as the container

138− , D.lastN

138+ , lastN

139 )

140where

141

@@ -157,17 +157,21 @@ import GHC.Ptr (Ptr(..))

157import GHC.Prim (touch#)

158import GHC.IO (IO(..))

159

160−import Streamly.Internal.Data.Fold.Types (Fold(..))

161−import Streamly.Internal.Data.Unfold.Types (Unfold(..))

162import Streamly.Internal.Data.Array.Storable.Foreign.Types (Array(..), length)

160import Streamly.Internal.Data.Array.Storable.Foreign.Types (Array(..), length)

161+import Streamly.Internal.Data.Fold.Types (Fold(..))

163import Streamly.Internal.Data.Stream.Serial (SerialT)

162import Streamly.Internal.Data.Stream.Serial (SerialT)

164import Streamly.Internal.Data.Stream.StreamK.Type (IsStream)

163import Streamly.Internal.Data.Stream.StreamK.Type (IsStream)

164+import Streamly.Internal.Data.Tuple.Strict (Tuple3'(..))

165+import Streamly.Internal.Data.Unfold.Types (Unfold(..))

165

166

167+import qualified Streamly.Internal.Data.Array.Storable.Foreign.Mut.Types as MA

166import qualified Streamly.Internal.Data.Array.Storable.Foreign.Types as A

168import qualified Streamly.Internal.Data.Array.Storable.Foreign.Types as A

169+import qualified Streamly.Internal.Data.Fold as FL

167import qualified Streamly.Internal.Data.Stream.Prelude as P

170import qualified Streamly.Internal.Data.Stream.Prelude as P

168import qualified Streamly.Internal.Data.Stream.Serial as Serial

171import qualified Streamly.Internal.Data.Stream.Serial as Serial

169import qualified Streamly.Internal.Data.Stream.StreamD as D

172import qualified Streamly.Internal.Data.Stream.StreamD as D

170import qualified Streamly.Internal.Data.Stream.StreamK as K

173import qualified Streamly.Internal.Data.Stream.StreamK as K

174+import qualified Streamly.Memory.Ring as RB

171

175

172-------------------------------------------------------------------------------

176-------------------------------------------------------------------------------

173-- Construction

177-- Construction

@@ -301,6 +305,29 @@ null arr = length arr == 0

301last :: Storable a => Array a -> Maybe a

305last :: Storable a => Array a -> Maybe a

302last arr = readIndex arr (length arr - 1)

306last arr = readIndex arr (length arr - 1)

303

307

308+-------------------------------------------------------------------------------

309+-- Folds with Array as the container

310+-------------------------------------------------------------------------------

311+

312+-- | Take last 'n' elements from the stream and discard the rest.

313+{-# INLINE lastN #-}

314+lastN :: (Storable a, MonadIO m) => Int -> Fold m a (Array a)

315+lastN n

316+ | n <= 0 = fmap (const mempty) FL.drain

317+ | otherwise = A.unsafeFreeze <$> Fold step initial done

318+ where

319+ step (Tuple3' rb rh i) a = do

320+ rh1 <- liftIO $ RB.unsafeInsert rb rh a

321+ return $ FL.Partial $ Tuple3' rb rh1 (i + 1)

322+ initial = fmap (\(a, b) -> Tuple3' a b (0 :: Int)) $ liftIO $ RB.new n

323+ done (Tuple3' rb rh i) = do

324+ arr <- liftIO $ MA.newArray n

325+ foldFunc i rh snoc' arr rb

326+ snoc' b a = liftIO $ MA.unsafeSnoc b a

327+ foldFunc i

328+ | i < n = RB.unsafeFoldRingM

329+ | otherwise = RB.unsafeFoldRingFullM

330+

304-------------------------------------------------------------------------------

331-------------------------------------------------------------------------------

305-- Random Access

332-- Random Access

306-------------------------------------------------------------------------------

333-------------------------------------------------------------------------------

src/Streamly/Internal/Data/Stream/StreamD.hs−24

@@ -299,8 +299,6 @@ module Streamly.Internal.Data.Stream.StreamD

299 , mkParallel

300 , mkParallelD

301 , newCallbackStream

302−

303− , lastN

304 )

302 )

305where

303where

306

304

@@ -337,13 +335,11 @@ import Streamly.Internal.Data.Time.Clock (Clock(Monotonic), getTime)

337import Streamly.Internal.Data.Time.Units

335import Streamly.Internal.Data.Time.Units

338 (MicroSecond64(..), fromAbsTime, toAbsTime, AbsTime)

336 (MicroSecond64(..), fromAbsTime, toAbsTime, AbsTime)

339import Streamly.Internal.Data.Unfold.Types (Unfold(..))

337import Streamly.Internal.Data.Unfold.Types (Unfold(..))

340−import Streamly.Internal.Data.Tuple.Strict (Tuple3'(..))

341import Streamly.Internal.Data.Stream.SVar (fromConsumer, pushToFold)

338import Streamly.Internal.Data.Stream.SVar (fromConsumer, pushToFold)

342

339

343import qualified Streamly.Internal.Data.IORef.Prim as Prim

340import qualified Streamly.Internal.Data.IORef.Prim as Prim

344import qualified Streamly.Internal.Data.Pipe.Types as Pipe

341import qualified Streamly.Internal.Data.Pipe.Types as Pipe

345import qualified Streamly.Internal.Data.Array.Storable.Foreign.Types as A

342import qualified Streamly.Internal.Data.Array.Storable.Foreign.Types as A

346−import qualified Streamly.Internal.Data.Array.Storable.Foreign.Mut.Types as MA

347import qualified Streamly.Internal.Data.Fold as FL

343import qualified Streamly.Internal.Data.Fold as FL

348import qualified Streamly.Memory.Ring as RB

344import qualified Streamly.Memory.Ring as RB

349import qualified Streamly.Internal.Data.Stream.StreamK as K

345import qualified Streamly.Internal.Data.Stream.StreamK as K

@@ -4819,26 +4815,6 @@ tapAsync f (Stream step1 state1) = Stream step TapInit

4819 Skip s -> Skip (TapDone s)

4815 Skip s -> Skip (TapDone s)

4820 Stop -> Stop

4816 Stop -> Stop

4821

4817

4822−-- XXX Exported from Array again as this fold is specific to Array

4823−-- | Take last 'n' elements from the stream and discard the rest.

4824−{-# INLINE lastN #-}

4825−lastN :: (Storable a, MonadIO m) => Int -> Fold m a (Array a)

4826−lastN n

4827− | n <= 0 = fmap (const mempty) FL.drain

4828− | otherwise = A.unsafeFreeze <$> Fold step initial done

4829− where

4830− step (Tuple3' rb rh i) a = do

4831− rh1 <- liftIO $ RB.unsafeInsert rb rh a

4832− return $ FL.Partial $ Tuple3' rb rh1 (i + 1)

4833− initial = fmap (\(a, b) -> Tuple3' a b (0 :: Int)) $ liftIO $ RB.new n

4834− done (Tuple3' rb rh i) = do

4835− arr <- liftIO $ MA.newArray n

4836− foldFunc i rh snoc' arr rb

4837− snoc' b a = liftIO $ MA.unsafeSnoc b a

4838− foldFunc i

4839− | i < n = RB.unsafeFoldRingM

4840− | otherwise = RB.unsafeFoldRingFullM

4841−

4842------------------------------------------------------------------------------

4818------------------------------------------------------------------------------

4843-- Time related

4819-- Time related

4844------------------------------------------------------------------------------

4820------------------------------------------------------------------------------