{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE DeriveAnyClass #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE ImportQualifiedPost #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE ViewPatterns #-}

module Data.OMap.Strict (
  HasOKey (okeyL),
  OMap (Empty, (:<|:), (:|>:)),
  null,
  size,
  empty,
  singleton,
  lookup,
  member,
  (!?),
  mapUnsafe,
  fromSet,
  fromFoldable,
  fromFoldableDuplicates,
  toMap,
  assocList,
  elems,
  toStrictSeq,
  toStrictSeqOKeys,
  toStrictSeqOfPairs,
  invariantHolds,
  invariantHolds',
  (|>),
  (<|),
  (<||),
  (||>),
  (|><),
  (><|),
  elem,
  extractKeys,
  adjust,
  filter,
  decodeOMap,
) where

import Cardano.Ledger.Binary (
  DecCBOR,
  Decoder,
  EncCBOR (encCBOR),
  decodeListLenOrIndef,
  decodeListLikeEnforceNoDuplicates,
  encodeStrictSeq,
 )
import Cardano.Ledger.Binary.Decoding (DecCBOR (decCBOR))
import Control.DeepSeq (NFData (..))
import Data.Aeson (ToJSON (..))
import Data.Default (Default (..))
import Data.Foldable qualified as F
import Data.Map.Strict qualified as Map
import Data.MapExtras qualified as MapE
import Data.Maybe (isJust)
import Data.Sequence.Strict qualified as SSeq
import Data.Set qualified as Set
import Data.Typeable (Typeable)
import GHC.Exts qualified as Exts
import GHC.Generics (Generic)
import Lens.Micro
import NoThunks.Class (NoThunks (..))
import Prelude hiding (elem, filter, lookup, null, seq)

-- | Class of types that can be mapped by a lens or a projection to an
-- Ord type.
--
-- For a type @V@, defines a lens from @V@ to and Ord type @K@.
class Ord k => HasOKey k v | v -> k where
  okeyL :: Lens' v k

-- | A general-purpose finite, insert-ordered, map that is strict in its
-- keys and values.
--
-- The strictness is enforced by the underlying strict `Map` that can
-- be looked-up by a projection or lens. and the ordering is maintained
-- by the constructing functions, leveraging `StrictSeq` to hold the
-- insert-order of the keys.
--
-- TODO: DecShareCBOR instance
data OMap k v = OMap
  { forall k v. OMap k v -> StrictSeq k
omSSeq :: !(SSeq.StrictSeq k)
  , forall k v. OMap k v -> Map k v
omMap :: !(Map.Map k v)
  }
  deriving ((forall x. OMap k v -> Rep (OMap k v) x)
-> (forall x. Rep (OMap k v) x -> OMap k v) -> Generic (OMap k v)
forall x. Rep (OMap k v) x -> OMap k v
forall x. OMap k v -> Rep (OMap k v) x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
forall k v x. Rep (OMap k v) x -> OMap k v
forall k v x. OMap k v -> Rep (OMap k v) x
$cfrom :: forall k v x. OMap k v -> Rep (OMap k v) x
from :: forall x. OMap k v -> Rep (OMap k v) x
$cto :: forall k v x. Rep (OMap k v) x -> OMap k v
to :: forall x. Rep (OMap k v) x -> OMap k v
Generic, OMap k v -> OMap k v -> Bool
(OMap k v -> OMap k v -> Bool)
-> (OMap k v -> OMap k v -> Bool) -> Eq (OMap k v)
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
forall k v. (Eq k, Eq v) => OMap k v -> OMap k v -> Bool
$c== :: forall k v. (Eq k, Eq v) => OMap k v -> OMap k v -> Bool
== :: OMap k v -> OMap k v -> Bool
$c/= :: forall k v. (Eq k, Eq v) => OMap k v -> OMap k v -> Bool
/= :: OMap k v -> OMap k v -> Bool
Eq)

instance (Show v, Ord k, Show k) => Show (OMap k v) where
  show :: OMap k v -> String
show = StrictSeq (k, v) -> String
forall a. Show a => a -> String
show (StrictSeq (k, v) -> String)
-> (OMap k v -> StrictSeq (k, v)) -> OMap k v -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. OMap k v -> StrictSeq (k, v)
forall k v. Ord k => OMap k v -> StrictSeq (k, v)
toStrictSeqOfPairs

deriving instance (NoThunks k, NoThunks v) => NoThunks (OMap k v)

deriving instance (NFData k, NFData v) => NFData (OMap k v)

-- | \(O(1)\).
empty :: OMap k v
empty :: forall k v. OMap k v
empty = StrictSeq k -> Map k v -> OMap k v
forall k v. StrictSeq k -> Map k v -> OMap k v
OMap StrictSeq k
forall a. StrictSeq a
SSeq.Empty Map k v
forall k a. Map k a
Map.empty

instance Default (OMap k v) where
  def :: OMap k v
def = OMap k v
forall k v. OMap k v
empty

-- | \(O(1)\). Shallow invariant using just `length` and `size`.
invariantHolds :: OMap k v -> Bool
invariantHolds :: forall k v. OMap k v -> Bool
invariantHolds (OMap StrictSeq k
sseq Map k v
kv) = StrictSeq k -> Int
forall a. StrictSeq a -> Int
SSeq.length StrictSeq k
sseq Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Map k v -> Int
forall k a. Map k a -> Int
Map.size Map k v
kv

-- | \(O(n \log n)\). Deep, costly invariant using membership check for each
-- value. By the pigeon-hole principle, this check is exhaustive.
invariantHolds' :: Ord k => OMap k v -> Bool
invariantHolds' :: forall k v. Ord k => OMap k v -> Bool
invariantHolds' omap :: OMap k v
omap@(OMap StrictSeq k
sseq Map k v
kv) =
  OMap k v -> Bool
forall k v. OMap k v -> Bool
invariantHolds OMap k v
omap Bool -> Bool -> Bool
&& (k -> Bool) -> StrictSeq k -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (\k
k -> Maybe v -> Bool
forall a. Maybe a -> Bool
isJust (Maybe v -> Bool) -> Maybe v -> Bool
forall a b. (a -> b) -> a -> b
$ k -> Map k v -> Maybe v
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup k
k Map k v
kv) StrictSeq k
sseq

-- | \(O(1)\).
null :: OMap k v -> Bool
null :: forall k v. OMap k v -> Bool
null (OMap StrictSeq k
sseq Map k v
_) = StrictSeq k -> Bool
forall a. StrictSeq a -> Bool
SSeq.null StrictSeq k
sseq

-- | \(O(1)\).
size :: OMap k v -> Int
size :: forall k v. OMap k v -> Int
size (OMap StrictSeq k
sseq Map k v
_) = StrictSeq k -> Int
forall a. StrictSeq a -> Int
SSeq.length StrictSeq k
sseq

-- | \(O(1)\). Strict in its arguments.
singleton :: HasOKey k v => v -> OMap k v
singleton :: forall k v. HasOKey k v => v -> OMap k v
singleton !v
v =
  let k :: k
k = v
v v -> Getting k v k -> k
forall s a. s -> Getting a s a -> a
^. Getting k v k
forall k v. HasOKey k v => Lens' v k
Lens' v k
okeyL
   in StrictSeq k -> Map k v -> OMap k v
forall k v. StrictSeq k -> Map k v -> OMap k v
OMap (k -> StrictSeq k
forall a. a -> StrictSeq a
SSeq.singleton k
k) (k -> v -> Map k v
forall k a. k -> a -> Map k a
Map.singleton k
k v
v)

-- | \(O(\log n)\). If the key is not present 'lookup' returns
-- 'Nothing'.
lookup :: Ord k => k -> OMap k v -> Maybe v
lookup :: forall k v. Ord k => k -> OMap k v -> Maybe v
lookup k
k (OMap StrictSeq k
_seq Map k v
kv) = k -> Map k v -> Maybe v
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup k
k Map k v
kv

-- | `flip`ed version of `lookup`
(!?) :: Ord k => OMap k v -> k -> Maybe v
!? :: forall k v. Ord k => OMap k v -> k -> Maybe v
(!?) = (k -> OMap k v -> Maybe v) -> OMap k v -> k -> Maybe v
forall a b c. (a -> b -> c) -> b -> a -> c
flip k -> OMap k v -> Maybe v
forall k v. Ord k => k -> OMap k v -> Maybe v
lookup

-- | \(O(\log n)\). Checks membership before cons'ing.
cons :: HasOKey k v => v -> OMap k v -> OMap k v
cons :: forall k v. HasOKey k v => v -> OMap k v -> OMap k v
cons v
v omap :: OMap k v
omap@(OMap StrictSeq k
sseq Map k v
kv)
  | k -> Map k v -> Bool
forall k a. Ord k => k -> Map k a -> Bool
Map.member k
k Map k v
kv = OMap k v
omap
  | Bool
otherwise = StrictSeq k -> Map k v -> OMap k v
forall k v. StrictSeq k -> Map k v -> OMap k v
OMap (k
k k -> StrictSeq k -> StrictSeq k
forall a. a -> StrictSeq a -> StrictSeq a
SSeq.<| StrictSeq k
sseq) (k -> v -> Map k v -> Map k v
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert k
k v
v Map k v
kv)
  where
    k :: k
k = v
v v -> Getting k v k -> k
forall s a. s -> Getting a s a -> a
^. Getting k v k
forall k v. HasOKey k v => Lens' v k
Lens' v k
okeyL

-- | \(O(\log n)\). Checks membership before cons'ing.
(<|) :: HasOKey k v => v -> OMap k v -> OMap k v
<| :: forall k v. HasOKey k v => v -> OMap k v -> OMap k v
(<|) = v -> OMap k v -> OMap k v
forall k v. HasOKey k v => v -> OMap k v -> OMap k v
cons

infixr 5 <|

-- | \(O(\log n)\). Checks membership before cons'ing. Overwrites a
-- duplicate.
cons' :: HasOKey k v => v -> OMap k v -> OMap k v
cons' :: forall k v. HasOKey k v => v -> OMap k v -> OMap k v
cons' v
v (OMap StrictSeq k
sseq Map k v
kv)
  | k -> Map k v -> Bool
forall k a. Ord k => k -> Map k a -> Bool
Map.member k
k Map k v
kv = StrictSeq k -> Map k v -> OMap k v
forall k v. StrictSeq k -> Map k v -> OMap k v
OMap StrictSeq k
sseq Map k v
kv'
  | Bool
otherwise = StrictSeq k -> Map k v -> OMap k v
forall k v. StrictSeq k -> Map k v -> OMap k v
OMap (k
k k -> StrictSeq k -> StrictSeq k
forall a. a -> StrictSeq a -> StrictSeq a
SSeq.<| StrictSeq k
sseq) Map k v
kv'
  where
    k :: k
k = v
v v -> Getting k v k -> k
forall s a. s -> Getting a s a -> a
^. Getting k v k
forall k v. HasOKey k v => Lens' v k
Lens' v k
okeyL
    kv' :: Map k v
kv' = k -> v -> Map k v -> Map k v
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert k
k v
v Map k v
kv

-- | \(O(\log n)\). Checks membership before cons'ing. Overwrites a
-- duplicate.
(<||) :: HasOKey k v => v -> OMap k v -> OMap k v
<|| :: forall k v. HasOKey k v => v -> OMap k v -> OMap k v
(<||) = v -> OMap k v -> OMap k v
forall k v. HasOKey k v => v -> OMap k v -> OMap k v
cons'

infixr 5 <||

-- TODO: export along with others that are hidden or remove them completely.

-- | \(O(\log n)\). Checks membership before snoc'ing.
snoc :: HasOKey k v => OMap k v -> v -> OMap k v
snoc :: forall k v. HasOKey k v => OMap k v -> v -> OMap k v
snoc omap :: OMap k v
omap@(OMap StrictSeq k
sseq Map k v
kv) v
v
  | k -> Map k v -> Bool
forall k a. Ord k => k -> Map k a -> Bool
Map.member k
k Map k v
kv = OMap k v
omap
  | Bool
otherwise = StrictSeq k -> Map k v -> OMap k v
forall k v. StrictSeq k -> Map k v -> OMap k v
OMap (StrictSeq k
sseq StrictSeq k -> k -> StrictSeq k
forall a. StrictSeq a -> a -> StrictSeq a
SSeq.|> k
k) (k -> v -> Map k v -> Map k v
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert k
k v
v Map k v
kv)
  where
    k :: k
k = v
v v -> Getting k v k -> k
forall s a. s -> Getting a s a -> a
^. Getting k v k
forall k v. HasOKey k v => Lens' v k
Lens' v k
okeyL

-- | \(O(\log n)\). Checks membership before snoc'ing.
(|>) :: HasOKey k v => OMap k v -> v -> OMap k v
|> :: forall k v. HasOKey k v => OMap k v -> v -> OMap k v
(|>) = OMap k v -> v -> OMap k v
forall k v. HasOKey k v => OMap k v -> v -> OMap k v
snoc

infixl 5 |>

-- | \(O(\log n)\). Checks membership before snoc'ing. Overwrites a
-- duplicate.
snoc' :: HasOKey k v => OMap k v -> v -> OMap k v
snoc' :: forall k v. HasOKey k v => OMap k v -> v -> OMap k v
snoc' (OMap StrictSeq k
sseq Map k v
kv) v
v
  | k -> Map k v -> Bool
forall k a. Ord k => k -> Map k a -> Bool
Map.member k
k Map k v
kv = StrictSeq k -> Map k v -> OMap k v
forall k v. StrictSeq k -> Map k v -> OMap k v
OMap StrictSeq k
sseq Map k v
kv'
  | Bool
otherwise = StrictSeq k -> Map k v -> OMap k v
forall k v. StrictSeq k -> Map k v -> OMap k v
OMap (StrictSeq k
sseq StrictSeq k -> k -> StrictSeq k
forall a. StrictSeq a -> a -> StrictSeq a
SSeq.|> k
k) Map k v
kv'
  where
    k :: k
k = v
v v -> Getting k v k -> k
forall s a. s -> Getting a s a -> a
^. Getting k v k
forall k v. HasOKey k v => Lens' v k
Lens' v k
okeyL
    kv' :: Map k v
kv' = k -> v -> Map k v -> Map k v
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert k
k v
v Map k v
kv

-- | \(O(\log n)\). Checks membership before snoc'ing. Overwrites a
-- duplicate.
(||>) :: HasOKey k v => OMap k v -> v -> OMap k v
||> :: forall k v. HasOKey k v => OMap k v -> v -> OMap k v
(||>) = OMap k v -> v -> OMap k v
forall k v. HasOKey k v => OMap k v -> v -> OMap k v
snoc'

infixl 5 ||>

-- | \(O(\log n)\).
uncons :: Ord k => OMap k v -> Maybe (v, OMap k v)
uncons :: forall k v. Ord k => OMap k v -> Maybe (v, OMap k v)
uncons (OMap StrictSeq k
sseq Map k v
kv) = case StrictSeq k
sseq of
  StrictSeq k
SSeq.Empty -> Maybe (v, OMap k v)
forall a. Maybe a
Nothing
  k
k SSeq.:<| StrictSeq k
ks ->
    case k -> Map k v -> Maybe v
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup k
k Map k v
kv of
      Just v
v -> (v, OMap k v) -> Maybe (v, OMap k v)
forall a. a -> Maybe a
Just (v
v, StrictSeq k -> Map k v -> OMap k v
forall k v. StrictSeq k -> Map k v -> OMap k v
OMap StrictSeq k
ks (k -> Map k v -> Map k v
forall k a. Ord k => k -> Map k a -> Map k a
Map.delete k
k Map k v
kv))
      Maybe v
Nothing -> String -> Maybe (v, OMap k v)
forall a. HasCallStack => String -> a
error String
"Invariant falsified! In OMap, key from sequence not found in corresponding map"

-- | \(O(\log n)\).
unsnoc :: Ord k => OMap k v -> Maybe (OMap k v, v)
unsnoc :: forall k v. Ord k => OMap k v -> Maybe (OMap k v, v)
unsnoc (OMap StrictSeq k
sseq Map k v
kv) = case StrictSeq k
sseq of
  StrictSeq k
SSeq.Empty -> Maybe (OMap k v, v)
forall a. Maybe a
Nothing
  StrictSeq k
ks SSeq.:|> k
k ->
    case k -> Map k v -> Maybe v
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup k
k Map k v
kv of
      Just v
v -> (OMap k v, v) -> Maybe (OMap k v, v)
forall a. a -> Maybe a
Just (StrictSeq k -> Map k v -> OMap k v
forall k v. StrictSeq k -> Map k v -> OMap k v
OMap StrictSeq k
ks (k -> Map k v -> Map k v
forall k a. Ord k => k -> Map k a -> Map k a
Map.delete k
k Map k v
kv), v
v)
      Maybe v
Nothing -> String -> Maybe (OMap k v, v)
forall a. HasCallStack => String -> a
error String
"Invariant falsified! In OMap, key from sequence not found in corresponding map"

-- | \(O(n \log n)\). Checks membership before snoc'ing.
-- De-duplicates the StrictSeq without overwriting.
-- Starts from the left or head, using `foldl'`
fromFoldable :: (Foldable f, HasOKey k v) => f v -> OMap k v
fromFoldable :: forall (f :: * -> *) k v.
(Foldable f, HasOKey k v) =>
f v -> OMap k v
fromFoldable = (OMap k v -> v -> OMap k v) -> OMap k v -> f v -> OMap k v
forall b a. (b -> a -> b) -> b -> f a -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
F.foldl' OMap k v -> v -> OMap k v
forall k v. HasOKey k v => OMap k v -> v -> OMap k v
snoc OMap k v
forall k v. OMap k v
empty

-- | \(O(n \log n)\). Checks membership before snoc'ing.
-- De-duplicates the StrictSeq and collects and returns the duplicates found.
-- Starts from the left or head, using `foldl'`
fromFoldableDuplicates :: (Foldable f, HasOKey k v, Ord v) => f v -> (Set.Set v, OMap k v)
fromFoldableDuplicates :: forall (f :: * -> *) k v.
(Foldable f, HasOKey k v, Ord v) =>
f v -> (Set v, OMap k v)
fromFoldableDuplicates = ((Set v, OMap k v) -> v -> (Set v, OMap k v))
-> (Set v, OMap k v) -> f v -> (Set v, OMap k v)
forall b a. (b -> a -> b) -> b -> f a -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
F.foldl' (Set v, OMap k v) -> v -> (Set v, OMap k v)
forall k v.
(HasOKey k v, Ord v) =>
(Set v, OMap k v) -> v -> (Set v, OMap k v)
snoc_ (Set v
forall a. Set a
Set.empty, OMap k v
forall k v. OMap k v
empty)
  where
    snoc_ :: (HasOKey k v, Ord v) => (Set.Set v, OMap k v) -> v -> (Set.Set v, OMap k v)
    snoc_ :: forall k v.
(HasOKey k v, Ord v) =>
(Set v, OMap k v) -> v -> (Set v, OMap k v)
snoc_ (Set v
duplicates, omap :: OMap k v
omap@(OMap StrictSeq k
sseq Map k v
kv)) v
v =
      let k :: k
k = v
v v -> Getting k v k -> k
forall s a. s -> Getting a s a -> a
^. Getting k v k
forall k v. HasOKey k v => Lens' v k
Lens' v k
okeyL
       in if k -> Map k v -> Bool
forall k a. Ord k => k -> Map k a -> Bool
Map.member k
k Map k v
kv
            then (v -> Set v -> Set v
forall a. Ord a => a -> Set a -> Set a
Set.insert v
v Set v
duplicates, OMap k v
omap)
            else (Set v
duplicates, StrictSeq k -> Map k v -> OMap k v
forall k v. StrictSeq k -> Map k v -> OMap k v
OMap (StrictSeq k
sseq StrictSeq k -> k -> StrictSeq k
forall a. StrictSeq a -> a -> StrictSeq a
SSeq.|> k
k) (k -> v -> Map k v -> Map k v
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert k
k v
v Map k v
kv))

-- | \(O(n \log n)\).
fromSet :: HasOKey k v => Set.Set v -> OMap k v
fromSet :: forall k v. HasOKey k v => Set v -> OMap k v
fromSet = Set v -> OMap k v
forall (f :: * -> *) k v.
(Foldable f, HasOKey k v) =>
f v -> OMap k v
fromFoldable

-- | \(O(1)\).
toMap :: OMap k v -> Map.Map k v
toMap :: forall k v. OMap k v -> Map k v
toMap = OMap k v -> Map k v
forall k v. OMap k v -> Map k v
omMap

-- | \(O(n \log n)\).
toStrictSeq :: Ord k => OMap k v -> SSeq.StrictSeq v
toStrictSeq :: forall k v. Ord k => OMap k v -> StrictSeq v
toStrictSeq (OMap StrictSeq k
sseq Map k v
kv) = StrictSeq k
sseq StrictSeq k -> (k -> v) -> StrictSeq v
forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> \k
k -> let !v :: v
v = Map k v
kv Map k v -> k -> v
forall k a. Ord k => Map k a -> k -> a
Map.! k
k in v
v

-- | \(O(1)\).
toStrictSeqOKeys :: OMap k v -> SSeq.StrictSeq k
toStrictSeqOKeys :: forall k v. OMap k v -> StrictSeq k
toStrictSeqOKeys = OMap k v -> StrictSeq k
forall k v. OMap k v -> StrictSeq k
omSSeq

-- | \(O(n \log n)\).
toStrictSeqOfPairs :: Ord k => OMap k v -> SSeq.StrictSeq (k, v)
toStrictSeqOfPairs :: forall k v. Ord k => OMap k v -> StrictSeq (k, v)
toStrictSeqOfPairs (OMap StrictSeq k
sseq Map k v
kv) = StrictSeq k
sseq StrictSeq k -> (k -> (k, v)) -> StrictSeq (k, v)
forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> \k
k -> let !v :: v
v = Map k v
kv Map k v -> k -> v
forall k a. Ord k => Map k a -> k -> a
Map.! k
k in (k
k, v
v)

-- | \(O(\log n)\). Key membership check.
member :: Ord k => k -> OMap k v -> Bool
member :: forall k v. Ord k => k -> OMap k v -> Bool
member k
k (OMap StrictSeq k
_sseq Map k v
kv) = k -> Map k v -> Bool
forall k a. Ord k => k -> Map k a -> Bool
Map.member k
k Map k v
kv

-- | \(O(\log n)\). Value membership check.
elem :: (HasOKey k v, Eq v) => v -> OMap k v -> Bool
elem :: forall k v. (HasOKey k v, Eq v) => v -> OMap k v -> Bool
elem v
v = (v -> Maybe v
forall a. a -> Maybe a
Just v
v Maybe v -> Maybe v -> Bool
forall a. Eq a => a -> a -> Bool
==) (Maybe v -> Bool) -> (OMap k v -> Maybe v) -> OMap k v -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. k -> OMap k v -> Maybe v
forall k v. Ord k => k -> OMap k v -> Maybe v
lookup (v
v v -> Getting k v k -> k
forall s a. s -> Getting a s a -> a
^. Getting k v k
forall k v. HasOKey k v => Lens' v k
Lens' v k
okeyL)

-- | \(O(n)\). Given a `Set` of @k@s, and an `OMap` @k@ @v@ return
-- a pair of `Map` and `OMap` where the @k@s in the `Set` have been
-- removed from the `OMap` and presented as a separate `Map`.
extractKeys :: Ord k => Set.Set k -> OMap k v -> (OMap k v, Map.Map k v)
extractKeys :: forall k v. Ord k => Set k -> OMap k v -> (OMap k v, Map k v)
extractKeys Set k
ks (OMap StrictSeq k
sseq Map k v
kv) =
  let (Map k v
kv', Map k v
extractedKv) = Map k v -> Set k -> (Map k v, Map k v)
forall k a. Ord k => Map k a -> Set k -> (Map k a, Map k a)
MapE.extractKeys Map k v
kv Set k
ks
      sseq' :: StrictSeq k
sseq' =
        (StrictSeq k -> k -> StrictSeq k)
-> StrictSeq k -> StrictSeq k -> StrictSeq k
forall b a. (b -> a -> b) -> b -> StrictSeq a -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
F.foldl'
          (\StrictSeq k
accum k
k -> if k -> Set k -> Bool
forall a. Ord a => a -> Set a -> Bool
Set.member k
k Set k
ks then StrictSeq k
accum else StrictSeq k
accum StrictSeq k -> k -> StrictSeq k
forall a. StrictSeq a -> a -> StrictSeq a
SSeq.|> k
k)
          StrictSeq k
forall a. StrictSeq a
SSeq.empty
          StrictSeq k
sseq
   in (StrictSeq k -> Map k v -> OMap k v
forall k v. StrictSeq k -> Map k v -> OMap k v
OMap StrictSeq k
sseq' Map k v
kv', Map k v
extractedKv)

-- | \(O(n)\). Like `Map.adjust`.
--
-- Returns the original `OMap` unaltered when the key does not exist.
--
-- If the key exists, then the function is applied to the value, but we need to consider
-- three possible cases:
--
--     1. The modified value's `okeyL` is unaltered
--            - we return omap with the adjusted value,
--     2. The modified value's `okeyL` is altered, but not a duplicate
--            - we return the omap with adjusted key (in place) and value
--     3. The modified value's `okeyL` is altered and is a duplicate
--            - we return the omap with the old key deleted from the sequence but
--              without inserting the new key since it is a duplicate, and
--              deleting old value and inserting the new value in place of its duplicate.
--
-- Examples:
--
-- >>> import Data.OMap.Strict
-- >>> import Lens.Micro
-- >>> instance HasOKey Int (Int, Char) where okeyL = _1
-- >>> let m = fromFoldable $ zip [1,2] ['a','b'] :: OMap Int (Int, Char)
-- >>> m
-- StrictSeq {fromStrict = fromList [(1,(1,'a')),(2,(2,'b'))]}
-- >>> let adjustingFn (k, v) = (k, succ v) -- Changes the value
-- >>> let overwritingAdjustingFn (k,v) = (succ k, v) -- Changes the `okeyL`.
-- >>> adjust adjustingFn 1 m
-- StrictSeq {fromStrict = fromList [(1,(1,'b')),(2,(2,'b'))]}
-- >>> adjust overwritingAdjustingFn  1 m
-- StrictSeq {fromStrict = fromList [(2,(2,'a'))]}
adjust :: HasOKey k v => (v -> v) -> k -> OMap k v -> OMap k v
adjust :: forall k v. HasOKey k v => (v -> v) -> k -> OMap k v -> OMap k v
adjust v -> v
f k
k omap :: OMap k v
omap@(OMap StrictSeq k
sseq Map k v
kv) =
  case k -> Map k v -> Maybe v
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup k
k Map k v
kv of
    Maybe v
Nothing -> OMap k v
omap
    Just v
v ->
      let v' :: v
v' = v -> v
f v
v
          k' :: k
k' = v
v' v -> Getting k v k -> k
forall s a. s -> Getting a s a -> a
^. Getting k v k
forall k v. HasOKey k v => Lens' v k
Lens' v k
okeyL
       in if k
k' k -> k -> Bool
forall a. Eq a => a -> a -> Bool
== k
k
            then StrictSeq k -> Map k v -> OMap k v
forall k v. StrictSeq k -> Map k v -> OMap k v
OMap StrictSeq k
sseq (k -> v -> Map k v -> Map k v
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert k
k v
v' Map k v
kv)
            else
              let kv' :: Map k v
kv' = k -> v -> Map k v -> Map k v
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert k
k' v
v' (Map k v -> Map k v) -> Map k v -> Map k v
forall a b. (a -> b) -> a -> b
$ k -> Map k v -> Map k v
forall k a. Ord k => k -> Map k a -> Map k a
Map.delete k
k Map k v
kv
                  (StrictSeq k
lseq, StrictSeq k
rseq) = case (k -> Bool) -> StrictSeq k -> (StrictSeq k, StrictSeq k)
forall a. (a -> Bool) -> StrictSeq a -> (StrictSeq a, StrictSeq a)
SSeq.spanl (k -> k -> Bool
forall a. Eq a => a -> a -> Bool
/= k
k) StrictSeq k
sseq of
                    (StrictSeq k
l, k
_ SSeq.:<| StrictSeq k
r) -> (StrictSeq k
l, StrictSeq k
r)
                    (StrictSeq k, StrictSeq k)
_ -> String -> (StrictSeq k, StrictSeq k)
forall a. HasCallStack => String -> a
error String
"Impossible: supplied key expected to be in the sequence"
               in case k -> Map k v -> Maybe v
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup k
k' Map k v
kv of
                    Maybe v
Nothing -> StrictSeq k -> Map k v -> OMap k v
forall k v. StrictSeq k -> Map k v -> OMap k v
OMap (StrictSeq k
lseq StrictSeq k -> StrictSeq k -> StrictSeq k
forall a. Semigroup a => a -> a -> a
<> (k
k' k -> StrictSeq k -> StrictSeq k
forall a. a -> StrictSeq a -> StrictSeq a
SSeq.:<| StrictSeq k
rseq)) Map k v
kv'
                    Just v
_ -> StrictSeq k -> Map k v -> OMap k v
forall k v. StrictSeq k -> Map k v -> OMap k v
OMap (StrictSeq k
lseq StrictSeq k -> StrictSeq k -> StrictSeq k
forall a. Semigroup a => a -> a -> a
<> StrictSeq k
rseq) Map k v
kv'

-- | This mapping function is only safe when the key stored in the new value matches the
-- key stored in the new value. This invariant is not checked for performance reasons
mapUnsafe :: (v1 -> v2) -> OMap k v1 -> OMap k v2
mapUnsafe :: forall v1 v2 k. (v1 -> v2) -> OMap k v1 -> OMap k v2
mapUnsafe v1 -> v2
f (OMap StrictSeq k
sseq Map k v1
kv) = StrictSeq k -> Map k v2 -> OMap k v2
forall k v. StrictSeq k -> Map k v -> OMap k v
OMap StrictSeq k
sseq ((v1 -> v2) -> Map k v1 -> Map k v2
forall a b k. (a -> b) -> Map k a -> Map k b
Map.map v1 -> v2
f Map k v1
kv)

-- | \(O(1)\)
pattern Empty :: OMap k v
pattern $mEmpty :: forall {r} {k} {v}. OMap k v -> ((# #) -> r) -> ((# #) -> r) -> r
$bEmpty :: forall k v. OMap k v
Empty <- (null -> True)
  where
    Empty = OMap k v
forall k v. OMap k v
empty

-- | \(O(\log n)\).
pattern (:<|:) :: (HasOKey k v, Ord k) => v -> OMap k v -> OMap k v
pattern x $m:<|: :: forall {r} {k} {v}.
(HasOKey k v, Ord k) =>
OMap k v -> (v -> OMap k v -> r) -> ((# #) -> r) -> r
$b:<|: :: forall k v. (HasOKey k v, Ord k) => v -> OMap k v -> OMap k v
:<|: xs <- (uncons -> Just (x, xs))
  where
    v
x :<|: OMap k v
xs = v
x v -> OMap k v -> OMap k v
forall k v. HasOKey k v => v -> OMap k v -> OMap k v
<| OMap k v
xs

infixr 5 :<|:

-- | \(O(\log n)\).
pattern (:|>:) :: (HasOKey k v, Ord k) => OMap k v -> v -> OMap k v
pattern xs $m:|>: :: forall {r} {k} {v}.
(HasOKey k v, Ord k) =>
OMap k v -> (OMap k v -> v -> r) -> ((# #) -> r) -> r
$b:|>: :: forall k v. (HasOKey k v, Ord k) => OMap k v -> v -> OMap k v
:|>: x <- (unsnoc -> Just (xs, x))
  where
    OMap k v
xs :|>: v
x = OMap k v
xs OMap k v -> v -> OMap k v
forall k v. HasOKey k v => OMap k v -> v -> OMap k v
|> v
x

infixl 5 :|>:

{-# COMPLETE Empty, (:|>:) #-}

{-# COMPLETE Empty, (:<|:) #-}

-- | \( O(n \log m) \). For every uncons-ed element from the sequence on the right,
-- check its membership in the sequence on the left, before snoc'ing it.
-- Preserve order. Remove duplicates from sequence on the right.
(|><) :: HasOKey k v => OMap k v -> OMap k v -> OMap k v
OMap k v
omapl |>< :: forall k v. HasOKey k v => OMap k v -> OMap k v -> OMap k v
|>< OMap k v
omapr = case OMap k v
omapr of
  OMap k v
Empty -> OMap k v
omapl
  v
r :<|: OMap k v
rs -> (OMap k v
omapl OMap k v -> v -> OMap k v
forall k v. HasOKey k v => OMap k v -> v -> OMap k v
|> v
r) OMap k v -> OMap k v -> OMap k v
forall k v. HasOKey k v => OMap k v -> OMap k v -> OMap k v
|>< OMap k v
rs

infixl 5 |><

-- | \( O(m \log n) \). For every unsnoc-ed element from the sequence on the left,
-- check its membership in the sequence on the right, before cons'ing it.
-- Preserve order. Remove duplicates from sequence on the left.
(><|) :: HasOKey k v => OMap k v -> OMap k v -> OMap k v
OMap k v
omapl ><| :: forall k v. HasOKey k v => OMap k v -> OMap k v -> OMap k v
><| OMap k v
omapr = case OMap k v
omapl of
  OMap k v
Empty -> OMap k v
omapr
  OMap k v
ls :|>: v
l -> OMap k v
ls OMap k v -> OMap k v -> OMap k v
forall k v. HasOKey k v => OMap k v -> OMap k v -> OMap k v
><| (v
l v -> OMap k v -> OMap k v
forall k v. HasOKey k v => v -> OMap k v -> OMap k v
<| OMap k v
omapr)

infixr 5 ><|

instance HasOKey k v => Exts.IsList (OMap k v) where
  type Item (OMap k v) = v
  fromList :: [Item (OMap k v)] -> OMap k v
fromList = [v] -> OMap k v
[Item (OMap k v)] -> OMap k v
forall (f :: * -> *) k v.
(Foldable f, HasOKey k v) =>
f v -> OMap k v
fromFoldable
  toList :: OMap k v -> [Item (OMap k v)]
toList = OMap k v -> [v]
OMap k v -> [Item (OMap k v)]
forall a. OMap k a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
F.toList

assocList :: Ord k => OMap k v -> [(k, v)]
assocList :: forall k v. Ord k => OMap k v -> [(k, v)]
assocList = StrictSeq (k, v) -> [(k, v)]
forall a. StrictSeq a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
F.toList (StrictSeq (k, v) -> [(k, v)])
-> (OMap k v -> StrictSeq (k, v)) -> OMap k v -> [(k, v)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. OMap k v -> StrictSeq (k, v)
forall k v. Ord k => OMap k v -> StrictSeq (k, v)
toStrictSeqOfPairs

elems :: Ord k => OMap k v -> [v]
elems :: forall k v. Ord k => OMap k v -> [v]
elems = StrictSeq v -> [v]
forall a. StrictSeq a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
F.toList (StrictSeq v -> [v])
-> (OMap k v -> StrictSeq v) -> OMap k v -> [v]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. OMap k v -> StrictSeq v
forall k v. Ord k => OMap k v -> StrictSeq v
toStrictSeq

instance (HasOKey k v, ToJSON v) => ToJSON (OMap k v) where
  toJSON :: OMap k v -> Value
toJSON = StrictSeq v -> Value
forall a. ToJSON a => a -> Value
toJSON (StrictSeq v -> Value)
-> (OMap k v -> StrictSeq v) -> OMap k v -> Value
forall b c a. (b -> c) -> (a -> b) -> a -> c
. OMap k v -> StrictSeq v
forall k v. Ord k => OMap k v -> StrictSeq v
toStrictSeq
  toEncoding :: OMap k v -> Encoding
toEncoding = StrictSeq v -> Encoding
forall a. ToJSON a => a -> Encoding
toEncoding (StrictSeq v -> Encoding)
-> (OMap k v -> StrictSeq v) -> OMap k v -> Encoding
forall b c a. (b -> c) -> (a -> b) -> a -> c
. OMap k v -> StrictSeq v
forall k v. Ord k => OMap k v -> StrictSeq v
toStrictSeq

instance HasOKey k v => Semigroup (OMap k v) where
  <> :: OMap k v -> OMap k v -> OMap k v
(<>) = OMap k v -> OMap k v -> OMap k v
forall k v. HasOKey k v => OMap k v -> OMap k v -> OMap k v
(|><)

instance HasOKey k v => Monoid (OMap k v) where
  mempty :: OMap k v
mempty = OMap k v
forall k v. OMap k v
empty

instance Ord k => Foldable (OMap k) where
  foldMap :: forall m a. Monoid m => (a -> m) -> OMap k a -> m
foldMap a -> m
f (OMap StrictSeq k
sseq Map k a
kv) = (k -> m) -> StrictSeq k -> m
forall m a. Monoid m => (a -> m) -> StrictSeq a -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
F.foldMap (\k
k -> a -> m
f (Map k a
kv Map k a -> k -> a
forall k a. Ord k => Map k a -> k -> a
Map.! k
k)) StrictSeq k
sseq
  {-# INLINEABLE foldMap #-}
  foldr :: forall a b. (a -> b -> b) -> b -> OMap k a -> b
foldr a -> b -> b
f b
z (OMap StrictSeq k
sseq Map k a
kv) = (k -> b -> b) -> b -> StrictSeq k -> b
forall a b. (a -> b -> b) -> b -> StrictSeq a -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
F.foldr (\k
k -> a -> b -> b
f (Map k a
kv Map k a -> k -> a
forall k a. Ord k => Map k a -> k -> a
Map.! k
k)) b
z StrictSeq k
sseq
  {-# INLINEABLE foldr #-}
  foldl :: forall b a. (b -> a -> b) -> b -> OMap k a -> b
foldl b -> a -> b
f b
z (OMap StrictSeq k
sseq Map k a
kv) = (b -> k -> b) -> b -> StrictSeq k -> b
forall b a. (b -> a -> b) -> b -> StrictSeq a -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
F.foldl (\b
acc k
k -> b -> a -> b
f b
acc (Map k a
kv Map k a -> k -> a
forall k a. Ord k => Map k a -> k -> a
Map.! k
k)) b
z StrictSeq k
sseq
  {-# INLINEABLE foldl #-}
  foldr' :: forall a b. (a -> b -> b) -> b -> OMap k a -> b
foldr' a -> b -> b
f b
z (OMap StrictSeq k
sseq Map k a
kv) = (k -> b -> b) -> b -> StrictSeq k -> b
forall a b. (a -> b -> b) -> b -> StrictSeq a -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
F.foldr' (\k
k -> a -> b -> b
f (Map k a
kv Map k a -> k -> a
forall k a. Ord k => Map k a -> k -> a
Map.! k
k)) b
z StrictSeq k
sseq
  {-# INLINEABLE foldr' #-}
  foldl' :: forall b a. (b -> a -> b) -> b -> OMap k a -> b
foldl' b -> a -> b
f b
z (OMap StrictSeq k
sseq Map k a
kv) = (b -> k -> b) -> b -> StrictSeq k -> b
forall b a. (b -> a -> b) -> b -> StrictSeq a -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
F.foldl' (\b
acc k
k -> b -> a -> b
f b
acc (Map k a
kv Map k a -> k -> a
forall k a. Ord k => Map k a -> k -> a
Map.! k
k)) b
z StrictSeq k
sseq
  {-# INLINEABLE foldl' #-}
  length :: forall a. OMap k a -> Int
length = Map k a -> Int
forall k a. Map k a -> Int
Map.size (Map k a -> Int) -> (OMap k a -> Map k a) -> OMap k a -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. OMap k a -> Map k a
forall k v. OMap k v -> Map k v
omMap
  {-# INLINE length #-}
  null :: forall a. OMap k a -> Bool
null = Map k a -> Bool
forall k a. Map k a -> Bool
Map.null (Map k a -> Bool) -> (OMap k a -> Map k a) -> OMap k a -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. OMap k a -> Map k a
forall k v. OMap k v -> Map k v
omMap
  {-# INLINE null #-}

instance (Typeable k, EncCBOR v, Ord k) => EncCBOR (OMap k v) where
  encCBOR :: OMap k v -> Encoding
encCBOR OMap k v
omap = (v -> Encoding) -> StrictSeq v -> Encoding
forall a. (a -> Encoding) -> StrictSeq a -> Encoding
encodeStrictSeq v -> Encoding
forall a. EncCBOR a => a -> Encoding
encCBOR (OMap k v -> StrictSeq v
forall k v. Ord k => OMap k v -> StrictSeq v
toStrictSeq OMap k v
omap)

instance (Typeable k, HasOKey k v, DecCBOR v, Eq v) => DecCBOR (OMap k v) where
  decCBOR :: forall s. Decoder s (OMap k v)
decCBOR = Decoder s v -> Decoder s (OMap k v)
forall k v s. HasOKey k v => Decoder s v -> Decoder s (OMap k v)
decodeOMap Decoder s v
forall s. Decoder s v
forall a s. DecCBOR a => Decoder s a
decCBOR

decodeOMap :: HasOKey k v => Decoder s v -> Decoder s (OMap k v)
decodeOMap :: forall k v s. HasOKey k v => Decoder s v -> Decoder s (OMap k v)
decodeOMap Decoder s v
decValue =
  Decoder s (Maybe Int)
-> (v -> OMap k v -> OMap k v)
-> (OMap k v -> (Int, OMap k v))
-> Decoder s v
-> Decoder s (OMap k v)
forall s a b c.
Monoid b =>
Decoder s (Maybe Int)
-> (a -> b -> b) -> (b -> (Int, c)) -> Decoder s a -> Decoder s c
decodeListLikeEnforceNoDuplicates
    Decoder s (Maybe Int)
forall s. Decoder s (Maybe Int)
decodeListLenOrIndef
    ((OMap k v -> v -> OMap k v) -> v -> OMap k v -> OMap k v
forall a b c. (a -> b -> c) -> b -> a -> c
flip OMap k v -> v -> OMap k v
forall k v. HasOKey k v => OMap k v -> v -> OMap k v
snoc)
    (\OMap k v
omap -> (OMap k v -> Int
forall k v. OMap k v -> Int
size OMap k v
omap, OMap k v
omap))
    Decoder s v
decValue

-- | \( O(n \log n) \)
filter :: Ord k => (v -> Bool) -> OMap k v -> OMap k v
filter :: forall k v. Ord k => (v -> Bool) -> OMap k v -> OMap k v
filter v -> Bool
f (OMap StrictSeq k
sseq Map k v
kv) =
  let kv' :: Map k v
kv' = (v -> Bool) -> Map k v -> Map k v
forall a k. (a -> Bool) -> Map k a -> Map k a
Map.filter v -> Bool
f Map k v
kv
      sseq' :: StrictSeq k
sseq' =
        (StrictSeq k -> k -> StrictSeq k)
-> StrictSeq k -> StrictSeq k -> StrictSeq k
forall b a. (b -> a -> b) -> b -> StrictSeq a -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
F.foldl' (\StrictSeq k
accum k
k -> if k -> Map k v -> Bool
forall k a. Ord k => k -> Map k a -> Bool
Map.member k
k Map k v
kv' then StrictSeq k
accum StrictSeq k -> k -> StrictSeq k
forall a. StrictSeq a -> a -> StrictSeq a
SSeq.:|> k
k else StrictSeq k
accum) StrictSeq k
forall a. StrictSeq a
SSeq.empty StrictSeq k
sseq
   in StrictSeq k -> Map k v -> OMap k v
forall k v. StrictSeq k -> Map k v -> OMap k v
OMap StrictSeq k
sseq' Map k v
kv'