module Optics.Extra.Internal.Vector
  ( ordinalNub
  ) where

import Data.IntSet (IntSet)
import qualified Data.IntSet as IntSet

-- | Return the the subset of given ordinals within a given bound and in order
-- of the first occurrence seen.
--
-- Bound: @0 <= x < l@
--
-- >>> ordinalNub 3 [-1,2,1,4,2,3]
-- [2,1]
ordinalNub ::
  Int   {- ^ strict upper bound -} ->
  [Int] {- ^ ordinals           -} ->
  [Int] {- ^ unique, in-bound ordinals, in order seen -}
ordinalNub :: Int -> [Int] -> [Int]
ordinalNub Int
l [Int]
xs = (Int -> (IntSet -> [Int]) -> IntSet -> [Int])
-> (IntSet -> [Int]) -> [Int] -> IntSet -> [Int]
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (Int -> Int -> (IntSet -> [Int]) -> IntSet -> [Int]
ordinalNubHelper Int
l) ([Int] -> IntSet -> [Int]
forall a b. a -> b -> a
const []) [Int]
xs IntSet
IntSet.empty

ordinalNubHelper :: Int -> Int -> (IntSet -> [Int]) -> (IntSet -> [Int])
ordinalNubHelper :: Int -> Int -> (IntSet -> [Int]) -> IntSet -> [Int]
ordinalNubHelper Int
l Int
x IntSet -> [Int]
next IntSet
seen
  | Bool
outOfBounds Bool -> Bool -> Bool
|| Bool
notUnique = IntSet -> [Int]
next IntSet
seen
  | Bool
otherwise                = Int
x Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
: IntSet -> [Int]
next (Int -> IntSet -> IntSet
IntSet.insert Int
x IntSet
seen)
  where
  outOfBounds :: Bool
outOfBounds = Int
x Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| Int
l Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
x
  notUnique :: Bool
notUnique   = Int
x Int -> IntSet -> Bool
`IntSet.member` IntSet
seen