Add intersperse

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 <emilypi@cohomolo.gy>
--- Stability	: Experimental
--- Portability	: DataTypeable, CPP
+-- Maintainer : Emily Pillmore <emilypi@cohomolo.gy>
+-- 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,18 +180,22 @@
 , 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
 
 import qualified Data.Foldable as Foldable
 import Data.Either (Either(..))
 import Data.Functor hiding (unzip)
 import Data.Int
 import Data.List.NonEmpty (NonEmpty(..))
 import qualified Data.List.NonEmpty as NonEmpty
@@ -201,7 +205,7 @@
 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' :: (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 #-}

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 <emilypi@cohomolo.gy>
--- Stability	: Experimental
--- Portability	: TypeFamilies
+-- Maintainer : Emily Pillmore <emilypi@cohomolo.gy>
+-- 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