-- |
--
-- Module      : Data.Vector.Ext
-- Copyright   : (c) 2025 Patrick Brisbin
-- License     : AGPL-3
-- Maintainer  : pbrisbin@gmail.com
-- Stability   : experimental
-- Portability : POSIX
module Data.Vector.Ext
  ( deleteAt
  , insertAt
  ) where

import Prelude

import Data.Vector (Vector)
import Data.Vector qualified as V

-- | Delete from a vector and shift all later elements left
deleteAt :: Int -> Vector a -> Vector a
deleteAt :: forall a. Int -> Vector a -> Vector a
deleteAt Int
n Vector a
vec
  | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 = Vector a
vec
  | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Vector a -> Int
forall a. Vector a -> Int
V.length Vector a
vec = Vector a
vec
  | Bool
otherwise = Vector (Int, a) -> Vector a
forall a. Vector (Int, a) -> Vector a
generateFrom (Vector (Int, a) -> Vector a) -> Vector (Int, a) -> Vector a
forall a b. (a -> b) -> a -> b
$ (Int -> a -> Maybe (Int, a)) -> Vector a -> Vector (Int, a)
forall a b. (Int -> a -> Maybe b) -> Vector a -> Vector b
V.imapMaybe Int -> a -> Maybe (Int, a)
shift Vector a
vec
 where
  shift :: Int -> a -> Maybe (Int, a)
shift Int
idx a
a
    | Int
idx Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
n = (Int, a) -> Maybe (Int, a)
forall a. a -> Maybe a
Just (Int
idx, a
a)
    | Int
idx Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
n = Maybe (Int, a)
forall a. Maybe a
Nothing
    | Bool
otherwise = (Int, a) -> Maybe (Int, a)
forall a. a -> Maybe a
Just (Int
idx Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1, a
a)

-- | Insert into a vector and shift all later elements right
insertAt :: Int -> a -> Vector a -> Vector a
insertAt :: forall a. Int -> a -> Vector a -> Vector a
insertAt Int
n a
v Vector a
vec
  | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 = Vector a
vec
  | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Vector a -> Int
forall a. Vector a -> Int
V.length Vector a
vec = Vector a
vec
  | Bool
otherwise = Vector (Int, a) -> Vector a
forall a. Vector (Int, a) -> Vector a
generateFrom (Vector (Int, a) -> Vector a) -> Vector (Int, a) -> Vector a
forall a b. (a -> b) -> a -> b
$ (Int -> a -> (Int, a)) -> Vector a -> Vector (Int, a)
forall a b. (Int -> a -> b) -> Vector a -> Vector b
V.imap Int -> a -> (Int, a)
shift Vector a
vec Vector (Int, a) -> Vector (Int, a) -> Vector (Int, a)
forall a. Semigroup a => a -> a -> a
<> (Int, a) -> Vector (Int, a)
forall a. a -> Vector a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int
n, a
v)
 where
  shift :: Int -> a -> (Int, a)
shift Int
idx a
a
    | Int
idx Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
n = (Int
idx Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1, a
a)
    | Bool
otherwise = (Int
idx, a
a)

generateFrom :: Vector (Int, a) -> Vector a
generateFrom :: forall a. Vector (Int, a) -> Vector a
generateFrom Vector (Int, a)
indexed = Int -> (Int -> a) -> Vector a
forall a. Int -> (Int -> a) -> Vector a
V.generate (Vector (Int, a) -> Int
forall a. Vector a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Vector (Int, a)
indexed) ((Int -> a) -> Vector a) -> (Int -> a) -> Vector a
forall a b. (a -> b) -> a -> b
$ \Int
idx ->
  a -> ((Int, a) -> a) -> Maybe (Int, a) -> a
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (Int -> a
forall {a} {a}. Show a => a -> a
badIndex Int
idx) (Int, a) -> a
forall a b. (a, b) -> b
snd (Maybe (Int, a) -> a) -> Maybe (Int, a) -> a
forall a b. (a -> b) -> a -> b
$ ((Int, a) -> Bool) -> Vector (Int, a) -> Maybe (Int, a)
forall a. (a -> Bool) -> Vector a -> Maybe a
V.find ((Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
idx) (Int -> Bool) -> ((Int, a) -> Int) -> (Int, a) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int, a) -> Int
forall a b. (a, b) -> a
fst) Vector (Int, a)
indexed
 where
  badIndex :: a -> a
badIndex a
idx =
    [Char] -> a
forall a. HasCallStack => [Char] -> a
error
      ([Char] -> a) -> [Char] -> a
forall a b. (a -> b) -> a -> b
$ [Char]
"Index "
        [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> a -> [Char]
forall a. Show a => a -> [Char]
show a
idx
        [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> [Char]
" is not present in vector of indexed values: "
        [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> [Int] -> [Char]
forall a. Show a => a -> [Char]
show [Int]
indexes

  indexes :: [Int]
  indexes :: [Int]
indexes = ((Int, a) -> Int) -> [(Int, a)] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map (Int, a) -> Int
forall a b. (a, b) -> a
fst ([(Int, a)] -> [Int]) -> [(Int, a)] -> [Int]
forall a b. (a -> b) -> a -> b
$ Vector (Int, a) -> [(Int, a)]
forall a. Vector a -> [a]
V.toList Vector (Int, a)
indexed