From cfe31c2dc6869466da16dcb61de39218db3941d4 Mon Sep 17 00:00:00 2001 From: flip111 Date: Thu, 27 Mar 2025 18:04:50 +0100 Subject: [PATCH] Add intersperse --- src/Data/Vector/NonEmpty.hs | 52 ++++++++++++++++++---- test/Main.hs | 87 ++++++++++++++++++++++++++++++++----- 2 files changed, 121 insertions(+), 18 deletions(-) diff --git a/src/Data/Vector/NonEmpty.hs b/src/Data/Vector/NonEmpty.hs index 9901fbf..c44eac6 100644 --- a/src/Data/Vector/NonEmpty.hs +++ b/src/Data/Vector/NonEmpty.hs @@ -1,13 +1,13 @@ {-# LANGUAGE NoImplicitPrelude #-} {-# LANGUAGE RankNTypes #-} -- | --- Module : Data.Vector.NonEmpty --- Copyright : (c) 2019-2024 Emily Pillmore --- License : BSD-style +-- Module : Data.Vector.NonEmpty +-- Copyright: (c) 2019-2025 Emily Pillmore +-- License : BSD-style -- --- Maintainer : Emily Pillmore --- Stability : Experimental --- Portability : DataTypeable, CPP +-- Maintainer : Emily Pillmore +-- Stability : Experimental +-- Portability: DataTypeable, CPP -- -- A library for non-empty boxed vectors (that is, polymorphic arrays capable of -- holding any Haskell value). Non-empty vectors come in two flavors: @@ -180,11 +180,15 @@ module Data.Vector.NonEmpty , scanl, scanl', scanl1, scanl1', iscanl, iscanl' , prescanr, prescanr', postscanr, postscanr' , scanr, scanr', scanr1, scanr1', iscanr, iscanr' + + -- ** Transformations +, intersperse ) where import Prelude ( Bool, Eq, Ord, Num, Enum - , (.), Ordering, max, uncurry, snd) + , (.), Ordering, max, uncurry, snd + , pure, (+), otherwise, (>=), (-), (*), ($), (<=)) import Control.Monad (Monad) import Control.Monad.ST @@ -201,7 +205,7 @@ import Data.Traversable (Traversable) import Data.Vector (Vector) import qualified Data.Vector as V import qualified Data.Vector.Generic as G -import Data.Vector.Mutable (MVector) +import Data.Vector.Mutable (MVector, unsafeWrite) import Data.Vector.NonEmpty.Internal @@ -2646,3 +2650,35 @@ iscanr f b = NonEmptyVector . V.iscanr f b . _neVec iscanr' :: (Int -> a -> b -> b) -> b -> NonEmptyVector a -> NonEmptyVector b iscanr' f b = NonEmptyVector . V.iscanr' f b . _neVec {-# INLINE iscanr' #-} + +-- | /O(n)/ The 'intersperse' function takes an element and a NonEmptyVector +-- and 'intersperses' that element between the elements of the NonEmptyVector. +-- +-- >>> intersperse 0 (unsafeFromList [1,2,3]) +-- [1,0,2,0,3] +-- +-- >>> intersperse 0 (singleton 1) +-- [1] +-- +intersperse :: a -> NonEmptyVector a -> NonEmptyVector a +intersperse sep nev@(NonEmptyVector v) + | V.length v <= 1 = nev + | otherwise = unsafeCreate $ do + let n = V.length v + let newLen = 2*n - 1 + mv <- V.unsafeThaw $ V.replicate newLen (V.unsafeHead v) -- Create mutable vector + + -- Fill the first element + unsafeWrite mv 0 (V.unsafeIndex v 0) + + -- Fill remaining elements with separators + let go i j + | j >= n = pure () + | otherwise = do + unsafeWrite mv i sep + unsafeWrite mv (i+1) (V.unsafeIndex v j) + go (i+2) (j+1) + + go 1 1 + pure mv +{-# INLINE intersperse #-} \ No newline at end of file diff --git a/test/Main.hs b/test/Main.hs index 85182f4..93f12fd 100644 --- a/test/Main.hs +++ b/test/Main.hs @@ -1,13 +1,13 @@ {-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE RankNTypes #-} -- | --- Module : Main (tests) --- Copyright : 2019-2023 (c) Emily Pillmore --- License : BSD +-- Module : Main (tests) +-- Copyright: 2019-2025 (c) Emily Pillmore +-- License : BSD -- --- Maintainer : Emily Pillmore --- Stability : Experimental --- Portability : TypeFamilies +-- Maintainer : Emily Pillmore +-- Stability : Experimental +-- Portability: TypeFamilies -- module Main ( main @@ -16,6 +16,7 @@ module Main import Data.Maybe import Data.Vector (Vector) +import qualified Data.Vector as V import Data.Vector.NonEmpty (NonEmptyVector) import qualified Data.Vector.NonEmpty as NEV @@ -24,10 +25,20 @@ import Test.Tasty import Test.Tasty.QuickCheck main :: IO () -main = defaultMain $ testGroup "NonEmptyVector constructor" - [ testProperty "prop_reverse" prop_reverse - , testProperty "prop_from_to_list" prop_from_to_list - , testProperty "prop_from_to_vec" prop_from_to_vec +main = defaultMain $ testGroup "NonEmptyVector tests" + [ testGroup "Basic operations" + [ testProperty "prop_reverse" prop_reverse + , testProperty "prop_from_to_list" prop_from_to_list + , testProperty "prop_from_to_vec" prop_from_to_vec + ] + , testGroup "Intersperse tests" + [ testProperty "prop_intersperse_length" prop_intersperse_length + , testProperty "prop_intersperse_first_last" prop_intersperse_first_last + , testProperty "prop_intersperse_separator" prop_intersperse_separator + , testProperty "prop_intersperse_reconstruction" prop_intersperse_reconstruction + , testProperty "prop_intersperse_singleton" prop_intersperse_singleton + , testProperty "prop_explicit_examples" prop_explicit_examples + ] ] genList :: Gen [Int] @@ -56,3 +67,59 @@ prop_from_to_vec = forAll genV $ \u -> NEV.fromVector (NEV.toVector t) == Just t && (NEV.toVector <$> NEV.fromVector u) == Just u + +-- Helper function to reduce duplication in intersperse tests +withRandomVectorAndSep :: (NonEmptyVector Int -> Int -> Bool) -> Property +withRandomVectorAndSep f = forAll genNEV $ \v -> + forAll (arbitrary :: Gen Int) $ \sep -> f v sep + +-- Intersperse properties + +-- For vectors with length > 1, interspersing should result in length 2*n - 1 +prop_intersperse_length :: Property +prop_intersperse_length = withRandomVectorAndSep $ \v sep -> + let result = NEV.intersperse sep v + originalLen = NEV.length v + in if originalLen > 1 + then NEV.length result == 2 * originalLen - 1 + else NEV.length result == originalLen + +-- The first and last elements of the result should match the original vector +prop_intersperse_first_last :: Property +prop_intersperse_first_last = withRandomVectorAndSep $ \v sep -> + let result = NEV.intersperse sep v + in NEV.head result == NEV.head v && NEV.last result == NEV.last v + +-- Every odd-indexed element should be the separator (for vectors with length > 1) +prop_intersperse_separator :: Property +prop_intersperse_separator = withRandomVectorAndSep $ \v sep -> + let result = NEV.intersperse sep v + in NEV.length v <= 1 || + all (\i -> result NEV.! (2*i+1) == sep) [0 .. NEV.length v - 2] + +-- If we remove the separator, we should get back the original vector +prop_intersperse_reconstruction :: Property +prop_intersperse_reconstruction = withRandomVectorAndSep $ \v sep -> + let result = NEV.intersperse sep v + resultVector = NEV.toVector result + reconstructed = V.ifilter (\i _ -> i `mod` 2 == 0) resultVector + in V.toList reconstructed == NEV.toList v + +-- For singleton vectors, interspersing should not change the vector +prop_intersperse_singleton :: Property +prop_intersperse_singleton = + forAll (arbitrary :: Gen Int) $ \x -> + forAll (arbitrary :: Gen Int) $ \sep -> + let singleton = NEV.singleton x + result = NEV.intersperse sep singleton + in result == singleton + +-- Explicit test cases for edge cases and examples combined into one property +prop_explicit_examples :: Property +prop_explicit_examples = + let normalCase = NEV.intersperse (0 :: Int) (NEV.unsafeFromList [1,2,3 :: Int]) + === NEV.unsafeFromList [1,0,2,0,3 :: Int] + + singletonCase = NEV.intersperse (0 :: Int) (NEV.singleton (5 :: Int)) + === NEV.singleton (5 :: Int) + in normalCase .&&. singletonCase \ No newline at end of file