Skip to content

Instantly share code, notes, and snippets.

@nomeata
Forked from sacundim/TakeR.hs
Last active December 19, 2015 19:39
Show Gist options
  • Save nomeata/6007713 to your computer and use it in GitHub Desktop.
Save nomeata/6007713 to your computer and use it in GitHub Desktop.
Fix benchmarking (avoid sharing)
{-# LANGUAGE RankNTypes, ScopedTypeVariables #-}
import Control.Monad.ST
import Data.Foldable as F
import Data.Array.ST
import Data.Sequence (Seq)
import qualified Data.Sequence as Seq
import Criterion.Main
main = defaultMain [
makeGroup "takeRNaive" takeRNaive (\n -> [1..n]),
makeGroup "takeRNaive2" takeRNaive2 (\n -> [1..n]),
makeGroup "takeRArray" takeRArray (\n -> [1..n]),
makeGroup "takeRArray2" takeRArray2 (\n -> [1..n]),
makeGroup "takeRIdiomatic" takeRIdiomatic (\n -> [1..n]),
makeGroup "takeRSeq" takeRSeq (\n -> Seq.fromList [1..n]),
makeGroup "takeRInternalSeq" takeRInternalSeq (\n -> [1..n])
]
{-# INLINE makeGroup #-}
makeGroup name takeR makeList =
bgroup name
[ bench "100000" $ whnf (takeR 100000 . makeList) 100000
, bench "100" $ whnf (takeR 100 . makeList) 100000
]
takeRNaive :: Int -> [a] -> [a]
takeRNaive n = reverse . take n . reverse
-- From: http://www.reddit.com/r/haskell/comments/1i64mt/on_taking_the_last_n_elements_of_a_list/cb1ecdp
takeRNaive2 :: Int -> [a] -> [a]
takeRNaive2 n xs = drop (length xs - n) xs
takeRArray :: forall a. Int -> [a] -> [a]
takeRArray n l | n <= 0 = []
takeRArray n l = runST stAction
where
stAction :: forall s. ST s [a]
stAction = do
buffer <- newArray_ (0, n-1)
i <- go (buffer :: STArray s Int a) 0 l
let s = min i n
sequence $ [ readArray buffer (j `mod` n) | j <- [i-s..i-1] ]
go buffer i [] = return i
go buffer i (x:xs) = writeArray buffer (i `mod` n) x >> go buffer (i+1) xs
takeRArray2 :: forall a. Int -> [a] -> [a]
takeRArray2 n l | n <= 0 = []
takeRArray2 (-1) l = []
takeRArray2 0 l = []
takeRArray2 n l = runST stAction
where
stAction :: forall s. ST s [a]
stAction = do
buffer <- newArray_ (0, n)
i <- go (buffer :: STArray s Int [a]) 0 l
let s = min i n
if s <= 0 then return [] else readArray buffer ((i-s) `mod` n)
go buffer i [] = return i
go buffer i l@(x:xs) = writeArray buffer (i `mod` n) l >> go buffer (i+1) xs
takeRIdiomatic :: Int -> [a] -> [a]
takeRIdiomatic n l = go (drop n l) l
where
go [] r = r
go (_:xs) (_:ys) = go xs ys
takeRSeq :: Int -> Seq a -> Seq a
{-# INLINE takeRSeq #-}
takeRSeq n xs = Seq.drop (Seq.length xs - n) xs
takeRInternalSeq :: Int -> [a] -> [a]
takeRInternalSeq n = F.toList . takeRSeq n . Seq.fromList
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment