NonEmpty.hs

{-# LANGUAGE NoImplicitPrelude #-}
{-# OPTIONS_GHC -Wall #-}

{- HLINT ignore "Avoid restricted function" -}

module Papa.Base.Data.List.NonEmpty (
  module P,
  head,
  tail,
  last,
  init,
) where

import Control.Lens (Lens')
import Data.Functor (Functor (fmap))
import Data.List.NonEmpty as P (
  NonEmpty ((:|)),
  group,
  group1,
  groupAllWith,
  groupAllWith1,
  groupBy,
  groupBy1,
  groupWith,
  groupWith1,
  inits,
  insert,
  iterate,
  nonEmpty,
  repeat,
  scanl,
  scanl1,
  scanr,
  scanr1,
  some1,
  tails,
  unfold,
  xor,
 )
import qualified Data.List.NonEmpty as NE
import Data.Semigroup ((<>))

-- | Lens to the head (first element) of a 'NonEmpty' list.
--
-- >>> import Control.Lens
-- >>> let ne = 1 :| [2,3,4]
-- >>> ne ^. head
-- 1
-- >>> ne & head .~ 99
-- 99 :| [2,3,4]
-- >>> ne & head %~ (\x -> x)
-- 1 :| [2,3,4]
{-# INLINE head #-}
head :: Lens' (NonEmpty a) a
head f (h :| t) = fmap (:| t) (f h)

-- | Lens to the tail (all elements except the first) of a 'NonEmpty' list.
--
-- >>> import Control.Lens
-- >>> let ne = 1 :| [2,3,4]
-- >>> ne ^. tail
-- [2,3,4]
-- >>> ne & tail .~ [99]
-- 1 :| [99]
-- >>> ne & tail .~ []
-- 1 :| []
{-# INLINE tail #-}
tail :: Lens' (NonEmpty a) [a]
tail f (h :| t) = fmap (h :|) (f t)

-- | Lens to the last element of a 'NonEmpty' list.
--
-- >>> import Control.Lens
-- >>> let ne = 1 :| [2,3,4]
-- >>> ne ^. last
-- 4
-- >>> ne & last .~ 99
-- 1 :| [2,3,99]
-- >>> (1 :| []) ^. last
-- 1
-- >>> (1 :| []) & last .~ 99
-- 99 :| []
{-# INLINE last #-}
last :: Lens' (NonEmpty a) a
last f (h :| []) = fmap (:| []) (f h)
last f (h :| (x : xs)) = fmap (\(y :| ys) -> h :| (y : ys)) (last f (x :| xs))

-- | Lens to the init (all elements except the last) of a 'NonEmpty' list.
--
-- >>> import Control.Lens
-- >>> let ne = 1 :| [2,3,4]
-- >>> ne ^. init
-- [1,2,3]
-- >>> ne & init .~ [99]
-- 99 :| [4]
-- >>> ne & init .~ []
-- 4 :| []
-- >>> (1 :| []) ^. init
-- []
-- >>> (1 :| []) & init .~ [10,20]
-- 10 :| [20,1]
{-# INLINE init #-}
init :: Lens' (NonEmpty a) [a]
init f ne = fmap rebuild (f (NE.init ne))
 where
  rebuild [] = NE.last ne :| []
  rebuild (x : xs) = x :| (xs <> [NE.last ne])